Why on-device

Everything on this page happens on the GPU’s own host. The target is an 8 GB Jetson Orin Nano — sm_87, aarch64, one shared pool of unified memory — and the point is that the whole lift → analyze → rewrite → patch loop runs on the edge box itself, against the kernels its own processes are loading. No cubin gets shipped to a workstation to come back patched a day later; the box that runs the kernel is the box that rewrites it.

That matters for the workloads the Orin actually runs: vision pipelines and quantized-LLM inference whose kernels arrive inside prebuilt binaries with no source you control and no ptxas invocation you can influence. sass2mlir works below all of that — on the SASS in the cubin, at the moment it loads.

The toolchain is the same one described in the overview; nothing here is an Orin-specific fork. The only differences are a build flag and a config file.

A build that fits in 8 GB

The full multi-architecture build compiles a decode table for all 22 supported SM architectures — the generated multi-SM table runs to roughly 270k lines, and table generation peaks above what 8 GB of unified memory can hold while the GPU driver and a desktop are also resident. The single-architecture build compiles only the sm_87 probe tables (opcodes/SM87/*.tsv) through union-gen — and -DSASS_MLIR=OFF drops the LLVM dependency entirely, building against the internal simplified-MLIR IR instead (see build options):

free -h
#                total        used        free
# Mem:           7.4Gi       1.1Gi       5.6Gi

cmake -B build -G Ninja \
  -DCMAKE_BUILD_TYPE=Release \
  -DSASS_ACTIVE_SMS=sm_87 \
  -DSASS_MLIR=OFF
# -- sass2mlir — active architectures: sm_87
# -- IR backend: internal (SASS_MLIR=OFF), LLVM not required
# -- union-gen: opcodes/SM87/*.tsv → sm_87 decode table (11,842 lines generated)
# -- skipping opcodes/SM70 ... opcodes/SM121a (not in SASS_ACTIVE_SMS)
# -- Configuring done
# -- Generating done

ninja -C build
# [404/417] Building CXX object lib/Passes/CMakeFiles/... 
# [412/417] Linking CXX executable bin/sass2mlir
# [417/417] Linking CXX shared module lib/sass2mlir/intercept.so

Cross-compiling from an x86 workstation is supported via the aarch64 toolchain file, but on a single Orin it buys little: the host-side TableGen and LLVM tools have to be built for the cross host first, and the single-arch table is what made the native build fit at all. Native is the simpler path here — one overnight build, then only incremental rebuilds of whatever pass you are working on.

Capturing kernels on the Orin

The workload: a llama.cpp-style quantized inference binary running entirely on the GPU. Its kernels live inside the prebuilt executable; sass2mlir intercept captures the cubins as the CUDA driver loads them, gated on the CUDA ELF header:

sass2mlir intercept \
  --out /data/captures/llama \
  --filter '*mul_mat*' \
  -- ./llama-cli -m q4_K_M.gguf -ngl 99 -p "The SASS instruction set"
# [intercept] sm_87: intercepted module (3 kernels matching '*mul_mat*')
# [intercept] wrote /data/captures/llama/mul_mat_q4_K.{cubin,code,mlir}
# [intercept] wrote /data/captures/llama/mul_mat_vec_q4_K.{cubin,code,mlir}

--filter matters more here than on a datacenter card: storage is an SD card, and an unfiltered capture of a full inference binary is hundreds of kernels. Each capture is the raw cubin, the .code stream, the lifted .mlir (dropped with --no-lift), and a run.csv row per launch:

kernel,grid,block,shared_mem,launches
mul_mat_q4_K,"(4096,1,1)","(256,1,1)",0,1821
mul_mat_vec_q4_K,"(2048,1,1)","(128,1,1)",0,336

Round-trip the capture on the device itself — the single-arch binary still carries the full verifier for the architecture it was built for:

sass2mlir roundtrip /data/captures/llama/mul_mat_q4_K.cubin --arch sm_87
# roundtrip: sm_87 mul_mat_q4_K: 1204/1204 instructions byte-identical (0 mismatches)

module {
  sass.func "mul_mat_q4_K" {
    %0 = builtin.unrealized_conversion_cast to !sass.reg     // live-in: kernel params
    ...
    // dequant inner loop: load the quantized block, widen, accumulate
    %52 = sass.LDG_0x385.E.128 guard(%PT) [%44], %0
            : (!sass.reg, !sass.reg) -> (!sass.reg)
    %57 = sass.LOP3_0x1ff.LUT guard(%PT) %52, 0x0f0f0f0f, %53
            : (!sass.reg, !sass.imm, !sass.reg) -> (!sass.reg)
    %63 = sass.IMAD_0x24e.WIDE guard(%PT) %57, %61, %62
            : (!sass.reg, !sass.reg, !sass.reg) -> (!sass.reg, !sass.reg)
    ...
    sass.BRA_0x942 guard(%38) ^bb2 : () -> ()                // row loop back-edge
    sass.exit
  }
}

Rewriting at load time

The online path is the offline one moved earlier in the pipeline. Offline, sass2mlir apply reads a cubin from disk, runs the plugin passes, and patches the file. Online, sass2mlir intercept --apply runs the same engine at the interception point: a kernel arrives as a cubin in flight, is lifted, transformed per the config, re-encoded, and patched — in the original kernel slot, ABI untouched — before the process ever launches it. The loop is identical: lift → passes → re-encode → patch. Only the source of bytes changed.

The configuration is a pass.toml targeted at sm_87:

# pass.toml — sm_87 on-device configuration
[[pass]]
name = "scalar-to-uniform"
sm   = ["sm_87"]
fn   = ["*mul_mat*"]
mode = "apply"           # applies real byte changes today

[[pass]]
name = "ffma-deftz"
sm   = ["sm_87"]
fn   = ["*mul_mat*"]
mode = "report"          # where would it fire; changes nothing

Exercised offline against the capture first — the discipline is the same as on the H100 path: --keep-mlir for inspection, and structural passes stay in report mode until they are driven to full application:

sass2mlir apply /data/captures/llama/mul_mat_q4_K.cubin \
  --passes build/lib/sass2mlir/passes \
  --config pass.toml \
  --keep-mlir /data/work/mul_mat \
  -o mul_mat_q4_K.patched.cubin
# lift: sm_87 mul_mat_q4_K — 1204 instructions
# scalar-to-uniform [apply]: 2 warp-uniform webs migrated (14 ops rewritten)
# ffma-deftz [report]: would fire at 61 sites — no bytes changed
# patch: .text rewritten in place, .nv.info updated — ABI unchanged

What scalar-to-uniform actually did to the inner loop — index math that is constant across the warp moves off the per-thread register file onto the uniform datapath (URn), freeing per-thread registers:

- IMAD.MOV.U32 R9, RZ, RZ, c[0x0][0x210]   // per-thread copy of a warp-uniform param
+ UMOV UR9, c[0x0][0x210]                    // one uniform-register move, warp-wide

Every kernel the config does not name stays byte-identical — that is the safety property that makes running this at load time sane. And because ffma-deftz is in report mode, the patched cubin differs from the capture in exactly the 14 ops scalar-to-uniform rewrote and nothing else.

Constraints of the edge

The Orin imposes three limits the datacenter path never sees.

• RAM. The 8 GB is unified memory shared with the GPU and the desktop stack. Table generation for the multi-SM build does not fit — that is the whole reason SASS_ACTIVE_SMS exists. At run time the footprint is small (one decode table, one kernel’s IR in memory at a time), but a resident model plus a desktop leaves little headroom for anything else large.
• Storage. Captures are per-kernel and verbose — the lifted .mlir is several times the size of the .code stream. On SD-card storage, pass --filter for the kernels you actually intend to work on, or --no-lift to capture bytes only and lift later.
• Iteration latency. Native rebuilds on six Cortex-A78AE cores are fine for the toolchain once, but not as an inner loop for pass development.

The practical loop, then:

1. Capture on the device — sass2mlir intercept against the real workload, filtered to the target kernels.
2. Develop on a workstation — copy the captures off, and iterate on the pass against the captured .mlir with apply --report-only, roundtrip, and regalloc --pinned as the oracles. This is where the fast compile-test cycle lives.
3. Ship back to the device — the pass .so (aarch64 — build the plugin natively on the Orin, or cross-compile just the plugin, not the whole toolchain) plus the pass.toml. The device runs the same engine that validated the pass on the workstation.