diff --git a/Project.toml b/Project.toml index ae9646f1..8f08964f 100644 --- a/Project.toml +++ b/Project.toml @@ -31,9 +31,9 @@ SPIRVIntrinsics = {path = "lib/intrinsics"} [compat] Adapt = "4" GPUArrays = "11.2.1" -GPUCompiler = "1.7.1" +GPUCompiler = "2" KernelAbstractions = "0.9.38" -LLVM = "9.1" +LLVM = "9.6" LinearAlgebra = "1" OpenCL_jll = "=2024.10.24" Preferences = "1" diff --git a/src/OpenCL.jl b/src/OpenCL.jl index 96e4889f..af17604a 100644 --- a/src/OpenCL.jl +++ b/src/OpenCL.jl @@ -37,6 +37,7 @@ include("compiler/capabilities.jl") include("compiler/compilation.jl") include("compiler/execution.jl") include("compiler/reflection.jl") +include("compiler/precompile.jl") # integrations and specialized functionality include("util.jl") diff --git a/src/compiler/compilation.jl b/src/compiler/compilation.jl index 3ca31795..348cdf8c 100644 --- a/src/compiler/compilation.jl +++ b/src/compiler/compilation.jl @@ -12,6 +12,32 @@ 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}) = OpenCL GPUCompiler.method_table_view(job::OpenCLCompilerJob) = @@ -125,18 +151,7 @@ function GPUCompiler.finish_linked_module!(@nospecialize(job::OpenCLCompilerJob) return 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}() @@ -157,7 +172,8 @@ end const SPIRV_VERSION = v"1.4" @noinline function _compiler_config(dev, backend; kernel=true, name=nothing, always_inline=false, - sub_group_size::Union{Nothing,Int}=_sub_group_size(dev), kwargs...) + sub_group_size::Union{Nothing,Int}=_sub_group_size(dev), + extensions::AbstractVector{<:AbstractString}=String[], kwargs...) supports_fp16 = "cl_khr_fp16" in dev.extensions supports_fp64 = "cl_khr_fp64" in dev.extensions @@ -165,26 +181,26 @@ const SPIRV_VERSION = v"1.4" error("Device does not support cl_intel_required_subgroup_size") end + spirv_ext = join(("+$ext" for ext in extensions), ",") + # create GPUCompiler objects target = SPIRVCompilerTarget(; version=SPIRV_VERSION, supports_fp16, supports_fp64, - validate=true, kwargs...) + validate=true, extensions=spirv_ext, kwargs...) params = OpenCLCompilerParams(; sub_group_size, features=device_features(dev), program_backend=backend) CompilerConfig(target, params; kernel, name, always_inline) end -# compile to executable machine code +# 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(@nospecialize(job::CompilerJob)) - compilations[] += 1 +function compile_to_obj(@nospecialize(job::CompilerJob)) + Threads.atomic_add!(compilations, 1) - # TODO: this creates a context; cache those. - obj, meta = JuliaContext() do ctx + JuliaContext() do ctx obj, meta = GPUCompiler.compile(:obj, job) - entry = LLVM.name(meta.entry) device_rng = StringAttribute("julia.opencl.rng", "") in collect(function_attributes(meta.entry)) - (; obj, entry, device_rng) end end @@ -289,9 +305,9 @@ function dump_artifacts(artifacts::Pair{String}...) return paths end -# link into an executable kernel -function link(@nospecialize(job::CompilerJob), compiled) - spirv_bitcode = compiled.obj +# 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) + spirv_bitcode = obj clc_source = nothing build_options = "" @@ -342,5 +358,5 @@ function link(@nospecialize(job::CompilerJob), compiled) msg *= "\nIf you think this is a bug, please file an issue and attach $(join(files, " and "))" error(msg) end - (; kernel=cl.Kernel(prog, compiled.entry), compiled.device_rng) + return cl.Kernel(prog, entry) end diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index f32433e0..07a9cadc 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -180,27 +180,57 @@ end const clfunction_lock = ReentrantLock() function clfunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT} - ctx = cl.context() - dev = cl.device() - Base.@lock clfunction_lock begin - # compile the function - cache = compiler_cache(ctx) + config = compiler_config(cl.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 = cl.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))) + 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 === nothing || res.obj === nothing || GPUCompiler.compile_hook[] !== nothing + compiled = compile_to_obj(job) + res = @something res GPUCompiler.cached_results(OpenCLResults, 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/compiler/precompile.jl b/src/compiler/precompile.jl new file mode 100644 index 00000000..00fa3207 --- /dev/null +++ b/src/compiler/precompile.jl @@ -0,0 +1,6 @@ +# Cache plumbing reached on every kernel launch. These specializations need no OpenCL device +# and keep the first launch from paying for GPUCompiler's generic results machinery itself. +let Job = OpenCLCompilerJob + precompile(Tuple{typeof(compile_or_lookup), Job}) + precompile(Tuple{typeof(GPUCompiler.cached_results), Type{OpenCLResults}, Job}) +end diff --git a/src/device/runtime.jl b/src/device/runtime.jl index ef2f6083..5a8cd6fa 100644 --- a/src/device/runtime.jl +++ b/src/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/Project.toml b/test/Project.toml index 64eacfcf..df290bb5 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,8 +1,10 @@ [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +CompilerCaching = "9db33cc3-5358-4881-8759-fa4194144afd" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" +GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55" IOCapture = "b5f81e59-6552-4d32-b1f0-c071b021bf89" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" diff --git a/test/execution.jl b/test/execution.jl index 7d29e821..9fddf0b3 100644 --- a/test/execution.jl +++ b/test/execution.jl @@ -167,3 +167,45 @@ end @test occursin("54321", ir) && !occursin("12345", ir) end end + +@testset "compilation cache" begin + mod = @eval module $(gensym()) + @noinline child() = return + kernel() = child() + end + + count() = OpenCL.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... + OpenCL.clfunction(mod.kernel, Tuple{}; name="custom") + @test count() == n+2 + # ... which are cached as well + OpenCL.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