`nx_f32_matmul.nx` source

1// nx_f32_matmul.nx -- bits-up f32 dot product + GEMM kernel.
2//
3// L7 composition brick.  Last load-bearing math primitive for an
4// end-to-end f32-mode transformer block.  Composes nx_f32_mul +
5// nx_f32_add over rows / columns.
6//
7// API:
8//   nx_f32_dot(a, b, n)         -- scalar dot product (sum of a[i]*b[i])
9//   nx_f32_matmul(a, b, c, m, k, n)
10//                                -- C[m,n] = A[m,k] @ B[k,n]
11//                                   row-major storage
12//
13// Bit-exact when all operands and intermediate sums are
14// representable f32 (e.g., small ints).  Otherwise round-to-
15// nearest-even accumulates per the L4 f32 ops.
16//
17// No libm.  Inner MAC uses nx_f32_hw SSE intrinsics (__f32_add/__f32_mul -> real addss/mulss on the FPU),
18// which are BIT-EXACT vs the bits-up software f32 (proven by nx_f32_hw_matmul_gate's hw==sw differential) but
19// replace dozens of integer ops per multiply with one instruction -- the hardware rung the LLM forward stands on.
20import "nx_syscalls.nx"
21import "nx_tier.nx"
22import "nx_f32.nx"
23import "nx_f32_hw.nx"
24
25const NX_F32_MM_OK:           nx_int = 0
26const NX_F32_MM_ERR_BAD_DIM:  nx_int = 1
27const NX_F32_MM_N_VERDICTS:   nx_int = 2
28
29func nx_f32_mm_verdict_is_valid(v: nx_int) -> nx_int {
30    if v < 0 { return 0 }
31    if v >= NX_F32_MM_N_VERDICTS { return 0 }
32    return 1
33}
34
35// Scalar dot product: sum_i a[i] * b[i], all f32 raw bits.
36
37func nx_f32_dot(a: *i64, b: *i64, n: nx_int) -> i64 {
38    var acc: i64 = 0
39    var i: nx_int = 0
40    while i < n {
41        acc = __f32_add(acc, __f32_mul(a[i], b[i]))
42        i = i + 1
43    }
44    return acc
45}
46
47// GEMM: C[m,n] = A[m,k] @ B[k,n], row-major.
48//
49// C[i,j] = sum_l A[i,l] * B[l,j]
50//
51// For typical LLM forward (single token), m=1, so the outer loop
52// runs once and this degenerates to k matrix-vector products.
53
54func nx_f32_matmul(a: *i64, b: *i64, c: *i64,
55                   m: nx_int, k: nx_int, n: nx_int) -> nx_int {
56    if m <= 0 { return NX_F32_MM_ERR_BAD_DIM }
57    if k <= 0 { return NX_F32_MM_ERR_BAD_DIM }
58    if n <= 0 { return NX_F32_MM_ERR_BAD_DIM }
59
60    var i: nx_int = 0
61    while i < m {
62        var j: nx_int = 0
63        while j < n {
64            var acc: i64 = 0
65            var l: nx_int = 0
66            while l < k {
67                let a_il: i64 = a[i * k + l]
68                let b_lj: i64 = b[l * n + j]
69                acc = __f32_add(acc, __f32_mul(a_il, b_lj))
70                l = l + 1
71            }
72            c[i * n + j] = acc
73            j = j + 1
74        }
75        i = i + 1
76    }
77    return NX_F32_MM_OK
78}
nx_f32_matmul.nx source

`nx_f32_matmul.nx` source
