From e485282604d91d7a6a233c2e1e47aa6aad45c931 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 10 Jul 2026 15:51:23 +0200 Subject: [PATCH 1/3] Adapt to GPUCompiler 2 --- Project.toml | 4 +- src/compiler/compilation.jl | 78 +++++++++++++++++++++++-------------- src/compiler/execution.jl | 59 +++++++++++++++++++++------- 3 files changed, 97 insertions(+), 44 deletions(-) diff --git a/Project.toml b/Project.toml index b71260e9..29523d7f 100644 --- a/Project.toml +++ b/Project.toml @@ -17,6 +17,7 @@ LLVM = "929cbde3-209d-540e-8aea-75f648917ca0" Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" NEO_jll = "700fe977-ac61-5f37-bbc8-c6c4b2b6a9fd" +PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Preferences = "21216c6a-2e73-6563-6e65-726566657250" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -37,11 +38,12 @@ Adapt = "4" CEnum = "0.4, 0.5" ExprTools = "0.1" GPUArrays = "11.2.1" -GPUCompiler = "1.6" +GPUCompiler = "2" GPUToolbox = "0.1, 0.2, 0.3, 1, 3" KernelAbstractions = "0.9.39" LLVM = "6, 7, 8, 9" NEO_jll = "=26.18.38308" +PrecompileTools = "1" Preferences = "1" SPIRVIntrinsics = "1" SPIRV_LLVM_Backend_jll = "22" diff --git a/src/compiler/compilation.jl b/src/compiler/compilation.jl index 142b75cb..5a398a65 100644 --- a/src/compiler/compilation.jl +++ b/src/compiler/compilation.jl @@ -4,6 +4,30 @@ struct oneAPICompilerParams <: AbstractCompilerParams end const oneAPICompilerConfig = CompilerConfig{SPIRVCompilerTarget, oneAPICompilerParams} const oneAPICompilerJob = CompilerJob{SPIRVCompilerTarget,oneAPICompilerParams} +""" + oneAPIResults + +Cached compilation results for a oneAPI kernel job, managed by +`GPUCompiler.cached_results`. Fields are populated through the compile pipeline: +`image` (SPIR-V bytes) + `entry` after codegen, and `kernels` after the session-local +link onto a Level Zero module. The first two are session-portable (cached through +precompilation); `kernels` is session-local and never populated during precompilation. +`image === nothing` identifies a job that has not been compiled yet. + +`kernels` is a small linear cache of `(ZeContext, ZeDevice, ZeKernel)` tuples. The cache +partition already covers everything that affects codegen via `GPUCompiler.cache_owner`, so +the only runtime-visible dimensions left are the Level Zero context and device that own the +linked `ZeKernel` (a `ZeModule` is built for a specific `(context, device)` pair). 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 or devices. +""" +mutable struct oneAPIResults + image::Union{Nothing, Vector{UInt8}} # SPIR-V binary + entry::Union{Nothing, String} + kernels::Vector{Tuple{ZeContext, ZeDevice, ZeKernel}} # session-local; linear-scanned + oneAPIResults() = new(nothing, nothing, Tuple{ZeContext, ZeDevice, ZeKernel}[]) +end + GPUCompiler.runtime_module(::oneAPICompilerJob) = oneAPI GPUCompiler.method_table_view(job::oneAPICompilerJob) = @@ -52,9 +76,13 @@ function GPUCompiler.finish_ir!(job::oneAPICompilerJob, mod::LLVM.Module, # When the device supports BFloat16 but the SPIR-V runtime doesn't accept # SPV_KHR_bfloat16, lower all bfloat types to i16 so the translator can - # handle the module without the extension. + # handle the module without the extension. Both conditions are read from the + # (device-independent) compiler config so this stays valid without a live + # device: `_compiler_config` sets `supports_bfloat16` from the device and adds + # the `SPV_KHR_bfloat16` extension iff the driver's SPIR-V runtime accepts it. + target = job.config.target if @static(isdefined(Core, :BFloat16) && isdefined(LLVM, :BFloatType)) && - _device_supports_bfloat16() && !_driver_supports_bfloat16_spirv() + target.supports_bfloat16 && !occursin("SPV_KHR_bfloat16", target.extensions) lower_bfloat_to_i16!(mod) end @@ -266,24 +294,13 @@ function eliminate_bf16_bitcasts!(mod::LLVM.Module, T_bf16::LLVMType, T_i16::LLV end -## compiler implementation (cache, configure, compile, and link) - -# cache of compilation caches, per device -const _compiler_caches = Dict{ZeDevice, Dict{Any, Any}}() -function compiler_cache(dev::ZeDevice) - cache = get(_compiler_caches, dev, nothing) - if cache === nothing - cache = Dict{Any, Any}() - _compiler_caches[dev] = 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}() const _compiler_configs = Dict{UInt, oneAPICompilerConfig}() function compiler_config(dev; kwargs...) - h = hash(dev, hash(kwargs)) + h = hash(dev.driver, hash(dev, hash(kwargs))) config = get(_compiler_configs, h, nothing) if config === nothing config = _compiler_config(dev; kwargs...) @@ -292,10 +309,10 @@ function compiler_config(dev; kwargs...) return config end # Whether the driver's SPIR-V runtime accepts the SPV_KHR_bfloat16 extension. -function _driver_supports_bfloat16_spirv() +function _driver_supports_bfloat16_spirv(dev=device()) return @static if isdefined(Core, :BFloat16) haskey( - oneL0.extension_properties(driver()), + oneL0.extension_properties(dev.driver), oneL0.ZE_BFLOAT16_CONVERSIONS_EXT_NAME ) else @@ -304,26 +321,29 @@ function _driver_supports_bfloat16_spirv() end @noinline function _compiler_config(dev; kernel=true, name=nothing, always_inline=false, kwargs...) - supports_fp16 = oneL0.module_properties(device()).fp16flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP16 == oneL0.ZE_DEVICE_MODULE_FLAG_FP16 - supports_fp64 = oneL0.module_properties(device()).fp64flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP64 == oneL0.ZE_DEVICE_MODULE_FLAG_FP64 + properties = oneL0.module_properties(dev) + supports_fp16 = properties.fp16flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP16 == oneL0.ZE_DEVICE_MODULE_FLAG_FP16 + supports_fp64 = properties.fp64flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP64 == oneL0.ZE_DEVICE_MODULE_FLAG_FP64 # Allow BFloat16 in IR if the device supports it (even if the SPIR-V runtime doesn't # advertise the extension). We lower bfloat→i16 in finish_ir! when needed. - supports_bfloat16 = _device_supports_bfloat16() + supports_bfloat16 = _device_supports_bfloat16(dev) extensions = String[] # Only add the SPIR-V extension if the runtime actually supports it - if _driver_supports_bfloat16_spirv() + if _driver_supports_bfloat16_spirv(dev) push!(extensions, "SPV_KHR_bfloat16") end + extensions_str = join(map(ext -> "+$ext", extensions), ",") # create GPUCompiler objects - target = SPIRVCompilerTarget(; extensions, supports_fp16, supports_fp64, supports_bfloat16, kwargs...) + target = SPIRVCompilerTarget(; extensions=extensions_str, supports_fp16, supports_fp64, supports_bfloat16, kwargs...) params = oneAPICompilerParams() CompilerConfig(target, params; kernel, name, always_inline) end -# compile to executable machine code -function compile(@nospecialize(job::CompilerJob)) +# run inference + LLVM codegen + SPIR-V emission. returns `(image, entry)`, both +# session-portable so they survive precompilation when stored on a cached `CodeInstance`. +function compile_to_obj(@nospecialize(job::CompilerJob)) # TODO: on 1.9, this actually creates a context. cache those. asm, meta = JuliaContext() do ctx GPUCompiler.compile(:obj, job) @@ -332,10 +352,8 @@ function compile(@nospecialize(job::CompilerJob)) (image=asm, entry=LLVM.name(meta.entry)) end -# link into an executable kernel -function link(@nospecialize(job::CompilerJob), compiled) - ctx = context() - dev = device() - mod = ZeModule(ctx, dev, compiled.image) - kernels(mod)[compiled.entry] +# link the SPIR-V bytes into a session-local `ZeKernel` on the given context and device. +function link_kernel(image::Vector{UInt8}, entry::String, ctx::ZeContext, dev::ZeDevice) + mod = ZeModule(ctx, dev, image) + kernels(mod)[entry] end diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index 9ce5bc5c..cc8d3310 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -241,26 +241,59 @@ end const zefunction_lock = ReentrantLock() function zefunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT} - dev = device() - Base.@lock zefunction_lock begin - # compile the function - cache = compiler_cache(dev) - source = methodinstance(F, tt) + ctx = context() + dev = device() config = compiler_config(dev; kwargs...)::oneAPICompilerConfig - fun = GPUCompiler.cached_compilation(cache, source, config, compile, link) + source = methodinstance(F, tt) + job = CompilerJob(source, config) + + res = compile_or_lookup(job)::oneAPIResults + + # Resolve the ZeKernel for the active context and device. Linear scan over the + # session-local cache; almost always n=1, so this is one `===`/`==` compare. + fun = nothing + @inbounds for (cached_ctx, cached_dev, cached_kernel) in res.kernels + if cached_ctx == ctx && cached_dev == dev + fun = cached_kernel + break + end + end + if fun === nothing + fun = link_kernel(res.image::Vector{UInt8}, res.entry::String, ctx, dev) + # 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, dev, fun)) + end + end # 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(fun, hash(f, hash(tt))) - kernel = get(_kernel_instances, h, nothing) - if kernel === nothing - # create the kernel state object - kernel = HostKernel{F,tt}(f, fun) - _kernel_instances[h] = kernel - end - return kernel::HostKernel{F,tt} + get!(_kernel_instances, h) do + HostKernel{F,tt}(f, fun) + 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). +# +# `image === nothing` identifies a `oneAPIResults` 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))::oneAPIResults + res = GPUCompiler.cached_results(oneAPIResults, job) + if res === nothing || res.image === nothing || GPUCompiler.compile_hook[] !== nothing + compiled = compile_to_obj(job) + res = @something res GPUCompiler.cached_results(oneAPIResults, job) + res.image = compiled.image + res.entry = compiled.entry end + return res end # cache of kernel instances From 1bd3510c2c8691c225c5d44ce5742d649e9d0cfb Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 10 Jul 2026 15:51:55 +0200 Subject: [PATCH 2/3] CI: test GPUCompiler main before release --- .buildkite/pipeline.yml | 24 +++++++++++++++++++++++- Project.toml | 3 +++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index b0690075..4618770a 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -1,5 +1,28 @@ steps: # Test supported Julia versions + - label: "Julia 1.10" + plugins: + - JuliaCI/julia#v1: + version: "1.10" + - JuliaCI/julia-coverage#v1: + dirs: + - src + - lib + - examples + agents: + queue: "oneapi" + commands: | + git clone --depth 1 --branch main https://github.com/JuliaGPU/GPUCompiler.jl GPUCompiler + julia --project -e ' + using Pkg + Pkg.develop(path="GPUCompiler")' + julia --project=deps deps/build_ci.jl + julia --project -e ' + using Pkg + Pkg.test(; test_args=["--quickfail"])' + if: build.message !~ /\[skip tests\]/ + timeout_in_minutes: 120 + - group: ":julia: Julia" key: "julia" steps: @@ -23,7 +46,6 @@ steps: matrix: setup: julia: - - "1.10" - "1.11" - "1.12" - "nightly" diff --git a/Project.toml b/Project.toml index 29523d7f..fa00db2e 100644 --- a/Project.toml +++ b/Project.toml @@ -31,6 +31,9 @@ oneAPI_Level_Zero_Headers_jll = "f4bc562b-d309-54f8-9efb-476e56f0410d" oneAPI_Level_Zero_Loader_jll = "13eca655-d68d-5b81-8367-6d99d727ab01" oneAPI_Support_jll = "b049733a-a71d-5ed3-8eba-7d323ac00b36" +[sources] +GPUCompiler = {url = "https://github.com/JuliaGPU/GPUCompiler.jl", rev = "main"} + [compat] AbstractFFTs = "1.5.0" AcceleratedKernels = "0.3.1, 0.4" From fc976ae89e420b44b670e1124a4c5e3ded7bd544 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 10 Jul 2026 15:52:20 +0200 Subject: [PATCH 3/3] Clean up cache migration --- src/array.jl | 6 +++--- src/compiler/precompile.jl | 37 +++++++++++++++++++++++++++++++++++++ src/device/runtime.jl | 3 --- src/oneAPI.jl | 3 +++ 4 files changed, 43 insertions(+), 6 deletions(-) create mode 100644 src/compiler/precompile.jl diff --git a/src/array.jl b/src/array.jl index 99f5bef1..569da017 100644 --- a/src/array.jl +++ b/src/array.jl @@ -28,17 +28,17 @@ function contains_eltype(T, X) return false end -function _device_supports_bfloat16() +function _device_supports_bfloat16(dev=device()) # check the driver extension first if haskey( - oneL0.extension_properties(driver()), + oneL0.extension_properties(dev.driver), oneL0.ZE_BFLOAT16_CONVERSIONS_EXT_NAME ) return true end # some drivers (e.g. older versions on PVC/Max) don't advertise the extension, # but the hardware supports BFloat16 natively. fall back to checking device ID. - dev_id = oneL0.properties(device()).deviceId + dev_id = oneL0.properties(dev).deviceId # Intel Data Center GPU Max (Ponte Vecchio): device IDs 0x0BD0-0x0BDB if 0x0BD0 <= dev_id <= 0x0BDB return true diff --git a/src/compiler/precompile.jl b/src/compiler/precompile.jl new file mode 100644 index 00000000..4a8493c6 --- /dev/null +++ b/src/compiler/precompile.jl @@ -0,0 +1,37 @@ +using PrecompileTools: @compile_workload + +# Warm up the GPUCompiler -> SPIR-V pipeline during precompilation so the first real +# `zefunction` call is cheap. This doesn't need a GPU: SPIR-V codegen is device-independent +# (the target only carries version/extension/capability knobs), and the workload never +# launches anything. It is gated on the SPIR-V LLVM back-end being available so the package +# still precompiles on platforms/toolchains without it. +if SPIRV_LLVM_Backend_jll.is_available() + @compile_workload begin + let + function _precompile_kernel(a) + @inbounds a[1] += 1.0f0 + return + end + + # Build a device-independent compiler config. `_compiler_config` normally derives + # these knobs from the device; here we use conservative, portable defaults (the + # workload only exercises the pipeline, it does not target a specific device). + target = SPIRVCompilerTarget(; extensions="", supports_fp16=true, + supports_fp64=true, supports_bfloat16=false) + params = oneAPICompilerParams() + config = CompilerConfig(target, params; kernel=true, name=nothing, + always_inline=false) + + tt = Tuple{oneDeviceArray{Float32,1,AS.CrossWorkgroup}} + source = methodinstance(typeof(_precompile_kernel), tt) + job = CompilerJob(source, config) + + # On Julia < 1.12, driving GPU compilation during precompilation can leak foreign + # MethodInstances into host native compilation; only run the full compile on 1.12+. + @static if VERSION >= v"1.12-" + # Exercise the launch-side cache path and serialize its portable SPIR-V image. + compile_or_lookup(job) + end + end + end +end diff --git a/src/device/runtime.jl b/src/device/runtime.jl index 1db2ab2f..2f48b0f1 100644 --- a/src/device/runtime.jl +++ b/src/device/runtime.jl @@ -3,9 +3,6 @@ ## Julia library -# reset the runtime cache from global scope, so that any change triggers recompilation -GPUCompiler.reset_runtime() - function signal_exception() return end diff --git a/src/oneAPI.jl b/src/oneAPI.jl index f36f8171..2abcb64e 100644 --- a/src/oneAPI.jl +++ b/src/oneAPI.jl @@ -77,6 +77,9 @@ include("sorting.jl") include("indexing.jl") export oneAPIBackend +# precompilation workload (warms up the SPIR-V compilation pipeline) +include("compiler/precompile.jl") + # Work around a deadlock in Pkg's parallel precompilation on Julia 1.10, where it does # not pass `loadable_exts` to `Base.compilecache` (the kwarg is accidentally commented # out in Pkg's precompilation.jl), so a worker precompiling an extension freely loads