From bc02d87f30886ba6368e2d6a951109baf017cef8 Mon Sep 17 00:00:00 2001 From: AntonOresten Date: Sat, 29 Aug 2026 13:59:11 +0200 Subject: [PATCH 1/2] Make the job the cache identity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A cuTile compilation is a `TileJob`: a MethodInstance and world, and a `TileConfig` (target architecture and bytecode version, hints, const-seeded argument types, kernel name). The job keys the results cache, is what the `@device_code_*` hook reports, and is what the reflection entry points take. Inference is partitioned by a single owner, `:cuTile`. cuTile's inference depends on neither the target nor the hints — the interpreter takes only a world — so every configuration of a kernel shares one CodeInstance, with const-seeded arguments as `const_entries` on it; a Symbol is interned, so the owner is identity-stable for Julia's inference engine. Codegen results (`CuTileResults`: CUBIN and per-context kernels) are stored on that CodeInstance per config, so they survive world bumps and persist into package images with it. `compile_or_lookup(job)` runs `compile(job)` on a miss, under a lock. This replaces cuTile's own cache layer — `TileCacheKey`, the `CacheView` results plumbing, `ensure_compiled` and the cached `emit_*!` chain — and with it the per-hint re-inference that layer caused. Reflection and launches now share one inference partition, so `@device_code_typed` no longer re-infers. `get_inferred` reads its const-specialized entry directly: CompilerCaching's `get_source(ci, argtypes)` takes the first `CachedResult` on the CodeInstance whatever its results type, and `TileResults` is attached to the same one. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XNHajvrhhKBBUF7w6cDiKr --- docs/src/lib/reflection.md | 7 + src/compiler/codegen/kernel.jl | 2 +- src/compiler/driver.jl | 255 +++++++++--------- src/compiler/reflection.jl | 344 ++++++++++++------------ src/cuTile.jl | 2 +- src/launch.jl | 233 +++++++--------- src/precompile.jl | 6 +- test/codegen/boundscheck_integration.jl | 2 +- test/codegen/reflection.jl | 36 +++ test/host/cache.jl | 8 +- test/host/inference_cache.jl | 25 +- 11 files changed, 454 insertions(+), 466 deletions(-) diff --git a/docs/src/lib/reflection.md b/docs/src/lib/reflection.md index 34df5fdb..82dca1f4 100644 --- a/docs/src/lib/reflection.md +++ b/docs/src/lib/reflection.md @@ -15,6 +15,13 @@ code_ptx code_sass ``` +## Jobs + +```@docs +TileJob +tile_job +``` + ## Intercepting a launch ```@docs diff --git a/src/compiler/codegen/kernel.jl b/src/compiler/codegen/kernel.jl index cd71985d..f99361c4 100644 --- a/src/compiler/codegen/kernel.jl +++ b/src/compiler/codegen/kernel.jl @@ -387,7 +387,7 @@ function emit_subprogram!(ctx::CGCtx, func, arg_types::Vector, if !haskey(ctx.cache, mi) error("Expected $func($(join(arg_types, ", "))) to be cached already by inference.") end - sci, _, _ = emit_structured!(ctx.cache, mi) + sci, _, _ = emit_structured(emit_julia(ctx.cache, mi)...) # 2b. Run the pass pipeline on subprogram IR sub_divby, sub_bounds = run_passes!(sci) diff --git a/src/compiler/driver.jl b/src/compiler/driver.jl index 0853c808..71883fda 100644 --- a/src/compiler/driver.jl +++ b/src/compiler/driver.jl @@ -1,12 +1,9 @@ -# Compilation hook for @device_code_* macros - intercepts compilations for reflection -const compile_hook = Ref{Union{Nothing,Function}}(nothing) - #============================================================================= Meta nodes and compilation hints =============================================================================# -# Compilation options for cache sharding. +# Codegen options, as consumed by `emit_tile`. # Hint fields (opt_level, num_ctas, occupancy, num_worker_warps) represent explicit # overrides only; `nothing` means "consult @compiler_options meta nodes in the IR # during compilation." @@ -19,6 +16,105 @@ const CGOpts = @NamedTuple{ bytecode_version::VersionNumber } +#============================================================================= + Compiler jobs + + A cuTile compilation is a `TileJob`: a MethodInstance and world, a target + (architecture and bytecode version), params (hints and const-seeded argument + types) and a kernel name. The job names a compilation: it keys the results + cache (see `compile_or_lookup` in launch.jl), is what the `@device_code_*` + hook reports, and is what the reflection entry points take. +=============================================================================# + +struct TileCompilerTarget + sm_arch::Union{VersionNumber, Nothing} + bytecode_version::VersionNumber +end + +struct TileCompilerParams + opt_level::Union{Int, Nothing} + num_ctas::Union{Int, Nothing} + occupancy::Union{Int, Nothing} + num_worker_warps::Union{Int, Nothing} + # `(Const(f), arg2, …)` for const-seeded inference; a Tuple, not the driver's + # Vector, so that jobs compare structurally (`===` on immutables). + const_argtypes::Union{Tuple, Nothing} +end + +# Everything about a compilation except what is being compiled: keys the codegen +# results of a kernel's CodeInstance, which outlives any one world. +struct TileConfig + target::TileCompilerTarget + params::TileCompilerParams + name::String +end + +""" + TileJob + +A cuTile compilation: a `MethodInstance` and world, and a configuration — the +target architecture and bytecode version, the compilation hints and const-seeded +argument types, and the kernel name. Built by [`tile_job`](@ref); taken by the +reflection entry points and reported to the `@device_code_*` macros. +""" +struct TileJob + source::MethodInstance + world::UInt + config::TileConfig +end + +# Compilation hook for `@device_code_*`: called with every launch's `TileJob`. +const compile_hook = Ref{Union{Nothing,Function}}(nothing) + +# Inference is partitioned by the owner stamped on every CodeInstance; codegen +# results are stored per job on that CodeInstance (`TileResults`). cuTile's +# inference depends on neither the target nor the hints — the interpreter takes +# only a world — so every Tile job shares one partition: one inference per +# kernel, N codegens. Const-seeded arguments are `const_entries` on that shared +# CodeInstance. A `Symbol` is interned, so the owner is identity-stable for +# Julia's inference engine without any interning of our own. +const TILE_CACHE_OWNER = :cuTile + +# CompilerCaching's handle on the inference partition. The results type is +# `Nothing`: codegen artifacts are stored per job by `compile_or_lookup`, not here. +inference_cache(world::UInt) = CacheView{Any, Nothing}(TILE_CACHE_OWNER, world) +inference_cache(job::TileJob) = inference_cache(job.world) + +function tile_job(mi::MethodInstance, world::UInt; + sm_arch::Union{VersionNumber, Nothing}=nothing, + opt_level::Union{Int, Nothing}=nothing, + num_ctas::Union{Int, Nothing}=nothing, + occupancy::Union{Int, Nothing}=nothing, + num_worker_warps::Union{Int, Nothing}=nothing, + bytecode_version::VersionNumber=cuTile.bytecode_version(), + const_argtypes=nothing, + name::Union{String, Nothing}=nothing) + target = TileCompilerTarget(sm_arch, bytecode_version) + params = TileCompilerParams(opt_level, num_ctas, occupancy, num_worker_warps, + const_argtypes === nothing ? nothing : Tuple(const_argtypes)) + config = TileConfig(target, params, something(name, sanitize_name(string(mi.def.name)))) + return TileJob(mi, world, config) +end + +# The same job for another architecture. +with_target(job::TileJob, sm_arch::VersionNumber) = + TileJob(job.source, job.world, + TileConfig(TileCompilerTarget(sm_arch, job.config.target.bytecode_version), + job.config.params, job.config.name)) + +# The job's const-seeded argument types in the driver's `Vector{Any}` form. +function job_const_argtypes(job::TileJob) + cats = job.config.params.const_argtypes + cats === nothing ? nothing : collect(Any, cats) +end + +job_opts(job::TileJob) = CGOpts((sm_arch=job.config.target.sm_arch, + opt_level=job.config.params.opt_level, + num_ctas=job.config.params.num_ctas, + occupancy=job.config.params.occupancy, + num_worker_warps=job.config.params.num_worker_warps, + bytecode_version=job.config.target.bytecode_version)) + """ process_meta!(ir::CC.IRCode) -> ir @@ -182,139 +278,34 @@ function emit_tile(sci::StructuredIRCode, rettype, kernel_meta::Dict{Symbol,Any} end -#============================================================================= - Cached compilation -=============================================================================# - -# Results struct for caching compilation phases -mutable struct CuTileResults - julia_ir::Any # (StructuredIRCode, rettype) - tile_bc::Any # Vector{UInt8} bytecode - cuda_bin::Any # Vector{UInt8} cubin - cuda_func::Any # CuFunction - tile_kernel::Any # TileKernel{F, tt} wrapper around cuda_func - CuTileResults() = new(nothing, nothing, nothing, nothing, nothing) -end - -""" - ensure_compiled(cache, mi, const_argtypes) -> (ci::CodeInstance, res::CuTileResults) - -Hot-path entry point: single-pass cache `lookup`, falling back to `typeinf!` on -miss. Hoists the `(ci, res)` pair so the downstream emit_*! chain doesn't have -to resolve them on every phase. """ -@inline function ensure_compiled(cache::CacheView{K,CuTileResults}, - mi::Core.MethodInstance, - const_argtypes::Union{Vector{Any}, Nothing}) where {K} - if const_argtypes === nothing - hit = lookup(cache, mi) - hit !== nothing && return hit - - ci = get_ci(cache, mi; const_argtypes) - return (ci, results(cache, ci)) - else - hit = lookup(cache, mi, const_argtypes) - hit !== nothing && return (hit[1], results(hit[2])) - - ci = get_ci(cache, mi; const_argtypes) - entry = @something specialization(cache, ci, const_argtypes) - return (ci, results(entry)) - end -end - -# Cached wrappers around the driver's emit_* functions: check/populate -# CuTileResults. Used by the launch path; reflection calls the driver directly. - -# The `(f, tt)` a launch compiles, with `Constant{T,V}` argument types restored -# from `const_argtypes` (the MethodInstance has them unwrapped), for the -# `@device_code_*` hook. -function hook_signature(mi::Core.MethodInstance, const_argtypes::Union{Vector{Any}, Nothing}) - ftype = mi.specTypes.parameters[1] - f = isdefined(ftype, :instance) ? ftype.instance : ftype - arg_types = collect(Any, mi.specTypes.parameters[2:end]) - if const_argtypes !== nothing - # const_argtypes is [Const(f), arg2, ...]; arg_types omits f. - for i in eachindex(arg_types) - if const_argtypes[i+1] isa CC.Const - arg_types[i] = typeof(Constant(const_argtypes[i+1].val)) - end - end - end - return f, Tuple{arg_types...} -end + emit_bytecode(job::TileJob) -> (; bytecode, kernel_meta) +Julia phase through Tile IR for a job: inference (cached in the job's +partition), structurization and bytecode emission (recomputed on every call). """ - emit_structured!(cache, mi, ci, res; const_argtypes=nothing) -> (StructuredIRCode, rettype, kernel_meta) - -Cached IR phase. Checks `res.julia_ir` and delegates to `emit_structured` on -cache miss. - -`ci` and `res` must come from `ensure_compiled(cache, mi, const_argtypes)` — -i.e. inference is already done. Reflection callers that want the lookup -implicit can use the convenience `emit_structured!(cache, mi; const_argtypes)` -overload. -""" -function emit_structured!(cache::CacheView, mi::Core.MethodInstance, - ci::Core.CodeInstance, res::CuTileResults; - const_argtypes::Union{Vector{Any}, Nothing}=nothing) - res.julia_ir !== nothing && return res.julia_ir - - # Compute fresh via driver - ir, rettype = emit_julia(cache, mi; const_argtypes) - result = emit_structured(ir, rettype) - res.julia_ir = result - return result -end - -# Convenience overload that resolves (ci, res) on the caller's behalf. Use -# from non-hot paths (e.g. sub-program compilation in codegen/kernel.jl); -# the production launch path goes through `ensure_compiled` once at the top -# and feeds (ci, res) through the chain explicitly. -function emit_structured!(cache::CacheView, mi::Core.MethodInstance; - const_argtypes::Union{Vector{Any}, Nothing}=nothing) - ci, res = ensure_compiled(cache, mi, const_argtypes) - return emit_structured!(cache, mi, ci, res; const_argtypes) -end - -""" - emit_tile!(cache, mi, ci, res; const_argtypes=nothing) -> Vector{UInt8} - -Cached code phase. Checks `res.tile_bc`; on cache miss, takes the IR from -`emit_structured!` and calls the driver's `emit_tile`. -""" -function emit_tile!(cache::CacheView, mi::Core.MethodInstance, - ci::Core.CodeInstance, res::CuTileResults; - const_argtypes::Union{Vector{Any}, Nothing}=nothing) - res.tile_bc !== nothing && return res.tile_bc - - sci, rettype, kernel_meta = emit_structured!(cache, mi, ci, res; const_argtypes) - key = cache.owner::TileCacheKey - opts = CGOpts((sm_arch=unpack_version(key.sm_arch), - opt_level=unpack_hint(key.opt_level), - num_ctas=unpack_hint(key.num_ctas), - occupancy=unpack_hint(key.occupancy), - num_worker_warps=unpack_hint(key.num_worker_warps), - bytecode_version=unpack_version(key.bytecode_version))) +function emit_bytecode(job::TileJob) + cache = inference_cache(job) + const_argtypes = job_const_argtypes(job) + ir, rettype = emit_julia(cache, job.source; const_argtypes) + sci, rettype, kernel_meta = emit_structured(ir, rettype) bytecode = emit_tile(sci, rettype, kernel_meta; - name=sanitize_name(string(mi.def.name)), - opts, cache, const_argtypes) + name=job.config.name, opts=job_opts(job), cache, const_argtypes) + return (; bytecode, kernel_meta) +end - # Dump bytecode if JULIA_CUTILE_DUMP_BYTECODE is set +# Dump bytecode to `$JULIA_CUTILE_DUMP_BYTECODE/.ln[.n].cutile`, if set. +function dump_bytecode(mi::MethodInstance, bytecode::Vector{UInt8}) dump_dir = get(ENV, "JULIA_CUTILE_DUMP_BYTECODE", nothing) - if dump_dir !== nothing - mkpath(dump_dir) - base_filename = basename(string(mi.def.file)) - base_filename = first(splitext(base_filename)) - dump_path = joinpath(dump_dir, "$(base_filename).ln$(mi.def.line).cutile") - counter = 1 - while isfile(dump_path) - counter += 1 - dump_path = joinpath(dump_dir, "$(base_filename).ln$(mi.def.line).$(counter).cutile") - end - println(stderr, "Dumping TILEIR bytecode to file: $dump_path") - write(dump_path, bytecode) + dump_dir === nothing && return + mkpath(dump_dir) + base_filename = first(splitext(basename(string(mi.def.file)))) + dump_path = joinpath(dump_dir, "$(base_filename).ln$(mi.def.line).cutile") + counter = 1 + while isfile(dump_path) + counter += 1 + dump_path = joinpath(dump_dir, "$(base_filename).ln$(mi.def.line).$(counter).cutile") end - - res.tile_bc = bytecode - return bytecode + println(stderr, "Dumping TILEIR bytecode to file: $dump_path") + write(dump_path, bytecode) end diff --git a/src/compiler/reflection.jl b/src/compiler/reflection.jl index c1354786..0d4ce65c 100644 --- a/src/compiler/reflection.jl +++ b/src/compiler/reflection.jl @@ -21,20 +21,6 @@ function disassemble_tileir(bytecode::Vector{UInt8}, version::VersionNumber; end end -""" - code_typed(f, argtypes; world, kwargs...) -> Vector{Any} - -Return typed code for a cuTile function. Analogous to `Base.code_typed`. -""" -function code_typed(@nospecialize(f), @nospecialize(argtypes); - world::UInt=Base.get_world_counter(), kwargs...) - stripped, const_argtypes = process_const_argtypes(f, argtypes) - mi = lookup_method_instance(f, stripped; world) - cache = CacheView{CuTileResults}(:cuTile, world) - ir, rettype = emit_julia(cache, mi; const_argtypes) - [ir => rettype] -end - """ code_ircode(mi::MethodInstance; world, always_inline=true) -> (IRCode, rettype) @@ -43,8 +29,7 @@ If always_inline=true (default), forces all functions to be inlined. """ function code_ircode(mi::MethodInstance; world::UInt=Base.get_world_counter(), always_inline::Bool=true) - cache = CacheView{CuTileResults}(:cuTile, world) - interp = cuTileInterpreter(cache; always_inline) + interp = cuTileInterpreter(inference_cache(world); always_inline) result = CC.typeinf_ircode(interp, mi, nothing) if result === nothing @@ -55,26 +40,6 @@ function code_ircode(mi::MethodInstance; world::UInt=Base.get_world_counter(), return ir, rettype end -""" - code_structured(f, argtypes; kwargs...) -> Vector{Pair{StructuredIRCode, DataType}} - -Return the structured IR for a cuTile function. -""" -function code_structured(@nospecialize(f), @nospecialize(argtypes); - world::UInt=Base.get_world_counter(), - optimize::Bool=true) - stripped, const_argtypes = process_const_argtypes(f, argtypes) - mi = lookup_method_instance(f, stripped; world) - cache = CacheView{CuTileResults}(:cuTile, world) - ir, rettype = emit_julia(cache, mi; const_argtypes) - sci, rettype, _ = emit_structured(ir, rettype) - if optimize - sci = copy(sci) - run_passes!(sci) - end - [sci => rettype] -end - """ process_const_argtypes(f, argtypes) -> (stripped, const_argtypes) @@ -112,57 +77,117 @@ end constant_eltype(::Type{Constant{T,V}}) where {T,V} = T constant_value(::Type{Constant{T,V}}) where {T,V} = V + +#============================================================================= + Jobs +=============================================================================# + +public tile_job, TileJob + +""" + tile_job(f, argtypes; sm_arch, opt_level, num_ctas, occupancy, num_worker_warps, + bytecode_version, world) -> TileJob + +The [`TileJob`](@ref) for compiling `f` on `argtypes` (which may contain +`Constant{T,V}` types) with the given compilation options. `sm_arch` may be +`nothing` for stages that do not run `tileiras`. +""" +function tile_job(@nospecialize(f), @nospecialize(argtypes); + world::UInt=Base.get_world_counter(), kwargs...) + stripped, const_argtypes = process_const_argtypes(f, argtypes) + mi = lookup_method_instance(f, stripped; world) + tile_job(mi, world; const_argtypes, kwargs...) +end + +# `(f, tt)` with `Constant` argument types restored, for printing headers. +function job_signature(job::TileJob) + mi = job.source + ftype = mi.specTypes.parameters[1] + f = isdefined(ftype, :instance) ? ftype.instance : ftype + arg_types = collect(Any, mi.specTypes.parameters[2:end]) + cats = job.config.params.const_argtypes + if cats !== nothing + # cats is [Const(f), arg2, ...]; arg_types omits f. + for i in eachindex(arg_types) + cats[i+1] isa CC.Const && (arg_types[i] = typeof(Constant(cats[i+1].val))) + end + end + return f, Tuple{arg_types...} +end + +# The architecture tileiras will target: the job's, or the active device's. +function resolve_target(job::TileJob, purpose::String) + sm_arch = job.config.target.sm_arch + sm_arch === nothing || return sm_arch + try + default_sm_arch() + catch + throw(ArgumentError("sm_arch must be specified when $purpose without a CUDA device")) + end +end + +# tileiras's -O level: the job's, else the kernel's hint, else 3. +resolve_opt_level(job::TileJob, kernel_meta, target) = + something(resolve_hint(job.config.params.opt_level, kernel_meta, :opt_level, target), 3) + + +#============================================================================= + Stages +=============================================================================# + +""" + code_typed(f, argtypes; world, kwargs...) -> Vector{Any} + +Return typed code for a cuTile function. Analogous to `Base.code_typed`. +""" +function code_typed(job::TileJob) + ir, rettype = emit_julia(inference_cache(job), job.source; + const_argtypes=job_const_argtypes(job)) + [ir => rettype] +end +code_typed(@nospecialize(f), @nospecialize(argtypes); kwargs...) = + code_typed(tile_job(f, argtypes; kwargs...)) + +""" + code_structured(f, argtypes; kwargs...) -> Vector{Pair{StructuredIRCode, DataType}} + +Return the structured IR for a cuTile function. +""" +function code_structured(job::TileJob; optimize::Bool=true) + ir, rettype = emit_julia(inference_cache(job), job.source; + const_argtypes=job_const_argtypes(job)) + sci, rettype, _ = emit_structured(ir, rettype) + if optimize + sci = copy(sci) + run_passes!(sci) + end + [sci => rettype] +end +code_structured(@nospecialize(f), @nospecialize(argtypes); optimize::Bool=true, kwargs...) = + code_structured(tile_job(f, argtypes; kwargs...); optimize) + """ code_tiled([io::IO], f, argtypes; sm_arch, opt_level, num_ctas, occupancy, num_worker_warps, remarks=false) Print the CUDA Tile IR for a Julia function as a textual MLIR representation. -Analogous to `code_llvm`/`code_native`. Calls the driver directly without -caching in CuTileResults, so reflection never pollutes the compilation cache. +Analogous to `code_llvm`. Set `remarks=true` to also run `tileiras` and print its optimization remarks. This requires `tileiras` 13.4 or newer. When no GPU is available, pass `sm_arch` explicitly. """ -function code_tiled(io::IO, @nospecialize(f), @nospecialize(argtypes); - sm_arch::Union{VersionNumber, Nothing}=nothing, - opt_level::Union{Int, Nothing}=nothing, - num_ctas::Union{Int, Nothing}=nothing, - occupancy::Union{Int, Nothing}=nothing, - num_worker_warps::Union{Int, Nothing}=nothing, - bytecode_version::VersionNumber=cuTile.bytecode_version(), - debuginfo::Bool=false, - remarks::Bool=false, - world::UInt=Base.get_world_counter()) - stripped, const_argtypes = process_const_argtypes(f, argtypes) - mi = lookup_method_instance(f, stripped; world) - - opts = CGOpts((sm_arch=sm_arch, opt_level=opt_level, num_ctas=num_ctas, occupancy=occupancy, - num_worker_warps=num_worker_warps, bytecode_version=bytecode_version)) - cache = CacheView{CuTileResults}(:cuTile, world) - ir, rettype = emit_julia(cache, mi; const_argtypes) - sci, rettype, kernel_meta = emit_structured(ir, rettype) - bytecode = emit_tile(sci, rettype, kernel_meta; - name=sanitize_name(string(mi.def.name)), - opts, cache, const_argtypes) +function code_tiled(io::IO, job::TileJob; debuginfo::Bool=false, remarks::Bool=false) + (; bytecode, kernel_meta) = emit_bytecode(job) + bytecode_version = job.config.target.bytecode_version print(io, disassemble_tileir(bytecode, bytecode_version; debuginfo)) if remarks tileiras_version() >= v"13.4" || throw(ArgumentError( "tileiras optimization remarks require tileiras 13.4 or newer")) validate_tileiras_target(bytecode_version) - target = if sm_arch === nothing - try - default_sm_arch() - catch - throw(ArgumentError( - "sm_arch must be specified when requesting remarks without a CUDA device")) - end - else - sm_arch - end - resolved_opt_level = something(resolve_hint(opt_level, kernel_meta, - :opt_level, target), 3) - _, text = run_tileiras(bytecode, target, resolved_opt_level; remarks=true) + target = resolve_target(job, "requesting remarks") + _, text = run_tileiras(bytecode, target, resolve_opt_level(job, kernel_meta, target); + remarks=true) if !isempty(text) println(io) println(io, "// tileiras optimization remarks") @@ -172,53 +197,30 @@ function code_tiled(io::IO, @nospecialize(f), @nospecialize(argtypes); end end end +code_tiled(io::IO, @nospecialize(f), @nospecialize(argtypes); + debuginfo::Bool=false, remarks::Bool=false, kwargs...) = + code_tiled(io, tile_job(f, argtypes; kwargs...); debuginfo, remarks) code_tiled(@nospecialize(f), @nospecialize(argtypes); kwargs...) = code_tiled(stdout, f, argtypes; kwargs...) """ - compile_to_cubin(f, argtypes; sm_arch, opt_level, num_ctas, occupancy, - num_worker_warps, bytecode_version, world) -> Vector{UInt8} + compile(job::TileJob) -> Vector{UInt8} -Compile a Julia function all the way to a CUBIN for reflection: emit Tile IR -bytecode (like [`code_tiled`](@ref)) and assemble it with `tileiras`. Like the -rest of the reflection entry points, this calls the driver directly without -caching in CuTileResults or the disk cache. +Compile a job to a CUBIN: Tile IR bytecode, assembled with `tileiras` (through +the disk cache). Uncached otherwise; `compile_or_lookup` caches for launches. """ -function compile_to_cubin(@nospecialize(f), @nospecialize(argtypes); - sm_arch::Union{VersionNumber, Nothing}=nothing, - opt_level::Union{Int, Nothing}=nothing, - num_ctas::Union{Int, Nothing}=nothing, - occupancy::Union{Int, Nothing}=nothing, - num_worker_warps::Union{Int, Nothing}=nothing, - bytecode_version::VersionNumber=cuTile.bytecode_version(), - world::UInt=Base.get_world_counter()) - target = if sm_arch === nothing - try - default_sm_arch() - catch - throw(ArgumentError( - "sm_arch must be specified when compiling without a CUDA device")) - end - else - sm_arch +function compile(job::TileJob) + # Report to `@device_code_*`. Launches get here through `invoke_frozen`, but + # the hook closure lives in the user's latest world — `invokelatest` it. + if compile_hook[] !== nothing + Base.invokelatest(compile_hook[], job) end - validate_tileiras_target(bytecode_version) - - stripped, const_argtypes = process_const_argtypes(f, argtypes) - mi = lookup_method_instance(f, stripped; world) - opts = CGOpts((sm_arch=target, opt_level=opt_level, num_ctas=num_ctas, - occupancy=occupancy, num_worker_warps=num_worker_warps, - bytecode_version=bytecode_version)) - cache = CacheView{CuTileResults}(:cuTile, world) - ir, rettype = emit_julia(cache, mi; const_argtypes) - sci, rettype, kernel_meta = emit_structured(ir, rettype) - bytecode = emit_tile(sci, rettype, kernel_meta; - name=sanitize_name(string(mi.def.name)), - opts, cache, const_argtypes) - resolved_opt_level = something(resolve_hint(opt_level, kernel_meta, - :opt_level, target), 3) - cubin, _ = run_tileiras(bytecode, target, resolved_opt_level) - return cubin + job = with_target(job, resolve_target(job, "compiling")) + validate_tileiras_target(job.config.target.bytecode_version) + (; bytecode, kernel_meta) = emit_bytecode(job) + dump_bytecode(job.source, bytecode) + sm_arch = job.config.target.sm_arch + return assemble(bytecode, sm_arch, resolve_opt_level(job, kernel_meta, sm_arch)) end """ @@ -231,14 +233,12 @@ compiler decision (thread mapping, CTA size, pipelining, synchronization) already made. When no GPU is available, pass `sm_arch` explicitly. !!! warning "Unstable" - PTX is an implementation detail of `tileiras`, not an interface. It is read - from an undocumented debug section of the CUBIN and may stop being produced - or recorded at any time; `code_ptx` will be removed with it. + The PTX is recorded by `tileiras` in an undocumented CUBIN section and may + go away. """ -function code_ptx(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) - cubin = compile_to_cubin(f, argtypes; kwargs...) - print(io, extract_ptx(cubin)) -end +code_ptx(io::IO, job::TileJob) = print(io, extract_ptx(compile(job))) +code_ptx(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) = + code_ptx(io, tile_job(f, argtypes; kwargs...)) code_ptx(@nospecialize(f), @nospecialize(argtypes); kwargs...) = code_ptx(stdout, f, argtypes; kwargs...) @@ -248,16 +248,12 @@ code_ptx(@nospecialize(f), @nospecialize(argtypes); kwargs...) = Print the SASS machine code that a Julia function compiles to, by assembling the Tile IR with `tileiras` and disassembling the resulting CUBIN with -`nvdisasm`. When no GPU is available, pass `sm_arch` explicitly. - -For the binary a launch actually loaded, use `CUDA.@device_code_sass`. Unlike -that, this needs no CUPTI subscription, so it also works under Nsight or -another active profiler. +`nvdisasm`. When no GPU is available, pass `sm_arch` explicitly. For the +binary a launch actually loaded, use `CUDA.@device_code_sass`. """ -function code_sass(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) - cubin = compile_to_cubin(f, argtypes; kwargs...) - print(io, disassemble_cubin(cubin)) -end +code_sass(io::IO, job::TileJob) = print(io, disassemble_cubin(compile(job))) +code_sass(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) = + code_sass(io, tile_job(f, argtypes; kwargs...)) code_sass(@nospecialize(f), @nospecialize(argtypes); kwargs...) = code_sass(stdout, f, argtypes; kwargs...) @@ -270,22 +266,23 @@ export @device_code_tiled public @device_code_typed, @device_code_structured public @device_code_ptx -# Following GPUCompiler's pattern for @device_code_* macros +# Install `inner_hook` as the compile hook for the duration of the expression, +# called once per distinct job. function emit_hooked_compilation(inner_hook, ex...) user_code = ex[end] user_kwargs = ex[1:end-1] quote - seen = Set{Tuple{Any,Any}}() - function outer_hook(f, tt) - if !in((f, tt), seen) - old_hook = $compile_hook[] - try - $compile_hook[] = nothing - $inner_hook(f, tt; $(map(esc, user_kwargs)...)) - finally - $compile_hook[] = old_hook - end - push!(seen, (f, tt)) + jobs = Set{TileJob}() + function outer_hook(job::TileJob) + job in jobs && return + push!(jobs, job) + # the user hook might invoke the compiler again, so disable the hook + old_hook = $compile_hook[] + try + $compile_hook[] = nothing + $inner_hook(job; $(map(esc, user_kwargs)...)) + finally + $compile_hook[] = old_hook end end @@ -296,14 +293,23 @@ function emit_hooked_compilation(inner_hook, ex...) $compile_hook[] = nothing end - if isempty(seen) + if isempty(jobs) error("no kernels executed while evaluating the given expression") end - nothing end end +function tile_hook(inner) + function (job::TileJob; io::IO=stdout, kwargs...) + f, tt = job_signature(job) + println(io, "// $f($(join(tt.parameters, ", ")))") + println(io) + inner(io, job; kwargs...) + println(io) + end +end + """ @device_code_tiled [io=stdout] [remarks=false] expression @@ -316,72 +322,52 @@ With `remarks=true`, also print `tileiras` optimization remarks for each kernel. ``` """ macro device_code_tiled(ex...) - function hook(f, tt; io::IO=stdout, kwargs...) - println(io, "// $f($(join(tt.parameters, ", ")))") - println(io) - code_tiled(io, f, Tuple(tt.parameters); kwargs...) - println(io) - end + hook = tile_hook((io, job; kwargs...) -> code_tiled(io, job; kwargs...)) emit_hooked_compilation(hook, ex...) end """ - @device_code_ptx [io=stdout] expression + @device_code_structured [io=stdout] expression -Print the PTX generated by `tileiras` for all kernels compiled while -evaluating the expression. Unstable, like [`code_ptx`](@ref). +Print the StructuredIRCode for all kernels compiled while evaluating the expression. # Example ```julia -@device_code_ptx @cuda backend=cuTile blocks=grid vadd(a, b, c) +@device_code_structured @cuda backend=cuTile blocks=grid vadd(a, b, c) ``` """ -macro device_code_ptx(ex...) - function hook(f, tt; io::IO=stdout, kwargs...) - println(io, "// $f($(join(tt.parameters, ", ")))") - println(io) - code_ptx(io, f, Tuple(tt.parameters); kwargs...) - println(io) - end +macro device_code_structured(ex...) + hook = tile_hook((io, job; kwargs...) -> println(io, first(only(code_structured(job; kwargs...))))) emit_hooked_compilation(hook, ex...) end """ - @device_code_structured [io=stdout] expression + @device_code_typed [io=stdout] expression -Print the StructuredIRCode for all kernels compiled while evaluating the expression. +Print the typed Julia IR for all kernels compiled while evaluating the expression. # Example ```julia -@device_code_structured @cuda backend=cuTile blocks=grid vadd(a, b, c) +@device_code_typed @cuda backend=cuTile blocks=grid vadd(a, b, c) ``` """ -macro device_code_structured(ex...) - function hook(f, tt; io::IO=stdout, kwargs...) - println(io, "// $f($(join(tt.parameters, ", ")))") - println(io) - sci, _ = only(code_structured(f, Tuple(tt.parameters); kwargs...)) - println(io, sci) - end +macro device_code_typed(ex...) + hook = tile_hook((io, job; kwargs...) -> println(io, first(only(code_typed(job))))) emit_hooked_compilation(hook, ex...) end """ - @device_code_typed [io=stdout] expression + @device_code_ptx [io=stdout] expression -Print the typed Julia IR for all kernels compiled while evaluating the expression. +Print the PTX generated by `tileiras` for all kernels compiled while +evaluating the expression. Unstable, like [`code_ptx`](@ref). # Example ```julia -@device_code_typed @cuda backend=cuTile blocks=grid vadd(a, b, c) +@device_code_ptx @cuda backend=cuTile blocks=grid vadd(a, b, c) ``` """ -macro device_code_typed(ex...) - function hook(f, tt; io::IO=stdout, kwargs...) - println(io, "// $f($(join(tt.parameters, ", ")))") - println(io) - ci, _ = only(code_typed(f, Tuple(tt.parameters); kwargs...)) - println(io, ci) - end +macro device_code_ptx(ex...) + hook = tile_hook((io, job; kwargs...) -> code_ptx(io, job)) emit_hooked_compilation(hook, ex...) end diff --git a/src/cuTile.jl b/src/cuTile.jl index 51b79867..afd109f6 100644 --- a/src/cuTile.jl +++ b/src/cuTile.jl @@ -22,7 +22,7 @@ using EnumX public BFloat16 import CompilerCaching -using CompilerCaching: CacheView, method_instance, match_method_instance, typeinf!, results, lookup, +using CompilerCaching: CacheView, method_instance, match_method_instance, typeinf!, specialization, get_source # Shared definitions diff --git a/src/launch.jl b/src/launch.jl index dda1ec00..b1908d58 100644 --- a/src/launch.jl +++ b/src/launch.jl @@ -3,11 +3,12 @@ # Compiles a Julia function with `TileArray` arguments to Tile IR bytecode, # runs `tileiras` to lower bytecode → CUBIN, loads the cubin into the active # CUDA context, and launches it via `cudacall`. Compilation is cached per -# `(MethodInstance, sm_arch, opt_level, num_ctas, occupancy, num_worker_warps, bytecode_version)`. +# `TileJob` on the kernel's CodeInstance. -using ObjectFile: readmeta, Sections -using CUDACore: CUDACore, CuArray, CuModule, CuFunction, cudacall, device, capability, - AbstractBackend, AbstractKernel, kernel_convert, kernel_compile +using ObjectFile: ObjectFile, readmeta, Sections +using CUDACore: CUDACore, CuArray, CuModule, CuFunction, CuContext, context, cudacall, + device, capability, AbstractBackend, AbstractKernel, kernel_convert, + kernel_compile using CUDA_Compiler_jll using GPUToolbox: LazyInitialized using Preferences: @load_preference @@ -80,53 +81,6 @@ CUDACore.kernel_compile(::TileBackend, f::F, tt::TT=Tuple{}; kwargs...) where {F cufunction(f, tt; kwargs...) -#============================================================================= - Cache sharding key. -=============================================================================# - -# Pack a `VersionNumber` into a `UInt16` as `(major << 8) | minor`. Lossy: drops -# patch/prerelease/build (none of which we need for SM architectures or Tile IR -# bytecode versions). Used by `TileCacheKey` to keep the owner isbits. -@inline pack_version(v::VersionNumber) = (UInt16(v.major) << 8) | UInt16(v.minor) -@inline unpack_version(x::UInt16) = VersionNumber(Int(x >> 8), Int(x & 0xff)) - -# isbits sentinel codec for `Union{Int, Nothing}` hint fields (`opt_level`, -# `num_ctas`, `occupancy`, `num_worker_warps`). `-1` is unused as a value, so we -# use it for `nothing`. -const _UNSET = -1 -@inline pack_hint(x::Union{Int, Nothing}) = x === nothing ? _UNSET : x -@inline unpack_hint(x::Int) = x == _UNSET ? nothing : x - -""" - TileCacheKey - -Owner stamped onto every launch-cached `CodeInstance`. `lookup` (called once -per `cufunction` via `ensure_compiled`) ccalls `jl_rettype_inferred` with the -owner as `Any`, which forces a heap box; making the owner `isbits` shrinks the -box from ~80 B (the previous `Tuple{Symbol, NamedTuple{…}}` shape, dominated by -`VersionNumber`'s non-isbits prerelease/build tuples) to ~32 B. - -`VersionNumber` fields are packed into `UInt16` and `Union{Int, Nothing}` -hint fields use `_UNSET` (`-1`) as the `nothing` sentinel. Decoding back to -the original types happens once per cache miss in `emit_binary!` / -`emit_tile!`, never on the hot lookup path. -""" -struct TileCacheKey - sm_arch::UInt16 - bytecode_version::UInt16 - opt_level::Int - num_ctas::Int - occupancy::Int - num_worker_warps::Int -end -TileCacheKey(sm_arch::VersionNumber, bytecode_version::VersionNumber, - opt_level::Union{Int, Nothing}, num_ctas::Union{Int, Nothing}, - occupancy::Union{Int, Nothing}, num_worker_warps::Union{Int, Nothing}) = - TileCacheKey(pack_version(sm_arch), pack_version(bytecode_version), - pack_hint(opt_level), pack_hint(num_ctas), pack_hint(occupancy), - pack_hint(num_worker_warps)) - - #============================================================================= Toolchain and target validation. =============================================================================# @@ -638,58 +592,85 @@ end =============================================================================# """ - emit_binary!(cache, mi, ci, res; const_argtypes=nothing) -> Vector{UInt8} + assemble(bytecode, sm_arch, opt_level) -> Vector{UInt8} -Cached binary phase: compile Tile IR bytecode to CUBIN using tileiras. +Assemble Tile IR bytecode to a CUBIN with `tileiras`, through the disk cache. """ -function emit_binary!(cache::CacheView, mi::Core.MethodInstance, - ci::Core.CodeInstance, res::CuTileResults; - const_argtypes::Union{Vector{Any}, Nothing}=nothing) - res.cuda_bin !== nothing && return res.cuda_bin - - bytecode = emit_tile!(cache, mi, ci, res; const_argtypes) - - sm_arch = unpack_version(cache.owner.sm_arch) - - # Resolve opt_level here (not in emit_tile) because it's a tileiras flag, not bytecode. - # num_ctas/occupancy/num_worker_warps are resolved in emit_tile because they're encoded in bytecode. - _, _, kernel_meta = emit_structured!(cache, mi, ci, res; const_argtypes) - opt_level = something(resolve_hint(unpack_hint(cache.owner.opt_level), - kernel_meta, :opt_level, sm_arch), 3) - +function assemble(bytecode::Vector{UInt8}, sm_arch::VersionNumber, opt_level::Int) # Cross-session cache of the tileiras output. The key covers every input # that changes the CUBIN: bytecode, sm_arch, opt_level and the tileiras # identity, so different compiler builds never collide. `bytecode_version` # is encoded in the bytecode itself, so it's covered transitively. CUBIN is # address-free, so it is always persistable. - res.cuda_bin = if disk_cache_enabled - ObjCache.get!(CUBIN_CACHE_NS, cubin_cache_fields(bytecode, sm_arch, opt_level)...; - schema=CUBIN_CACHE_SCHEMA, persistable=true) do - first(run_tileiras(bytecode, sm_arch, opt_level)) - end - else + disk_cache_enabled || return first(run_tileiras(bytecode, sm_arch, opt_level)) + return ObjCache.get!(CUBIN_CACHE_NS, cubin_cache_fields(bytecode, sm_arch, opt_level)...; + schema=CUBIN_CACHE_SCHEMA, persistable=true) do first(run_tileiras(bytecode, sm_arch, opt_level)) end +end - return res.cuda_bin +# The compilation results of a job: its CUBIN, shared by every context, and the +# session-local kernels linked from it. +mutable struct CuTileResults + cuda_bin::Union{Nothing, Vector{UInt8}} + # linear-scanned by context; usually holds a single entry + kernels::Vector{Tuple{CuContext, Any}} + CuTileResults() = new(nothing, Tuple{CuContext, Any}[]) +end + +# Results for every configuration of a kernel, attached to its CodeInstance by +# CompilerCaching (and persisted with it into package images). One CodeInstance +# serves every target, hint and constant configuration of a kernel, so entries are +# keyed by config; the CodeInstance itself outlives any one world. +mutable struct TileResults + entries::Vector{Pair{TileConfig, CuTileResults}} + TileResults() = new(Pair{TileConfig, CuTileResults}[]) +end + +# The results struct for `job`, or `nothing` while its kernel has not been inferred. +# Once it has, an empty struct is created on first access and returned ever after. +function cached_results(job::TileJob) + ci = get(inference_cache(job), job.source, nothing) + ci === nothing && return nothing + results = CompilerCaching.results(TileResults, ci) + for (config, res) in results.entries + # configs are immutable, so `===` compares them structurally + config === job.config && return res + end + res = CuTileResults() + push!(results.entries, job.config => res) + return res end +const compile_lock = ReentrantLock() + """ - link(cache, mi, ci, res) -> CuFunction + compile_or_lookup(job::TileJob) -> CuTileResults -GPU-side link phase: load the CUBIN cached in `res.cuda_bin` (produced by -[`compile`](@ref)) onto the active CUDA device and return the resulting -`CuFunction`. +The cached compilation results for `job`, running the compiler on a miss. The +`compile_hook` check forces the compile path so that `@device_code_*` observe +the compilation even on a hit. No CUDA context required. """ -function link(cache::CacheView, mi::Core.MethodInstance, - ci::Core.CodeInstance, res::CuTileResults) - res.cuda_func !== nothing && return res.cuda_func +function compile_or_lookup(job::TileJob)::CuTileResults + Base.@lock compile_lock begin + res = cached_results(job) + if res === nothing || res.cuda_bin === nothing || compile_hook[] !== nothing + cubin = compile(job) + res = @something res cached_results(job) + res.cuda_bin = cubin + end + return res + end +end + +""" + link(job::TileJob, res::CuTileResults) -> CuFunction - kernel_name = sanitize_name(string(mi.def.name)) +Load the job's CUBIN onto the active CUDA context. +""" +function link(job::TileJob, res::CuTileResults) cumod = CuModule(res.cuda_bin::Vector{UInt8}) - cufunc = CuFunction(cumod, kernel_name) - res.cuda_func = cufunc - return cufunc + return CuFunction(cumod, job.config.name) end @@ -698,7 +679,7 @@ end Mirrors the `cufunction(f, tt) -> HostKernel` pattern in CUDACore. Once obtained, calling `(::TileKernel)(args...; blocks=…)` skips the MI lookup - and CompilerCaching dispatch — only argument flatten + `cudacall` runs. + and cache dispatch — only argument flatten + `cudacall` runs. =============================================================================# """ @@ -723,10 +704,9 @@ argument types (i.e. after `cuTileconvert`/`Adapt.adapt(KernelAdaptor(), …)`). Compilation is cached; calling `cufunction` repeatedly with the same `(f, tt, opts)` is O(1) after the first compile. -Mirrors `CUDACore.cufunction` but produces a [`TileKernel`](@ref). Caching -is delegated to CompilerCaching: the resulting `TileKernel` is stored in -the `CuTileResults` attached to the underlying Julia `CodeInstance`, so -invalidation rides on Julia's normal CI lifecycle. +Mirrors `CUDACore.cufunction` but produces a [`TileKernel`](@ref). Results are +stored on the kernel's Julia `CodeInstance`, so invalidation rides on Julia's +normal CI lifecycle. """ function cufunction(@nospecialize(f), tt::Type{<:Tuple}=Tuple{}; sm_arch::Union{VersionNumber, Nothing}=nothing, @@ -738,9 +718,6 @@ function cufunction(@nospecialize(f), tt::Type{<:Tuple}=Tuple{}; resolved_sm_arch = sm_arch !== nothing ? sm_arch : default_sm_arch() bytecode_version = check_tile_ir_support(resolved_sm_arch) - key = TileCacheKey(resolved_sm_arch, bytecode_version, opt_level, num_ctas, occupancy, - num_worker_warps) - # Single pass over `tt.parameters`: build the unwrapped argtypes tuple # (Constant{T,V} → T for MI lookup) and the const_argtypes vector # (Constant{T,V} → CC.Const(V) for inference) together. cufunction @@ -751,21 +728,15 @@ function cufunction(@nospecialize(f), tt::Type{<:Tuple}=Tuple{}; # invalidated by any package that defines methods on Base.Compiler hooks # like `OptimizationParams(::AbstractInterpreter)`. To reuse precompiled # native code, run the pipeline in the world captured at __init__. - invoke_frozen(cufunction_compile, f, tt, argtypes, const_argtypes, key)::TileKernel{Core.Typeof(f), tt} + opts = (; sm_arch=resolved_sm_arch, bytecode_version, opt_level, num_ctas, occupancy, + num_worker_warps, name) + invoke_frozen(cufunction_compile, f, tt, argtypes, const_argtypes, opts)::TileKernel{Core.Typeof(f), tt} end -""" - compile(f, argtypes, const_argtypes, key) -> (cache, mi, ci, res) - -Host-side compile phase: run inference, codegen, bytecode emission, and -`tileiras` to produce a CUBIN. Returns the compilation cache state needed -by [`link`](@ref) to load the result onto the GPU. No CUDA context required. -""" -function compile(@nospecialize(f), @nospecialize(argtypes), - const_argtypes::Union{Vector{Any}, Nothing}, - key::TileCacheKey) - validate_tile_ir_target(unpack_version(key.sm_arch), - unpack_version(key.bytecode_version)) +# The job of a launch: the kernel's MethodInstance for the unwrapped argument +# types, in the current world, with the launch's options. +function launch_job(@nospecialize(f), @nospecialize(argtypes), + const_argtypes::Union{Vector{Any}, Nothing}; kwargs...) world = Base.get_world_counter() mi = method_instance(f, argtypes; world) mi === nothing && throw(MethodError(f, argtypes)) @@ -773,41 +744,31 @@ function compile(@nospecialize(f), @nospecialize(argtypes), sig = Base.signature_type(f, argtypes) mi = CC.specialize_method(mi.def, sig, mi.sparam_vals)::Core.MethodInstance end - - cache = CacheView{CuTileResults}(key, world) - - # Single resolution of (ci, res) up front — threaded through the emit_*! - # chain so each phase only does its own short-circuit, not redundant - # cache lookups. The cached compilation results are attached to the - # underlying `CodeInstance` via CompilerCaching; the `TileKernel` wrapper - # rides along in the same `CuTileResults`, so kernel-instance lifecycle - # follows the CI's instead of needing a separate global Dict. - ci, res = ensure_compiled(cache, mi, const_argtypes) - - # Report the compilation to the `@device_code_*` hook. This runs through - # `invoke_frozen`, but the hook closure lives in the user's latest world — - # `invokelatest` it. - if compile_hook[] !== nothing - Base.invokelatest(compile_hook[], hook_signature(mi, const_argtypes)...) - end - - emit_binary!(cache, mi, ci, res; const_argtypes) - return cache, mi, ci, res + return tile_job(mi, world; const_argtypes, kwargs...) end # Inner compilation routine; called via `invoke_frozen` so its method dispatches # happen in the world captured at __init__, reusing precompiled native code # even when later-loaded packages would otherwise have invalidated it. function cufunction_compile(@nospecialize(f), @nospecialize(tt), @nospecialize(argtypes), - const_argtypes::Union{Vector{Any}, Nothing}, - key::TileCacheKey) - cache, mi, ci, res = compile(f, argtypes, const_argtypes, key) - - cufunc = link(cache, mi, ci, res) - - res.tile_kernel !== nothing && return res.tile_kernel::TileKernel{Core.Typeof(f), tt} - kernel = TileKernel{Core.Typeof(f), tt}(f, cufunc) - res.tile_kernel = kernel + const_argtypes::Union{Vector{Any}, Nothing}, opts::NamedTuple) + validate_tile_ir_target(opts.sm_arch, opts.bytecode_version) + job = launch_job(f, argtypes, const_argtypes; opts...) + res = compile_or_lookup(job) + + # Resolve the kernel for the active context. `CuFunction`s are session-local + # handles, so they live in the results struct's linear cache rather than being + # persisted; the scan is almost always over a single entry. + ctx = context() + for (cached_ctx, cached_kernel) in res.kernels + cached_ctx === ctx && return cached_kernel::TileKernel{Core.Typeof(f), tt} + end + kernel = TileKernel{Core.Typeof(f), tt}(f, link(job, res)) + # don't cache session-local handles while generating output: the results struct + # is serialized into the package image along with its CodeInstance. + if ccall(:jl_generating_output, Cint, ()) != 1 + push!(res.kernels, (ctx, kernel)) + end return kernel end diff --git a/src/precompile.jl b/src/precompile.jl index d613251b..33e4a9e5 100644 --- a/src/precompile.jl +++ b/src/precompile.jl @@ -36,22 +36,20 @@ import REPL return end - # Drive the host-side `compile` so the precompile workload reaches + # Drive the host-side `compile_or_lookup` so the precompile workload reaches # `tileiras` → CUBIN without needing a CUDA context. The runtime path # tacks `link` on top to load the CUBIN onto the GPU. function precompile_kernel(@nospecialize(f), @nospecialize(tt)) # Without `tileiras` (e.g. a pre-13.2 CUDA_Compiler_jll) kernels cannot # compile; still allow the package itself to precompile. tileiras_available() || return - argtypes, const_argtypes = unwrap_argtypes(f, tt) bv = bytecode_version() for sm_arch in [v"8.0", v"8.6", v"8.7", v"8.9", v"10.0", v"11.0", v"12.0", v"12.1"] requirement = tile_ir_requirement(sm_arch) requirement === nothing && continue bv < requirement[2] && continue - key = TileCacheKey(sm_arch, bv, nothing, nothing, nothing, nothing) - compile(f, argtypes, const_argtypes, key) + compile_or_lookup(tile_job(f, Tuple(tt.parameters); sm_arch, bytecode_version=bv)) end return end diff --git a/test/codegen/boundscheck_integration.jl b/test/codegen/boundscheck_integration.jl index 3c68a932..36de3e15 100644 --- a/test/codegen/boundscheck_integration.jl +++ b/test/codegen/boundscheck_integration.jl @@ -61,7 +61,7 @@ function test_bytecode(@nospecialize(f), @nospecialize(argtypes)) stripped, const_argtypes = ct.process_const_argtypes(f, argtypes) world = Base.get_world_counter() mi = ct.lookup_method_instance(f, stripped; world) - cache = ct.CacheView{ct.CuTileResults}(:cuTile, world) + cache = ct.inference_cache(world) ir, rettype = ct.emit_julia(cache, mi; const_argtypes) sci, rettype, kernel_meta = ct.emit_structured(ir, rettype) empty!(sci.line_map) diff --git a/test/codegen/reflection.jl b/test/codegen/reflection.jl index 95ad5b0e..4216618e 100644 --- a/test/codegen/reflection.jl +++ b/test/codegen/reflection.jl @@ -87,6 +87,42 @@ if ct.tileiras_available() @test occursin(r"\.target\s+sm_100", output) @test occursin(".text.reflect_vadd", output) end + + @testset "compile(job)" begin + job = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0") + cubin = ct.compile(job) + @test cubin isa Vector{UInt8} && cubin[1:4] == b"\x7fELF" + ptx = ct.extract_ptx(cubin) + @test startswith(ptx, ".version") && occursin(".entry reflect_vadd", ptx) + @test ptx == ct.extract_ptx(IOBuffer(cubin)) == sprint(ct.code_ptx, job) + @test_throws ct.ObjectFile.MagicMismatch ct.extract_ptx(UInt8[0x7f, 0x45, 0x4c, 0x46, 2, 1, 1, 0]) + end + + @testset "jobs" begin + job = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0") + @test job isa ct.TileJob && job.config.name == "reflect_vadd" + @test job === ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0") + @test job !== ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0", opt_level=1) + @test ct.job_signature(job) == (reflect_vadd, TT3) + @test sprint(show, only(ct.code_typed(job))) == sprint(show, only(ct.code_typed(reflect_vadd, TT3))) + end + + @testset "shared inference" begin + # Inference is independent of target and hints, so every Tile job shares + # one partition and one CodeInstance; codegen results are per job. + job1 = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0") + job2 = ct.tile_job(reflect_vadd, TT3; sm_arch=v"8.9", opt_level=1) + res1, res2 = ct.compile_or_lookup(job1), ct.compile_or_lookup(job2) + ci1, ci2 = (get(ct.inference_cache(job), job.source, nothing) for job in (job1, job2)) + @test ci1 === ci2 !== nothing + @test ci1.owner === ct.TILE_CACHE_OWNER + @test res1 !== res2 && res1.cuda_bin != res2.cuda_bin + # results are keyed by config, not job: they survive a world bump + bump() = nothing + job3 = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0") + @test job3.world > job1.world && job3.config === job1.config + @test ct.compile_or_lookup(job3) === res1 + end end if ct.tileiras_version() >= v"13.4" diff --git a/test/host/cache.jl b/test/host/cache.jl index a42eedf2..1f5f2343 100644 --- a/test/host/cache.jl +++ b/test/host/cache.jl @@ -33,11 +33,11 @@ const CUBIN_CACHE_PRELUDE = quote return end - key() = ct.TileCacheKey(SM_ARCH, ct.bytecode_version(), OPT_LEVEL, - nothing, nothing, nothing) + vadd_job() = ct.tile_job(cached_vadd, TT; sm_arch=SM_ARCH, opt_level=OPT_LEVEL) function compile_vadd() - _, _, _, res = ct.compile(cached_vadd, TT, nothing, key()) - return res + job = vadd_job() + res = ct.compile_or_lookup(job) + return (; cuda_bin=res.cuda_bin, tile_bc=ct.emit_bytecode(job).bytecode) end cubin_key(res) = OC.keyhash(ct.CUBIN_CACHE_SCHEMA, ct.cubin_cache_fields(res.tile_bc, diff --git a/test/host/inference_cache.jl b/test/host/inference_cache.jl index 98f3aca6..ff8b7ee5 100644 --- a/test/host/inference_cache.jl +++ b/test/host/inference_cache.jl @@ -1,18 +1,27 @@ @testset "Inference cache" begin choose_type(flag, x) = flag ? x + one(x) : Float32(x) world = Base.get_world_counter() - cache = ct.CacheView{ct.CuTileResults}(gensym(:inference_cache), world) + cache = ct.inference_cache(world) mi = ct.lookup_method_instance(choose_type, Tuple{Bool, Int32}; world) - ci, generic = ct.ensure_compiled(cache, mi, nothing) - @test ct.ensure_compiled(cache, mi, nothing) === (ci, generic) + ci = ct.get_ci(cache, mi) + @test ct.get_ci(cache, mi) === ci _, generic_rt = ct.get_inferred(cache, ci, mi) @test generic_rt == Union{Int32, Float32} + # Attach codegen results before looking up specialized inference source. + # Both result types live on this CI; lookup must use the inference view. + generic_job = ct.tile_job(mi, world; sm_arch=v"10.0", bytecode_version=v"13.2") + generic = ct.cached_results(generic_job) specialized = ct.CuTileResults[] for (flag, expected_rt) in ((true, Int32), (false, Float32)) argtypes = Any[ct.CC.Const(choose_type), ct.CC.Const(flag), Int32] - specialized_ci, res = ct.ensure_compiled(cache, mi, argtypes) + specialized_ci = ct.get_ci(cache, mi; const_argtypes=argtypes) + entry = ct.specialization(cache, specialized_ci, argtypes) + @test entry isa ct.CompilerCaching.SpecializedResult{Nothing} + job = ct.tile_job(mi, world; sm_arch=v"10.0", bytecode_version=v"13.2", + const_argtypes=argtypes) + res = ct.cached_results(job) @test specialized_ci === ci @test res !== generic @test all(previous -> previous !== res, specialized) @@ -22,12 +31,12 @@ ir, rt = ct.get_inferred(cache, ci, mi; const_argtypes=argtypes) @test ir isa ct.CC.IRCode @test rt === expected_rt - structured = ct.emit_structured!(cache, mi, ci, res; const_argtypes=argtypes) + structured = ct.emit_structured(ir, rt) @test structured[2] === expected_rt # An equivalent argument vector must reuse the compiled results. - @test ct.ensure_compiled(cache, mi, copy(argtypes)) === (ci, res) - @test ct.emit_structured!(cache, mi, ci, res; - const_argtypes=argtypes) === structured + @test ct.get_ci(cache, mi; const_argtypes=copy(argtypes)) === ci + @test ct.specialization(cache, ci, copy(argtypes)) === entry + @test ct.cached_results(job) === res end end From 1eca5d941c2fa5d3cd43bf8fb540c7ab6f71d70f Mon Sep 17 00:00:00 2001 From: AntonOresten Date: Fri, 4 Sep 2026 18:06:58 +0200 Subject: [PATCH 2/2] Fix TileJob target caching and reflection API --- src/compiler/reflection.jl | 6 ++++++ src/launch.jl | 4 ++++ test/codegen/reflection.jl | 25 +++++++++++++++++++++++++ 3 files changed, 35 insertions(+) diff --git a/src/compiler/reflection.jl b/src/compiler/reflection.jl index 0d4ce65c..ddd97b64 100644 --- a/src/compiler/reflection.jl +++ b/src/compiler/reflection.jl @@ -167,6 +167,7 @@ code_structured(@nospecialize(f), @nospecialize(argtypes); optimize::Bool=true, code_structured(tile_job(f, argtypes; kwargs...); optimize) """ + code_tiled([io::IO], job::TileJob; debuginfo=false, remarks=false) code_tiled([io::IO], f, argtypes; sm_arch, opt_level, num_ctas, occupancy, num_worker_warps, remarks=false) @@ -200,6 +201,7 @@ end code_tiled(io::IO, @nospecialize(f), @nospecialize(argtypes); debuginfo::Bool=false, remarks::Bool=false, kwargs...) = code_tiled(io, tile_job(f, argtypes; kwargs...); debuginfo, remarks) +code_tiled(job::TileJob; kwargs...) = code_tiled(stdout, job; kwargs...) code_tiled(@nospecialize(f), @nospecialize(argtypes); kwargs...) = code_tiled(stdout, f, argtypes; kwargs...) @@ -224,6 +226,7 @@ function compile(job::TileJob) end """ + code_ptx([io::IO], job::TileJob) code_ptx([io::IO], f, argtypes; sm_arch, opt_level, num_ctas, occupancy, num_worker_warps) @@ -239,10 +242,12 @@ already made. When no GPU is available, pass `sm_arch` explicitly. code_ptx(io::IO, job::TileJob) = print(io, extract_ptx(compile(job))) code_ptx(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) = code_ptx(io, tile_job(f, argtypes; kwargs...)) +code_ptx(job::TileJob) = code_ptx(stdout, job) code_ptx(@nospecialize(f), @nospecialize(argtypes); kwargs...) = code_ptx(stdout, f, argtypes; kwargs...) """ + code_sass([io::IO], job::TileJob) code_sass([io::IO], f, argtypes; sm_arch, opt_level, num_ctas, occupancy, num_worker_warps) @@ -254,6 +259,7 @@ binary a launch actually loaded, use `CUDA.@device_code_sass`. code_sass(io::IO, job::TileJob) = print(io, disassemble_cubin(compile(job))) code_sass(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) = code_sass(io, tile_job(f, argtypes; kwargs...)) +code_sass(job::TileJob) = code_sass(stdout, job) code_sass(@nospecialize(f), @nospecialize(argtypes); kwargs...) = code_sass(stdout, f, argtypes; kwargs...) diff --git a/src/launch.jl b/src/launch.jl index b1908d58..da057b5b 100644 --- a/src/launch.jl +++ b/src/launch.jl @@ -652,6 +652,10 @@ The cached compilation results for `job`, running the compiler on a miss. The the compilation even on a hit. No CUDA context required. """ function compile_or_lookup(job::TileJob)::CuTileResults + # A targetless job is useful for pre-tileiras reflection, but a cached CUBIN + # must be identified by the architecture it was assembled for. Resolve the + # target before looking up results so it becomes part of the cache identity. + job = with_target(job, resolve_target(job, "compiling")) Base.@lock compile_lock begin res = cached_results(job) if res === nothing || res.cuda_bin === nothing || compile_hook[] !== nothing diff --git a/test/codegen/reflection.jl b/test/codegen/reflection.jl index 4216618e..ae2740c5 100644 --- a/test/codegen/reflection.jl +++ b/test/codegen/reflection.jl @@ -9,6 +9,16 @@ function reflect_vadd(a, b, c) return end +function capture_stdout(f) + mktemp() do _, io + redirect_stdout(io) do + f() + end + seekstart(io) + read(io, String) + end +end + @testset "code_typed" begin @test @filecheck begin @check "get_tile_block_id" @@ -86,6 +96,9 @@ if ct.tileiras_available() output = sprint(io -> ct.code_sass(io, reflect_vadd, TT3; sm_arch=v"10.0")) @test occursin(r"\.target\s+sm_100", output) @test occursin(".text.reflect_vadd", output) + job = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0") + stdout_output = capture_stdout(() -> ct.code_sass(job)) + @test stdout_output == output end @testset "compile(job)" begin @@ -95,6 +108,8 @@ if ct.tileiras_available() ptx = ct.extract_ptx(cubin) @test startswith(ptx, ".version") && occursin(".entry reflect_vadd", ptx) @test ptx == ct.extract_ptx(IOBuffer(cubin)) == sprint(ct.code_ptx, job) + stdout_ptx = capture_stdout(() -> ct.code_ptx(job)) + @test stdout_ptx == ptx @test_throws ct.ObjectFile.MagicMismatch ct.extract_ptx(UInt8[0x7f, 0x45, 0x4c, 0x46, 2, 1, 1, 0]) end @@ -105,6 +120,8 @@ if ct.tileiras_available() @test job !== ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0", opt_level=1) @test ct.job_signature(job) == (reflect_vadd, TT3) @test sprint(show, only(ct.code_typed(job))) == sprint(show, only(ct.code_typed(reflect_vadd, TT3))) + stdout_tiled = capture_stdout(() -> ct.code_tiled(job)) + @test stdout_tiled == sprint(ct.code_tiled, job) end @testset "shared inference" begin @@ -122,6 +139,14 @@ if ct.tileiras_available() job3 = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0") @test job3.world > job1.world && job3.config === job1.config @test ct.compile_or_lookup(job3) === res1 + + # A targetless cached job is normalized before lookup, so its CUBIN is + # stored under the architecture it was assembled for. + if CUDA.functional() && !isempty(CUDA.devices()) + targetless = ct.tile_job(reflect_vadd, TT3) + targeted = ct.with_target(targetless, ct.default_sm_arch()) + @test ct.compile_or_lookup(targetless) === ct.compile_or_lookup(targeted) + end end end