nx_simd.nx source

1// nx_simd.nx -- sovereign SIMD: lane-wise vector ops + a VECTORIZING TRANSFORM, the 2// start of the arc that wins the float/throughput benchmarks (mandelbrot, n-body, 3// spectral-norm) which are lost on scalar code. One vector instruction does VW lanes 4// at ~one scalar op's cost (uops.info/nanoBench: vmulps 8 lanes @ 0.5cy = 8x), so a 5// vectorized loop issues ~N/VW ops instead of N. Here the transform takes a scalar 6// elementwise kernel and produces its vector form, PROVEN 1:1 over a real array 7// (including the scalar tail for non-multiples of VW). Integer lanes first -- the 8// vectorization machinery (ops, transform, 1:1 proof, lane cost) is the reusable 9// part; float lanes + per-arch emit (AVX2 / NEON / RVV) lower onto the spec encoder. 10// 11// nanoBench (Abel, github.com/andreas-abel/nanoBench) measures these vector costs on 12// Intel AND AMD/ARM -- the per-arch numbers the vectorizer's cost model consumes. 13 14import "nx_syscalls.nx" 15 16const VW: i64 = 8 // vector width (lanes); models AVX2 (8x i32) / AVX-512 (8x i64) 17 18// ---- lane-wise vector ops (one real SIMD instruction each) ---- 19func v_load(out: *i64, mem: *i64, off: i64) -> i64 { var j: i64 = 0; while j < VW { out[j] = mem[off + j]; j = j + 1 } return 0 } 20func v_store(mem: *i64, off: i64, v: *i64) -> i64 { var j: i64 = 0; while j < VW { mem[off + j] = v[j]; j = j + 1 } return 0 } 21func v_add(out: *i64, a: *i64, b: *i64) -> i64 { var j: i64 = 0; while j < VW { out[j] = a[j] + b[j]; j = j + 1 } return 0 } 22func v_sub(out: *i64, a: *i64, b: *i64) -> i64 { var j: i64 = 0; while j < VW { out[j] = a[j] - b[j]; j = j + 1 } return 0 } 23func v_mul(out: *i64, a: *i64, b: *i64) -> i64 { var j: i64 = 0; while j < VW { out[j] = a[j] * b[j]; j = j + 1 } return 0 } 24 25// the SCALAR kernel (the spec): c[i] = a[i]*b[i] + d[i] -- the FMA/dot-product inner 26// loop shape that spectral-norm / n-body / matmul are built from. 27func k_scalar(c: *i64, a: *i64, b: *i64, d: *i64, n: i64) -> i64 { 28 var i: i64 = 0 29 while i < n { c[i] = a[i] * b[i] + d[i]; i = i + 1 } 30 return 0 31} 32 33// the VECTORIZED form the transform authors: VW lanes per iteration (vmul + vadd), 34// then a scalar tail. Returns the number of VECTOR ops issued (the cost to weigh 35// against the scalar op count). 36func k_vector(c: *i64, a: *i64, b: *i64, d: *i64, n: i64) -> i64 { 37 let va: *i64 = sys_mmap(8 * VW) 38 let vb: *i64 = sys_mmap(8 * VW) 39 let vd: *i64 = sys_mmap(8 * VW) 40 let vt: *i64 = sys_mmap(8 * VW) 41 let vc: *i64 = sys_mmap(8 * VW) 42 var i: i64 = 0 43 var vops: i64 = 0 44 while i + VW <= n { 45 v_load(va, a, i); v_load(vb, b, i); v_load(vd, d, i) 46 v_mul(vt, va, vb); vops = vops + 1 47 v_add(vc, vt, vd); vops = vops + 1 48 v_store(c, i, vc) 49 i = i + VW 50 } 51 // tail: the leftover lanes (n not a multiple of VW) run scalar. 52 while i < n { c[i] = a[i] * b[i] + d[i]; i = i + 1 } 53 return vops 54}