nx_blas_i64.nx source
↩ module page · 238 lines · 8860 B
1// nx_blas_i64.nx -- BLAS Level-1/2/3 reference kernels on i64+Q10.
2//
3// First kernel module on the sovereign-from-bits-up ML stack.
4// Iteration I3 per nxc2/docs/MULTIMODAL_ORCHESTRATOR_ROADMAP.md.
5//
6// Per the "world class without the comfyui annoyance but the power
7// and more" framing: we ship the CORRECT version first (slow, simple,
8// auditable). AlphaTensor-discovered (Fawzi 2022) / AlphaEvolve-
9// proposed (Romera-Paredes 2025) variants land later as alternative
10// implementations of the same kernel, gated by nx_numeric_oracle
11// verdict against this reference.
12//
13// Layout: row-major contiguous tensors only (v1). Strided / view
14// support lands when permute becomes a hot path.
15//
16// Operations shipped:
17// gemm -- C = alpha * A @ B + beta * C (Level 3)
18// matmul -- gemm with alpha=1, beta=0
19// dot -- scalar = sum(x[i] * y[i]) (Level 1)
20// axpy -- y = alpha * x + y (Level 1)
21// scale -- x = alpha * x (Level 1, BLAS-1 _scal)
22// gemv -- y = alpha * A @ x + beta * y (Level 2)
23//
24// Q-format convention: alpha and beta are Q10 scaling factors;
25// tensor elements are plain i64 (NOT pre-scaled). Result element =
26// (raw_dot * alpha_q10 + raw_c * beta_q10) / 1024.
27//
28// genealogy_id: blis_van_zee_2015 + goto_2008_matmul +
29// alphatensor_fawzi_2022 + strassen_1969 +
30// ieee_754_kahan_1965
31// lineage_id: substrate_blas_i64_v1
32
33// nx_safety_envelope:
34// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
35// sil_target: SIL1
36// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
37// verdict: NOT_YET_EVALUATED
38
39import "nx_syscalls.nx"
40import "nx_tier.nx"
41import "nx_tensor.nx"
42
43const NX_BLAS_Q10: nx_int = 1024
44
45// ===== Sealed-enum: GemmVerdict ===================================
46
47const NX_BLAS_OK: nx_int = 0
48const NX_BLAS_ERR_BAD_DTYPE: nx_int = 1 // not I64
49const NX_BLAS_ERR_BAD_NDIM: nx_int = 2 // matmul needs ndim==2
50const NX_BLAS_ERR_SHAPE_MISMATCH: nx_int = 3 // K dimension or M/N mismatch
51const NX_BLAS_ERR_NOT_CONTIGUOUS: nx_int = 4 // v1 row-major only
52const NX_BLAS_N_VERDICTS: nx_int = 5
53
54func nx_blas_verdict_is_valid(v: nx_int) -> nx_int {
55 if v < 0 { return 0 }
56 if v >= NX_BLAS_N_VERDICTS { return 0 }
57 return 1
58}
59
60// ===== Internal: validate two matrices for matmul =================
61
62func _blas_check_matmul(a: *NxTensor, b: *NxTensor, c: *NxTensor) -> nx_int {
63 if a.dtype != NX_DT_I64 { return NX_BLAS_ERR_BAD_DTYPE }
64 if b.dtype != NX_DT_I64 { return NX_BLAS_ERR_BAD_DTYPE }
65 if c.dtype != NX_DT_I64 { return NX_BLAS_ERR_BAD_DTYPE }
66
67 if a.ndim != 2 { return NX_BLAS_ERR_BAD_NDIM }
68 if b.ndim != 2 { return NX_BLAS_ERR_BAD_NDIM }
69 if c.ndim != 2 { return NX_BLAS_ERR_BAD_NDIM }
70
71 // a: M x K, b: K x N, c: M x N
72 if a.shape[1] != b.shape[0] { return NX_BLAS_ERR_SHAPE_MISMATCH }
73 if c.shape[0] != a.shape[0] { return NX_BLAS_ERR_SHAPE_MISMATCH }
74 if c.shape[1] != b.shape[1] { return NX_BLAS_ERR_SHAPE_MISMATCH }
75
76 if nx_t_is_contiguous(a) == 0 { return NX_BLAS_ERR_NOT_CONTIGUOUS }
77 if nx_t_is_contiguous(b) == 0 { return NX_BLAS_ERR_NOT_CONTIGUOUS }
78 if nx_t_is_contiguous(c) == 0 { return NX_BLAS_ERR_NOT_CONTIGUOUS }
79
80 return NX_BLAS_OK
81}
82
83// ===== gemm =======================================================
84//
85// C[m, n] = (alpha_q10 * sum_k(A[m, k] * B[k, n]) + beta_q10 * C[m, n]) / Q10
86//
87// All-i64 arithmetic; integer overflow is the caller's responsibility
88// for now (Q-format substrate; bounded magnitudes). Future iteration
89// adds Kahan-style compensated accumulation when F32 lands.
90
91func nx_blas_gemm(a: *NxTensor, b: *NxTensor, c: *NxTensor,
92 alpha_q10: nx_int, beta_q10: nx_int) -> nx_int {
93 let pre: nx_int = _blas_check_matmul(a, b, c)
94 if pre != NX_BLAS_OK { return pre }
95
96 let M: nx_int = a.shape[0]
97 let K: nx_int = a.shape[1]
98 let N: nx_int = b.shape[1]
99
100 let pa: *i64 = a.storage as *i64
101 let pb: *i64 = b.storage as *i64
102 let pc: *i64 = c.storage as *i64
103
104 var m: nx_int = 0
105 while m < M {
106 var n: nx_int = 0
107 while n < N {
108 // Accumulate dot(A[m, :], B[:, n]) in i64
109 var acc: nx_int = 0
110 var k: nx_int = 0
111 while k < K {
112 acc = acc + pa[m * K + k] * pb[k * N + n]
113 k = k + 1
114 }
115 let c_old: nx_int = pc[m * N + n]
116 // Q10 weighted combine
117 pc[m * N + n] = (acc * alpha_q10 + c_old * beta_q10) / NX_BLAS_Q10
118 n = n + 1
119 }
120 m = m + 1
121 }
122 return NX_BLAS_OK
123}
124
125// ===== matmul =====================================================
126//
127// C = A @ B. Convenience wrapper for gemm with alpha=1, beta=0.
128// C is OVERWRITTEN (its prior contents are ignored).
129
130func nx_blas_matmul(a: *NxTensor, b: *NxTensor, c: *NxTensor) -> nx_int {
131 return nx_blas_gemm(a, b, c, NX_BLAS_Q10, 0)
132}
133
134// ===== dot ========================================================
135//
136// scalar = sum_i(x[i] * y[i])
137// x, y: 1-D i64 tensors of identical length. Returns the scalar in
138// out[0]; returns sealed verdict.
139
140func nx_blas_dot(x: *NxTensor, y: *NxTensor, out: *i64) -> nx_int {
141 if x.dtype != NX_DT_I64 { return NX_BLAS_ERR_BAD_DTYPE }
142 if y.dtype != NX_DT_I64 { return NX_BLAS_ERR_BAD_DTYPE }
143 if x.ndim != 1 { return NX_BLAS_ERR_BAD_NDIM }
144 if y.ndim != 1 { return NX_BLAS_ERR_BAD_NDIM }
145 if x.shape[0] != y.shape[0] { return NX_BLAS_ERR_SHAPE_MISMATCH }
146 if nx_t_is_contiguous(x) == 0 { return NX_BLAS_ERR_NOT_CONTIGUOUS }
147 if nx_t_is_contiguous(y) == 0 { return NX_BLAS_ERR_NOT_CONTIGUOUS }
148
149 let px: *i64 = x.storage as *i64
150 let py: *i64 = y.storage as *i64
151 var acc: nx_int = 0
152 var i: nx_int = 0
153 while i < x.shape[0] {
154 acc = acc + px[i] * py[i]
155 i = i + 1
156 }
157 out[0] = acc
158 return NX_BLAS_OK
159}
160
161// ===== axpy =======================================================
162//
163// y = alpha_q10 * x + y (Q10 scaling)
164// Both 1-D, same length. Mutates y in place.
165
166func nx_blas_axpy(alpha_q10: nx_int, x: *NxTensor, y: *NxTensor) -> nx_int {
167 if x.dtype != NX_DT_I64 { return NX_BLAS_ERR_BAD_DTYPE }
168 if y.dtype != NX_DT_I64 { return NX_BLAS_ERR_BAD_DTYPE }
169 if x.ndim != 1 { return NX_BLAS_ERR_BAD_NDIM }
170 if y.ndim != 1 { return NX_BLAS_ERR_BAD_NDIM }
171 if x.shape[0] != y.shape[0] { return NX_BLAS_ERR_SHAPE_MISMATCH }
172 if nx_t_is_contiguous(x) == 0 { return NX_BLAS_ERR_NOT_CONTIGUOUS }
173 if nx_t_is_contiguous(y) == 0 { return NX_BLAS_ERR_NOT_CONTIGUOUS }
174
175 let px: *i64 = x.storage as *i64
176 let py: *i64 = y.storage as *i64
177 var i: nx_int = 0
178 while i < x.shape[0] {
179 py[i] = (alpha_q10 * px[i]) / NX_BLAS_Q10 + py[i]
180 i = i + 1
181 }
182 return NX_BLAS_OK
183}
184
185// ===== scale ======================================================
186//
187// x = alpha_q10 * x (Q10 scaling). In-place.
188
189func nx_blas_scale(alpha_q10: nx_int, x: *NxTensor) -> nx_int {
190 if x.dtype != NX_DT_I64 { return NX_BLAS_ERR_BAD_DTYPE }
191 if nx_t_is_contiguous(x) == 0 { return NX_BLAS_ERR_NOT_CONTIGUOUS }
192 let px: *i64 = x.storage as *i64
193 var i: nx_int = 0
194 while i < x.numel {
195 px[i] = (alpha_q10 * px[i]) / NX_BLAS_Q10
196 i = i + 1
197 }
198 return NX_BLAS_OK
199}
200
201// ===== gemv =======================================================
202//
203// y = alpha_q10 * A @ x + beta_q10 * y (Level 2)
204// A: M x N, x: N, y: M
205
206func nx_blas_gemv(a: *NxTensor, x: *NxTensor, y: *NxTensor,
207 alpha_q10: nx_int, beta_q10: nx_int) -> nx_int {
208 if a.dtype != NX_DT_I64 { return NX_BLAS_ERR_BAD_DTYPE }
209 if x.dtype != NX_DT_I64 { return NX_BLAS_ERR_BAD_DTYPE }
210 if y.dtype != NX_DT_I64 { return NX_BLAS_ERR_BAD_DTYPE }
211 if a.ndim != 2 { return NX_BLAS_ERR_BAD_NDIM }
212 if x.ndim != 1 { return NX_BLAS_ERR_BAD_NDIM }
213 if y.ndim != 1 { return NX_BLAS_ERR_BAD_NDIM }
214 if a.shape[1] != x.shape[0] { return NX_BLAS_ERR_SHAPE_MISMATCH }
215 if a.shape[0] != y.shape[0] { return NX_BLAS_ERR_SHAPE_MISMATCH }
216 if nx_t_is_contiguous(a) == 0 { return NX_BLAS_ERR_NOT_CONTIGUOUS }
217 if nx_t_is_contiguous(x) == 0 { return NX_BLAS_ERR_NOT_CONTIGUOUS }
218 if nx_t_is_contiguous(y) == 0 { return NX_BLAS_ERR_NOT_CONTIGUOUS }
219
220 let M: nx_int = a.shape[0]
221 let N: nx_int = a.shape[1]
222 let pa: *i64 = a.storage as *i64
223 let px: *i64 = x.storage as *i64
224 let py: *i64 = y.storage as *i64
225
226 var m: nx_int = 0
227 while m < M {
228 var acc: nx_int = 0
229 var n: nx_int = 0
230 while n < N {
231 acc = acc + pa[m * N + n] * px[n]
232 n = n + 1
233 }
234 py[m] = (acc * alpha_q10 + py[m] * beta_q10) / NX_BLAS_Q10
235 m = m + 1
236 }
237 return NX_BLAS_OK
238}