From dfbf7cc80ddf436443ce5c56cc3328f5a07269c4 Mon Sep 17 00:00:00 2001 From: AntonOresten Date: Fri, 21 Aug 2026 16:05:02 +0200 Subject: [PATCH 1/9] Add PTX and SASS reflection --- docs/src/lib/reflection.md | 8 +++ docs/src/man/debugging.md | 39 ++++++++++- src/compiler/reflection.jl | 131 +++++++++++++++++++++++++++++++++++ src/launch.jl | 136 +++++++++++++++++++++++++++++++++++++ test/codegen/reflection.jl | 17 +++++ 5 files changed, 329 insertions(+), 2 deletions(-) diff --git a/docs/src/lib/reflection.md b/docs/src/lib/reflection.md index fc478368..cdd03fc9 100644 --- a/docs/src/lib/reflection.md +++ b/docs/src/lib/reflection.md @@ -11,6 +11,8 @@ code_tiled code_typed code_ircode code_structured +code_ptx +code_sass ``` ## Intercepting a launch @@ -19,4 +21,10 @@ code_structured @device_code_tiled @device_code_typed @device_code_structured +@device_code_ptx +@device_code_sass ``` + +`CUDA.@device_code_sass` also 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..9e28c8a2 100644 --- a/docs/src/man/debugging.md +++ b/docs/src/man/debugging.md @@ -49,6 +49,31 @@ device is available. 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: `sm_arch` may be omitted when a CUDA device is +available, and must be given explicitly otherwise. + +The PTX is recovered from a debug section `tileiras` embeds in the CUBIN, in +which comment lines are blanked; the code itself is unaffected. + +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. + + ## Intercepting a launch `@device_code_*` macros intercept compilation during a kernel launch, deriving @@ -73,17 +98,27 @@ 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: +Five are available, corresponding to successive stages of the pipeline: | 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_structured` | Structured IR (after control-flow structurization) | +| `ct.@device_code_ptx` | PTX generated by `tileiras` | +| `ct.@device_code_sass` | SASS disassembly of the compiled CUBIN | 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. +flow, then emitted Tile IR, then the PTX it lowers to and the machine code +that assembles to. + +`CUDA.@device_code_sass` also 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. The cuTile variant instead recompiles each intercepted signature, +which requires no CUPTI subscription and therefore also works under Nsight or +another active profiler. ## Dumping bytecode diff --git a/src/compiler/reflection.jl b/src/compiler/reflection.jl index b54287f3..5ae347d3 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 @@ -174,6 +175,86 @@ end code_tiled(@nospecialize(f), @nospecialize(argtypes); kwargs...) = code_tiled(stdout, f, argtypes; kwargs...) +""" + compile_to_cubin(f, argtypes; sm_arch, opt_level, num_ctas, occupancy, + num_worker_warps, bytecode_version, world) -> Vector{UInt8} + +Compile a Julia function all the way to a CUBIN for reflection: emit Tile IR +bytecode (like [`code_tiled`](@ref)) and assemble it with `tileiras`. Like the +rest of the reflection entry points, this calls the driver directly without +caching in CuTileResults or the disk cache. +""" +function compile_to_cubin(@nospecialize(f), @nospecialize(argtypes); + sm_arch::Union{VersionNumber, Nothing}=nothing, + opt_level::Union{Int, Nothing}=nothing, + num_ctas::Union{Int, Nothing}=nothing, + occupancy::Union{Int, Nothing}=nothing, + num_worker_warps::Union{Int, Nothing}=nothing, + bytecode_version::VersionNumber=cuTile.bytecode_version(), + world::UInt=Base.get_world_counter()) + target = if sm_arch === nothing + try + default_sm_arch() + catch + throw(ArgumentError( + "sm_arch must be specified when compiling without a CUDA device")) + end + else + sm_arch + end + validate_tileiras_target(bytecode_version) + + stripped, const_argtypes = process_const_argtypes(f, argtypes) + mi = lookup_method_instance(f, stripped; world) + opts = CGOpts((sm_arch=target, opt_level=opt_level, num_ctas=num_ctas, + occupancy=occupancy, num_worker_warps=num_worker_warps, + bytecode_version=bytecode_version)) + cache = CacheView{CuTileResults}(:cuTile, world) + ir, rettype = emit_julia(cache, mi; const_argtypes) + sci, rettype, kernel_meta = emit_structured(ir, rettype) + bytecode = emit_tile(sci, rettype, kernel_meta; + name=sanitize_name(string(mi.def.name)), + opts, cache, const_argtypes) + resolved_opt_level = something(resolve_hint(opt_level, kernel_meta, + :opt_level, target), 3) + cubin, _ = run_tileiras(bytecode, target, resolved_opt_level) + return cubin +end + +""" + code_ptx([io::IO], f, argtypes; sm_arch, opt_level, num_ctas, occupancy, + num_worker_warps) + +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. When no GPU is available, pass `sm_arch` explicitly. + +The PTX is recovered from the `.nv_debug_ptx_txt` section the assembler embeds +in the CUBIN when compiling with line info, so header comments are blanked. +""" +function code_ptx(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) + cubin = compile_to_cubin(f, argtypes; kwargs...) + print(io, extract_ptx(cubin)) +end +code_ptx(@nospecialize(f), @nospecialize(argtypes); kwargs...) = + code_ptx(stdout, f, argtypes; kwargs...) + +""" + code_sass([io::IO], f, argtypes; sm_arch, opt_level, num_ctas, occupancy, + num_worker_warps) + +Print the SASS machine code that a Julia function compiles to, by assembling +the Tile IR with `tileiras` and disassembling the resulting CUBIN with +`nvdisasm`. When no GPU is available, pass `sm_arch` explicitly. +""" +function code_sass(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) + cubin = compile_to_cubin(f, argtypes; kwargs...) + print(io, disassemble_cubin(cubin)) +end +code_sass(@nospecialize(f), @nospecialize(argtypes); kwargs...) = + code_sass(stdout, f, argtypes; kwargs...) + #============================================================================= Device code reflection macros @@ -181,6 +262,7 @@ code_tiled(@nospecialize(f), @nospecialize(argtypes); kwargs...) = export @device_code_tiled public @device_code_typed, @device_code_structured +public @device_code_ptx, @device_code_sass # Following GPUCompiler's pattern for @device_code_* macros function emit_hooked_compilation(inner_hook, ex...) @@ -237,6 +319,55 @@ macro device_code_tiled(ex...) emit_hooked_compilation(hook, ex...) end +""" + @device_code_ptx [io=stdout] expression + +Print the PTX generated by `tileiras` for all kernels compiled while +evaluating the expression. + +# Example +```julia +@device_code_ptx @cuda backend=cuTile blocks=grid vadd(a, b, c) +``` +""" +macro device_code_ptx(ex...) + function hook(f, tt; io::IO=stdout, kwargs...) + println(io, "// $f($(join(tt.parameters, ", ")))") + println(io) + code_ptx(io, f, Tuple(tt.parameters); kwargs...) + println(io) + end + emit_hooked_compilation(hook, ex...) +end + +""" + @device_code_sass [io=stdout] expression + +Print the SASS machine code for all kernels compiled while evaluating the +expression, by recompiling each intercepted signature and disassembling the +result with `nvdisasm`. Works without a CUPTI subscription, so it can be used +under Nsight or another active profiler. + +Note that `CUDA.@device_code_sass` also works for cuTile kernels: it +intercepts module loads at the driver level (CUPTI), showing the binary that +was actually loaded — along with Julia source locations, since `tileiras` +compiles with line info. + +# Example +```julia +@device_code_sass @cuda backend=cuTile blocks=grid vadd(a, b, c) +``` +""" +macro device_code_sass(ex...) + function hook(f, tt; io::IO=stdout, kwargs...) + println(io, "// $f($(join(tt.parameters, ", ")))") + println(io) + code_sass(io, f, Tuple(tt.parameters); kwargs...) + println(io) + end + emit_hooked_compilation(hook, ex...) +end + """ @device_code_structured [io=stdout] expression diff --git a/src/launch.jl b/src/launch.jl index 19e8b29a..e15b26b2 100644 --- a/src/launch.jl +++ b/src/launch.jl @@ -404,6 +404,142 @@ function run_tileiras(bytecode::Vector{UInt8}, sm_arch::VersionNumber, end end +# Minimal ELF64 section lookup. Binutils does not recognize the CUBIN machine +# type, so locate the section by parsing the headers directly. Returns the +# section contents, or `nothing` if absent or not a little-endian ELF64 file. +function elf_section(elf::Vector{UInt8}, name::String) + read_at(T, pos) = reinterpret(T, @view elf[pos+1:pos+sizeof(T)])[] + length(elf) >= 64 || return nothing + (view(elf, 1:4) == b"\x7fELF" && elf[5] == 2 && elf[6] == 1) || return nothing + shoff = Int(read_at(UInt64, 0x28)) + shentsize = Int(read_at(UInt16, 0x3a)) + shnum = Int(read_at(UInt16, 0x3c)) + shstrndx = Int(read_at(UInt16, 0x3e)) + (shstrndx < shnum && shoff + shnum * shentsize <= length(elf)) || return nothing + strtab = Int(read_at(UInt64, shoff + shstrndx * shentsize + 0x18)) + for i in 0:shnum-1 + header = shoff + i * shentsize + name_pos = strtab + Int(read_at(UInt32, header)) + name_end = findnext(iszero, elf, name_pos + 1) + name_end === nothing && continue + String(elf[name_pos+1:name_end-1]) == name || continue + offset = Int(read_at(UInt64, header + 0x18)) + size = Int(read_at(UInt64, header + 0x20)) + offset + size <= length(elf) || return nothing + return elf[offset+1:offset+size] + end + return nothing +end + +""" + extract_ptx(cubin::Vector{UInt8}) -> String + +The PTX source embedded in `cubin`. [`run_tileiras`](@ref) compiles with +`--lineinfo`, which makes the assembler record the PTX it consumed in a +`.nv_debug_ptx_txt` section (used by `cuda-gdb` to display PTX-level source). +The section stores one NUL-terminated entry per PTX line, with comment and +`.loc` lines blanked and leading whitespace stripped; runs of blank lines are +collapsed and the original indentation is reconstructed to match what +`tileiras` wrote before embedding: one tab per brace level in PTX code with +labels and debug-`.section` contents flush-left, and — in the +`.nv_intermediate_source_section` that CUDA 13.4+ emits — two-space +scaffolding around a verbatim Tile IR block, identical to +[`code_tiled`](@ref) output. +""" +function extract_ptx(cubin::Vector{UInt8}) + section = elf_section(cubin, ".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(replace(String(section), '\0' => '\n'), '\n') + io = IOBuffer() + depth = 0 + mlir_base = -1 # depth of `.nv_intermediate_source_section`, or -1 + src_base = -1 # depth of `.source_begin`, or -1 outside the block + section_base = -1 # depth of a debug `.section` block, or -1 + pending_section = false + blank_run = true # suppress leading blanks; collapse runs + for line in lines + if isempty(line) + blank_run || println(io) + blank_run = true + continue + end + blank_run = false + closer = startswith(line, '}') || startswith(line, ')') + d = max(closer ? depth - 1 : depth, 0) + opens = count(==('{'), line) + # Parentheses nest only as function-header parameter lists: a line + # ending in `(` opens one, a line starting with `)` closes it. + depth += opens - count(==('}'), line) + + (endswith(line, '(') ? 1 : 0) - (startswith(line, ')') ? 1 : 0) + + indent = if mlir_base >= 0 + src_base >= 0 ? " "^max(d - src_base, 0) : + " "^max(d - mlir_base, 0) + elseif pending_section || startswith(line, ".section") || + (section_base >= 0 && closer) + # NVVM writes the `.section` directive and its braces with one + # tab, the block contents flush-left. + "\t" + elseif section_base >= 0 || occursin(r"^[\w\$]+:$", line) + "" + else + "\t"^d + end + println(io, indent, line) + + if startswith(line, ".nv_intermediate_source_section") + mlir_base = d + elseif mlir_base >= 0 + line == ".source_begin" && (src_base = d) + line == ".source_end" && (src_base = -1) + depth <= mlir_base && (mlir_base = -1; src_base = -1) + elseif pending_section || startswith(line, ".section") + if opens > 0 + section_base = d + pending_section = false + else + pending_section = startswith(line, ".section") + end + elseif section_base >= 0 && depth <= section_base + section_base = -1 + end + end + return rstrip(String(take!(io)), '\n') * "\n" +end + +""" + 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) diff --git a/test/codegen/reflection.jl b/test/codegen/reflection.jl index 049b5a90..95ad5b0e 100644 --- a/test/codegen/reflection.jl +++ b/test/codegen/reflection.jl @@ -72,6 +72,23 @@ 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) + end +end + if ct.tileiras_version() >= v"13.4" @testset "code_tiled remarks" begin output = sprint(io -> ct.code_tiled(io, reflect_vadd, TT3; From 03d25ca2253c39e63eb45ceedaff11f228fecc72 Mon Sep 17 00:00:00 2001 From: AntonOresten Date: Wed, 26 Aug 2026 18:12:01 +0200 Subject: [PATCH 2/9] Use ObjectFile.jl for the CUBIN section lookup and stop reformatting PTX --- Project.toml | 2 + docs/src/man/debugging.md | 3 +- src/compiler/reflection.jl | 4 +- src/launch.jl | 114 ++++++------------------------------- 4 files changed, 21 insertions(+), 102 deletions(-) diff --git a/Project.toml b/Project.toml index 9429a68e..d3563ab5 100644 --- a/Project.toml +++ b/Project.toml @@ -17,6 +17,7 @@ EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56" GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" 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" @@ -43,6 +44,7 @@ GPUArrays = "11" GPUToolbox = "3" IRStructurizer = "0.6.3" Microfloats = "0.2" +ObjectFile = "0.5" PrecompileTools = "1" Preferences = "1.5" REPL = "1" diff --git a/docs/src/man/debugging.md b/docs/src/man/debugging.md index 9e28c8a2..89148ecd 100644 --- a/docs/src/man/debugging.md +++ b/docs/src/man/debugging.md @@ -66,8 +66,7 @@ code, disassembled with `nvdisasm`. Both compile the Tile IR with `tileiras` but do not need a GPU: `sm_arch` may be omitted when a CUDA device is available, and must be given explicitly otherwise. -The PTX is recovered from a debug section `tileiras` embeds in the CUBIN, in -which comment lines are blanked; the code itself is unaffected. +The PTX is read from a debug section `tileiras` embeds in the CUBIN. 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 diff --git a/src/compiler/reflection.jl b/src/compiler/reflection.jl index 5ae347d3..aad1f1ad 100644 --- a/src/compiler/reflection.jl +++ b/src/compiler/reflection.jl @@ -230,8 +230,8 @@ thread-level SIMT program the tile-level kernel is lowered to, with every compiler decision (thread mapping, CTA size, pipelining, synchronization) already made. When no GPU is available, pass `sm_arch` explicitly. -The PTX is recovered from the `.nv_debug_ptx_txt` section the assembler embeds -in the CUBIN when compiling with line info, so header comments are blanked. +The PTX is read from the `.nv_debug_ptx_txt` section the assembler embeds in +the CUBIN when compiling with line info. """ function code_ptx(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) cubin = compile_to_cubin(f, argtypes; kwargs...) diff --git a/src/launch.jl b/src/launch.jl index e15b26b2..76536a8c 100644 --- a/src/launch.jl +++ b/src/launch.jl @@ -5,6 +5,7 @@ # CUDA context, and launches it via `cudacall`. Compilation is cached per # `(MethodInstance, sm_arch, opt_level, num_ctas, occupancy, num_worker_warps, bytecode_version)`. +using ObjectFile: readmeta, Sections using CUDACore: CUDACore, CuArray, CuModule, CuFunction, cudacall, device, capability, AbstractBackend, AbstractKernel, kernel_convert, kernel_compile using CUDA_Compiler_jll @@ -404,110 +405,27 @@ function run_tileiras(bytecode::Vector{UInt8}, sm_arch::VersionNumber, end end -# Minimal ELF64 section lookup. Binutils does not recognize the CUBIN machine -# type, so locate the section by parsing the headers directly. Returns the -# section contents, or `nothing` if absent or not a little-endian ELF64 file. -function elf_section(elf::Vector{UInt8}, name::String) - read_at(T, pos) = reinterpret(T, @view elf[pos+1:pos+sizeof(T)])[] - length(elf) >= 64 || return nothing - (view(elf, 1:4) == b"\x7fELF" && elf[5] == 2 && elf[6] == 1) || return nothing - shoff = Int(read_at(UInt64, 0x28)) - shentsize = Int(read_at(UInt16, 0x3a)) - shnum = Int(read_at(UInt16, 0x3c)) - shstrndx = Int(read_at(UInt16, 0x3e)) - (shstrndx < shnum && shoff + shnum * shentsize <= length(elf)) || return nothing - strtab = Int(read_at(UInt64, shoff + shstrndx * shentsize + 0x18)) - for i in 0:shnum-1 - header = shoff + i * shentsize - name_pos = strtab + Int(read_at(UInt32, header)) - name_end = findnext(iszero, elf, name_pos + 1) - name_end === nothing && continue - String(elf[name_pos+1:name_end-1]) == name || continue - offset = Int(read_at(UInt64, header + 0x18)) - size = Int(read_at(UInt64, header + 0x20)) - offset + size <= length(elf) || return nothing - return elf[offset+1:offset+size] - end - return nothing -end - """ - extract_ptx(cubin::Vector{UInt8}) -> String - -The PTX source embedded in `cubin`. [`run_tileiras`](@ref) compiles with -`--lineinfo`, which makes the assembler record the PTX it consumed in a -`.nv_debug_ptx_txt` section (used by `cuda-gdb` to display PTX-level source). -The section stores one NUL-terminated entry per PTX line, with comment and -`.loc` lines blanked and leading whitespace stripped; runs of blank lines are -collapsed and the original indentation is reconstructed to match what -`tileiras` wrote before embedding: one tab per brace level in PTX code with -labels and debug-`.section` contents flush-left, and — in the -`.nv_intermediate_source_section` that CUDA 13.4+ emits — two-space -scaffolding around a verbatim Tile IR block, identical to -[`code_tiled`](@ref) output. + 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(cubin::Vector{UInt8}) - section = elf_section(cubin, ".nv_debug_ptx_txt") +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(replace(String(section), '\0' => '\n'), '\n') - io = IOBuffer() - depth = 0 - mlir_base = -1 # depth of `.nv_intermediate_source_section`, or -1 - src_base = -1 # depth of `.source_begin`, or -1 outside the block - section_base = -1 # depth of a debug `.section` block, or -1 - pending_section = false - blank_run = true # suppress leading blanks; collapse runs - for line in lines - if isempty(line) - blank_run || println(io) - blank_run = true - continue - end - blank_run = false - closer = startswith(line, '}') || startswith(line, ')') - d = max(closer ? depth - 1 : depth, 0) - opens = count(==('{'), line) - # Parentheses nest only as function-header parameter lists: a line - # ending in `(` opens one, a line starting with `)` closes it. - depth += opens - count(==('}'), line) + - (endswith(line, '(') ? 1 : 0) - (startswith(line, ')') ? 1 : 0) - - indent = if mlir_base >= 0 - src_base >= 0 ? " "^max(d - src_base, 0) : - " "^max(d - mlir_base, 0) - elseif pending_section || startswith(line, ".section") || - (section_base >= 0 && closer) - # NVVM writes the `.section` directive and its braces with one - # tab, the block contents flush-left. - "\t" - elseif section_base >= 0 || occursin(r"^[\w\$]+:$", line) - "" - else - "\t"^d - end - println(io, indent, line) - - if startswith(line, ".nv_intermediate_source_section") - mlir_base = d - elseif mlir_base >= 0 - line == ".source_begin" && (src_base = d) - line == ".source_end" && (src_base = -1) - depth <= mlir_base && (mlir_base = -1; src_base = -1) - elseif pending_section || startswith(line, ".section") - if opens > 0 - section_base = d - pending_section = false - else - pending_section = startswith(line, ".section") - end - elseif section_base >= 0 && depth <= section_base - section_base = -1 - end - end - return rstrip(String(take!(io)), '\n') * "\n" + 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 From 0808d8cd1956f2e04a9f88436e44ea46f3aaae68 Mon Sep 17 00:00:00 2001 From: AntonOresten Date: Sat, 29 Aug 2026 13:46:32 +0200 Subject: [PATCH 3/9] Document PTX reflection as unstable, drop @device_code_sass PTX is an implementation detail of `tileiras`, read from an undocumented CUBIN section; `code_ptx`/`@device_code_ptx` say so and go away with it. cuTile's `@device_code_sass` duplicated `CUDA.@device_code_sass` with a narrower scope (only cuTile kernels, by recompiling each signature). The CUDA macro already shows every loaded module, cuTile's included; `code_sass` remains for inspecting a signature without a device or under a profiler. The `@device_code_*` hook now fires once from the launch's compile entry point rather than from inside the cached emit chain, so the chain checks its own cached field before recursing and subprogram compilation needs no hook suppression. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XNHajvrhhKBBUF7w6cDiKr --- docs/src/lib/reflection.md | 7 ++-- docs/src/man/debugging.md | 26 +++++++------- src/compiler/codegen/kernel.jl | 10 +----- src/compiler/driver.jl | 65 +++++++++++++--------------------- src/compiler/reflection.jl | 42 ++++++---------------- src/launch.jl | 19 +++++----- 6 files changed, 62 insertions(+), 107 deletions(-) diff --git a/docs/src/lib/reflection.md b/docs/src/lib/reflection.md index cdd03fc9..34df5fdb 100644 --- a/docs/src/lib/reflection.md +++ b/docs/src/lib/reflection.md @@ -22,9 +22,8 @@ code_sass @device_code_typed @device_code_structured @device_code_ptx -@device_code_sass ``` -`CUDA.@device_code_sass` also 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. +`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 89148ecd..2e3daf8f 100644 --- a/docs/src/man/debugging.md +++ b/docs/src/man/debugging.md @@ -66,7 +66,10 @@ code, disassembled with `nvdisasm`. Both compile the Tile IR with `tileiras` but do not need a GPU: `sm_arch` may be omitted when a CUDA device is available, and must be given explicitly otherwise. -The PTX is read from a debug section `tileiras` embeds in the CUBIN. +!!! 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 @@ -97,27 +100,24 @@ intercepted kernel: ct.@device_code_tiled remarks=true @cuda backend=cuTile blocks=grid kernel(args...) ``` -Five are available, corresponding to successive stages of the pipeline: +Four are available, corresponding to successive stages of the pipeline: | 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_structured` | Structured IR (after control-flow structurization) | -| `ct.@device_code_ptx` | PTX generated by `tileiras` | -| `ct.@device_code_sass` | SASS disassembly of the compiled CUBIN | +| `ct.@device_code_ptx` | PTX generated by `tileiras` (unstable, as above) | 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, then the PTX it lowers to and the machine code -that assembles to. - -`CUDA.@device_code_sass` also 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. The cuTile variant instead recompiles each intercepted signature, -which requires no CUPTI subscription and therefore also works under Nsight or -another active profiler. +flow, then emitted Tile IR, then the PTX it lowers to. + +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..cd71985d 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!(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..0853c808 100644 --- a/src/compiler/driver.jl +++ b/src/compiler/driver.jl @@ -222,16 +222,32 @@ to resolve them on every phase. end 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. +# Cached wrappers around the driver's emit_* functions: check/populate +# CuTileResults. Used by the launch path; reflection calls the driver directly. + +# The `(f, tt)` a launch compiles, with `Constant{T,V}` argument types restored +# from `const_argtypes` (the MethodInstance has them unwrapped), for the +# `@device_code_*` hook. +function hook_signature(mi::Core.MethodInstance, const_argtypes::Union{Vector{Any}, Nothing}) + ftype = mi.specTypes.parameters[1] + f = isdefined(ftype, :instance) ? ftype.instance : ftype + arg_types = collect(Any, mi.specTypes.parameters[2:end]) + if const_argtypes !== nothing + # const_argtypes is [Const(f), arg2, ...]; arg_types omits f. + for i in eachindex(arg_types) + if const_argtypes[i+1] isa CC.Const + arg_types[i] = typeof(Constant(const_argtypes[i+1].val)) + end + end + end + return f, Tuple{arg_types...} +end """ emit_structured!(cache, mi, ci, res; const_argtypes=nothing) -> (StructuredIRCode, rettype, kernel_meta) -Cached IR phase. Invokes the compile hook (for `@device_code_*` macros), -checks `res.julia_ir`, and delegates to `emit_structured` on cache miss. +Cached IR phase. Checks `res.julia_ir` and delegates to `emit_structured` on +cache miss. `ci` and `res` must come from `ensure_compiled(cache, mi, const_argtypes)` — i.e. inference is already done. Reflection callers that want the lookup @@ -241,32 +257,6 @@ overload. 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 - 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 @@ -289,20 +279,15 @@ end """ emit_tile!(cache, mi, ci, res; const_argtypes=nothing) -> Vector{UInt8} -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. +Cached code phase. Checks `res.tile_bc`; on cache miss, takes the IR from +`emit_structured!` and calls the driver's `emit_tile`. """ function emit_tile!(cache::CacheView, mi::Core.MethodInstance, ci::Core.CodeInstance, res::CuTileResults; const_argtypes::Union{Vector{Any}, Nothing}=nothing) - # 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 + sci, rettype, kernel_meta = emit_structured!(cache, mi, ci, res; const_argtypes) key = cache.owner::TileCacheKey opts = CGOpts((sm_arch=unpack_version(key.sm_arch), opt_level=unpack_hint(key.opt_level), diff --git a/src/compiler/reflection.jl b/src/compiler/reflection.jl index aad1f1ad..c1354786 100644 --- a/src/compiler/reflection.jl +++ b/src/compiler/reflection.jl @@ -230,8 +230,10 @@ thread-level SIMT program the tile-level kernel is lowered to, with every compiler decision (thread mapping, CTA size, pipelining, synchronization) already made. When no GPU is available, pass `sm_arch` explicitly. -The PTX is read from the `.nv_debug_ptx_txt` section the assembler embeds in -the CUBIN when compiling with line info. +!!! warning "Unstable" + PTX is an implementation detail of `tileiras`, not an interface. It is read + from an undocumented debug section of the CUBIN and may stop being produced + or recorded at any time; `code_ptx` will be removed with it. """ function code_ptx(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) cubin = compile_to_cubin(f, argtypes; kwargs...) @@ -247,6 +249,10 @@ code_ptx(@nospecialize(f), @nospecialize(argtypes); kwargs...) = Print the SASS machine code that a Julia function compiles to, by assembling the Tile IR with `tileiras` and disassembling the resulting CUBIN with `nvdisasm`. When no GPU is available, pass `sm_arch` explicitly. + +For the binary a launch actually loaded, use `CUDA.@device_code_sass`. Unlike +that, this needs no CUPTI subscription, so it also works under Nsight or +another active profiler. """ function code_sass(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) cubin = compile_to_cubin(f, argtypes; kwargs...) @@ -262,7 +268,7 @@ code_sass(@nospecialize(f), @nospecialize(argtypes); kwargs...) = export @device_code_tiled public @device_code_typed, @device_code_structured -public @device_code_ptx, @device_code_sass +public @device_code_ptx # Following GPUCompiler's pattern for @device_code_* macros function emit_hooked_compilation(inner_hook, ex...) @@ -323,7 +329,7 @@ end @device_code_ptx [io=stdout] expression Print the PTX generated by `tileiras` for all kernels compiled while -evaluating the expression. +evaluating the expression. Unstable, like [`code_ptx`](@ref). # Example ```julia @@ -340,34 +346,6 @@ macro device_code_ptx(ex...) emit_hooked_compilation(hook, ex...) end -""" - @device_code_sass [io=stdout] expression - -Print the SASS machine code for all kernels compiled while evaluating the -expression, by recompiling each intercepted signature and disassembling the -result with `nvdisasm`. Works without a CUPTI subscription, so it can be used -under Nsight or another active profiler. - -Note that `CUDA.@device_code_sass` also works for cuTile kernels: it -intercepts module loads at the driver level (CUPTI), showing the binary that -was actually loaded — along with Julia source locations, since `tileiras` -compiles with line info. - -# Example -```julia -@device_code_sass @cuda backend=cuTile blocks=grid vadd(a, b, c) -``` -""" -macro device_code_sass(ex...) - function hook(f, tt; io::IO=stdout, kwargs...) - println(io, "// $f($(join(tt.parameters, ", ")))") - println(io) - code_sass(io, f, Tuple(tt.parameters); kwargs...) - println(io) - end - emit_hooked_compilation(hook, ex...) -end - """ @device_code_structured [io=stdout] expression diff --git a/src/launch.jl b/src/launch.jl index 76536a8c..dda1ec00 100644 --- a/src/launch.jl +++ b/src/launch.jl @@ -645,18 +645,15 @@ Cached binary phase: compile Tile IR bytecode to CUBIN using tileiras. 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 + bytecode = emit_tile!(cache, mi, ci, res; const_argtypes) + sm_arch = unpack_version(cache.owner.sm_arch) # Resolve opt_level here (not in emit_tile) because it's a tileiras flag, not bytecode. # num_ctas/occupancy/num_worker_warps are resolved in emit_tile because they're encoded in bytecode. - _, _, kernel_meta = res.julia_ir + _, _, kernel_meta = emit_structured!(cache, mi, ci, res; const_argtypes) opt_level = something(resolve_hint(unpack_hint(cache.owner.opt_level), kernel_meta, :opt_level, sm_arch), 3) @@ -787,9 +784,13 @@ function compile(@nospecialize(f), @nospecialize(argtypes), # 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). + # Report the compilation to the `@device_code_*` hook. This runs through + # `invoke_frozen`, but the hook closure lives in the user's latest world — + # `invokelatest` it. + if compile_hook[] !== nothing + Base.invokelatest(compile_hook[], hook_signature(mi, const_argtypes)...) + end + emit_binary!(cache, mi, ci, res; const_argtypes) return cache, mi, ci, res end From 3f5a2c3508f2f23687350d6bc6e34a21bab2f618 Mon Sep 17 00:00:00 2001 From: AntonOresten Date: Tue, 8 Sep 2026 13:20:46 +0200 Subject: [PATCH 4/9] Use compilation jobs for caching and reflection Share inference across targets and hints, storing compiled binaries per configuration on the generic or constant-specialized inference result. Cache loaded kernels per CUDA context and synchronize result publication without holding a lock during compilation or module loading. Pass the launch job to task-scoped reflection hooks so inspection preserves its target, constants, and options, including on cache hits. Co-authored-by: Tim Besard --- docs/src/lib/reflection.md | 7 + docs/src/man/debugging.md | 23 +- src/compiler/codegen/kernel.jl | 2 +- src/compiler/driver.jl | 502 +++++++++++++----------- src/compiler/reflection.jl | 346 +++++++--------- src/cuTile.jl | 4 +- src/launch.jl | 304 +++++++------- src/precompile.jl | 6 +- test/codegen/boundscheck_integration.jl | 15 +- test/codegen/operations.jl | 2 +- test/codegen/reflection.jl | 140 +++++++ test/device/compilation.jl | 60 +++ test/host/cache.jl | 18 +- test/host/inference_cache.jl | 53 ++- 14 files changed, 828 insertions(+), 654 deletions(-) create mode 100644 test/device/compilation.jl diff --git a/docs/src/lib/reflection.md b/docs/src/lib/reflection.md index 34df5fdb..82dca1f4 100644 --- a/docs/src/lib/reflection.md +++ b/docs/src/lib/reflection.md @@ -15,6 +15,13 @@ code_ptx code_sass ``` +## Jobs + +```@docs +TileJob +tile_job +``` + ## Intercepting a launch ```@docs diff --git a/docs/src/man/debugging.md b/docs/src/man/debugging.md index 2e3daf8f..84b8803c 100644 --- a/docs/src/man/debugging.md +++ b/docs/src/man/debugging.md @@ -36,6 +36,20 @@ 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. +`code_tiled` compiles for the active CUDA device by default, so its output +matches what a launch emits. Without a device, pass `sm_arch` explicitly; the +Tile IR is otherwise still generated, but `ct.@compiler_options` hints are +ignored. Explicit hint keywords still apply. + +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: @@ -104,14 +118,15 @@ Four are available, corresponding to successive stages of the pipeline: | 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_structured` | Structured IR (after control-flow structurization) | +| `ct.@device_code_tiled` | Final Tile IR (MLIR textual format) | | `ct.@device_code_ptx` | PTX generated by `tileiras` (unstable, as above) | -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, then the PTX it lowers to. +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 printed 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 diff --git a/src/compiler/codegen/kernel.jl b/src/compiler/codegen/kernel.jl index cd71985d..f99361c4 100644 --- a/src/compiler/codegen/kernel.jl +++ b/src/compiler/codegen/kernel.jl @@ -387,7 +387,7 @@ function emit_subprogram!(ctx::CGCtx, func, arg_types::Vector, if !haskey(ctx.cache, mi) error("Expected $func($(join(arg_types, ", "))) to be cached already by inference.") end - sci, _, _ = emit_structured!(ctx.cache, mi) + sci, _, _ = emit_structured(emit_julia(ctx.cache, mi)...) # 2b. Run the pass pipeline on subprogram IR sub_divby, sub_bounds = run_passes!(sci) diff --git a/src/compiler/driver.jl b/src/compiler/driver.jl index 0853c808..b12da4e6 100644 --- a/src/compiler/driver.jl +++ b/src/compiler/driver.jl @@ -1,145 +1,227 @@ -# Compilation hook for @device_code_* macros - intercepts compilations for reflection -const compile_hook = Ref{Union{Nothing,Function}}(nothing) - +using Base.ScopedValues: ScopedValue, with #============================================================================= - 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 + # `nothing` for a job created without a CUDA device: the stages before + # `tileiras` still run, but architecture-dependent `@compiler_options` + # hints cannot be resolved and are ignored. + sm_arch::Union{VersionNumber, Nothing} 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: + +- `sm_arch`: the target architecture. Defaults to the active CUDA device's, or + `nothing` without one; then only the stages before `tileiras` are available + and architecture-dependent `@compiler_options` hints are ignored. +- `bytecode_version`: the Tile IR bytecode version to emit. +- `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] +function tile_job(mi::MethodInstance, world::UInt; + const_argtypes::Union{Tuple, Nothing}=nothing, + sm_arch::Union{VersionNumber, Nothing}=nothing, + bytecode_version::VersionNumber=cuTile.bytecode_version(), + 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) + sm_arch = @something sm_arch device_sm_arch() Some(nothing) + sm_arch === nothing || 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 + +# The architecture `tileiras` assembles a job for. +function target_arch(job::TileJob) + sm_arch = job.config.target.sm_arch + sm_arch === nothing && throw(ArgumentError( + "the job has no target architecture (it was created without a CUDA device); " * + "pass `sm_arch` explicitly")) + return sm_arch +end + +# 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 meta + return f, Tuple{arg_types...} end -""" - resolve_hint(explicit, kernel_meta, key, sm_arch) - -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=", something(target.sm_arch, "nothing"), + ", 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 phases + Compilation hook + + `@device_code_*` observe compilations through a task-scoped hook, called + with the job of every kernel that is compiled or launched while it is set. =============================================================================# +const compile_hook = ScopedValue{Union{Nothing, Function}}(nothing) + +# 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 = compile_hook[] + hook === nothing || Base.invokelatest(hook, job) + return +end + + +#============================================================================= + 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,162 +231,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, opt_level) +end - return bytecode +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 + +""" + compile(job::TileJob) -> Vector{UInt8} + +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. +""" +function compile(job::TileJob) + sm_arch = target_arch(job) + validate_tileiras_target(job.config.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 compilation + Meta nodes and compilation hints =============================================================================# -# Results struct for caching compilation phases -mutable struct CuTileResults - julia_ir::Any # (StructuredIRCode, rettype) - tile_bc::Any # Vector{UInt8} bytecode - cuda_bin::Any # Vector{UInt8} cubin - cuda_func::Any # CuFunction - tile_kernel::Any # TileKernel{F, tt} wrapper around cuda_func - CuTileResults() = new(nothing, nothing, nothing, nothing, nothing) -end - """ - ensure_compiled(cache, mi, const_argtypes) -> (ci::CodeInstance, res::CuTileResults) + process_meta!(ir::CC.IRCode) -> ir -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. +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`). """ -@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)) +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 end -# Cached wrappers around the driver's emit_* functions: check/populate -# CuTileResults. Used by the launch path; reflection calls the driver directly. +""" + extract_meta(ir::CC.IRCode) -> Dict{Symbol, Any} -# The `(f, tt)` a launch compiles, with `Constant{T,V}` argument types restored -# from `const_argtypes` (the MethodInstance has them unwrapped), for the -# `@device_code_*` hook. -function hook_signature(mi::Core.MethodInstance, const_argtypes::Union{Vector{Any}, Nothing}) - ftype = mi.specTypes.parameters[1] - f = isdefined(ftype, :instance) ? ftype.instance : ftype - arg_types = collect(Any, mi.specTypes.parameters[2:end]) - if const_argtypes !== nothing - # const_argtypes is [Const(f), arg2, ...]; arg_types omits f. - for i in eachindex(arg_types) - if const_argtypes[i+1] isa CC.Const - arg_types[i] = typeof(Constant(const_argtypes[i+1].val)) - end +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 f, Tuple{arg_types...} + return meta end """ - emit_structured!(cache, mi, ci, res; const_argtypes=nothing) -> (StructuredIRCode, rettype, kernel_meta) - -Cached IR phase. Checks `res.julia_ir` and delegates to `emit_structured` on -cache miss. + resolve_hint(explicit, kernel_meta, key, sm_arch) -`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. +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_structured!(cache::CacheView, mi::Core.MethodInstance, - ci::Core.CodeInstance, res::CuTileResults; - const_argtypes::Union{Vector{Any}, Nothing}=nothing) - res.julia_ir !== nothing && return res.julia_ir - - # Compute fresh via driver - ir, rettype = emit_julia(cache, mi; const_argtypes) - result = emit_structured(ir, rettype) - res.julia_ir = result - return result -end - -# Convenience overload that resolves (ci, res) on the caller's behalf. Use -# from non-hot paths (e.g. sub-program compilation in codegen/kernel.jl); -# the production launch path goes through `ensure_compiled` once at the top -# and feeds (ci, res) through the chain explicitly. -function emit_structured!(cache::CacheView, mi::Core.MethodInstance; - const_argtypes::Union{Vector{Any}, Nothing}=nothing) - ci, res = ensure_compiled(cache, mi, const_argtypes) - return emit_structured!(cache, mi, ci, res; const_argtypes) +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 -""" - emit_tile!(cache, mi, ci, res; const_argtypes=nothing) -> Vector{UInt8} +# Dump bytecode to `$JULIA_CUTILE_DUMP_BYTECODE/.ln[.n].cutile`, if set. +const bytecode_dump_lock = ReentrantLock() -Cached code phase. Checks `res.tile_bc`; on cache miss, takes the IR from -`emit_structured!` and calls the driver's `emit_tile`. -""" -function emit_tile!(cache::CacheView, mi::Core.MethodInstance, - ci::Core.CodeInstance, res::CuTileResults; - const_argtypes::Union{Vector{Any}, Nothing}=nothing) - res.tile_bc !== nothing && return res.tile_bc - - sci, rettype, kernel_meta = emit_structured!(cache, mi, ci, res; const_argtypes) - key = cache.owner::TileCacheKey - opts = CGOpts((sm_arch=unpack_version(key.sm_arch), - opt_level=unpack_hint(key.opt_level), - num_ctas=unpack_hint(key.num_ctas), - occupancy=unpack_hint(key.occupancy), - num_worker_warps=unpack_hint(key.num_worker_warps), - bytecode_version=unpack_version(key.bytecode_version))) - 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 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) @@ -314,7 +360,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 c1354786..69c05a76 100644 --- a/src/compiler/reflection.jl +++ b/src/compiler/reflection.jl @@ -21,20 +21,6 @@ function disassemble_tileir(bytecode::Vector{UInt8}, version::VersionNumber; end end -""" - code_typed(f, argtypes; world, kwargs...) -> Vector{Any} - -Return typed code for a cuTile function. Analogous to `Base.code_typed`. -""" -function code_typed(@nospecialize(f), @nospecialize(argtypes); - world::UInt=Base.get_world_counter(), kwargs...) - stripped, const_argtypes = process_const_argtypes(f, argtypes) - mi = lookup_method_instance(f, stripped; world) - cache = CacheView{CuTileResults}(:cuTile, world) - ir, rettype = emit_julia(cache, mi; const_argtypes) - [ir => rettype] -end - """ code_ircode(mi::MethodInstance; world, always_inline=true) -> (IRCode, rettype) @@ -43,8 +29,7 @@ If always_inline=true (default), forces all functions to be inlined. """ function code_ircode(mi::MethodInstance; world::UInt=Base.get_world_counter(), always_inline::Bool=true) - cache = CacheView{CuTileResults}(:cuTile, world) - interp = cuTileInterpreter(cache; always_inline) + interp = cuTileInterpreter(inference_cache(world); always_inline) result = CC.typeinf_ircode(interp, mi, nothing) if result === nothing @@ -55,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. """ @@ -91,78 +57,88 @@ 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 + + +#============================================================================= + Stages +=============================================================================# + """ - code_tiled([io::IO], f, argtypes; sm_arch, opt_level, num_ctas, occupancy, - num_worker_warps, remarks=false) + 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). +""" +function code_typed(job::TileJob) + ir, rettype = emit_julia(job) + [ir => rettype] +end +code_typed(@nospecialize(f), @nospecialize(argtypes); kwargs...) = + code_typed(tile_job(f, argtypes; kwargs...)) + +""" + 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). +""" +function code_structured(job::TileJob; optimize::Bool=true) + ir, rettype = emit_julia(job) + sci, rettype, _ = emit_structured(ir, rettype) + if optimize + sci = copy(sci) + run_passes!(sci) + end + [sci => rettype] +end +code_structured(@nospecialize(f), @nospecialize(argtypes); optimize::Bool=true, kwargs...) = + code_structured(tile_job(f, argtypes; kwargs...); optimize) + +""" + code_tiled([io::IO], 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): +without a CUDA device, pass `sm_arch` explicitly to resolve +architecture-dependent `@compiler_options` hints as a launch would. 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, and a target architecture. """ -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, target_arch(job), opt_level; remarks=true) if !isempty(text) println(io) println(io, "// tileiras optimization remarks") @@ -172,92 +148,48 @@ 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...) """ - compile_to_cubin(f, argtypes; sm_arch, opt_level, num_ctas, occupancy, - num_worker_warps, bytecode_version, world) -> Vector{UInt8} - -Compile a Julia function all the way to a CUBIN for reflection: emit Tile IR -bytecode (like [`code_tiled`](@ref)) and assemble it with `tileiras`. Like the -rest of the reflection entry points, this calls the driver directly without -caching in CuTileResults or the disk cache. -""" -function compile_to_cubin(@nospecialize(f), @nospecialize(argtypes); - sm_arch::Union{VersionNumber, Nothing}=nothing, - opt_level::Union{Int, Nothing}=nothing, - num_ctas::Union{Int, Nothing}=nothing, - occupancy::Union{Int, Nothing}=nothing, - num_worker_warps::Union{Int, Nothing}=nothing, - bytecode_version::VersionNumber=cuTile.bytecode_version(), - world::UInt=Base.get_world_counter()) - target = if sm_arch === nothing - try - default_sm_arch() - catch - throw(ArgumentError( - "sm_arch must be specified when compiling without a CUDA device")) - end - else - sm_arch - end - validate_tileiras_target(bytecode_version) - - stripped, const_argtypes = process_const_argtypes(f, argtypes) - mi = lookup_method_instance(f, stripped; world) - opts = CGOpts((sm_arch=target, opt_level=opt_level, num_ctas=num_ctas, - occupancy=occupancy, num_worker_warps=num_worker_warps, - bytecode_version=bytecode_version)) - cache = CacheView{CuTileResults}(:cuTile, world) - ir, rettype = emit_julia(cache, mi; const_argtypes) - sci, rettype, kernel_meta = emit_structured(ir, rettype) - bytecode = emit_tile(sci, rettype, kernel_meta; - name=sanitize_name(string(mi.def.name)), - opts, cache, const_argtypes) - resolved_opt_level = something(resolve_hint(opt_level, kernel_meta, - :opt_level, target), 3) - cubin, _ = run_tileiras(bytecode, target, resolved_opt_level) - return cubin -end - -""" - code_ptx([io::IO], f, argtypes; sm_arch, opt_level, num_ctas, occupancy, - num_worker_warps) + 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. When no GPU is available, pass `sm_arch` explicitly. +already made. Keyword arguments are those of [`tile_job`](@ref); when no GPU is +available, pass `sm_arch` explicitly. !!! warning "Unstable" - PTX is an implementation detail of `tileiras`, not an interface. It is read - from an undocumented debug section of the CUBIN and may stop being produced - or recorded at any time; `code_ptx` will be removed with it. + The PTX is recorded by `tileiras` in an undocumented CUBIN section and may + go away. """ -function code_ptx(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) - cubin = compile_to_cubin(f, argtypes; kwargs...) - print(io, extract_ptx(cubin)) -end +code_ptx(io::IO, job::TileJob) = print(io, extract_ptx(compile(job))) +code_ptx(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) = + code_ptx(io, tile_job(f, argtypes; kwargs...)) +code_ptx(job::TileJob) = code_ptx(stdout, job) code_ptx(@nospecialize(f), @nospecialize(argtypes); kwargs...) = code_ptx(stdout, f, argtypes; kwargs...) """ - code_sass([io::IO], f, argtypes; sm_arch, opt_level, num_ctas, occupancy, - num_worker_warps) + 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`. When no GPU is available, pass `sm_arch` explicitly. - -For the binary a launch actually loaded, use `CUDA.@device_code_sass`. Unlike -that, this needs no CUPTI subscription, so it also works under Nsight or -another active profiler. +`nvdisasm`. Keyword arguments are those of [`tile_job`](@ref); when no GPU is +available, pass `sm_arch` explicitly. For the binary a launch actually loaded, +use `CUDA.@device_code_sass`. """ -function code_sass(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) - cubin = compile_to_cubin(f, argtypes; kwargs...) - print(io, disassemble_cubin(cubin)) -end +code_sass(io::IO, job::TileJob) = print(io, disassemble_cubin(compile(job))) +code_sass(io::IO, @nospecialize(f), @nospecialize(argtypes); kwargs...) = + code_sass(io, tile_job(f, argtypes; kwargs...)) +code_sass(job::TileJob) = code_sass(stdout, job) code_sass(@nospecialize(f), @nospecialize(argtypes); kwargs...) = code_sass(stdout, f, argtypes; kwargs...) @@ -270,40 +202,48 @@ export @device_code_tiled public @device_code_typed, @device_code_structured public @device_code_ptx -# Following GPUCompiler's pattern for @device_code_* macros +# Install `inner_hook` as the compile hook for the duration of the expression, +# called once per distinct job. function emit_hooked_compilation(inner_hook, ex...) user_code = ex[end] user_kwargs = ex[1:end-1] quote - seen = Set{Tuple{Any,Any}}() - function outer_hook(f, tt) - if !in((f, tt), seen) - old_hook = $compile_hook[] - try - $compile_hook[] = nothing - $inner_hook(f, tt; $(map(esc, user_kwargs)...)) - finally - $compile_hook[] = old_hook + # The job set is shared with any child task compiling under the hook. + jobs = Set{TileJob}() + jobs_lock = ReentrantLock() + function outer_hook(job::TileJob) + Base.@lock jobs_lock begin + job in jobs && return + push!(jobs, job) + # the user hook might invoke the compiler again, so disable the hook + $with($compile_hook => nothing) do + $inner_hook(job; $(map(esc, user_kwargs)...)) end - push!(seen, (f, tt)) end end - try - $compile_hook[] = outer_hook + $with($compile_hook => outer_hook) do $(esc(user_code)) - finally - $compile_hook[] = nothing end - if isempty(seen) + if isempty(jobs) error("no kernels executed while evaluating the given expression") end - nothing end end +# 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 + """ @device_code_tiled [io=stdout] [remarks=false] expression @@ -316,72 +256,52 @@ With `remarks=true`, also print `tileiras` optimization remarks for each kernel. ``` """ macro device_code_tiled(ex...) - function hook(f, tt; io::IO=stdout, kwargs...) - println(io, "// $f($(join(tt.parameters, ", ")))") - println(io) - code_tiled(io, f, Tuple(tt.parameters); kwargs...) - println(io) - end + hook = tile_hook((io, job; kwargs...) -> code_tiled(io, job; kwargs...)) emit_hooked_compilation(hook, ex...) end """ - @device_code_ptx [io=stdout] expression + @device_code_structured [io=stdout] [optimize=true] expression -Print the PTX generated by `tileiras` for all kernels compiled while -evaluating the expression. Unstable, like [`code_ptx`](@ref). +Print the StructuredIRCode for all kernels compiled while evaluating the expression. # Example ```julia -@device_code_ptx @cuda backend=cuTile blocks=grid vadd(a, b, c) +@device_code_structured @cuda backend=cuTile blocks=grid vadd(a, b, c) ``` """ -macro device_code_ptx(ex...) - function hook(f, tt; io::IO=stdout, kwargs...) - println(io, "// $f($(join(tt.parameters, ", ")))") - println(io) - code_ptx(io, f, Tuple(tt.parameters); kwargs...) - println(io) - end +macro device_code_structured(ex...) + hook = tile_hook((io, job; kwargs...) -> println(io, first(only(code_structured(job; kwargs...))))) emit_hooked_compilation(hook, ex...) end """ - @device_code_structured [io=stdout] expression + @device_code_typed [io=stdout] expression -Print the StructuredIRCode for all kernels compiled while evaluating the expression. +Print the typed Julia IR for all kernels compiled while evaluating the expression. # Example ```julia -@device_code_structured @cuda backend=cuTile blocks=grid vadd(a, b, c) +@device_code_typed @cuda backend=cuTile blocks=grid vadd(a, b, c) ``` """ -macro device_code_structured(ex...) - function hook(f, tt; io::IO=stdout, kwargs...) - println(io, "// $f($(join(tt.parameters, ", ")))") - println(io) - sci, _ = only(code_structured(f, Tuple(tt.parameters); kwargs...)) - println(io, sci) - end +macro device_code_typed(ex...) + hook = tile_hook((io, job) -> println(io, first(only(code_typed(job))))) emit_hooked_compilation(hook, ex...) end """ - @device_code_typed [io=stdout] expression + @device_code_ptx [io=stdout] expression -Print the typed Julia IR for all kernels compiled while evaluating the expression. +Print the PTX generated by `tileiras` for all kernels compiled while +evaluating the expression. Unstable, like [`code_ptx`](@ref). # Example ```julia -@device_code_typed @cuda backend=cuTile blocks=grid vadd(a, b, c) +@device_code_ptx @cuda backend=cuTile blocks=grid vadd(a, b, c) ``` """ -macro device_code_typed(ex...) - function hook(f, tt; io::IO=stdout, kwargs...) - println(io, "// $f($(join(tt.parameters, ", ")))") - println(io) - ci, _ = only(code_typed(f, Tuple(tt.parameters); kwargs...)) - println(io, ci) - end +macro device_code_ptx(ex...) + hook = tile_hook((io, job) -> code_ptx(io, job)) emit_hooked_compilation(hook, ex...) end diff --git a/src/cuTile.jl b/src/cuTile.jl index 51b79867..6e4e96e1 100644 --- a/src/cuTile.jl +++ b/src/cuTile.jl @@ -22,8 +22,8 @@ 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 # Shared definitions include("shapes.jl") diff --git a/src/launch.jl b/src/launch.jl index dda1ec00..59d25d4a 100644 --- a/src/launch.jl +++ b/src/launch.jl @@ -2,12 +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 ObjectFile: readmeta, Sections -using CUDACore: CUDACore, CuArray, CuModule, CuFunction, cudacall, device, capability, - AbstractBackend, AbstractKernel, kernel_convert, kernel_compile +using ObjectFile: ObjectFile, readmeta, Sections +using CUDACore: CUDACore, CuArray, CuModule, CuFunction, CuContext, context, cudacall, + device, capability, AbstractBackend, AbstractKernel, kernel_convert, + kernel_compile using CUDA_Compiler_jll using GPUToolbox: LazyInitialized using Preferences: @load_preference @@ -80,53 +81,6 @@ CUDACore.kernel_compile(::TileBackend, f::F, tt::TT=Tuple{}; kwargs...) where {F cufunction(f, tt; kwargs...) -#============================================================================= - Cache sharding key. -=============================================================================# - -# Pack a `VersionNumber` into a `UInt16` as `(major << 8) | minor`. Lossy: drops -# patch/prerelease/build (none of which we need for SM architectures or Tile IR -# bytecode versions). Used by `TileCacheKey` to keep the owner isbits. -@inline pack_version(v::VersionNumber) = (UInt16(v.major) << 8) | UInt16(v.minor) -@inline unpack_version(x::UInt16) = VersionNumber(Int(x >> 8), Int(x & 0xff)) - -# isbits sentinel codec for `Union{Int, Nothing}` hint fields (`opt_level`, -# `num_ctas`, `occupancy`, `num_worker_warps`). `-1` is unused as a value, so we -# use it for `nothing`. -const _UNSET = -1 -@inline pack_hint(x::Union{Int, Nothing}) = x === nothing ? _UNSET : x -@inline unpack_hint(x::Int) = x == _UNSET ? nothing : x - -""" - TileCacheKey - -Owner stamped onto every launch-cached `CodeInstance`. `lookup` (called once -per `cufunction` via `ensure_compiled`) ccalls `jl_rettype_inferred` with the -owner as `Any`, which forces a heap box; making the owner `isbits` shrinks the -box from ~80 B (the previous `Tuple{Symbol, NamedTuple{…}}` shape, dominated by -`VersionNumber`'s non-isbits prerelease/build tuples) to ~32 B. - -`VersionNumber` fields are packed into `UInt16` and `Union{Int, Nothing}` -hint fields use `_UNSET` (`-1`) as the `nothing` sentinel. Decoding back to -the original types happens once per cache miss in `emit_binary!` / -`emit_tile!`, never on the hot lookup path. -""" -struct TileCacheKey - sm_arch::UInt16 - bytecode_version::UInt16 - opt_level::Int - num_ctas::Int - occupancy::Int - num_worker_warps::Int -end -TileCacheKey(sm_arch::VersionNumber, bytecode_version::VersionNumber, - opt_level::Union{Int, Nothing}, num_ctas::Union{Int, Nothing}, - occupancy::Union{Int, Nothing}, num_worker_warps::Union{Int, Nothing}) = - TileCacheKey(pack_version(sm_arch), pack_version(bytecode_version), - pack_hint(opt_level), pack_hint(num_ctas), pack_hint(occupancy), - pack_hint(num_worker_warps)) - - #============================================================================= Toolchain and target validation. =============================================================================# @@ -150,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) @@ -558,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 @@ -573,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) @@ -602,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) @@ -629,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 @@ -638,58 +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) - res.cuda_bin !== nothing && return res.cuda_bin - - bytecode = emit_tile!(cache, mi, ci, res; const_argtypes) - - sm_arch = unpack_version(cache.owner.sm_arch) - - # Resolve opt_level here (not in emit_tile) because it's a tileiras flag, not bytecode. - # num_ctas/occupancy/num_worker_warps are resolved in emit_tile because they're encoded in bytecode. - _, _, kernel_meta = emit_structured!(cache, mi, ci, res; const_argtypes) - opt_level = something(resolve_hint(unpack_hint(cache.owner.opt_level), - kernel_meta, :opt_level, sm_arch), 3) - +function assemble(bytecode::Vector{UInt8}, sm_arch::VersionNumber, opt_level::Int) # Cross-session cache of the tileiras output. The key covers every input # that changes the CUBIN: bytecode, sm_arch, opt_level and the tileiras # identity, so different compiler builds never collide. `bytecode_version` # is encoded in the bytecode itself, so it's covered transitively. CUBIN is # address-free, so it is always persistable. - res.cuda_bin = if disk_cache_enabled - ObjCache.get!(CUBIN_CACHE_NS, cubin_cache_fields(bytecode, sm_arch, opt_level)...; - schema=CUBIN_CACHE_SCHEMA, persistable=true) do - first(run_tileiras(bytecode, sm_arch, opt_level)) - end - else + disk_cache_enabled || return first(run_tileiras(bytecode, sm_arch, opt_level)) + return ObjCache.get!(CUBIN_CACHE_NS, cubin_cache_fields(bytecode, sm_arch, opt_level)...; + schema=CUBIN_CACHE_SCHEMA, persistable=true) do first(run_tileiras(bytecode, sm_arch, opt_level)) end +end - return res.cuda_bin +# The 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 -GPU-side link phase: load the CUBIN cached in `res.cuda_bin` (produced by -[`compile`](@ref)) onto the active CUDA device and return the resulting -`CuFunction`. +The results slot of a job, created empty on first access. Runs inference if the +job's source has not been inferred yet. """ -function link(cache::CacheView, mi::Core.MethodInstance, - ci::Core.CodeInstance, res::CuTileResults) - res.cuda_func !== nothing && return res.cuda_func +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 - 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 +""" + 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 + +""" + link(job::TileJob, res::TileCompilerResults) -> CuFunction + +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 @@ -698,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. =============================================================================# """ @@ -723,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, @@ -735,37 +729,23 @@ 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) + sm_arch = @something sm_arch default_sm_arch() - # 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)) @@ -773,41 +753,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) - - # Report the compilation to the `@device_code_*` hook. This runs through - # `invoke_frozen`, but the hook closure lives in the user's latest world — - # `invokelatest` it. - if compile_hook[] !== nothing - Base.invokelatest(compile_hook[], hook_signature(mi, const_argtypes)...) - end - - emit_binary!(cache, mi, ci, res; const_argtypes) - return cache, mi, ci, res + return tile_job(mi, world; const_argtypes, kwargs...) end # Inner compilation routine; called via `invoke_frozen` so its method dispatches # happen in the world captured at __init__, reusing precompiled native code # even when later-loaded packages would otherwise have invalidated it. function cufunction_compile(@nospecialize(f), @nospecialize(tt), @nospecialize(argtypes), - const_argtypes::Union{Vector{Any}, Nothing}, - key::TileCacheKey) - cache, mi, ci, res = compile(f, argtypes, const_argtypes, key) - - cufunc = link(cache, mi, ci, res) - - res.tile_kernel !== nothing && return res.tile_kernel::TileKernel{Core.Typeof(f), tt} - kernel = TileKernel{Core.Typeof(f), tt}(f, cufunc) - res.tile_kernel = kernel + const_argtypes::Union{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 @@ -948,6 +920,12 @@ Returns e.g. `v"12.0"` for compute capability 12.0. """ default_sm_arch() = capability(device()) +# The compute capability of the current CUDA device, or `nothing` without one. +function device_sm_arch() + CUDACore.functional() && CUDACore.ndevices() > 0 || return nothing + return default_sm_arch() +end + #============================================================================= Version reporting diff --git a/src/precompile.jl b/src/precompile.jl index d613251b..33e4a9e5 100644 --- a/src/precompile.jl +++ b/src/precompile.jl @@ -36,22 +36,20 @@ import REPL return end - # Drive the host-side `compile` so the precompile workload reaches + # Drive the host-side `compile_or_lookup` so the precompile workload reaches # `tileiras` → CUBIN without needing a CUDA context. The runtime path # tacks `link` on top to load the CUBIN onto the GPU. function precompile_kernel(@nospecialize(f), @nospecialize(tt)) # Without `tileiras` (e.g. a pre-13.2 CUDA_Compiler_jll) kernels cannot # compile; still allow the package itself to precompile. tileiras_available() || return - argtypes, const_argtypes = unwrap_argtypes(f, tt) bv = bytecode_version() for sm_arch in [v"8.0", v"8.6", v"8.7", v"8.9", v"10.0", v"11.0", v"12.0", v"12.1"] requirement = tile_ir_requirement(sm_arch) requirement === nothing && continue bv < requirement[2] && continue - key = TileCacheKey(sm_arch, bv, nothing, nothing, nothing, nothing) - compile(f, argtypes, const_argtypes, key) + compile_or_lookup(tile_job(f, Tuple(tt.parameters); sm_arch, bytecode_version=bv)) end return end diff --git a/test/codegen/boundscheck_integration.jl b/test/codegen/boundscheck_integration.jl index 3c68a932..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..c3baa9a6 100644 --- a/test/codegen/operations.jl +++ b/test/codegen/operations.jl @@ -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)) diff --git a/test/codegen/reflection.jl b/test/codegen/reflection.jl index 95ad5b0e..159203b8 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" @@ -86,6 +104,128 @@ if ct.tileiras_available() output = sprint(io -> ct.code_sass(io, reflect_vadd, TT3; sm_arch=v"10.0")) @test occursin(r"\.target\s+sm_100", output) @test occursin(".text.reflect_vadd", output) + job = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0") + stdout_output = capture_stdout(() -> ct.code_sass(job)) + @test stdout_output == output + end + + @testset "compile(job)" begin + 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) + + # A job created without a target can run every stage before tileiras. + targetless = ct.TileJob(job.source, nothing, job.world, + ct.TileConfig(ct.TileCompilerTarget(nothing, job.config.target.bytecode_version), + job.config.params, job.config.name)) + @test occursin("addf", sprint(ct.code_tiled, targetless)) + @test_throws ArgumentError ct.compile(targetless) + @test_throws ArgumentError ct.code_tiled(devnull, targetless; remarks=true) + + # Incompatible targets are rejected when the job is built. + @test_throws ArgumentError ct.tile_job(reflect_vadd, TT3; sm_arch=v"7.5") + end + + @testset "shared inference" begin + # Inference is independent of target and hints, so every Tile job shares + # one partition and one CodeInstance; codegen results are per job. + job1 = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0") + job2 = ct.tile_job(reflect_vadd, TT3; sm_arch=v"8.9", opt_level=1) + res1, res2 = ct.compile_or_lookup(job1), ct.compile_or_lookup(job2) + ci1, ci2 = 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 + + # The default target is the active device's, so reflection matches a launch. + if CUDA.functional() && !isempty(CUDA.devices()) + job = ct.tile_job(reflect_vadd, TT3) + @test job.config.target.sm_arch == ct.default_sm_arch() + @test job === ct.tile_job(reflect_vadd, TT3; sm_arch=ct.default_sm_arch()) + end + end + + @testset "compile hook" begin + job = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0") + ct.compile_or_lookup(job) + # the hook observes hits as well as misses, once per distinct job + seen = ct.TileJob[] + ct.with(ct.compile_hook => (job -> push!(seen, job))) do + ct.compile_or_lookup(job) + ct.compile_or_lookup(job) + end + @test seen == [job, job] + @test ct.compile_hook[] === nothing + + # 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_typed io=outer begin + ct.@device_code_typed 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 ct.compile_hook[] === nothing + @test_throws ErrorException ct.@device_code_typed error("reflection failed") + @test ct.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 ct.with(ct.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 diff --git a/test/device/compilation.jl b/test/device/compilation.jl new file mode 100644 index 00000000..b4c066c3 --- /dev/null +++ b/test/device/compilation.jl @@ -0,0 +1,60 @@ +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 "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 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 From b58e8a74ae104bf7a875004bc96365b79e75bbdf Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 8 Sep 2026 13:50:32 +0200 Subject: [PATCH 5/9] Share GPUCompiler's reflection protocol cuTile keeps its own compiler, but reports compilations through GPUCompiler.compile_hook instead of a hook of its own, and builds its @device_code_* macros on GPUCompiler.emit_hooked_compilation. TileJob implements GPUCompiler.code_typed and code_warntype, so GPUCompiler's (and CUDA.jl's) @device_code_typed and @device_code_warntype observe cuTile launches too, mixed-backend expressions included. Requires GPUCompiler 2.6 for the backend-agnostic hook. --- Project.toml | 2 ++ src/compiler/driver.jl | 11 ++++------- src/compiler/reflection.jl | 40 ++++++++++---------------------------- src/cuTile.jl | 4 ++++ test/Project.toml | 1 + test/codegen/reflection.jl | 22 +++++++++++++++------ test/runtests.jl | 2 ++ 7 files changed, 39 insertions(+), 43 deletions(-) diff --git a/Project.toml b/Project.toml index d3563ab5..b9260424 100644 --- a/Project.toml +++ b/Project.toml @@ -15,6 +15,7 @@ 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" @@ -41,6 +42,7 @@ CompilerCaching = "0.5" DLFP8Types = "0.1.0" EnumX = "1.0" GPUArrays = "11" +GPUCompiler = "2.6" GPUToolbox = "3" IRStructurizer = "0.6.3" Microfloats = "0.2" diff --git a/src/compiler/driver.jl b/src/compiler/driver.jl index b12da4e6..52366bc6 100644 --- a/src/compiler/driver.jl +++ b/src/compiler/driver.jl @@ -1,5 +1,3 @@ -using Base.ScopedValues: ScopedValue, with - #============================================================================= Compiler jobs @@ -132,16 +130,15 @@ end #============================================================================= Compilation hook - `@device_code_*` observe compilations through a task-scoped hook, called - with the job of every kernel that is compiled or launched while it is set. + `@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. =============================================================================# -const compile_hook = ScopedValue{Union{Nothing, Function}}(nothing) - # 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 = compile_hook[] + hook = GPUCompiler.compile_hook[] hook === nothing || Base.invokelatest(hook, job) return end diff --git a/src/compiler/reflection.jl b/src/compiler/reflection.jl index 69c05a76..d9994162 100644 --- a/src/compiler/reflection.jl +++ b/src/compiler/reflection.jl @@ -99,6 +99,14 @@ end code_typed(@nospecialize(f), @nospecialize(argtypes); kwargs...) = code_typed(tile_job(f, argtypes; kwargs...)) +# GPUCompiler's reflection protocol, for its `@device_code_*` macros: the job +# selects cuTile's inference partition (and const-seeded source) for the +# typed IR, and its interpreter for the warntype view of the generic source. +GPUCompiler.code_typed(job::TileJob) = code_typed(job) +GPUCompiler.code_warntype(io::IO, job::TileJob; kwargs...) = + GPUCompiler.code_warntype_by_type(io, job.source.specTypes; + interp=cuTileInterpreter(inference_cache(job)), kwargs...) + """ code_structured(job::TileJob; optimize=true) -> Vector{Pair{StructuredIRCode, DataType}} code_structured(f, argtypes; optimize=true, kwargs...) @@ -202,36 +210,8 @@ export @device_code_tiled public @device_code_typed, @device_code_structured public @device_code_ptx -# Install `inner_hook` as the compile hook for the duration of the expression, -# called once per distinct job. -function emit_hooked_compilation(inner_hook, ex...) - user_code = ex[end] - user_kwargs = ex[1:end-1] - quote - # The job set is shared with any child task compiling under the hook. - jobs = Set{TileJob}() - jobs_lock = ReentrantLock() - function outer_hook(job::TileJob) - Base.@lock jobs_lock begin - job in jobs && return - push!(jobs, job) - # the user hook might invoke the compiler again, so disable the hook - $with($compile_hook => nothing) do - $inner_hook(job; $(map(esc, user_kwargs)...)) - end - end - end - - $with($compile_hook => outer_hook) do - $(esc(user_code)) - end - - if isempty(jobs) - error("no kernels executed while evaluating the given expression") - end - nothing - end -end +# cuTile's stage macros are GPUCompiler's, over cuTile's `code_*` functions. +const emit_hooked_compilation = GPUCompiler.emit_hooked_compilation # A hook printing a signature header around `inner(io, job; kwargs...)`. function tile_hook(inner) diff --git a/src/cuTile.jl b/src/cuTile.jl index 6e4e96e1..d07aa72f 100644 --- a/src/cuTile.jl +++ b/src/cuTile.jl @@ -25,6 +25,10 @@ import CompilerCaching 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/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/reflection.jl b/test/codegen/reflection.jl index 159203b8..623ddba3 100644 --- a/test/codegen/reflection.jl +++ b/test/codegen/reflection.jl @@ -182,14 +182,24 @@ if ct.tileiras_available() @testset "compile hook" begin job = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0") ct.compile_or_lookup(job) - # the hook observes hits as well as misses, once per distinct job + # GPUCompiler's hook observes hits as well as misses, once per distinct job seen = ct.TileJob[] - ct.with(ct.compile_hook => (job -> push!(seen, job))) do + 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 ct.compile_hook[] === nothing + @test GPUCompiler.compile_hook[] === nothing + + # GPUCompiler's own macros accept Tile jobs through the shared protocol + typed = GPUCompiler.@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) # Nested macros restore the outer hook; child tasks inherit it, and # repeated jobs are printed only once in each scope. @@ -201,16 +211,16 @@ if ct.tileiras_available() end end @test String(take!(outer)) == String(take!(inner)) != "" - @test ct.compile_hook[] === nothing + @test GPUCompiler.compile_hook[] === nothing @test_throws ErrorException ct.@device_code_typed error("reflection failed") - @test ct.compile_hook[] === nothing + @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 ct.with(ct.compile_hook => (_ -> (put!(entered, nothing); take!(resume)))) do + late = @async with(GPUCompiler.compile_hook => (_ -> (put!(entered, nothing); take!(resume)))) do ct.compile_or_lookup(job) end take!(entered) 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 From c1629d168b00825fca59fefd709a86189ad84461 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 8 Sep 2026 13:52:07 +0200 Subject: [PATCH 6/9] Use GPUCompiler's @device_code_typed and @device_code_warntype With compilations reported through the shared hook, the Julia-level stages need no cuTile-specific macros: re-export GPUCompiler's, which behave like their CUDA.jl counterparts (@device_code_typed returns a dictionary from job to typed code), and keep cuTile's own macros for its own stages. --- docs/src/lib/reflection.md | 4 +++- docs/src/man/debugging.md | 12 ++++++++++-- src/compiler/reflection.jl | 23 ++++++----------------- test/codegen/reflection.jl | 12 +++++++----- test/device/core.jl | 10 +++++----- 5 files changed, 31 insertions(+), 30 deletions(-) diff --git a/docs/src/lib/reflection.md b/docs/src/lib/reflection.md index 82dca1f4..5780f849 100644 --- a/docs/src/lib/reflection.md +++ b/docs/src/lib/reflection.md @@ -26,11 +26,13 @@ tile_job ```@docs @device_code_tiled -@device_code_typed @device_code_structured @device_code_ptx ``` +`@device_code_typed` and `@device_code_warntype` are GPUCompiler's, re-exported; +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 84b8803c..e4c45178 100644 --- a/docs/src/man/debugging.md +++ b/docs/src/man/debugging.md @@ -114,15 +114,23 @@ intercepted kernel: ct.@device_code_tiled remarks=true @cuda backend=cuTile blocks=grid kernel(args...) ``` -Four are available, corresponding to successive stages of the pipeline: +Five are available, corresponding to successive stages of the pipeline: | Macro | Output | |-------|--------| -| `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) | | `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 GPUCompiler's macros, re-exported: 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. +Like their CUDA.jl counterparts, `@device_code_typed` returns a dictionary +from job to typed code rather than printing. + 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 printed once per macro invocation. Hooks are scoped to the current task diff --git a/src/compiler/reflection.jl b/src/compiler/reflection.jl index d9994162..9ddbd56e 100644 --- a/src/compiler/reflection.jl +++ b/src/compiler/reflection.jl @@ -207,8 +207,12 @@ code_sass(@nospecialize(f), @nospecialize(argtypes); kwargs...) = =============================================================================# export @device_code_tiled -public @device_code_typed, @device_code_structured -public @device_code_ptx +public @device_code_structured, @device_code_ptx + +# 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 # cuTile's stage macros are GPUCompiler's, over cuTile's `code_*` functions. const emit_hooked_compilation = GPUCompiler.emit_hooked_compilation @@ -255,21 +259,6 @@ macro device_code_structured(ex...) emit_hooked_compilation(hook, ex...) end -""" - @device_code_typed [io=stdout] expression - -Print the typed Julia IR for all kernels compiled while evaluating the expression. - -# Example -```julia -@device_code_typed @cuda backend=cuTile blocks=grid vadd(a, b, c) -``` -""" -macro device_code_typed(ex...) - hook = tile_hook((io, job) -> println(io, first(only(code_typed(job))))) - emit_hooked_compilation(hook, ex...) -end - """ @device_code_ptx [io=stdout] expression diff --git a/test/codegen/reflection.jl b/test/codegen/reflection.jl index 623ddba3..db08bc5d 100644 --- a/test/codegen/reflection.jl +++ b/test/codegen/reflection.jl @@ -191,8 +191,10 @@ if ct.tileiras_available() @test seen == [job, job] @test GPUCompiler.compile_hook[] === nothing - # GPUCompiler's own macros accept Tile jobs through the shared protocol - typed = GPUCompiler.@device_code_typed ct.compile_or_lookup(job) + # 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 @@ -204,15 +206,15 @@ if ct.tileiras_available() # 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_typed io=outer begin - ct.@device_code_typed io=inner ct.compile_or_lookup(job) + 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_typed error("reflection failed") + @test_throws ErrorException ct.@device_code_structured error("reflection failed") @test GPUCompiler.compile_hook[] === nothing 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 From d9240d90ca9ef1cf38fa9ad8cfe5b6159c99a16a Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 8 Sep 2026 14:01:14 +0200 Subject: [PATCH 7/9] Preserve job specialization in shared reflection Display warntype from the inferred TileJob source, including constant arguments. Let CUDA PTX reflection dispatch to cuTile, and filter non-Tile jobs from cuTile-specific stage macros. Test constant-specialized reflection and warm mixed CUDA/cuTile launches through the shared macros. --- docs/src/lib/reflection.md | 2 +- docs/src/man/debugging.md | 12 ++++++--- src/compiler/reflection.jl | 18 +++++++------ test/codegen/reflection.jl | 12 +++++++++ test/device/compilation.jl | 53 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 84 insertions(+), 13 deletions(-) diff --git a/docs/src/lib/reflection.md b/docs/src/lib/reflection.md index 5780f849..8ad493b3 100644 --- a/docs/src/lib/reflection.md +++ b/docs/src/lib/reflection.md @@ -30,7 +30,7 @@ tile_job @device_code_ptx ``` -`@device_code_typed` and `@device_code_warntype` are GPUCompiler's, re-exported; +`@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 diff --git a/docs/src/man/debugging.md b/docs/src/man/debugging.md index e4c45178..a2e7e315 100644 --- a/docs/src/man/debugging.md +++ b/docs/src/man/debugging.md @@ -114,7 +114,7 @@ intercepted kernel: ct.@device_code_tiled remarks=true @cuda backend=cuTile blocks=grid kernel(args...) ``` -Five are available, corresponding to successive stages of the pipeline: +The following macros are available: | Macro | Output | |-------|--------| @@ -124,16 +124,20 @@ Five are available, corresponding to successive stages of the pipeline: | `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 GPUCompiler's macros, re-exported: cuTile reports its +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. -Like their CUDA.jl counterparts, `@device_code_typed` returns a dictionary +`@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 printed once per macro invocation. Hooks are scoped to the current task +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 diff --git a/src/compiler/reflection.jl b/src/compiler/reflection.jl index 9ddbd56e..3e8c2b6e 100644 --- a/src/compiler/reflection.jl +++ b/src/compiler/reflection.jl @@ -99,13 +99,13 @@ end code_typed(@nospecialize(f), @nospecialize(argtypes); kwargs...) = code_typed(tile_job(f, argtypes; kwargs...)) -# GPUCompiler's reflection protocol, for its `@device_code_*` macros: the job -# selects cuTile's inference partition (and const-seeded source) for the -# typed IR, and its interpreter for the warntype view of the generic source. GPUCompiler.code_typed(job::TileJob) = code_typed(job) -GPUCompiler.code_warntype(io::IO, job::TileJob; kwargs...) = - GPUCompiler.code_warntype_by_type(io, job.source.specTypes; - interp=cuTileInterpreter(inference_cache(job)), kwargs...) +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}} @@ -181,6 +181,7 @@ 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...) @@ -214,8 +215,9 @@ public @device_code_structured, @device_code_ptx using GPUCompiler: @device_code_typed, @device_code_warntype public @device_code_typed, @device_code_warntype -# cuTile's stage macros are GPUCompiler's, over cuTile's `code_*` functions. -const emit_hooked_compilation = GPUCompiler.emit_hooked_compilation +# 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) # A hook printing a signature header around `inner(io, job; kwargs...)`. function tile_hook(inner) diff --git a/test/codegen/reflection.jl b/test/codegen/reflection.jl index db08bc5d..7eb35650 100644 --- a/test/codegen/reflection.jl +++ b/test/codegen/reflection.jl @@ -203,6 +203,18 @@ if ct.tileiras_available() @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() diff --git a/test/device/compilation.jl b/test/device/compilation.jl index b4c066c3..3a6b09a5 100644 --- a/test/device/compilation.jl +++ b/test/device/compilation.jl @@ -58,3 +58,56 @@ end @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 From 170244dde1d19b88b2a3efa7d865aed4bd599f6f Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 9 Sep 2026 14:47:25 +0200 Subject: [PATCH 8/9] Always target the current device unless an architecture is given Jobs could be built without a target, deferring the architecture to codegen and leaving reflection with a different notion of "current target" than launches. Give every job a target: `sm_arch` defaults to the current CUDA device, as a launch, and `bytecode_version` to the preference or the newest version the selected `tileiras` accepts. The pair is validated when the job is built, so an older bytecode version needs an explicit `sm_arch` on devices it cannot target, and without a CUDA device `sm_arch` must be given. Launches now reject an explicit `sm_arch` other than the device's, since `tileiras` emits architecture-specific code that only runs there. The CUDA 13.1 CI job is dropped: 13.1 bytecode cannot target the A100 it runs on, so no job in the suite can be built there. 13.1 bytecode emission stays covered by the explicit tests on newer toolchains, which now pin `sm_arch=v"10.0"` alongside the bytecode version. --- .buildkite/pipeline.yml | 1 - docs/src/man/debugging.md | 22 +++++++++++---------- src/compiler/driver.jl | 40 ++++++++++++++++---------------------- src/compiler/reflection.jl | 17 ++++++++-------- src/launch.jl | 25 +++++++++++++----------- test/codegen/operations.jl | 4 ++-- test/codegen/reflection.jl | 31 +++++++++++++++-------------- test/device/compilation.jl | 14 +++++++++++++ 8 files changed, 83 insertions(+), 71 deletions(-) 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/docs/src/man/debugging.md b/docs/src/man/debugging.md index a2e7e315..d99ae32a 100644 --- a/docs/src/man/debugging.md +++ b/docs/src/man/debugging.md @@ -36,10 +36,13 @@ 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. -`code_tiled` compiles for the active CUDA device by default, so its output -matches what a launch emits. Without a device, pass `sm_arch` explicitly; the -Tile IR is otherwise still generated, but `ct.@compiler_options` hints are -ignored. Explicit hint keywords still apply. +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: @@ -58,9 +61,8 @@ 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 @@ -77,8 +79,7 @@ ct.code_sass(vadd, argtypes; sm_arch=v"10.0") 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: `sm_arch` may be omitted when a CUDA device is -available, and must be given explicitly otherwise. +but do not need a GPU. !!! warning "PTX reflection is unstable" PTX is an implementation detail of `tileiras`, not an interface, and may @@ -87,7 +88,8 @@ available, and must be given explicitly otherwise. 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. +exact architecture. Launches therefore always compile for the active device, +and reject an explicit `sm_arch` that differs from it. ## Intercepting a launch diff --git a/src/compiler/driver.jl b/src/compiler/driver.jl index 52366bc6..958df1a2 100644 --- a/src/compiler/driver.jl +++ b/src/compiler/driver.jl @@ -8,10 +8,7 @@ public TileJob, tile_job struct TileCompilerTarget - # `nothing` for a job created without a CUDA device: the stages before - # `tileiras` still run, but architecture-dependent `@compiler_options` - # hints cannot be resolved and are ignored. - sm_arch::Union{VersionNumber, Nothing} + sm_arch::VersionNumber bytecode_version::VersionNumber end @@ -58,10 +55,11 @@ 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: -- `sm_arch`: the target architecture. Defaults to the active CUDA device's, or - `nothing` without one; then only the stages before `tileiras` are available - and architecture-dependent `@compiler_options` hints are ignored. -- `bytecode_version`: the Tile IR bytecode version to emit. +- `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. @@ -69,29 +67,25 @@ argument types. Keyword arguments configure the compilation: function tile_job(mi::MethodInstance, world::UInt; const_argtypes::Union{Tuple, Nothing}=nothing, sm_arch::Union{VersionNumber, Nothing}=nothing, - bytecode_version::VersionNumber=cuTile.bytecode_version(), + 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) - sm_arch = @something sm_arch device_sm_arch() Some(nothing) - sm_arch === nothing || validate_tile_ir_target(sm_arch, bytecode_version) + if sm_arch === nothing + CUDACore.functional() || throw(ArgumentError( + "No CUDA device to target; pass `sm_arch`")) + sm_arch = device_sm_arch() + end + 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 -# The architecture `tileiras` assembles a job for. -function target_arch(job::TileJob) - sm_arch = job.config.target.sm_arch - sm_arch === nothing && throw(ArgumentError( - "the job has no target architecture (it was created without a CUDA device); " * - "pass `sm_arch` explicitly")) - return sm_arch -end - # 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) @@ -116,7 +110,7 @@ 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=", something(target.sm_arch, "nothing"), + print(io, "; sm_arch=", target.sm_arch, ", bytecode_version=v\"", target.bytecode_version, "\"") for field in fieldnames(TileCompilerParams) hint = getfield(params, field) @@ -267,8 +261,8 @@ the object cache). Reports the job to the `@device_code_*` hook. Uncached; `compile_or_lookup` caches for launches. """ function compile(job::TileJob) - sm_arch = target_arch(job) - validate_tileiras_target(job.config.target.bytecode_version) + (; 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) diff --git a/src/compiler/reflection.jl b/src/compiler/reflection.jl index 3e8c2b6e..4f2da872 100644 --- a/src/compiler/reflection.jl +++ b/src/compiler/reflection.jl @@ -132,11 +132,11 @@ code_structured(@nospecialize(f), @nospecialize(argtypes); optimize::Bool=true, Print the CUDA Tile IR for a Julia function as a textual MLIR representation. Analogous to `code_llvm`. Keyword arguments are those of [`tile_job`](@ref): -without a CUDA device, pass `sm_arch` explicitly to resolve -architecture-dependent `@compiler_options` hints as a launch would. +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, and a target architecture. +This requires `tileiras` 13.4 or newer. """ function code_tiled(io::IO, job::TileJob; debuginfo::Bool=false, remarks::Bool=false) (; bytecode, opt_level) = emit_tile(job) @@ -146,7 +146,7 @@ function code_tiled(io::IO, job::TileJob; debuginfo::Bool=false, remarks::Bool=f tileiras_version() >= v"13.4" || throw(ArgumentError( "tileiras optimization remarks require tileiras 13.4 or newer")) validate_tileiras_target(bytecode_version) - _, text = run_tileiras(bytecode, target_arch(job), 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") @@ -170,8 +170,8 @@ code_tiled(@nospecialize(f), @nospecialize(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); when no GPU is -available, pass `sm_arch` explicitly. +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 @@ -191,9 +191,8 @@ code_ptx(@nospecialize(f), @nospecialize(argtypes); kwargs...) = Print the SASS machine code that a Julia function compiles to, by assembling the Tile IR with `tileiras` and disassembling the resulting CUBIN with -`nvdisasm`. Keyword arguments are those of [`tile_job`](@ref); when no GPU is -available, pass `sm_arch` explicitly. For the binary a launch actually loaded, -use `CUDA.@device_code_sass`. +`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...) = diff --git a/src/launch.jl b/src/launch.jl index 59d25d4a..f72c7810 100644 --- a/src/launch.jl +++ b/src/launch.jl @@ -729,7 +729,17 @@ 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) - sm_arch = @something sm_arch default_sm_arch() + # `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 # Unwrap Constant arguments for method lookup and seed their values into inference. argtypes, const_argtypes = unwrap_argtypes(f, tt) @@ -913,18 +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()) - -# The compute capability of the current CUDA device, or `nothing` without one. -function device_sm_arch() - CUDACore.functional() && CUDACore.ndevices() > 0 || return nothing - return default_sm_arch() -end +device_sm_arch() = capability(device()) #============================================================================= diff --git a/test/codegen/operations.jl b/test/codegen/operations.jl index c3baa9a6..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,)) @@ -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 7eb35650..435f987e 100644 --- a/test/codegen/reflection.jl +++ b/test/codegen/reflection.jl @@ -143,23 +143,31 @@ if ct.tileiras_available() @test occursin("Constant{Int64, 16}", sprint(show, const_job)) @test ct.compile_or_lookup(const_job) === ct.compile_or_lookup(const_job) - # A job created without a target can run every stage before tileiras. - targetless = ct.TileJob(job.source, nothing, job.world, - ct.TileConfig(ct.TileCompilerTarget(nothing, job.config.target.bytecode_version), - job.config.params, job.config.name)) - @test occursin("addf", sprint(ct.code_tiled, targetless)) - @test_throws ArgumentError ct.compile(targetless) - @test_throws ArgumentError ct.code_tiled(devnull, targetless; remarks=true) + # 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"8.9", opt_level=1) + 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 @@ -170,13 +178,6 @@ if ct.tileiras_available() job3 = ct.tile_job(reflect_vadd, TT3; sm_arch=v"10.0") @test job3.world > job1.world && job3.config === job1.config @test ct.compile_or_lookup(job3) === res1 - - # The default target is the active device's, so reflection matches a launch. - if CUDA.functional() && !isempty(CUDA.devices()) - job = ct.tile_job(reflect_vadd, TT3) - @test job.config.target.sm_arch == ct.default_sm_arch() - @test job === ct.tile_job(reflect_vadd, TT3; sm_arch=ct.default_sm_arch()) - end end @testset "compile hook" begin diff --git a/test/device/compilation.jl b/test/device/compilation.jl index 3a6b09a5..1ad87f9d 100644 --- a/test/device/compilation.jl +++ b/test/device/compilation.jl @@ -18,6 +18,20 @@ using CUDA @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() From b5f687caf4f8f573ee02612e73be1ee3d462602b Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 9 Sep 2026 14:47:25 +0200 Subject: [PATCH 9/9] Infer directly in the Julia-level reflection stages `code_typed(f, argtypes)` and `code_structured(f, argtypes)` built a `TileJob`, whose default bytecode version probes `tileiras`, before running inference. Without target keywords these stages only need the inferred code, so run inference directly and leave the job to the target-dependent stages. Neither a device nor `tileiras` is needed for them anymore. --- src/compiler/reflection.jl | 36 ++++++++++++++++++++++++++++++------ 1 file changed, 30 insertions(+), 6 deletions(-) diff --git a/src/compiler/reflection.jl b/src/compiler/reflection.jl index 4f2da872..bffd75d9 100644 --- a/src/compiler/reflection.jl +++ b/src/compiler/reflection.jl @@ -80,6 +80,17 @@ function tile_job(@nospecialize(f), @nospecialize(argtypes); 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 @@ -90,14 +101,19 @@ end 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). +Keyword arguments are those of [`tile_job`](@ref); without any, no toolchain +or device is needed. """ function code_typed(job::TileJob) ir, rettype = emit_julia(job) [ir => rettype] end -code_typed(@nospecialize(f), @nospecialize(argtypes); kwargs...) = - code_typed(tile_job(f, argtypes; kwargs...)) +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) @@ -112,10 +128,20 @@ end 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). +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) @@ -123,8 +149,6 @@ function code_structured(job::TileJob; optimize::Bool=true) end [sci => rettype] end -code_structured(@nospecialize(f), @nospecialize(argtypes); optimize::Bool=true, kwargs...) = - code_structured(tile_job(f, argtypes; kwargs...); optimize) """ code_tiled([io::IO], job::TileJob; debuginfo=false, remarks=false)