diff --git a/Project.toml b/Project.toml index 944049498..0a0e63842 100644 --- a/Project.toml +++ b/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" diff --git a/src/pocl/compiler/compilation.jl b/src/pocl/compiler/compilation.jl index 8aee85447..20a405e4a 100644 --- a/src/pocl/compiler/compilation.jl +++ b/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 diff --git a/src/pocl/compiler/execution.jl b/src/pocl/compiler/execution.jl index 595e7ff1a..3fdea393c 100644 --- a/src/pocl/compiler/execution.jl +++ b/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 diff --git a/src/pocl/device/runtime.jl b/src/pocl/device/runtime.jl index bec58123e..f99bde239 100644 --- a/src/pocl/device/runtime.jl +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index a655acb53..221d35bfc 100644 --- a/test/runtests.jl +++ b/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)