Skip to content

Add PTX and SASS reflection - #302

Draft
AntonOresten wants to merge 1 commit into
mainfrom
ao/code-ptx-sass
Draft

Add PTX and SASS reflection#302
AntonOresten wants to merge 1 commit 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) baselines with a compiler-generated generic fallback, all from one 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.

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.

1 participant