nx_f32_matmul_block_gate.nx source

1// nx_f32_matmul_block_gate.nx -- referee for the 4x4 register-blocked matmul_t (nx_f32_matmul_t.mmt_block, 2026-09-15). 2// The blocked kernel claims BIT-EXACT equality with mmt_range (the serial reference every other matmul_t path is judged 3// against) on every shape: the interior 4x4 blocks, the ragged right edge (n not a multiple of 4), the ragged bottom edge 4// (m not a multiple of 4), both edges at once, m=1 decode rows, k below and above the block, and the cross-encoder's own 5// shapes (T x 384 x 1536, T x 1536 x 384, T x 384 x 384). Equality is asserted on every output cell of every shape -- 6// a tolerance would be the wrong oracle for a change that only re-interleaves independent sums. The fills are a seeded 7// integer LCG mapped to small f32 values so every product and sum is reproducible on any host. The gate also TIMES the 8// reference and the blocked kernel on the cross-encoder FFN shape and prints both as values (never a threshold: a 9// storming host would make a speed tooth flaky, and the number is the evidence the plan row cites). license_tier: ORIGINAL 10import "nx_f32_matmul_t.nx" 11import "nx_f32_cvt.nx" 12import "nx_syscalls.nx" 13import "nx_gate_verdict.nx" 14 15const G_I64: i64 = 8 16const G_SEED_A: i64 = 7 17const G_SEED_B: i64 = 11 18const G_LCG_MUL: i64 = 1103515245 19const G_LCG_ADD: i64 = 12345 20const G_LCG_MASK: i64 = 2147483647 21const G_FILL_MOD: i64 = 16 22const G_FILL_OFF: i64 = 8 23const G_T: i64 = 300 // a typical nfcorpus pair length 24const G_D: i64 = 384 25const G_I: i64 = 1536 26const G_SHAPES: i64 = 12 27const G_S_W: i64 = 3 28const G_USEC: i64 = 1000000 29const G_NSEC_PER_USEC: i64 = 1000 30const G_PERMIL: i64 = 1000 31 32func g_lcg(s: i64) -> i64 { var v: i64 = s * G_LCG_MUL + G_LCG_ADD; v = v & G_LCG_MASK; return v } 33func g_fill(p: *i64, count: i64, seed: i64) -> i64 { 34 var s: i64 = seed 35 var i: i64 = 0 36 while i < count { s = g_lcg(s); p[i] = nx_i32_to_f32((s - (s / G_FILL_MOD) * G_FILL_MOD) - G_FILL_OFF); i = i + 1 } 37 return 0 38} 39func g_same(a: *i64, b: *i64, count: i64) -> i64 { var i: i64 = 0; while i < count { if a[i] != b[i] { return 0 } i = i + 1 } return 1 } 40func g_diff_cells(a: *i64, b: *i64, count: i64) -> i64 { var d: i64 = 0; var i: i64 = 0; while i < count { if a[i] != b[i] { d = d + 1 } i = i + 1 } return d } 41func g_now_us() -> i64 { let ts: *i64 = sys_mmap(2 * G_I64) as *i64; sys_clock_gettime_mono(ts); let v: i64 = ts[0] * G_USEC + ts[1] / G_NSEC_PER_USEC; sys_munmap(ts as *u8, 2 * G_I64); return v } 42// one shape: fill, run both, compare every cell; returns the differing cell count (0 = bit-exact) 43func g_shape(m: i64, k: i64, n: i64, seed: i64) -> i64 { 44 let A: *i64 = sys_mmap(m * k * G_I64) as *i64 45 let B: *i64 = sys_mmap(n * k * G_I64) as *i64 46 let C1: *i64 = sys_mmap(m * n * G_I64) as *i64 47 let C2: *i64 = sys_mmap(m * n * G_I64) as *i64 48 g_fill(A, m * k, seed) 49 g_fill(B, n * k, seed + G_SEED_B) 50 mmt_range(A, B, C1, k, n, 0, m * n) 51 nx_f32_matmul_t_blocked(A, B, C2, m, k, n) 52 let d: i64 = g_diff_cells(C1, C2, m * n) 53 sys_munmap(A as *u8, m * k * G_I64); sys_munmap(B as *u8, n * k * G_I64); sys_munmap(C1 as *u8, m * n * G_I64); sys_munmap(C2 as *u8, m * n * G_I64) 54 return d 55} 56 57func main() -> i64 { 58 let ctr: *i64 = gv_ctr() 59 gv_head("=== NX-F32 MATMUL BLOCK GATE -- the 4x4 register-blocked matmul_t is bit-exact with the serial reference on interior blocks, both ragged edges, decode rows and the cross-encoder shapes; timed on the FFN shape ===" as *u8) 60 gv_check_eq("T0-exact-4x4-single-block" as *u8, g_shape(4, 4, 4, G_SEED_A), 0, ctr) 61 gv_check_eq("T1-interior-only-8x16x12" as *u8, g_shape(8, 16, 12, G_SEED_A + 1), 0, ctr) 62 gv_check_eq("T2-ragged-right-edge-n-not-a-multiple-of-4" as *u8, g_shape(8, 16, 13, G_SEED_A + 2), 0, ctr) 63 gv_check_eq("T3-ragged-bottom-edge-m-not-a-multiple-of-4" as *u8, g_shape(7, 16, 12, G_SEED_A + 3), 0, ctr) 64 gv_check_eq("T4-both-edges-ragged" as *u8, g_shape(5, 7, 3, G_SEED_A + 4), 0, ctr) 65 gv_check_eq("T5-decode-row-m-equals-1" as *u8, g_shape(1, G_D, G_D, G_SEED_A + 5), 0, ctr) 66 gv_check_eq("T6-one-column-n-equals-1-the-classifier" as *u8, g_shape(1, G_D, 1, G_SEED_A + 6), 0, ctr) 67 gv_check_eq("T7-k-of-one" as *u8, g_shape(6, 1, 6, G_SEED_A + 7), 0, ctr) 68 gv_check_eq("T8-cross-encoder-qkv-shape" as *u8, g_shape(G_T, G_D, G_D, G_SEED_A + 8), 0, ctr) 69 gv_check_eq("T9-cross-encoder-ffn-up-shape" as *u8, g_shape(G_T, G_D, G_I, G_SEED_A + 9), 0, ctr) 70 gv_check_eq("T10-cross-encoder-ffn-down-shape" as *u8, g_shape(G_T, G_I, G_D, G_SEED_A + 10), 0, ctr) 71 gv_check_eq("T11-llm-prefill-shape-8x896x512" as *u8, g_shape(8, 896, 512, G_SEED_A + 11), 0, ctr) 72 // NEG-CONTROL: the comparer must be able to see a difference (a perturbed cell) 73 let m: i64 = 5 74 let pa: *i64 = sys_mmap(m * G_I64) as *i64 75 let pb: *i64 = sys_mmap(m * G_I64) as *i64 76 g_fill(pa, m, G_SEED_A); g_fill(pb, m, G_SEED_A) 77 pb[2] = nx_i32_to_f32(99) 78 gv_check_eq("T12-neg-control-the-cell-comparer-sees-one-perturbed-cell" as *u8, g_diff_cells(pa, pb, m), 1, ctr) 79 // timing on the FFN-up shape: values, not a threshold 80 let A: *i64 = sys_mmap(G_T * G_D * G_I64) as *i64 81 let B: *i64 = sys_mmap(G_I * G_D * G_I64) as *i64 82 let C1: *i64 = sys_mmap(G_T * G_I * G_I64) as *i64 83 let C2: *i64 = sys_mmap(G_T * G_I * G_I64) as *i64 84 g_fill(A, G_T * G_D, G_SEED_A); g_fill(B, G_I * G_D, G_SEED_B) 85 let t0: i64 = g_now_us() 86 mmt_range(A, B, C1, G_D, G_I, 0, G_T * G_I) 87 let t1: i64 = g_now_us() 88 nx_f32_matmul_t_blocked(A, B, C2, G_T, G_D, G_I) 89 let t2: i64 = g_now_us() 90 gv_check_eq("T13-timed-ffn-run-is-bit-exact-too" as *u8, g_diff_cells(C1, C2, G_T * G_I), 0, ctr) 91 var speed_permil: i64 = 0 92 if t2 - t1 > 0 { speed_permil = ((t1 - t0) * G_PERMIL) / (t2 - t1) } 93 gv_values_head() 94 gv_kv("ffn_shape_macs" as *u8, G_T * G_D * G_I) 95 gv_kv("reference_us" as *u8, t1 - t0) 96 gv_kv("blocked_us" as *u8, t2 - t1) 97 gv_kv("blocked_speedup_permil" as *u8, speed_permil) 98 return gv_verdict("nx_f32_matmul_block_gate" as *u8, ctr, "the 4x4 register-blocked matmul_t is bit-exact with the serial reference on every shape class and the cross-encoder shapes; its speed is printed as a value" as *u8) 99}