Lower host references through CUDA loader slots

Author: maleadt

CUDACore/src/compiler/compilation.jl

@@ -18,7 +18,10 @@ const CUDACompilerConfig = CompilerConfig{PTXCompilerTarget, CUDACompilerParams}
 const AnyCUDAJob = CompilerJob{PTXCompilerTarget, <:AbstractCUDACompilerParams}
 
 GPUCompiler.runtime_module(@nospecialize(job::AnyCUDAJob)) = CUDACore
-GPUCompiler.supports_host_reference_patching(@nospecialize(job::AnyCUDAJob)) = true
+function GPUCompiler.lower_host_references!(@nospecialize(job::AnyCUDAJob), mod::LLVM.Module,
+                                            refs::GPUCompiler.HostReferences)
+    GPUCompiler.emit_host_reference_slots!(mod, refs)
+end
 
 # filter out functions from libdevice and cudadevrt
 GPUCompiler.isintrinsic(@nospecialize(job::AnyCUDAJob), fn::String) =
@@ -525,8 +528,9 @@ function link_kernel(@nospecialize(job::CompilerJob), image::Vector{UInt8}, entr
     mod = CuModule(image)
     roots = Any[]
     for (name, ref) in refs.slots
-        CuGlobal{UInt}(mod, name)[] = GPUCompiler.resolve_host_reference(ref)
-        ref isa GPUCompiler.JuliaObjectRef && push!(roots, ref.value)
+        slot = CuGlobal{UInt}(mod, name)
+        slot[] = GPUCompiler.resolve_host_reference(ref)
+        ref isa GPUCompiler.JuliaValueRef && push!(roots, ref.value)
     end
     return CuFunction(mod, entry), roots
 end

CUDACore/src/precompile.jl

@@ -45,7 +45,8 @@ end
 # kernel launch infrastructure
 let CUDACompilerJob = CompilerJob{PTXCompilerTarget, CUDACompilerParams}
     precompile(Tuple{typeof(cufunction), typeof(identity), Type{Tuple{Nothing}}})
-    precompile(Tuple{typeof(link_kernel), CUDACompilerJob, Vector{UInt8}, String})
+    precompile(Tuple{typeof(link_kernel), CUDACompilerJob, Vector{UInt8}, String,
+                     GPUCompiler.HostReferences})
 
     # GPUCompiler 2.0 caching pipeline (specialized for CUDACore's results struct)
     precompile(Tuple{typeof(compile_or_lookup), CUDACompilerJob})

test/core/codegen.jl

@@ -256,15 +256,26 @@ end
 ############################################################################################
 
 @testset "host reference patching" begin
-    @eval begin
+    runtime_refs = @eval module $(gensym())
+        using ..CUDACore
         const host_reference_type_tag = CUDACore.GPUCompiler.Runtime.type_tag
         host_reference_kernel(out) = (out[1] = host_reference_type_tag(Val(:float32)); return)
     end
 
     out = CUDA.zeros(UInt, 1)
-    @cuda threads=1 host_reference_kernel(out)
+    @cuda threads=1 runtime_refs.host_reference_kernel(out)
     expected = UInt(unsafe_load(cglobal(:jl_float32_type, Ptr{UInt})))
     @test Array(out)[] == expected
+
+    value_refs = @eval module $(gensym())
+        const host_reference_symbol = Symbol("value#global")
+        host_reference_symbol_kernel(out, name) =
+            (out[1] = UInt(name === host_reference_symbol); return)
+    end
+
+    out = CUDA.zeros(UInt, 1)
+    @cuda threads=1 value_refs.host_reference_symbol_kernel(out, Symbol("value#global"))
+    @test Array(out)[] == 1
 end
 
 ############################################################################################