pocl: define the device runtime as overlays so sysimages load

Project.toml

@@ -49,8 +49,8 @@ oneAPIExt = "oneAPI"
 AMDGPU = "2"
 Adapt = "0.4, 1.0, 2.0, 3.0, 4"
 Atomix = "1.5"
-GPUCompiler = "2.18"
-GPUToolbox = "3.4.1"
+GPUCompiler = "2.19"
+GPUToolbox = "3.5"
 KernelInterface = "0.4"
 LLVM = "10"
 LinearAlgebra = "1.6"
@@ -61,7 +61,7 @@ Printf = "<0.0.1, 1"
 Random = "1"
 Random123 = "1.7.1"
 RandomNumbers = "1.6.0"
-SPIRVIntrinsics = "1.4"
+SPIRVIntrinsics = "1.6"
 SPIRV_LLVM_Backend_jll = "23"
 SPIRV_Tools_jll = "2024.4, 2025.1"
 SparseArrays = "<0.0.1, 1.6"

src/pocl/compiler/compilation.jl

@@ -79,8 +79,11 @@ function GPUCompiler.finish_module!(
         ft = typeof(f)
         tt = Tuple{}
 
-        # create a deferred compilation job for `initialize_rng_state`
-        src = methodinstance(ft, tt, job.world)
+        # create a deferred compilation job for `initialize_rng_state`. it is an overlay, so
+        # look it up in the job's method tables.
+        src = GPUCompiler.generic_methodinstance(
+            ft, tt, job.world; method_table_view = GPUCompiler.method_table_view(job)
+        )
         cfg = CompilerConfig(job.config; kernel = false, name = nothing)
         job = CompilerJob(src, cfg, job.world)
         id = length(GPUCompiler.deferred_codegen_jobs) + 1
@@ -107,10 +110,9 @@ function GPUCompiler.finish_module!(
             end
             fptr = call!(builder, deferred_codegen_ft, deferred_codegen, [ConstantInt(id)])
 
-            # call the `initialize_rng_state` function
-            rt = Core.Compiler.return_type(f, tt)
-            llvm_rt = convert(LLVMType, rt)
-            llvm_ft = LLVM.FunctionType(llvm_rt)
+            # call the `initialize_rng_state` function, which returns nothing. inferring its
+            # return type natively would leave host code for it in sysimages.
+            llvm_ft = LLVM.FunctionType(LLVM.VoidType())
             fptr = inttoptr!(builder, fptr, LLVM.PointerType(llvm_ft))
             call!(builder, llvm_ft, fptr)
             br!(builder, top_bb)

src/pocl/device/random.jl

@@ -4,24 +4,26 @@ using Random
 import RandomNumbers
 
 # local memory with the actual seed, per subgroup, set by `initialize_rng_state`` or overridden by calling `seed!`
-@inline function global_random_keys()
+@device_function @inline function global_random_keys()
     n = get_num_sub_groups()
     ptr = random_keys()::LLVMPtr{UInt32, AS.Workgroup}
     return CLDeviceArray{UInt32, 1, AS.Workgroup}((n,), ptr)
 end
 
 # local memory with per-subgroup counters, incremented when generating numbers
-@inline function global_random_counters()
+@device_function @inline function global_random_counters()
     n = get_num_sub_groups()
     ptr = random_counters()::LLVMPtr{UInt32, AS.Workgroup}
     return CLDeviceArray{UInt32, 1, AS.Workgroup}((n,), ptr)
 end
 
-# initialization function, called automatically at the start of each kernel
-function initialize_rng_state()
+# initialization function, called automatically at the start of each kernel.
+# it must return nothing, as the kernel calls it as a void function.
+@device_function function initialize_rng_state()
     subgroup_id = get_sub_group_id()
     @inbounds global_random_keys()[subgroup_id] = kernel_state().random_seed
-    return @inbounds global_random_counters()[subgroup_id] = 0
+    @inbounds global_random_counters()[subgroup_id] = 0
+    return
 end
 
 # generators

src/pocl/device/runtime.jl

@@ -20,15 +20,19 @@ end
     return :(reinterpret(LLVMPtr{$T, AS.CrossWorkgroup}, kernel_state().exception_info + $offset))
 end
 
+# device-only functions are overlays, with a host method that throws. sysimage builds compile
+# host methods, and a sysimage that references device intrinsics fails to load
+# (JuliaGPU/GPUCompiler.jl#611).
+
 # the record is shared by all work-items on the device, so these need device scope, which
 # UnsafeAtomics' primitives emit as given (SPIRVIntrinsics' atomics are relaxed)
-@inline function atomic_store_device!(ptr::LLVMPtr{Int32, AS.CrossWorkgroup}, val::Int32)
+@device_function @inline function atomic_store_device!(ptr::LLVMPtr{Int32, AS.CrossWorkgroup}, val::Int32)
     UnsafeAtomics.Internal.llvm_store!(
         ptr, val, Val(:release), Val(:device), Val(false), Val(sizeof(Int32)), Val(())
     )
     return
 end
-@inline function atomic_cas_device!(ptr::LLVMPtr{Int32, AS.CrossWorkgroup}, cmp::Int32, val::Int32)
+@device_function @inline function atomic_cas_device!(ptr::LLVMPtr{Int32, AS.CrossWorkgroup}, cmp::Int32, val::Int32)
     (; old) = UnsafeAtomics.Internal.llvm_cmpxchg!(
         ptr, cmp, val, Val(:acq_rel), Val(:acquire), Val(:device),
         Val(false), Val(false), Val(sizeof(Int32)), Val(())
@@ -40,7 +44,7 @@ end
 # would be interleaved. it can do so over several calls (e.g., a quirk printing the exception,
 # and GPUCompiler printing a backtrace). returns 1 when the output was just claimed, 2 when
 # it was claimed by an earlier call of this work-item, and 0 when it belongs to another.
-@noinline function claim_output()
+@device_function @noinline function claim_output()
     lock = exception_field(Val(:output_lock))
     me = (get_global_id(1), get_global_id(2), get_global_id(3), 0)
     state = atomic_cas_device!(lock, Int32(0), Int32(1))
@@ -55,7 +59,7 @@ end
     return state == 2 && unsafe_load(exception_field(Val(:owner))) == me ? 2 : 0
 end
 
-function report_exception(ex)
+@device_function function report_exception(ex)
     claim = claim_output()
     if claim == 1
         SPIRVIntrinsics.@printf(
@@ -68,7 +72,7 @@ function report_exception(ex)
     return
 end
 
-function report_exception_name(ex)
+@device_function function report_exception_name(ex)
     claim = claim_output()
     if claim == 1
         SPIRVIntrinsics.@printf(
@@ -81,19 +85,21 @@ function report_exception_name(ex)
     return
 end
 
-function report_exception_frame(idx, func, file, line)
+@device_function function report_exception_frame(idx, func, file, line)
     if claim_output() != 0
         SPIRVIntrinsics.@printf(" [%d] %s at %s:%d\n", idx, func, file, line)
     end
     return
 end
 
 # the work-item that threw stops executing after this
-function signal_exception()
+@device_function function signal_exception()
     atomic_store_device!(exception_field(Val(:status)), Int32(1))
     return
 end
 
+# these use no device intrinsics, so they can stay host methods. as overlays, inference could
+# constant-fold calls to them by running the host method, which throws.
 malloc(sz) = C_NULL
 
 report_oom(sz) = return
@@ -107,7 +113,7 @@ struct KernelState
     exception_info::UInt64
 end
 
-@inline @generated kernel_state() = GPUCompiler.kernel_state_value(KernelState)
+@device_function @inline @generated kernel_state() = GPUCompiler.kernel_state_value(KernelState)
 
 ## intrinsics for adding and accessing additional kernel arguments
 
@@ -132,5 +138,5 @@ end
 end
 
 for name in [:random_keys, :random_counters]
-    @eval @inline $name() = additional_arg_value(LLVMPtr{UInt32, AS.Workgroup}, Val($(QuoteNode(name))))
+    @eval @device_function @inline $name() = additional_arg_value(LLVMPtr{UInt32, AS.Workgroup}, Val($(QuoteNode(name))))
 end

src/precompile.jl

@@ -8,15 +8,11 @@
     @inbounds A[i] = fma(2, B[i], 1)
 end
 
-# on Julia 1.11, the launch leaks GPUCompiler's runtime into the package image, which then
-# fails to link on Windows
-const launch_in_workload = !(Sys.iswindows() && v"1.11-" <= VERSION < v"1.12-")
-
 # whether the workload below launched its kernel, checked by the tests
 const precompiled_launch = Ref(false)
 
 PrecompileTools.@setup_workload begin
-    if launch_in_workload && POCL.nanoOpenCL.pocl_standalone_jll.is_available() &&
+    if POCL.nanoOpenCL.pocl_standalone_jll.is_available() &&
             POCL.SPIRV_LLVM_Backend_jll.is_available() && POCL.SPIRV_Tools_jll.is_available()
         try
             # keep PoCL's kernel cache out of the user's cache directory, and don't have it

test/Project.toml

@@ -2,6 +2,7 @@
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 FileCheck = "4e644321-382b-4b05-b0b6-5d23c3d944fb"
+GPUToolbox = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"

test/runtests.jl

@@ -1,6 +1,7 @@
 using KernelAbstractions
 using Random
 using Test
+import GPUToolbox
 
 include("quality_assurance.jl")
 include("linenumbers.jl")
@@ -23,8 +24,7 @@ end
 @testset "Precompilation" begin
     # without it, the first kernel launch on the CPU back-end takes many seconds
     workloads = Base.JLOptions().use_compiled_modules == 1 &&
-        KernelAbstractions.PrecompileTools.workload_enabled(KernelAbstractions) &&
-        KernelAbstractions.launch_in_workload
+        KernelAbstractions.PrecompileTools.workload_enabled(KernelAbstractions)
     @test KernelAbstractions.precompiled_launch[] skip = !workloads
 end
 
@@ -559,6 +559,12 @@ end
     @test occursin("throw_boundserror", out)
 end
 
+# sysimage builds compile methods for the host, which must not reach device-only code, or
+# the sysimage references device intrinsics and fails to load (JuliaGPU/GPUCompiler.jl#611)
+@testset "host leaks" begin
+    @test isempty(GPUToolbox.Overlays.host_leaks(KernelAbstractions))
+end
+
 @testset "POCL compilation cache" begin
     mod = @eval module $(gensym())
     @noinline child() = return