main.bend
bend-ml-tensor-array@0.1.17.0 · 1,827 lines · 0 definitions
1# bend-ml-tensor-array: tensors over a flat Array<F32>, with the shape in the type.
2#
3# import bend-ml-tensor-array@0.1.17.0/main.bend as TA
4#
5# Mat<r, c> stores r*c numbers in an Array<F32>, row by row (index i*c + j). The
6# dimensions are erased type parameters: the same guarantee as bend-ml-tensor
7# (a product with wrong dimensions does not compile), but ~50x faster than lists:
8# Array.get/set by index, and products in parallel blocks of rows.
9#
10# An Array is affine (a single owner): every operation that reads an operand returns it
11# together with the result (as Array.get does), in records MM, AR, ... Use Mat.clone
12# when you need two copies.
13#
14# What is NOT proved: the invariant "the Array has capacity >= r*c" holds because
15# the constructors (Mat.zeros, Mat.of_list) allocate the right size; and F32 is
16# validated by tests against PyTorch (reference/test_tensor_array.py), not by proof.
17
18import Base
19
20# ---------------------------------------------------------------------
21# helper types
22# ---------------------------------------------------------------------
23
24type R is Type:
25 R{a: Array<F32>, b: Array<F32>, x: F32}
26
27type G is Type:
28 G{a: Array<F32>, b: Array<F32>, c: Array<F32>}
29
30# A dot product with four partial sums. Each step adds its product to a0 and rotates the sums
31# (a1, a2, a3, a0), so each sum gets every fourth product and the processor can work on the next
32# addition before the last one has finished; at the end the sums are added in pairs. The order of
33# the additions is fixed, so the result is the same on every run and every thread count, but it is
34# not the order of a single running sum (NOTES.md, exp. 17).
35def dotr(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, a0: F32, a1: F32, a2: F32, a3: F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> R:
36 match k ra rb:
37 case 0n Tuple{a, x} Tuple{b, y}:
38 R{a, b, ((a0 + a1 : F32) + (a2 + a3 : F32) : F32)}
39 case 1n+p Tuple{a, +x} Tuple{b, +y}:
40 dotr(p, (ia + sa : U32), (ib + sb : U32), sa, sb, a1, a2, a3, (a0 + (x * y : F32) : F32), Array.get(F32, a, (ia + sa : U32)), Array.get(F32, b, (ib + sb : U32)))
41
42def dotg(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, acc: F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> R:
43 dotr(k, ia, ib, sa, sb, acc, 0.0, 0.0, 0.0, ra, rb)
44
45# writes one cell of C: c[idx] = x, or c[idx] = c[idx] + x when ac (an accumulating product, 0.1.10.0). x is computed
46# exactly as without ac and added once, so C + A·B has the bits of adding a separately computed A·B to C.
47
48def cwr_add(+idx: U32, +x: F32, r: Array<F32> & F32) -> Array<F32>:
49 match r:
50 case Tuple{c, +v}:
51 Array.set(F32, c, idx, (v + x : F32))
52
53def cwr(ac: Bool, +idx: U32, +x: F32, c: Array<F32>) -> Array<F32>:
54 match ac:
55 case True{}:
56 cwr_add(idx, x, Array.get(F32, c, idx))
57 case False{}:
58 Array.set(F32, c, idx, x)
59
60# row i of C, columns j.. ; r is the product of the current column (already computed)
61
62def gcols(+ac: Bool, mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, +cidx: U32, c: Array<F32>, r: R) -> G:
63 match mleft r:
64 case 0n R{a, b, +x}:
65 G{a, b, cwr(ac, cidx, x, c)}
66 case 1n+p R{a, b, +x}:
67 gcols(ac, p, kn, ia0, (ib0 + bcol : U32), sa, sb, bcol, (cidx + 1 : U32), cwr(ac, cidx, x, c), dotg(kn, ia0, (ib0 + bcol : U32), sa, sb, 0.0, Array.get(F32, a, ia0), Array.get(F32, b, (ib0 + bcol : U32))))
68
69def grows(+ac: Bool, nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +mu: U32, g: G) -> G:
70 match nleft g:
71 case 0n G{a, b, c}:
72 G{a, b, c}
73 case 1n+p G{a, b, c}:
74 grows(ac, p, mm1, kn, (i + 1 : U32), arow, sa, sb, bcol, mu, gcols(ac, mm1, kn, (i * arow : U32), 0, sa, sb, bcol, (i * mu : U32), c, dotg(kn, (i * arow : U32), 0, sa, sb, 0.0, Array.get(F32, a, (i * arow : U32)), Array.get(F32, b, 0))))
75
76# C (n x m) = A · B with k steps, or C + A · B when ac; c is the output array (with capacity n*m)
77
78def gemm_ac(+ac: Bool, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
79 grows(ac, n, Nat.sub(m, 1n), k, 0, arow, sa, sb, bcol, U32.from_nat(m), G{a, b, c})
80
81def gemm(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
82 gemm_ac(False{}, n, m, k, arow, sa, sb, bcol, a, b, c)
83
84
85# ---------------------------------------------------------------------
86# parallel gemm: 2^d blocks of rows of C, each with its own copy of A and B,
87# returning lists that are written into C at the end
88# ---------------------------------------------------------------------
89
90type S3 is Type:
91 S3{a: Array<F32>, b: Array<F32>, out: List<&2, F32>}
92
93def lcols(mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, out: List<&2, F32>, r: R) -> S3:
94 match mleft r:
95 case 0n R{a, b, +x}:
96 S3{a, b, x <> out}
97 case 1n+p R{a, b, +x}:
98 lcols(p, kn, ia0, (ib0 + bcol : U32), sa, sb, bcol, x <> out, dotg(kn, ia0, (ib0 + bcol : U32), sa, sb, 0.0, Array.get(F32, a, ia0), Array.get(F32, b, (ib0 + bcol : U32))))
99
100def lrows(nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, st: S3) -> S3:
101 match nleft st:
102 case 0n S3{a, b, out}:
103 S3{a, b, out}
104 case 1n+p S3{a, b, out}:
105 lrows(p, mm1, kn, (i + 1 : U32), arow, sa, sb, bcol, lcols(mm1, kn, (i * arow : U32), 0, sa, sb, bcol, out, dotg(kn, (i * arow : U32), 0, sa, sb, 0.0, Array.get(F32, a, (i * arow : U32)), Array.get(F32, b, 0))))
106
107def lblock_fin(st: S3) -> List<&2, F32>:
108 match st:
109 case S3{a, b, out}:
110 List.reverse(&2, F32, out)
111
112def lblock(cnt: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>) -> List<&2, F32>:
113 lblock_fin(lrows(cnt, mm1, kn, i, arow, sa, sb, bcol, S3{a, b, Nil{}}))
114
115def lpar(d: Nat, +cnt: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, ra: Array<F32> & Array<F32>, rb: Array<F32> & Array<F32>) -> List<&2, F32>:
116 match d ra rb:
117 case 0n Tuple{a, a2} Tuple{b, b2}:
118 lblock(cnt, mm1, kn, i, arow, sa, sb, bcol, a, b)
119 case 1n++q Tuple{a, a2} Tuple{b, b2}:
120 x y = lpar(q, Nat.div(cnt, 2n), mm1, kn, i, arow, sa, sb, bcol, Array.clone(F32, a), Array.clone(F32, b)) lpar(q, Nat.sub(cnt, Nat.div(cnt, 2n)), mm1, kn, (i + U32.from_nat(Nat.div(cnt, 2n)) : U32), arow, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2))
121 List.append(&2, F32, x, y)
122
123def write_l(xs: List<&2, F32>, +idx: U32, c: Array<F32>) -> Array<F32>:
124 match xs:
125 case Nil{}:
126 c
127 case Con{h, t}:
128 write_l(t, (idx + 1 : U32), Array.set(F32, c, idx, h))
129
130# c[idx..] = list, or c[idx..] += list when ac
131
132def write_la(+ac: Bool, xs: List<&2, F32>, +idx: U32, c: Array<F32>) -> Array<F32>:
133 match xs:
134 case Nil{}:
135 c
136 case Con{+h, t}:
137 write_la(ac, t, (idx + 1 : U32), cwr(ac, idx, h, c))
138
139def gp3(+ac: Bool, a: Array<F32>, b: Array<F32>, c: Array<F32>, xs: List<&2, F32>) -> G:
140 G{a, b, write_la(ac, xs, 0, c)}
141
142def gp2(d: Nat, +n: Nat, +mm1: Nat, +kn: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, c: Array<F32>, pa: Array<F32> & Array<F32>, pb: Array<F32> & Array<F32>) -> G:
143 match pa pb:
144 case Tuple{a1, a2} Tuple{b1, b2}:
145 gp3(False{}, a1, b1, c, lpar(d, n, mm1, kn, 0, arow, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2)))
146
147# C (n x m) = A · B, with 2^d blocks of rows in parallel; returns A, B (intact) and C
148
149def gemm_par(d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
150 gp2(d, n, Nat.sub(m, 1n), k, arow, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b))
151
152
153# ---------------------------------------------------------------------
154# parallel gemm by COLUMNS, for n = 1 (matrix · vector): each block computes a
155# range of columns of C, with its own copy of the operands
156# ---------------------------------------------------------------------
157
158def cblock_fin(st: S3) -> List<&2, F32>:
159 match st:
160 case S3{a, b, out}:
161 List.reverse(&2, F32, out)
162
163# cnt columns starting at j0 (cnt >= 1); row 0 of A is at ia0
164
165def cblock_go(cnt1: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>) -> List<&2, F32>:
166 cblock_fin(lcols(cnt1, kn, ia0, (j0 * bcol : U32), sa, sb, bcol, Nil{}, dotg(kn, ia0, (j0 * bcol : U32), sa, sb, 0.0, Array.get(F32, a, ia0), Array.get(F32, b, (j0 * bcol : U32)))))
167
168def cblock(cnt: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>) -> List<&2, F32>:
169 match cnt:
170 case 0n:
171 Nil{}
172 case 1n+p:
173 cblock_go(p, kn, ia0, j0, sa, sb, bcol, a, b)
174
175def cpar(d: Nat, +cnt: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, ra: Array<F32> & Array<F32>, rb: Array<F32> & Array<F32>) -> List<&2, F32>:
176 match d ra rb:
177 case 0n Tuple{a, a2} Tuple{b, b2}:
178 cblock(cnt, kn, ia0, j0, sa, sb, bcol, a, b)
179 case 1n++q Tuple{a, a2} Tuple{b, b2}:
180 x y = cpar(q, Nat.div(cnt, 2n), kn, ia0, j0, sa, sb, bcol, Array.clone(F32, a), Array.clone(F32, b)) cpar(q, Nat.sub(cnt, Nat.div(cnt, 2n)), kn, ia0, (j0 + U32.from_nat(Nat.div(cnt, 2n)) : U32), sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2))
181 List.append(&2, F32, x, y)
182
183def cp2(+ac: Bool, d: Nat, +m: Nat, +kn: Nat, +sa: U32, +sb: U32, +bcol: U32, c: Array<F32>, pa: Array<F32> & Array<F32>, pb: Array<F32> & Array<F32>) -> G:
184 match pa pb:
185 case Tuple{a1, a2} Tuple{b1, b2}:
186 gp3(ac, a1, b1, c, cpar(d, m, kn, 0, 0, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2)))
187
188# C (1 x m) = A (1 x k) · B, with 2^d blocks of columns in parallel. Since 0.1.12.0 no product calls it (one row runs
189# the scalar loop, see gemm_any3): it clones both operands at every node. Kept for callers of earlier versions.
190
191def gemm_par_cols(+ac: Bool, d: Nat, +m: Nat, +k: Nat, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
192 cp2(ac, d, m, k, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b))
193
194# ---------------------------------------------------------------------
195# element-wise operations (in place, by index)
196# ---------------------------------------------------------------------
197
198# ---------------------------------------------------------------------
199# gemm: C[i][j] = sum_p A[i*arow + p*sa] * B[p*sb + j*bcol]
200# (nn: arow=K, sa=1, sb=M, bcol=1 | nt: B stored m x k: sb=1, bcol=K | tn: A stored k x n: arow=1, sa=N)
201# ---------------------------------------------------------------------
202
203# acc += A[ia + t*sa] * B[ib + t*sb] for k steps; ra, rb already carry the first pair
204
205type P2 is Type:
206 P2{a: Array<F32>, b: Array<F32>}
207
208type AL is Type:
209 AL{c: Array<F32>, l: List<&2, F32>}
210
211# list of m elements starting at idx (reads with Array.get, returning the array)
212
213def read_l(mleft: Nat, +idx: U32, acc: List<&2, F32>, rc: Array<F32> & F32) -> AL:
214 match mleft rc:
215 case 0n Tuple{c, +v}:
216 AL{c, List.reverse(&2, F32, v <> acc)}
217 case 1n+p Tuple{c, +v}:
218 read_l(p, (idx + 1 : U32), v <> acc, Array.get(F32, c, (idx + 1 : U32)))
219
220# writes the list starting at idx
221
222def addl(xs: List<&2, F32>, +idx: U32, rc: Array<F32> & F32) -> Array<F32>:
223 match xs rc:
224 case Nil{} Tuple{c, v}:
225 c
226 case Con{+h, t} Tuple{c, +v}:
227 addl(t, (idx + 1 : U32), Array.get(F32, Array.set(F32, c, idx, (v + h : F32)), (idx + 1 : U32)))
228
229# adds a bias (a list of m) to each of the n rows of c (n x m)
230
231def bias_rows(nleft: Nat, +i: U32, +mu: U32, +bl: List<&2, F32>, c: Array<F32>) -> Array<F32>:
232 match nleft:
233 case 0n:
234 c
235 case 1n+p:
236 bias_rows(p, (i + 1 : U32), mu, bl, addl(bl, (i * mu : U32), Array.get(F32, c, (i * mu : U32))))
237
238# in-place relu on the first n elements
239
240def relu_l(nleft: Nat, +t: U32, rc: Array<F32> & F32) -> Array<F32>:
241 match nleft rc:
242 case 0n Tuple{c, v}:
243 c
244 case 1n+p Tuple{c, +v}:
245 relu_l(p, (t + 1 : U32), Array.get(F32, Array.set(F32, c, t, F32.max(v, 0.0)), (t + 1 : U32)))
246
247def pos2(b: Bool) -> F32:
248 match b:
249 case True{}:
250 1.0
251 case False{}:
252 0.0
253
254def pos(+y: F32) -> F32:
255 pos2(F32.is_gt(y, 0.0))
256
257# d[t] *= (h[t] > 0): the relu gradient; returns both arrays
258
259def mask_l(nleft: Nat, +t: U32, rd: Array<F32> & F32, rh: Array<F32> & F32) -> P2:
260 match nleft rd rh:
261 case 0n Tuple{d, x} Tuple{h, y}:
262 P2{d, h}
263 case 1n+p Tuple{d, +x} Tuple{h, +y}:
264 mask_l(p, (t + 1 : U32), Array.get(F32, Array.set(F32, d, t, (x * pos(y) : F32)), (t + 1 : U32)), Array.get(F32, h, (t + 1 : U32)))
265
266# w[t] -= lr * dw[t]
267
268def maxl2(xs: List<&2, F32>, cur: F32) -> F32:
269 match xs:
270 case Nil{}:
271 cur
272 case Con{+h, t}:
273 maxl2(t, F32.max(cur, h))
274
275def maxl(xs: List<&2, F32>) -> F32:
276 match xs:
277 case Nil{}:
278 0.0
279 case Con{+h, t}:
280 maxl2(t, h)
281
282def expl(xs: List<&2, F32>, +m: F32) -> List<&2, F32>:
283 match xs:
284 case Nil{}:
285 Nil{}
286 case Con{+h, t}:
287 F32.exp((h - m : F32)) <> expl(t, m)
288
289def suml(xs: List<&2, F32>, acc: F32) -> F32:
290 match xs:
291 case Nil{}:
292 acc
293 case Con{+h, t}:
294 suml(t, (acc + h : F32))
295
296def divl(xs: List<&2, F32>, +s: F32) -> List<&2, F32>:
297 match xs:
298 case Nil{}:
299 Nil{}
300 case Con{+h, t}:
301 (h / s : F32) <> divl(t, s)
302
303def softmax2(+es: List<&2, F32>) -> List<&2, F32>:
304 divl(es, suml(es, 0.0))
305
306def softmax(+xs: List<&2, F32>) -> List<&2, F32>:
307 softmax2(expl(xs, maxl(xs)))
308
309# p[label]
310
311def sgd_l(nleft: Nat, +t: U32, +lr: F32, rw: Array<F32> & F32, rd: Array<F32> & F32) -> P2:
312 match nleft rw rd:
313 case 0n Tuple{w, x} Tuple{d, y}:
314 P2{w, d}
315 case 1n+p Tuple{w, +x} Tuple{d, +y}:
316 sgd_l(p, (t + 1 : U32), lr, Array.get(F32, Array.set(F32, w, t, (x - (lr * y : F32) : F32)), (t + 1 : U32)), Array.get(F32, d, (t + 1 : U32)))
317
318# ---------------------------------------------------------------------
319# softmax + cross-entropy per row (lists of 10 elements: negligible cost)
320# ---------------------------------------------------------------------
321
322def pick(ps: List<&2, F32>, n: Nat) -> F32:
323 match ps n:
324 case Con{+h, t} 0n:
325 h
326 case Con{h, t} 1n+q:
327 pick(t, q)
328 case Nil{} _:
329 0.0
330
331# (p - onehot(label)) / n
332
333def hit(b: Bool) -> F32:
334 match b:
335 case True{}:
336 1.0
337 case False{}:
338 0.0
339
340def grad_l(ps: List<&2, F32>, +label: Nat, +j: Nat, +nf: F32) -> List<&2, F32>:
341 match ps:
342 case Nil{}:
343 Nil{}
344 case Con{+h, t}:
345 (((h - hit(Nat.is_eq(label, j)) : F32) / nf) : F32) <> grad_l(t, label, 1n+j, nf)
346
347type L2 is Type:
348 L2{loss: F32, c: Array<F32>}
349
350def ce_row2(+ps: List<&2, F32>, +label: Nat, +nf: F32, +idx: U32, loss: F32, c: Array<F32>) -> L2:
351 L2{(loss - F32.log(F32.max(pick(ps, label), 0.000000000001)) : F32), write_l(grad_l(ps, label, 0n, nf), idx, c)}
352
353def ce_row(al: AL, +label: Nat, +nf: F32, +idx: U32, loss: F32) -> L2:
354 match al:
355 case AL{c, l}:
356 ce_row2(softmax(l), label, nf, idx, loss, c)
357
358# walks the rows of z2 (n x m): r carries the accumulated loss and the array; at the end the array has become the gradient
359
360type Mat<-r: Nat, -c: Nat> is Type:
361 Mat{d: Array<F32>}
362
363def ce_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +nf: F32, +i: U32, r: L2) -> L2:
364 match labels r:
365 case Nil{} L2{loss, c}:
366 L2{loss, c}
367 case Con{+y, t} L2{loss, c}:
368 ce_rows(t, mm1, mu, nf, (i + 1 : U32), ce_row(read_l(mm1, (i * mu : U32), Nil{}, Array.get(F32, c, (i * mu : U32))), y, nf, (i * mu : U32), loss))
369
370# ---- the same rows over the array itself: no lists ----
371# The list version above reads each row into a list and walks it five times. This one walks the row in place with tail
372# loops, with the same operations in the same order (the max from the first cell, exp(x - max) summed from +0.0 left to
373# right, p = e / s, the loss term read before the row is overwritten, (p - onehot) / n written back), so the same bits;
374# MetalBend measured 3.4x fewer cycles on 64 rows of 49,152 (2026-10-25). The list version stays as the reference
375# tests/ce_array.bend compares it with.
376
377type AF is Type:
378 AF{c: Array<F32>, x: F32}
379
380# the max of k + 1 cells from idx, from m (the first cell's value, already read in r)
381def cea_max(k: Nat, +idx: U32, +m: F32, r: Array<F32> & F32) -> AF:
382 match k r:
383 case 0n Tuple{c, +x}:
384 AF{c, F32.max(m, x)}
385 case 1n+p Tuple{c, +x}:
386 cea_max(p, (idx + 1 : U32), F32.max(m, x), Array.get(F32, c, (idx + 1 : U32)))
387
388# cells idx..idx + k become exp(cell - m); returns their sum from s, added left to right
389def cea_exp(k: Nat, +idx: U32, +m: F32, +s: F32, r: Array<F32> & F32) -> AF:
390 match k r:
391 case 0n Tuple{c, +x}:
392 +e = F32.exp((x - m : F32))
393 AF{Array.set(F32, c, idx, e), (s + e : F32)}
394 case 1n+p Tuple{c, +x}:
395 +e = F32.exp((x - m : F32))
396 cea_exp(p, (idx + 1 : U32), m, (s + e : F32), Array.get(F32, Array.set(F32, c, idx, e), (idx + 1 : U32)))
397
398# cell idx (column j) becomes ((e / s) - onehot) / nf, for k + 1 cells
399def cea_grad(k: Nat, +idx: U32, +j: Nat, +label: Nat, +s: F32, +nf: F32, r: Array<F32> & F32) -> Array<F32>:
400 match k r:
401 case 0n Tuple{c, +x}:
402 Array.set(F32, c, idx, ((((x / s : F32) - hit(Nat.is_eq(label, j)) : F32) / nf) : F32))
403 case 1n+p Tuple{c, +x}:
404 cea_grad(p, (idx + 1 : U32), 1n+j, label, s, nf, Array.get(F32, Array.set(F32, c, idx, ((((x / s : F32) - hit(Nat.is_eq(label, j)) : F32) / nf) : F32)), (idx + 1 : U32)))
405
406def cea_p2(+s: F32, r: Array<F32> & F32) -> AF:
407 match r:
408 case Tuple{c, +e}:
409 AF{c, (e / s : F32)}
410
411# p[label] = e[label] / s, or 0.0 for a label past the row (pick's Nil case)
412def cea_p(inside: Bool, +o: U32, +label: Nat, +s: F32, c: Array<F32>) -> AF:
413 match inside:
414 case True{}:
415 cea_p2(s, Array.get(F32, c, (o + U32.from_nat(label) : U32)))
416 case False{}:
417 AF{c, 0.0}
418
419def cea_row3(+mm1: Nat, +o: U32, +label: Nat, +s: F32, +nf: F32, +loss: F32, q: AF) -> L2:
420 match q:
421 case AF{c, +p}:
422 L2{(loss - F32.log(F32.max(p, 0.000000000001)) : F32), cea_grad(mm1, o, 0n, label, s, nf, Array.get(F32, c, o))}
423
424def cea_row2(+mm1: Nat, +o: U32, +label: Nat, +nf: F32, +loss: F32, q: AF) -> L2:
425 match q:
426 case AF{c, +s}:
427 cea_row3(mm1, o, label, s, nf, loss, cea_p(Nat.is_le(label, mm1), o, label, s, c))
428
429def cea_row1(+mm1: Nat, +o: U32, +label: Nat, +nf: F32, +loss: F32, q: AF) -> L2:
430 match q:
431 case AF{c, +m}:
432 cea_row2(mm1, o, label, nf, loss, cea_exp(mm1, o, m, 0.0, Array.get(F32, c, o)))
433
434# a row's max starts from its first cell (maxl), then folds over the k + 1 cells from it
435def cea_row0(+mm1: Nat, +o: U32, +label: Nat, +nf: F32, +loss: F32, r: Array<F32> & F32) -> L2:
436 match r:
437 case Tuple{c, +x0}:
438 cea_row1(mm1, o, label, nf, loss, cea_max(mm1, o, x0, Array.get(F32, c, o)))
439
440# one row per label, in order, as ce_rows
441def cea_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +nf: F32, +i: U32, r: L2) -> L2:
442 match labels r:
443 case Nil{} L2{loss, c}:
444 L2{loss, c}
445 case Con{+y, t} L2{+loss, c}:
446 cea_rows(t, mm1, mu, nf, (i + 1 : U32), cea_row0(mm1, (i * mu : U32), y, nf, loss, Array.get(F32, c, (i * mu : U32))))
447
448# ---- helper: hits per row ----
449
450type CR is Type:
451 CR{c: Array<F32>, n: Nat}
452
453def gt_head(xs: List<&2, F32>, +v: F32) -> Bool:
454 match xs:
455 case Nil{}:
456 False{}
457 case Con{+h, t}:
458 F32.is_gt(h, v)
459
460def argmax2(xs: List<&2, F32>, g: Bool, +i: Nat, +bi: Nat, +bv: F32) -> Nat:
461 match xs g:
462 case Nil{} _:
463 bi
464 case Con{+h, +t} True{}:
465 argmax2(t, gt_head(t, h), 1n+i, i, h)
466 case Con{h, +t} False{}:
467 argmax2(t, gt_head(t, bv), 1n+i, bi, bv)
468
469def argmax(xs: List<&2, F32>) -> Nat:
470 match xs:
471 case Nil{}:
472 0n
473 case Con{+h, +t}:
474 argmax2(t, gt_head(t, h), 1n, 0n, h)
475
476def one(b: Bool) -> Nat:
477 match b:
478 case True{}:
479 1n
480 case False{}:
481 0n
482
483def count_rows3(+y: Nat, n: Nat, al: AL) -> CR:
484 match al:
485 case AL{c, l}:
486 CR{c, Nat.add(n, one(Nat.is_eq(argmax(l), y)))}
487
488def count_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +i: U32, r: CR) -> CR:
489 match labels r:
490 case Nil{} CR{c, n}:
491 CR{c, n}
492 case Con{+y, t} CR{c, n}:
493 count_rows(t, mm1, mu, (i + 1 : U32), count_rows3(y, n, read_l(mm1, (i * mu : U32), Nil{}, Array.get(F32, c, (i * mu : U32)))))
494
495
496# =====================================================================
497# Typed API: Mat<r, c> (the dimensions enter the type)
498# =====================================================================
499
500# =====================================================================
501# Array capacity: a PROVED depth for the number of slots
502# =====================================================================
503#
504# cap_depth(n) is the d such that an Array with 2^d slots holds n elements. It is computed by
505# structural recursion on n (not with U32.log2) so that the law cap_ok below can prove
506# `n <= 2^cap_depth(n)`. What stays trusted is only that Array.new(T, d, v) allocates 2^d slots.
507
508# a type chosen by a Bool: rewriting through it refutes True == False
509def BD(b: Bool, t: Type, f: Type) -> Type:
510 match b:
511 case True{}:
512 t
513 case False{}:
514 f
515
516# a <= b, by direct recursion
517def leb(a: Nat, b: Nat) -> Bool:
518 match a b:
519 case 0n _:
520 True{}
521 case 1n+p 0n:
522 False{}
523 case 1n+p 1n+q:
524 leb(p, q)
525
526def pow2(d: Nat) -> Nat:
527 match d:
528 case 0n:
529 1n
530 case 1n++q:
531 Nat.add(pow2(q), pow2(q))
532
533# smallest power of two >= n, by structural recursion on n: returns (depth, 2^depth)
534def cap_pick(ok: Bool, d: Nat, +p: Nat) -> Nat & Nat:
535 match ok:
536 case True{}:
537 (d, p)
538 case False{}:
539 (1n+d, Nat.add(p, p))
540
541def cap_step(r: Nat & Nat, +n: Nat) -> Nat & Nat:
542 match r:
543 case Tuple{+d, +p}:
544 cap_pick(leb(n, p), d, p)
545
546def cap_pair(n: Nat) -> Nat & Nat:
547 match n:
548 case 0n:
549 (0n, 1n)
550 case 1n++m:
551 cap_step(cap_pair(m), 1n+m)
552
553def fst_nat(r: Nat & Nat) -> Nat:
554 match r:
555 case Tuple{d, p}:
556 d
557
558def snd_nat(r: Nat & Nat) -> Nat:
559 match r:
560 case Tuple{d, p}:
561 p
562
563# leb is Base's Nat.is_le, case for case. leb walks its arguments one step at a time in the compiled program; Base's
564# comparison is one native comparison there (0.1.13.0).
565def le_leb(a: Nat, b: Nat) -> {Nat.is_le(a, b) == leb(a, b) : Bool}:
566 match a b:
567 case 0n 0n:
568 {==}
569 case 0n 1n+q:
570 {==}
571 case 1n+p 0n:
572 {==}
573 case 1n+p 1n+q:
574 le_leb(p, q)
575
576def leb_of(+a: Nat, +b: Nat, e: {True{} == Nat.is_le(a, b) : Bool}) -> {True{} == leb(a, b) : Bool}:
577 %le_leb(a, b) : {True{} == _ : Bool}
578 e
579
580# the same pair by doubling (0.1.10.0): p doubles until n <= p, log2(n) steps in a tail loop, where cap_pair takes n
581# steps of a recursion that is not a tail call, each with a comparison of up to n steps. Its fuel is n, which always
582# suffices (2^n >= n); if it ran out, the answer would be cap_pair's, so the law below needs no arithmetic about how
583# many doublings there are. ok is Nat.is_le(n, p), computed by the caller (match only inspects a parameter); with leb
584# it was n steps in all, 152 ms for n = 786,432 on one GPU lane (0.1.13.0).
585def cap_dbl(f: Nat, ok: Bool, +n: Nat, +d: Nat, +p: Nat) -> Nat & Nat:
586 match f ok:
587 case _ True{}:
588 (d, p)
589 case 0n False{}:
590 cap_pair(n)
591 case 1n+g False{}:
592 cap_dbl(g, Nat.is_le(n, Nat.add(p, p)), n, 1n+d, Nat.add(p, p))
593
594# the depth d such that an Array with 2^d slots holds n elements (the smallest such d)
595def cap_depth(+n: Nat) -> Nat:
596 fst_nat(cap_dbl(n, Nat.is_le(n, 1n), n, 0n, 1n))
597
598# ---- lemmas ----
599
600# a <= b -> a <= 1 + b
601def leb_succ(a: Nat, b: Nat, e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, 1n+b) : Bool}:
602 match a b:
603 case 0n _:
604 {==}
605 case 1n+p 0n:
606 %e : BD(_, Unit, {True{} == leb(1n+p, 1n) : Bool})
607 Unit{}
608 case 1n+p 1n+q:
609 leb_succ(p, q, e)
610
611# a <= b -> a <= c + b
612def leb_up_l(+a: Nat, +b: Nat, c: Nat, +e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, Nat.add(c, b)) : Bool}:
613 match c:
614 case 0n:
615 e
616 case 1n++c2:
617 leb_succ(a, Nat.add(c2, b), leb_up_l(a, b, c2, e))
618
619# a <= b -> a <= b + c
620def leb_up(a: Nat, b: Nat, c: Nat, e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, Nat.add(b, c)) : Bool}:
621 match a b:
622 case 0n _:
623 {==}
624 case 1n+p 0n:
625 %e : BD(_, Unit, {True{} == leb(1n+p, Nat.add(0n, c)) : Bool})
626 Unit{}
627 case 1n+p 1n+q:
628 leb_up(p, q, c, e)
629
630# a <= b and c <= d -> a + c <= b + d
631def leb_add(a: Nat, b: Nat, c: Nat, d: Nat, e1: {True{} == leb(a, b) : Bool}, e2: {True{} == leb(c, d) : Bool}) -> {True{} == leb(Nat.add(a, c), Nat.add(b, d)) : Bool}:
632 match a b:
633 case 0n _:
634 leb_up_l(c, d, b, e2)
635 case 1n+p 0n:
636 %e1 : BD(_, Unit, {True{} == leb(Nat.add(1n+p, c), Nat.add(0n, d)) : Bool})
637 Unit{}
638 case 1n+p 1n+q:
639 leb_add(p, q, c, d, e1, e2)
640
641# 1 <= 2^d
642def pow2_pos(d: Nat) -> {True{} == leb(1n, pow2(d)) : Bool}:
643 match d:
644 case 0n:
645 {==}
646 case 1n++q:
647 leb_up(1n, pow2(q), pow2(q), pow2_pos(q))
648
649# 1 <= p when p = 2^d
650def pos_p(d: Nat, p: Nat, i1: {p == pow2(d) : Nat}) -> {True{} == leb(1n, p) : Bool}:
651 %Equal.sym(Nat, p, pow2(d), i1) : {True{} == leb(1n, _) : Bool}
652 pow2_pos(d)
653
654# p + p == 2^(1+d) when p == 2^d
655def step_eq(d: Nat, p: Nat, i1: {p == pow2(d) : Nat}) -> {Nat.add(p, p) == pow2(1n+d) : Nat}:
656 %i1 : {Nat.add(p, p) == Nat.add(_, _) : Nat}
657 {==}
658
659# one step: given the invariant for m, the invariant for 1+m holds whichever way the comparison goes
660def step_ok(ok: Bool, +d: Nat, +p: Nat, +m: Nat, +e: {ok == leb(1n+m, p) : Bool}, +i1: {p == pow2(d) : Nat}, +i2: {True{} == leb(m, p) : Bool}) -> {snd_nat(cap_pick(ok, d, p)) == pow2(fst_nat(cap_pick(ok, d, p))) : Nat} & {True{} == leb(1n+m, snd_nat(cap_pick(ok, d, p))) : Bool}:
661 match ok:
662 case True{}:
663 (i1, e)
664 case False{}:
665 (step_eq(d, p, i1), leb_add(1n, p, m, p, pos_p(d, p, i1), i2))
666
667
668# the invariant for cap_step on a pair, opened
669def step_pair_ok(+m: Nat, r: Nat & Nat, i1: {snd_nat(r) == pow2(fst_nat(r)) : Nat}, i2: {True{} == leb(m, snd_nat(r)) : Bool}) -> {snd_nat(cap_step(r, 1n+m)) == pow2(fst_nat(cap_step(r, 1n+m))) : Nat} & {True{} == leb(1n+m, snd_nat(cap_step(r, 1n+m))) : Bool}:
670 match r:
671 case Tuple{+d, +p}:
672 step_ok(leb(1n+m, p), d, p, m, {==}, i1, i2)
673
674def cap_succ(+m: Nat, ih: {snd_nat(cap_pair(m)) == pow2(fst_nat(cap_pair(m))) : Nat} & {True{} == leb(m, snd_nat(cap_pair(m))) : Bool}) -> {snd_nat(cap_pair(1n+m)) == pow2(fst_nat(cap_pair(1n+m))) : Nat} & {True{} == leb(1n+m, snd_nat(cap_pair(1n+m))) : Bool}:
675 match ih:
676 case Tuple{i1, i2}:
677 step_pair_ok(m, cap_pair(m), i1, i2)
678
679# the invariant of cap_pair: the second component is 2^(first) and is >= n
680def cap_inv(n: Nat) -> {snd_nat(cap_pair(n)) == pow2(fst_nat(cap_pair(n))) : Nat} & {True{} == leb(n, snd_nat(cap_pair(n))) : Bool}:
681 match n:
682 case 0n:
683 ({==}, {==})
684 case 1n++m:
685 cap_succ(m, cap_inv(m))
686
687# the invariant of cap_dbl: whichever way it ends (n <= p found, or the fuel out and cap_pair's answer), the second
688# component is 2^(first) and is >= n
689def dbl_inv(f: Nat, ok: Bool, +n: Nat, +d: Nat, +p: Nat, e: {ok == Nat.is_le(n, p) : Bool}, i1: {p == pow2(d) : Nat}) -> {snd_nat(cap_dbl(f, ok, n, d, p)) == pow2(fst_nat(cap_dbl(f, ok, n, d, p))) : Nat} & {True{} == leb(n, snd_nat(cap_dbl(f, ok, n, d, p))) : Bool}:
690 match f ok:
691 case _ True{}:
692 (i1, leb_of(n, p, e))
693 case 0n False{}:
694 cap_inv(n)
695 case 1n+g False{}:
696 dbl_inv(g, Nat.is_le(n, Nat.add(p, p)), n, 1n+d, Nat.add(p, p), {==}, step_eq(d, p, i1))
697
698def dbl_final(n: Nat, pr: {snd_nat(cap_dbl(n, Nat.is_le(n, 1n), n, 0n, 1n)) == pow2(fst_nat(cap_dbl(n, Nat.is_le(n, 1n), n, 0n, 1n))) : Nat} & {True{} == leb(n, snd_nat(cap_dbl(n, Nat.is_le(n, 1n), n, 0n, 1n))) : Bool}) -> {True{} == leb(n, pow2(cap_depth(n))) : Bool}:
699 match pr:
700 case Tuple{i1, i2}:
701 %i1 : {True{} == leb(n, _) : Bool}
702 i2
703
704# LAW (cap_ok): an Array with 2^cap_depth(n) slots always has room for n elements.
705law cap_ok:
706 for +n: Nat
707 {True{} == leb(n, pow2(cap_depth(n))) : Bool}
708
709def cap_ok(n):
710 dbl_final(n, dbl_inv(n, Nat.is_le(n, 1n), n, 0n, 1n, {==}, {==}))
711
712# r x c of zeros, with the exact capacity (smallest power of 2)
713
714def Mat.zeros(+r: Nat, +c: Nat) -> Mat<r, c>:
715 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0)}
716
717def Mat.fill(+r: Nat, +c: Nat, +v: F32) -> Mat<r, c>:
718 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), v)}
719
720def fill_list(xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
721 match xs:
722 case Nil{}:
723 a
724 case Con{h, t}:
725 fill_list(t, (i + 1 : U32), Array.set(F32, a, i, h))
726
727def len_is(xs: List<&2, F32>, n: Nat) -> Bool:
728 match xs n:
729 case Nil{} 0n:
730 True{}
731 case Con{h, t} 1n+p:
732 len_is(t, p)
733 case _ _:
734 False{}
735
736def Mat.of_list2(-r: Nat, -c: Nat, +rr: Nat, +cc: Nat, ok: Bool, xs: List<&2, F32>) -> Maybe<&1, Mat<r, c>>:
737 match ok:
738 case True{}:
739 Some{Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(rr, cc)), 0.0))}}
740 case False{}:
741 None{}
742
743# Mat<r, c> from r*c numbers in row order; None if the size does not match
744
745def Mat.of_list(+r: Nat, +c: Nat, +xs: List<&2, F32>) -> Maybe<&1, Mat<r, c>>:
746 Mat.of_list2(r, c, r, c, len_is(xs, Nat.mul(r, c)), xs)
747
748# reads the first n numbers, returning the matrix
749
750# Mat<r, c> without checking the list size (a short list leaves the rest at zero, a long one is cut
751# by the capacity): only for when the size is known by construction. Prefer Mat.of_list.
752
753def Mat.from_list(+r: Nat, +c: Nat, xs: List<&2, F32>) -> Mat<r, c>:
754 Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0))}
755
756# writes the numbers of xs into m starting at flat index i (row-major); numbers past the
757# capacity are dropped. For loading a large matrix block by block without building one big list.
758def Mat.fill_at(-r: Nat, -c: Nat, +i: U32, xs: List<&2, F32>, m: Mat<r, c>) -> Mat<r, c>:
759 match m:
760 case Mat{d}:
761 Mat{fill_list(xs, i, d)}
762
763type ML<-r: Nat, -c: Nat> is Type:
764 ML{m: Mat<r, c>, l: List<&2, F32>}
765
766def Mat.to_list2(-r: Nat, -c: Nat, al: AL) -> ML<r, c>:
767 match al:
768 case AL{d, l}:
769 ML{Mat{d}, l}
770
771# two independent copies
772
773def Mat.to_list(+r: Nat, +c: Nat, m: Mat<r, c>) -> ML<r, c>:
774 match m:
775 case Mat{d}:
776 Mat.to_list2(r, c, read_l(Nat.sub(Nat.mul(r, c), 1n), 0, Nil{}, Array.get(F32, d, 0)))
777
778type MM2<-r: Nat, -c: Nat> is Type:
779 MM2{a: Mat<r, c>, b: Mat<r, c>}
780
781def Mat.clone2(-r: Nat, -c: Nat, p: Array<F32> & Array<F32>) -> MM2<r, c>:
782 match p:
783 case Tuple{x, y}:
784 MM2{Mat{x}, Mat{y}}
785
786# ---- products: C (n x m) = A · B, with 2^par blocks of rows in parallel ----
787
788# C = A(n x k) · B(k x m)
789
790def Mat.clone(-r: Nat, -c: Nat, m: Mat<r, c>) -> MM2<r, c>:
791 match m:
792 case Mat{d}:
793 Mat.clone2(r, c, Array.clone(F32, d))
794
795type MMul<-n: Nat, -k: Nat, -m: Nat> is Type:
796 MMul{a: Mat<n, k>, b: Mat<k, m>, c: Mat<n, m>}
797
798# C = A(n x k) · Bᵀ, with B stored m x k
799
800type MMulNT<-n: Nat, -k: Nat, -m: Nat> is Type:
801 MMulNT{a: Mat<n, k>, b: Mat<m, k>, c: Mat<n, m>}
802
803# C = Aᵀ · B, with A stored k x n
804
805type MMulTN<-n: Nat, -k: Nat, -m: Nat> is Type:
806 MMulTN{a: Mat<k, n>, b: Mat<k, m>, c: Mat<n, m>}
807
808# ---------------------------------------------------------------------
809# gemm in tiles of 8 (0.1.8.0): for products with 2 or more rows, A's rows and B's columns are packed once into cells of
810# 8 F32 (a row of k numbers is ceil(k/8) cells, zeros after the last), so the inner loop reads 8 numbers per array
811# access instead of one. The order of the additions is that of `dotr`: product p of a dot goes to partial sum p mod 4,
812# in order, and the four sums are combined as dotr combines them after k steps (its rotation depends on k mod 4). The
813# padding adds +0 to a sum that starts at +0, which changes nothing, so the results are the same bits as the scalar
814# loop's. Products with one row (matrix · vector) keep the scalar loop: packing B would cost as much as the product.
815# ---------------------------------------------------------------------
816
817type T8 is Data:
818 T8{v0: F32, v1: F32, v2: F32, v3: F32, v4: F32, v5: F32, v6: F32, v7: F32}
819
820def t8_zero() -> T8:
821 T8{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}
822
823def pick0(b: Bool, +v: F32) -> F32:
824 match b:
825 case True{}:
826 v
827 case False{}:
828 0.0
829
830type PK is Type:
831 PK{a: Array<F32>, t: Array<T8>}
832
833# one cell: the 8 numbers of a at idx, idx + st, ..., those at positions q0 + q >= lim (past the end of the line) as 0
834def pk_g7(+v0: F32, +v1: F32, +v2: F32, +v3: F32, +v4: F32, +v5: F32, +v6: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
835 match r:
836 case Tuple{a, +v}:
837 PK{a, Array.set(T8, t, o, T8{v0, v1, v2, v3, v4, v5, v6, pick0(U32.is_lt((q0 + 7 : U32), lim), v)})}
838
839def pk_g6(+v0: F32, +v1: F32, +v2: F32, +v3: F32, +v4: F32, +v5: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
840 match r:
841 case Tuple{a, +v}:
842 pk_g7(v0, v1, v2, v3, v4, v5, pick0(U32.is_lt((q0 + 6 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 7 : U32) : U32)))
843
844def pk_g5(+v0: F32, +v1: F32, +v2: F32, +v3: F32, +v4: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
845 match r:
846 case Tuple{a, +v}:
847 pk_g6(v0, v1, v2, v3, v4, pick0(U32.is_lt((q0 + 5 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 6 : U32) : U32)))
848
849def pk_g4(+v0: F32, +v1: F32, +v2: F32, +v3: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
850 match r:
851 case Tuple{a, +v}:
852 pk_g5(v0, v1, v2, v3, pick0(U32.is_lt((q0 + 4 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 5 : U32) : U32)))
853
854def pk_g3(+v0: F32, +v1: F32, +v2: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
855 match r:
856 case Tuple{a, +v}:
857 pk_g4(v0, v1, v2, pick0(U32.is_lt((q0 + 3 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 4 : U32) : U32)))
858
859def pk_g2(+v0: F32, +v1: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
860 match r:
861 case Tuple{a, +v}:
862 pk_g3(v0, v1, pick0(U32.is_lt((q0 + 2 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 3 : U32) : U32)))
863
864def pk_g1(+v0: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
865 match r:
866 case Tuple{a, +v}:
867 pk_g2(v0, pick0(U32.is_lt((q0 + 1 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 2 : U32) : U32)))
868
869def pk_g0(+q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
870 match r:
871 case Tuple{a, +v}:
872 pk_g1(pick0(U32.is_lt((q0 + 0 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 1 : U32) : U32)))
873
874# the cells of one line, cell o.. of t; q0 counts the line's numbers before the cell
875def pk_line(tl: Nat, +q0: U32, +lim: U32, +base: U32, +st: U32, +o: U32, r: PK) -> PK:
876 match tl r:
877 case 0n r:
878 r
879 case 1n+q PK{a, t}:
880 pk_line(q, (q0 + 8 : U32), lim, (base + (st * 8 : U32) : U32), st, (o + 1 : U32), pk_g0(q0, lim, base, st, o, t, Array.get(F32, a, base)))
881
882# lines count lines; source line i starts at i * lstep (in a), and goes to cell row o (cells o * tku .. in t)
883def pk_lines(cnt: Nat, +tk: Nat, +tku: U32, +ku: U32, +i: U32, +o: U32, +lstep: U32, +st: U32, r: PK) -> PK:
884 match cnt:
885 case 0n:
886 r
887 case 1n+p:
888 pk_lines(p, tk, tku, ku, (i + 1 : U32), (o + 1 : U32), lstep, st, pk_line(tk, 0, ku, (i * lstep : U32), st, (o * tku : U32), r))
889
890def pk_pack(+cnt: Nat, +tk: Nat, +k: Nat, +lstep: U32, +st: U32, a: Array<F32>) -> PK:
891 pk_lines(cnt, tk, U32.from_nat(tk), U32.from_nat(k), 0, 0, lstep, st, PK{a, Array.new(T8, cap_depth(Nat.mul(cnt, tk)), t8_zero())})
892
893# the dot of two lines of cells: tk cells from ia (in a) and ib (in b); ra, rb hold the arrays and the cells at ia, ib,
894# already read. Product q of a cell goes to sum q mod 4, the first four before the last four, as dotr adds them.
895def t8_dot(tk: Nat, +ia: U32, +ib: U32, s0: F32, s1: F32, s2: F32, s3: F32, ra: Array<T8> & T8, rb: Array<T8> & T8) -> TD:
896 match tk ra rb:
897 case 0n Tuple{a, x} Tuple{b, y}:
898 TD{a, b, s0, s1, s2, s3}
899 case 1n+p Tuple{a, T8{x0, x1, x2, x3, x4, x5, x6, x7}} Tuple{b, T8{y0, y1, y2, y3, y4, y5, y6, y7}}:
900 t8_dot(p, (ia + 1 : U32), (ib + 1 : U32), ((s0 + (x0 * y0 : F32) : F32) + (x4 * y4 : F32) : F32), ((s1 + (x1 * y1 : F32) : F32) + (x5 * y5 : F32) : F32), ((s2 + (x2 * y2 : F32) : F32) + (x6 * y6 : F32) : F32), ((s3 + (x3 * y3 : F32) : F32) + (x7 * y7 : F32) : F32), Array.get(T8, a, (ia + 1 : U32)), Array.get(T8, b, (ib + 1 : U32)))
901
902type TD is Type:
903 TD{a: Array<T8>, b: Array<T8>, s0: F32, s1: F32, s2: F32, s3: F32}
904
905# dotr's final sum after k products: the sums sit rotated by k mod 4 when its loop ends
906def t8_comb(+rot: Nat, +s0: F32, +s1: F32, +s2: F32, +s3: F32) -> F32:
907 match rot:
908 case 0n:
909 ((s0 + s1 : F32) + (s2 + s3 : F32) : F32)
910 case 1n:
911 ((s1 + s2 : F32) + (s3 + s0 : F32) : F32)
912 case 2n:
913 ((s2 + s3 : F32) + (s0 + s1 : F32) : F32)
914 case _:
915 ((s3 + s0 : F32) + (s1 + s2 : F32) : F32)
916
917type TR is Type:
918 TR{a: Array<T8>, b: Array<T8>, x: F32}
919
920def t8_fin(+rot: Nat, d: TD) -> TR:
921 match d:
922 case TD{a, b, +s0, +s1, +s2, +s3}:
923 TR{a, b, t8_comb(rot, s0, s1, s2, s3)}
924
925def t8_cell(+tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, a: Array<T8>, b: Array<T8>) -> TR:
926 t8_fin(rot, t8_dot(tk, (i * tku : U32), (j * tku : U32), 0.0, 0.0, 0.0, 0.0, Array.get(T8, a, (i * tku : U32)), Array.get(T8, b, (j * tku : U32))))
927
928type TG is Type:
929 TG{a: Array<T8>, b: Array<T8>, c: Array<F32>}
930
931def t8_set(+ac: Bool, +idx: U32, c: Array<F32>, r: TR) -> TG:
932 match r:
933 case TR{a, b, +x}:
934 TG{a, b, cwr(ac, idx, x, c)}
935
936def t8_cols(+ac: Bool, mleft: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, +cidx: U32, g: TG) -> TG:
937 match mleft g:
938 case 0n g:
939 g
940 case 1n+p TG{a, b, c}:
941 t8_cols(ac, p, tk, tku, rot, i, (j + 1 : U32), (cidx + 1 : U32), t8_set(ac, cidx, c, t8_cell(tk, tku, rot, i, j, a, b)))
942
943def t8_rows(+ac: Bool, nleft: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +mu: U32, g: TG) -> TG:
944 match nleft:
945 case 0n:
946 g
947 case 1n+p:
948 t8_rows(ac, p, m, tk, tku, rot, (i + 1 : U32), mu, t8_cols(ac, m, tk, tku, rot, i, 0, (i * mu : U32), g))
949
950# rows i0 .. i0+cnt-1 of C as a list, in order (the parallel version's leaves)
951type TL is Type:
952 TL{a: Array<T8>, b: Array<T8>, out: List<&2, F32>}
953
954# cells j.. of row i, m - j of them; r is cell j, already computed
955def t8_lcols(mleft: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, out: List<&2, F32>, r: TR) -> TL:
956 match mleft r:
957 case 0n TR{a, b, +x}:
958 TL{a, b, x <> out}
959 case 1n+p TR{a, b, +x}:
960 t8_lcols(p, tk, tku, rot, i, (j + 1 : U32), x <> out, t8_cell(tk, tku, rot, i, (j + 1 : U32), a, b))
961
962def t8_lrows(nleft: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, l: TL) -> TL:
963 match nleft l:
964 case 0n l:
965 l
966 case 1n+p TL{a, b, out}:
967 t8_lrows(p, m, tk, tku, rot, (i + 1 : U32), t8_lcols(Nat.sub(m, 1n), tk, tku, rot, i, 0, out, t8_cell(tk, tku, rot, i, 0, a, b)))
968
969def t8_lfin(l: TL) -> List<&2, F32>:
970 match l:
971 case TL{a, b, out}:
972 List.reverse(&2, F32, out)
973
974type PKK is Type:
975 PKK{a: Array<F32>, b: Array<F32>, ta: Array<T8>, tb: Array<T8>}
976
977def t8_pk2(pa: PK, pb: PK) -> PKK:
978 match pa pb:
979 case PK{a, ta} PK{b, tb}:
980 PKK{a, b, ta, tb}
981
982def t8_g(a: Array<F32>, b: Array<F32>, g: TG) -> G:
983 match g:
984 case TG{ta, tb, c}:
985 G{a, b, c}
986
987def t8_out_seq(+ac: Bool, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array<F32>, p: PKK) -> G:
988 match p:
989 case PKK{a, b, ta, tb}:
990 t8_g(a, b, t8_rows(ac, n, m, tk, tku, rot, 0, U32.from_nat(m), TG{ta, tb, c}))
991
992# ---- the parallel path (0.1.9.0): the larger operand is packed into one piece per block, so a block owns its piece;
993# the other operand is copied once per node of the fork tree (2^d - 1 copies), instead of both at every node ----
994
995type PT is Type:
996 PLeaf{t: Array<T8>}
997 PNode{l: PT, r: PT}
998
999type PKT is Type:
1000 PKT{a: Array<F32>, t: PT}
1001
1002def pkt_leaf(r: PK) -> PKT:
1003 match r:
1004 case PK{a, t}:
1005 PKT{a, PLeaf{t}}
1006
1007def pkt_then(r: PKT, k: Array<F32> -> PT -> PKT) -> PKT:
1008 match r:
1009 case PKT{a, t}:
1010 k(a, t)
1011
1012# lines i .. i+cnt-1 of a, packed into 2^d pieces (halves of the line count, as the fork tree splits them)
1013def pk_tree(+d: Nat, +cnt: Nat, +i: U32, +tk: Nat, +tku: U32, +ku: U32, +lstep: U32, +st: U32, a: Array<F32>) -> PKT:
1014 match d:
1015 case 0n:
1016 pkt_leaf(pk_lines(cnt, tk, tku, ku, i, 0, lstep, st, PK{a, Array.new(T8, cap_depth(Nat.mul(cnt, tk)), t8_zero())}))
1017 case 1n+q:
1018 pkt_then(pk_tree(q, Nat.div(cnt, 2n), i, tk, tku, ku, lstep, st, a), a2 => l => pkt_then(pk_tree(q, Nat.sub(cnt, Nat.div(cnt, 2n)), (i + U32.from_nat(Nat.div(cnt, 2n)) : U32), tk, tku, ku, lstep, st, a2), a3 => r => PKT{a3, PNode{l, r}}))
1019
1020def t8_with2(-R: Type, p: Array<T8> & Array<T8>, k: Array<T8> -> Array<T8> -> R) -> R:
1021 match p:
1022 case Tuple{x, y}:
1023 k(x, y)
1024
1025# split by rows of C: the pieces are A's rows, B is shared; the result is C's rows in order
1026def t8_rpar(t: PT, +cnt: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, b: Array<T8>) -> List<&2, F32>:
1027 match t:
1028 case PLeaf{a}:
1029 t8_lfin(t8_lrows(cnt, m, tk, tku, rot, 0, TL{a, b, Nil{}}))
1030 case PNode{l, r}:
1031 t8_with2(List<&2, F32>, Array.clone(T8, b), b1 => b2 => (
1032 x y = t8_rpar(l, Nat.div(cnt, 2n), m, tk, tku, rot, b1) t8_rpar(r, Nat.sub(cnt, Nat.div(cnt, 2n)), m, tk, tku, rot, b2)
1033 List.append(&2, F32, x, y)))
1034
1035# a block of C's columns j0 .. j0+w-1 for every row, row by row
1036type CB is Data:
1037 CB{j0: U32, w: Nat, xs: List<&2, F32>}
1038
1039# split by columns of C: the pieces are B's columns, A is shared
1040def t8_cpar(t: PT, +cnt: Nat, +j0: U32, +n: Nat, +tk: Nat, +tku: U32, +rot: Nat, a: Array<T8>) -> List<&2, CB>:
1041 match t:
1042 case PLeaf{b}:
1043 [CB{j0, cnt, t8_lfin(t8_lrows(n, cnt, tk, tku, rot, 0, TL{a, b, Nil{}}))}]
1044 case PNode{l, r}:
1045 t8_with2(List<&2, CB>, Array.clone(T8, a), a1 => a2 => (
1046 x y = t8_cpar(l, Nat.div(cnt, 2n), j0, n, tk, tku, rot, a1) t8_cpar(r, Nat.sub(cnt, Nat.div(cnt, 2n)), (j0 + U32.from_nat(Nat.div(cnt, 2n)) : U32), n, tk, tku, rot, a2)
1047 List.append(&2, CB, x, y)))
1048
1049type RW is Type:
1050 RW{c: Array<F32>, xs: List<&2, F32>}
1051
1052# w numbers of xs into c from idx; what is left of xs comes back
1053def cb_row(+ac: Bool, w: Nat, +idx: U32, c: Array<F32>, xs: List<&2, F32>) -> RW:
1054 match w xs:
1055 case 1n+p Con{+h, t}:
1056 cb_row(ac, p, (idx + 1 : U32), cwr(ac, idx, h, c), t)
1057 case _ _:
1058 RW{c, xs}
1059
1060def cb_rows(+ac: Bool, n: Nat, +i: U32, +j0: U32, +w: Nat, +mu: U32, r: RW) -> Array<F32>:
1061 match n r:
1062 case 0n RW{c, xs}:
1063 c
1064 case 1n+p RW{c, xs}:
1065 cb_rows(ac, p, (i + 1 : U32), j0, w, mu, cb_row(ac, w, ((i * mu : U32) + j0 : U32), c, xs))
1066
1067def cb_write(+ac: Bool, bs: List<&2, CB>, +n: Nat, +mu: U32, c: Array<F32>) -> Array<F32>:
1068 match bs:
1069 case Nil{}:
1070 c
1071 case Con{CB{+j0, +w, xs}, t}:
1072 cb_write(ac, t, n, mu, cb_rows(ac, n, 0, j0, w, mu, RW{c, xs}))
1073
1074def t8_pr2(+ac: Bool, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array<F32>, ra: PKT, rb: PK) -> G:
1075 match ra rb:
1076 case PKT{a, t} PK{b, tb}:
1077 G{a, b, write_la(ac, t8_rpar(t, n, m, tk, tku, rot, tb), 0, c)}
1078
1079def t8_pc2(+ac: Bool, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array<F32>, ra: PK, rb: PKT) -> G:
1080 match ra rb:
1081 case PK{a, ta} PKT{b, t}:
1082 G{a, b, cb_write(ac, t8_cpar(t, m, 0, n, tk, tku, rot, ta), n, U32.from_nat(m), c)}
1083
1084# 2^d blocks along the larger of n (rows of C, pieces of A) and m (columns of C, pieces of B)
1085def t8_par2(+ac: Bool, byrows: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +tk: Nat, +tku: U32, +rot: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1086 match byrows:
1087 case True{}:
1088 t8_pr2(ac, n, m, tk, tku, rot, c, pk_tree(d, n, 0, tk, tku, U32.from_nat(k), arow, sa, a), pk_pack(m, tk, k, bcol, sb, b))
1089 case False{}:
1090 t8_pc2(ac, n, m, tk, tku, rot, c, pk_pack(n, tk, k, arow, sa, a), pk_tree(d, m, 0, tk, tku, U32.from_nat(k), bcol, sb, b))
1091
1092def gemm_t8b(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +tk: Nat, +tku: U32, +rot: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1093 match d:
1094 case 0n:
1095 t8_out_seq(ac, n, m, tk, tku, rot, c, t8_pk2(pk_pack(n, tk, k, arow, sa, a), pk_pack(m, tk, k, bcol, sb, b)))
1096 case 1n+q:
1097 t8_par2(ac, Nat.is_le(m, n), 1n+q, n, m, k, tk, tku, rot, arow, sa, sb, bcol, a, b, c)
1098
1099# C (n x m) = A · B (or C + A · B when ac) in cells of 8, sequential (d = 0) or with 2^d blocks of rows in parallel; A
1100# and B come back intact
1101def gemm_t8_ac(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1102 gemm_t8b(ac, d, n, m, k, Nat.div(Nat.add(k, 7n), 8n), U32.from_nat(Nat.div(Nat.add(k, 7n), 8n)), Nat.mod(k, 4n), arow, sa, sb, bcol, a, b, c)
1103
1104def gemm_t8(d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1105 gemm_t8_ac(False{}, d, n, m, k, arow, sa, sb, bcol, a, b, c)
1106
1107# one row (matrix · vector): the scalar loop at every depth; two or more rows: cells of 8. Up to 0.1.11.0 one row in
1108# parallel ran gemm_par_cols, which clones both operands at every node of the fork tree: a 50,257 x 768 weight matrix at
1109# depth 4 ran out of the default heap (5.1 GB) and 8192 x 2048 took 124 ms against 10 ms sequentially. Packing the
1110# matrix into cells of 8 instead costs more than the product, which reads each weight once (64 ms against 24 at
1111# 50,257 x 768). Matrix · vector in parallel is Bands.matvec. Measured 2026-10-16 (NOTES.md).
1112def gemm_any3(+ac: Bool, one: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1113 match one:
1114 case True{}:
1115 gemm_ac(ac, n, m, k, arow, sa, sb, bcol, a, b, c)
1116 case False{}:
1117 gemm_t8_ac(ac, d, n, m, k, arow, sa, sb, bcol, a, b, c)
1118
1119def gemm_any2(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1120 gemm_any3(ac, Nat.is_le(n, 1n), d, n, m, k, arow, sa, sb, bcol, a, b, c)
1121
1122def gemm_any(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1123 gemm_any2(ac, d, n, m, k, arow, sa, sb, bcol, a, b, c)
1124
1125def mmul_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMul<n, k, m>:
1126 match g:
1127 case G{a, b, c}:
1128 MMul{Mat{a}, Mat{b}, Mat{c}}
1129
1130def mmul_nt_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulNT<n, k, m>:
1131 match g:
1132 case G{a, b, c}:
1133 MMulNT{Mat{a}, Mat{b}, Mat{c}}
1134
1135def mmul_tn_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulTN<n, k, m>:
1136 match g:
1137 case G{a, b, c}:
1138 MMulTN{Mat{a}, Mat{b}, Mat{c}}
1139
1140# ---- the split depth of a product (0.1.11.0): par is an upper bound, and a product too small to gain splits 0 times ----
1141#
1142# A product of n x k by k x m splits into 2^par blocks only when it does at least 2^23 multiply-adds and each cell of C
1143# sums at least 64 of them; otherwise it runs sequentially. Measured on 16 threads (2026-10-16, 18 shapes): below 2^23 the
1144# split costs more than it saves, and with a short inner dimension (a weight gradient over a few tokens) every split
1145# was 2-3.6x slower, because each block hands its rows of C back as a list that is then written into C, a cost per cell
1146# of C that k multiply-adds per cell do not pay for. The result does not depend on the depth: the same bits at every
1147# depth (tests/t8.bend, tests/acc.bend, tests/split.bend).
1148def split_pick(big: Bool, d: Nat) -> Nat:
1149 match big:
1150 case True{}:
1151 d
1152 case False{}:
1153 0n
1154
1155def split_depth(d: Nat, +n: Nat, +m: Nat, +k: Nat) -> Nat:
1156 split_pick(Bool.and(Nat.is_le(64n, k), Nat.is_le(8388608n, Nat.mul(Nat.mul(n, m), k))), d)
1157
1158def mm_a(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMul<n, k, m>:
1159 mmul_fin(n, k, m, gemm_any(False{}, d, n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0)))
1160
1161# C = A · B (n x k) · (k x m); par = log2 of the most parallel blocks (0 = sequential); split_depth picks 0 or par
1162
1163def Mat.matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<k, m>) -> MMul<n, k, m>:
1164 match a b:
1165 case Mat{x} Mat{y}:
1166 mm_a(n, k, m, split_depth(par, n, m, k), x, y)
1167
1168# C + A · B, into C (n x m); A and B come back with the result, as Mat.matmul returns them. Each cell of A · B is
1169# computed as Mat.matmul computes it and added once to C's cell: the same bits as Mat.add(c, Mat.matmul(a, b)'s c),
1170# without allocating the product or walking it a second time.
1171
1172def Mat.matmul_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<k, m>, c: Mat<n, m>) -> MMul<n, k, m>:
1173 match a b c:
1174 case Mat{x} Mat{y} Mat{z}:
1175 mmul_fin(n, k, m, gemm_any(True{}, split_depth(par, n, m, k), n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, x, y, z))
1176
1177def mm_nt(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMulNT<n, k, m>:
1178 mmul_nt_fin(n, k, m, gemm_any(False{}, d, n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0)))
1179
1180# C = A · Bᵀ (n x k) · (m x k)ᵀ: for weights stored with one row per output
1181
1182def Mat.matmul_nt(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<m, k>) -> MMulNT<n, k, m>:
1183 match a b:
1184 case Mat{x} Mat{y}:
1185 mm_nt(n, k, m, split_depth(par, n, m, k), x, y)
1186
1187# C + A · Bᵀ, into C (n x m)
1188
1189def Mat.matmul_nt_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<m, k>, c: Mat<n, m>) -> MMulNT<n, k, m>:
1190 match a b c:
1191 case Mat{x} Mat{y} Mat{z}:
1192 mmul_nt_fin(n, k, m, gemm_any(True{}, split_depth(par, n, m, k), n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), x, y, z))
1193
1194def mm_tn(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMulTN<n, k, m>:
1195 mmul_tn_fin(n, k, m, gemm_any(False{}, d, n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0)))
1196
1197# C = Aᵀ · B (k x n)ᵀ · (k x m): weight gradient, X^T · dY
1198
1199def Mat.matmul_tn(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<k, n>, b: Mat<k, m>) -> MMulTN<n, k, m>:
1200 match a b:
1201 case Mat{x} Mat{y}:
1202 mm_tn(n, k, m, split_depth(par, n, m, k), x, y)
1203
1204# C + Aᵀ · B, into C (n x m): a weight gradient added into the gradient it accumulates, dW += Xᵀ · dY
1205
1206def Mat.matmul_tn_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<k, n>, b: Mat<k, m>, c: Mat<n, m>) -> MMulTN<n, k, m>:
1207 match a b c:
1208 case Mat{x} Mat{y} Mat{z}:
1209 mmul_tn_fin(n, k, m, gemm_any(True{}, split_depth(par, n, m, k), n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, x, y, z))
1210
1211# ---- a precise product for long dots (0.1.14.0) ----
1212#
1213# Mat.matmul_precise is Mat.matmul with k summed in blocks of 1024 (the last one 1..1024): block b's columns of A and
1214# rows of B are copied out, the first block's product written into C, each next one's added to it as matmul_acc adds,
1215# so every cell is ((x0 + x1) + x2) .., x_b the block's dot with the usual four partial sums. A dot of 49,152 products
1216# as four running sums of 12,288 lost the small terms (SmolLM2's head: its gradient 1.8x further from float64 than
1217# PyTorch's float32, 0.03x in blocks; BendTorch NOTES, 2026-10-18). Opt-in: it copies the slices and runs one product
1218# a block, and the default products keep their speed and bits. Up to 1024 steps it is Mat.matmul itself.
1219
1220type SL is Type:
1221 SL{a: Array<F32>, s: Array<F32>}
1222
1223# d[di..] = src[si..], k numbers; r holds src[si], already read
1224def sl_row(k: Nat, +si: U32, +di: U32, d: Array<F32>, r: Array<F32> & F32) -> SL:
1225 match k r:
1226 case 0n Tuple{a, x}: SL{a, d}
1227 case 1n+p Tuple{a, +x}: sl_row(p, (si + 1 : U32), (di + 1 : U32), Array.set(F32, d, di, x), Array.get(F32, a, (si + 1 : U32)))
1228
1229# rows i.. of a slice: w numbers a row from src[i*st + o], into d[i*w]
1230def sl_rows(r: Nat, +i: U32, +w: Nat, +w32: U32, +st: U32, +o: U32, sl: SL) -> SL:
1231 match r sl:
1232 case 0n sl: sl
1233 case 1n+p SL{a, d}: sl_rows(p, (i + 1 : U32), w, w32, st, o, sl_row(w, ((i * st : U32) + o : U32), (i * w32 : U32), d, Array.get(F32, a, ((i * st : U32) + o : U32))))
1234
1235def blk_len(+bi: U32, +k: U32) -> Nat:
1236 U32.to_nat(Bool.pick(U32, U32.is_lt((k - (bi * 1024 : U32) : U32), 1024), (k - (bi * 1024 : U32) : U32), 1024))
1237
1238def mp_s4(a: Array<F32>, b: Array<F32>, g: G) -> G:
1239 match g:
1240 case G{x, y, c}: G{a, b, c}
1241
1242def mp_s3(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +w: Nat, +w32: U32, +m32: U32, a: Array<F32>, as: Array<F32>, c: Array<F32>, sb: SL) -> G:
1243 match sb:
1244 case SL{b, bs}: mp_s4(a, b, gemm_any(ac, split_depth(par, n, m, w), n, m, w, w32, 1, m32, 1, as, bs, c))
1245
1246def mp_s2(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +w: Nat, +w32: U32, +m32: U32, +o: U32, c: Array<F32>, b: Array<F32>, sa: SL) -> G:
1247 match sa:
1248 case SL{a, as}: mp_s3(ac, par, n, m, w, w32, m32, a, as, c, sl_row(Nat.mul(w, m), (o * m32 : U32), 0, Array.new(F32, cap_depth(Nat.mul(w, m)), 0.0), Array.get(F32, b, (o * m32 : U32))))
1249
1250# block bi: A's columns o..o+w (A is n x k) and B's rows o..o+w (B is k x m), multiplied into C (added when ac)
1251def mp_step(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +k32: U32, +bi: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
1252 mp_s2(ac, par, n, m, blk_len(bi, k32), U32.from_nat(blk_len(bi, k32)), U32.from_nat(m), (bi * 1024 : U32), c, b, sl_rows(n, 0, blk_len(bi, k32), U32.from_nat(blk_len(bi, k32)), k32, (bi * 1024 : U32), SL{a, Array.new(F32, cap_depth(Nat.mul(n, blk_len(bi, k32))), 0.0)}))
1253
1254def mp_go(nb: Nat, +bi: U32, +par: Nat, +n: Nat, +m: Nat, +k32: U32, g: G) -> G:
1255 match nb g:
1256 case 0n g: g
1257 case 1n++q G{a, b, c}: mp_go(q, (bi + 1 : U32), par, n, m, k32, mp_step(U32.is_lt(0, bi), par, n, m, k32, bi, a, b, c))
1258
1259def mm_precise2(long: Bool, +n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<k, m>) -> MMul<n, k, m>:
1260 match long a b:
1261 case True{} Mat{x} Mat{y}: mmul_fin(n, k, m, mp_go(U32.to_nat(U32.div((U32.from_nat(k) + 1023 : U32), 1024)), 0, par, n, m, U32.from_nat(k), G{x, y, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0)}))
1262 case False{} a b: Mat.matmul(n, k, m, par, a, b)
1263
1264# C = A · B (n x k) · (k x m), k summed in blocks of 1024 (above); opt-in
1265def Mat.matmul_precise(+n: Nat, +k: Nat, +m: Nat, +par: Nat, a: Mat<n, k>, b: Mat<k, m>) -> MMul<n, k, m>:
1266 mm_precise2(Nat.is_lt(1024n, k), n, k, m, par, a, b)
1267
1268# ---- element-wise and training operations ----
1269
1270# adds the bias (one row, m numbers) to each row of C (n x m)
1271
1272type MBias<-n: Nat, -m: Nat> is Type:
1273 MBias{c: Mat<n, m>, b: Mat<1n, m>}
1274
1275def add_row2(-n: Nat, -m: Nat, +nn: Nat, +mm: Nat, cd: Array<F32>, al: AL) -> MBias<n, m>:
1276 match al:
1277 case AL{bd, bl}:
1278 MBias{Mat{bias_rows(nn, 0, U32.from_nat(mm), bl, cd)}, Mat{bd}}
1279
1280def Mat.add_row(+n: Nat, +m: Nat, c: Mat<n, m>, b: Mat<1n, m>) -> MBias<n, m>:
1281 match c b:
1282 case Mat{cd} Mat{bd}:
1283 add_row2(n, m, n, m, cd, read_l(Nat.sub(m, 1n), 0, Nil{}, Array.get(F32, bd, 0)))
1284
1285# pair of matrices of the same shape (the result of an operation that reads two operands)
1286
1287type MPair<-r: Nat, -c: Nat> is Type:
1288 MPair{a: Mat<r, c>, b: Mat<r, c>}
1289
1290def mpair_fin(-r: Nat, -c: Nat, p: P2) -> MPair<r, c>:
1291 match p:
1292 case P2{a, b}:
1293 MPair{Mat{a}, Mat{b}}
1294
1295def relu_fin(-r: Nat, -c: Nat, d: Array<F32>) -> Mat<r, c>:
1296 Mat{d}
1297
1298def Mat.relu(+r: Nat, +c: Nat, a: Mat<r, c>) -> Mat<r, c>:
1299 match a:
1300 case Mat{d}:
1301 relu_fin(r, c, relu_l(Nat.mul(r, c), 0, Array.get(F32, d, 0)))
1302
1303# (dy * (h > 0), h): the gradient that passes through the relu
1304
1305def Mat.relu_bwd(+r: Nat, +c: Nat, dy: Mat<r, c>, h: Mat<r, c>) -> MPair<r, c>:
1306 match dy h:
1307 case Mat{d} Mat{x}:
1308 mpair_fin(r, c, mask_l(Nat.mul(r, c), 0, Array.get(F32, d, 0), Array.get(F32, x, 0)))
1309
1310# (w - lr * dw, dw): one SGD step
1311
1312def Mat.sgd(+r: Nat, +c: Nat, +lr: F32, w: Mat<r, c>, dw: Mat<r, c>) -> MPair<r, c>:
1313 match w dw:
1314 case Mat{x} Mat{y}:
1315 mpair_fin(r, c, sgd_l(Nat.mul(r, c), 0, lr, Array.get(F32, x, 0), Array.get(F32, y, 0)))
1316
1317# (a + b, b): element-wise sum (a - (-1)*b is exactly a + b)
1318
1319def Mat.add(+r: Nat, +c: Nat, a: Mat<r, c>, b: Mat<r, c>) -> MPair<r, c>:
1320 match a b:
1321 case Mat{x} Mat{y}:
1322 mpair_fin(r, c, sgd_l(Nat.mul(r, c), 0, (0.0 - 1.0 : F32), Array.get(F32, x, 0), Array.get(F32, y, 0)))
1323
1324# sum of each column (n x m -> 1 x m), returning the matrix
1325
1326type MColSum<-n: Nat, -m: Nat> is Type:
1327 MColSum{a: Mat<n, m>, s: Mat<1n, m>}
1328
1329def colsum_fin(-n: Nat, -m: Nat, g: G) -> MColSum<n, m>:
1330 match g:
1331 case G{ones, a, s}:
1332 MColSum{Mat{a}, Mat{s}}
1333
1334def Mat.col_sums(+n: Nat, +m: Nat, a: Mat<n, m>) -> MColSum<n, m>:
1335 match a:
1336 case Mat{d}:
1337 colsum_fin(n, m, gemm(1n, m, n, 0, 1, U32.from_nat(m), 1, Array.new(F32, cap_depth(n), 1.0), d, Array.new(F32, cap_depth(m), 0.0)))
1338
1339# s + the sum of each column, into s (1 x m): a bias gradient added into its accumulator, with the bits of
1340# Mat.add(s, Mat.col_sums(a)'s s)
1341
1342def Mat.col_sums_acc(+n: Nat, +m: Nat, a: Mat<n, m>, s: Mat<1n, m>) -> MColSum<n, m>:
1343 match a s:
1344 case Mat{d} Mat{z}:
1345 colsum_fin(n, m, gemm_ac(True{}, 1n, m, n, 0, 1, U32.from_nat(m), 1, Array.new(F32, cap_depth(n), 1.0), d, z))
1346
1347# mean cross-entropy and its gradient on the logits: (softmax - one-hot) / n.
1348# labels has one entry per row (n); the gradient replaces the logits.
1349
1350type MCE<-n: Nat, -c: Nat> is Type:
1351 MCE{loss: F32, grad: Mat<n, c>}
1352
1353def ce_fin(-n: Nat, -c: Nat, +nf: F32, r: L2) -> MCE<n, c>:
1354 match r:
1355 case L2{loss, d}:
1356 MCE{(loss / nf : F32), Mat{d}}
1357
1358def Mat.softmax_ce(+n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> MCE<n, c>:
1359 match z:
1360 case Mat{d}:
1361 ce_fin(n, c, F32.from_nat(n), cea_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), F32.from_nat(n), 0, L2{0.0, d}))
1362
1363# argmax hits per row against labels
1364
1365type MHits<-n: Nat, -c: Nat> is Type:
1366 MHits{z: Mat<n, c>, hits: Nat}
1367
1368def hits_fin(-n: Nat, -c: Nat, r: CR) -> MHits<n, c>:
1369 match r:
1370 case CR{d, k}:
1371 MHits{Mat{d}, k}
1372
1373def Mat.count_correct(+n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> MHits<n, c>:
1374 match z:
1375 case Mat{d}:
1376 hits_fin(n, c, count_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), 0, CR{d, 0n}))
1377
1378# a row as a list (returning the matrix) and the other way around
1379
1380type MRow<-r: Nat, -c: Nat> is Type:
1381 MRow{m: Mat<r, c>, row: List<&2, F32>}
1382
1383def row_fin(-r: Nat, -c: Nat, al: AL) -> MRow<r, c>:
1384 match al:
1385 case AL{d, l}:
1386 MRow{Mat{d}, l}
1387
1388def Mat.read_row(+r: Nat, +c: Nat, m: Mat<r, c>, +i: U32) -> MRow<r, c>:
1389 match m:
1390 case Mat{d}:
1391 row_fin(r, c, read_l(Nat.sub(c, 1n), (i * U32.from_nat(c) : U32), Nil{}, Array.get(F32, d, (i * U32.from_nat(c) : U32))))
1392
1393def Mat.write_row(-r: Nat, +c: Nat, m: Mat<r, c>, +i: U32, xs: List<&2, F32>) -> Mat<r, c>:
1394 match m:
1395 case Mat{d}:
1396 Mat{write_l(xs, (i * U32.from_nat(c) : U32), d)}
1397
1398# ---- variants that check the number of labels ----
1399
1400# does the list of labels have exactly n entries (one per row)?
1401def labels_len_is(xs: List<&2, Nat>, n: Nat) -> Bool:
1402 match xs n:
1403 case Nil{} 0n:
1404 True{}
1405 case Con{h, t} 1n+p:
1406 labels_len_is(t, p)
1407 case _ _:
1408 False{}
1409
1410def ce_checked(ok: Bool, +n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> Maybe<&1, MCE<n, c>>:
1411 match ok:
1412 case True{}:
1413 Some{Mat.softmax_ce(n, c, z, labels)}
1414 case False{}:
1415 None{}
1416
1417# like Mat.softmax_ce, but returns None unless labels has exactly n entries
1418def Mat.softmax_ce_checked(+n: Nat, +c: Nat, z: Mat<n, c>, +labels: List<&2, Nat>) -> Maybe<&1, MCE<n, c>>:
1419 ce_checked(labels_len_is(labels, n), n, c, z, labels)
1420
1421def hits_checked(ok: Bool, +n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> Maybe<&1, MHits<n, c>>:
1422 match ok:
1423 case True{}:
1424 Some{Mat.count_correct(n, c, z, labels)}
1425 case False{}:
1426 None{}
1427
1428# like Mat.count_correct, but returns None unless labels has exactly n entries
1429def Mat.count_correct_checked(+n: Nat, +c: Nat, z: Mat<n, c>, +labels: List<&2, Nat>) -> Maybe<&1, MHits<n, c>>:
1430 hits_checked(labels_len_is(labels, n), n, c, z, labels)
1431
1432# ---------------------------------------------------------------------
1433# Bands<r, c>: a matrix cut into bands of rows, for matrix · vector in parallel
1434# without copying the weights.
1435#
1436# Mat.matmul_nt with n = 1 can only run in parallel by cloning the whole weight
1437# matrix for every task, and in matrix · vector each weight is read once, so the
1438# copies cost more than the products (NOTES.md, exp. 9). Bands keeps each band in
1439# its own Array: a task takes its band and nothing is copied except the input vector.
1440#
1441# A node with r rows has children with half(r) and r - half(r) rows. The split is
1442# part of the type, so the row count of every band follows from r, and the law
1443# half_cover says that the two halves always cover the r rows exactly.
1444# ---------------------------------------------------------------------
1445
1446# n / 2 rounded down. A tail loop with an accumulator: a non-tail recursion here would put a
1447# continuation on every node of the band tree and keep the bands from running in parallel.
1448def half.go(n: Nat, acc: Nat) -> Nat:
1449 match n:
1450 case 0n:
1451 acc
1452 case 1n+p:
1453 match p:
1454 case 0n:
1455 acc
1456 case 1n+q:
1457 half.go(q, 1n+acc)
1458
1459def half(n: Nat) -> Nat:
1460 half.go(n, 0n)
1461
1462# a <= a
1463def leb_refl(+a: Nat) -> {True{} == leb(a, a) : Bool}:
1464 match a:
1465 case 0n:
1466 {==}
1467 case 1n+p:
1468 leb_refl(p)
1469
1470# 1 + (a + b) == a + (1 + b)
1471def add_succ_r(+a: Nat, -b: Nat) -> {1n+Nat.add(a, b) == Nat.add(a, 1n+b) : Nat}:
1472 match a:
1473 case 0n:
1474 {==}
1475 case 1n+p:
1476 %add_succ_r(p, b) : {2n+Nat.add(p, b) == 1n+_ : Nat}
1477 {==}
1478
1479# a + 0 == a
1480def add_zero_r(+a: Nat) -> {Nat.add(a, 0n) == a : Nat}:
1481 match a:
1482 case 0n:
1483 {==}
1484 case 1n+p:
1485 %add_zero_r(p) : {1n+Nat.add(p, 0n) == 1n+_ : Nat}
1486 {==}
1487
1488# half.go(n, acc) <= n + acc
1489def half_go_le(+n: Nat, +acc: Nat) -> {True{} == leb(half.go(n, acc), Nat.add(n, acc)) : Bool}:
1490 match n:
1491 case 0n:
1492 leb_refl(acc)
1493 case 1n+p:
1494 match p:
1495 case 0n:
1496 leb_succ(acc, acc, leb_refl(acc))
1497 case 1n+q:
1498 %Equal.sym(Nat, 1n+Nat.add(q, acc), Nat.add(q, 1n+acc), add_succ_r(q, acc)) : {True{} == leb(half.go(q, 1n+acc), 1n+_) : Bool}
1499 leb_succ(half.go(q, 1n+acc), Nat.add(q, 1n+acc), half_go_le(q, 1n+acc))
1500
1501# half(n) <= n
1502def half_le(+n: Nat) -> {True{} == leb(half(n), n) : Bool}:
1503 %add_zero_r(n) : {True{} == leb(half(n), _) : Bool}
1504 half_go_le(n, 0n)
1505
1506# a <= r -> r == a + (r - a)
1507def add_sub(+a: Nat, +r: Nat, e: {True{} == leb(a, r) : Bool}) -> {r == Nat.add(a, Nat.sub(r, a)) : Nat}:
1508 match a r:
1509 case 0n 0n:
1510 {==}
1511 case 0n 1n+q:
1512 {==}
1513 case 1n+p 0n:
1514 %e : BD(_, Unit, {0n == Nat.add(1n+p, Nat.sub(0n, 1n+p)) : Nat})
1515 Unit{}
1516 case 1n+p 1n+q:
1517 %add_sub(p, q, e) : {1n+q == 1n+_ : Nat}
1518 {==}
1519
1520# LAW (half_cover): the two bands of a node, half(r) rows and r - half(r) rows,
1521# add up to exactly r rows: no row is lost and none is counted twice.
1522law half_cover:
1523 for +n: Nat
1524 {n == Nat.add(half(n), Nat.sub(n, half(n))) : Nat}
1525
1526def half_cover(n):
1527 add_sub(half(n), n, half_le(n))
1528
1529type Bands<-r: Nat, -c: Nat> is Type:
1530 BLeaf{m: Mat<r, c>}
1531 BNode{x: Bands<half(r), c>, y: Bands<Nat.sub(r, half(r)), c>}
1532
1533# r x c of zeros in 2^d bands (d = 0: a single band, the same as a Mat)
1534def Bands.zeros(+d: Nat, +r: Nat, +c: Nat) -> Bands<r, c>:
1535 match d:
1536 case 0n:
1537 BLeaf{Mat.zeros(r, c)}
1538 case 1n+q:
1539 BNode{Bands.zeros(q, half(r), c), Bands.zeros(q, Nat.sub(r, half(r)), c)}
1540
1541# the row counts of the two bands follow the row count of the node: if rr == r at run
1542# time, the same holds for both halves (the bands carry r only in their type)
1543def eq_half(-a: Nat, -b: Nat, e: {a == b : Nat}) -> {half(a) == half(b) : Nat}:
1544 %e : {half(a) == half(_) : Nat}
1545 {==}
1546
1547def eq_rest(-a: Nat, -b: Nat, e: {a == b : Nat}) -> {Nat.sub(a, half(a)) == Nat.sub(b, half(b)) : Nat}:
1548 %e : {Nat.sub(a, half(a)) == Nat.sub(_, half(_)) : Nat}
1549 {==}
1550
1551# does a block of n numbers starting at flat index i end before the second band of a
1552# node with r rows? does it start in the second band?
1553def bfill_l(+r: Nat, +c: Nat, +i: Nat, +n: Nat) -> Bool:
1554 Nat.is_le(Nat.add(i, n), Nat.mul(half(r), c))
1555
1556def bfill_r(+r: Nat, +c: Nat, +i: Nat) -> Bool:
1557 Nat.is_le(Nat.mul(half(r), c), i)
1558
1559# writes the first n numbers of xs at index i (and nothing past them)
1560def fill_n(n: Nat, xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
1561 match n xs:
1562 case 1n+p Con{h, t}:
1563 fill_n(p, t, (i + 1 : U32), Array.set(F32, a, i, h))
1564 case _ _:
1565 a
1566
1567def fill_band(-r: Nat, -c: Nat, n: Nat, xs: List<&2, F32>, +i: U32, m: Mat<r, c>) -> Mat<r, c>:
1568 match m:
1569 case Mat{d}:
1570 Mat{fill_n(n, xs, i, d)}
1571
1572# writes the first n numbers of xs at flat index i; l and rt say which bands they touch
1573# (bfill_l, bfill_r); rr is the row count r at run time. A block that crosses into the second
1574# band is not copied: the first band writes only its part, the second gets the rest by List.drop.
1575def bfill(+c: Nat, -r: Nat, b: Bands<r, c>, l: Bool, rt: Bool, +rr: Nat, -e: {rr == r : Nat}, +i: Nat, +n: Nat, +xs: List<&2, F32>) -> Bands<r, c>:
1576 match b l rt:
1577 case BLeaf{m} _ _:
1578 BLeaf{fill_band(r, c, n, xs, U32.from_nat(i), m)}
1579 case BNode{x, y} True{} _:
1580 BNode{bfill(c, half(r), x, bfill_l(half(rr), c, i, n), bfill_r(half(rr), c, i), half(rr), eq_half(rr, r, e), i, n, xs), y}
1581 case BNode{x, y} False{} True{}:
1582 BNode{x, bfill(c, Nat.sub(r, half(r)), y, bfill_l(Nat.sub(rr, half(rr)), c, Nat.sub(i, Nat.mul(half(rr), c)), n), bfill_r(Nat.sub(rr, half(rr)), c, Nat.sub(i, Nat.mul(half(rr), c))), Nat.sub(rr, half(rr)), eq_rest(rr, r, e), Nat.sub(i, Nat.mul(half(rr), c)), n, xs)}
1583 case BNode{x, y} False{} False{}:
1584 BNode{bfill(c, half(r), x, bfill_l(half(rr), c, i, Nat.sub(Nat.mul(half(rr), c), i)), bfill_r(half(rr), c, i), half(rr), eq_half(rr, r, e), i, Nat.sub(Nat.mul(half(rr), c), i), xs), bfill(c, Nat.sub(r, half(r)), y, bfill_l(Nat.sub(rr, half(rr)), c, 0n, Nat.sub(n, Nat.sub(Nat.mul(half(rr), c), i))), bfill_r(Nat.sub(rr, half(rr)), c, 0n), Nat.sub(rr, half(rr)), eq_rest(rr, r, e), 0n, Nat.sub(n, Nat.sub(Nat.mul(half(rr), c), i)), List.drop(&2, F32, xs, Nat.sub(Nat.mul(half(rr), c), i)))}
1585
1586def bfill_top(ok: Bool, +r: Nat, +c: Nat, b: Bands<r, c>, +i: Nat, +n: Nat, +xs: List<&2, F32>) -> Bands<r, c>:
1587 match ok:
1588 case True{}:
1589 bfill(c, r, b, bfill_l(r, c, i, n), bfill_r(r, c, i), r, {==}, i, n, xs)
1590 case False{}:
1591 bfill(c, r, b, bfill_l(r, c, i, Nat.sub(Nat.mul(r, c), i)), bfill_r(r, c, i), r, {==}, i, Nat.sub(Nat.mul(r, c), i), xs)
1592
1593# writes the n numbers of xs into b starting at flat index i (row-major, as Mat.fill_at);
1594# numbers past r*c are dropped. For loading a large matrix block by block.
1595def Bands.fill_at(+r: Nat, +c: Nat, +i: Nat, +n: Nat, +xs: List<&2, F32>, b: Bands<r, c>) -> Bands<r, c>:
1596 bfill_top(Nat.is_le(Nat.add(i, n), Nat.mul(r, c)), r, c, b, i, n, xs)
1597
1598# ---- y = B · x for B with r rows of c numbers (a weight matrix stored out x in) ----
1599
1600# B (intact) and y = B · x, typed like Mat.matmul_nt: x is 1 x c, y is 1 x r
1601type BV<-r: Nat, -c: Nat> is Type:
1602 BV{b: Bands<r, c>, y: Mat<1n, r>}
1603
1604# the bands and y as a list in row order (the result of Bands.matvec_l)
1605type BL<-r: Nat, -c: Nat> is Type:
1606 BL{b: Bands<r, c>, y: List<&2, F32>}
1607
1608def bv_g2(-r: Nat, -c: Nat, w: Array<F32>, al: AL) -> BL<r, c>:
1609 match al:
1610 case AL{d, l}:
1611 BL{BLeaf{Mat{w}}, l}
1612
1613# the band's 1 + m rows came out in an Array (one number per row); read them back as a list
1614def bv_g(-r: Nat, -c: Nat, +m: Nat, g: G) -> BL<r, c>:
1615 match g:
1616 case G{a, b, y}:
1617 bv_g2(r, c, b, read_l(m, 0, Nil{}, Array.get(F32, y, 0)))
1618
1619# the rows of one band, one dot product each, with the gemm of Mat.matmul_nt (n = 1), so the
1620# results are identical. Writing the results into an Array and reading them back is faster than
1621# building the list row by row (NOTES.md, exp. 14).
1622def bv_gemm(-r: Nat, +c: Nat, rows: Nat, +rr: Nat, x: Array<F32>, w: Array<F32>) -> BL<r, c>:
1623 match rows:
1624 case 0n:
1625 BL{BLeaf{Mat{w}}, Nil{}}
1626 case 1n+p:
1627 bv_g(r, c, p, gemm(1n, rr, c, U32.from_nat(c), 1, 1, U32.from_nat(c), x, w, Array.new(F32, cap_depth(rr), 0.0)))
1628
1629def bv_mat(-r: Nat, +c: Nat, +rr: Nat, x: Array<F32>, m: Mat<r, c>) -> BL<r, c>:
1630 match m:
1631 case Mat{w}:
1632 bv_gemm(r, c, rr, rr, x, w)
1633
1634def bv_join(-r: Nat, -c: Nat, p: BL<half(r), c>, q: BL<Nat.sub(r, half(r)), c>) -> BL<r, c>:
1635 match p q:
1636 case BL{bx, yx} BL{by, yy}:
1637 BL{BNode{bx, by}, List.append(&2, F32, yx, yy)}
1638
1639# the two bands of a node run in parallel, each with its own copy of x (c numbers);
1640# xp is a pair of copies of x (a leaf uses one). xp comes after the erased e: up to Bend 2.0.36 an erased
1641# parameter at the end of the list turned the parallel let into two sequential calls (NOTES.md, exp. 11; fixed in
1642# 2.0.37 by bendlang/bend#1377). The order is kept, and harmless.
1643def bmv(+c: Nat, -r: Nat, b: Bands<r, c>, +rr: Nat, -e: {rr == r : Nat}, xp: Array<F32> & Array<F32>) -> BL<r, c>:
1644 match b xp:
1645 case BLeaf{m} Tuple{x, x2}:
1646 bv_mat(r, c, rr, x, m)
1647 case BNode{bx, by} Tuple{x1, x2}:
1648 p q = bmv(c, half(r), bx, half(rr), eq_half(rr, r, e), Array.clone(F32, x1)) bmv(c, Nat.sub(r, half(r)), by, Nat.sub(rr, half(rr)), eq_rest(rr, r, e), Array.clone(F32, x2))
1649 bv_join(r, c, p, q)
1650
1651def bv_fin(+r: Nat, -c: Nat, bl: BL<r, c>) -> BV<r, c>:
1652 match bl:
1653 case BL{b, y}:
1654 BV{b, Mat.from_list(1n, r, y)}
1655
1656# y = B · x for B (r x c, one row per output) and x (1 x c); returns B intact and y (1 x r).
1657# The same product as Mat.matmul_nt(1n, c, r, ..), and the same numbers, bit for bit.
1658def Bands.matvec(+r: Nat, +c: Nat, b: Bands<r, c>, x: Mat<1n, c>) -> BV<r, c>:
1659 match x:
1660 case Mat{xa}:
1661 bv_fin(r, c, bmv(c, r, b, r, {==}, Array.clone(F32, xa)))
1662
1663# the same product with x and y as lists (x: c numbers, y: r numbers in row order), without the
1664# Mat conversions of Bands.matvec: for callers that already hold the vector as a list
1665def Bands.matvec_l(+r: Nat, +c: Nat, b: Bands<r, c>, +x: List<&2, F32>) -> BL<r, c>:
1666 bmv(c, r, b, r, {==}, Array.clone(F32, fill_list(x, 0, Array.new(F32, cap_depth(c), 0.0))))
1667
1668# ---- reading rows and the whole matrix ----
1669
1670type BRow<-r: Nat, -c: Nat> is Type:
1671 BRow{b: Bands<r, c>, row: List<&2, F32>}
1672
1673def brow_fin(-r: Nat, -c: Nat, al: AL) -> BRow<r, c>:
1674 match al:
1675 case AL{d, l}:
1676 BRow{BLeaf{Mat{d}}, l}
1677
1678# n numbers starting at flat index i of one band (a band may have no rows)
1679def brow_leaf2(-r: Nat, -c: Nat, n: Nat, +i: U32, d: Array<F32>) -> BRow<r, c>:
1680 match n:
1681 case 0n:
1682 BRow{BLeaf{Mat{d}}, Nil{}}
1683 case 1n+p:
1684 brow_fin(r, c, read_l(p, i, Nil{}, Array.get(F32, d, i)))
1685
1686def brow_leaf(-r: Nat, -c: Nat, m: Mat<r, c>, +i: U32, n: Nat) -> BRow<r, c>:
1687 match m:
1688 case Mat{d}:
1689 brow_leaf2(r, c, n, i, d)
1690
1691def brow_l(-r: Nat, -c: Nat, p: BRow<half(r), c>, y: Bands<Nat.sub(r, half(r)), c>) -> BRow<r, c>:
1692 match p:
1693 case BRow{bx, row}:
1694 BRow{BNode{bx, y}, row}
1695
1696def brow_r(-r: Nat, -c: Nat, x: Bands<half(r), c>, q: BRow<Nat.sub(r, half(r)), c>) -> BRow<r, c>:
1697 match q:
1698 case BRow{by, row}:
1699 BRow{BNode{x, by}, row}
1700
1701# row i; left says whether i is in the first band of the node (i < half(rr))
1702def brow(+c: Nat, -r: Nat, b: Bands<r, c>, left: Bool, +rr: Nat, -e: {rr == r : Nat}, +i: Nat) -> BRow<r, c>:
1703 match b left:
1704 case BLeaf{m} _:
1705 brow_leaf(r, c, m, U32.from_nat(Nat.mul(i, c)), c)
1706 case BNode{x, y} True{}:
1707 brow_l(r, c, brow(c, half(r), x, Nat.is_lt(i, half(half(rr))), half(rr), eq_half(rr, r, e), i), y)
1708 case BNode{x, y} False{}:
1709 brow_r(r, c, x, brow(c, Nat.sub(r, half(r)), y, Nat.is_lt(Nat.sub(i, half(rr)), half(Nat.sub(rr, half(rr)))), Nat.sub(rr, half(rr)), eq_rest(rr, r, e), Nat.sub(i, half(rr))))
1710
1711# row i (c numbers), for example an embedding lookup
1712def Bands.read_row(+r: Nat, +c: Nat, b: Bands<r, c>, +i: Nat) -> BRow<r, c>:
1713 brow(c, r, b, Nat.is_lt(i, half(r)), r, {==}, i)
1714
1715def bl_join(-r: Nat, -c: Nat, p: BRow<half(r), c>, q: BRow<Nat.sub(r, half(r)), c>) -> BRow<r, c>:
1716 match p q:
1717 case BRow{bx, xs} BRow{by, ys}:
1718 BRow{BNode{bx, by}, List.append(&2, F32, xs, ys)}
1719
1720def blist(+c: Nat, -r: Nat, b: Bands<r, c>, +rr: Nat, -e: {rr == r : Nat}) -> BRow<r, c>:
1721 match b:
1722 case BLeaf{m}:
1723 brow_leaf(r, c, m, 0, Nat.mul(rr, c))
1724 case BNode{x, y}:
1725 bl_join(r, c, blist(c, half(r), x, half(rr), eq_half(rr, r, e)), blist(c, Nat.sub(r, half(r)), y, Nat.sub(rr, half(rr)), eq_rest(rr, r, e)))
1726
1727# all r*c numbers in row order (in BRow.row)
1728def Bands.to_list(+r: Nat, +c: Nat, b: Bands<r, c>) -> BRow<r, c>:
1729 blist(c, r, b, r, {==})
1730
1731# ---- the list path returns exactly r numbers ----
1732#
1733# The numbers do not matter here, only how many there are: each lemma below counts the list a
1734# function builds, following the function's own recursion.
1735
1736def bl_len(-r: Nat, -c: Nat, bl: BL<r, c>) -> Nat:
1737 match bl:
1738 case BL{b, y}:
1739 List.length(&2, F32, y)
1740
1741# s == a + (1 + b) -> s == 1 + (a + b)
1742def succ_out(-s: Nat, +a: Nat, +b: Nat, e: {s == Nat.add(a, 1n+b) : Nat}) -> {s == 1n+Nat.add(a, b) : Nat}:
1743 %Equal.sym(Nat, 1n+Nat.add(a, b), Nat.add(a, 1n+b), add_succ_r(a, b)) : {s == _ : Nat}
1744 e
1745
1746# s == 1 + (p + (1 + l)) -> s == 2 + (p + l)
1747def len_step(-s: Nat, +p: Nat, +l: Nat, e: {s == 1n+Nat.add(p, 1n+l) : Nat}) -> {s == 2n+Nat.add(p, l) : Nat}:
1748 %Equal.sym(Nat, 1n+Nat.add(p, l), Nat.add(p, 1n+l), add_succ_r(p, l)) : {s == 1n+_ : Nat}
1749 e
1750
1751# s == n + 0 -> s == n
1752def zero_out(-s: Nat, +n: Nat, e: {s == Nat.add(n, 0n) : Nat}) -> {s == n : Nat}:
1753 %add_zero_r(n) : {s == _ : Nat}
1754 e
1755
1756# reversing onto acc adds the lengths
1757def rev_len(+xs: List<&2, F32>, +acc: List<&2, F32>) -> {List.length(&2, F32, List.reverse.go(&2, F32, xs, acc)) == Nat.add(List.length(&2, F32, xs), List.length(&2, F32, acc)) : Nat}:
1758 match xs:
1759 case Nil{}:
1760 {==}
1761 case Con{h, t}:
1762 succ_out(List.length(&2, F32, List.reverse.go(&2, F32, t, h <> acc)), List.length(&2, F32, t), List.length(&2, F32, acc), rev_len(t, h <> acc))
1763
1764# appending adds the lengths
1765def app_len(+xs: List<&2, F32>, +ys: List<&2, F32>) -> {List.length(&2, F32, List.append(&2, F32, xs, ys)) == Nat.add(List.length(&2, F32, xs), List.length(&2, F32, ys)) : Nat}:
1766 match xs:
1767 case Nil{}:
1768 {==}
1769 case Con{h, t}:
1770 %app_len(t, ys) : {1n+List.length(&2, F32, List.append(&2, F32, t, ys)) == 1n+_ : Nat}
1771 {==}
1772
1773# a leaf's list: read_l reads mleft + 1 numbers onto acc and reverses them. By induction on
1774# mleft: with 0 left it reads one number (rev_len counts the reversed list); with 1 + p left it
1775# reads one and recurses, and len_step moves the 1 out of the sum.
1776law leaf_len:
1777 for -r: Nat
1778 for -c: Nat
1779 for w: Array<F32>
1780 for +mleft: Nat
1781 for +idx: U32
1782 for +acc: List<&2, F32>
1783 for rc: Array<F32> & F32
1784 {bl_len(r, c, bv_g2(r, c, w, read_l(mleft, idx, acc, rc))) == Nat.add(1n+mleft, List.length(&2, F32, acc)) : Nat}
1785
1786def leaf_len(r, c, w, mleft, idx, acc, rc):
1787 match mleft rc:
1788 case 0n Tuple{d, +v}:
1789 zero_out(List.length(&2, F32, List.reverse.go(&2, F32, v <> acc, Nil{})), 1n+List.length(&2, F32, acc), rev_len(v <> acc, Nil{}))
1790 case 1n+p Tuple{d, +v}:
1791 len_step(bl_len(r, c, bv_g2(r, c, w, read_l(p, (idx + 1 : U32), v <> acc, Array.get(F32, d, (idx + 1 : U32))))), p, List.length(&2, F32, acc), leaf_len(r, c, w, p, (idx + 1 : U32), v <> acc, Array.get(F32, d, (idx + 1 : U32))))
1792
1793# whatever the gemm wrote, reading 1 + p rows back gives 1 + p numbers (G has one constructor, so
1794# matching g opens it, and leaf_len counts the list)
1795def g_len(-r: Nat, -c: Nat, +p: Nat, g: G) -> {bl_len(r, c, bv_g(r, c, p, g)) == 1n+p : Nat}:
1796 match g:
1797 case G{a, b, y}:
1798 zero_out(bl_len(r, c, bv_g2(r, c, b, read_l(p, 0, Nil{}, Array.get(F32, y, 0)))), 1n+p, leaf_len(r, c, b, p, 0, Nil{}, Array.get(F32, y, 0)))
1799
1800# LAW (band_len): the product of one band with `rows` rows gives exactly `rows` numbers.
1801# 0 rows give the empty list; 1 + p rows go through the gemm, and g_len counts what is read back.
1802law band_len:
1803 for -r: Nat
1804 for +c: Nat
1805 for +rows: Nat
1806 for x: Array<F32>
1807 for w: Array<F32>
1808 {bl_len(r, c, bv_gemm(r, c, rows, rows, x, w)) == rows : Nat}
1809
1810def band_len(r, c, rows, x, w):
1811 match rows:
1812 case 0n:
1813 {==}
1814 case 1n+p:
1815 g_len(r, c, p, gemm(1n, 1n+p, c, U32.from_nat(c), 1, 1, U32.from_nat(c), x, w, Array.new(F32, cap_depth(1n+p), 0.0)))
1816
1817# joining two halves of h1 and h2 numbers gives h1 + h2 numbers. (The statement for the whole tree,
1818# "Bands.matvec_l gives r numbers", would apply this to the results of the two recursive calls and
1819# to the induction hypotheses about the same calls; that uses each band's Array twice, which Bend's
1820# affine rules refuse. It stays open, covered by tests: NOTES.md, v3.1.)
1821def join_len(-r: Nat, -c: Nat, -h1: Nat, -h2: Nat, p: BL<half(r), c>, q: BL<Nat.sub(r, half(r)), c>, hp: {bl_len(half(r), c, p) == h1 : Nat}, hq: {bl_len(Nat.sub(r, half(r)), c, q) == h2 : Nat}) -> {bl_len(r, c, bv_join(r, c, p, q)) == Nat.add(h1, h2) : Nat}:
1822 match p q:
1823 case BL{bx, +yx} BL{by, +yy}:
1824 %hp : {List.length(&2, F32, List.append(&2, F32, yx, yy)) == Nat.add(_, h2) : Nat}
1825 %hq : {List.length(&2, F32, List.append(&2, F32, yx, yy)) == Nat.add(List.length(&2, F32, yx), _) : Nat}
1826 app_len(yx, yy)
1827