`nx_mma_asm.nx` source

1// nx_mma_asm.nx -- sovereign PER-VENDOR tensor-core MMA ASSEMBLER (RE-derived encoders). ONE lib; each vendor
2// is a set of encoder funcs sharing the same disassemble->decode->encode method. Pure funcs, no main.
3//
4// NVIDIA Blackwell sm_120 DONE (RE-verified bit-exact vs nvcc -arch=sm_120 -cubin -> cuobjdump -sass,
5// bench/mma_probe.cu + mma_probe2.cu, 10 samples): HMMA bf16 + IMMA s8 (integer = deterministic = the exceed).
6//   low64  = opcode | Rd<<16 | Ra<<24 | Rb<<32
7//   high64 = Rc | typebits<<8 | control<<40   (Rc=RZ=0xff when no accumulator addend)
8// AMD MFMA/WMMA + Intel DPAS = future SIBLING encoders added here (same method, swap the instruction).
9// This is the per-device R4 (tensor-core RE) backend lib for the "S-class exceed for ANY DEVICE" ladder.
10// license_tier: ORIGINAL
11import "nx_syscalls.nx"
12
13// ---- NVIDIA Blackwell sm_120 (cracked) ----
14const NV_HMMA_16816: i64 = 0x723c   // HMMA.16816.F32.BF16  (opcode bits[0:15])
15const NV_IMMA_16816: i64 = 0x7237   // IMMA.16816.S8.S8     (integer, deterministic)
16const NV_QMMA_16832: i64 = 0x727a   // QMMA.16832.F32.E4M3  (FP8 -- the highest-throughput Blackwell dtype / float ceiling)
17const NV_TB_HMMA:    i64 = 0x0418   // typebits (high[8:39]) for HMMA.BF16
18const NV_TB_HMMA_F16: i64 = 0x0018  // HMMA.F16 -- SAME opcode 0x723c, dtype in typebits (BF16 sets bit10=0x0400)
19const NV_TB_IMMA:    i64 = 0x4054   // typebits for IMMA.16816 (k16, .ROW/.COL + type/layout)
20const NV_TB_IMMA_K32: i64 = 0x405c  // IMMA.16832 (m16n8k32, wider-K int8 -- the deterministic exceed op) vs k16 0x4054
21const NV_TB_QMMA_E4M3: i64 = 0x002c // QMMA.16832 E4M3.E4M3 (FP8) typebits
22const NV_RZ:         i64 = 0xff     // RZ accumulator register (Rc when no addend)
23
24func nv_mma_low(opcode: i64, rd: i64, ra: i64, rb: i64) -> i64 { return opcode | (rd<<16) | (ra<<24) | (rb<<32) }
25func nv_mma_high(rc: i64, typebits: i64, control: i64) -> i64 { return rc | (typebits<<8) | (control<<40) }
26
27// SASS control/scheduler word (the high64[40:63] field, = control<<40 in nv_mma_high). RE-DECODED by differential
28// analysis + VERIFIED bit-exact across the cuobjdump samples (mma_probe2.cu). Layout (within the 24-bit word):
29//   [0:3] stall · [4] yield · [5:7] write-barrier(7=none) · [8:10] read-barrier(7=none) · [11] fixed=1 ·
30//   [13:18] wait-barrier mask (6 bits, one per scoreboard barrier) · [19:22] reuse flags.
31// This completes the C0 frontier: the encoder can now emit COMPLETE, correctly-SCHEDULED instructions from
32// semantic fields (stall/yield/barriers/reuse), not just operand-correct ones.
33func nv_control(stall: i64, yield: i64, wbar: i64, rbar: i64, wait_mask: i64, reuse: i64) -> i64 {
34    return (stall & 0xf) | ((yield & 1) << 4) | ((wbar & 7) << 5) | ((rbar & 7) << 8) | (1 << 11) | ((wait_mask & 0x3f) << 13) | ((reuse & 0xf) << 19)
35}
nx_mma_asm.nx source

`nx_mma_asm.nx` source
