nx_sparse_tensor.nx source
↩ module page · 248 lines · 8775 B
1// nx_sparse_tensor.nx -- compressed sparse row tensor + sparse matmul.
2//
3// FIRST algo-led-not-brute-force kernel per the
4// min-hardware-floor cardinal. Post-softmax attention is mostly zero;
5// LLM weights after pruning are mostly zero; world-model frame deltas
6// are mostly zero. Skipping those zeros is 10-100x speedup on the
7// substrate AND 10-100x memory reduction -- both essential for the
8// "$50 SBC beats Genie 3 in coherence" trophy.
9//
10// CSR (Compressed Sparse Row) format:
11//
12// row_ptr[m+1] -- row_ptr[i+1] - row_ptr[i] = nnz in row i
13// col_idx[nnz] -- column index of each non-zero
14// values[nnz] -- the non-zero values themselves
15//
16// For an M x N dense matrix with NNZ non-zeros:
17// dense storage = M * N * 8 bytes
18// CSR storage = (M+1) * 8 + NNZ * 16 bytes
19// crossover = NNZ < M * N / 2 (sparsity > 50%)
20// typical attention = NNZ ~ M * N / 100 (>99% sparse)
21//
22// Operations shipped:
23// dense_to_csr -- materialise a dense tensor as CSR
24// csr_to_dense -- back-convert (audit + numeric_oracle gate)
25// csr_matmul_dense -- C[m, n] = sum over non-zero entries of row m
26// times dense column n. Skips ZERO ROWS entirely.
27// O(NNZ * N) instead of O(M * K * N).
28//
29// Future variants (gated on oracle):
30// AlphaTensor 2022 / AlphaEvolve 2025 sparse-matmul algorithms
31// register here as alternative dispatch entries.
32//
33// genealogy_id: tinkham_1969_sparse_matrices + bell_garland_2009_csr +
34// pissanetzky_1984 + alphatensor_2022
35// lineage_id: substrate_sparse_tensor_v1
36
37// nx_safety_envelope:
38// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
39// sil_target: SIL1
40// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
41// verdict: NOT_YET_EVALUATED
42
43import "nx_syscalls.nx"
44import "nx_tier.nx"
45import "nx_tensor.nx"
46const NX_MAGIC_1024: i64 = 1024
47
48// ===== Sealed-enum: SparseFormat ===================================
49//
50// Closed taxonomy. v1 ships CSR; COO and CSC queued.
51
52const NX_SP_FMT_CSR: nx_int = 0 // Compressed Sparse Row
53const NX_SP_FMT_COO: nx_int = 1 // Coordinate (queued)
54const NX_SP_FMT_CSC: nx_int = 2 // Compressed Sparse Column (queued)
55const NX_SP_FMT_BSR: nx_int = 3 // Block Sparse Row (queued; matches Goto's block-tiling)
56const NX_SP_FMT_N_KINDS: nx_int = 4
57
58func nx_sp_fmt_is_valid(f: nx_int) -> nx_int {
59 if f < 0 { return 0 }
60 if f >= NX_SP_FMT_N_KINDS { return 0 }
61 return 1
62}
63
64// ===== Sealed-enum: SparseVerdict =================================
65
66const NX_SP_OK: nx_int = 0
67const NX_SP_ERR_BAD_DTYPE: nx_int = 1
68const NX_SP_ERR_BAD_NDIM: nx_int = 2
69const NX_SP_ERR_SHAPE_MISMATCH: nx_int = 3
70const NX_SP_ERR_NNZ_OVERFLOW: nx_int = 4
71const NX_SP_ERR_UNSUPPORTED: nx_int = 5
72const NX_SP_N_VERDICTS: nx_int = 6
73
74func nx_sp_verdict_is_valid(v: nx_int) -> nx_int {
75 if v < 0 { return 0 }
76 if v >= NX_SP_N_VERDICTS { return 0 }
77 return 1
78}
79
80// ===== CSR tensor struct ===========================================
81
82struct NxSparseTensor {
83 format: nx_int,
84 n_rows: nx_int,
85 n_cols: nx_int,
86 nnz: nx_int,
87 dtype: nx_int,
88 row_ptr: *i64, // [n_rows + 1]
89 col_idx: *i64, // [nnz]
90 values: *i64 // [nnz]
91}
92
93const NX_SP_BYTES: nx_int = 64 // 8 fields * 8
94
95// ===== Allocator ==================================================
96//
97// Caller specifies maximum NNZ. Buffers sized for the worst case;
98// nnz field tracks actual. This is the substrate convention
99// (no realloc; predictable footprint for MCU).
100
101func nx_sp_alloc(n_rows: nx_int, n_cols: nx_int, max_nnz: nx_int,
102 dtype: nx_int) -> *NxSparseTensor {
103 let s: *NxSparseTensor = (sys_mmap(NX_SP_BYTES)) as *NxSparseTensor
104 s.format = NX_SP_FMT_CSR
105 s.n_rows = n_rows
106 s.n_cols = n_cols
107 s.nnz = 0
108 s.dtype = dtype
109 s.row_ptr = (sys_mmap((n_rows + 1) * NX_SIZEOF_NX_INT)) as *i64
110 s.col_idx = (sys_mmap(max_nnz * NX_SIZEOF_NX_INT)) as *i64
111 s.values = (sys_mmap(max_nnz * NX_SIZEOF_NX_INT)) as *i64
112 var i: nx_int = 0
113 while i <= n_rows {
114 s.row_ptr[i] = 0
115 i = i + 1
116 }
117 return s
118}
119
120// ===== Convert dense -> CSR =======================================
121//
122// O(M*N) scan; counts non-zeros first, then fills. Caller's max_nnz
123// must accommodate the actual non-zero count or we return
124// NX_SP_ERR_NNZ_OVERFLOW.
125
126func nx_sp_dense_to_csr(t: *NxTensor, s: *NxSparseTensor) -> nx_int {
127 if t.dtype != NX_DT_I64 { return NX_SP_ERR_BAD_DTYPE }
128 if t.ndim != 2 { return NX_SP_ERR_BAD_NDIM }
129 if t.shape[0] != s.n_rows { return NX_SP_ERR_SHAPE_MISMATCH }
130 if t.shape[1] != s.n_cols { return NX_SP_ERR_SHAPE_MISMATCH }
131
132 let pt: *i64 = t.storage as *i64
133 var nnz: nx_int = 0
134 var m: nx_int = 0
135 while m < s.n_rows {
136 s.row_ptr[m] = nnz
137 var n: nx_int = 0
138 while n < s.n_cols {
139 let v: nx_int = pt[m * s.n_cols + n]
140 if v != 0 {
141 s.col_idx[nnz] = n
142 s.values[nnz] = v
143 nnz = nnz + 1
144 }
145 n = n + 1
146 }
147 m = m + 1
148 }
149 s.row_ptr[s.n_rows] = nnz
150 s.nnz = nnz
151 return NX_SP_OK
152}
153
154// ===== Convert CSR -> dense =======================================
155//
156// Required for oracle verification. Allocates output via nx_tensor;
157// fills zeros, then scatter non-zeros from CSR storage.
158
159func nx_sp_csr_to_dense(s: *NxSparseTensor, out: *NxTensor) -> nx_int {
160 if out.dtype != NX_DT_I64 { return NX_SP_ERR_BAD_DTYPE }
161 if out.ndim != 2 { return NX_SP_ERR_BAD_NDIM }
162 if out.shape[0] != s.n_rows { return NX_SP_ERR_SHAPE_MISMATCH }
163 if out.shape[1] != s.n_cols { return NX_SP_ERR_SHAPE_MISMATCH }
164
165 nx_t_fill_zero(out)
166 let po: *i64 = out.storage as *i64
167 var m: nx_int = 0
168 while m < s.n_rows {
169 var k: nx_int = s.row_ptr[m]
170 let stop: nx_int = s.row_ptr[m + 1]
171 while k < stop {
172 let n: nx_int = s.col_idx[k]
173 po[m * s.n_cols + n] = s.values[k]
174 k = k + 1
175 }
176 m = m + 1
177 }
178 return NX_SP_OK
179}
180
181// ===== Sparse matmul: csr A (MxK) * dense B (KxN) -> dense C (MxN)
182//
183// Standard sparse-times-dense matmul. For each output row m, walk
184// only the non-zeros of A's row m. For each non-zero A[m, k], do
185// the AXPY: C[m, :] += A_value * B[k, :].
186//
187// Complexity: O(NNZ_A * N) instead of dense O(M * K * N).
188// Speedup ~= K / avg_nnz_per_row. At 99% sparsity, K=1024 -> 100x.
189//
190// MIN-HARDWARE WIN: memory bandwidth is the bottleneck on small SoCs;
191// CSR cuts the bandwidth by the sparsity factor.
192
193func nx_sp_csr_matmul_dense(a: *NxSparseTensor, b: *NxTensor, c: *NxTensor) -> nx_int {
194 if a.dtype != NX_DT_I64 { return NX_SP_ERR_BAD_DTYPE }
195 if b.dtype != NX_DT_I64 { return NX_SP_ERR_BAD_DTYPE }
196 if c.dtype != NX_DT_I64 { return NX_SP_ERR_BAD_DTYPE }
197 if b.ndim != 2 { return NX_SP_ERR_BAD_NDIM }
198 if c.ndim != 2 { return NX_SP_ERR_BAD_NDIM }
199 if a.n_cols != b.shape[0] { return NX_SP_ERR_SHAPE_MISMATCH }
200 if c.shape[0] != a.n_rows { return NX_SP_ERR_SHAPE_MISMATCH }
201 if c.shape[1] != b.shape[1] { return NX_SP_ERR_SHAPE_MISMATCH }
202
203 nx_t_fill_zero(c)
204 let pb: *i64 = b.storage as *i64
205 let pc: *i64 = c.storage as *i64
206 let N: nx_int = b.shape[1]
207 let K: nx_int = a.n_cols
208
209 var m: nx_int = 0
210 while m < a.n_rows {
211 var k_idx: nx_int = a.row_ptr[m]
212 let stop: nx_int = a.row_ptr[m + 1]
213 while k_idx < stop {
214 let k: nx_int = a.col_idx[k_idx]
215 let a_val: nx_int = a.values[k_idx]
216 // c[m, :] += a_val * b[k, :]
217 var n: nx_int = 0
218 while n < N {
219 pc[m * N + n] = pc[m * N + n] + a_val * pb[k * N + n]
220 n = n + 1
221 }
222 k_idx = k_idx + 1
223 }
224 m = m + 1
225 }
226 return NX_SP_OK
227}
228
229// ===== Sparsity-ratio measurement =================================
230//
231// Q10 fraction of dense entries that are non-zero. For tracking the
232// algorithmic-win actually realised. 0 = all zero; 1024 = all non-zero.
233
234func nx_sp_sparsity_ratio_q10(s: *NxSparseTensor) -> nx_int {
235 let denom: nx_int = s.n_rows * s.n_cols
236 if denom <= 0 { return 0 }
237 return (s.nnz * NX_MAGIC_1024) / denom
238}
239
240// ===== Speedup estimate (vs dense matmul) =========================
241//
242// Q10 multiplier. CSR matmul does NNZ * N work; dense does M * K * N.
243// Speedup = (M * K * N) / (NNZ * N) = (M * K) / NNZ.
244
245func nx_sp_estimated_speedup_q10(s: *NxSparseTensor) -> nx_int {
246 if s.nnz <= 0 { return 0 }
247 return (s.n_rows * s.n_cols * NX_MAGIC_1024) / s.nnz
248}