main.bend
bend-ml-tensor-array@0.1.2.0 · 676 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.2.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# smallest d with 2^d >= n (n >= 2)
31
32def dotg(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, acc: F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> R:
33 match k ra rb:
34 case 0n Tuple{a, x} Tuple{b, y}:
35 R{a, b, acc}
36 case 1n+p Tuple{a, +x} Tuple{b, +y}:
37 dotg(p, (ia + sa : U32), (ib + sb : U32), sa, sb, (acc + (x * y : F32) : F32), Array.get(F32, a, (ia + sa : U32)), Array.get(F32, b, (ib + sb : U32)))
38
39# row i of C, columns j.. ; r is the product of the current column (already computed)
40
41def gcols(mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, +cidx: U32, c: Array<F32>, r: R) -> G:
42 match mleft r:
43 case 0n R{a, b, +x}:
44 G{a, b, Array.set(F32, c, cidx, x)}
45 case 1n+p R{a, b, +x}:
46 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))))
47
48def grows(nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +mu: U32, g: G) -> G:
49 match nleft g:
50 case 0n G{a, b, c}:
51 G{a, b, c}
52 case 1n+p G{a, b, c}:
53 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))))
54
55# C (n x m) = A · B with k steps; c is the output array (with capacity n*m)
56
57def 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:
58 grows(n, Nat.sub(m, 1n), k, 0, arow, sa, sb, bcol, U32.from_nat(m), G{a, b, c})
59
60
61# ---------------------------------------------------------------------
62# parallel gemm: 2^d blocks of rows of C, each with its own copy of A and B,
63# returning lists that are written into C at the end
64# ---------------------------------------------------------------------
65
66type S3 is Type:
67 S3{a: Array<F32>, b: Array<F32>, out: List<&2, F32>}
68
69def lcols(mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, out: List<&2, F32>, r: R) -> S3:
70 match mleft r:
71 case 0n R{a, b, +x}:
72 S3{a, b, x <> out}
73 case 1n+p R{a, b, +x}:
74 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))))
75
76def lrows(nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, st: S3) -> S3:
77 match nleft st:
78 case 0n S3{a, b, out}:
79 S3{a, b, out}
80 case 1n+p S3{a, b, out}:
81 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))))
82
83def lblock_fin(st: S3) -> List<&2, F32>:
84 match st:
85 case S3{a, b, out}:
86 List.reverse(&2, F32, out)
87
88def 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>:
89 lblock_fin(lrows(cnt, mm1, kn, i, arow, sa, sb, bcol, S3{a, b, Nil{}}))
90
91def 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>:
92 match d ra rb:
93 case 0n Tuple{a, a2} Tuple{b, b2}:
94 lblock(cnt, mm1, kn, i, arow, sa, sb, bcol, a, b)
95 case 1n++q Tuple{a, a2} Tuple{b, b2}:
96 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))
97 List.append(&2, F32, x, y)
98
99def write_l(xs: List<&2, F32>, +idx: U32, c: Array<F32>) -> Array<F32>:
100 match xs:
101 case Nil{}:
102 c
103 case Con{h, t}:
104 write_l(t, (idx + 1 : U32), Array.set(F32, c, idx, h))
105
106# c[idx..] += list
107
108def gp3(a: Array<F32>, b: Array<F32>, c: Array<F32>, xs: List<&2, F32>) -> G:
109 G{a, b, write_l(xs, 0, c)}
110
111def 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:
112 match pa pb:
113 case Tuple{a1, a2} Tuple{b1, b2}:
114 gp3(a1, b1, c, lpar(d, n, mm1, kn, 0, arow, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2)))
115
116# C (n x m) = A · B, with 2^d blocks of rows in parallel; returns A, B (intact) and C
117
118def 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:
119 gp2(d, n, Nat.sub(m, 1n), k, arow, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b))
120
121
122# ---------------------------------------------------------------------
123# parallel gemm by COLUMNS, for n = 1 (matrix · vector): each block computes a
124# range of columns of C, with its own copy of the operands
125# ---------------------------------------------------------------------
126
127def cblock_fin(st: S3) -> List<&2, F32>:
128 match st:
129 case S3{a, b, out}:
130 List.reverse(&2, F32, out)
131
132# cnt columns starting at j0 (cnt >= 1); row 0 of A is at ia0
133
134def 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>:
135 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)))))
136
137def cblock(cnt: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>) -> List<&2, F32>:
138 match cnt:
139 case 0n:
140 Nil{}
141 case 1n+p:
142 cblock_go(p, kn, ia0, j0, sa, sb, bcol, a, b)
143
144def 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>:
145 match d ra rb:
146 case 0n Tuple{a, a2} Tuple{b, b2}:
147 cblock(cnt, kn, ia0, j0, sa, sb, bcol, a, b)
148 case 1n++q Tuple{a, a2} Tuple{b, b2}:
149 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))
150 List.append(&2, F32, x, y)
151
152def 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:
153 match pa pb:
154 case Tuple{a1, a2} Tuple{b1, b2}:
155 gp3(a1, b1, c, cpar(d, m, kn, 0, 0, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2)))
156
157# C (1 x m) = A (1 x k) · B, with 2^d blocks of columns in parallel
158
159def 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:
160 cp2(d, m, k, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b))
161
162# ---------------------------------------------------------------------
163# element-wise operations (in place, by index)
164# ---------------------------------------------------------------------
165
166def depth(n: Nat) -> Nat:
167 Nat.add(U32.log2(U32.from_nat(Nat.sub(n, 1n))), 1n)
168
169# ---------------------------------------------------------------------
170# gemm: C[i][j] = sum_p A[i*arow + p*sa] * B[p*sb + j*bcol]
171# (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)
172# ---------------------------------------------------------------------
173
174# acc += A[ia + t*sa] * B[ib + t*sb] for k steps; ra, rb already carry the first pair
175
176type P2 is Type:
177 P2{a: Array<F32>, b: Array<F32>}
178
179type AL is Type:
180 AL{c: Array<F32>, l: List<&2, F32>}
181
182# list of m elements starting at idx (reads with Array.get, returning the array)
183
184def read_l(mleft: Nat, +idx: U32, acc: List<&2, F32>, rc: Array<F32> & F32) -> AL:
185 match mleft rc:
186 case 0n Tuple{c, +v}:
187 AL{c, List.reverse(&2, F32, v <> acc)}
188 case 1n+p Tuple{c, +v}:
189 read_l(p, (idx + 1 : U32), v <> acc, Array.get(F32, c, (idx + 1 : U32)))
190
191# writes the list starting at idx
192
193def addl(xs: List<&2, F32>, +idx: U32, rc: Array<F32> & F32) -> Array<F32>:
194 match xs rc:
195 case Nil{} Tuple{c, v}:
196 c
197 case Con{+h, t} Tuple{c, +v}:
198 addl(t, (idx + 1 : U32), Array.get(F32, Array.set(F32, c, idx, (v + h : F32)), (idx + 1 : U32)))
199
200# adds a bias (a list of m) to each of the n rows of c (n x m)
201
202def bias_rows(nleft: Nat, +i: U32, +mu: U32, +bl: List<&2, F32>, c: Array<F32>) -> Array<F32>:
203 match nleft:
204 case 0n:
205 c
206 case 1n+p:
207 bias_rows(p, (i + 1 : U32), mu, bl, addl(bl, (i * mu : U32), Array.get(F32, c, (i * mu : U32))))
208
209# in-place relu on the first n elements
210
211def relu_l(nleft: Nat, +t: U32, rc: Array<F32> & F32) -> Array<F32>:
212 match nleft rc:
213 case 0n Tuple{c, v}:
214 c
215 case 1n+p Tuple{c, +v}:
216 relu_l(p, (t + 1 : U32), Array.get(F32, Array.set(F32, c, t, F32.max(v, 0.0)), (t + 1 : U32)))
217
218def pos2(b: Bool) -> F32:
219 match b:
220 case True{}:
221 1.0
222 case False{}:
223 0.0
224
225def pos(+y: F32) -> F32:
226 pos2(F32.is_gt(y, 0.0))
227
228# d[t] *= (h[t] > 0): the relu gradient; returns both arrays
229
230def mask_l(nleft: Nat, +t: U32, rd: Array<F32> & F32, rh: Array<F32> & F32) -> P2:
231 match nleft rd rh:
232 case 0n Tuple{d, x} Tuple{h, y}:
233 P2{d, h}
234 case 1n+p Tuple{d, +x} Tuple{h, +y}:
235 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)))
236
237# w[t] -= lr * dw[t]
238
239def maxl2(xs: List<&2, F32>, cur: F32) -> F32:
240 match xs:
241 case Nil{}:
242 cur
243 case Con{+h, t}:
244 maxl2(t, F32.max(cur, h))
245
246def maxl(xs: List<&2, F32>) -> F32:
247 match xs:
248 case Nil{}:
249 0.0
250 case Con{+h, t}:
251 maxl2(t, h)
252
253def expl(xs: List<&2, F32>, +m: F32) -> List<&2, F32>:
254 match xs:
255 case Nil{}:
256 Nil{}
257 case Con{+h, t}:
258 F32.exp((h - m : F32)) <> expl(t, m)
259
260def suml(xs: List<&2, F32>, acc: F32) -> F32:
261 match xs:
262 case Nil{}:
263 acc
264 case Con{+h, t}:
265 suml(t, (acc + h : F32))
266
267def divl(xs: List<&2, F32>, +s: F32) -> List<&2, F32>:
268 match xs:
269 case Nil{}:
270 Nil{}
271 case Con{+h, t}:
272 (h / s : F32) <> divl(t, s)
273
274def softmax2(+es: List<&2, F32>) -> List<&2, F32>:
275 divl(es, suml(es, 0.0))
276
277def softmax(+xs: List<&2, F32>) -> List<&2, F32>:
278 softmax2(expl(xs, maxl(xs)))
279
280# p[label]
281
282def sgd_l(nleft: Nat, +t: U32, +lr: F32, rw: Array<F32> & F32, rd: Array<F32> & F32) -> P2:
283 match nleft rw rd:
284 case 0n Tuple{w, x} Tuple{d, y}:
285 P2{w, d}
286 case 1n+p Tuple{w, +x} Tuple{d, +y}:
287 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)))
288
289# ---------------------------------------------------------------------
290# softmax + cross-entropy per row (lists of 10 elements: negligible cost)
291# ---------------------------------------------------------------------
292
293def pick(ps: List<&2, F32>, n: Nat) -> F32:
294 match ps n:
295 case Con{+h, t} 0n:
296 h
297 case Con{h, t} 1n+q:
298 pick(t, q)
299 case Nil{} _:
300 0.0
301
302# (p - onehot(label)) / n
303
304def hit(b: Bool) -> F32:
305 match b:
306 case True{}:
307 1.0
308 case False{}:
309 0.0
310
311def grad_l(ps: List<&2, F32>, +label: Nat, +j: Nat, +nf: F32) -> List<&2, F32>:
312 match ps:
313 case Nil{}:
314 Nil{}
315 case Con{+h, t}:
316 (((h - hit(Nat.is_eq(label, j)) : F32) / nf) : F32) <> grad_l(t, label, 1n+j, nf)
317
318type L2 is Type:
319 L2{loss: F32, c: Array<F32>}
320
321def ce_row2(+ps: List<&2, F32>, +label: Nat, +nf: F32, +idx: U32, loss: F32, c: Array<F32>) -> L2:
322 L2{(loss - F32.log(F32.max(pick(ps, label), 0.000000000001)) : F32), write_l(grad_l(ps, label, 0n, nf), idx, c)}
323
324def ce_row(al: AL, +label: Nat, +nf: F32, +idx: U32, loss: F32) -> L2:
325 match al:
326 case AL{c, l}:
327 ce_row2(softmax(l), label, nf, idx, loss, c)
328
329# 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
330
331type Mat<-r: Nat, -c: Nat> is Type:
332 Mat{d: Array<F32>}
333
334# smallest d with 2^d >= n (n >= 2); n < 2 gives 0
335
336def ce_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +nf: F32, +i: U32, r: L2) -> L2:
337 match labels r:
338 case Nil{} L2{loss, c}:
339 L2{loss, c}
340 case Con{+y, t} L2{loss, c}:
341 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))
342
343# ---- helper: hits per row ----
344
345type CR is Type:
346 CR{c: Array<F32>, n: Nat}
347
348def gt_head(xs: List<&2, F32>, +v: F32) -> Bool:
349 match xs:
350 case Nil{}:
351 False{}
352 case Con{+h, t}:
353 F32.is_gt(h, v)
354
355def argmax2(xs: List<&2, F32>, g: Bool, +i: Nat, +bi: Nat, +bv: F32) -> Nat:
356 match xs g:
357 case Nil{} _:
358 bi
359 case Con{+h, +t} True{}:
360 argmax2(t, gt_head(t, h), 1n+i, i, h)
361 case Con{h, +t} False{}:
362 argmax2(t, gt_head(t, bv), 1n+i, bi, bv)
363
364def argmax(xs: List<&2, F32>) -> Nat:
365 match xs:
366 case Nil{}:
367 0n
368 case Con{+h, +t}:
369 argmax2(t, gt_head(t, h), 1n, 0n, h)
370
371def one(b: Bool) -> Nat:
372 match b:
373 case True{}:
374 1n
375 case False{}:
376 0n
377
378def count_rows3(+y: Nat, n: Nat, al: AL) -> CR:
379 match al:
380 case AL{c, l}:
381 CR{c, Nat.add(n, one(Nat.is_eq(argmax(l), y)))}
382
383def count_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +i: U32, r: CR) -> CR:
384 match labels r:
385 case Nil{} CR{c, n}:
386 CR{c, n}
387 case Con{+y, t} CR{c, n}:
388 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)))))
389
390
391# =====================================================================
392# Typed API: Mat<r, c> (the dimensions enter the type)
393# =====================================================================
394
395def cap_depth2(small: Bool, n: Nat) -> Nat:
396 match small:
397 case True{}:
398 0n
399 case False{}:
400 depth(n)
401
402def cap_depth(+n: Nat) -> Nat:
403 cap_depth2(Nat.is_lt(n, 2n), n)
404
405# r x c of zeros, with the exact capacity (smallest power of 2)
406
407def Mat.zeros(+r: Nat, +c: Nat) -> Mat<r, c>:
408 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0)}
409
410def Mat.fill(+r: Nat, +c: Nat, +v: F32) -> Mat<r, c>:
411 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), v)}
412
413def fill_list(xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
414 match xs:
415 case Nil{}:
416 a
417 case Con{h, t}:
418 fill_list(t, (i + 1 : U32), Array.set(F32, a, i, h))
419
420def len_is(xs: List<&2, F32>, n: Nat) -> Bool:
421 match xs n:
422 case Nil{} 0n:
423 True{}
424 case Con{h, t} 1n+p:
425 len_is(t, p)
426 case _ _:
427 False{}
428
429def Mat.of_list2(-r: Nat, -c: Nat, +rr: Nat, +cc: Nat, ok: Bool, xs: List<&2, F32>) -> Maybe<&1, Mat<r, c>>:
430 match ok:
431 case True{}:
432 Some{Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(rr, cc)), 0.0))}}
433 case False{}:
434 None{}
435
436# Mat<r, c> from r*c numbers in row order; None if the size does not match
437
438def Mat.of_list(+r: Nat, +c: Nat, +xs: List<&2, F32>) -> Maybe<&1, Mat<r, c>>:
439 Mat.of_list2(r, c, r, c, len_is(xs, Nat.mul(r, c)), xs)
440
441# reads the first n numbers, returning the matrix
442
443# Mat<r, c> without checking the list size (a short list leaves the rest at zero, a long one is cut
444# by the capacity): only for when the size is known by construction. Prefer Mat.of_list.
445
446def Mat.from_list(+r: Nat, +c: Nat, xs: List<&2, F32>) -> Mat<r, c>:
447 Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0))}
448
449type ML<-r: Nat, -c: Nat> is Type:
450 ML{m: Mat<r, c>, l: List<&2, F32>}
451
452def Mat.to_list2(-r: Nat, -c: Nat, al: AL) -> ML<r, c>:
453 match al:
454 case AL{d, l}:
455 ML{Mat{d}, l}
456
457# two independent copies
458
459def Mat.to_list(+r: Nat, +c: Nat, m: Mat<r, c>) -> ML<r, c>:
460 match m:
461 case Mat{d}:
462 Mat.to_list2(r, c, read_l(Nat.sub(Nat.mul(r, c), 1n), 0, Nil{}, Array.get(F32, d, 0)))
463
464type MM2<-r: Nat, -c: Nat> is Type:
465 MM2{a: Mat<r, c>, b: Mat<r, c>}
466
467def Mat.clone2(-r: Nat, -c: Nat, p: Array<F32> & Array<F32>) -> MM2<r, c>:
468 match p:
469 case Tuple{x, y}:
470 MM2{Mat{x}, Mat{y}}
471
472# ---- products: C (n x m) = A · B, with 2^par blocks of rows in parallel ----
473
474# C = A(n x k) · B(k x m)
475
476def Mat.clone(-r: Nat, -c: Nat, m: Mat<r, c>) -> MM2<r, c>:
477 match m:
478 case Mat{d}:
479 Mat.clone2(r, c, Array.clone(F32, d))
480
481type MMul<-n: Nat, -k: Nat, -m: Nat> is Type:
482 MMul{a: Mat<n, k>, b: Mat<k, m>, c: Mat<n, m>}
483
484# C = A(n x k) · Bᵀ, with B stored m x k
485
486type MMulNT<-n: Nat, -k: Nat, -m: Nat> is Type:
487 MMulNT{a: Mat<n, k>, b: Mat<m, k>, c: Mat<n, m>}
488
489# C = Aᵀ · B, with A stored k x n
490
491type MMulTN<-n: Nat, -k: Nat, -m: Nat> is Type:
492 MMulTN{a: Mat<k, n>, b: Mat<k, m>, c: Mat<n, m>}
493
494def gemm_any3(cols: 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:
495 match cols:
496 case True{}:
497 gemm_par_cols(d, m, k, sa, sb, bcol, a, b, c)
498 case False{}:
499 gemm_par(d, n, m, k, arow, sa, sb, bcol, a, b, c)
500
501def 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:
502 match d:
503 case 0n:
504 gemm(n, m, k, arow, sa, sb, bcol, a, b, c)
505 case 1n+q:
506 gemm_any3(Nat.is_eq(n, 1n), 1n+q, n, m, k, arow, sa, sb, bcol, a, b, c)
507
508def 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:
509 gemm_any2(d, n, m, k, arow, sa, sb, bcol, a, b, c)
510
511def mmul_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMul<n, k, m>:
512 match g:
513 case G{a, b, c}:
514 MMul{Mat{a}, Mat{b}, Mat{c}}
515
516def mmul_nt_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulNT<n, k, m>:
517 match g:
518 case G{a, b, c}:
519 MMulNT{Mat{a}, Mat{b}, Mat{c}}
520
521def mmul_tn_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulTN<n, k, m>:
522 match g:
523 case G{a, b, c}:
524 MMulTN{Mat{a}, Mat{b}, Mat{c}}
525
526def mm_a(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMul<n, k, m>:
527 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)))
528
529# C = A · B (n x k) · (k x m); par = log2 of the number of parallel blocks (0 = sequential)
530
531def Mat.matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<k, m>) -> MMul<n, k, m>:
532 match a b:
533 case Mat{x} Mat{y}:
534 mm_a(n, k, m, par, x, y)
535
536def mm_nt(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMulNT<n, k, m>:
537 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)))
538
539# C = A · Bᵀ (n x k) · (m x k)ᵀ: for weights stored with one row per output
540
541def Mat.matmul_nt(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<m, k>) -> MMulNT<n, k, m>:
542 match a b:
543 case Mat{x} Mat{y}:
544 mm_nt(n, k, m, par, x, y)
545
546def mm_tn(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMulTN<n, k, m>:
547 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)))
548
549# C = Aᵀ · B (k x n)ᵀ · (k x m): weight gradient, X^T · dY
550
551def Mat.matmul_tn(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<k, n>, b: Mat<k, m>) -> MMulTN<n, k, m>:
552 match a b:
553 case Mat{x} Mat{y}:
554 mm_tn(n, k, m, par, x, y)
555
556# ---- element-wise and training operations ----
557
558# adds the bias (one row, m numbers) to each row of C (n x m)
559
560type MBias<-n: Nat, -m: Nat> is Type:
561 MBias{c: Mat<n, m>, b: Mat<1n, m>}
562
563def add_row2(-n: Nat, -m: Nat, +nn: Nat, +mm: Nat, cd: Array<F32>, al: AL) -> MBias<n, m>:
564 match al:
565 case AL{bd, bl}:
566 MBias{Mat{bias_rows(nn, 0, U32.from_nat(mm), bl, cd)}, Mat{bd}}
567
568def Mat.add_row(+n: Nat, +m: Nat, c: Mat<n, m>, b: Mat<1n, m>) -> MBias<n, m>:
569 match c b:
570 case Mat{cd} Mat{bd}:
571 add_row2(n, m, n, m, cd, read_l(Nat.sub(m, 1n), 0, Nil{}, Array.get(F32, bd, 0)))
572
573# pair of matrices of the same shape (the result of an operation that reads two operands)
574
575type MPair<-r: Nat, -c: Nat> is Type:
576 MPair{a: Mat<r, c>, b: Mat<r, c>}
577
578def mpair_fin(-r: Nat, -c: Nat, p: P2) -> MPair<r, c>:
579 match p:
580 case P2{a, b}:
581 MPair{Mat{a}, Mat{b}}
582
583def relu_fin(-r: Nat, -c: Nat, d: Array<F32>) -> Mat<r, c>:
584 Mat{d}
585
586def Mat.relu(+r: Nat, +c: Nat, a: Mat<r, c>) -> Mat<r, c>:
587 match a:
588 case Mat{d}:
589 relu_fin(r, c, relu_l(Nat.mul(r, c), 0, Array.get(F32, d, 0)))
590
591# (dy * (h > 0), h): the gradient that passes through the relu
592
593def Mat.relu_bwd(+r: Nat, +c: Nat, dy: Mat<r, c>, h: Mat<r, c>) -> MPair<r, c>:
594 match dy h:
595 case Mat{d} Mat{x}:
596 mpair_fin(r, c, mask_l(Nat.mul(r, c), 0, Array.get(F32, d, 0), Array.get(F32, x, 0)))
597
598# (w - lr * dw, dw): one SGD step
599
600def Mat.sgd(+r: Nat, +c: Nat, +lr: F32, w: Mat<r, c>, dw: Mat<r, c>) -> MPair<r, c>:
601 match w dw:
602 case Mat{x} Mat{y}:
603 mpair_fin(r, c, sgd_l(Nat.mul(r, c), 0, lr, Array.get(F32, x, 0), Array.get(F32, y, 0)))
604
605# (a + b, b): element-wise sum (a - (-1)*b is exactly a + b)
606
607def Mat.add(+r: Nat, +c: Nat, a: Mat<r, c>, b: Mat<r, c>) -> MPair<r, c>:
608 match a b:
609 case Mat{x} Mat{y}:
610 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)))
611
612# sum of each column (n x m -> 1 x m), returning the matrix
613
614type MColSum<-n: Nat, -m: Nat> is Type:
615 MColSum{a: Mat<n, m>, s: Mat<1n, m>}
616
617def colsum_fin(-n: Nat, -m: Nat, g: G) -> MColSum<n, m>:
618 match g:
619 case G{ones, a, s}:
620 MColSum{Mat{a}, Mat{s}}
621
622def Mat.col_sums(+n: Nat, +m: Nat, a: Mat<n, m>) -> MColSum<n, m>:
623 match a:
624 case Mat{d}:
625 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)))
626
627# mean cross-entropy and its gradient on the logits: (softmax - one-hot) / n.
628# labels has one entry per row (n); the gradient replaces the logits.
629
630type MCE<-n: Nat, -c: Nat> is Type:
631 MCE{loss: F32, grad: Mat<n, c>}
632
633def ce_fin(-n: Nat, -c: Nat, +nf: F32, r: L2) -> MCE<n, c>:
634 match r:
635 case L2{loss, d}:
636 MCE{(loss / nf : F32), Mat{d}}
637
638def Mat.softmax_ce(+n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> MCE<n, c>:
639 match z:
640 case Mat{d}:
641 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}))
642
643# argmax hits per row against labels
644
645type MHits<-n: Nat, -c: Nat> is Type:
646 MHits{z: Mat<n, c>, hits: Nat}
647
648def hits_fin(-n: Nat, -c: Nat, r: CR) -> MHits<n, c>:
649 match r:
650 case CR{d, k}:
651 MHits{Mat{d}, k}
652
653def Mat.count_correct(+n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> MHits<n, c>:
654 match z:
655 case Mat{d}:
656 hits_fin(n, c, count_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), 0, CR{d, 0n}))
657
658# a row as a list (returning the matrix) and the other way around
659
660type MRow<-r: Nat, -c: Nat> is Type:
661 MRow{m: Mat<r, c>, row: List<&2, F32>}
662
663def row_fin(-r: Nat, -c: Nat, al: AL) -> MRow<r, c>:
664 match al:
665 case AL{d, l}:
666 MRow{Mat{d}, l}
667
668def Mat.read_row(+r: Nat, +c: Nat, m: Mat<r, c>, +i: U32) -> MRow<r, c>:
669 match m:
670 case Mat{d}:
671 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))))
672
673def Mat.write_row(-r: Nat, +c: Nat, m: Mat<r, c>, +i: U32, xs: List<&2, F32>) -> Mat<r, c>:
674 match m:
675 case Mat{d}:
676 Mat{write_l(xs, (i * U32.from_nat(c) : U32), d)}