diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index ebb6c28d1..72222251f 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -14,6 +14,7 @@ steps: julia --project -e ' println("--- :julia: Instantiating project") using Pkg + Pkg.add(url="https://github.com/JuliaGPU/GPUCompiler.jl", rev="main") Pkg.instantiate() Pkg.build() @@ -35,14 +36,18 @@ steps: plugins: - JuliaCI/julia#v1: version: "1.10" - - JuliaCI/julia-test#v1: - JuliaCI/julia-coverage#v1: codecov: true agents: queue: "rocm" rocmgpu: "*" if: build.message !~ /\[skip tests\]/ - command: "julia --project -e 'using Pkg; Pkg.update()'" + command: | + git clone --depth 1 --branch main https://github.com/JuliaGPU/GPUCompiler.jl GPUCompiler + julia --project -e ' + using Pkg + Pkg.develop(path="GPUCompiler") + Pkg.test(; coverage=true)' timeout_in_minutes: 90 env: JULIA_NUM_THREADS: 4 @@ -132,13 +137,16 @@ steps: plugins: - JuliaCI/julia#v1: version: "1.10" - - JuliaCI/julia-test#v1: - test_args: "enzyme" agents: queue: "rocm" rocmgpu: "*" if: build.message !~ /\[skip tests\]/ - command: "julia --project -e 'using Pkg; Pkg.update()'" + command: | + git clone --depth 1 --branch main https://github.com/JuliaGPU/GPUCompiler.jl GPUCompiler + julia --project -e ' + using Pkg + Pkg.develop(path="GPUCompiler") + Pkg.test(; test_args=["enzyme"])' timeout_in_minutes: 45 env: JULIA_NUM_THREADS: 4 @@ -168,10 +176,11 @@ steps: plugins: - JuliaCI/julia#v1: version: "1.10" - - JuliaCI/julia-test#v1: - run_tests: false command: | + git clone --depth 1 --branch main https://github.com/JuliaGPU/GPUCompiler.jl GPUCompiler julia --project -e ' + using Pkg + Pkg.develop(path="GPUCompiler") using AMDGPU @assert !AMDGPU.functional()' agents: diff --git a/Project.toml b/Project.toml index 58afe83a5..0c9c73c40 100644 --- a/Project.toml +++ b/Project.toml @@ -49,6 +49,9 @@ AMDGPUEnzymeCoreExt = "EnzymeCore" AMDGPUSparseMatricesCSRExt = "SparseMatricesCSR" AMDGPUSpecialFunctionsExt = "SpecialFunctions" +[sources] +GPUCompiler = {url = "https://github.com/JuliaGPU/GPUCompiler.jl", rev = "main"} + [compat] AbstractFFTs = "1.0" AcceleratedKernels = "0.3.1, 0.4" @@ -60,7 +63,7 @@ ChainRulesCore = "1" EnzymeCore = "0.8" ExprTools = "0.1" GPUArrays = "11.3.1" -GPUCompiler = "1" +GPUCompiler = "2" GPUToolbox = "0.1.0, 0.2, 0.3, 1, 2, 3" KernelAbstractions = "0.9.2" LLD_jll = "15, 16, 17, 18, 19, 20, 21.1" diff --git a/src/compiler/codegen.jl b/src/compiler/codegen.jl index f859d7a32..e4cc59c4e 100644 --- a/src/compiler/codegen.jl +++ b/src/compiler/codegen.jl @@ -12,15 +12,33 @@ end const HIPCompilerConfig = CompilerConfig{GCNCompilerTarget, HIPCompilerParams} const HIPCompilerJob = CompilerJob{GCNCompilerTarget, HIPCompilerParams} -const _hip_compiler_cache = Dict{HIP.HIPDevice, Dict{Any, HIP.HIPFunction}}() +""" + HIPResults + +Cached compilation results for a HIP kernel job, managed by +`GPUCompiler.cached_results`. Session-portable artifacts (the lld-linked shared +object `obj`, the entry-point name `entry`, and the detected `global_hostcalls`) +are populated after codegen and persist across sessions (e.g. through package +precompilation). The session-local `functions` are `HIPFunction` handles linked +onto a specific device; they are device-specific and never populated during +precompilation. `obj === nothing` identifies a job that has not been compiled yet. + +`functions` is a small linear cache of `(HIPDevice, HIPFunction)` pairs, matching the +old per-device cache semantics; the scan is almost always over a single entry. +""" +mutable struct HIPResults + # session-portable artifacts + obj::Union{Nothing,Vector{UInt8}} # lld-linked shared object + entry::Union{Nothing,String} + global_hostcalls::Vector{Symbol} + # session-local handles (never populated during precompilation) + functions::Vector{Tuple{HIP.HIPDevice,HIP.HIPFunction}} + HIPResults() = new(nothing, nothing, Symbol[], Tuple{HIP.HIPDevice,HIP.HIPFunction}[]) +end # hash(fun, hash(f, hash(tt))) => HIPKernel const _kernel_instances = Dict{UInt, Runtime.HIPKernel}() -function compiler_cache(dev::HIP.HIPDevice) - get!(() -> Dict{UInt, Any}(), _hip_compiler_cache, dev) -end - GPUCompiler.runtime_module(@nospecialize(::HIPCompilerJob)) = AMDGPU GPUCompiler.method_table(@nospecialize(::HIPCompilerJob)) = AMDGPU.method_table @@ -168,12 +186,33 @@ The following kwargs are supported: function hipfunction(f::F, tt::TT = Tuple{}; kwargs...) where {F <: Core.Function, TT} Base.@lock hipfunction_lock begin dev = AMDGPU.device() - cache = compiler_cache(dev) config = compiler_config(dev; kwargs...) source = methodinstance(F, tt) - fun = GPUCompiler.cached_compilation( - cache, source, config, hipcompile, hiplink) + job = CompilerJob(source, config) + res = compile_or_lookup(job) + + # Resolve the `HIPFunction` for the active device. This is a session-local + # handle, so it lives in the results struct's linear cache rather than being + # persisted; the scan is almost always over a single entry, matching the old + # per-device cache (`==` compare, as `HIPDevice` was the Dict key before). + fun = nothing + for (cached_dev, cached_fun) in res.functions + if cached_dev == dev + fun = cached_fun + break + end + end + if fun === nothing + fun = hiplink(job, res.obj::Vector{UInt8}, res.entry::String, + res.global_hostcalls) + # Don't cache session-local handles while generating output: 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.functions, (dev, fun)) + end + end h = hash(fun, hash(f, hash(tt))) kernel = get!(_kernel_instances, h) do @@ -183,6 +222,25 @@ function hipfunction(f::F, tt::TT = Tuple{}; kwargs...) where {F <: Core.Functio end end +# Look up the cached compilation artifacts for `job`, running the compiler on a miss. +# +# Storage is managed by `GPUCompiler.cached_results`: Julia's integrated code cache on +# 1.11+ (which also persists artifacts through precompilation), or a session-local store +# on 1.10. `obj === nothing` identifies a freshly-created `HIPResults` that hasn't been +# compiled yet; the `compile_hook` check additionally forces the compile path so that +# reflection consumers (`@device_code_*`) observe the compilation even on a cache hit. +function compile_or_lookup(@nospecialize(job::CompilerJob))::HIPResults + res = GPUCompiler.cached_results(HIPResults, job) + if res === nothing || res.obj === nothing || GPUCompiler.compile_hook[] !== nothing + compiled = hipcompile(job) + res = @something res GPUCompiler.cached_results(HIPResults, job) + res.obj = compiled.obj + res.entry = compiled.entry + res.global_hostcalls = compiled.global_hostcalls + end + return res +end + function create_executable(obj) lld = if AMDGPU.lld_artifact `$(LLD_jll.lld()) -flavor gnu` @@ -250,8 +308,8 @@ function hipcompile(@nospecialize(job::CompilerJob)) (; obj=create_executable(codeunits(obj)), entry, global_hostcalls) end -function hiplink(@nospecialize(job::CompilerJob), compiled) - (; obj, entry, global_hostcalls) = compiled +# link a compiled shared object into a session-local `HIPFunction` on the active device. +function hiplink(@nospecialize(job::CompilerJob), obj, entry, global_hostcalls) mod = HIP.HIPModule(obj) HIP.HIPFunction(mod, entry, global_hostcalls) end diff --git a/src/device/runtime.jl b/src/device/runtime.jl index dcabf27cb..0084d40a2 100644 --- a/src/device/runtime.jl +++ b/src/device/runtime.jl @@ -2,9 +2,6 @@ using Core: LLVMPtr ## GPU runtime library -# reset the runtime cache from global scope, so that any change triggers recompilation -GPUCompiler.reset_runtime() - @inline @generated kernel_state() = GPUCompiler.kernel_state_value(AMDGPU.KernelState) @generated function llvm_atomic_cas(ptr::LLVMPtr{T,A}, cmp::T, val::T) where {T, A} diff --git a/src/precompile.jl b/src/precompile.jl index 36e7fa457..8355f5fc6 100644 --- a/src/precompile.jl +++ b/src/precompile.jl @@ -18,9 +18,9 @@ if :AMDGPU in LLVM.backends() end # Build a device-free compiler config for a baseline GCN target. - # `gfx1030` (RDNA2, wavefront 32) is a portable baseline that exercises - # the full pipeline; the cached *code* is reused regardless of the - # actual device's ISA at runtime (only the kernel binary differs). + # `gfx1030` (RDNA2, wavefront 32) is a representative baseline that exercises + # the full pipeline. GPU artifacts remain correctly partitioned by target ISA; + # the reusable benefit here is precompiled host-side compiler machinery. # # NOTE: the ISA must be RDNA/CDNA, not pre-RDNA. The `wavefrontsize*` # LLVM features only exist on gfx10+, so pairing them with e.g. gfx900 @@ -52,9 +52,9 @@ if :AMDGPU in LLVM.backends() # MIs into native compilation, causing LLVM errors. Guard like CUDA.jl. @static if VERSION >= v"1.12-" if !instrumented - GPUCompiler.JuliaContext() do ctx - GPUCompiler.compile(:obj, job) - end + # Exercise the same compile-or-lookup path used by kernel launches, and + # attach its artifact to the package-image CI. + Compiler.compile_or_lookup(job) # The compile above runs during precompilation, when ROCm # discovery (`__init__`) has NOT run, so `libdevice_libs` is @@ -72,14 +72,17 @@ end # Kernel launch infrastructure that the workload above cannot reach, because it # requires a live device (mirrors CUDA.jl's explicit precompile directives: -# `cufunction`, `link`, and `actual_compilation`). -precompile(Tuple{typeof(Compiler.hipfunction), typeof(identity), Type{Tuple{Nothing}}}) -precompile(Tuple{typeof(GPUCompiler.actual_compilation), - Dict{Any, HIP.HIPFunction}, Core.MethodInstance, UInt64, - Compiler.HIPCompilerConfig, typeof(Compiler.hipcompile), typeof(Compiler.hiplink)}) -precompile(Tuple{typeof(Compiler.hiplink), Compiler.HIPCompilerJob, - NamedTuple{(:obj, :entry, :global_hostcalls), - Tuple{Vector{UInt8}, String, Vector{Symbol}}}}) +# `cufunction`, `link_kernel`, `compile_or_lookup`, and `cached_results`). +let HIPCompilerJob = Compiler.HIPCompilerJob + precompile(Tuple{typeof(Compiler.hipfunction), typeof(identity), Type{Tuple{Nothing}}}) + precompile(Tuple{typeof(Compiler.hiplink), HIPCompilerJob, + Vector{UInt8}, String, Vector{Symbol}}) + + # GPUCompiler 2.0 caching pipeline (specialized for AMDGPU's results struct) + precompile(Tuple{typeof(Compiler.compile_or_lookup), HIPCompilerJob}) + precompile(Tuple{typeof(GPUCompiler.cached_results), + Type{Compiler.HIPResults}, HIPCompilerJob}) +end # Hot entry points of the bundled ROCm libraries, mirroring CUDA.jl's per-library # precompile directives. These compile the (GPU-free) Julia wrappers so the first