sass2mlir: a byte-exact compiler for NVIDIA machine code

Machine code you can read but never edit

SASS — NVIDIA’s native GPU instruction set — is undocumented, encoding-dense, and changes with every architecture. The vendor tooling stops at a read-only disassembler: you can look at what ptxas emitted, but there is no assembler to put modified code back. Once a kernel ships inside a cubin, its machine code is effectively frozen.

sass2mlir closes the loop. It lifts SASS into a first-class MLIR dialect, analyzes and rewrites it with real compiler infrastructure — CFG recovery, SSA, liveness, register allocation, conversion passes — and lowers it back to bytes that are identical to the input: lower(lift(bytes)) == bytes, per instruction, gated against the full corpus.

The mechanism is a two-path contract. Every lifted instruction carries its original 128-bit encoding as an authoritative $raw attribute; unmodified instructions lower verbatim from it. The moment a pass edits an instruction, its $raw is invalidated and the structural assembler (nvisa_encode) must reproduce the exact bytes from the semantic fields alone — so an edit can never silently corrupt its neighbors, and any byte that changes in the output is exactly the transform you applied.

Probe-derived, not hand-written

Nobody hand-wrote the instruction tables. Opcode field layouts are discovered by a systematic bit-toggle probing methodology, reduced to per-architecture tables (opcodes/SM*/*.tsv), and compiled by union-gen into both the disassembler dispatch and the MLIR dialect definition. Adding an architecture is a data problem, not a rewrite.

22 SM architectures are supported today — Volta (sm_70) through Blackwell (sm_121a).

An MLIR dialect for machine code

Each SASS instruction becomes one op named for its mnemonic and opcode (sass.<MNEMONIC>_0x<opcode>) — 770 per-mnemonic ops generated straight from the probe tables — carrying its SSA operands and results, an inline guard(%p) predicate, modifiers glued to the name, and the authoritative 128-bit raw encoding that makes lowering exact. Operands are structured and lossless: registers, immediates, memory descriptors, TMA/tensor descriptors, constant-bank references, special registers — 38 decoder operand variants modeled without loss, over register types with hardware semantics (!sass.reg, !sass.pred, and the uniform-datapath !sass.ureg, !sass.upred).

module {
  sass.func "lifted" {
    %0 = builtin.unrealized_conversion_cast to !sass.reg     // kernel live-in values
    ...
    // fused multiply-add: predicated, negated first operand, one result
    %4706 = sass.FFMA_0x223.FTZ guard(%4380) -%4550, %4464, %4551
              : (!sass.reg, !sass.reg, !sass.reg) -> (!sass.reg)

    // warpgroup matrix-multiply: two results, four operands
    %5717:2 = sass.HGMMA_0xdf0.64x8x16.F32 guard(%425) %5716, %5590, %5716, %5716
              : (!sass.reg, !sass.reg, !sass.reg, !sass.reg) -> (!sass.reg, !sass.reg)

    sass.UTMALDG_0x5b4.3D guard(%266) %31250, %31086 : (!sass.reg, !sass.reg) -> ()  // TMA load
    sass.STSM_0x844.16.M88.4 %33430, %33553 : (!sass.reg, !sass.reg) -> ()           // smem store

    cf.br ^bb1(%4783, %4751, ... : !sass.reg, ...)   // block args are SSA phis
    sass.exit
  }
}

Undocumented instructions pass through untouched

sass2mlir models semantics only for instructions it has positively identified. Anything it does not recognize — hidden, undocumented, or an arch-specific opcode with no registered form — is carried by the generic sass.inst op: the original 128-bit encoding in $raw, and nothing more. Such instructions lower in place, byte-for-byte — preserved exactly, never reverse-engineered.

// A recognized instruction — a structured, per-mnemonic op:
%r = sass.LDC_0xb82 : () -> (!sass.reg) {
       operands = [#sass.reg_ref<R, dst, ...>, #sass.constmem<...>],
       raw = #sass.raw<0xdf00ff017b82, 0xfe20000000800>, ... }

// An unrecognized (hidden / undocumented) instruction — generic carrier, raw only:
%r = sass.inst : () -> (...) {
       mnemonic = "…", opcode_val = 0x… : ui16, operands = [],
       raw = #sass.raw<0x…, 0x…>, ... }

A real compiler over recovered SSA

Lifting is only useful if you can reason about what you lifted. On top of the dialect sits standard — and some not-so-standard — compiler infrastructure:

• CFG recovery from branch targets, into native MLIR blocks.
• SSA construction (Braun-style), with φ-nodes as block arguments.
• Def-use with reaching definitions and per-block liveness.
• Uniformity analysis — which per-thread values are actually scalar across a warp. This is the analysis that unlocks the uniform-datapath rewrites below, and it is conservative by construction: any uncertainty marks a value non-uniform, because a spurious “uniform” would be unsound.

Register allocation is Chaitin-Briggs graph coloring over the recovered SSA, extended to the uniform register and predicate classes (UR/UP). It runs in two modes: pinned — must reproduce byte-identical SASS, the hard corpus-wide oracle that proves the liveness and interference model right — and free — real recoloring, validated by SSA contributor-set equivalence.

From a flat instruction stream to blocks

CFG recovery is the first analysis, and everything above it — SSA, liveness, register allocation — depends on getting it right.

Every terminator falls into one of five kinds: BRA/BRX/JMP are unconditional with no guard (or an always-true PT guard), or conditional under any other guard; RET returns; EXIT ends the thread; and BSSY/BSYNC form a convergence pair — BSSY pushes a convergence point that the matching BSYNC branches to when the warp reconverges. Branch targets are absolute, resolved per instruction from target = addr + 16 + (rel << scale), where addr is the branching instruction’s own offset and rel/scale come out of its decoded operand.

Leaders fall out of one rule, applied once over the whole flat stream: the entry instruction, every branch target, and the instruction immediately after every terminator. Leaders partition the stream into blocks; each block’s successor edges are read straight off its terminator kind — a conditional branch gets a taken edge and a fall-through edge, an unconditional branch or a BSYNC reconvergence gets one, RET/EXIT get none.

Recovery only reads the lifted module — it builds the block structure without touching any op’s $raw — so the byte-exact round-trip stays intact by construction, not by a rule the CFG pass has to remember to honor. It’s also cross-checked against nvdisasm -bbcfg as an independent oracle: the recovered block boundaries have to agree with NVIDIA’s own basic-block disassembly on the same cubin.

Converting the graph, not just editing inside it

Most of the conversion passes described above are peephole rewrites: they match an instruction or a small window inside a block and replace it, leaving the CFG’s shape untouched. warp-specialization is different — it converts the control-flow graph completely, rather than editing inside it.

Starting from the single recovered CFG, it splits the function body into a producer subgraph — the loads, MMA, and index math — and a consumer subgraph — control, reconvergence, and epilogue — divides the register budget between the two halves, and combines them back into one warp-specialized graph stitched with mbarrier handoff edges where the producer signals data-ready and the consumer waits on it. The demo above walks that flow end to end: stream → single CFG → producer/consumer split → recombined graph.

Apply: passes as plugins

Every SASS→SASS transform is a .so plugin with a small stable ABI, loaded with dlopen. sass2mlir apply lifts a cubin, runs a folder of passes, re-encodes, and patches the cubin in place — .text plus the updated .nv.info/relocations — so the output is a single self-contained binary in the original kernel’s slot.

Two safety properties are load-bearing. Passes are targeted per-architecture and per-function by config (sm sm_90* / fn *flash_fwd_kernel* globs), so every kernel a pass doesn’t name stays byte-identical. And the default manifest is report-only: structural passes show where they would fire without mutating, because once bytes change, round-trip identity can no longer be the safety net — turning a pass on is an explicit, per-pass decision.

Kernels don’t have to come from disk. sass2mlir intercept runs any process under a capture shim and intercepts cubins in flight — gated on the CUDA ELF header — harvesting every launched kernel: raw bytes, lifted MLIR, and a per-launch run.csv with grid/block/shared-memory config. With --apply it rewrites the kernels at load time instead of just capturing them. sass2mlir harvest <cubin> is the offline equivalent.

Regression-gated at corpus scale

Every headline number is held by a coverage tool in continuous regression, not asserted once.

Decode is validated across the 13-architecture local corpus; the decoder itself supports all 22 listed architectures. The assembler didn’t start at byte-identity — a coverage tool measured encode(decode(x)) == x at 86.4% of the corpus, and successive fixes drove it to 91.2% and then to full byte-identity, where continuous corpus regression now holds it.

Recompiling FlashAttention for Hopper

FlashAttention’s official forward kernel is the flagship target. As shipped, it asks for 255 registers per thread — the hardware maximum — which caps occupancy at 12.5% and forces 85 register spills to local memory. It moves data with cp.async (48 LDGSTS) + 144 LDSM and multiplies with 512 per-warp HMMA tensor-core ops. All correct — and all leaving Hopper’s newest hardware idle.

The root cause is an ISA gate, not the source. The kernel was built for sm_90, but Hopper’s headline instructions — warpgroup HGMMA (wgmma) and TMA — are only emittable for the architecture-specific sm_90a target. Built for plain sm_90, ptxas cannot emit them and falls back to Ampere-class code on an H100 that supports better. That is precisely the gap sass2mlir works in: below the source compiler, on the SASS it actually emitted — re-optimize for the exact GPU you have, without waiting on upstream.

- HMMA.16816.F32.F16  R8, R24, R40, R8   // per-warp MMA, operands in regs
+ HGMMA.64x64x16.F32.F16   // warpgroup async MMA, smem descriptors

- LDGSTS.E.128  [R12], [R4]   // cp.async global→shared
+ UTMALDG.2D  [desc], [R2]    // TMA bulk load + mbarrier

The rewrite is an ordered pipeline of six conversion passes, and the order is load-bearing — free the registers first, lift address and index math onto the uniform datapath, restructure into warpgroups, then convert the tensor-core and load paths that depend on that structure:

1. spill-removal — eliminate the STL/LDL spill traffic.
2. scalar-to-uniform — migrate closed warp-uniform integer webs to the uniform datapath (URn, UPn).
3. warp-specialization — restructure the monolithic warp into producer/consumer warpgroups.
4. hmma-to-hgmma — per-warp HMMA → warpgroup async HGMMA.
5. ldsm-to-tma — cp.async + ldsm → TMA (UTMALDG) + mbarrier.
6. ffma-deftz (opt-in) — drop the FTZ default where it isn’t needed.

Why HGMMA needs the full round trip

HMMA and HGMMA aren’t two spellings of one instruction — they’re different execution models, so “use HGMMA” can’t be a textual swap. Everything feeding the multiply changes with it:

At the SASS level the matmul is fused into registers, schedule, and sync — it can’t be safely edited in place. Lifted to MLIR it becomes a typed op that a conversion pass can legally re-express in the warpgroup paradigm, then lower back to sm_90a SASS.

The output is reinjected into the original cubin’s kernel slot — rewritten .text plus updated .nv.info — with the kernel ABI untouched: the patched kernel reads its parameters from the exact same const-bank offsets the official one did. Same launch site, same host code, one self-contained binary. Validated end-to-end in an unmodified flash_attn process via sass2mlir intercept.

On-device rewriting: Orin (sm_87)

The same framework runs on the GPU’s host — including an 8 GB Jetson Orin Nano. A build scoped with -DSASS_ACTIVE_SMS=sm_87 compiles only that architecture’s decode table instead of the full ~270k-line multi-SM set — and with -DSASS_MLIR=OFF the LLVM dependency drops out too, replaced by the internal simplified-MLIR IR. Small enough for low-RAM edge targets.

That enables the online path: sass2mlir intercept --apply captures kernels as they load and rewrites them on the fly — the same lift → passes → re-encode → patch loop, applied to a process’s kernels at load time instead of a file on disk. An Orin-specific kernel-optimization case study (the sm_87 counterpart of the H100 work above) is in progress; its numbers land here as the sweeps complete.

Measured speedups

Each device runs the same test: the official FlashAttention forward kernel as shipped (baseline) against the sass2mlir-patched cubin dropped into the exact same kernel slot — same ABI, same launch, only the machine code changed. Times are GPU kernel milliseconds per iteration; speedup is baseline ÷ sass2mlir.

NVIDIA H100 NVL

The win grows with context length. At short sequences (≤2k) the patched kernel runs a few percent behind (0.92–0.98×), then it pulls ahead to 1.23–1.24× from 8k tokens up — the long-context regime that matters.

NVIDIA GeForce RTX 5070 Laptop GPU

On this Blackwell laptop part the trade-off inverts relative to the H100: the patched kernel is fastest at short context (1.12–1.15× up to 512 tokens) and runs a few percent behind baseline once sequences grow past 1k (~0.96–0.99×) — a different occupancy and memory regime, and a candidate for the scheduling passes on the roadmap.

Jetson Orin Nano (sm_87)

Numbers landing — the sm_87 sweep is being added here.

Passes the framework is built for

Through its v1.x series, the FP8 GEMM library DeepGEMM piped every JIT-compiled cubin through a Python post-pass: cuobjdump-disassemble it, regex-match the FFMA hex in the FP8 dequantization epilogue, then mmap the .so to clear the operand-cache reuse bit and set the warp-scheduler yield bit on alternating instructions — a raw two-field bit-flip in the sm_90 control word, because ptxas left long FFMA chains with reuse set and yield clear, stalling the wgmma pipeline. (nvcc 12.9 has since internalized the pass, but the point stands.) Scheduling-aware SASS transforms like this one either live inside closed ptxas, or ship as byte-level patches with no instruction semantics — they can’t compose, and they break on the next toolkit revision. In sass2mlir, reuse, yield, stall count, and barriers are first-class control-code attributes, so the same transform is an attribute edit on the IR, not a hex patch on the binary.

The transform pattern — model the semantics, ship it as an apply plugin, and validate at corpus scale — generalizes past the six FlashAttention passes above. The following are capabilities of this AGPL-3.0 framework as the passes land, not shipped results yet:

• Security GPU antivirus
  GPU function blacklisting
  PROBLEM Malicious kernels (e.g. the Membar+Load cross-MIG L2 side channel) ship as machine code with no source to inspect.
  LEVER The lifter reads the post-ptxas cubin, so a pass signature-matches the attack primitive and refuses to launch (or neutralizes) the function.
• Scheduling ~10% throughput
  Natural Yield Strategy
  PROBLEM Compilers scatter the warp-scheduler yield flag on a blind every-7–8-instruction heuristic, stalling warps and killing the reuse cache.
  LEVER The yield bit is a per-instruction control-code field the dialect models, so a one-field edit lets warps yield only when they actually stall.
• Scheduling up to 26%
  Instruction scheduling
  PROBLEM Overlapping memory with compute needs reorderings whose real dependencies hide in control-code barriers/stall-counts and undocumented hardware rules.
  LEVER SSA def-use, liveness, and the control-code fields are all first-class IR, so the pass computes which reorderings are legal and throughput-positive.

1. GPU function blacklisting → a GPU antivirus. sass2mlir lifts a cubin to a semantic IR after ptxas — the code that actually executes, not source an attacker won’t provide — so a pass can pattern-match the instruction-level signature of a known attack primitive and blacklist the offending function: refuse to launch it, or neutralize it in place. The reference attack class is the Membar+Load side channel described in “Behind Bars: A Side-Channel Attack on NVIDIA MIG Cache Partitioning Using Memory Barriers” (USENIX Security 2026) — MEMBAR requests issued in one MIG instance measurably slow load instructions in another, despite MIG’s physical L2 partitioning, letting a sender modulate MEMBAR traffic and a receiver time LD.STRONG.GPU to recover a bit. The signature is exactly at the instruction level (MEMBAR modulation plus LD.STRONG.GPU timing loops) in the post-ptxas cubin — the layer the lifter already reads.
2. Natural Yield Strategy. Volta/Turing-class SASS encodes each instruction with a control code that includes a 1-bit warp-scheduler yield flag: yield=1 keeps issuing from the current warp, yield=0 prefers switching warps at the cost of a cycle and the register reuse cache. Compilers scatter yield=0 roughly every 7-8 instructions as a blind periodic heuristic. The Natural strategy from Yan, Wang, and Chu, “Optimizing Batched Winograd Convolution on GPUs” (PPoPP ‘20) instead lets a warp yield only when it actually stalls — the paper reports about 10% higher throughput from tuning the yield flag alone, 1.09x over NVCC and 1.11x over cuDNN on the main loop. It’s exactly a per-instruction control-code field the dialect already models, so the rewrite is a single-field edit per instruction on the byte-exact re-encode path — and it’s literally what DeepGEMM’s hand patch above was doing, made composable.
3. Instruction scheduling: legal reordering + pipeline saturation. Reordering SASS is normally treated as unsafe because the real dependencies aren’t in the instructions — they’re spread across the control code (wait-barrier mask, read/write barriers, stall count) plus undocumented hardware rules, and the goal is to overlap memory (LDG/LDGSTS/STG) with compute (FFMA/IMAD) so the pipeline stays saturated when there aren’t enough eligible warps to hide latency by switching. CuAsmRL (He and Yoneki, CGO ‘25) trains an RL agent to mutate a -O3 SASS schedule against measured throughput, reporting up to 26% and 9% average improvement over -O3. sass2mlir already surfaces the full dependency set this needs as first-class IR — SSA def-use, per-block liveness, basic blocks, and the control-code fields themselves — so a scheduling pass can compute which reorderings are legal and which are throughput-positive directly from the model, instead of assembling and running each candidate to find out.

Each of these gets the same treatment as scalar-to-uniform: model the transform, ship it as an apply plugin, and validate it at corpus scale before it’s called done.

Where this stands

Working today: byte-exact lift/lower across 22 architectures, the analyses and register allocator, the apply engine with six FlashAttention passes, on-device capture and rewriting, and the probe-derived opcode tables — all regression-gated at corpus scale. In progress: driving the structural passes to full application, the H100 and Orin benchmark sweeps, and the roadmap passes above. The method is the subject of a paper in preparation.