Use GPUToolbox's Float64-free math overrides on devices without fp64

Project.toml

@@ -9,6 +9,7 @@ projects = ["lib/intrinsics", "test", "docs", "res"]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
+GPUToolbox = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -32,6 +33,7 @@ SPIRVIntrinsics = {path = "lib/intrinsics"}
 Adapt = "4"
 GPUArrays = "11.2.1"
 GPUCompiler = "2.7"
+GPUToolbox = "3.1"
 KernelAbstractions = "0.9.38"
 LLVM = "9.6"
 LinearAlgebra = "1"

src/OpenCL.jl

@@ -1,6 +1,7 @@
 module OpenCL
 
 using GPUCompiler
+import GPUToolbox
 using LLVM, LLVM.Interop
 using SPIRV_LLVM_Backend_jll, SPIRV_Tools_jll, spirv2clc_jll
 using Adapt

src/compiler/compilation.jl

@@ -40,8 +40,16 @@ end
 
 GPUCompiler.runtime_module(::CompilerJob{<:Any,OpenCLCompilerParams}) = OpenCL
 
-GPUCompiler.method_table_view(job::OpenCLCompilerJob) =
-    GPUCompiler.StackedMethodTable(job.world, method_table, SPIRVIntrinsics.method_table)
+function GPUCompiler.method_table_view(job::OpenCLCompilerJob)
+    if job.config.target.supports_fp64
+        parent = SPIRVIntrinsics.method_table
+    else
+        # keep single-precision math that Base computes in Float64 out of double precision
+        parent = GPUCompiler.StackedMethodTable(job.world, SPIRVIntrinsics.method_table,
+                                                GPUToolbox.Overlays.float64_overrides)
+    end
+    GPUCompiler.StackedMethodTable(job.world, method_table, parent)
+end
 
 # filter out OpenCL built-ins
 # TODO: eagerly lower these using the translator API

test/intrinsics.jl

@@ -1,4 +1,5 @@
 using SIMD
+import GPUCompiler
 
 function call_on_device(f, args...)
     function kernel(res, f, args...)
@@ -330,3 +331,21 @@ end # if cl.sub_groups_supported(cl.device())
 end
 
 end
+
+@testset "devices without Float64" begin
+    # Base computes some single-precision math in double precision; on devices that lack
+    # it, GPUToolbox's overrides are used instead. Compile for such a device by disabling
+    # Float64 support (which makes the compiler reject any use of it).
+    config = OpenCL.compiler_config(cl.device())
+    target = GPUCompiler.SPIRVCompilerTarget(;
+        (field => getfield(config.target, field) for field in fieldnames(GPUCompiler.SPIRVCompilerTarget))...,
+        supports_fp64=false)
+    config = GPUCompiler.CompilerConfig(config; target)
+    kernel(out, f, args...) = (@inbounds out[] = f(args...); return)
+    for (f, T, args) in ((div, Float32, (Float32, Float32)), (sind, Float32, (Float32,)),
+                         (inv, ComplexF32, (ComplexF32,)))
+        tt = Tuple{CLDeviceArray{T,0,AS.CrossWorkgroup}, typeof(f), args...}
+        job = GPUCompiler.CompilerJob(GPUCompiler.methodinstance(typeof(kernel), tt), config)
+        @test GPUCompiler.JuliaContext(_ -> GPUCompiler.compile(:llvm, job)) !== nothing
+    end
+end