Module codegen_sass 

codegen_sass — the foundation of our OWN sm_86 SASS assembler (the MaxAs / turingas lever).

ptxas is the only PTX→SASS path and it is closed-source: it fixes our register allocation and, crucially, the per-instruction CONTROL CODES (stall counts, yields, read/write barriers, the operand reuse cache) that decide how well a hot loop hides latency. That last 5-20% cuBLAS keeps lives there. To take it we must read AND write SASS ourselves. This module is that capability: locate the kernel’s machine code in the cubin ELF, decode every 128-bit instruction’s control field (verified bit-exact against nvdisasm on sm_86), re-encode it, and patch the bytes back. A scheduler/optimizer on top is the follow-on; this is the assembler core.

Ampere (sm_80/86) 128-bit instruction = (lo:u64, hi:u64). The control field sits in the high word:

stall   bits [41:44]  yield bit [45]  wr_barrier [46:48]  rd_barrier [49:51]  wait_mask [52:57]  reuse [58:61]

(barrier index 7 = “none”.) Verified: S2R/S2UR/memory ops set a write barrier; pure ALU ops do not.

Structs

Control

The per-instruction control codes — where the GPU’s static scheduling lives.

DepGraph

A register-dependency DAG over a straight-line SASS region: deps[i] lists earlier instructions that i must stay after (RAW/WAR/WAW hazards + barrier ordering). Any reordering that respects these edges preserves semantics — this is what the scheduler is free to permute to hide latency.

SassInsn

One decoded SASS instruction: its byte offset within the .text section and the raw 128-bit encoding.

SassReport

What the COMPILER learns about the machine code ptxas produced for one of its kernels — the SASS-awareness the backend needs to reason about the ptxas ceiling (register pressure/occupancy, tensor-core utilisation, reuse-cache usage). Every field is read from the real cubin, not guessed.

Enums

SassClass

Execution-unit class of a SASS opcode — what a scheduler must know to place it and set control codes.

Functions

analyze_source

The compiler emits cu_src, compiles it with nvcc to a cubin, and returns the SASS report — the backend “knowing” the machine code its codegen ultimately becomes. Requires nvcc + nvdisasm; returns None otherwise.

assign_minimal_stalls

CONTROL-CODE ASSIGNER (perf): compute the MINIMAL stall per instruction from the fixed-latency dependency timing, keeping ptxas’s memory barriers. We walk the instruction order tracking the cycle each register’s fixed-latency result becomes ready; each instruction issues at the max of (prev issue + 1) and its producers’ ready cycles; the stall between two instructions is exactly that gap (clamped to [1,15]). This is the tightest schedule that still satisfies every fixed-latency hazard — if ptxas over-stalled anywhere, ours is faster; if it was already minimal, ours matches. Correctness is guarded on-device (a too-tight stall that violates a hazard we don’t model shows up as a wrong result). Returns the patched cubin bytes.

assign_minimal_stalls_opt

As above, but floor_ptxas selects whether to floor our computed stall at ptxas’s (sound) or trust our CALIBRATED model outright (false — the true test of whether the measured latencies are complete).

assign_stalls_by_latency

CONTROL-CODE ASSIGNER (correctness-first, stalls). Recompute every instruction’s STALL count from our own latency model — stall = fixed_latency(class) — which is provably >= the cycles any fixed-latency consumer needs (the immediate next instruction already issues stall cycles later, so later consumers have even more). We KEEP ptxas’s write/read barriers + wait masks (the variable-latency memory / cp.async tracking, the fragile part) untouched. If the kernel stays numerically correct with OUR stalls, the stall model is validated on-device — the first half of a working control-code assigner. Returns the patched cubin bytes.

build_dependency_graph

Build the dependency DAG from per-instruction (reads, writes, is_barrier) triples. Barriers (BAR/BSYNC/ control) force a total order across themselves: nothing moves past a barrier in either direction.

classify_opcode

Classify an sm_86 SASS base mnemonic (the token before the first .). Covers the full 81-opcode inventory extracted from real machine code (our GEMM + cuBLAS) — this is the ISA a SASS scheduler operates over.

dependency_graph_from_cubin

Build the dependency DAG directly from a cubin, using nvdisasm to name each opcode (so classification — hence the read/write sets — is precise). Returns the graph paired with the parsed instructions. This is the mnemonic-aware, SOUND path the scheduler uses.

find_text_section

Locate the .text.<kernel> section in a cubin ELF64: returns (file_offset, size_in_bytes).

instructions

Parse every 128-bit instruction out of the kernel’s .text section of a cubin.

list_schedule

LIST SCHEDULER — the scheduler engine. Given the sound dependency DAG and per-instruction latencies, produce a new instruction order that respects every hazard while placing latency-critical work early so its result is ready by the time a consumer needs it (classic critical-path list scheduling). Returns the new order as original-instruction indices. This is what a re-emitter writes back to close the gap ptxas’s schedule leaves.

nvdisasm

Disassemble a cubin file to SASS text via nvdisasm (for inspection / labelling).

optimize_reuse

The reuse-cache optimizer pass — our first SASS optimization ON TOP of ptxas. The operand collector can REUSE a source register read by the previous instruction instead of re-reading the register file, saving RF bandwidth (the bottleneck under our 222-register, 16.7%-occupancy kernel). ptxas already sets many reuse flags; this finds the ones it MISSED: consecutive instructions that read the same register into the same operand slot (Ra→reuse bit 0 = hi[58], Rb→reuse bit 1 = hi[59]) but where the flag is unset. RZ is skipped (reusing the zero register is pointless). Returns the patched cubin bytes and the number of flags added. The result MUST be validated on-device (varied-fill oracle) — the operand→slot mapping is verified there, never trusted blind.

patch_hi

Patch instruction idx’s high word (control codes live here) in-place in the cubin bytes.

precise_dependency_graph

Build a PRECISE dependency DAG from a cubin by parsing nvdisasm’s exact operands per instruction (not the conservative range approximation). This gives the scheduler maximal legal freedom — the graph has only the TRUE hazards, so a reschedule can hide latency ptxas left exposed. Returns (insns, classes, mnemonics, DAG).

precise_operands

PRECISE operand decode from the nvdisasm TEXT of one instruction (everything after the mnemonic, before the ;). nvdisasm already prints the exact registers, so we parse them rather than reverse-engineer the binary operand fields. Returns (writes, reads) as register numbers (RZ/255 dropped). Rules, verified against real SASS: operand 0 is the WRITE destination UNLESS it is a memory ref [...] (a store — then it’s a read address and nothing is written); every other operand is a READ; registers inside [...] are address reads; a .64/.128 width on an address adds the consecutive address registers; accumulate ops (IMMA/FFMA) naturally read their destination because it reappears as a later operand. Groups {R8,R9,...} expand.