Skip to content

Add PTX reflection (unstable) - #302

Draft
AntonOresten wants to merge 3 commits into
mainfrom
ao/code-ptx-sass
Draft

Add PTX reflection (unstable)#302
AntonOresten wants to merge 3 commits into
mainfrom
ao/code-ptx-sass

Conversation

@AntonOresten

@AntonOresten AntonOresten commented Aug 21, 2026

Copy link
Copy Markdown
Collaborator

Draft: the functionality is complete and tested, but the launch-interception half deliberately vendors machinery that could be unified with GPUCompiler/CUDATools

I was pleasantly surprised to see that CUDA.@device_code_sass works even on the cuTile backend, and even more so that PTX is an interceptable intermediate of tileiras, 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 tileiras generates. ct.code_ptx / ct.code_sass compile a signature the same way code_tiled does and show the resulting PTX and SASS; ct.@device_code_ptx / ct.@device_code_sass intercept kernel launches, completing the pipeline walk: @device_code_typed@device_code_structured@device_code_tiled@device_code_ptx@device_code_sass.

ct.code_ptx(vadd, argtypes; sm_arch=v"10.0")   # works without a GPU: tileiras is an offline compiler
ct.code_sass(vadd, argtypes)                   # sm_arch defaults to the active device

How the PTX is obtained

tileiras offers no way to emit PTX, and its /tmp/mlir-kernels-*.ptx intermediates are unlinked on success. But run_tileiras always passes --lineinfo, which makes the assembler embed the complete PTX it consumed in a .nv_debug_ptx_txt section of the CUBIN (for cuda-gdb). code_ptx recovers 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_section embedding the Tile IR verbatim, so code_ptx output contains the same Tile IR code_tiled shows.

SASS comes from nvdisasm, which ships next to tileiras in CUDA_Compiler_jll (the tileiras preference override resolves a sibling binary, mirroring tileirdisasm discovery).

Both entry points call the compiler driver directly and never touch CuTileResults or 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_txt section is undocumented, but it's what cuda-gdb reads 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, shimming ptxas) stop working the day tileiras compiles in-process. If a toolkit ever stops embedding it, extract_ptx fails loudly and the tests flag it on the bump; the only assumption on our side is that run_tileiras keeps passing --lineinfo. The eventual clean fix is a --keep-style flag on tileiras itself (nvcc precedent), which would make the extraction redundant.

Interop with CUDATools

Two findings shaped the API:

  • CUDA.@device_code_sass already 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_sass is still included because it's semantically different.
  • CUDA.code_sass(f, types) does not work (InvalidIRError): the signature form hard-constructs a GPUCompiler CompilerJob and 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/_sass currently ride cuTile's private compile_hook, a mirror of GPUCompiler's pattern. That duplication can collapse, because GPUCompiler's hook machinery is target-agnostic — compile_hook receives a CompilerJob, and @device_code_native's hook just calls the generic code_native(io, job):

  1. cuTile speaks CompilerJob (needs no upstream changes): define TileCompilerTarget <: AbstractCompilerTarget + params carrying CGOpts, fire GPUCompiler.compile_hook[] from emit_structured!, and overload GPUCompiler.code_native(io, ::CompilerJob{TileCompilerTarget}) → this PR's compile_to_cubin + extract_ptx. CUDA.@device_code_ptx then 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.
  2. GPUCompiler blesses the extension point: compile_hook and the code_*(job) generics are unexported internals; a small docs+CI PR would make the pattern supportable.
  3. CUDATools, complementary: a CUPTI-side embedded-PTX mode — for every loaded module, print .nv_debug_ptx_txt when present — gives loaded-truth PTX for any producer, the same duality code_sass already has between its CompilerJob and CUPTI forms. A ~25-line PoC (mirroring code_sass(::Callable)) is verified working against a cuTile launch; extract_ptx would relocate next to disassemble_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_txt debug section and parses it properly.

@AntonOresten

AntonOresten commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator Author

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 ptxas subprocess, so an ACF is just an extra flag (ptx_options="--apply-controls=…"). tileiras assembles in-process and exposes no such hook, which made cuTile kernels look unreachable — but the embedded PTX from this PR is exactly the input ptxas consumed, so the loop becomes: run_tileirasextract_ptx → standalone ptxas --apply-controls → load the resulting CUBIN in place of tileiras's.

@maleadt

maleadt commented Aug 26, 2026

Copy link
Copy Markdown
Member

which makes the assembler embed the complete PTX it consumed in a .nv_debug_ptx_txt section of the CUBIN (for cuda-gdb)

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.

@maleadt

maleadt commented Aug 26, 2026

Copy link
Copy Markdown
Member

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 disassemble_cubin from CUDATools?

I'm also not particularly happy that we have two semantics of @device_code_sass now, the one here only disassembling cuTile kernels, while the CUDA one disassembling all...

@AntonOresten

Copy link
Copy Markdown
Collaborator Author

Addressed all three, and took a step further:

  • PTX unstable: documented on code_ptx and in the debugging manual as an implementation detail of tileiras, read from an undocumented debug section, removed with it.
  • Two @device_code_sass: cuTile's @device_code_ptx/@device_code_sass macros are gone. Instead, cuTile compilations are now GPUCompiler CompilerJobs (TileCompilerTarget/TileCompilerParams), and launches report their job to GPUCompiler.compile_hook from the single compile entry point, like GPUCompiler.compile does. So CUDA.@device_code_typed, @device_code_ptx (= code_native(::TileJob)) and @device_code_sass (CUPTI, the binary actually loaded) work on @cuda backend=cuTile launches with one semantics each; cuTile's remaining @device_code_tiled/typed/structured are thin wrappers over GPUCompiler.emit_hooked_compilation. compile_to_cubin became GPUCompiler.compile(:obj|:asm, ::TileJob).
  • disassemble_cubin: it's the same nvdisasm shell-out plus your @code_native-style cleanup, so yes. But cuTile only depends on CUDACore, and there's a more useful version of the question. cuTile's code_sass(f, tt; sm_arch) needs no device: tileiras and nvdisasm are host binaries, so it shows e.g. sm_100 SASS on an sm_89 box. CUDATools' code_sass(io, job::CompilerJob) still loads the module to capture the cubin via CUPTI, which predates CUDACore running ptxas itself; CUDACore.compile(job).image already is the cubin (checked: for a cuTile kernel, the CUPTI-captured module and the host cubin disassemble to an identical instruction stream). If that became disassemble_cubin(io, CUDACore.compile(job).image), it would be device-free for any arch and any backend that defines CUDACore.compile(job). cuTile would overload it for TileJob and drop its own code_sass/disassemble_cubin. (Today CUDA.code_sass(io, ::TileJob) dispatches to that generic method and fails inside CUDACore.compile.) Happy to do that as a CUDA.jl PR if you agree; kept cuTile's wrapper until then.

Not in this PR, but related to GPUCompiler: get_interpreter(job) still infers under cuTile's reflection partition, and launches still use cuTile's own cache key. Making the job the cache identity via cached_results/cache_owner (what CUDACore 6.3 did) is a follow-up; it also removes cuTile's per-hint re-inference (previous attempts: #95 (comment), JuliaGPU/GPUCompiler.jl#917). It will also need a small GPUCompiler change: can_persist_results requires a relocation lowering, which an LLVM-free backend never has, though its results are trivially session-portable.

@maleadt

maleadt commented Aug 28, 2026

Copy link
Copy Markdown
Member

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.

@AntonOresten AntonOresten changed the title Add PTX and SASS reflection Add PTX reflection (unstable) Aug 29, 2026
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants