Add PTX reflection (unstable) - #302
Conversation
|
Interestingly this opens up the possibility of autotuning using CompileIQ (https://nvidia.github.io/CompileIQ/stable/index.html). There's precedence for this as they have a Triton example (https://nvidia.github.io/CompileIQ/stable/triton_example.html), difference being that Triton emits PTX as an explicit stage and assembles it with an external |
Yeah I don't know we want to rely on this, especially with that string parsing code being as messy as it is. Tile IR had been advertised to explicitly "bypass" PTX, so it may not be there to stay. |
faa8e2a to
26e7259
Compare
|
OK yeah much cleaner with ObjectFile. I'd also document that PTX reflection is unstable, and may disappear at any point (i.e. when NVIDIA likes it). And maybe we should re-use I'm also not particularly happy that we have two semantics of |
|
Addressed all three, and took a step further:
Not in this PR, but related to GPUCompiler: |
|
That's cute, but integration with GPUCompiler.jl (which should really be LLVMCompiler.jl or so, we just couldn't figure out a name) feels like the wrong thing to do. We share none of the compilation-related functionality, since GPUCompiler.jl is entirely targeted on integrating with Julia's codegen, which we don't reuse here. That's not to say there isn't room for such a package, possibly with some code lifted from GPUCompiler.jl, but I didn't want to think about such an abstraction at a point in time where there's only a single user, cuTile.jl. |
8ea591d to
29723b5
Compare
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XNHajvrhhKBBUF7w6cDiKr
I was pleasantly surprised to see that
CUDA.@device_code_sassworks even on the cuTile backend, and even more so that PTX is an interceptable intermediate oftileiras, as this could make cuTile good for e.g. per-(arch × shape-class) improvable baselines, from a single kernel definition.This PR adds two new stages to cuTile's reflection surface: what
tileirasgenerates.ct.code_ptx/ct.code_sasscompile a signature the same waycode_tileddoes and show the resulting PTX and SASS;ct.@device_code_ptx/ct.@device_code_sassintercept kernel launches, completing the pipeline walk:@device_code_typed→@device_code_structured→@device_code_tiled→@device_code_ptx→@device_code_sass.How the PTX is obtained
tileirasoffers no way to emit PTX, and its/tmp/mlir-kernels-*.ptxintermediates are unlinked on success. Butrun_tileirasalways passes--lineinfo, which makes the assembler embed the complete PTX it consumed in a.nv_debug_ptx_txtsection of the CUBIN (forcuda-gdb).code_ptxrecovers it from there — a small ELF section parse, no extra tileiras invocation, and it would work retroactively on any disk-cached CUBIN. The section stores NUL-separated lines with comments/.locs blanked and indentation stripped, so the printer reconstructs the original layout; verified byte-identical (modulo the blanked lines) against shim-captured genuine tileiras output. A nice property on 13.4: the PTX carries an.nv_intermediate_source_sectionembedding the Tile IR verbatim, so code_ptx output contains the same Tile IRcode_tiledshows.SASS comes from
nvdisasm, which ships next totileirasin CUDA_Compiler_jll (thetileiraspreference override resolves a sibling binary, mirroringtileirdisasmdiscovery).Both entry points call the compiler driver directly and never touch
CuTileResultsor the disk cache, like the rest of the reflection API. Everything is public rather than exported: CUDA.jl exports the same four names.Stability
The
.nv_debug_ptx_txtsection is undocumented, but it's whatcuda-gdbreads for PTX-level debugging, so breaking it could mean breaking their own debugger. It's also a property of the output artifact rather than of tileiras' process tree: the other routes to the PTX (racing the temp files, shimmingptxas) stop working the day tileiras compiles in-process. If a toolkit ever stops embedding it,extract_ptxfails loudly and the tests flag it on the bump; the only assumption on our side is thatrun_tileiraskeeps passing--lineinfo. The eventual clean fix is a--keep-style flag ontileirasitself (nvccprecedent), which would make the extraction redundant.Interop with CUDATools
Two findings shaped the API:
CUDA.@device_code_sassalready works on cuTile kernels. It intercepts module loads at the driver level (CUPTI), which is backend-agnostic by construction, and shows Julia source locations, since tileiras compiles with line info.ct.@device_code_sassis still included because it's semantically different.CUDA.code_sass(f, types)does not work (InvalidIRError): the signature form hard-constructs a GPUCompilerCompilerJoband compiles through the LLVM pipeline, where cuTile's intrinsics don't exist. The signature forms here fill that gap, including compiling for other architectures (sm_arch=v"10.0") from a machine with no GPU.Potential unification roadmap
ct.@device_code_ptx/_sasscurrently ride cuTile's privatecompile_hook, a mirror of GPUCompiler's pattern. That duplication can collapse, because GPUCompiler's hook machinery is target-agnostic —compile_hookreceives aCompilerJob, and@device_code_native's hook just calls the genericcode_native(io, job):CompilerJob(needs no upstream changes): defineTileCompilerTarget <: AbstractCompilerTarget+ params carryingCGOpts, fireGPUCompiler.compile_hook[]fromemit_structured!, and overloadGPUCompiler.code_native(io, ::CompilerJob{TileCompilerTarget})→ this PR'scompile_to_cubin+extract_ptx.CUDA.@device_code_ptxthen works on cuTile kernels verbatim, mixed-backend expressions included, and cuTile's private hook + macro copies get deleted. GPUCompiler is already in the dependency closure via CUDACore.compile_hookand thecode_*(job)generics are unexported internals; a small docs+CI PR would make the pattern supportable..nv_debug_ptx_txtwhen present — gives loaded-truth PTX for any producer, the same dualitycode_sassalready has between itsCompilerJoband CUPTI forms. A ~25-line PoC (mirroringcode_sass(::Callable)) is verified working against a cuTile launch;extract_ptxwould relocate next todisassemble_cubin.Until (1) lands, this PR's vendored hook is the working proof of concept.
Made with Claude Code
The pointer that CUBINs carry their PTX came from Patrick Toulmé's cuTile on Blackwell post, whose dump script recovers it with
strings; this PR identifies the underlying.nv_debug_ptx_txtdebug section and parses it properly.