Switch POCL backend to CompilerCaching.jl

Project.toml

@@ -22,6 +22,9 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 pocl_standalone_jll = "54f56a70-6062-5590-a942-1226658f6c83"
 
+[sources]
+GPUCompiler = {url = "https://github.com/JuliaGPU/GPUCompiler.jl", rev = "tb/compilercaching"}
+
 [weakdeps]
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -36,9 +39,9 @@ SparseArraysExt = "SparseArrays"
 Adapt = "0.4, 1.0, 2.0, 3.0, 4"
 Atomix = "0.1, 1"
 EnzymeCore = "0.7, 0.8.1"
-GPUCompiler = "1.13.3"
+GPUCompiler = "2"
 InteractiveUtils = "1.6"
-LLVM = "9.4.1"
+LLVM = "9.9"
 LinearAlgebra = "1.6"
 MacroTools = "0.5"
 PrecompileTools = "1"

src/pocl/compiler/compilation.jl

@@ -4,6 +4,32 @@ struct OpenCLCompilerParams <: AbstractCompilerParams end
 const OpenCLCompilerConfig = CompilerConfig{SPIRVCompilerTarget, OpenCLCompilerParams}
 const OpenCLCompilerJob = CompilerJob{SPIRVCompilerTarget, OpenCLCompilerParams}
 
+"""
+    OpenCLResults
+
+Cached compilation results for an OpenCL kernel job, managed by
+`GPUCompiler.cached_results`. Fields are populated through the compile pipeline:
+`obj` (SPIR-V bytes) + `entry` + `device_rng` after codegen, and `kernels` after the
+session-local link onto an OpenCL context. The first three are session-portable
+(cached through precompilation, except when GPUCompiler marks the job
+session-dependent and wipes its entries before image serialization); `kernels` is
+session-local and never populated during precompilation. `obj === nothing`
+identifies a job that has not been compiled yet.
+
+`kernels` is a small linear cache of `(cl.Context, cl.Kernel)` pairs. The cache partition
+already covers everything that affects codegen via `GPUCompiler.cache_owner`, so the only
+runtime-visible dimension left is the OpenCL context that owns the linked `cl.Kernel`.
+A linear scan with `===` is fastest in the common case (n=1) and stays cheap for the
+rare workload that bounces between a handful of contexts on the same device.
+"""
+mutable struct OpenCLResults
+    obj::Union{Nothing, Vector{UInt8}}                   # SPIR-V binary
+    entry::Union{Nothing, String}
+    device_rng::Bool
+    kernels::Vector{Tuple{cl.Context, cl.Kernel}}        # session-local; linear-scanned
+    OpenCLResults() = new(nothing, nothing, false, Tuple{cl.Context, cl.Kernel}[])
+end
+
 GPUCompiler.runtime_module(::CompilerJob{<:Any, OpenCLCompilerParams}) = POCL
 
 GPUCompiler.method_table_view(job::OpenCLCompilerJob) = GPUCompiler.StackedMethodTable(job.world, method_table, SPIRVIntrinsics.method_table)
@@ -105,18 +131,7 @@ function GPUCompiler.finish_linked_module!(@nospecialize(job::OpenCLCompilerJob)
 end
 
 
-## compiler implementation (cache, configure, compile, and link)
-
-# cache of compilation caches, per context
-const _compiler_caches = Dict{cl.Context, Dict{Any, Any}}()
-function compiler_cache(ctx::cl.Context)
-    cache = get(_compiler_caches, ctx, nothing)
-    if cache === nothing
-        cache = Dict{Any, Any}()
-        _compiler_caches[ctx] = cache
-    end
-    return cache
-end
+## compiler implementation (configure, compile, and link)
 
 # cache of compiler configurations, per device (but additionally configurable via kwargs)
 const _toolchain = Ref{Any}()
@@ -141,10 +156,13 @@ end
     return CompilerConfig(target, params; kernel, name, always_inline)
 end
 
-# compile to executable machine code
-function compile(@nospecialize(job::CompilerJob))
-    # TODO: this creates a context; cache those.
-    return obj, meta = JuliaContext() do ctx
+# run inference + LLVM codegen + SPIR-V emission. returns `(obj, entry, device_rng)`,
+# all session-portable so they survive precompilation when stored on a cached `CodeInstance`.
+const compilations = Threads.Atomic{Int}(0)
+function compile_to_obj(@nospecialize(job::CompilerJob))
+    compilations[] += 1
+
+    return JuliaContext() do ctx
         obj, meta = GPUCompiler.compile(:obj, job)
 
         entry = LLVM.name(meta.entry)
@@ -154,13 +172,13 @@ function compile(@nospecialize(job::CompilerJob))
     end
 end
 
-# link into an executable kernel
-function link(@nospecialize(job::CompilerJob), compiled)
+# link the SPIR-V bytes into a session-local `cl.Kernel` on the active context.
+function link_kernel(@nospecialize(job::CompilerJob), obj::Vector{UInt8}, entry::String)
     prog = if "cl_khr_il_program" in device().extensions
-        cl.Program(compiled.obj, context())
+        cl.Program(obj, context())
     else
         error("Your device does not support SPIR-V, which is currently required for native execution.")
     end
     cl.build!(prog)
-    return (; kernel = cl.Kernel(prog, compiled.entry), compiled.device_rng)
+    return cl.Kernel(prog, entry)
 end

src/pocl/compiler/execution.jl

@@ -195,27 +195,56 @@ end
 const clfunction_lock = ReentrantLock()
 
 function clfunction(f::F, tt::TT = Tuple{}; kwargs...) where {F, TT}
-    ctx = context()
-    dev = device()
-
     Base.@lock clfunction_lock begin
-        # compile the function
-        cache = compiler_cache(ctx)
+        config = compiler_config(device(); kwargs...)::OpenCLCompilerConfig
         source = methodinstance(F, tt)
-        config = compiler_config(dev; kwargs...)::OpenCLCompilerConfig
-        linked = GPUCompiler.cached_compilation(cache, source, config, compile, link)
-
-        # create a callable object that captures the function instance. we don't need to think
-        # about world age here, as GPUCompiler already does and will return a different object
-        h = hash(linked.kernel, hash(f, hash(tt)))
-        kernel = get(_kernel_instances, h, nothing)
+        job = CompilerJob(source, config)
+
+        res = compile_or_lookup(job)::OpenCLResults
+
+        # Resolve the cl.Kernel for the active context. Linear scan over the
+        # session-local cache; almost always n=1, so this is one `===` compare.
+        ctx = context()
+        kernel = nothing
+        @inbounds for (cached_ctx, cached_kernel) in res.kernels
+            if cached_ctx === ctx
+                kernel = cached_kernel
+                break
+            end
+        end
         if kernel === nothing
-            # create the kernel state object
-            kernel = HostKernel{F, tt}(f, linked.kernel, linked.device_rng)
-            _kernel_instances[h] = kernel
+            kernel = link_kernel(job, res.obj::Vector{UInt8}, res.entry::String)
+            # Don't cache session-local kernel handles while precompiling: the
+            # results struct is serialized into the package image along with its
+            # CodeInstance, and the handles would come back dangling.
+            if ccall(:jl_generating_output, Cint, ()) != 1
+                push!(res.kernels, (ctx, kernel))
+            end
         end
-        return kernel::HostKernel{F, tt}
+
+        h = hash(kernel, hash(f, hash(tt)))
+        return get!(_kernel_instances, h) do
+            HostKernel{F, tt}(f, kernel, res.device_rng)
+        end::HostKernel{F, tt}
+    end
+end
+
+# Look up cached compile artifacts for `job`, compiling on miss. Storage is managed
+# by `GPUCompiler.cached_results` (Julia's integrated code cache on 1.11+, which also
+# persists artifacts through precompilation; a session-local store on 1.10).
+#
+# `obj === nothing` identifies an `OpenCLResults` that hasn't been compiled yet. The
+# `compile_hook` check additionally forces the compile path so reflection-style
+# consumers (`@device_code_*`) observe the compilation even on a cache hit.
+function compile_or_lookup(@nospecialize(job::CompilerJob))::OpenCLResults
+    res = GPUCompiler.cached_results(OpenCLResults, job)
+    if res.obj === nothing || GPUCompiler.compile_hook[] !== nothing
+        compiled = compile_to_obj(job)
+        res.obj = compiled.obj
+        res.entry = compiled.entry
+        res.device_rng = compiled.device_rng
     end
+    return res
 end
 
 # cache of kernel instances

src/pocl/device/runtime.jl

@@ -1,6 +1,3 @@
-# reset the runtime cache from global scope, so that any change triggers recompilation
-GPUCompiler.reset_runtime()
-
 signal_exception() = return
 
 malloc(sz) = C_NULL

test/runtests.jl

@@ -11,6 +11,50 @@ end
 KernelAbstractions.versioninfo(POCLBackend())
 @info "Configuration" pocl = KernelAbstractions.POCL.nanoOpenCL.pocl_standalone_jll.libpocl
 
+import KernelAbstractions.POCL: POCL, @opencl, @device_code_llvm
+
+@testset "POCL compilation cache" begin
+    mod = @eval module $(gensym())
+    @noinline child() = return
+    kernel() = child()
+    end
+
+    count() = POCL.compilations[]
+    launch() = @opencl mod.kernel()
+
+    # the initial launch compiles
+    n = count()
+    Base.invokelatest(launch)
+    @test count() == n + 1
+
+    # a second launch hits the cache
+    Base.invokelatest(launch)
+    @test count() == n + 1
+
+    # jobs differing only in codegen-level settings get their own artifacts...
+    POCL.clfunction(mod.kernel, Tuple{}; name = "custom")
+    @test count() == n + 2
+    # ... which are cached as well
+    POCL.clfunction(mod.kernel, Tuple{}; name = "custom")
+    @test count() == n + 2
+
+    # reflection observes already-compiled kernels (by forcing recompilation,
+    # which must leave the cached entry in a usable state)
+    @test !isempty(sprint(io -> (@device_code_llvm io = io Base.invokelatest(launch))))
+    n = count()
+    Base.invokelatest(launch)
+    @test count() == n
+
+    # redefining the kernel recompiles...
+    @eval mod kernel() = (child(); child())
+    Base.invokelatest(launch)
+    @test count() == n + 1
+    # ... as does redefining a callee
+    @eval mod @noinline child() = nothing
+    Base.invokelatest(launch)
+    @test count() == n + 2
+end
+
 @testset "CPU back-end" begin
     struct CPUBackendArray{T, N, A} end # Fake and unused
     Testsuite.testsuite(CPU, "CPU", Base, Array, CPUBackendArray)