Add PTX and SASS reflection - #302
Draft
AntonOresten wants to merge 1 commit into
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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) baselines with a compiler-generated generic fallback, all from one kernel definitionThis 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.