Adapt to GPUCompiler 2

Author: maleadt

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"

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

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

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

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

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

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