The sass2mlir command set

One binary, thirteen subcommands. The general shape:

sass2mlir <command> <input> [flags]

The core loop is two commands: lift turns a cubin’s machine code into MLIR, apply runs pass plugins and patches the cubin in place. Around them: harvest to get kernels out of files, intercept to get them out of running processes, the analyses to see what you lifted, and roundtrip/crosscheck to prove nothing broke.

Global conventions used by every command that takes a cubin or MLIR module:

• --fn <glob> — restrict to kernels/functions matching the glob (e.g. --fn '*flash_fwd*'). Repeatable. Default: all functions.
• --arch <sm> — assert the input architecture (e.g. --arch sm_90a); the tool refuses to lift a cubin whose SM doesn’t match.
• -o <path> — output file or directory. Default: stdout for single outputs, alongside the input for multi-file outputs.
• -q — quiet: suppress the per-function progress lines, print only failures.

Bytes ↔ MLIR

sass2mlir lift  kernel.code -o kernel.mlir      # raw SASS bytes → sass dialect MLIR
sass2mlir lift  app.cubin   -o lifted/          # every kernel in a cubin
sass2mlir lower kernel.mlir -o kernel.code      # MLIR → raw SASS bytes

lift decodes the 128-bit instruction stream into sass.<MNEMONIC>_<opcode> ops with SSA operands and the authoritative $raw encoding on each op, and recovers control flow into MLIR blocks. lower re-encodes: untouched ops emit their $raw verbatim; any op a pass modified is re-encoded structurally from its fields by nvisa_encode.

Unrecognized (hidden/undocumented) instructions lift to the generic sass.inst carrier op — $raw only — and lower back in place, byte-for-byte.

The byte-identity check

sass2mlir roundtrip app.cubin
# roundtrip: sm_90a flash_fwd_kernel: 1042/1042 instructions byte-identical (0 mismatches)

Lift → print → reparse → lower → compare bytes against the input. This is the safety net the whole framework stands on; it exits non-zero on the first mismatched instruction and prints its offset, expected bytes, and actual bytes.

Structural re-encode audit

sass2mlir crosscheck kernel.mlir

Re-encodes every instruction from its structural fields alone — ignoring $raw — and compares against the $raw it actually carries. Where roundtrip proves the $raw path is exact, crosscheck proves the assembler path is: it’s the measurement that drove encode coverage from 86.4% of the corpus to full byte-identity. Any diff here means a field layout the encoder and decoder disagree on.

Cubin → .code + .mlir per kernel

sass2mlir harvest app.cubin -o out/
ls out/
# flash_fwd_kernel.code   flash_fwd_kernel.mlir
# softmax_kernel.code     softmax_kernel.mlir

Parses the cubin ELF, and for each kernel writes the raw instruction stream (.code) and the lifted module (.mlir). This is the offline half of kernel acquisition — use it when the cubin is already on disk. For kernels that only exist inside a running process, see intercept below.

cfg · defuse · liveness · uniform

All four take a lifted module and a function glob, and print to stdout. They only read the IR — none of them can invalidate a $raw.

sass2mlir cfg kernel.mlir --fn my_kernel
# bb0 (entry): 0x0000–0x0150  → succ: bb1 (taken), bb2 (fall-through)
# bb1:         0x0160–0x01a0  → succ: bb2
# bb2:         0x01b0–0x02f0  → succ: (exit)

sass2mlir cfg kernel.mlir --fn my_kernel --dot -o cfg.dot   # Graphviz output

sass2mlir defuse kernel.mlir --fn my_kernel
# %41 = sass.FFMA_0x223.FTZ ...   defs: {%41}  uses: {%27, %14, %40}
# reaching defs into bb2: {%38 (bb0), %41 (bb1), ...}

sass2mlir liveness kernel.mlir --fn my_kernel
# bb0: live-in {R2, R3, UR4}  live-out {R4, R5, P0} ...

sass2mlir uniform kernel.mlir --fn my_kernel
# uniform:   UR4 (blockDim), UR6 (param base), UP0 (bounds check) ...
# non-uniform: %27 (lane id), %41 (FFMA result) ...

cfg resolves absolute branch targets per instruction and classifies every terminator (unconditional / conditional / RET / EXIT / BSSY–BSYNC convergence pairs). uniform is conservative by construction: any uncertainty marks a value non-uniform, because a spurious “uniform” would be unsound — it’s the analysis the scalar-to-uniform pass trusts.

Chaitin-Briggs over recovered SSA

sass2mlir regalloc kernel.mlir --fn my_kernel --pinned
# regalloc(pinned): my_kernel — byte-identical allocation reproduced

sass2mlir regalloc kernel.mlir --fn my_kernel --free -o reallocated.mlir
# regalloc(free): my_kernel — 168 regs → 154 regs, contributor sets equivalent

Graph-coloring allocation over the recovered SSA, extended to the uniform register and predicate classes (UR/UP). Two modes:

• --pinned — must reproduce the byte-identical allocation the hardware was given. The hard corpus-wide oracle: it proves the liveness and interference model right, because any error shows up as a wrong register in the output.
• --free — real recoloring, validated by SSA contributor-set equivalence. Writes a new module with -o.

The engine: lift → passes → re-encode → patch

sass2mlir apply app.cubin \
  --passes build/lib/sass2mlir/passes \
  --config flash.toml \
  -o app.patched.cubin

Lifts the cubin, loads every .so plugin in --passes, applies them per the config, re-encodes, and patches the cubin in place — rewritten .text plus updated .nv.info and relocations — so the output is a self-contained binary in the original kernel slot, ABI untouched.

The config targets passes per architecture and per function; everything a pass doesn’t name stays byte-identical:

# flash.toml
[[pass]]
name = "scalar-to-uniform"
sm   = ["sm_90*"]
fn   = ["*flash_fwd_kernel*"]
mode = "apply"          # default is "report"

[[pass]]
name = "hmma-to-hgmma"
sm   = ["sm_90a"]
fn   = ["*flash_fwd_kernel*"]
mode = "report"         # shows where it would fire, mutates nothing

• mode = "report" (the default when mode is omitted) — dry-run: print every site the pass would rewrite, change no bytes.
• mode = "apply" — actually transform. Turning a pass on is an explicit, per-pass decision, because once bytes change, round-trip identity can no longer be the safety net.

Flags: --passes <dir> (plugin directory), --config <toml> (pass selection/targeting), --report-only (force every pass to report mode, ignoring the config), --keep-mlir <dir> (write the lifted and transformed modules for inspection).

Live processes: capture and rewrite in flight

intercept runs any command under sass2mlir’s capture shim — no manual LD_PRELOAD, no environment variables. Cubins are intercepted in flight, gated on the CUDA ELF header, as the target process loads them:

sass2mlir intercept --out ./captures -- python train.py

ls captures/
# flash_fwd_kernel.cubin  flash_fwd_kernel.code  flash_fwd_kernel.mlir  run.csv

Per kernel: the raw cubin, the .code stream, the lifted .mlir, and a run.csv row per launch with grid/block/shared-memory configuration — the same artifacts harvest produces offline, for kernels that never sit still on disk.

• --out <dir> — capture directory (default ./sass2mlir-out).
• --filter <glob> — only capture kernels matching the glob (e.g. --filter '*flash*'). Default: all.
• --no-lift — capture bytes only, skip the .mlir.

With --apply, intercept stops being observational and becomes the online engine: each kernel is lifted, transformed, re-encoded, and patched at load time — the same pass pipeline as the offline apply, applied to a process’s kernels as they arrive instead of a file on disk:

sass2mlir intercept --apply \
  --passes build/lib/sass2mlir/passes \
  --config jetson.toml \
  -- python infer.py

This is how the Jetson on-device path works, and how a patched kernel reaches a process you can’t or don’t want to modify — nothing of the workflow sits on the hot path after load. The full walkthroughs: Capturing & lifting kernels and On-device: Jetson Orin.

Constant banks and kernel overlay

# Shift constant-bank offsets after an ABI change:
sass2mlir remap kernel.mlir --const-bank-shift 0x210:0x40 -o remapped.mlir

# Overlay a patched kernel into a different cubin's slot:
sass2mlir remap patched.mlir --overlay target.cubin --slot flash_fwd_kernel -o target.patched.cubin

Two jobs: rewrite constant-bank references (when kernel parameters move, every c[0x0][...] offset referencing them has to move too), and splice a module into an existing cubin’s kernel slot, fixing up .nv.info and relocations. apply calls the same machinery for its final patch step.

Scripting against sass2mlir

Everything prints machine-greppable lines on stdout and human detail on stderr, so sass2mlir roundtrip app.cubin -q && echo OK does what you expect in CI.