Skip to content

Make the job the cache identity #308

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
AntonOresten wants to merge 2 commits into from
Closed

Make the job the cache identity #308

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bc02d87
Make the job the cache identity
AntonOresten Aug 29, 2026
1eca5d9
Fix TileJob target caching and reflection API
AntonOresten Sep 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions docs/src/lib/reflection.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,13 @@ code_ptx
code_sass
```

## Jobs

```@docs
TileJob
tile_job
```

## Intercepting a launch

```@docs
Expand Down
2 changes: 1 addition & 1 deletion src/compiler/codegen/kernel.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
255 changes: 123 additions & 132 deletions src/compiler/driver.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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."
Expand All @@ -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

Expand Down Expand Up @@ -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/<file>.ln<line>[.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
Loading