nx_matmul_tile_kat.nx source
↩ module page · 119 lines · 4702 B
1// nx_matmul_tile_kat.nx -- proves the m>1 COLUMN-TILED matmul is BIT-EXACT against the existing kernel,
2// and measures what hoisting the column actually buys.
3//
4// WHY: nx_batchscale_kat measured the matmul path FLAT across m (us_per_row 4705/4056/4295/4191 at
5// m=1/2/4/8, speedup 1.12x) -- the signature of each row re-streaming the weight column. mmt_range walks
6// OUTPUT CELLS, so column B[*,j] is read once PER ROW. mmt_tile_cols hoists that load: read the column
7// once, accumulate against all m rows.
8//
9// THE EQUALITY CLAIM IS NOT A TOLERANCE CLAIM. For any (i,j) both kernels start at f32 zero and add over
10// kk ASCENDING -- identical order, so identical rounding. Only the interleaving across i differs. Anything
11// other than BIT-EQUALITY here means the tiled kernel changed the arithmetic, and that is a REJECT, not a
12// rounding difference to be waved through.
13//
14// T4 is the load-bearing negative control: a DELIBERATELY WRONG tiled result must be caught by the same
15// comparison, so a passing T1-T3 cannot be an artefact of comparing a buffer with itself.
16// license_tier: ORIGINAL No hw writes (Rule 26). expect_exit: 0
17import "nx_f32_matmul_t.nx"
18import "nx_f32_cvt.nx"
19import "nx_fmt.nx"
20
21const TK: i64 = 896
22const TN: i64 = 512
23const TM: i64 = 8
24
25func tk_lcg(s: i64) -> i64 {
26 var v: i64 = s * 1103515245 + 12345
27 v = v & 2147483647
28 return v
29}
30
31func tk_fill(p: *i64, count: i64, seed: i64) -> i64 {
32 var s: i64 = seed
33 var i: i64 = 0
34 while i < count {
35 s = tk_lcg(s)
36 p[i] = nx_i32_to_f32((s - (s / 16) * 16) - 8)
37 i = i + 1
38 }
39 return 0
40}
41
42func tk_same(a: *i64, b: *i64, count: i64) -> i64 {
43 var i: i64 = 0
44 while i < count { if a[i] != b[i] { return 0 } i = i + 1 }
45 return 1
46}
47
48func tk_nl() -> i64 { fmt_puts("\n" as *u8); return 0 }
49
50func tk_t(name: *u8, cond: i64, ctr: *i64) {
51 if cond == 1 { fmt_puts(" ok " as *u8); ctr[0] = ctr[0] + 1 }
52 else { fmt_puts(" FAIL " as *u8) }
53 fmt_puts(name)
54 tk_nl()
55 ctr[1] = ctr[1] + 1
56}
57
58func main() -> i64 {
59 var ctr: *i64 = sys_mmap(64) as *i64
60 ctr[0] = 0
61 ctr[1] = 0
62
63 let A: *i64 = sys_mmap(TM * TK * 8) as *i64
64 let B: *i64 = sys_mmap(TK * TN * 8) as *i64
65 let Cref: *i64 = sys_mmap(TM * TN * 8) as *i64
66 let Ctil: *i64 = sys_mmap(TM * TN * 8) as *i64
67 tk_fill(A, TM * TK, 20260731)
68 tk_fill(B, TK * TN, 90210)
69
70 fmt_puts("=== nx_matmul_tile_kat -- column-tiled m>1 matmul: bit-exact? and how much faster? ===" as *u8); tk_nl()
71 fmt_puts("shape m=" as *u8); fmt_putn(TM); fmt_puts(" k=" as *u8); fmt_putn(TK)
72 fmt_puts(" n=" as *u8); fmt_putn(TN); tk_nl()
73
74 let t0: i64 = sys_now_us()
75 nx_f32_matmul_t(A, B, Cref, TM, TK, TN)
76 let us_ref: i64 = sys_now_us() - t0
77
78 let t1: i64 = sys_now_us()
79 nx_f32_matmul_t_tiled(A, B, Ctil, TM, TK, TN)
80 let us_til: i64 = sys_now_us() - t1
81
82 tk_t("T1 tiled == existing kernel, BIT-EXACT on all m*n cells" as *u8,
83 tk_same(Cref, Ctil, TM * TN), ctr)
84
85 tk_t("T2 NON-VACUITY: the reference actually computed something (not all zero)" as *u8,
86 Cref[0] != 0, ctr)
87
88 tk_t("T3 m=1 also matches (the tiled path must not regress the decode case)" as *u8,
89 tk_m1_ok(A, B, Cref, Ctil), ctr)
90
91 // NEG-CONTROL: corrupt one cell and prove the SAME comparison rejects it.
92 Ctil[TM * TN / 2] = Ctil[TM * TN / 2] + 1
93 tk_t("T4 NEG-CONTROL: a single corrupted cell IS caught (T1 is not vacuous)" as *u8,
94 tk_same(Cref, Ctil, TM * TN) == 0, ctr)
95
96 tk_nl()
97 fmt_puts(" existing_us=" as *u8); fmt_putn(us_ref)
98 fmt_puts(" tiled_us=" as *u8); fmt_putn(us_til)
99 var sx100: i64 = 0
100 if us_til > 0 { sx100 = us_ref * 100 / us_til }
101 fmt_puts(" speedup_x100=" as *u8); fmt_putn(sx100); tk_nl()
102 fmt_puts(" (weight traffic m*k*n -> k*n; at m=" as *u8); fmt_putn(TM)
103 fmt_puts(" that is " as *u8); fmt_putn(TM); fmt_puts("x fewer weight reads)" as *u8); tk_nl()
104 fmt_puts("envelope: ONE timing run on a shared host -- one window is not a rate. The BIT-EXACT result" as *u8); tk_nl()
105 fmt_puts(" is the durable claim; the speedup number is indicative and should be re-read under load." as *u8); tk_nl()
106
107 fmt_puts("MATMUL-TILE-KAT " as *u8); fmt_putn(ctr[0]); fmt_puts("/" as *u8); fmt_putn(ctr[1])
108 if ctr[0] == ctr[1] { fmt_puts(" GREEN" as *u8); tk_nl(); return 0 }
109 fmt_puts(" RED" as *u8); tk_nl()
110 return 1
111}
112
113func tk_m1_ok(A: *i64, B: *i64, r1: *i64, r2: *i64) -> i64 {
114 let c1: *i64 = sys_mmap(TN * 8) as *i64
115 let c2: *i64 = sys_mmap(TN * 8) as *i64
116 nx_f32_matmul_t(A, B, c1, 1, TK, TN)
117 nx_f32_matmul_t_tiled(A, B, c2, 1, TK, TN)
118 return tk_same(c1, c2, TN)
119}