Generated-code inspection¶

Import these helpers from tirx_harness.dump_kernel. They accept a compiled TVM executable or runtime module and expose CUDA source, Parallel Thread Execution (PTX) assembly, native GPU assembly (SASS), and compiler resource information. See the compilation example.

Extract code¶

tirx_harness.dump_kernel.dump_module(mod, *, arch: str | None = None, ptx: bool = False, sass: bool = True, lineinfo: bool = True, function: str | None = None, outdir: str | None = None, name: str = 'kernel', quiet: bool = True) → DumpResult¶

Dump CUDA / PTX / SASS from a compiled module (the tvm_tirx.compile output).

mod is the Executable returned by tvm_tirx.compile (or a runtime Module). CUDA is always returned; ptx/sass toggle those stages, which shell out to nvcc/ptxas/cuobjdump. arch defaults to the local GPU (B200 -> sm_100a). With outdir set, artifacts are written there as <name>.{cu,ptx,cubin,sass,ptxas.log} and their paths returned in .paths; otherwise a temp dir is used and removed before returning.

If a requested stage fails (bad arch, unknown function, etc.) its field stays None and a message is appended to .errors (.ok is then False); the failure is printed to stderr only when quiet=False.

Example:

ex = tvm_tirx.compile(tvm_tirx.IRModule({"main": pf}), target="cuda", pipeline="tirx")
d = dump_module(ex, ptx=True)
assert d.ok, d.errors
print(d.sass)                 # SASS text
print(d.ptxas["registers"])   # 153

The requested assembly stages require CUDA toolkit executables. Set arch explicitly when inspecting code for a GPU other than the local device. sass=True by default; use ptx=False, sass=False to extract only CUDA source, or call dump_cuda directly.

tirx_harness.dump_kernel.dump_cuda(mod) → str¶

Extract generated CUDA source from a compiled module (the tvm_tirx.compile output).

Accepts the Executable returned by tvm_tirx.compile (uses its .mod) or a runtime Module directly. The device CUDA lives in an imported submodule (the top-level host module doesn’t carry it), so walk the import tree and return the first source that actually contains a __global__ entry.

Results¶

class tirx_harness.dump_kernel.DumpResult¶

Result of dump_module() — text of each requested artifact + metadata.

cuda is always populated; ptx/sass only when requested and the nvcc/cuobjdump step succeeded — if a requested stage fails, its field stays None and a message is appended to errors (check d.ok / d.errors to tell “not requested” from “requested but failed”). ptxas holds the parsed register/spill/smem report and ptxas_log its raw text. paths maps artifact -> written file, populated only for files actually written (requires outdir).

cuda: str¶

arch: str¶

symbols: list[str] = []¶

ptx: str | None = None¶

sass: str | None = None¶

ptxas: dict¶

ptxas_log: str | None = None¶

paths: dict[str, str]¶

errors: list[str] = []¶

property ok: bool¶: True if no requested artifact stage failed.

ptxas is the CUDA assembler’s resource report. Its smem entry describes static shared memory; an absent entry does not imply zero dynamic shared-memory use. Check ok and errors before consuming a requested artifact. A failure to extract the initial CUDA source raises ValueError directly.

Inspect existing text¶

tirx_harness.dump_kernel.extract_symbols(cuda_src: str) → list[str]¶

Return the __global__ kernel symbol names in generated CUDA source.

The actual entry symbol comes from the PrimFunc’s global_symbol attribute (e.g. _kernel_kernel), NOT the IRModule key — so read it back from the emitted __global__ void [__launch_bounds__(...)] <name>( signature rather than guessing <key>_kernel.

tirx_harness.dump_kernel.parse_ptxas(log: str) → dict¶

Extract registers / spill / smem / barriers from a ptxas -v log.

Caveats: only the first entry function’s report is parsed (single-kernel modules are the common case; DumpResult.symbols lists all kernels). The smem key is present only for static shared memory — kernels using extern __shared__ dynamic SMEM (most TIRx GEMM/attention kernels) don’t report it, so its absence does not mean “0 bytes SMEM”.