nx_native_matmul_bench.nx source

1// nx_native_matmul_bench.nx -- matmul-shaped native micro-bench. 2// 3// Most LLM-relevant micro-bench in the paired suite. llama.cpp's 4// transformer forward is dominated by GEMM kernels; this measures 5// how Nishi's native codegen handles the same access pattern at a 6// tractable size. 7// 8// Workload: C = A * B for 128x128 i64 matrices. ~2M multiply-adds. 9// Returns (C[0][0] & 0xFF) as exit code. 10// 11// Compares against gcc -O0 (no opt) and gcc -O2 (auto-vectorized 12// matmul) on the same machine for paired comparison. 13 14import "nx_syscalls.nx" 15 16func main() -> i64 { 17 let N: i64 = 256 // 256x256 matrix = 16.7M multiply-adds 18 19 let bufA: *u8 = sys_mmap(N * N * 8) 20 let bufB: *u8 = sys_mmap(N * N * 8) 21 let bufC: *u8 = sys_mmap(N * N * 8) 22 let A: *i64 = bufA as *i64 23 let B: *i64 = bufB as *i64 24 let C: *i64 = bufC as *i64 25 26 // Initialize A[i][j] = (i + j) & 255; B[i][j] = (i * j + 1) & 255 27 var i: i64 = 0 28 while i < N { 29 var j: i64 = 0 30 while j < N { 31 A[i * N + j] = (i + j) & 255 32 B[i * N + j] = (i * j + 1) & 255 33 C[i * N + j] = 0 34 j = j + 1 35 } 36 i = i + 1 37 } 38 39 // C = A * B (triple-loop matmul; standard ijk ordering) 40 var ii: i64 = 0 41 while ii < N { 42 var jj: i64 = 0 43 while jj < N { 44 var sum: i64 = 0 45 var k: i64 = 0 46 while k < N { 47 sum = sum + A[ii * N + k] * B[k * N + jj] 48 k = k + 1 49 } 50 C[ii * N + jj] = sum 51 jj = jj + 1 52 } 53 ii = ii + 1 54 } 55 56 return C[0] 57}