main.bend
bend-ml-tensor-array@0.1.8.0 · 1,486 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.8.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# row i of C, columns j.. ; r is the product of the current column (already computed)
46
47def gcols(mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, +cidx: U32, c: Array<F32>, r: R) -> G:
48 match mleft r:
49 case 0n R{a, b, +x}:
50 G{a, b, Array.set(F32, c, cidx, x)}
51 case 1n+p R{a, b, +x}:
52 gcols(p, kn, ia0, (ib0 + bcol : U32), sa, sb, bcol, (cidx + 1 : U32), Array.set(F32, c, cidx, x), dotg(kn, ia0, (ib0 + bcol : U32), sa, sb, 0.0, Array.get(F32, a, ia0), Array.get(F32, b, (ib0 + bcol : U32))))
53
54def grows(nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +mu: U32, g: G) -> G:
55 match nleft g:
56 case 0n G{a, b, c}:
57 G{a, b, c}
58 case 1n+p G{a, b, c}:
59 grows(p, mm1, kn, (i + 1 : U32), arow, sa, sb, bcol, mu, gcols(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))))
60
61# C (n x m) = A · B with k steps; c is the output array (with capacity n*m)
62
63def 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:
64 grows(n, Nat.sub(m, 1n), k, 0, arow, sa, sb, bcol, U32.from_nat(m), G{a, b, c})
65
66
67# ---------------------------------------------------------------------
68# parallel gemm: 2^d blocks of rows of C, each with its own copy of A and B,
69# returning lists that are written into C at the end
70# ---------------------------------------------------------------------
71
72type S3 is Type:
73 S3{a: Array<F32>, b: Array<F32>, out: List<&2, F32>}
74
75def lcols(mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, out: List<&2, F32>, r: R) -> S3:
76 match mleft r:
77 case 0n R{a, b, +x}:
78 S3{a, b, x <> out}
79 case 1n+p R{a, b, +x}:
80 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))))
81
82def lrows(nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, st: S3) -> S3:
83 match nleft st:
84 case 0n S3{a, b, out}:
85 S3{a, b, out}
86 case 1n+p S3{a, b, out}:
87 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))))
88
89def lblock_fin(st: S3) -> List<&2, F32>:
90 match st:
91 case S3{a, b, out}:
92 List.reverse(&2, F32, out)
93
94def 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>:
95 lblock_fin(lrows(cnt, mm1, kn, i, arow, sa, sb, bcol, S3{a, b, Nil{}}))
96
97def 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>:
98 match d ra rb:
99 case 0n Tuple{a, a2} Tuple{b, b2}:
100 lblock(cnt, mm1, kn, i, arow, sa, sb, bcol, a, b)
101 case 1n++q Tuple{a, a2} Tuple{b, b2}:
102 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))
103 List.append(&2, F32, x, y)
104
105def write_l(xs: List<&2, F32>, +idx: U32, c: Array<F32>) -> Array<F32>:
106 match xs:
107 case Nil{}:
108 c
109 case Con{h, t}:
110 write_l(t, (idx + 1 : U32), Array.set(F32, c, idx, h))
111
112# c[idx..] += list
113
114def gp3(a: Array<F32>, b: Array<F32>, c: Array<F32>, xs: List<&2, F32>) -> G:
115 G{a, b, write_l(xs, 0, c)}
116
117def 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:
118 match pa pb:
119 case Tuple{a1, a2} Tuple{b1, b2}:
120 gp3(a1, b1, c, lpar(d, n, mm1, kn, 0, arow, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2)))
121
122# C (n x m) = A · B, with 2^d blocks of rows in parallel; returns A, B (intact) and C
123
124def 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:
125 gp2(d, n, Nat.sub(m, 1n), k, arow, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b))
126
127
128# ---------------------------------------------------------------------
129# parallel gemm by COLUMNS, for n = 1 (matrix · vector): each block computes a
130# range of columns of C, with its own copy of the operands
131# ---------------------------------------------------------------------
132
133def cblock_fin(st: S3) -> List<&2, F32>:
134 match st:
135 case S3{a, b, out}:
136 List.reverse(&2, F32, out)
137
138# cnt columns starting at j0 (cnt >= 1); row 0 of A is at ia0
139
140def 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>:
141 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)))))
142
143def cblock(cnt: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>) -> List<&2, F32>:
144 match cnt:
145 case 0n:
146 Nil{}
147 case 1n+p:
148 cblock_go(p, kn, ia0, j0, sa, sb, bcol, a, b)
149
150def 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>:
151 match d ra rb:
152 case 0n Tuple{a, a2} Tuple{b, b2}:
153 cblock(cnt, kn, ia0, j0, sa, sb, bcol, a, b)
154 case 1n++q Tuple{a, a2} Tuple{b, b2}:
155 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))
156 List.append(&2, F32, x, y)
157
158def cp2(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:
159 match pa pb:
160 case Tuple{a1, a2} Tuple{b1, b2}:
161 gp3(a1, b1, c, cpar(d, m, kn, 0, 0, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2)))
162
163# C (1 x m) = A (1 x k) · B, with 2^d blocks of columns in parallel
164
165def gemm_par_cols(d: Nat, +m: Nat, +k: Nat, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
166 cp2(d, m, k, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b))
167
168# ---------------------------------------------------------------------
169# element-wise operations (in place, by index)
170# ---------------------------------------------------------------------
171
172# ---------------------------------------------------------------------
173# gemm: C[i][j] = sum_p A[i*arow + p*sa] * B[p*sb + j*bcol]
174# (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)
175# ---------------------------------------------------------------------
176
177# acc += A[ia + t*sa] * B[ib + t*sb] for k steps; ra, rb already carry the first pair
178
179type P2 is Type:
180 P2{a: Array<F32>, b: Array<F32>}
181
182type AL is Type:
183 AL{c: Array<F32>, l: List<&2, F32>}
184
185# list of m elements starting at idx (reads with Array.get, returning the array)
186
187def read_l(mleft: Nat, +idx: U32, acc: List<&2, F32>, rc: Array<F32> & F32) -> AL:
188 match mleft rc:
189 case 0n Tuple{c, +v}:
190 AL{c, List.reverse(&2, F32, v <> acc)}
191 case 1n+p Tuple{c, +v}:
192 read_l(p, (idx + 1 : U32), v <> acc, Array.get(F32, c, (idx + 1 : U32)))
193
194# writes the list starting at idx
195
196def addl(xs: List<&2, F32>, +idx: U32, rc: Array<F32> & F32) -> Array<F32>:
197 match xs rc:
198 case Nil{} Tuple{c, v}:
199 c
200 case Con{+h, t} Tuple{c, +v}:
201 addl(t, (idx + 1 : U32), Array.get(F32, Array.set(F32, c, idx, (v + h : F32)), (idx + 1 : U32)))
202
203# adds a bias (a list of m) to each of the n rows of c (n x m)
204
205def bias_rows(nleft: Nat, +i: U32, +mu: U32, +bl: List<&2, F32>, c: Array<F32>) -> Array<F32>:
206 match nleft:
207 case 0n:
208 c
209 case 1n+p:
210 bias_rows(p, (i + 1 : U32), mu, bl, addl(bl, (i * mu : U32), Array.get(F32, c, (i * mu : U32))))
211
212# in-place relu on the first n elements
213
214def relu_l(nleft: Nat, +t: U32, rc: Array<F32> & F32) -> Array<F32>:
215 match nleft rc:
216 case 0n Tuple{c, v}:
217 c
218 case 1n+p Tuple{c, +v}:
219 relu_l(p, (t + 1 : U32), Array.get(F32, Array.set(F32, c, t, F32.max(v, 0.0)), (t + 1 : U32)))
220
221def pos2(b: Bool) -> F32:
222 match b:
223 case True{}:
224 1.0
225 case False{}:
226 0.0
227
228def pos(+y: F32) -> F32:
229 pos2(F32.is_gt(y, 0.0))
230
231# d[t] *= (h[t] > 0): the relu gradient; returns both arrays
232
233def mask_l(nleft: Nat, +t: U32, rd: Array<F32> & F32, rh: Array<F32> & F32) -> P2:
234 match nleft rd rh:
235 case 0n Tuple{d, x} Tuple{h, y}:
236 P2{d, h}
237 case 1n+p Tuple{d, +x} Tuple{h, +y}:
238 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)))
239
240# w[t] -= lr * dw[t]
241
242def maxl2(xs: List<&2, F32>, cur: F32) -> F32:
243 match xs:
244 case Nil{}:
245 cur
246 case Con{+h, t}:
247 maxl2(t, F32.max(cur, h))
248
249def maxl(xs: List<&2, F32>) -> F32:
250 match xs:
251 case Nil{}:
252 0.0
253 case Con{+h, t}:
254 maxl2(t, h)
255
256def expl(xs: List<&2, F32>, +m: F32) -> List<&2, F32>:
257 match xs:
258 case Nil{}:
259 Nil{}
260 case Con{+h, t}:
261 F32.exp((h - m : F32)) <> expl(t, m)
262
263def suml(xs: List<&2, F32>, acc: F32) -> F32:
264 match xs:
265 case Nil{}:
266 acc
267 case Con{+h, t}:
268 suml(t, (acc + h : F32))
269
270def divl(xs: List<&2, F32>, +s: F32) -> List<&2, F32>:
271 match xs:
272 case Nil{}:
273 Nil{}
274 case Con{+h, t}:
275 (h / s : F32) <> divl(t, s)
276
277def softmax2(+es: List<&2, F32>) -> List<&2, F32>:
278 divl(es, suml(es, 0.0))
279
280def softmax(+xs: List<&2, F32>) -> List<&2, F32>:
281 softmax2(expl(xs, maxl(xs)))
282
283# p[label]
284
285def sgd_l(nleft: Nat, +t: U32, +lr: F32, rw: Array<F32> & F32, rd: Array<F32> & F32) -> P2:
286 match nleft rw rd:
287 case 0n Tuple{w, x} Tuple{d, y}:
288 P2{w, d}
289 case 1n+p Tuple{w, +x} Tuple{d, +y}:
290 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)))
291
292# ---------------------------------------------------------------------
293# softmax + cross-entropy per row (lists of 10 elements: negligible cost)
294# ---------------------------------------------------------------------
295
296def pick(ps: List<&2, F32>, n: Nat) -> F32:
297 match ps n:
298 case Con{+h, t} 0n:
299 h
300 case Con{h, t} 1n+q:
301 pick(t, q)
302 case Nil{} _:
303 0.0
304
305# (p - onehot(label)) / n
306
307def hit(b: Bool) -> F32:
308 match b:
309 case True{}:
310 1.0
311 case False{}:
312 0.0
313
314def grad_l(ps: List<&2, F32>, +label: Nat, +j: Nat, +nf: F32) -> List<&2, F32>:
315 match ps:
316 case Nil{}:
317 Nil{}
318 case Con{+h, t}:
319 (((h - hit(Nat.is_eq(label, j)) : F32) / nf) : F32) <> grad_l(t, label, 1n+j, nf)
320
321type L2 is Type:
322 L2{loss: F32, c: Array<F32>}
323
324def ce_row2(+ps: List<&2, F32>, +label: Nat, +nf: F32, +idx: U32, loss: F32, c: Array<F32>) -> L2:
325 L2{(loss - F32.log(F32.max(pick(ps, label), 0.000000000001)) : F32), write_l(grad_l(ps, label, 0n, nf), idx, c)}
326
327def ce_row(al: AL, +label: Nat, +nf: F32, +idx: U32, loss: F32) -> L2:
328 match al:
329 case AL{c, l}:
330 ce_row2(softmax(l), label, nf, idx, loss, c)
331
332# 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
333
334type Mat<-r: Nat, -c: Nat> is Type:
335 Mat{d: Array<F32>}
336
337def ce_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +nf: F32, +i: U32, r: L2) -> L2:
338 match labels r:
339 case Nil{} L2{loss, c}:
340 L2{loss, c}
341 case Con{+y, t} L2{loss, c}:
342 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))
343
344# ---- helper: hits per row ----
345
346type CR is Type:
347 CR{c: Array<F32>, n: Nat}
348
349def gt_head(xs: List<&2, F32>, +v: F32) -> Bool:
350 match xs:
351 case Nil{}:
352 False{}
353 case Con{+h, t}:
354 F32.is_gt(h, v)
355
356def argmax2(xs: List<&2, F32>, g: Bool, +i: Nat, +bi: Nat, +bv: F32) -> Nat:
357 match xs g:
358 case Nil{} _:
359 bi
360 case Con{+h, +t} True{}:
361 argmax2(t, gt_head(t, h), 1n+i, i, h)
362 case Con{h, +t} False{}:
363 argmax2(t, gt_head(t, bv), 1n+i, bi, bv)
364
365def argmax(xs: List<&2, F32>) -> Nat:
366 match xs:
367 case Nil{}:
368 0n
369 case Con{+h, +t}:
370 argmax2(t, gt_head(t, h), 1n, 0n, h)
371
372def one(b: Bool) -> Nat:
373 match b:
374 case True{}:
375 1n
376 case False{}:
377 0n
378
379def count_rows3(+y: Nat, n: Nat, al: AL) -> CR:
380 match al:
381 case AL{c, l}:
382 CR{c, Nat.add(n, one(Nat.is_eq(argmax(l), y)))}
383
384def count_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +i: U32, r: CR) -> CR:
385 match labels r:
386 case Nil{} CR{c, n}:
387 CR{c, n}
388 case Con{+y, t} CR{c, n}:
389 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)))))
390
391
392# =====================================================================
393# Typed API: Mat<r, c> (the dimensions enter the type)
394# =====================================================================
395
396# =====================================================================
397# Array capacity: a PROVED depth for the number of slots
398# =====================================================================
399#
400# cap_depth(n) is the d such that an Array with 2^d slots holds n elements. It is computed by
401# structural recursion on n (not with U32.log2) so that the law cap_ok below can prove
402# `n <= 2^cap_depth(n)`. What stays trusted is only that Array.new(T, d, v) allocates 2^d slots.
403
404# a type chosen by a Bool: rewriting through it refutes True == False
405def BD(b: Bool, t: Type, f: Type) -> Type:
406 match b:
407 case True{}:
408 t
409 case False{}:
410 f
411
412# a <= b, by direct recursion
413def leb(a: Nat, b: Nat) -> Bool:
414 match a b:
415 case 0n _:
416 True{}
417 case 1n+p 0n:
418 False{}
419 case 1n+p 1n+q:
420 leb(p, q)
421
422def pow2(d: Nat) -> Nat:
423 match d:
424 case 0n:
425 1n
426 case 1n++q:
427 Nat.add(pow2(q), pow2(q))
428
429# smallest power of two >= n, by structural recursion on n: returns (depth, 2^depth)
430def cap_pick(ok: Bool, d: Nat, +p: Nat) -> Nat & Nat:
431 match ok:
432 case True{}:
433 (d, p)
434 case False{}:
435 (1n+d, Nat.add(p, p))
436
437def cap_step(r: Nat & Nat, +n: Nat) -> Nat & Nat:
438 match r:
439 case Tuple{+d, +p}:
440 cap_pick(leb(n, p), d, p)
441
442def cap_pair(n: Nat) -> Nat & Nat:
443 match n:
444 case 0n:
445 (0n, 1n)
446 case 1n++m:
447 cap_step(cap_pair(m), 1n+m)
448
449def fst_nat(r: Nat & Nat) -> Nat:
450 match r:
451 case Tuple{d, p}:
452 d
453
454def snd_nat(r: Nat & Nat) -> Nat:
455 match r:
456 case Tuple{d, p}:
457 p
458
459# the depth d such that an Array with 2^d slots holds n elements
460def cap_depth(n: Nat) -> Nat:
461 fst_nat(cap_pair(n))
462
463# ---- lemmas ----
464
465# a <= b -> a <= 1 + b
466def leb_succ(a: Nat, b: Nat, e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, 1n+b) : Bool}:
467 match a b:
468 case 0n _:
469 {==}
470 case 1n+p 0n:
471 %e : BD(_, Unit, {True{} == leb(1n+p, 1n) : Bool})
472 Unit{}
473 case 1n+p 1n+q:
474 leb_succ(p, q, e)
475
476# a <= b -> a <= c + b
477def leb_up_l(+a: Nat, +b: Nat, c: Nat, +e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, Nat.add(c, b)) : Bool}:
478 match c:
479 case 0n:
480 e
481 case 1n++c2:
482 leb_succ(a, Nat.add(c2, b), leb_up_l(a, b, c2, e))
483
484# a <= b -> a <= b + c
485def leb_up(a: Nat, b: Nat, c: Nat, e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, Nat.add(b, c)) : Bool}:
486 match a b:
487 case 0n _:
488 {==}
489 case 1n+p 0n:
490 %e : BD(_, Unit, {True{} == leb(1n+p, Nat.add(0n, c)) : Bool})
491 Unit{}
492 case 1n+p 1n+q:
493 leb_up(p, q, c, e)
494
495# a <= b and c <= d -> a + c <= b + d
496def 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}:
497 match a b:
498 case 0n _:
499 leb_up_l(c, d, b, e2)
500 case 1n+p 0n:
501 %e1 : BD(_, Unit, {True{} == leb(Nat.add(1n+p, c), Nat.add(0n, d)) : Bool})
502 Unit{}
503 case 1n+p 1n+q:
504 leb_add(p, q, c, d, e1, e2)
505
506# 1 <= 2^d
507def pow2_pos(d: Nat) -> {True{} == leb(1n, pow2(d)) : Bool}:
508 match d:
509 case 0n:
510 {==}
511 case 1n++q:
512 leb_up(1n, pow2(q), pow2(q), pow2_pos(q))
513
514# 1 <= p when p = 2^d
515def pos_p(d: Nat, p: Nat, i1: {p == pow2(d) : Nat}) -> {True{} == leb(1n, p) : Bool}:
516 %Equal.sym(Nat, p, pow2(d), i1) : {True{} == leb(1n, _) : Bool}
517 pow2_pos(d)
518
519# p + p == 2^(1+d) when p == 2^d
520def step_eq(d: Nat, p: Nat, i1: {p == pow2(d) : Nat}) -> {Nat.add(p, p) == pow2(1n+d) : Nat}:
521 %i1 : {Nat.add(p, p) == Nat.add(_, _) : Nat}
522 {==}
523
524# one step: given the invariant for m, the invariant for 1+m holds whichever way the comparison goes
525def 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}:
526 match ok:
527 case True{}:
528 (i1, e)
529 case False{}:
530 (step_eq(d, p, i1), leb_add(1n, p, m, p, pos_p(d, p, i1), i2))
531
532
533# the invariant for cap_step on a pair, opened
534def 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}:
535 match r:
536 case Tuple{+d, +p}:
537 step_ok(leb(1n+m, p), d, p, m, {==}, i1, i2)
538
539def 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}:
540 match ih:
541 case Tuple{i1, i2}:
542 step_pair_ok(m, cap_pair(m), i1, i2)
543
544# the invariant of cap_pair: the second component is 2^(first) and is >= n
545def 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}:
546 match n:
547 case 0n:
548 ({==}, {==})
549 case 1n++m:
550 cap_succ(m, cap_inv(m))
551
552def cap_final(n: Nat, pr: {snd_nat(cap_pair(n)) == pow2(fst_nat(cap_pair(n))) : Nat} & {True{} == leb(n, snd_nat(cap_pair(n))) : Bool}) -> {True{} == leb(n, pow2(cap_depth(n))) : Bool}:
553 match pr:
554 case Tuple{i1, i2}:
555 %i1 : {True{} == leb(n, _) : Bool}
556 i2
557
558# LAW (cap_ok): an Array with 2^cap_depth(n) slots always has room for n elements.
559law cap_ok:
560 for +n: Nat
561 {True{} == leb(n, pow2(cap_depth(n))) : Bool}
562
563def cap_ok(n):
564 cap_final(n, cap_inv(n))
565
566# r x c of zeros, with the exact capacity (smallest power of 2)
567
568def Mat.zeros(+r: Nat, +c: Nat) -> Mat<r, c>:
569 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0)}
570
571def Mat.fill(+r: Nat, +c: Nat, +v: F32) -> Mat<r, c>:
572 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), v)}
573
574def fill_list(xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
575 match xs:
576 case Nil{}:
577 a
578 case Con{h, t}:
579 fill_list(t, (i + 1 : U32), Array.set(F32, a, i, h))
580
581def len_is(xs: List<&2, F32>, n: Nat) -> Bool:
582 match xs n:
583 case Nil{} 0n:
584 True{}
585 case Con{h, t} 1n+p:
586 len_is(t, p)
587 case _ _:
588 False{}
589
590def Mat.of_list2(-r: Nat, -c: Nat, +rr: Nat, +cc: Nat, ok: Bool, xs: List<&2, F32>) -> Maybe<&1, Mat<r, c>>:
591 match ok:
592 case True{}:
593 Some{Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(rr, cc)), 0.0))}}
594 case False{}:
595 None{}
596
597# Mat<r, c> from r*c numbers in row order; None if the size does not match
598
599def Mat.of_list(+r: Nat, +c: Nat, +xs: List<&2, F32>) -> Maybe<&1, Mat<r, c>>:
600 Mat.of_list2(r, c, r, c, len_is(xs, Nat.mul(r, c)), xs)
601
602# reads the first n numbers, returning the matrix
603
604# Mat<r, c> without checking the list size (a short list leaves the rest at zero, a long one is cut
605# by the capacity): only for when the size is known by construction. Prefer Mat.of_list.
606
607def Mat.from_list(+r: Nat, +c: Nat, xs: List<&2, F32>) -> Mat<r, c>:
608 Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0))}
609
610# writes the numbers of xs into m starting at flat index i (row-major); numbers past the
611# capacity are dropped. For loading a large matrix block by block without building one big list.
612def Mat.fill_at(-r: Nat, -c: Nat, +i: U32, xs: List<&2, F32>, m: Mat<r, c>) -> Mat<r, c>:
613 match m:
614 case Mat{d}:
615 Mat{fill_list(xs, i, d)}
616
617type ML<-r: Nat, -c: Nat> is Type:
618 ML{m: Mat<r, c>, l: List<&2, F32>}
619
620def Mat.to_list2(-r: Nat, -c: Nat, al: AL) -> ML<r, c>:
621 match al:
622 case AL{d, l}:
623 ML{Mat{d}, l}
624
625# two independent copies
626
627def Mat.to_list(+r: Nat, +c: Nat, m: Mat<r, c>) -> ML<r, c>:
628 match m:
629 case Mat{d}:
630 Mat.to_list2(r, c, read_l(Nat.sub(Nat.mul(r, c), 1n), 0, Nil{}, Array.get(F32, d, 0)))
631
632type MM2<-r: Nat, -c: Nat> is Type:
633 MM2{a: Mat<r, c>, b: Mat<r, c>}
634
635def Mat.clone2(-r: Nat, -c: Nat, p: Array<F32> & Array<F32>) -> MM2<r, c>:
636 match p:
637 case Tuple{x, y}:
638 MM2{Mat{x}, Mat{y}}
639
640# ---- products: C (n x m) = A · B, with 2^par blocks of rows in parallel ----
641
642# C = A(n x k) · B(k x m)
643
644def Mat.clone(-r: Nat, -c: Nat, m: Mat<r, c>) -> MM2<r, c>:
645 match m:
646 case Mat{d}:
647 Mat.clone2(r, c, Array.clone(F32, d))
648
649type MMul<-n: Nat, -k: Nat, -m: Nat> is Type:
650 MMul{a: Mat<n, k>, b: Mat<k, m>, c: Mat<n, m>}
651
652# C = A(n x k) · Bᵀ, with B stored m x k
653
654type MMulNT<-n: Nat, -k: Nat, -m: Nat> is Type:
655 MMulNT{a: Mat<n, k>, b: Mat<m, k>, c: Mat<n, m>}
656
657# C = Aᵀ · B, with A stored k x n
658
659type MMulTN<-n: Nat, -k: Nat, -m: Nat> is Type:
660 MMulTN{a: Mat<k, n>, b: Mat<k, m>, c: Mat<n, m>}
661
662# ---------------------------------------------------------------------
663# 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
664# 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
665# 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,
666# in order, and the four sums are combined as dotr combines them after k steps (its rotation depends on k mod 4). The
667# padding adds +0 to a sum that starts at +0, which changes nothing, so the results are the same bits as the scalar
668# loop's. Products with one row (matrix · vector) keep the scalar loop: packing B would cost as much as the product.
669# ---------------------------------------------------------------------
670
671type T8 is Data:
672 T8{v0: F32, v1: F32, v2: F32, v3: F32, v4: F32, v5: F32, v6: F32, v7: F32}
673
674def t8_zero() -> T8:
675 T8{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}
676
677def pick0(b: Bool, +v: F32) -> F32:
678 match b:
679 case True{}:
680 v
681 case False{}:
682 0.0
683
684type PK is Type:
685 PK{a: Array<F32>, t: Array<T8>}
686
687# 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
688def 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:
689 match r:
690 case Tuple{a, +v}:
691 PK{a, Array.set(T8, t, o, T8{v0, v1, v2, v3, v4, v5, v6, pick0(U32.is_lt((q0 + 7 : U32), lim), v)})}
692
693def 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:
694 match r:
695 case Tuple{a, +v}:
696 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)))
697
698def 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:
699 match r:
700 case Tuple{a, +v}:
701 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)))
702
703def 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:
704 match r:
705 case Tuple{a, +v}:
706 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)))
707
708def 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:
709 match r:
710 case Tuple{a, +v}:
711 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)))
712
713def pk_g2(+v0: F32, +v1: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
714 match r:
715 case Tuple{a, +v}:
716 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)))
717
718def pk_g1(+v0: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
719 match r:
720 case Tuple{a, +v}:
721 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)))
722
723def pk_g0(+q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array<T8>, r: Array<F32> & F32) -> PK:
724 match r:
725 case Tuple{a, +v}:
726 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)))
727
728# the cells of one line, cell o.. of t; q0 counts the line's numbers before the cell
729def pk_line(tl: Nat, +q0: U32, +lim: U32, +base: U32, +st: U32, +o: U32, r: PK) -> PK:
730 match tl r:
731 case 0n r:
732 r
733 case 1n+q PK{a, t}:
734 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)))
735
736# lines count lines; line i starts at i * lstep (in a), its cells at i * tku (in t)
737def pk_lines(cnt: Nat, +tk: Nat, +tku: U32, +ku: U32, +i: U32, +lstep: U32, +st: U32, r: PK) -> PK:
738 match cnt:
739 case 0n:
740 r
741 case 1n+p:
742 pk_lines(p, tk, tku, ku, (i + 1 : U32), lstep, st, pk_line(tk, 0, ku, (i * lstep : U32), st, (i * tku : U32), r))
743
744def pk_pack(+cnt: Nat, +tk: Nat, +k: Nat, +lstep: U32, +st: U32, a: Array<F32>) -> PK:
745 pk_lines(cnt, tk, U32.from_nat(tk), U32.from_nat(k), 0, lstep, st, PK{a, Array.new(T8, cap_depth(Nat.mul(cnt, tk)), t8_zero())})
746
747type TD is Type:
748 TD{a: Array<T8>, b: Array<T8>, s0: F32, s1: F32, s2: F32, s3: F32}
749
750# product q of a cell goes to sum q mod 4, the first four before the last four, as dotr adds them
751def t8_step(s0: F32, s1: F32, s2: F32, s3: F32, ra: Array<T8> & T8, rb: Array<T8> & T8) -> TD:
752 match ra rb:
753 case Tuple{a2, T8{x0, x1, x2, x3, x4, x5, x6, x7}} Tuple{b2, T8{y0, y1, y2, y3, y4, y5, y6, y7}}:
754 TD{a2, b2, ((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)}
755
756def t8_dot(tk: Nat, +ia: U32, +ib: U32, d: TD) -> TD:
757 match tk d:
758 case 0n d:
759 d
760 case 1n+p TD{a, b, s0, s1, s2, s3}:
761 t8_dot(p, (ia + 1 : U32), (ib + 1 : U32), t8_step(s0, s1, s2, s3, Array.get(T8, a, ia), Array.get(T8, b, ib)))
762
763# dotr's final sum after k products: the sums sit rotated by k mod 4 when its loop ends
764def t8_comb(+rot: Nat, +s0: F32, +s1: F32, +s2: F32, +s3: F32) -> F32:
765 match rot:
766 case 0n:
767 ((s0 + s1 : F32) + (s2 + s3 : F32) : F32)
768 case 1n:
769 ((s1 + s2 : F32) + (s3 + s0 : F32) : F32)
770 case 2n:
771 ((s2 + s3 : F32) + (s0 + s1 : F32) : F32)
772 case _:
773 ((s3 + s0 : F32) + (s1 + s2 : F32) : F32)
774
775type TR is Type:
776 TR{a: Array<T8>, b: Array<T8>, x: F32}
777
778def t8_fin(+rot: Nat, d: TD) -> TR:
779 match d:
780 case TD{a, b, +s0, +s1, +s2, +s3}:
781 TR{a, b, t8_comb(rot, s0, s1, s2, s3)}
782
783def t8_cell(+tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, a: Array<T8>, b: Array<T8>) -> TR:
784 t8_fin(rot, t8_dot(tk, (i * tku : U32), (j * tku : U32), TD{a, b, 0.0, 0.0, 0.0, 0.0}))
785
786type TG is Type:
787 TG{a: Array<T8>, b: Array<T8>, c: Array<F32>}
788
789def t8_set(+idx: U32, c: Array<F32>, r: TR) -> TG:
790 match r:
791 case TR{a, b, x}:
792 TG{a, b, Array.set(F32, c, idx, x)}
793
794def t8_cols(mleft: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, +cidx: U32, g: TG) -> TG:
795 match mleft g:
796 case 0n g:
797 g
798 case 1n+p TG{a, b, c}:
799 t8_cols(p, tk, tku, rot, i, (j + 1 : U32), (cidx + 1 : U32), t8_set(cidx, c, t8_cell(tk, tku, rot, i, j, a, b)))
800
801def t8_rows(nleft: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +mu: U32, g: TG) -> TG:
802 match nleft:
803 case 0n:
804 g
805 case 1n+p:
806 t8_rows(p, m, tk, tku, rot, (i + 1 : U32), mu, t8_cols(m, tk, tku, rot, i, 0, (i * mu : U32), g))
807
808# rows i0 .. i0+cnt-1 of C as a list, in order (the parallel version's leaves)
809type TL is Type:
810 TL{a: Array<T8>, b: Array<T8>, out: List<&2, F32>}
811
812# cells j.. of row i, m - j of them; r is cell j, already computed
813def t8_lcols(mleft: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, out: List<&2, F32>, r: TR) -> TL:
814 match mleft r:
815 case 0n TR{a, b, +x}:
816 TL{a, b, x <> out}
817 case 1n+p TR{a, b, +x}:
818 t8_lcols(p, tk, tku, rot, i, (j + 1 : U32), x <> out, t8_cell(tk, tku, rot, i, (j + 1 : U32), a, b))
819
820def t8_lrows(nleft: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, l: TL) -> TL:
821 match nleft l:
822 case 0n l:
823 l
824 case 1n+p TL{a, b, out}:
825 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)))
826
827def t8_lfin(l: TL) -> List<&2, F32>:
828 match l:
829 case TL{a, b, out}:
830 List.reverse(&2, F32, out)
831
832def t8_lpar(d: Nat, +cnt: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, ra: Array<T8> & Array<T8>, rb: Array<T8> & Array<T8>) -> List<&2, F32>:
833 match d ra rb:
834 case 0n Tuple{a, a2} Tuple{b, b2}:
835 t8_lfin(t8_lrows(cnt, m, tk, tku, rot, i, TL{a, b, Nil{}}))
836 case 1n++q Tuple{a, a2} Tuple{b, b2}:
837 x y = t8_lpar(q, Nat.div(cnt, 2n), m, tk, tku, rot, i, Array.clone(T8, a), Array.clone(T8, b)) t8_lpar(q, Nat.sub(cnt, Nat.div(cnt, 2n)), m, tk, tku, rot, (i + U32.from_nat(Nat.div(cnt, 2n)) : U32), Array.clone(T8, a2), Array.clone(T8, b2))
838 List.append(&2, F32, x, y)
839
840type PKK is Type:
841 PKK{a: Array<F32>, b: Array<F32>, ta: Array<T8>, tb: Array<T8>}
842
843def t8_pk2(pa: PK, pb: PK) -> PKK:
844 match pa pb:
845 case PK{a, ta} PK{b, tb}:
846 PKK{a, b, ta, tb}
847
848def t8_g(a: Array<F32>, b: Array<F32>, g: TG) -> G:
849 match g:
850 case TG{ta, tb, c}:
851 G{a, b, c}
852
853def t8_out_seq(+n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array<F32>, p: PKK) -> G:
854 match p:
855 case PKK{a, b, ta, tb}:
856 t8_g(a, b, t8_rows(n, m, tk, tku, rot, 0, U32.from_nat(m), TG{ta, tb, c}))
857
858def t8_out_par(d: Nat, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array<F32>, p: PKK) -> G:
859 match p:
860 case PKK{a, b, ta, tb}:
861 G{a, b, write_l(t8_lpar(d, n, m, tk, tku, rot, 0, Array.clone(T8, ta), Array.clone(T8, tb)), 0, c)}
862
863def gemm_t8b(d: Nat, +n: Nat, +m: Nat, +k: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array<F32>, p: PKK) -> G:
864 match d:
865 case 0n:
866 t8_out_seq(n, m, tk, tku, rot, c, p)
867 case 1n+q:
868 t8_out_par(1n+q, n, m, tk, tku, rot, c, p)
869
870# C (n x m) = A · B in cells of 8, sequential (d = 0) or with 2^d blocks of rows in parallel; A and B come back intact
871def 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:
872 gemm_t8b(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), c, t8_pk2(pk_pack(n, Nat.div(Nat.add(k, 7n), 8n), k, arow, sa, a), pk_pack(m, Nat.div(Nat.add(k, 7n), 8n), k, bcol, sb, b)))
873
874# one row (matrix · vector): the scalar loop, by columns in parallel when d > 0; two or more rows: cells of 8
875def gemm_any3(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:
876 match one d:
877 case True{} 0n:
878 gemm(n, m, k, arow, sa, sb, bcol, a, b, c)
879 case True{} 1n+q:
880 gemm_par_cols(1n+q, m, k, sa, sb, bcol, a, b, c)
881 case False{} d:
882 gemm_t8(d, n, m, k, arow, sa, sb, bcol, a, b, c)
883
884def gemm_any2(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:
885 gemm_any3(Nat.is_le(n, 1n), d, n, m, k, arow, sa, sb, bcol, a, b, c)
886
887def gemm_any(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:
888 gemm_any2(d, n, m, k, arow, sa, sb, bcol, a, b, c)
889
890def mmul_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMul<n, k, m>:
891 match g:
892 case G{a, b, c}:
893 MMul{Mat{a}, Mat{b}, Mat{c}}
894
895def mmul_nt_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulNT<n, k, m>:
896 match g:
897 case G{a, b, c}:
898 MMulNT{Mat{a}, Mat{b}, Mat{c}}
899
900def mmul_tn_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulTN<n, k, m>:
901 match g:
902 case G{a, b, c}:
903 MMulTN{Mat{a}, Mat{b}, Mat{c}}
904
905def mm_a(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMul<n, k, m>:
906 mmul_fin(n, k, m, gemm_any(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)))
907
908# C = A · B (n x k) · (k x m); par = log2 of the number of parallel blocks (0 = sequential)
909
910def Mat.matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<k, m>) -> MMul<n, k, m>:
911 match a b:
912 case Mat{x} Mat{y}:
913 mm_a(n, k, m, par, x, y)
914
915def mm_nt(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMulNT<n, k, m>:
916 mmul_nt_fin(n, k, m, gemm_any(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)))
917
918# C = A · Bᵀ (n x k) · (m x k)ᵀ: for weights stored with one row per output
919
920def Mat.matmul_nt(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<m, k>) -> MMulNT<n, k, m>:
921 match a b:
922 case Mat{x} Mat{y}:
923 mm_nt(n, k, m, par, x, y)
924
925def mm_tn(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMulTN<n, k, m>:
926 mmul_tn_fin(n, k, m, gemm_any(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)))
927
928# C = Aᵀ · B (k x n)ᵀ · (k x m): weight gradient, X^T · dY
929
930def Mat.matmul_tn(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<k, n>, b: Mat<k, m>) -> MMulTN<n, k, m>:
931 match a b:
932 case Mat{x} Mat{y}:
933 mm_tn(n, k, m, par, x, y)
934
935# ---- element-wise and training operations ----
936
937# adds the bias (one row, m numbers) to each row of C (n x m)
938
939type MBias<-n: Nat, -m: Nat> is Type:
940 MBias{c: Mat<n, m>, b: Mat<1n, m>}
941
942def add_row2(-n: Nat, -m: Nat, +nn: Nat, +mm: Nat, cd: Array<F32>, al: AL) -> MBias<n, m>:
943 match al:
944 case AL{bd, bl}:
945 MBias{Mat{bias_rows(nn, 0, U32.from_nat(mm), bl, cd)}, Mat{bd}}
946
947def Mat.add_row(+n: Nat, +m: Nat, c: Mat<n, m>, b: Mat<1n, m>) -> MBias<n, m>:
948 match c b:
949 case Mat{cd} Mat{bd}:
950 add_row2(n, m, n, m, cd, read_l(Nat.sub(m, 1n), 0, Nil{}, Array.get(F32, bd, 0)))
951
952# pair of matrices of the same shape (the result of an operation that reads two operands)
953
954type MPair<-r: Nat, -c: Nat> is Type:
955 MPair{a: Mat<r, c>, b: Mat<r, c>}
956
957def mpair_fin(-r: Nat, -c: Nat, p: P2) -> MPair<r, c>:
958 match p:
959 case P2{a, b}:
960 MPair{Mat{a}, Mat{b}}
961
962def relu_fin(-r: Nat, -c: Nat, d: Array<F32>) -> Mat<r, c>:
963 Mat{d}
964
965def Mat.relu(+r: Nat, +c: Nat, a: Mat<r, c>) -> Mat<r, c>:
966 match a:
967 case Mat{d}:
968 relu_fin(r, c, relu_l(Nat.mul(r, c), 0, Array.get(F32, d, 0)))
969
970# (dy * (h > 0), h): the gradient that passes through the relu
971
972def Mat.relu_bwd(+r: Nat, +c: Nat, dy: Mat<r, c>, h: Mat<r, c>) -> MPair<r, c>:
973 match dy h:
974 case Mat{d} Mat{x}:
975 mpair_fin(r, c, mask_l(Nat.mul(r, c), 0, Array.get(F32, d, 0), Array.get(F32, x, 0)))
976
977# (w - lr * dw, dw): one SGD step
978
979def Mat.sgd(+r: Nat, +c: Nat, +lr: F32, w: Mat<r, c>, dw: Mat<r, c>) -> MPair<r, c>:
980 match w dw:
981 case Mat{x} Mat{y}:
982 mpair_fin(r, c, sgd_l(Nat.mul(r, c), 0, lr, Array.get(F32, x, 0), Array.get(F32, y, 0)))
983
984# (a + b, b): element-wise sum (a - (-1)*b is exactly a + b)
985
986def Mat.add(+r: Nat, +c: Nat, a: Mat<r, c>, b: Mat<r, c>) -> MPair<r, c>:
987 match a b:
988 case Mat{x} Mat{y}:
989 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)))
990
991# sum of each column (n x m -> 1 x m), returning the matrix
992
993type MColSum<-n: Nat, -m: Nat> is Type:
994 MColSum{a: Mat<n, m>, s: Mat<1n, m>}
995
996def colsum_fin(-n: Nat, -m: Nat, g: G) -> MColSum<n, m>:
997 match g:
998 case G{ones, a, s}:
999 MColSum{Mat{a}, Mat{s}}
1000
1001def Mat.col_sums(+n: Nat, +m: Nat, a: Mat<n, m>) -> MColSum<n, m>:
1002 match a:
1003 case Mat{d}:
1004 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)))
1005
1006# mean cross-entropy and its gradient on the logits: (softmax - one-hot) / n.
1007# labels has one entry per row (n); the gradient replaces the logits.
1008
1009type MCE<-n: Nat, -c: Nat> is Type:
1010 MCE{loss: F32, grad: Mat<n, c>}
1011
1012def ce_fin(-n: Nat, -c: Nat, +nf: F32, r: L2) -> MCE<n, c>:
1013 match r:
1014 case L2{loss, d}:
1015 MCE{(loss / nf : F32), Mat{d}}
1016
1017def Mat.softmax_ce(+n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> MCE<n, c>:
1018 match z:
1019 case Mat{d}:
1020 ce_fin(n, c, F32.from_nat(n), ce_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), F32.from_nat(n), 0, L2{0.0, d}))
1021
1022# argmax hits per row against labels
1023
1024type MHits<-n: Nat, -c: Nat> is Type:
1025 MHits{z: Mat<n, c>, hits: Nat}
1026
1027def hits_fin(-n: Nat, -c: Nat, r: CR) -> MHits<n, c>:
1028 match r:
1029 case CR{d, k}:
1030 MHits{Mat{d}, k}
1031
1032def Mat.count_correct(+n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> MHits<n, c>:
1033 match z:
1034 case Mat{d}:
1035 hits_fin(n, c, count_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), 0, CR{d, 0n}))
1036
1037# a row as a list (returning the matrix) and the other way around
1038
1039type MRow<-r: Nat, -c: Nat> is Type:
1040 MRow{m: Mat<r, c>, row: List<&2, F32>}
1041
1042def row_fin(-r: Nat, -c: Nat, al: AL) -> MRow<r, c>:
1043 match al:
1044 case AL{d, l}:
1045 MRow{Mat{d}, l}
1046
1047def Mat.read_row(+r: Nat, +c: Nat, m: Mat<r, c>, +i: U32) -> MRow<r, c>:
1048 match m:
1049 case Mat{d}:
1050 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))))
1051
1052def Mat.write_row(-r: Nat, +c: Nat, m: Mat<r, c>, +i: U32, xs: List<&2, F32>) -> Mat<r, c>:
1053 match m:
1054 case Mat{d}:
1055 Mat{write_l(xs, (i * U32.from_nat(c) : U32), d)}
1056
1057# ---- variants that check the number of labels ----
1058
1059# does the list of labels have exactly n entries (one per row)?
1060def labels_len_is(xs: List<&2, Nat>, n: Nat) -> Bool:
1061 match xs n:
1062 case Nil{} 0n:
1063 True{}
1064 case Con{h, t} 1n+p:
1065 labels_len_is(t, p)
1066 case _ _:
1067 False{}
1068
1069def ce_checked(ok: Bool, +n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> Maybe<&1, MCE<n, c>>:
1070 match ok:
1071 case True{}:
1072 Some{Mat.softmax_ce(n, c, z, labels)}
1073 case False{}:
1074 None{}
1075
1076# like Mat.softmax_ce, but returns None unless labels has exactly n entries
1077def Mat.softmax_ce_checked(+n: Nat, +c: Nat, z: Mat<n, c>, +labels: List<&2, Nat>) -> Maybe<&1, MCE<n, c>>:
1078 ce_checked(labels_len_is(labels, n), n, c, z, labels)
1079
1080def hits_checked(ok: Bool, +n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> Maybe<&1, MHits<n, c>>:
1081 match ok:
1082 case True{}:
1083 Some{Mat.count_correct(n, c, z, labels)}
1084 case False{}:
1085 None{}
1086
1087# like Mat.count_correct, but returns None unless labels has exactly n entries
1088def Mat.count_correct_checked(+n: Nat, +c: Nat, z: Mat<n, c>, +labels: List<&2, Nat>) -> Maybe<&1, MHits<n, c>>:
1089 hits_checked(labels_len_is(labels, n), n, c, z, labels)
1090
1091# ---------------------------------------------------------------------
1092# Bands<r, c>: a matrix cut into bands of rows, for matrix · vector in parallel
1093# without copying the weights.
1094#
1095# Mat.matmul_nt with n = 1 can only run in parallel by cloning the whole weight
1096# matrix for every task, and in matrix · vector each weight is read once, so the
1097# copies cost more than the products (NOTES.md, exp. 9). Bands keeps each band in
1098# its own Array: a task takes its band and nothing is copied except the input vector.
1099#
1100# A node with r rows has children with half(r) and r - half(r) rows. The split is
1101# part of the type, so the row count of every band follows from r, and the law
1102# half_cover says that the two halves always cover the r rows exactly.
1103# ---------------------------------------------------------------------
1104
1105# n / 2 rounded down. A tail loop with an accumulator: a non-tail recursion here would put a
1106# continuation on every node of the band tree and keep the bands from running in parallel.
1107def half.go(n: Nat, acc: Nat) -> Nat:
1108 match n:
1109 case 0n:
1110 acc
1111 case 1n+p:
1112 match p:
1113 case 0n:
1114 acc
1115 case 1n+q:
1116 half.go(q, 1n+acc)
1117
1118def half(n: Nat) -> Nat:
1119 half.go(n, 0n)
1120
1121# a <= a
1122def leb_refl(+a: Nat) -> {True{} == leb(a, a) : Bool}:
1123 match a:
1124 case 0n:
1125 {==}
1126 case 1n+p:
1127 leb_refl(p)
1128
1129# 1 + (a + b) == a + (1 + b)
1130def add_succ_r(+a: Nat, -b: Nat) -> {1n+Nat.add(a, b) == Nat.add(a, 1n+b) : Nat}:
1131 match a:
1132 case 0n:
1133 {==}
1134 case 1n+p:
1135 %add_succ_r(p, b) : {2n+Nat.add(p, b) == 1n+_ : Nat}
1136 {==}
1137
1138# a + 0 == a
1139def add_zero_r(+a: Nat) -> {Nat.add(a, 0n) == a : Nat}:
1140 match a:
1141 case 0n:
1142 {==}
1143 case 1n+p:
1144 %add_zero_r(p) : {1n+Nat.add(p, 0n) == 1n+_ : Nat}
1145 {==}
1146
1147# half.go(n, acc) <= n + acc
1148def half_go_le(+n: Nat, +acc: Nat) -> {True{} == leb(half.go(n, acc), Nat.add(n, acc)) : Bool}:
1149 match n:
1150 case 0n:
1151 leb_refl(acc)
1152 case 1n+p:
1153 match p:
1154 case 0n:
1155 leb_succ(acc, acc, leb_refl(acc))
1156 case 1n+q:
1157 %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}
1158 leb_succ(half.go(q, 1n+acc), Nat.add(q, 1n+acc), half_go_le(q, 1n+acc))
1159
1160# half(n) <= n
1161def half_le(+n: Nat) -> {True{} == leb(half(n), n) : Bool}:
1162 %add_zero_r(n) : {True{} == leb(half(n), _) : Bool}
1163 half_go_le(n, 0n)
1164
1165# a <= r -> r == a + (r - a)
1166def add_sub(+a: Nat, +r: Nat, e: {True{} == leb(a, r) : Bool}) -> {r == Nat.add(a, Nat.sub(r, a)) : Nat}:
1167 match a r:
1168 case 0n 0n:
1169 {==}
1170 case 0n 1n+q:
1171 {==}
1172 case 1n+p 0n:
1173 %e : BD(_, Unit, {0n == Nat.add(1n+p, Nat.sub(0n, 1n+p)) : Nat})
1174 Unit{}
1175 case 1n+p 1n+q:
1176 %add_sub(p, q, e) : {1n+q == 1n+_ : Nat}
1177 {==}
1178
1179# LAW (half_cover): the two bands of a node, half(r) rows and r - half(r) rows,
1180# add up to exactly r rows: no row is lost and none is counted twice.
1181law half_cover:
1182 for +n: Nat
1183 {n == Nat.add(half(n), Nat.sub(n, half(n))) : Nat}
1184
1185def half_cover(n):
1186 add_sub(half(n), n, half_le(n))
1187
1188type Bands<-r: Nat, -c: Nat> is Type:
1189 BLeaf{m: Mat<r, c>}
1190 BNode{x: Bands<half(r), c>, y: Bands<Nat.sub(r, half(r)), c>}
1191
1192# r x c of zeros in 2^d bands (d = 0: a single band, the same as a Mat)
1193def Bands.zeros(+d: Nat, +r: Nat, +c: Nat) -> Bands<r, c>:
1194 match d:
1195 case 0n:
1196 BLeaf{Mat.zeros(r, c)}
1197 case 1n+q:
1198 BNode{Bands.zeros(q, half(r), c), Bands.zeros(q, Nat.sub(r, half(r)), c)}
1199
1200# the row counts of the two bands follow the row count of the node: if rr == r at run
1201# time, the same holds for both halves (the bands carry r only in their type)
1202def eq_half(-a: Nat, -b: Nat, e: {a == b : Nat}) -> {half(a) == half(b) : Nat}:
1203 %e : {half(a) == half(_) : Nat}
1204 {==}
1205
1206def eq_rest(-a: Nat, -b: Nat, e: {a == b : Nat}) -> {Nat.sub(a, half(a)) == Nat.sub(b, half(b)) : Nat}:
1207 %e : {Nat.sub(a, half(a)) == Nat.sub(_, half(_)) : Nat}
1208 {==}
1209
1210# does a block of n numbers starting at flat index i end before the second band of a
1211# node with r rows? does it start in the second band?
1212def bfill_l(+r: Nat, +c: Nat, +i: Nat, +n: Nat) -> Bool:
1213 Nat.is_le(Nat.add(i, n), Nat.mul(half(r), c))
1214
1215def bfill_r(+r: Nat, +c: Nat, +i: Nat) -> Bool:
1216 Nat.is_le(Nat.mul(half(r), c), i)
1217
1218# writes the first n numbers of xs at index i (and nothing past them)
1219def fill_n(n: Nat, xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
1220 match n xs:
1221 case 1n+p Con{h, t}:
1222 fill_n(p, t, (i + 1 : U32), Array.set(F32, a, i, h))
1223 case _ _:
1224 a
1225
1226def fill_band(-r: Nat, -c: Nat, n: Nat, xs: List<&2, F32>, +i: U32, m: Mat<r, c>) -> Mat<r, c>:
1227 match m:
1228 case Mat{d}:
1229 Mat{fill_n(n, xs, i, d)}
1230
1231# writes the first n numbers of xs at flat index i; l and rt say which bands they touch
1232# (bfill_l, bfill_r); rr is the row count r at run time. A block that crosses into the second
1233# band is not copied: the first band writes only its part, the second gets the rest by List.drop.
1234def 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>:
1235 match b l rt:
1236 case BLeaf{m} _ _:
1237 BLeaf{fill_band(r, c, n, xs, U32.from_nat(i), m)}
1238 case BNode{x, y} True{} _:
1239 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}
1240 case BNode{x, y} False{} True{}:
1241 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)}
1242 case BNode{x, y} False{} False{}:
1243 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)))}
1244
1245def bfill_top(ok: Bool, +r: Nat, +c: Nat, b: Bands<r, c>, +i: Nat, +n: Nat, +xs: List<&2, F32>) -> Bands<r, c>:
1246 match ok:
1247 case True{}:
1248 bfill(c, r, b, bfill_l(r, c, i, n), bfill_r(r, c, i), r, {==}, i, n, xs)
1249 case False{}:
1250 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)
1251
1252# writes the n numbers of xs into b starting at flat index i (row-major, as Mat.fill_at);
1253# numbers past r*c are dropped. For loading a large matrix block by block.
1254def Bands.fill_at(+r: Nat, +c: Nat, +i: Nat, +n: Nat, +xs: List<&2, F32>, b: Bands<r, c>) -> Bands<r, c>:
1255 bfill_top(Nat.is_le(Nat.add(i, n), Nat.mul(r, c)), r, c, b, i, n, xs)
1256
1257# ---- y = B · x for B with r rows of c numbers (a weight matrix stored out x in) ----
1258
1259# B (intact) and y = B · x, typed like Mat.matmul_nt: x is 1 x c, y is 1 x r
1260type BV<-r: Nat, -c: Nat> is Type:
1261 BV{b: Bands<r, c>, y: Mat<1n, r>}
1262
1263# the bands and y as a list in row order (the result of Bands.matvec_l)
1264type BL<-r: Nat, -c: Nat> is Type:
1265 BL{b: Bands<r, c>, y: List<&2, F32>}
1266
1267def bv_g2(-r: Nat, -c: Nat, w: Array<F32>, al: AL) -> BL<r, c>:
1268 match al:
1269 case AL{d, l}:
1270 BL{BLeaf{Mat{w}}, l}
1271
1272# the band's 1 + m rows came out in an Array (one number per row); read them back as a list
1273def bv_g(-r: Nat, -c: Nat, +m: Nat, g: G) -> BL<r, c>:
1274 match g:
1275 case G{a, b, y}:
1276 bv_g2(r, c, b, read_l(m, 0, Nil{}, Array.get(F32, y, 0)))
1277
1278# the rows of one band, one dot product each, with the gemm of Mat.matmul_nt (n = 1), so the
1279# results are identical. Writing the results into an Array and reading them back is faster than
1280# building the list row by row (NOTES.md, exp. 14).
1281def bv_gemm(-r: Nat, +c: Nat, rows: Nat, +rr: Nat, x: Array<F32>, w: Array<F32>) -> BL<r, c>:
1282 match rows:
1283 case 0n:
1284 BL{BLeaf{Mat{w}}, Nil{}}
1285 case 1n+p:
1286 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)))
1287
1288def bv_mat(-r: Nat, +c: Nat, +rr: Nat, x: Array<F32>, m: Mat<r, c>) -> BL<r, c>:
1289 match m:
1290 case Mat{w}:
1291 bv_gemm(r, c, rr, rr, x, w)
1292
1293def bv_join(-r: Nat, -c: Nat, p: BL<half(r), c>, q: BL<Nat.sub(r, half(r)), c>) -> BL<r, c>:
1294 match p q:
1295 case BL{bx, yx} BL{by, yy}:
1296 BL{BNode{bx, by}, List.append(&2, F32, yx, yy)}
1297
1298# the two bands of a node run in parallel, each with its own copy of x (c numbers);
1299# 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
1300# parameter at the end of the list turned the parallel let into two sequential calls (NOTES.md, exp. 11; fixed in
1301# 2.0.37 by bendlang/bend#1377). The order is kept, and harmless.
1302def bmv(+c: Nat, -r: Nat, b: Bands<r, c>, +rr: Nat, -e: {rr == r : Nat}, xp: Array<F32> & Array<F32>) -> BL<r, c>:
1303 match b xp:
1304 case BLeaf{m} Tuple{x, x2}:
1305 bv_mat(r, c, rr, x, m)
1306 case BNode{bx, by} Tuple{x1, x2}:
1307 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))
1308 bv_join(r, c, p, q)
1309
1310def bv_fin(+r: Nat, -c: Nat, bl: BL<r, c>) -> BV<r, c>:
1311 match bl:
1312 case BL{b, y}:
1313 BV{b, Mat.from_list(1n, r, y)}
1314
1315# 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).
1316# The same product as Mat.matmul_nt(1n, c, r, ..), and the same numbers, bit for bit.
1317def Bands.matvec(+r: Nat, +c: Nat, b: Bands<r, c>, x: Mat<1n, c>) -> BV<r, c>:
1318 match x:
1319 case Mat{xa}:
1320 bv_fin(r, c, bmv(c, r, b, r, {==}, Array.clone(F32, xa)))
1321
1322# the same product with x and y as lists (x: c numbers, y: r numbers in row order), without the
1323# Mat conversions of Bands.matvec: for callers that already hold the vector as a list
1324def Bands.matvec_l(+r: Nat, +c: Nat, b: Bands<r, c>, +x: List<&2, F32>) -> BL<r, c>:
1325 bmv(c, r, b, r, {==}, Array.clone(F32, fill_list(x, 0, Array.new(F32, cap_depth(c), 0.0))))
1326
1327# ---- reading rows and the whole matrix ----
1328
1329type BRow<-r: Nat, -c: Nat> is Type:
1330 BRow{b: Bands<r, c>, row: List<&2, F32>}
1331
1332def brow_fin(-r: Nat, -c: Nat, al: AL) -> BRow<r, c>:
1333 match al:
1334 case AL{d, l}:
1335 BRow{BLeaf{Mat{d}}, l}
1336
1337# n numbers starting at flat index i of one band (a band may have no rows)
1338def brow_leaf2(-r: Nat, -c: Nat, n: Nat, +i: U32, d: Array<F32>) -> BRow<r, c>:
1339 match n:
1340 case 0n:
1341 BRow{BLeaf{Mat{d}}, Nil{}}
1342 case 1n+p:
1343 brow_fin(r, c, read_l(p, i, Nil{}, Array.get(F32, d, i)))
1344
1345def brow_leaf(-r: Nat, -c: Nat, m: Mat<r, c>, +i: U32, n: Nat) -> BRow<r, c>:
1346 match m:
1347 case Mat{d}:
1348 brow_leaf2(r, c, n, i, d)
1349
1350def brow_l(-r: Nat, -c: Nat, p: BRow<half(r), c>, y: Bands<Nat.sub(r, half(r)), c>) -> BRow<r, c>:
1351 match p:
1352 case BRow{bx, row}:
1353 BRow{BNode{bx, y}, row}
1354
1355def brow_r(-r: Nat, -c: Nat, x: Bands<half(r), c>, q: BRow<Nat.sub(r, half(r)), c>) -> BRow<r, c>:
1356 match q:
1357 case BRow{by, row}:
1358 BRow{BNode{x, by}, row}
1359
1360# row i; left says whether i is in the first band of the node (i < half(rr))
1361def brow(+c: Nat, -r: Nat, b: Bands<r, c>, left: Bool, +rr: Nat, -e: {rr == r : Nat}, +i: Nat) -> BRow<r, c>:
1362 match b left:
1363 case BLeaf{m} _:
1364 brow_leaf(r, c, m, U32.from_nat(Nat.mul(i, c)), c)
1365 case BNode{x, y} True{}:
1366 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)
1367 case BNode{x, y} False{}:
1368 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))))
1369
1370# row i (c numbers), for example an embedding lookup
1371def Bands.read_row(+r: Nat, +c: Nat, b: Bands<r, c>, +i: Nat) -> BRow<r, c>:
1372 brow(c, r, b, Nat.is_lt(i, half(r)), r, {==}, i)
1373
1374def bl_join(-r: Nat, -c: Nat, p: BRow<half(r), c>, q: BRow<Nat.sub(r, half(r)), c>) -> BRow<r, c>:
1375 match p q:
1376 case BRow{bx, xs} BRow{by, ys}:
1377 BRow{BNode{bx, by}, List.append(&2, F32, xs, ys)}
1378
1379def blist(+c: Nat, -r: Nat, b: Bands<r, c>, +rr: Nat, -e: {rr == r : Nat}) -> BRow<r, c>:
1380 match b:
1381 case BLeaf{m}:
1382 brow_leaf(r, c, m, 0, Nat.mul(rr, c))
1383 case BNode{x, y}:
1384 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)))
1385
1386# all r*c numbers in row order (in BRow.row)
1387def Bands.to_list(+r: Nat, +c: Nat, b: Bands<r, c>) -> BRow<r, c>:
1388 blist(c, r, b, r, {==})
1389
1390# ---- the list path returns exactly r numbers ----
1391#
1392# The numbers do not matter here, only how many there are: each lemma below counts the list a
1393# function builds, following the function's own recursion.
1394
1395def bl_len(-r: Nat, -c: Nat, bl: BL<r, c>) -> Nat:
1396 match bl:
1397 case BL{b, y}:
1398 List.length(&2, F32, y)
1399
1400# s == a + (1 + b) -> s == 1 + (a + b)
1401def succ_out(-s: Nat, +a: Nat, +b: Nat, e: {s == Nat.add(a, 1n+b) : Nat}) -> {s == 1n+Nat.add(a, b) : Nat}:
1402 %Equal.sym(Nat, 1n+Nat.add(a, b), Nat.add(a, 1n+b), add_succ_r(a, b)) : {s == _ : Nat}
1403 e
1404
1405# s == 1 + (p + (1 + l)) -> s == 2 + (p + l)
1406def len_step(-s: Nat, +p: Nat, +l: Nat, e: {s == 1n+Nat.add(p, 1n+l) : Nat}) -> {s == 2n+Nat.add(p, l) : Nat}:
1407 %Equal.sym(Nat, 1n+Nat.add(p, l), Nat.add(p, 1n+l), add_succ_r(p, l)) : {s == 1n+_ : Nat}
1408 e
1409
1410# s == n + 0 -> s == n
1411def zero_out(-s: Nat, +n: Nat, e: {s == Nat.add(n, 0n) : Nat}) -> {s == n : Nat}:
1412 %add_zero_r(n) : {s == _ : Nat}
1413 e
1414
1415# reversing onto acc adds the lengths
1416def 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}:
1417 match xs:
1418 case Nil{}:
1419 {==}
1420 case Con{h, t}:
1421 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))
1422
1423# appending adds the lengths
1424def 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}:
1425 match xs:
1426 case Nil{}:
1427 {==}
1428 case Con{h, t}:
1429 %app_len(t, ys) : {1n+List.length(&2, F32, List.append(&2, F32, t, ys)) == 1n+_ : Nat}
1430 {==}
1431
1432# a leaf's list: read_l reads mleft + 1 numbers onto acc and reverses them. By induction on
1433# mleft: with 0 left it reads one number (rev_len counts the reversed list); with 1 + p left it
1434# reads one and recurses, and len_step moves the 1 out of the sum.
1435law leaf_len:
1436 for -r: Nat
1437 for -c: Nat
1438 for w: Array<F32>
1439 for +mleft: Nat
1440 for +idx: U32
1441 for +acc: List<&2, F32>
1442 for rc: Array<F32> & F32
1443 {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}
1444
1445def leaf_len(r, c, w, mleft, idx, acc, rc):
1446 match mleft rc:
1447 case 0n Tuple{d, +v}:
1448 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{}))
1449 case 1n+p Tuple{d, +v}:
1450 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))))
1451
1452# whatever the gemm wrote, reading 1 + p rows back gives 1 + p numbers (G has one constructor, so
1453# matching g opens it, and leaf_len counts the list)
1454def g_len(-r: Nat, -c: Nat, +p: Nat, g: G) -> {bl_len(r, c, bv_g(r, c, p, g)) == 1n+p : Nat}:
1455 match g:
1456 case G{a, b, y}:
1457 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)))
1458
1459# LAW (band_len): the product of one band with `rows` rows gives exactly `rows` numbers.
1460# 0 rows give the empty list; 1 + p rows go through the gemm, and g_len counts what is read back.
1461law band_len:
1462 for -r: Nat
1463 for +c: Nat
1464 for +rows: Nat
1465 for x: Array<F32>
1466 for w: Array<F32>
1467 {bl_len(r, c, bv_gemm(r, c, rows, rows, x, w)) == rows : Nat}
1468
1469def band_len(r, c, rows, x, w):
1470 match rows:
1471 case 0n:
1472 {==}
1473 case 1n+p:
1474 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)))
1475
1476# joining two halves of h1 and h2 numbers gives h1 + h2 numbers. (The statement for the whole tree,
1477# "Bands.matvec_l gives r numbers", would apply this to the results of the two recursive calls and
1478# to the induction hypotheses about the same calls; that uses each band's Array twice, which Bend's
1479# affine rules refuse. It stays open, covered by tests: NOTES.md, v3.1.)
1480def 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}:
1481 match p q:
1482 case BL{bx, +yx} BL{by, +yy}:
1483 %hp : {List.length(&2, F32, List.append(&2, F32, yx, yy)) == Nat.add(_, h2) : Nat}
1484 %hq : {List.length(&2, F32, List.append(&2, F32, yx, yy)) == Nat.add(List.length(&2, F32, yx), _) : Nat}
1485 app_len(yx, yy)
1486