nx_natmm_bench.nx source
↩ module page · 47 lines · 1934 B
1// nx_natmm_bench.nx -- ISOLATED native Q8 matmul: serial (_st) vs fork-join (_pteam).
2// No LLM, no load, no block ops -- just the kernel, looped, so external best-of-3
3// timing is clean. MODE 0=serial, 1=pteam. NG=n (outputs), MK=k (hidden).
4// Answers: does native fork-join speed up the memory-bound Q8 matmul, isolated?
5// license_tier: ORIGINAL expect_exit: 0
6import "nx_syscalls.nx"
7import "nx_f32.nx"
8import "nx_f32_lazy_weight.nx"
9
10const NX_MM_MODE: i64 = 1 // 0 = serial (_st) ; 1 = fork-join (_pteam)
11const NX_MM_ITER: i64 = 30
12const NX_MM_NG: i64 = 151936 // n = vocab (lm_head) -- the biggest matmul
13const NX_MM_MK: i64 = 896 // k = hidden
14
15func main() -> i64 {
16 let nblk: i64 = NX_MM_MK / 32
17 let bytes: i64 = NX_MM_NG * nblk * 34 // Q8_0: 34 B/block, nblk/row, NG rows
18 let w: *u8 = sys_mmap(bytes)
19 // touch every page + keep d/quants small (avoid f16 inf/nan): all bytes = 1
20 let n8: i64 = bytes / 8
21 let wi: *i64 = w as *i64
22 var i: i64 = 0
23 while i < n8 { wi[i] = 0x0101010101010101; i = i + 1 }
24 let W: *NxF32LazyWeight = nx_f32_lazy_weight_new_q8_0(w, 0, NX_MM_NG, NX_MM_MK)
25
26 let A: *i64 = sys_mmap(NX_MM_MK * 8) as *i64 // m=1 x k activation (f32 bits)
27 i = 0
28 while i < NX_MM_MK { A[i] = 0x3F800000; i = i + 1 } // 1.0
29 let C: *i64 = sys_mmap(NX_MM_NG * 8) as *i64 // m=1 x n output
30
31 sys_write(1, "GO\n" as *u8, 3)
32 var it: i64 = 0
33 while it < NX_MM_ITER {
34 if NX_MM_MODE == 0 { _lw_q8_0_matmul_st(W, A, C, 1, NX_MM_MK, NX_MM_NG) }
35 else { _lw_q8_0_matmul_pteam(W, A, C, 1, NX_MM_MK, NX_MM_NG) }
36 it = it + 1
37 }
38 // print a digit of the checksum so the loop stays live
39 var s: i64 = C[0] + C[NX_MM_NG - 1]
40 if s < 0 { s = 0 - s }
41 let o: *u8 = sys_mmap(16)
42 o[0] = (48 + (s - (s / 10) * 10)) as u8
43 o[1] = 10 as u8
44 sys_write(1, o, 2)
45 sys_exit(0)
46 return 0
47}