diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index afe871a8..7134974e 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -53,7 +53,6 @@ steps: - "13.4" - "13.3" - "13.2" - - "13.1" commands: | echo -e "[CUDA_Compiler_jll]\nversion = \"{{matrix.cuda}}\"" >LocalPreferences.toml diff --git a/Project.toml b/Project.toml index 9429a68e..b9260424 100644 --- a/Project.toml +++ b/Project.toml @@ -15,8 +15,10 @@ CUDA_Tile_jll = "2068806d-a867-5dbd-af0e-42c2eb5d895d" CompilerCaching = "9db33cc3-5358-4881-8759-fa4194144afd" EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56" GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" +GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55" GPUToolbox = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc" IRStructurizer = "93e32bba-5bb8-402b-805d-ffb066edee93" +ObjectFile = "d8793406-e978-5875-9003-1fc021f44a92" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Preferences = "21216c6a-2e73-6563-6e65-726566657250" REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb" @@ -40,9 +42,11 @@ CompilerCaching = "0.5" DLFP8Types = "0.1.0" EnumX = "1.0" GPUArrays = "11" +GPUCompiler = "2.6" GPUToolbox = "3" IRStructurizer = "0.6.3" Microfloats = "0.2" +ObjectFile = "0.5" PrecompileTools = "1" Preferences = "1.5" REPL = "1" diff --git a/docs/src/lib/reflection.md b/docs/src/lib/reflection.md index fc478368..8ad493b3 100644 --- a/docs/src/lib/reflection.md +++ b/docs/src/lib/reflection.md @@ -11,12 +11,28 @@ code_tiled code_typed code_ircode code_structured +code_ptx +code_sass +``` + +## Jobs + +```@docs +TileJob +tile_job ``` ## Intercepting a launch ```@docs @device_code_tiled -@device_code_typed @device_code_structured +@device_code_ptx ``` + +`@device_code_typed` and `@device_code_warntype` are shared with GPUCompiler; +they work for cuTile kernels through the shared compile hook. + +`CUDA.@device_code_sass` works for cuTile kernels: it intercepts module loads +at the driver level (via CUPTI), so it captures any backend's kernels without +backend-specific support. diff --git a/docs/src/man/debugging.md b/docs/src/man/debugging.md index 27b7557b..d99ae32a 100644 --- a/docs/src/man/debugging.md +++ b/docs/src/man/debugging.md @@ -36,6 +36,23 @@ Spelling out those types is only worth it when you have no GPU, since `code_tiled` does not need CUDA.jl. Otherwise let the launch site derive them for you with `ct.@device_code_tiled`, described next. +By default `code_tiled` compiles for the active CUDA device with the +toolchain's bytecode version, so its output matches what a launch emits, +including architecture-dependent `ct.@compiler_options` hints. Pass `sm_arch` +to target another architecture, or to inspect code without a CUDA device, and +`bytecode_version` to emit an older version. Combinations the bytecode does +not support are rejected, so an older `bytecode_version` may need an explicit +`sm_arch` on older devices. + +To inspect several stages with the same compilation options, create a job: + +```julia +job = ct.tile_job(vadd, argtypes; sm_arch=v"10.0", num_ctas=2) +ct.code_typed(job) +ct.code_structured(job) +ct.code_tiled(job) +``` + With a CUDA 13.4 or newer `tileiras`, pass `remarks=true` to compile the Tile IR and print optimization diagnostics after it: @@ -44,9 +61,35 @@ ct.code_tiled(vadd, argtypes; remarks=true, sm_arch=v"10.0") ``` The remarks report successful and failed optimizations, including tensor-core -selection and memory alignment issues. `sm_arch` may be omitted when a CUDA -device is available. Remarks are generated afresh for reflection and are not -read from or written to the compilation cache. +selection and memory alignment issues. Remarks are generated afresh for +reflection and are not read from or written to the compilation cache. + + +## Inspecting PTX and SASS + +The output of `tileiras` can be inspected with `ct.code_ptx` and +`ct.code_sass`, taking the same signature as `ct.code_tiled`: + +```julia +ct.code_ptx(vadd, argtypes; sm_arch=v"10.0") +ct.code_sass(vadd, argtypes; sm_arch=v"10.0") +``` + +`code_ptx` shows the thread-level SIMT program the tile-level kernel is +lowered to, with every compiler decision (thread mapping, CTA size, +pipelining, synchronization) already made; `code_sass` shows the final machine +code, disassembled with `nvdisasm`. Both compile the Tile IR with `tileiras` +but do not need a GPU. + +!!! warning "PTX reflection is unstable" + PTX is an implementation detail of `tileiras`, not an interface, and may + stop being produced or recorded at any time. `code_ptx` reads it from an + undocumented debug section of the CUBIN and will be removed with it. + +Note that `tileiras` always generates architecture-specific code: targeting +`sm_arch=v"10.0"` produces `sm_100a` PTX and SASS, which only runs on that +exact architecture. Launches therefore always compile for the active device, +and reject an explicit `sm_arch` that differs from it. ## Intercepting a launch @@ -73,17 +116,37 @@ intercepted kernel: ct.@device_code_tiled remarks=true @cuda backend=cuTile blocks=grid kernel(args...) ``` -Three are available, corresponding to successive stages of the pipeline: +The following macros are available: | Macro | Output | |-------|--------| -| `ct.@device_code_tiled` | Final Tile IR (MLIR textual format) | -| `ct.@device_code_typed` | Typed Julia IR after overlay resolution | +| `ct.@device_code_warntype` | Typed Julia IR with type-instability highlighting | +| `ct.@device_code_typed` | Typed Julia IR after overlay resolution, returned per job | | `ct.@device_code_structured` | Structured IR (after control-flow structurization) | - -Read top to bottom, they run backwards through the pipeline: Julia IR with Tile -IR intrinsics substituted by the overlay method table, then structured control -flow, then emitted Tile IR. +| `ct.@device_code_tiled` | Final Tile IR (MLIR textual format) | +| `ct.@device_code_ptx` | PTX generated by `tileiras` (unstable, as above) | + +The first two are shared with GPUCompiler: cuTile reports its +compilations through the same hook as the LLVM backends, so +`CUDA.@device_code_typed` and `CUDA.@device_code_warntype` cover cuTile +launches as well, including expressions that launch kernels of both kinds. +`@device_code_typed` returns a dictionary +from job to typed code rather than printing. +Both Julia-level views use cuTile's optimized inferred code, including constant +argument specializations. `CUDA.@device_code_ptx` also includes cuTile launches. +The cuTile-specific stage macros only inspect cuTile jobs and ignore other backends; +they report an error when the expression contains no cuTile compilations. + +The macros use the launch's actual target, constant arguments, and compilation +hints, including when the compiled kernel is already cached. Each distinct +job is inspected once per macro invocation. Hooks are scoped to the current task +and inherited by child tasks; nested macros restore the enclosing hook. + +For the machine code, `CUDA.@device_code_sass` works for cuTile kernels: it +intercepts module loads at the driver level (via CUPTI), showing the binary +that was actually loaded — complete with Julia source locations, since +`tileiras` compiles with line info. Under Nsight or another active profiler, +which conflicts with CUPTI, use `ct.code_sass` on the signature instead. ## Dumping bytecode diff --git a/src/compiler/codegen/kernel.jl b/src/compiler/codegen/kernel.jl index 9b7b21ca..f99361c4 100644 --- a/src/compiler/codegen/kernel.jl +++ b/src/compiler/codegen/kernel.jl @@ -387,15 +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 - # Suppress compile_hook to avoid @device_code_tiled treating - # region bodies (e.g. reduce combiners) as standalone entries. - old_hook = compile_hook[] - compile_hook[] = nothing - sci, _, _ = try - emit_structured!(ctx.cache, mi) - finally - compile_hook[] = old_hook - end + 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 16aa4a7e..958df1a2 100644 --- a/src/compiler/driver.jl +++ b/src/compiler/driver.jl @@ -1,145 +1,218 @@ -# Compilation hook for @device_code_* macros - intercepts compilations for reflection -const compile_hook = Ref{Union{Nothing,Function}}(nothing) - - #============================================================================= - Meta nodes and compilation hints + Compiler jobs + + TileJob identifies the source and configuration used by compilation, + caching, and reflection. =============================================================================# -# Compilation options for cache sharding. -# 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." -const CGOpts = @NamedTuple{ - sm_arch::Union{VersionNumber, Nothing}, - opt_level::Union{Int, Nothing}, - num_ctas::Union{Int, Nothing}, - occupancy::Union{Int, Nothing}, - num_worker_warps::Union{Int, Nothing}, +public TileJob, tile_job + +struct TileCompilerTarget + sm_arch::VersionNumber bytecode_version::VersionNumber -} +end + +# Compilation hints; `nothing` defers to the kernel's `@compiler_options`. +struct TileCompilerParams + opt_level::Union{Int, Nothing} + num_ctas::Union{Int, Nothing} + occupancy::Union{Int, Nothing} + num_worker_warps::Union{Int, Nothing} +end + +# Compilation results are keyed by config within each inference result. +struct TileConfig + target::TileCompilerTarget + params::TileCompilerParams + name::String +end """ - process_meta!(ir::CC.IRCode) -> ir + TileJob -Move `:meta` expression nodes from `ir.stmts` into `ir.meta`, mirroring -Julia's `process_meta!` in `Compiler/src/optimize.jl`. This normalizes IR -from `inflate_ir` (which leaves meta as stmts) to match the `typeinf_ircode` -path (which already extracts meta via `convert_to_ircode`). +A compilation request containing a `MethodInstance`, its world age and any +constant arguments, plus the target architecture, bytecode version, compilation +hints, and kernel name. Construct jobs with [`tile_job`](@ref) and pass them to +the reflection functions to inspect that configuration. + +Jobs are immutable and compare structurally, so equal jobs are `===`. """ -function process_meta!(ir::CC.IRCode) - for i in 1:length(ir.stmts) - stmt = ir.stmts[i][:stmt] - if stmt isa Expr && stmt.head === :meta - push!(ir.meta, stmt) - @static if VERSION >= v"1.12-" - ir.stmts[i][:stmt] = nothing - else - CC.setindex!(ir.stmts[i], nothing, :stmt) - end - end - end - return ir +struct TileJob + source::MethodInstance + # `(Const(f), arg2, …)` seeding const-propagating inference, or `nothing` + # for the generic inferred source. A Tuple rather than CompilerCaching's + # Vector so that jobs compare structurally. + const_argtypes::Union{Tuple, Nothing} + world::UInt + config::TileConfig end """ - extract_meta(ir::CC.IRCode) -> Dict{Symbol, Any} - -Extract cuTile meta nodes from IRCode. Meta nodes are inserted by `@compiler_options` -and survive through lowering/optimization. After `process_meta!` normalization, -all meta nodes reside in `ir.meta`. + tile_job(f, argtypes; world=Base.get_world_counter(), kwargs...) -> TileJob + tile_job(mi::MethodInstance, world; const_argtypes=nothing, kwargs...) -> TileJob + +Create a [`TileJob`](@ref) for `f` and `argtypes` (which may contain +`Constant{T,V}` types), or for a method instance with the given const-seeded +argument types. Keyword arguments configure the compilation: + +- `bytecode_version`: the Tile IR bytecode version to emit; defaults to + [`bytecode_version`](@ref)`()`, i.e. the `bytecode_version` preference or + the newest version the selected `tileiras` accepts. +- `sm_arch`: the target architecture; defaults to the current CUDA device's. + Combinations the bytecode does not support are rejected. +- `opt_level`, `num_ctas`, `occupancy`, `num_worker_warps`: compilation hints + overriding the kernel's `@compiler_options`. +- `name`: the kernel's name in the bytecode; defaults to the method's. """ -function extract_meta(ir::CC.IRCode) - meta = Dict{Symbol, Any}() - for expr in ir.meta - if expr isa Expr && expr.head === :meta && length(expr.args) >= 3 && expr.args[1] === :cuTile - meta[expr.args[2]::Symbol] = expr.args[3] - end +function tile_job(mi::MethodInstance, world::UInt; + const_argtypes::Union{Tuple, Nothing}=nothing, + sm_arch::Union{VersionNumber, Nothing}=nothing, + bytecode_version::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, + name::Union{String, Nothing}=nothing) + if sm_arch === nothing + CUDACore.functional() || throw(ArgumentError( + "No CUDA device to target; pass `sm_arch`")) + sm_arch = device_sm_arch() end - return meta + bytecode_version = @something bytecode_version cuTile.bytecode_version() + validate_tile_ir_target(sm_arch, bytecode_version) + target = TileCompilerTarget(sm_arch, bytecode_version) + params = TileCompilerParams(opt_level, num_ctas, occupancy, num_worker_warps) + config = TileConfig(target, params, @something name sanitize_name(string(mi.def.name))) + return TileJob(mi, const_argtypes, world, config) end -""" - resolve_hint(explicit, kernel_meta, key, sm_arch) +# The job's const-seeded argument types in CompilerCaching's `Vector{Any}` form. +const_argtypes_vector(job::TileJob) = + job.const_argtypes === nothing ? nothing : collect(Any, job.const_argtypes) + +# `(f, tt)` with `Constant` argument types restored. +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]) + if job.const_argtypes !== nothing + # const_argtypes is (Const(f), arg2, ...); arg_types omits f. + for i in eachindex(arg_types) + cat = job.const_argtypes[i+1] + cat isa CC.Const && (arg_types[i] = typeof(Constant(cat.val))) + end + end + return f, Tuple{arg_types...} +end -Resolve a hint value with precedence: explicit kwarg > @compiler_options meta > nothing. -""" -function resolve_hint(explicit, kernel_meta::Dict{Symbol, Any}, key::Symbol, - sm_arch::Union{VersionNumber, Nothing}) - val = if explicit !== nothing - explicit - elseif haskey(kernel_meta, key) && sm_arch !== nothing - resolve(kernel_meta[key], sm_arch) - else - nothing +function Base.show(io::IO, job::TileJob) + f, tt = job_signature(job) + (; target, params, name) = job.config + print(io, "TileJob(", f, "(", join(tt.parameters, ", "), ")") + print(io, "; sm_arch=", target.sm_arch, + ", bytecode_version=v\"", target.bytecode_version, "\"") + for field in fieldnames(TileCompilerParams) + hint = getfield(params, field) + hint === nothing || print(io, ", ", field, "=", hint) end - validate_hint(key, val) - return val + name == sanitize_name(string(job.source.def.name)) || print(io, ", name=", repr(name)) + print(io, ")") +end + + +#============================================================================= + Compilation hook + + `@device_code_*` macros, cuTile's and GPUCompiler's alike, observe + compilations through `GPUCompiler.compile_hook`, called with the job of every + kernel that is compiled or launched while it is set. +=============================================================================# + +# Launches run in the frozen world (`invoke_frozen`), but the hook closure +# lives in the user's latest one, hence `invokelatest`. +function run_compile_hook(job::TileJob) + hook = GPUCompiler.compile_hook[] + hook === nothing || Base.invokelatest(hook, job) + return end #============================================================================= - Compilation phases + Inference + + Inference is shared across targets and hints. Constant arguments select a + SpecializedResult on the generic CodeInstance. CompilerCaching attaches + JobResults to the corresponding inference result. =============================================================================# +const TILE_CACHE_OWNER = :cuTile + +inference_cache(world::UInt) = CacheView{JobResults}(TILE_CACHE_OWNER, world) +inference_cache(job::TileJob) = inference_cache(job.world) + """ - get_ci(cache, mi; const_argtypes=nothing) -> CodeInstance + infer(cache, mi) -> CodeInstance + infer(cache, mi, argtypes::Vector{Any}) -> SpecializedResult + infer(job::TileJob) -> Union{CodeInstance, SpecializedResult} -Ensure inference is done and return the CodeInstance. Runs `typeinf!` which is -a no-op when already cached. When `const_argtypes` is provided, also ensures -the const-specialized entry exists. +The inference result of a method instance in `cache`, or of a job: its +CodeInstance, or the const-seeded entry on it for `argtypes`. Runs inference on +a miss. """ -function get_ci(cache::CacheView, mi::Core.MethodInstance; - const_argtypes::Union{Vector{Any}, Nothing}=nothing) - # Ensure CI exists +function infer(cache::CacheView, mi::MethodInstance) ci = get(cache, mi, nothing) - if ci === nothing - interp = cuTileInterpreter(cache) - ci = typeinf!(interp, mi) - ci === nothing && error("Inference failed for $mi") - end - - # Run const-prop inference, if needed - if const_argtypes !== nothing - interp = cuTileInterpreter(cache) - typeinf!(cache, interp, mi, const_argtypes) - end - + ci === nothing || return ci + ci = typeinf!(cuTileInterpreter(cache), mi) + ci === nothing && error("Inference failed for $mi") return ci end +function infer(cache::CacheView, mi::MethodInstance, argtypes::Vector{Any}) + ci = infer(cache, mi) + entry = specialization(cache, ci, argtypes) + entry === nothing || return entry + entry = typeinf!(cache, cuTileInterpreter(cache), mi, argtypes) + entry === nothing && error("Inference failed for $mi on $argtypes") + return entry +end -# Get the inferred source and return type from a CodeInstance. -function get_inferred(cache::CacheView{K,V}, ci::Core.CodeInstance, - mi::Core.MethodInstance; const_argtypes::Union{Vector{Any}, - Nothing}=nothing) where {K,V} - rettype = CC.widenconst(ci.rettype) - if const_argtypes === nothing - src = @something get_source(ci) - else - entry = @something specialization(cache, ci, const_argtypes) - src = @something get_source(entry) - rettype = CC.widenconst(entry.rettype) - end - ir = CC.inflate_ir(src, mi) - return ir, rettype +function infer(job::TileJob) + cache = inference_cache(job) + job.const_argtypes === nothing && return infer(cache, job.source) + return infer(cache, job.source, const_argtypes_vector(job)) end +inferred_rettype(ci::Core.CodeInstance) = CC.widenconst(ci.rettype) +inferred_rettype(entry::SpecializedResult) = CC.widenconst(entry.rettype) + + +#============================================================================= + Stages + + emit_tile(job) runs inference, structurization, and bytecode generation. + compile(job) assembles the result with tileiras. Only inference is cached + within these stages; launch.jl caches the resulting CUBIN. +=============================================================================# + """ - emit_julia(cache, mi; const_argtypes=nothing) -> (IRCode, rettype) + emit_julia(job::TileJob) -> (IRCode, rettype) + emit_julia(cache, mi::MethodInstance) -> (IRCode, rettype) -Julia phase: run inference and return IRCode. +Julia phase: the inferred, optimized IR of a job, or of a callee (subprogram) +method instance in `cache`. """ -function emit_julia(cache::CacheView, mi::Core.MethodInstance; - const_argtypes::Union{Vector{Any}, Nothing}=nothing) - ci = get_ci(cache, mi; const_argtypes) - get_inferred(cache, ci, mi; const_argtypes) +function emit_julia(mi::MethodInstance, inferred) + src = @something get_source(inferred) error("No inferred source for $mi") + return CC.inflate_ir(src, mi), inferred_rettype(inferred) end +emit_julia(job::TileJob) = emit_julia(job.source, infer(job)) +emit_julia(cache::CacheView, mi::MethodInstance) = emit_julia(mi, infer(cache, mi)) """ emit_structured(ir::IRCode, rettype) -> (StructuredIRCode, rettype, kernel_meta) -Structurize IRCode into StructuredIRCode. Pure transformation, no caching. +Structurize IRCode into StructuredIRCode. """ function emit_structured(ir::CC.IRCode, rettype) process_meta!(ir) @@ -149,177 +222,126 @@ function emit_structured(ir::CC.IRCode, rettype) end """ - emit_tile(sci, rettype, kernel_meta; name, opts, cache, const_argtypes) -> Vector{UInt8} + emit_tile(job::TileJob, sci, rettype, kernel_meta) -> (; bytecode, opt_level) + emit_tile(job::TileJob) -> (; bytecode, opt_level) -Generate Tile IR bytecode from StructuredIRCode. Pure computation, no caching. -`cache` is needed for subprogram compilation inside `emit_kernel!`. +Tile IR phase: generate bytecode from StructuredIRCode, resolving the job's +hints against the kernel's `@compiler_options`. Also returns the resolved +`tileiras` optimization level, which is a flag to the assembler rather than +part of the bytecode. The one-argument form runs the preceding stages first. """ -function emit_tile(sci::StructuredIRCode, rettype, kernel_meta::Dict{Symbol,Any}; - name::String, - opts::CGOpts, - cache::CacheView, - const_argtypes::Union{Vector{Any}, Nothing}=nothing) - # Resolve hints: launch()/code_tiled() kwargs > @compiler_options meta > defaults - resolved_num_ctas = resolve_hint(opts.num_ctas, kernel_meta, :num_ctas, opts.sm_arch) - resolved_occupancy = resolve_hint(opts.occupancy, kernel_meta, :occupancy, opts.sm_arch) - resolved_num_worker_warps = resolve_hint(opts.num_worker_warps, kernel_meta, - :num_worker_warps, opts.sm_arch) - - # Generate Tile IR bytecode - bytecode = write_bytecode!(1; version=opts.bytecode_version) do writer, func_buf +function emit_tile(job::TileJob, sci::StructuredIRCode, rettype, kernel_meta::Dict{Symbol,Any}) + (; target, params, name) = job.config + hint(key, explicit) = resolve_hint(explicit, kernel_meta, key, target.sm_arch) + num_ctas = hint(:num_ctas, params.num_ctas) + occupancy = hint(:occupancy, params.occupancy) + num_worker_warps = hint(:num_worker_warps, params.num_worker_warps) + opt_level = something(hint(:opt_level, params.opt_level), 3) + + bytecode = write_bytecode!(1; version=target.bytecode_version) do writer, func_buf emit_kernel!(writer, func_buf, sci, rettype; - name, - sm_arch = opts.sm_arch, - num_ctas = resolved_num_ctas, - occupancy = resolved_occupancy, - num_worker_warps = resolved_num_worker_warps, - cache, - const_argtypes - ) + name, sm_arch=target.sm_arch, num_ctas, occupancy, num_worker_warps, + cache=inference_cache(job), + const_argtypes=const_argtypes_vector(job)) end - - return bytecode + return (; bytecode, opt_level) 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) +function emit_tile(job::TileJob) + ir, rettype = emit_julia(job) + sci, rettype, kernel_meta = emit_structured(ir, rettype) + return emit_tile(job, sci, rettype, kernel_meta) end """ - ensure_compiled(cache, mi, const_argtypes) -> (ci::CodeInstance, res::CuTileResults) + compile(job::TileJob) -> Vector{UInt8} -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. +Compile a job to a CUBIN: Tile IR bytecode, assembled with `tileiras` (through +the object cache). Reports the job to the `@device_code_*` hook. Uncached; +`compile_or_lookup` caches for launches. """ -@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 +function compile(job::TileJob) + (; sm_arch, bytecode_version) = job.config.target + validate_tileiras_target(bytecode_version) + run_compile_hook(job) + (; bytecode, opt_level) = emit_tile(job) + dump_bytecode(job.source, bytecode) + return assemble(bytecode, sm_arch, opt_level) end -# Cached wrappers around the driver's emit_* functions. These check/populate -# CuTileResults and are used by the production pipeline (launch → CUDAExt). -# Reflection APIs (code_tiled, code_structured, etc.) bypass this layer and -# call the driver directly, so they never pollute the compilation cache. -""" - emit_structured!(cache, mi, ci, res; const_argtypes=nothing) -> (StructuredIRCode, rettype, kernel_meta) +#============================================================================= + Meta nodes and compilation hints +=============================================================================# -Cached IR phase. Invokes the compile hook (for `@device_code_*` macros), -checks `res.julia_ir`, and delegates to `emit_structured` on cache miss. +""" + process_meta!(ir::CC.IRCode) -> ir -`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. +Move `:meta` expression nodes from `ir.stmts` into `ir.meta`, mirroring +Julia's `process_meta!` in `Compiler/src/optimize.jl`. This normalizes IR +from `inflate_ir` (which leaves meta as stmts) to match the `typeinf_ircode` +path (which already extracts meta via `convert_to_ircode`). """ -function emit_structured!(cache::CacheView, mi::Core.MethodInstance, - ci::Core.CodeInstance, res::CuTileResults; - const_argtypes::Union{Vector{Any}, Nothing}=nothing) - # Invoke compile hook if set (for @device_code_* reflection). - # Pass (f, tt) tuple to enable direct use with reflection utilities. - # Reconstruct Constant{T,V} types from const_argtypes so that code_tiled - # can recover const-seeded arguments (MI specTypes have them unwrapped). - if compile_hook[] !== 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, - # so arg_types[i] corresponds to const_argtypes[i+1]. - for i in eachindex(arg_types) - if const_argtypes[i+1] isa CC.Const - val = const_argtypes[i+1].val - arg_types[i] = typeof(Constant(val)) - end +function process_meta!(ir::CC.IRCode) + for i in 1:length(ir.stmts) + stmt = ir.stmts[i][:stmt] + if stmt isa Expr && stmt.head === :meta + push!(ir.meta, stmt) + @static if VERSION >= v"1.12-" + ir.stmts[i][:stmt] = nothing + else + CC.setindex!(ir.stmts[i], nothing, :stmt) end end - tt = Tuple{arg_types...} - # `cufunction` runs the codegen pipeline through `invoke_frozen` (so it - # can reuse precompiled native code), but the hook closure was defined - # at the user's latest world — invoke it via `invokelatest` so it - # dispatches there and not in the frozen world. - Base.invokelatest(compile_hook[], f, tt) end - - 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 + return ir 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) +""" + extract_meta(ir::CC.IRCode) -> Dict{Symbol, Any} + +Extract cuTile meta nodes from IRCode. Meta nodes are inserted by `@compiler_options` +and survive through lowering/optimization. After `process_meta!` normalization, +all meta nodes reside in `ir.meta`. +""" +function extract_meta(ir::CC.IRCode) + meta = Dict{Symbol, Any}() + for expr in ir.meta + if expr isa Expr && expr.head === :meta && length(expr.args) >= 3 && expr.args[1] === :cuTile + meta[expr.args[2]::Symbol] = expr.args[3] + end + end + return meta end """ - emit_tile!(cache, mi, ci, res; const_argtypes=nothing) -> Vector{UInt8} + resolve_hint(explicit, kernel_meta, key, sm_arch) -Cached code phase. Delegates to `emit_structured!` for the IR phase, checks -`res.tile_bc`, and calls the driver's `emit_tile` on cache miss. +Resolve a hint value with precedence: explicit kwarg > @compiler_options meta > nothing. +Meta hints depend on the architecture, so they are skipped without one. """ -function emit_tile!(cache::CacheView, mi::Core.MethodInstance, - ci::Core.CodeInstance, res::CuTileResults; - const_argtypes::Union{Vector{Any}, Nothing}=nothing) - # Delegate to cached IR phase — this also fires `compile_hook` for - # `@device_code_*` reflection, which must run on every launch. - ir_result = emit_structured!(cache, mi, ci, res; const_argtypes) - - res.tile_bc !== nothing && return res.tile_bc - - # Compute bytecode via driver - sci, rettype, kernel_meta = ir_result - 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))) - bytecode = emit_tile(sci, rettype, kernel_meta; - name=sanitize_name(string(mi.def.name)), - opts, cache, const_argtypes) - - # Dump bytecode if JULIA_CUTILE_DUMP_BYTECODE is set +function resolve_hint(explicit, kernel_meta::Dict{Symbol, Any}, key::Symbol, + sm_arch::Union{VersionNumber, Nothing}) + val = if explicit !== nothing + explicit + elseif haskey(kernel_meta, key) && sm_arch !== nothing + resolve(kernel_meta[key], sm_arch) + else + nothing + end + validate_hint(key, val) + return val +end + +# Dump bytecode to `$JULIA_CUTILE_DUMP_BYTECODE/.ln[.n].cutile`, if set. +const bytecode_dump_lock = ReentrantLock() + +function dump_bytecode(mi::MethodInstance, bytecode::Vector{UInt8}) dump_dir = get(ENV, "JULIA_CUTILE_DUMP_BYTECODE", nothing) - if dump_dir !== nothing + dump_dir === nothing && return + Base.@lock bytecode_dump_lock begin mkpath(dump_dir) - base_filename = basename(string(mi.def.file)) - base_filename = first(splitext(base_filename)) + 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) @@ -329,7 +351,5 @@ function emit_tile!(cache::CacheView, mi::Core.MethodInstance, println(stderr, "Dumping TILEIR bytecode to file: $dump_path") write(dump_path, bytecode) end - - res.tile_bc = bytecode - return bytecode + return end diff --git a/src/compiler/reflection.jl b/src/compiler/reflection.jl index b54287f3..bffd75d9 100644 --- a/src/compiler/reflection.jl +++ b/src/compiler/reflection.jl @@ -4,6 +4,7 @@ export code_tiled public code_typed, code_ircode, code_structured +public code_ptx, code_sass function disassemble_tileir(bytecode::Vector{UInt8}, version::VersionNumber; debuginfo::Bool=false)::String @@ -20,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) @@ -42,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 @@ -54,31 +40,12 @@ 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) -Split `Constant{T,V}` types from argtypes for method lookup, and build a -`const_argtypes` vector with `CC.Const(V)` entries for const-seeded inference. +Split `Constant{T,V}` types from argtypes for method lookup, and build the +`(Const(f), args...)` tuple with `CC.Const(V)` entries seeding const-prop +inference (see `TileJob`). Returns `(stripped, nothing)` when no Constant types are present. """ @@ -90,78 +57,120 @@ function process_const_argtypes(@nospecialize(f), @nospecialize(argtypes)) T <: Constant ? constant_eltype(T) : T end stripped = Tuple{stripped_params...} - const_argtypes = if has_consts - cats = Any[CC.Const(f)] - for T in params - if T <: Constant - push!(cats, CC.Const(constant_value(T))) - elseif CC.isconstType(T) - push!(cats, CC.Const(T.parameters[1])) - else - push!(cats, T) - end + has_consts || return stripped, nothing + const_argtypes = map(params) do T + if T <: Constant + CC.Const(constant_value(T)) + elseif CC.isconstType(T) + CC.Const(T.parameters[1]) + else + T end - cats - else - nothing end - return stripped, const_argtypes + return stripped, (CC.Const(f), const_argtypes...) end constant_eltype(::Type{Constant{T,V}}) where {T,V} = T constant_value(::Type{Constant{T,V}}) where {T,V} = V +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 + +# The Julia-level stages need only inference: without target keywords, they +# skip the job (and the toolchain probe behind its default bytecode version). +function emit_julia(@nospecialize(f), @nospecialize(argtypes); world::UInt) + stripped, const_argtypes = process_const_argtypes(f, argtypes) + mi = lookup_method_instance(f, stripped; world) + cache = inference_cache(world) + inferred = const_argtypes === nothing ? infer(cache, mi) : + infer(cache, mi, collect(Any, const_argtypes)) + return emit_julia(mi, inferred) +end + + +#============================================================================= + Stages +=============================================================================# + +""" + code_typed(job::TileJob) -> Vector{Pair{IRCode, DataType}} + code_typed(f, argtypes; kwargs...) -> Vector{Pair{IRCode, DataType}} + +Return typed code for a cuTile function. Analogous to `Base.code_typed`. +Keyword arguments are those of [`tile_job`](@ref); without any, no toolchain +or device is needed. """ - code_tiled([io::IO], f, argtypes; sm_arch, opt_level, num_ctas, occupancy, - num_worker_warps, remarks=false) +function code_typed(job::TileJob) + ir, rettype = emit_julia(job) + [ir => rettype] +end +function code_typed(@nospecialize(f), @nospecialize(argtypes); + world::UInt=Base.get_world_counter(), kwargs...) + isempty(kwargs) || return code_typed(tile_job(f, argtypes; world, kwargs...)) + ir, rettype = emit_julia(f, argtypes; world) + [ir => rettype] +end + +GPUCompiler.code_typed(job::TileJob) = code_typed(job) +function GPUCompiler.code_warntype(io::IO, job::TileJob; debuginfo::Symbol=:default) + inferred = infer(job) + src = @something get_source(inferred) error("No inferred source for $(job.source)") + GPUCompiler.code_warntype(io, src, inferred_rettype(inferred); debuginfo) +end + + +""" + code_structured(job::TileJob; optimize=true) -> Vector{Pair{StructuredIRCode, DataType}} + code_structured(f, argtypes; optimize=true, kwargs...) + +Return the structured IR for a cuTile function, after the optimization passes +unless `optimize=false`. Keyword arguments are those of [`tile_job`](@ref); +without any, no toolchain or device is needed. +""" +function code_structured(job::TileJob; optimize::Bool=true) + ir, rettype = emit_julia(job) + code_structured(ir, rettype; optimize) +end +function code_structured(@nospecialize(f), @nospecialize(argtypes); + world::UInt=Base.get_world_counter(), optimize::Bool=true, kwargs...) + isempty(kwargs) || return code_structured(tile_job(f, argtypes; world, kwargs...); optimize) + ir, rettype = emit_julia(f, argtypes; world) + code_structured(ir, rettype; optimize) +end +function code_structured(ir::CC.IRCode, rettype; optimize::Bool=true) + sci, rettype, _ = emit_structured(ir, rettype) + if optimize + sci = copy(sci) + run_passes!(sci) + end + [sci => rettype] +end + +""" + code_tiled([io::IO], job::TileJob; debuginfo=false, remarks=false) + code_tiled([io::IO], f, argtypes; debuginfo=false, remarks=false, kwargs...) 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`. Keyword arguments are those of [`tile_job`](@ref): +by default the code targets the current CUDA device, as a launch would. To +inspect the job of an actual launch, use [`@device_code_tiled`](@ref). 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. +This requires `tileiras` 13.4 or newer. """ -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, opt_level) = emit_tile(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) + _, text = run_tileiras(bytecode, job.config.target.sm_arch, opt_level; remarks=true) if !isempty(text) println(io) println(io, "// tileiras optimization remarks") @@ -171,48 +180,76 @@ 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(job::TileJob; kwargs...) = code_tiled(stdout, job; kwargs...) code_tiled(@nospecialize(f), @nospecialize(argtypes); kwargs...) = code_tiled(stdout, f, argtypes; kwargs...) +""" + code_ptx([io::IO], job::TileJob) + code_ptx([io::IO], f, argtypes; kwargs...) + +Print the PTX that `tileiras` generates for a Julia function. This shows the +thread-level SIMT program the tile-level kernel is lowered to, with every +compiler decision (thread mapping, CTA size, pipelining, synchronization) +already made. Keyword arguments are those of [`tile_job`](@ref); to inspect +the job of an actual launch, use [`@device_code_ptx`](@ref). + +!!! warning "Unstable" + The PTX is recorded by `tileiras` in an undocumented CUBIN section and may + go away. +""" +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) +GPUCompiler.code_native(io::IO, job::TileJob) = code_ptx(io, 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; 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`. Keyword arguments are those of [`tile_job`](@ref). For the 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...) + #============================================================================= Device code reflection macros =============================================================================# export @device_code_tiled -public @device_code_typed, @device_code_structured - -# Following GPUCompiler's pattern for @device_code_* macros -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)) - end - end +public @device_code_structured, @device_code_ptx - try - $compile_hook[] = outer_hook - $(esc(user_code)) - finally - $compile_hook[] = nothing - end +# The Julia-level stages are inspected with GPUCompiler's macros, which cover +# every backend through the shared hook; cuTile only adds its own stages. +using GPUCompiler: @device_code_typed, @device_code_warntype +public @device_code_typed, @device_code_warntype - if isempty(seen) - error("no kernels executed while evaluating the given expression") - end +# Tile-specific stages ignore other back-ends sharing the compile hook. +emit_hooked_compilation(hook, ex...) = + GPUCompiler.emit_hooked_compilation(hook, ex...; job_filter=job -> job isa TileJob) - nothing +# A hook printing a signature header around `inner(io, job; kwargs...)`. +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 @@ -228,17 +265,12 @@ 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_structured [io=stdout] expression + @device_code_structured [io=stdout] [optimize=true] expression Print the StructuredIRCode for all kernels compiled while evaluating the expression. @@ -248,31 +280,22 @@ Print the StructuredIRCode for all kernels compiled while evaluating the express ``` """ 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 + hook = tile_hook((io, job; kwargs...) -> println(io, first(only(code_structured(job; kwargs...))))) 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) -> code_ptx(io, job)) emit_hooked_compilation(hook, ex...) end diff --git a/src/cuTile.jl b/src/cuTile.jl index 51b79867..d07aa72f 100644 --- a/src/cuTile.jl +++ b/src/cuTile.jl @@ -22,8 +22,12 @@ using EnumX public BFloat16 import CompilerCaching -using CompilerCaching: CacheView, method_instance, match_method_instance, typeinf!, results, lookup, - specialization, get_source +using CompilerCaching: CacheView, SpecializedResult, method_instance, match_method_instance, + typeinf!, specialization, results, get_source + +# cuTile has its own compiler, but shares GPUCompiler's reflection protocol: the +# compile hook behind the `@device_code_*` macros, and the `code_*` generics. +import GPUCompiler # Shared definitions include("shapes.jl") diff --git a/src/launch.jl b/src/launch.jl index 19e8b29a..f72c7810 100644 --- a/src/launch.jl +++ b/src/launch.jl @@ -2,11 +2,13 @@ # # 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)`. +# CUDA context, and launches it via `cudacall`. Compilation results are cached +# per job configuration on the kernel's inference result. -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 @@ -79,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. =============================================================================# @@ -149,7 +104,7 @@ end const compiler_timeout = parse_compiler_timeout(@load_preference("compiler_timeout_seconds", nothing)) -# Cross-session CUBIN cache (see `emit_binary!`). The `disk_cache` preference gates it; +# Cross-session CUBIN cache (see `assemble`). The `disk_cache` preference gates it; # location and capacity follow Julia's object cache (`JULIA_OBJCACHE_PATH`, # `JULIA_OBJCACHE_CAPACITY`, `JULIA_OBJCACHE=0`). const disk_cache_enabled = let setting = @load_preference("disk_cache", true) @@ -404,6 +359,59 @@ function run_tileiras(bytecode::Vector{UInt8}, sm_arch::VersionNumber, end end +""" + extract_ptx(io::IO) -> String + extract_ptx(cubin::AbstractVector{UInt8}) -> String + +The PTX recorded in a CUBIN. [`run_tileiras`](@ref) compiles with +`--lineinfo`, which makes the assembler embed the PTX it consumed in a +`.nv_debug_ptx_txt` section (one NUL-terminated entry per line, used by +`cuda-gdb` for PTX-level source display). +""" +function extract_ptx(io::IO) + # CUBINs are ELF64 files with a machine type binutils does not know; + # ObjectFile.jl only needs the section headers. + handle = only(readmeta(io)) + section = findfirst(Sections(handle), ".nv_debug_ptx_txt") + section === nothing && error( + "CUBIN does not embed PTX (no .nv_debug_ptx_txt section); " * + "it was compiled without line info") + lines = split(String(read(section)), '\0'; keepempty=false) + return join(lines, '\n') * "\n" +end +extract_ptx(cubin::AbstractVector{UInt8}) = extract_ptx(IOBuffer(cubin)) + +""" + nvdisasm_cmd(args...) -> Cmd + +Construct a Cmd to invoke `nvdisasm` with `args`. With the `tileiras` +preference set, requires `nvdisasm` from the same toolkit; otherwise uses +`CUDA_Compiler_jll`. +""" +function nvdisasm_cmd(args...) + if tileiras_override !== nothing + disasm = joinpath(dirname(tileiras_override), "nvdisasm") + isfile(disasm) || error("no `nvdisasm` next to $(tileiras_path())") + return `$disasm $args` + end + CUDA_Compiler_jll.is_available() && isdefined(CUDA_Compiler_jll, :nvdisasm) || + error("CUDA_Compiler_jll does not provide `nvdisasm`") + return `$(CUDA_Compiler_jll.nvdisasm()) $args` +end + +""" + disassemble_cubin(cubin::Vector{UInt8}) -> String + +Disassemble a CUBIN to SASS using `nvdisasm`. +""" +function disassemble_cubin(cubin::Vector{UInt8}) + mktempdir() do dir + path = joinpath(dir, "kernel.cubin") + write(path, cubin) + read(nvdisasm_cmd(path), String) + end +end + const TILEIRAS_VERSION_REGEX = r"V(\d+\.\d+\.\d+)" function parse_tileiras_version(log::AbstractString) @@ -504,7 +512,7 @@ end The architecture-family name and the minimum bytecode version Tile IR requires on a device of compute capability `cap`, or `nothing` if Tile IR is not supported on that capability at all. Pure (no device access) so the gate logic -in [`check_tile_ir_support`] can be unit-tested without a GPU. +in [`validate_tile_ir_target`](@ref) can be unit-tested without a GPU. """ function tile_ir_requirement(cap::VersionNumber) if cap >= v"10.0" # Blackwell @@ -519,18 +527,11 @@ function tile_ir_requirement(cap::VersionNumber) end """ - check_tile_ir_support(sm_arch) - -Validate that the selected bytecode version supports Tile IR on `sm_arch`. -Returns the bytecode version cuTile should emit, provided it meets the target's -minimum requirement (Blackwell ≥ v13.1, Hopper ≥ v13.3, Ampere/Ada ≥ v13.2). -""" -function check_tile_ir_support(sm_arch::VersionNumber) - version = bytecode_version() validate_tile_ir_target(sm_arch, version) - return version -end +Validate that bytecode `version` supports Tile IR on `sm_arch`: Blackwell needs +v13.1+, Hopper v13.3+, Ampere/Ada v13.2+. +""" function validate_tile_ir_target(sm_arch::VersionNumber, version::VersionNumber) validate_bytecode_version(version) requirement = tile_ir_requirement(sm_arch) @@ -548,18 +549,18 @@ end =============================================================================# """ - unwrap_argtypes(f, tt) -> (argtypes::Type{<:Tuple}, const_argtypes::Union{Vector{Any},Nothing}) + unwrap_argtypes(f, tt) -> (argtypes::Type{<:Tuple}, const_argtypes::Union{Tuple,Nothing}) Compile-time-specialized derivation of: - `argtypes::Type{<:Tuple}` — concrete dispatch tuple for `method_instance(f, argtypes)`, with `Constant{T,V}` slots unwrapped to `T`. -- `const_argtypes::Vector{Any}` — `[CC.Const(f), ...args]` with `Constant{T,V}` slots - replaced by `CC.Const(V)`, for const-prop inference. `nothing` when no `Constant` - arguments are present (skips the const-seeding pipeline entirely). +- `const_argtypes` — `(CC.Const(f), args...)` with `Constant{T,V}` slots replaced by + `CC.Const(V)`, seeding const-prop inference (see `TileJob`). `nothing` when no + `Constant` arguments are present (skips the const-seeding pipeline entirely). `@generated` so the unwrapped `Tuple` type and the `Constant`-vs-not branching -fold to constants at the call site. Only the `Vector{Any}` allocation and the -`CC.Const(...)` boxes for runtime values survive to runtime. +fold to constants at the call site. Only the `CC.Const(f)` box for the runtime +function value survives to runtime. """ @generated function unwrap_argtypes(@nospecialize(f), ::Type{TT}) where TT <: Tuple unwrapped = map(t -> t <: Constant ? constant_eltype(t) : t, TT.parameters) @@ -575,7 +576,7 @@ fold to constants at the call site. Only the `Vector{Any}` allocation and the push!(cats_exprs, t) end end - return :(($argtypes_T, Any[$(cats_exprs...)])) + return :(($argtypes_T, ($(cats_exprs...),))) end @@ -584,61 +585,106 @@ 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) - # Recurse first — emit_structured! at the bottom of the chain fires - # `compile_hook` for `@device_code_*` reflection, which must run on every - # launch even when downstream artifacts are fully cached. - bytecode = emit_tile!(cache, mi, ci, res; const_argtypes) - - res.cuda_bin !== nothing && return res.cuda_bin - - 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 = res.julia_ir - 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 compiled form of one configuration of a kernel: its CUBIN, shared by every +# CUDA context and persisted with the inference result, and the session-local +# kernels linked from it. +mutable struct TileCompilerResults + cubin::Union{Nothing, Vector{UInt8}} + # linear-scanned by context; usually holds a single entry + kernels::Vector{Tuple{CuContext, Any}} + TileCompilerResults() = new(nothing, Tuple{CuContext, Any}[]) end +# The results of every job of a kernel, attached by CompilerCaching to the +# kernel's inference result (and persisted with it into package images). One +# inferred source serves every target and hint configuration, so entries are +# keyed by config; the inference result itself outlives any one world. +mutable struct JobResults + entries::Vector{Pair{TileConfig, TileCompilerResults}} + JobResults() = new(Pair{TileConfig, TileCompilerResults}[]) +end + +# Guards result lookup and publication, including CUBINs and linked kernels. +# Compilation and module loading run outside the lock. +const job_results_lock = ReentrantLock() + +job_results(cache::CacheView, ci::Core.CodeInstance) = results(cache, ci) +job_results(::CacheView, entry::SpecializedResult) = results(entry) + """ - link(cache, mi, ci, res) -> CuFunction + cached_results(job::TileJob) -> TileCompilerResults + +The results slot of a job, created empty on first access. Runs inference if the +job's source has not been inferred yet. +""" +function cached_results(job::TileJob) + table = job_results(inference_cache(job), infer(job)) + Base.@lock job_results_lock begin + for (config, res) in table.entries + # configs are immutable, so `===` compares them structurally + config === job.config && return res + end + res = TileCompilerResults() + push!(table.entries, job.config => res) + return res + end +end + +""" + compile_or_lookup(job::TileJob) -> TileCompilerResults + +The compilation results of `job`, compiling on a miss. Distinct jobs compile +concurrently; the `@device_code_*` hook observes hits as well as misses. No +CUDA context required. +""" +function compile_or_lookup(job::TileJob) + res = cached_results(job) + cubin = Base.@lock job_results_lock res.cubin + if cubin === nothing + cubin = compile(job) + Base.@lock job_results_lock begin + # Another task may have compiled the same job in the meantime. + res.cubin === nothing && (res.cubin = cubin) + end + else + run_compile_hook(job) + end + return res +end + +# Caller holds job_results_lock. +function find_kernel(res::TileCompilerResults, ctx::CuContext) + for (cached_ctx, kernel) in res.kernels + cached_ctx === ctx && return kernel + end + return nothing +end -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`. """ -function link(cache::CacheView, mi::Core.MethodInstance, - ci::Core.CodeInstance, res::CuTileResults) - res.cuda_func !== nothing && return res.cuda_func + link(job::TileJob, res::TileCompilerResults) -> CuFunction - kernel_name = sanitize_name(string(mi.def.name)) - cumod = CuModule(res.cuda_bin::Vector{UInt8}) - cufunc = CuFunction(cumod, kernel_name) - res.cuda_func = cufunc - return cufunc +Load the job's CUBIN onto the active CUDA context. +""" +function link(job::TileJob, res::TileCompilerResults) + cumod = CuModule(res.cubin::Vector{UInt8}) + return CuFunction(cumod, job.config.name) end @@ -647,7 +693,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. =============================================================================# """ @@ -672,10 +718,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, @@ -684,37 +729,33 @@ function cufunction(@nospecialize(f), tt::Type{<:Tuple}=Tuple{}; occupancy::Union{Int, Nothing}=nothing, num_worker_warps::Union{Int, Nothing}=nothing, name::Union{String, Nothing}=nothing) - 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) + # `tileiras` generates architecture-specific code (`sm_100a`), so a kernel + # only runs on the exact architecture it was compiled for. Reflection may + # target other architectures; execution cannot. + device_arch = device_sm_arch() + if sm_arch === nothing + sm_arch = device_arch + elseif sm_arch != device_arch + throw(ArgumentError( + "Cannot execute code compiled for $(format_sm_arch(sm_arch)) on $(device()) " * + "(compute capability $device_arch); omit `sm_arch` or pass the device's")) + end - # 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 - # specializes on `tt`, so this loop unrolls per kernel signature. + # Unwrap Constant arguments for method lookup and seed their values into inference. argtypes, const_argtypes = unwrap_argtypes(f, tt) # The compilation pipeline (typeinf!, codegen, bytecode emission) gets # 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, 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{Tuple, Nothing}; kwargs...) world = Base.get_world_counter() mi = method_instance(f, argtypes; world) mi === nothing && throw(MethodError(f, argtypes)) @@ -722,37 +763,33 @@ 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) - - # Always walk the emit chain (each phase short-circuits on its own cached - # field, but `emit_structured!` also fires `compile_hook` for reflection, - # which has to run on every launch even when the cube/cufunc is cached). - 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{Tuple, Nothing}, opts::NamedTuple) + 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() + cached_kernel = Base.@lock job_results_lock find_kernel(res, ctx) + cached_kernel === nothing || return cached_kernel::TileKernel{Core.Typeof(f), tt} + 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 inference result. + if ccall(:jl_generating_output, Cint, ()) != 1 + Base.@lock job_results_lock begin + cached_kernel = find_kernel(res, ctx) + cached_kernel === nothing || return cached_kernel::TileKernel{Core.Typeof(f), tt} + push!(res.kernels, (ctx, kernel)) + end + end return kernel end @@ -886,12 +923,11 @@ function launch(@nospecialize(f), grid, args...; end """ - default_sm_arch() -> VersionNumber + device_sm_arch() -> VersionNumber -Get the compute capability of the current CUDA device as a VersionNumber. -Returns e.g. `v"12.0"` for compute capability 12.0. +The compute capability of the current CUDA device, e.g. `v"12.0"`. """ -default_sm_arch() = capability(device()) +device_sm_arch() = capability(device()) #============================================================================= 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/Project.toml b/test/Project.toml index 5e763e1d..009b9668 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -3,6 +3,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" DLFP8Types = "f4c16678-4a16-415b-82ef-ed337c5d6c7c" FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" FileCheck = "4e644321-382b-4b05-b0b6-5d23c3d944fb" +GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55" IRStructurizer = "93e32bba-5bb8-402b-805d-ffb066edee93" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" diff --git a/test/codegen/boundscheck_integration.jl b/test/codegen/boundscheck_integration.jl index 3c68a932..5bf5ee03 100644 --- a/test/codegen/boundscheck_integration.jl +++ b/test/codegen/boundscheck_integration.jl @@ -57,19 +57,14 @@ function structured_if_count(@nospecialize(f), @nospecialize(argtypes)) for block in IRStructurizer.eachblock(sci))) end +# Bytecode of `f`, stripped of everything that differs between equivalent +# kernels: the kernel name and the source locations. 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) - ir, rettype = ct.emit_julia(cache, mi; const_argtypes) + job = ct.tile_job(f, argtypes; name="boundscheck_test") + ir, rettype = ct.emit_julia(job) sci, rettype, kernel_meta = ct.emit_structured(ir, rettype) empty!(sci.line_map) - opts = ct.CGOpts((sm_arch=nothing, opt_level=nothing, num_ctas=nothing, - occupancy=nothing, num_worker_warps=nothing, - bytecode_version=ct.bytecode_version())) - return ct.emit_tile(sci, rettype, kernel_meta; - name="boundscheck_test", opts, cache, const_argtypes) + return ct.emit_tile(job, sci, rettype, kernel_meta).bytecode end if Base.JLOptions().check_bounds == 2 diff --git a/test/codegen/operations.jl b/test/codegen/operations.jl index 56c70f2e..86a0d425 100644 --- a/test/codegen/operations.jl +++ b/test/codegen/operations.jl @@ -2210,7 +2210,7 @@ end # Older toolchains must reject the op cleanly. @test_throws "v13.2" code_tiled(Tuple{ct.TileArray{Float32,1,Int32,spec1d}, ct.TileArray{Float32,1,Int32,spec1d}}; - bytecode_version=v"13.1") do y, x + sm_arch=v"10.0", bytecode_version=v"13.1") do y, x pid = ct.bid(1) ty = ct.load(y, pid, (16,)) tx = ct.load(x, pid, (16,)) @@ -2714,7 +2714,7 @@ end end @test_throws "BFloat16 requires Tile IR bytecode ≥ 13.3" code_tiled( Tuple{ct.TileArray{ct.BFloat16,1,Int32,spec_bf16}}; - sm_arch=v"9.0", + sm_arch=v"10.0", bytecode_version=v"13.2") do arr indices = ct.arange(16; dtype=Int) ct.atomic_add(arr, indices, ct.BFloat16(1)) @@ -3325,7 +3325,7 @@ end @check_label "entry" @check_count 1 "make_token" code_tiled(Tuple{ct.TileArray{Float32,1,Int32,spec1d}}; - bytecode_version=v"13.1") do a + sm_arch=v"10.0", bytecode_version=v"13.1") do a bid = ct.bid(1) tile = ct.load(a, bid, (16,)) @check "print_tko" diff --git a/test/codegen/reflection.jl b/test/codegen/reflection.jl index 049b5a90..435f987e 100644 --- a/test/codegen/reflection.jl +++ b/test/codegen/reflection.jl @@ -9,6 +9,24 @@ function reflect_vadd(a, b, c) return end +function reflect_vadd_n(a, b, c, n) + pid = ct.bid(1) + tile_a = ct.load(a; index=pid, shape=(n,)) + tile_b = ct.load(b; index=pid, shape=(n,)) + ct.store(c; index=pid, tile=tile_a + tile_b) + 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" @@ -72,6 +90,170 @@ end end end +if ct.tileiras_available() + @testset "code_ptx" begin + @test @filecheck begin + @check ".visible .entry reflect_vadd" + @check "add.rn.f32" + @check "ret;" + ct.code_ptx(reflect_vadd, TT3; sm_arch=v"10.0") + end + end + + @testset "code_sass" begin + 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 + 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) + 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 + + @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.const_argtypes === nothing + @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, job) == "TileJob(reflect_vadd($(join(TT3.parameters, ", "))); " * + "sm_arch=10.0.0, bytecode_version=v\"$(ct.bytecode_version())\")" + @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) + + # Constant arguments seed inference; the job restores them in its signature. + const_tt = Tuple{TT3.parameters..., ct.Constant{Int, 16}} + const_job = ct.tile_job(reflect_vadd_n, const_tt; sm_arch=v"10.0") + @test const_job.const_argtypes == (ct.CC.Const(reflect_vadd_n), TT3.parameters..., ct.CC.Const(16)) + @test ct.job_signature(const_job) == (reflect_vadd_n, const_tt) + @test const_job === ct.tile_job(reflect_vadd_n, const_tt; sm_arch=v"10.0") + @test occursin("Constant{Int64, 16}", sprint(show, const_job)) + @test ct.compile_or_lookup(const_job) === ct.compile_or_lookup(const_job) + + # Without an architecture, jobs target the device. + if CUDA.functional() + device_arch = ct.device_sm_arch() + requirement = ct.tile_ir_requirement(device_arch) + if requirement !== nothing && ct.bytecode_version() >= requirement[2] + @test ct.tile_job(reflect_vadd, TT3).config.target.sm_arch == device_arch + else + @test_throws ArgumentError ct.tile_job(reflect_vadd, TT3) + end + else + @test_throws ArgumentError ct.tile_job(reflect_vadd, TT3) + end + old_job = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0", bytecode_version=v"13.1") + @test old_job.config.target == ct.TileCompilerTarget(v"10.0", v"13.1") + + # Incompatible targets are rejected when the job is built. + @test_throws ArgumentError ct.tile_job(reflect_vadd, TT3; sm_arch=v"7.5") + @test_throws "v13.2" ct.tile_job(reflect_vadd, TT3; sm_arch=v"8.0", bytecode_version=v"13.1") + 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"12.0", opt_level=1) + res1, res2 = ct.compile_or_lookup(job1), ct.compile_or_lookup(job2) + ci1, ci2 = ct.infer(job1), ct.infer(job2) + @test ci1 === ci2 + @test ci1.owner === ct.TILE_CACHE_OWNER + @test res1 !== res2 && res1.cubin != res2.cubin + # 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 + + @testset "compile hook" begin + job = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0") + ct.compile_or_lookup(job) + # GPUCompiler's hook observes hits as well as misses, once per distinct job + seen = ct.TileJob[] + with(GPUCompiler.compile_hook => (job -> push!(seen, job))) do + ct.compile_or_lookup(job) + ct.compile_or_lookup(job) + end + @test seen == [job, job] + @test GPUCompiler.compile_hook[] === nothing + + # GPUCompiler's macros accept Tile jobs through the shared protocol; cuTile + # re-exports the Julia-level ones rather than defining its own + @test ct.var"@device_code_typed" === GPUCompiler.var"@device_code_typed" + typed = ct.@device_code_typed ct.compile_or_lookup(job) + @test collect(keys(typed)) == [job] + @test sprint(show, only(typed[job])) == sprint(show, only(ct.code_typed(job))) + warntype = sprint() do io + GPUCompiler.@device_code_warntype io=io ct.compile_or_lookup(job) + end + @test occursin("reflect_vadd", warntype) && occursin("Body::Nothing", warntype) + @test_throws ArgumentError GPUCompiler.@device_code_llvm ct.compile_or_lookup(job) + + # Warntype must inspect the const-seeded source, just like typed reflection. + select_type(n) = n == 16 ? 1 : 1.0f0 + const_job = ct.tile_job(select_type, Tuple{ct.Constant{Int,16}}; sm_arch=v"10.0") + @test occursin("Body::Int", sprint(GPUCompiler.code_warntype, const_job)) + float_job = ct.tile_job(select_type, Tuple{ct.Constant{Int,32}}; sm_arch=v"10.0") + @test occursin("Body::Float32", sprint(GPUCompiler.code_warntype, float_job)) + shape_job = ct.tile_job(reflect_vadd_n, Tuple{TT3.parameters..., ct.Constant{Int,16}}; + sm_arch=v"10.0") + @test occursin("Body::Nothing", sprint(GPUCompiler.code_warntype, shape_job)) + @test sprint(GPUCompiler.code_native, job) == sprint(ct.code_ptx, job) + + + # Nested macros restore the outer hook; child tasks inherit it, and + # repeated jobs are printed only once in each scope. + outer, inner = IOBuffer(), IOBuffer() + ct.@device_code_structured io=outer begin + ct.@device_code_structured io=inner ct.compile_or_lookup(job) + @sync for _ in 1:2 + @async ct.compile_or_lookup(job) + end + end + @test String(take!(outer)) == String(take!(inner)) != "" + @test GPUCompiler.compile_hook[] === nothing + @test_throws ErrorException ct.@device_code_structured error("reflection failed") + @test GPUCompiler.compile_hook[] === nothing + end + + @testset "concurrent cache misses" begin + job = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0", name="concurrent_vadd") + ct.cached_results(job) + entered, resume = Channel{Nothing}(1), Channel{Nothing}(1) + late = @async with(GPUCompiler.compile_hook => (_ -> (put!(entered, nothing); take!(resume)))) do + ct.compile_or_lookup(job) + end + take!(entered) + first_res = try + # A different configuration can compile while the first is suspended. + other = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0", name="concurrent_other") + @test !isempty(ct.compile_or_lookup(other).cubin) + ct.compile_or_lookup(job) + finally + put!(resume, nothing) + end + first_cubin = first_res.cubin + last_res = fetch(late) + @test last_res === first_res + @test last_res.cubin === first_cubin + end +end + if ct.tileiras_version() >= v"13.4" @testset "code_tiled remarks" begin output = sprint(io -> ct.code_tiled(io, reflect_vadd, TT3; diff --git a/test/device/compilation.jl b/test/device/compilation.jl new file mode 100644 index 00000000..1ad87f9d --- /dev/null +++ b/test/device/compilation.jl @@ -0,0 +1,127 @@ +using CUDA + +@testset "job reflection uses launch options" begin + function hinted_kernel(a, n) + ct.store(a, 1, fill(1.0f0, (n,))) + return + end + a = CUDA.zeros(Float32, 16) + converted = (ct.cuTileconvert(a), ct.Constant(16)) + tt = Tuple{map(Core.Typeof, converted)...} + job = ct.tile_job(hinted_kernel, tt; occupancy=4, name="hinted_launch") + + # Populate the launch cache, then reflect the warm launch with the same options. + ct.cufunction(hinted_kernel, tt; occupancy=4, name="hinted_launch") + output = IOBuffer() + ct.@device_code_tiled io=output @cuda backend=cuTile occupancy=4 name="hinted_launch" hinted_kernel(a, ct.Constant(16)) + @test occursin(sprint(ct.code_tiled, job), String(take!(output))) + @test Array(a) == ones(Float32, 16) +end + +@testset "launches target the device" begin + function device_kernel(a) + ct.store(a, 1, fill(1.0f0, (16,))) + return + end + a = CUDA.zeros(Float32, 16) + tt = Tuple{Core.Typeof(ct.cuTileconvert(a))} + device_arch = ct.device_sm_arch() + @test ct.cufunction(device_kernel, tt) === ct.cufunction(device_kernel, tt; sm_arch=device_arch) + # tileiras generates architecture-specific code, so other targets cannot run here. + other_arch = device_arch == v"10.0" ? v"12.0" : v"10.0" + @test_throws "Cannot execute" ct.cufunction(device_kernel, tt; sm_arch=other_arch) +end + +@testset "linked kernels are cached per context" begin + context_kernel() = nothing + ctx = CUDA.context() + job = ct.tile_job(context_kernel, Tuple{}) + res = ct.compile_or_lookup(job) + cubin = res.cubin + + # All tasks must converge on the same linked kernel for this context. + tasks = [Threads.@spawn CUDA.context!(ctx) do + ct.cufunction(context_kernel) + end for _ in 1:8] + kernels = fetch.(tasks) + kernel = first(kernels) + @test all(k -> k === kernel, kernels) + @test ct.cufunction(context_kernel) === kernel + kernel() + CUDA.synchronize() + + exclusive = CUDA.attribute(CUDA.device(), CUDA.DEVICE_ATTRIBUTE_COMPUTE_MODE) == + CUDA.COMPUTEMODE_EXCLUSIVE_PROCESS + if !exclusive + CUDA.CuContext(CUDA.device()) do other_ctx + CUDA.context!(other_ctx) do + other_kernel = ct.cufunction(context_kernel) + @test other_kernel !== kernel + @test ct.cufunction(context_kernel) === other_kernel + stream = CUDA.CuStream() + try + other_kernel(; stream) + CUDA.synchronize(stream) + finally + finalize(stream) + finalize(other_kernel.fun.mod) + end + end + end + @test ct.cufunction(context_kernel) === kernel + @test ct.compile_or_lookup(job).cubin === cubin + end +end + +@testset "mixed-backend reflection" begin + function tile_kernel(a) + ct.store(a, 1, fill(1.0f0, (16,))) + return + end + function cuda_kernel(a) + a[1] = 2.0f0 + return + end + a = CUDA.zeros(Float32, 16) + + # Warm both launch caches before inspecting the same expression. + @cuda cuda_kernel(a) + @cuda backend=cuTile tile_kernel(a) + typed = CUDA.@device_code_typed begin + @cuda cuda_kernel(a) + @cuda backend=cuTile tile_kernel(a) + end + @test length(typed) == 2 + @test count(job -> job isa ct.TileJob, keys(typed)) == 1 + @test count(job -> job isa GPUCompiler.CompilerJob, keys(typed)) == 1 + + warntype = sprint() do io + CUDA.@device_code_warntype io=io begin + @cuda cuda_kernel(a) + @cuda backend=cuTile tile_kernel(a) + end + end + @test occursin("cuda_kernel", warntype) + @test occursin("tile_kernel", warntype) + + ptx = sprint() do io + CUDA.@device_code_ptx io=io begin + @cuda cuda_kernel(a) + @cuda backend=cuTile tile_kernel(a) + end + end + @test occursin("cuda_kernel", ptx) + @test occursin("tile_kernel", ptx) + + tiled = sprint() do io + ct.@device_code_tiled io=io begin + @cuda cuda_kernel(a) + @cuda backend=cuTile tile_kernel(a) + end + end + @test occursin("tile_kernel", tiled) + @test !occursin("cuda_kernel", tiled) + @test_throws "no kernels executed" ct.@device_code_tiled @cuda cuda_kernel(a) + @test GPUCompiler.compile_hook[] === nothing + CUDA.synchronize() +end diff --git a/test/device/core.jl b/test/device/core.jl index ffca735c..3e0323ad 100644 --- a/test/device/core.jl +++ b/test/device/core.jl @@ -202,16 +202,16 @@ end # invalidations String(take!(buf)) end - # @device_code_typed: check typed Julia IR output + # @device_code_typed: GPUCompiler's macro, returning the typed Julia IR per job @test @filecheck begin - @check "// reflect_vadd" @check "get_tile_block_id" @check "load_partition_view" @check "addf" @check "store_partition_view" - buf = IOBuffer() - ct.@device_code_typed io=buf @cuda backend=cuTile blocks=cld(n, 16) reflect_vadd(a, b, c) - String(take!(buf)) + typed = ct.@device_code_typed @cuda backend=cuTile blocks=cld(n, 16) reflect_vadd(a, b, c) + job = only(keys(typed)) + @test job isa ct.TileJob && ct.job_signature(job)[1] === reflect_vadd + string(first(only(typed[job]))) end # @device_code_tiled with debug info diff --git a/test/host/cache.jl b/test/host/cache.jl index a42eedf2..50a4c3c5 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 (; cubin=res.cubin, tile_bc=ct.emit_tile(job).bytecode) end cubin_key(res) = OC.keyhash(ct.CUBIN_CACHE_SCHEMA, ct.cubin_cache_fields(res.tile_bc, @@ -91,14 +91,14 @@ digest(out) = match(r"CUBIN_SHA=([0-9a-f]+)", out) ok, out = @with_objcache dir "JULIA_OBJCACHE" => "1" begin @test OC.enabled() res = compile_vadd() - @test res.cuda_bin isa Vector{UInt8} && !isempty(res.cuda_bin) + @test res.cubin isa Vector{UInt8} && !isempty(res.cubin) k = cubin_key(res) t0 = time() while OC.get(ct.CUBIN_CACHE_NS, k) === nothing && time() - t0 < 10 sleep(0.01) end - @test OC.get(ct.CUBIN_CACHE_NS, k) == res.cuda_bin - println("CUBIN_SHA=", bytes2hex(OC.keyhash(0, res.cuda_bin))) + @test OC.get(ct.CUBIN_CACHE_NS, k) == res.cubin + println("CUBIN_SHA=", bytes2hex(OC.keyhash(0, res.cubin))) end check_child(ok, out) sha1 = digest(out) @@ -115,8 +115,8 @@ digest(out) = match(r"CUBIN_SHA=([0-9a-f]+)", out) error("tileiras invoked despite a warm object cache") end res = compile_vadd() - @test res.cuda_bin isa Vector{UInt8} && !isempty(res.cuda_bin) - @test bytes2hex(OC.keyhash(0, res.cuda_bin)) == $expected_sha + @test res.cubin isa Vector{UInt8} && !isempty(res.cubin) + @test bytes2hex(OC.keyhash(0, res.cubin)) == $expected_sha end check_child(ok, out) end diff --git a/test/host/inference_cache.jl b/test/host/inference_cache.jl index 98f3aca6..c2990aa5 100644 --- a/test/host/inference_cache.jl +++ b/test/host/inference_cache.jl @@ -1,33 +1,50 @@ @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) - _, generic_rt = ct.get_inferred(cache, ci, mi) + ci = ct.infer(cache, mi) + @test ci isa Core.CodeInstance && ci.owner === ct.TILE_CACHE_OWNER + @test ct.infer(cache, mi) === ci + _, generic_rt = ct.emit_julia(cache, mi) @test generic_rt == Union{Int32, Float32} - specialized = ct.CuTileResults[] + # Every job of a kernel shares its inference result; results are stored per + # config on the generic CodeInstance, or on the const-seeded entry. + generic_job = ct.tile_job(mi, world; sm_arch=v"10.0", bytecode_version=v"13.2") + @test ct.infer(generic_job) === ci + generic = ct.cached_results(generic_job) + @test ct.cached_results(generic_job) === generic + hinted_job = ct.tile_job(mi, world; sm_arch=v"10.0", bytecode_version=v"13.2", opt_level=1) + @test ct.infer(hinted_job) === ci + @test ct.cached_results(hinted_job) !== generic + + specialized = ct.TileCompilerResults[] 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) - @test specialized_ci === ci - @test res !== generic - @test all(previous -> previous !== res, specialized) - push!(specialized, res) + const_argtypes = (ct.CC.Const(choose_type), ct.CC.Const(flag), Int32) + job = ct.tile_job(mi, world; const_argtypes, sm_arch=v"10.0", bytecode_version=v"13.2") + entry = ct.infer(job) + @test entry isa ct.SpecializedResult{ct.JobResults} + @test entry === ct.infer(cache, mi, collect(Any, const_argtypes)) + @test ct.infer(cache, mi) === ci - # Source and return type must come from the same specialization. - ir, rt = ct.get_inferred(cache, ci, mi; const_argtypes=argtypes) + ir, rt = ct.emit_julia(job) @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 + res = ct.cached_results(job) + @test res !== generic + @test all(previous -> previous !== res, specialized) + push!(specialized, res) + + # An equivalent job must reuse the inference result and the results slot. + equivalent = ct.tile_job(mi, world; sm_arch=v"10.0", bytecode_version=v"13.2", + const_argtypes=(ct.CC.Const(choose_type), ct.CC.Const(flag), Int32)) + @test equivalent === job + @test ct.infer(equivalent) === entry + @test ct.cached_results(equivalent) === res end end diff --git a/test/runtests.jl b/test/runtests.jl index d6ff45ea..3e03117f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -95,6 +95,8 @@ const init_code = quote using CUDA using cuTile import cuTile as ct + import GPUCompiler + using Base.ScopedValues: with using FileCheck end