bend_tensors.bend

bend-tensors@0.0.0.2 · 2,220 lines · 306 definitions

1# bend_tensors: dense linear algebra with shapes in the types.
2#
3# A Vec(n) holds 4 * 2^n floats and a Mat(r, c) 4 * 2^r rows of 4 * 2^c
4# floats, both as balanced trees of 4x4 register blocks. The tree is the
5# fork tree: every op below a node is a parallel let, and the matrix
6# product recurses on the tree, cache-oblivious, down to a 64-multiply
7# kernel. A shape mismatch is a checker error, not a runtime one.
8#
9# Nothing here matches on a computed value, recurses mutually, or opens
10# an Array: those are the rules the compiler makes fast code from. Read
11# AGENTS.md before editing. This file is assembled from gen/ by
12# gen/assemble.py; edit the parts, not this file. bend_tensors.bend is
13# the release (CPU BLAS); bend_tensors_gpu.bend adds the GPU paths.
14
15import Base
16
17# Base's F32.max and F32.show take + parameters, which a template slot
18# does not accept; these take plain ones.
19def F32.max.of(a: F32, b: F32) -> F32:
20  F32.max(a, b)
21
22def F32.show.of(x: F32) -> String:
23  F32.show(x)
24
25# Blocks
26# ======
27#
28# A Blk is a 4x4 block of F32 in row-major order, a Vec4 four F32. Both
29# are Data records of scalars: they ride in registers, and the trees
30# built from them are shared through + binders with no copy. A block
31# product is 64 fused multiply-adds with no memory traffic: the kernel
32# every larger op bottoms out in.
33
34type Blk is Data:
35  Blk{a00: F32, a01: F32, a02: F32, a03: F32,
36      a10: F32, a11: F32, a12: F32, a13: F32,
37      a20: F32, a21: F32, a22: F32, a23: F32,
38      a30: F32, a31: F32, a32: F32, a33: F32}
39
40type Vec4 is Data:
41  Vec4{x0: F32, x1: F32, x2: F32, x3: F32}
42
43# Vec4
44# ----
45
46def Vec4.full(+x: F32) -> Vec4:
47  Vec4{x, x, x, x}
48
49def Vec4.zeros() -> Vec4:
50  Vec4{0.0, 0.0, 0.0, 0.0}
51
52# x, x + 1, x + 2, x + 3
53def Vec4.arange(+x: F32) -> Vec4:
54  Vec4{x, (x + 1.0 : F32), (x + 2.0 : F32), (x + 3.0 : F32)}
55
56def Vec4.map(~f: F32 -> F32, v: Vec4) -> Vec4:
57  Vec4{x0, x1, x2, x3} = v
58  Vec4{f(x0), f(x1), f(x2), f(x3)}
59
60def Vec4.zip(~f: F32 -> F32 -> F32, a: Vec4, b: Vec4) -> Vec4:
61  Vec4{x0, x1, x2, x3} = a
62  Vec4{y0, y1, y2, y3} = b
63  Vec4{f(x0, y0), f(x1, y1), f(x2, y2), f(x3, y3)}
64
65def Vec4.add(a: Vec4, b: Vec4) -> Vec4:
66  Vec4{x0, x1, x2, x3} = a
67  Vec4{y0, y1, y2, y3} = b
68  Vec4{(x0 + y0 : F32), (x1 + y1 : F32), (x2 + y2 : F32), (x3 + y3 : F32)}
69
70def Vec4.sub(a: Vec4, b: Vec4) -> Vec4:
71  Vec4{x0, x1, x2, x3} = a
72  Vec4{y0, y1, y2, y3} = b
73  Vec4{(x0 - y0 : F32), (x1 - y1 : F32), (x2 - y2 : F32), (x3 - y3 : F32)}
74
75def Vec4.mul(a: Vec4, b: Vec4) -> Vec4:
76  Vec4{x0, x1, x2, x3} = a
77  Vec4{y0, y1, y2, y3} = b
78  Vec4{(x0 * y0 : F32), (x1 * y1 : F32), (x2 * y2 : F32), (x3 * y3 : F32)}
79
80def Vec4.scale(+k: F32, v: Vec4) -> Vec4:
81  Vec4{x0, x1, x2, x3} = v
82  Vec4{(k * x0 : F32), (k * x1 : F32), (k * x2 : F32), (k * x3 : F32)}
83
84# the sum, as (x0 + x1) + (x2 + x3): every reduction in this file
85# pairs its terms the same way, so laws can equate them
86def Vec4.sum(v: Vec4) -> F32:
87  Vec4{x0, x1, x2, x3} = v
88  ((x0 + x1) + (x2 + x3) : F32)
89
90def Vec4.dot(a: Vec4, b: Vec4) -> F32:
91  Vec4{x0, x1, x2, x3} = a
92  Vec4{y0, y1, y2, y3} = b
93  ((x0 * y0 + x1 * y1) + (x2 * y2 + x3 * y3) : F32)
94
95def Vec4.max(v: Vec4) -> F32:
96  Vec4{x0, x1, x2, x3} = v
97  F32.max(F32.max(x0, x1), F32.max(x2, x3))
98
99def Vec4.min(v: Vec4) -> F32:
100  Vec4{x0, x1, x2, x3} = v
101  F32.min(F32.min(x0, x1), F32.min(x2, x3))
102
103# the index (0-3) and value of the largest element; the first wins ties
104def Vec4.argmax.pick(c: Bool, +i: U32, +x: F32, +j: U32, +y: F32) -> U32 & F32:
105  match c:
106    case True{}:
107      (i, x)
108    case False{}:
109      (j, y)
110
111def Vec4.argmax.two(ix: U32 & F32, jy: U32 & F32) -> U32 & F32:
112  (i, x) = ix
113  (j, y) = jy
114  +x2 = x
115  +y2 = y
116  Vec4.argmax.pick((x2 >= y2 : F32), i, x2, j, y2)
117
118def Vec4.argmax(v: Vec4) -> U32 & F32:
119  Vec4{x0, x1, x2, x3} = v
120  Vec4.argmax.two(Vec4.argmax.two((0, x0), (1, x1)), Vec4.argmax.two((2, x2), (3, x3)))
121
122def Vec4.to_list(v: Vec4, acc: List<&2, F32>) -> List<&2, F32>:
123  Vec4{x0, x1, x2, x3} = v
124  x0 <> x1 <> x2 <> x3 <> acc
125
126# the first four of xs, padded with 0.0, and the rest
127def Vec4.from_list(xs: List<&2, F32>) -> Vec4:
128  match xs:
129    case x0 <> x1 <> x2 <> x3 <> rest:
130      Vec4{x0, x1, x2, x3}
131    case x0 <> x1 <> x2 <> []:
132      Vec4{x0, x1, x2, 0.0}
133    case x0 <> x1 <> []:
134      Vec4{x0, x1, 0.0, 0.0}
135    case x0 <> []:
136      Vec4{x0, 0.0, 0.0, 0.0}
137    case []:
138      Vec4{0.0, 0.0, 0.0, 0.0}
139
140def Vec4.show(v: Vec4) -> String:
141  Vec4{x0, x1, x2, x3} = v
142  F32.show(x0) ++ " " ++ F32.show(x1) ++ " " ++ F32.show(x2) ++ " " ++ F32.show(x3)
143
144# Blk
145# ---
146
147def Blk.full(+x: F32) -> Blk:
148  Blk{x, x, x, x,
149      x, x, x, x,
150      x, x, x, x,
151      x, x, x, x}
152
153def Blk.zeros() -> Blk:
154  Blk.full(0.0)
155
156def Blk.eye() -> Blk:
157  Blk{1.0, 0.0, 0.0, 0.0,
158      0.0, 1.0, 0.0, 0.0,
159      0.0, 0.0, 1.0, 0.0,
160      0.0, 0.0, 0.0, 1.0}
161
162def Blk.map(~f: F32 -> F32, m: Blk) -> Blk:
163  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
164  Blk{f(a00), f(a01), f(a02), f(a03),
165      f(a10), f(a11), f(a12), f(a13),
166      f(a20), f(a21), f(a22), f(a23),
167      f(a30), f(a31), f(a32), f(a33)}
168
169def Blk.zip(~f: F32 -> F32 -> F32, m: Blk, n: Blk) -> Blk:
170  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
171  Blk{b00, b01, b02, b03, b10, b11, b12, b13, b20, b21, b22, b23, b30, b31, b32, b33} = n
172  Blk{f(a00, b00), f(a01, b01), f(a02, b02), f(a03, b03),
173      f(a10, b10), f(a11, b11), f(a12, b12), f(a13, b13),
174      f(a20, b20), f(a21, b21), f(a22, b22), f(a23, b23),
175      f(a30, b30), f(a31, b31), f(a32, b32), f(a33, b33)}
176
177def Blk.add(m: Blk, n: Blk) -> Blk:
178  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
179  Blk{b00, b01, b02, b03, b10, b11, b12, b13, b20, b21, b22, b23, b30, b31, b32, b33} = n
180  Blk{(a00 + b00 : F32), (a01 + b01 : F32), (a02 + b02 : F32), (a03 + b03 : F32),
181      (a10 + b10 : F32), (a11 + b11 : F32), (a12 + b12 : F32), (a13 + b13 : F32),
182      (a20 + b20 : F32), (a21 + b21 : F32), (a22 + b22 : F32), (a23 + b23 : F32),
183      (a30 + b30 : F32), (a31 + b31 : F32), (a32 + b32 : F32), (a33 + b33 : F32)}
184
185def Blk.sub(m: Blk, n: Blk) -> Blk:
186  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
187  Blk{b00, b01, b02, b03, b10, b11, b12, b13, b20, b21, b22, b23, b30, b31, b32, b33} = n
188  Blk{(a00 - b00 : F32), (a01 - b01 : F32), (a02 - b02 : F32), (a03 - b03 : F32),
189      (a10 - b10 : F32), (a11 - b11 : F32), (a12 - b12 : F32), (a13 - b13 : F32),
190      (a20 - b20 : F32), (a21 - b21 : F32), (a22 - b22 : F32), (a23 - b23 : F32),
191      (a30 - b30 : F32), (a31 - b31 : F32), (a32 - b32 : F32), (a33 - b33 : F32)}
192
193def Blk.mul(m: Blk, n: Blk) -> Blk:
194  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
195  Blk{b00, b01, b02, b03, b10, b11, b12, b13, b20, b21, b22, b23, b30, b31, b32, b33} = n
196  Blk{(a00 * b00 : F32), (a01 * b01 : F32), (a02 * b02 : F32), (a03 * b03 : F32),
197      (a10 * b10 : F32), (a11 * b11 : F32), (a12 * b12 : F32), (a13 * b13 : F32),
198      (a20 * b20 : F32), (a21 * b21 : F32), (a22 * b22 : F32), (a23 * b23 : F32),
199      (a30 * b30 : F32), (a31 * b31 : F32), (a32 * b32 : F32), (a33 * b33 : F32)}
200
201def Blk.scale(+k: F32, m: Blk) -> Blk:
202  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
203  Blk{(k * a00 : F32), (k * a01 : F32), (k * a02 : F32), (k * a03 : F32),
204      (k * a10 : F32), (k * a11 : F32), (k * a12 : F32), (k * a13 : F32),
205      (k * a20 : F32), (k * a21 : F32), (k * a22 : F32), (k * a23 : F32),
206      (k * a30 : F32), (k * a31 : F32), (k * a32 : F32), (k * a33 : F32)}
207
208def Blk.transpose(m: Blk) -> Blk:
209  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
210  Blk{a00, a10, a20, a30,
211      a01, a11, a21, a31,
212      a02, a12, a22, a32,
213      a03, a13, a23, a33}
214
215# m * n: 64 multiply-adds, each cell as (a0 b0 + a1 b1) + (a2 b2 + a3 b3)
216def Blk.matmul(m: Blk, n: Blk) -> Blk:
217  Blk{+a00, +a01, +a02, +a03, +a10, +a11, +a12, +a13, +a20, +a21, +a22, +a23, +a30, +a31, +a32, +a33} = m
218  Blk{+b00, +b01, +b02, +b03, +b10, +b11, +b12, +b13, +b20, +b21, +b22, +b23, +b30, +b31, +b32, +b33} = n
219  Blk{((a00 * b00 + a01 * b10) + (a02 * b20 + a03 * b30) : F32), ((a00 * b01 + a01 * b11) + (a02 * b21 + a03 * b31) : F32), ((a00 * b02 + a01 * b12) + (a02 * b22 + a03 * b32) : F32), ((a00 * b03 + a01 * b13) + (a02 * b23 + a03 * b33) : F32),
220      ((a10 * b00 + a11 * b10) + (a12 * b20 + a13 * b30) : F32), ((a10 * b01 + a11 * b11) + (a12 * b21 + a13 * b31) : F32), ((a10 * b02 + a11 * b12) + (a12 * b22 + a13 * b32) : F32), ((a10 * b03 + a11 * b13) + (a12 * b23 + a13 * b33) : F32),
221      ((a20 * b00 + a21 * b10) + (a22 * b20 + a23 * b30) : F32), ((a20 * b01 + a21 * b11) + (a22 * b21 + a23 * b31) : F32), ((a20 * b02 + a21 * b12) + (a22 * b22 + a23 * b32) : F32), ((a20 * b03 + a21 * b13) + (a22 * b23 + a23 * b33) : F32),
222      ((a30 * b00 + a31 * b10) + (a32 * b20 + a33 * b30) : F32), ((a30 * b01 + a31 * b11) + (a32 * b21 + a33 * b31) : F32), ((a30 * b02 + a31 * b12) + (a32 * b22 + a33 * b32) : F32), ((a30 * b03 + a31 * b13) + (a32 * b23 + a33 * b33) : F32)}
223
224# m * v: each row dotted with v
225def Blk.matvec(m: Blk, v: Vec4) -> Vec4:
226  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
227  Vec4{+y0, +y1, +y2, +y3} = v
228  Vec4{((a00 * y0 + a01 * y1) + (a02 * y2 + a03 * y3) : F32), ((a10 * y0 + a11 * y1) + (a12 * y2 + a13 * y3) : F32), ((a20 * y0 + a21 * y1) + (a22 * y2 + a23 * y3) : F32), ((a30 * y0 + a31 * y1) + (a32 * y2 + a33 * y3) : F32)}
229
230# row i's sum, as (a_i0 + a_i1) + (a_i2 + a_i3)
231def Blk.sum_rows(m: Blk) -> Vec4:
232  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
233  Vec4{((a00 + a01) + (a02 + a03) : F32), ((a10 + a11) + (a12 + a13) : F32), ((a20 + a21) + (a22 + a23) : F32), ((a30 + a31) + (a32 + a33) : F32)}
234
235# the sum of every cell, as the sum of the row sums
236def Blk.sum(m: Blk) -> F32:
237  Vec4.sum(Blk.sum_rows(m))
238
239# column j's sum, as (a_0j + a_1j) + (a_2j + a_3j): the row sums of the transpose
240def Blk.sum_cols(m: Blk) -> Vec4:
241  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
242  Vec4{((a00 + a10) + (a20 + a30) : F32), ((a01 + a11) + (a21 + a31) : F32), ((a02 + a12) + (a22 + a32) : F32), ((a03 + a13) + (a23 + a33) : F32)}
243
244def Blk.max_rows(m: Blk) -> Vec4:
245  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
246  Vec4{F32.max(F32.max(a00, a01), F32.max(a02, a03)), F32.max(F32.max(a10, a11), F32.max(a12, a13)), F32.max(F32.max(a20, a21), F32.max(a22, a23)), F32.max(F32.max(a30, a31), F32.max(a32, a33))}
247
248# v added to every row
249def Blk.add_row(m: Blk, v: Vec4) -> Blk:
250  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
251  Vec4{+y0, +y1, +y2, +y3} = v
252  Blk{(a00 + y0 : F32), (a01 + y1 : F32), (a02 + y2 : F32), (a03 + y3 : F32),
253      (a10 + y0 : F32), (a11 + y1 : F32), (a12 + y2 : F32), (a13 + y3 : F32),
254      (a20 + y0 : F32), (a21 + y1 : F32), (a22 + y2 : F32), (a23 + y3 : F32),
255      (a30 + y0 : F32), (a31 + y1 : F32), (a32 + y2 : F32), (a33 + y3 : F32)}
256
257# v subtracted from every row
258def Blk.sub_row(m: Blk, v: Vec4) -> Blk:
259  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
260  Vec4{+y0, +y1, +y2, +y3} = v
261  Blk{(a00 - y0 : F32), (a01 - y1 : F32), (a02 - y2 : F32), (a03 - y3 : F32),
262      (a10 - y0 : F32), (a11 - y1 : F32), (a12 - y2 : F32), (a13 - y3 : F32),
263      (a20 - y0 : F32), (a21 - y1 : F32), (a22 - y2 : F32), (a23 - y3 : F32),
264      (a30 - y0 : F32), (a31 - y1 : F32), (a32 - y2 : F32), (a33 - y3 : F32)}
265
266# row i subtracted by v_i (the column vector v broadcast along rows)
267def Blk.sub_col(m: Blk, v: Vec4) -> Blk:
268  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
269  Vec4{+y0, +y1, +y2, +y3} = v
270  Blk{(a00 - y0 : F32), (a01 - y0 : F32), (a02 - y0 : F32), (a03 - y0 : F32),
271      (a10 - y1 : F32), (a11 - y1 : F32), (a12 - y1 : F32), (a13 - y1 : F32),
272      (a20 - y2 : F32), (a21 - y2 : F32), (a22 - y2 : F32), (a23 - y2 : F32),
273      (a30 - y3 : F32), (a31 - y3 : F32), (a32 - y3 : F32), (a33 - y3 : F32)}
274
275# row i scaled by v_i
276def Blk.scale_rows(m: Blk, v: Vec4) -> Blk:
277  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
278  Vec4{+y0, +y1, +y2, +y3} = v
279  Blk{(y0 * a00 : F32), (y0 * a01 : F32), (y0 * a02 : F32), (y0 * a03 : F32),
280      (y1 * a10 : F32), (y1 * a11 : F32), (y1 * a12 : F32), (y1 * a13 : F32),
281      (y2 * a20 : F32), (y2 * a21 : F32), (y2 * a22 : F32), (y2 * a23 : F32),
282      (y3 * a30 : F32), (y3 * a31 : F32), (y3 * a32 : F32), (y3 * a33 : F32)}
283
284# the 16 cells in row-major order, before acc
285def Blk.to_list(m: Blk, acc: List<&2, F32>) -> List<&2, F32>:
286  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
287  a00 <> a01 <> a02 <> a03 <> a10 <> a11 <> a12 <> a13 <> a20 <> a21 <> a22 <> a23 <> a30 <> a31 <> a32 <> a33 <> acc
288
289# a block from its four rows
290def Blk.from_vec4s(v0: Vec4, v1: Vec4, v2: Vec4, v3: Vec4) -> Blk:
291  Vec4{a00, a01, a02, a03} = v0
292  Vec4{a10, a11, a12, a13} = v1
293  Vec4{a20, a21, a22, a23} = v2
294  Vec4{a30, a31, a32, a33} = v3
295  Blk{a00, a01, a02, a03,
296      a10, a11, a12, a13,
297      a20, a21, a22, a23,
298      a30, a31, a32, a33}
299
300# a block from four rows of four; a short row pads with 0.0
301def Blk.from_rows(r0: List<&2, F32>, r1: List<&2, F32>, r2: List<&2, F32>, r3: List<&2, F32>) -> Blk:
302  Blk.from_vec4s(Vec4.from_list(r0), Vec4.from_list(r1), Vec4.from_list(r2), Vec4.from_list(r3))
303
304def Blk.row(m: Blk, i: Nat) -> Vec4:
305  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m
306  match i:
307    case 0n:
308      Vec4{a00, a01, a02, a03}
309    case 1n:
310      Vec4{a10, a11, a12, a13}
311    case 2n:
312      Vec4{a20, a21, a22, a23}
313    case _:
314      Vec4{a30, a31, a32, a33}
315
316# Masks
317# -----
318#
319# k counts the valid elements (or rows, or columns) of a block; the
320# rest are padding, set to 0.0 or left out of a reduction. A k of 0
321# only reaches a reduction whose result the caller discards.
322
323# elements at index >= k set to 0.0
324def Vec4.mask(k: Nat, v: Vec4) -> Vec4:
325  match k v:
326    case 0n Vec4{x0, x1, x2, x3}:
327      Vec4{0.0, 0.0, 0.0, 0.0}
328    case 1n Vec4{x0, x1, x2, x3}:
329      Vec4{x0, 0.0, 0.0, 0.0}
330    case 2n Vec4{x0, x1, x2, x3}:
331      Vec4{x0, x1, 0.0, 0.0}
332    case 3n Vec4{x0, x1, x2, x3}:
333      Vec4{x0, x1, x2, 0.0}
334    case _ v2:
335      v2
336
337# the largest of the first k elements
338def Vec4.max_upto(k: Nat, v: Vec4) -> F32:
339  match k v:
340    case 1n Vec4{x0, x1, x2, x3}:
341      x0
342    case 2n Vec4{x0, x1, x2, x3}:
343      F32.max(x0, x1)
344    case 3n Vec4{x0, x1, x2, x3}:
345      F32.max(F32.max(x0, x1), x2)
346    case _ v2:
347      Vec4.max(v2)
348
349# the smallest of the first k elements
350def Vec4.min_upto(k: Nat, v: Vec4) -> F32:
351  match k v:
352    case 1n Vec4{x0, x1, x2, x3}:
353      x0
354    case 2n Vec4{x0, x1, x2, x3}:
355      F32.min(x0, x1)
356    case 3n Vec4{x0, x1, x2, x3}:
357      F32.min(F32.min(x0, x1), x2)
358    case _ v2:
359      Vec4.min(v2)
360
361# the index and value of the largest of the first k elements
362def Vec4.argmax_upto(k: Nat, v: Vec4) -> U32 & F32:
363  match k v:
364    case 1n Vec4{x0, x1, x2, x3}:
365      (0, x0)
366    case 2n Vec4{x0, x1, x2, x3}:
367      Vec4.argmax.two((0, x0), (1, x1))
368    case 3n Vec4{x0, x1, x2, x3}:
369      Vec4.argmax.two(Vec4.argmax.two((0, x0), (1, x1)), (2, x2))
370    case _ v2:
371      Vec4.argmax(v2)
372
373# rows at index >= k set to 0.0
374def Blk.mask_rows(k: Nat, m: Blk) -> Blk:
375  match k m:
376    case 0n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
377      Blk{0.0, 0.0, 0.0, 0.0,
378      0.0, 0.0, 0.0, 0.0,
379      0.0, 0.0, 0.0, 0.0,
380      0.0, 0.0, 0.0, 0.0}
381    case 1n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
382      Blk{a00, a01, a02, a03,
383      0.0, 0.0, 0.0, 0.0,
384      0.0, 0.0, 0.0, 0.0,
385      0.0, 0.0, 0.0, 0.0}
386    case 2n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
387      Blk{a00, a01, a02, a03,
388      a10, a11, a12, a13,
389      0.0, 0.0, 0.0, 0.0,
390      0.0, 0.0, 0.0, 0.0}
391    case 3n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
392      Blk{a00, a01, a02, a03,
393      a10, a11, a12, a13,
394      a20, a21, a22, a23,
395      0.0, 0.0, 0.0, 0.0}
396    case _ m2:
397      m2
398
399# columns at index >= k set to 0.0
400def Blk.mask_cols(k: Nat, m: Blk) -> Blk:
401  match k m:
402    case 0n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
403      Blk{0.0, 0.0, 0.0, 0.0,
404      0.0, 0.0, 0.0, 0.0,
405      0.0, 0.0, 0.0, 0.0,
406      0.0, 0.0, 0.0, 0.0}
407    case 1n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
408      Blk{a00, 0.0, 0.0, 0.0,
409      a10, 0.0, 0.0, 0.0,
410      a20, 0.0, 0.0, 0.0,
411      a30, 0.0, 0.0, 0.0}
412    case 2n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
413      Blk{a00, a01, 0.0, 0.0,
414      a10, a11, 0.0, 0.0,
415      a20, a21, 0.0, 0.0,
416      a30, a31, 0.0, 0.0}
417    case 3n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
418      Blk{a00, a01, a02, 0.0,
419      a10, a11, a12, 0.0,
420      a20, a21, a22, 0.0,
421      a30, a31, a32, 0.0}
422    case _ m2:
423      m2
424
425# rows at index >= ki and columns at index >= kj set to 0.0
426def Blk.mask(ki: Nat, kj: Nat, m: Blk) -> Blk:
427  Blk.mask_cols(kj, Blk.mask_rows(ki, m))
428
429# each row's largest among its first k columns
430def Blk.max_rows_upto(k: Nat, m: Blk) -> Vec4:
431  match k m:
432    case 1n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
433      Vec4{a00, a10, a20, a30}
434    case 2n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
435      Vec4{F32.max(a00, a01), F32.max(a10, a11), F32.max(a20, a21), F32.max(a30, a31)}
436    case 3n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}:
437      Vec4{F32.max(F32.max(a00, a01), a02), F32.max(F32.max(a10, a11), a12), F32.max(F32.max(a20, a21), a22), F32.max(F32.max(a30, a31), a32)}
438    case _ m2:
439      Blk.max_rows(m2)
440
441# Vectors
442# =======
443#
444# VecT(n) is a balanced binary tree of depth n over Vec4 leaves: 4 * 2^n
445# floats. The depth is a live Nat, so every op matches on it first and
446# then opens the constructor the depth implies, as Base does for Word.
447# Every recursion below a node is a parallel let: the tree is the fork
448# tree, balanced by construction.
449
450law VecT:
451  for n: Nat
452  Data
453
454type VecT.Leaf is Data:
455  VecL{v: Vec4}
456
457type VecT.Node<-p: Nat> is Data:
458  VecN{l: VecT(p), r: VecT(p)}
459
460def VecT(n):
461  match n:
462    case 0n:
463      VecT.Leaf
464    case 1n+p:
465      VecT.Node<p>
466
467# the element count of a VecT(n): 4 * 2^n
468def VecT.size(+n: Nat) -> Nat:
469  match n:
470    case 0n:
471      4n
472    case 1n+p:
473      +h = VecT.size(p)
474      Nat.add(h, h)
475
476def VecT.full(+n: Nat, +x: F32) -> VecT(n):
477  match n:
478    case 0n:
479      VecL{Vec4.full(x)}
480    case 1n+p:
481      l r = VecT.full(p, x) VecT.full(p, x)
482      VecN{l, r}
483
484def VecT.zeros(+n: Nat) -> VecT(n):
485  VecT.full(n, 0.0)
486
487def VecT.ones(+n: Nat) -> VecT(n):
488  VecT.full(n, 1.0)
489
490# x, x + 1, x + 2, ..
491def VecT.arange.from(+n: Nat, +x: F32) -> VecT(n):
492  match n:
493    case 0n:
494      VecL{Vec4.arange(x)}
495    case 1n+p:
496      +h = F32.from_nat(VecT.size(p))
497      l r = VecT.arange.from(p, x) VecT.arange.from(p, (x + h : F32))
498      VecN{l, r}
499
500# 0, 1, 2, ..
501def VecT.arange(+n: Nat) -> VecT(n):
502  VecT.arange.from(n, 0.0)
503
504def VecT.map(~f: F32 -> F32, +n: Nat, v: VecT(n)) -> VecT(n):
505  match n:
506    case 0n:
507      VecL{x} = v
508      VecL{Vec4.map(~f, x)}
509    case 1n+p:
510      VecN{l, r} = v
511      a b = VecT.map(~f, p, l) VecT.map(~f, p, r)
512      VecN{a, b}
513
514def VecT.zip(~f: F32 -> F32 -> F32, +n: Nat, a: VecT(n), b: VecT(n)) -> VecT(n):
515  match n:
516    case 0n:
517      VecL{x} = a
518      VecL{y} = b
519      VecL{Vec4.zip(~f, x, y)}
520    case 1n+p:
521      VecN{al, ar} = a
522      VecN{bl, br} = b
523      l r = VecT.zip(~f, p, al, bl) VecT.zip(~f, p, ar, br)
524      VecN{l, r}
525
526def VecT.add(+n: Nat, a: VecT(n), b: VecT(n)) -> VecT(n):
527  VecT.zip(~F32.add, n, a, b)
528
529def VecT.sub(+n: Nat, a: VecT(n), b: VecT(n)) -> VecT(n):
530  VecT.zip(~F32.sub, n, a, b)
531
532# elementwise product
533def VecT.mul(+n: Nat, a: VecT(n), b: VecT(n)) -> VecT(n):
534  VecT.zip(~F32.mul, n, a, b)
535
536def VecT.exp(+n: Nat, v: VecT(n)) -> VecT(n):
537  VecT.map(~F32.exp, n, v)
538
539def VecT.scale(+n: Nat, +k: F32, v: VecT(n)) -> VecT(n):
540  match n:
541    case 0n:
542      VecL{x} = v
543      VecL{Vec4.scale(k, x)}
544    case 1n+p:
545      VecN{l, r} = v
546      a b = VecT.scale(p, k, l) VecT.scale(p, k, r)
547      VecN{a, b}
548
549# k added to every element
550def VecT.shift(+n: Nat, +k: F32, v: VecT(n)) -> VecT(n):
551  match n:
552    case 0n:
553      VecL{x} = v
554      Vec4{x0, x1, x2, x3} = x
555      VecL{Vec4{(x0 + k : F32), (x1 + k : F32), (x2 + k : F32), (x3 + k : F32)}}
556    case 1n+p:
557      VecN{l, r} = v
558      a b = VecT.shift(p, k, l) VecT.shift(p, k, r)
559      VecN{a, b}
560
561# the sum, left half plus right half
562def VecT.sum(+n: Nat, v: VecT(n)) -> F32:
563  match n:
564    case 0n:
565      VecL{x} = v
566      Vec4.sum(x)
567    case 1n+p:
568      VecN{l, r} = v
569      a b = VecT.sum(p, l) VecT.sum(p, r)
570      (a + b : F32)
571
572def VecT.dot(+n: Nat, a: VecT(n), b: VecT(n)) -> F32:
573  match n:
574    case 0n:
575      VecL{x} = a
576      VecL{y} = b
577      Vec4.dot(x, y)
578    case 1n+p:
579      VecN{al, ar} = a
580      VecN{bl, br} = b
581      l r = VecT.dot(p, al, bl) VecT.dot(p, ar, br)
582      (l + r : F32)
583
584def VecT.max(+n: Nat, v: VecT(n)) -> F32:
585  match n:
586    case 0n:
587      VecL{x} = v
588      Vec4.max(x)
589    case 1n+p:
590      VecN{l, r} = v
591      a b = VecT.max(p, l) VecT.max(p, r)
592      F32.max(a, b)
593
594def VecT.min(+n: Nat, v: VecT(n)) -> F32:
595  match n:
596    case 0n:
597      VecL{x} = v
598      Vec4.min(x)
599    case 1n+p:
600      VecN{l, r} = v
601      a b = VecT.min(p, l) VecT.min(p, r)
602      F32.min(a, b)
603
604# the two halves' winners, the right one's index offset by the half's size
605def VecT.argmax.two(+off: U32, ix: U32 & F32, jy: U32 & F32) -> U32 & F32:
606  (i, x) = ix
607  (j, y) = jy
608  Vec4.argmax.two((i, x), ((j + off : U32), y))
609
610# the index and value of the largest element; the first wins ties
611def VecT.argmax(+n: Nat, v: VecT(n)) -> U32 & F32:
612  match n:
613    case 0n:
614      VecL{x} = v
615      Vec4.argmax(x)
616    case 1n+p:
617      VecN{l, r} = v
618      a b = VecT.argmax(p, l) VecT.argmax(p, r)
619      VecT.argmax.two(U32.from_nat(VecT.size(p)), a, b)
620
621# the elements in order, before acc
622# the tree is read borrowed (+v, matched in place), never consumed: a
623# reader that consumes a shared tree makes the compiler copy it, and
624# bend 2.0.16 miscompiles that copy when the depth is a literal (the
625# bias of a model: memory faults in the loop; issues/shared_vec_loop_*)
626def VecT.to_list.go(+n: Nat, +v: VecT(n), acc: List<&2, F32>) -> List<&2, F32>:
627  match n v:
628    case 0n VecL{x}:
629      Vec4.to_list(x, acc)
630    case 1n+p VecN{l, r}:
631      VecT.to_list.go(p, l, VecT.to_list.go(p, r, acc))
632
633def VecT.to_list(+n: Nat, +v: VecT(n)) -> List<&2, F32>:
634  VecT.to_list.go(n, v, [])
635
636# the first 4 * 2^n elements of xs, padded with 0.0
637def VecT.from_list(+n: Nat, +xs: List<&2, F32>) -> VecT(n):
638  match n:
639    case 0n:
640      VecL{Vec4.from_list(xs)}
641    case 1n+p:
642      +h = VecT.size(p)
643      l r = VecT.from_list(p, List.take(&2, F32, xs, h)) VecT.from_list(p, List.drop(&2, F32, xs, h))
644      VecN{l, r}
645
646def VecT.softmax.fin(+n: Nat, +e: VecT(n)) -> VecT(n):
647  VecT.scale(n, F32.div(1.0, VecT.sum(n, e)), e)
648
649# exp(v - max v) / sum
650def VecT.softmax(+n: Nat, +v: VecT(n)) -> VecT(n):
651  VecT.softmax.fin(n, VecT.exp(n, VecT.shift(n, (0.0 - VecT.max(n, v) : F32), v)))
652
653def VecT.show(+n: Nat, v: VecT(n)) -> String:
654  List.show(~&2, ~F32, ~F32.show.of, VecT.to_list(n, v))
655
656# Matrices
657# ========
658#
659# MatT(r, c) has 4 * 2^r rows and 4 * 2^c columns of F32, as a tree of Blk
660# leaves. While both depths are positive a node is a quadrant split
661# (MatQ); when only rows remain it splits top and bottom (MatR); when
662# only columns remain, left and right (MatC). The block product recurses
663# on that tree, so it is cache-oblivious: every level halves the working
664# set, down to a 4x4 kernel in registers. Operands are + and shared, never
665# copied. The shape is the type: MatT.matmul only accepts MatT(r, k) and
666# MatT(k, c), and answers MatT(r, c).
667
668law MatT:
669  for r: Nat
670  for c: Nat
671  Data
672
673type MatT.Leaf is Data:
674  MatB{b: Blk}
675
676type MatT.Rows<-p: Nat, -c: Nat> is Data:
677  MatR{top: MatT(p, c), bot: MatT(p, c)}
678
679type MatT.Cols<-r: Nat, -q: Nat> is Data:
680  MatC{left: MatT(r, q), right: MatT(r, q)}
681
682type MatT.Quad<-p: Nat, -q: Nat> is Data:
683  MatQ{tl: MatT(p, q), tr: MatT(p, q), bl: MatT(p, q), br: MatT(p, q)}
684
685def MatT(r, c):
686  match r c:
687    case 0n 0n:
688      MatT.Leaf
689    case 1n+p 0n:
690      MatT.Rows<p, 0n>
691    case 0n 1n+q:
692      MatT.Cols<0n, q>
693    case 1n+p 1n+q:
694      MatT.Quad<p, q>
695
696def MatT.full(+r: Nat, +c: Nat, +x: F32) -> MatT(r, c):
697  match r c:
698    case 0n 0n:
699      MatB{Blk.full(x)}
700    case 1n+p 0n:
701      t b = MatT.full(p, 0n, x) MatT.full(p, 0n, x)
702      MatR{t, b}
703    case 0n 1n+q:
704      l r2 = MatT.full(0n, q, x) MatT.full(0n, q, x)
705      MatC{l, r2}
706    case 1n+p 1n+q:
707      a b c2 d = MatT.full(p, q, x) MatT.full(p, q, x) MatT.full(p, q, x) MatT.full(p, q, x)
708      MatQ{a, b, c2, d}
709
710def MatT.zeros(+r: Nat, +c: Nat) -> MatT(r, c):
711  MatT.full(r, c, 0.0)
712
713def MatT.ones(+r: Nat, +c: Nat) -> MatT(r, c):
714  MatT.full(r, c, 1.0)
715
716# the identity of MatT(n, n)
717def MatT.eye(+n: Nat) -> MatT(n, n):
718  match n:
719    case 0n:
720      MatB{Blk.eye()}
721    case 1n+p:
722      a b c d = MatT.eye(p) MatT.zeros(p, p) MatT.zeros(p, p) MatT.eye(p)
723      MatQ{a, b, c, d}
724
725def MatT.map(~f: F32 -> F32, +r: Nat, +c: Nat, m: MatT(r, c)) -> MatT(r, c):
726  match r c:
727    case 0n 0n:
728      MatB{b} = m
729      MatB{Blk.map(~f, b)}
730    case 1n+p 0n:
731      MatR{t, b} = m
732      x y = MatT.map(~f, p, 0n, t) MatT.map(~f, p, 0n, b)
733      MatR{x, y}
734    case 0n 1n+q:
735      MatC{l, r2} = m
736      x y = MatT.map(~f, 0n, q, l) MatT.map(~f, 0n, q, r2)
737      MatC{x, y}
738    case 1n+p 1n+q:
739      MatQ{a, b, c2, d} = m
740      w x y z = MatT.map(~f, p, q, a) MatT.map(~f, p, q, b) MatT.map(~f, p, q, c2) MatT.map(~f, p, q, d)
741      MatQ{w, x, y, z}
742
743def MatT.zip(~f: F32 -> F32 -> F32, +r: Nat, +c: Nat, m: MatT(r, c), n: MatT(r, c)) -> MatT(r, c):
744  match r c:
745    case 0n 0n:
746      MatB{x} = m
747      MatB{y} = n
748      MatB{Blk.zip(~f, x, y)}
749    case 1n+p 0n:
750      MatR{mt, mb} = m
751      MatR{nt, nb} = n
752      x y = MatT.zip(~f, p, 0n, mt, nt) MatT.zip(~f, p, 0n, mb, nb)
753      MatR{x, y}
754    case 0n 1n+q:
755      MatC{ml, mr} = m
756      MatC{nl, nr} = n
757      x y = MatT.zip(~f, 0n, q, ml, nl) MatT.zip(~f, 0n, q, mr, nr)
758      MatC{x, y}
759    case 1n+p 1n+q:
760      MatQ{ma, mb, mc, md} = m
761      MatQ{na, nb, nc, nd} = n
762      w x y z = MatT.zip(~f, p, q, ma, na) MatT.zip(~f, p, q, mb, nb) MatT.zip(~f, p, q, mc, nc) MatT.zip(~f, p, q, md, nd)
763      MatQ{w, x, y, z}
764
765def MatT.add(+r: Nat, +c: Nat, m: MatT(r, c), n: MatT(r, c)) -> MatT(r, c):
766  MatT.zip(~F32.add, r, c, m, n)
767
768def MatT.sub(+r: Nat, +c: Nat, m: MatT(r, c), n: MatT(r, c)) -> MatT(r, c):
769  MatT.zip(~F32.sub, r, c, m, n)
770
771# elementwise (Hadamard) product; the matrix product is MatT.matmul
772def MatT.mul(+r: Nat, +c: Nat, m: MatT(r, c), n: MatT(r, c)) -> MatT(r, c):
773  MatT.zip(~F32.mul, r, c, m, n)
774
775def MatT.exp(+r: Nat, +c: Nat, m: MatT(r, c)) -> MatT(r, c):
776  MatT.map(~F32.exp, r, c, m)
777
778def MatT.scale(+r: Nat, +c: Nat, +k: F32, m: MatT(r, c)) -> MatT(r, c):
779  match r c:
780    case 0n 0n:
781      MatB{b} = m
782      MatB{Blk.scale(k, b)}
783    case 1n+p 0n:
784      MatR{t, b} = m
785      x y = MatT.scale(p, 0n, k, t) MatT.scale(p, 0n, k, b)
786      MatR{x, y}
787    case 0n 1n+q:
788      MatC{l, r2} = m
789      x y = MatT.scale(0n, q, k, l) MatT.scale(0n, q, k, r2)
790      MatC{x, y}
791    case 1n+p 1n+q:
792      MatQ{a, b, c2, d} = m
793      w x y z = MatT.scale(p, q, k, a) MatT.scale(p, q, k, b) MatT.scale(p, q, k, c2) MatT.scale(p, q, k, d)
794      MatQ{w, x, y, z}
795
796def MatT.transpose(+r: Nat, +c: Nat, m: MatT(r, c)) -> MatT(c, r):
797  match r c:
798    case 0n 0n:
799      MatB{b} = m
800      MatB{Blk.transpose(b)}
801    case 1n+p 0n:
802      MatR{t, b} = m
803      x y = MatT.transpose(p, 0n, t) MatT.transpose(p, 0n, b)
804      MatC{x, y}
805    case 0n 1n+q:
806      MatC{l, r2} = m
807      x y = MatT.transpose(0n, q, l) MatT.transpose(0n, q, r2)
808      MatR{x, y}
809    case 1n+p 1n+q:
810      MatQ{a, b, c2, d} = m
811      w x y z = MatT.transpose(p, q, a) MatT.transpose(p, q, c2) MatT.transpose(p, q, b) MatT.transpose(p, q, d)
812      MatQ{w, x, y, z}
813
814# the four quadrant sums of a product node, in one balanced fork
815def MatT.add4(+p: Nat, +q: Nat, a0: MatT(p, q), b0: MatT(p, q), a1: MatT(p, q), b1: MatT(p, q),
816             a2: MatT(p, q), b2: MatT(p, q), a3: MatT(p, q), b3: MatT(p, q)) -> MatT(1n+p, 1n+q):
817  w x y z = MatT.add(p, q, a0, b0) MatT.add(p, q, a1, b1) MatT.add(p, q, a2, b2) MatT.add(p, q, a3, b3)
818  MatQ{w, x, y, z}
819
820def MatT.add2r(+p: Nat, a0: MatT(p, 0n), b0: MatT(p, 0n), a1: MatT(p, 0n), b1: MatT(p, 0n)) -> MatT(1n+p, 0n):
821  x y = MatT.add(p, 0n, a0, b0) MatT.add(p, 0n, a1, b1)
822  MatR{x, y}
823
824def MatT.add2c(+q: Nat, a0: MatT(0n, q), b0: MatT(0n, q), a1: MatT(0n, q), b1: MatT(0n, q)) -> MatT(0n, 1n+q):
825  x y = MatT.add(0n, q, a0, b0) MatT.add(0n, q, a1, b1)
826  MatC{x, y}
827
828# a * b for square operands of one depth: the same recursion with one
829# depth riding instead of three, which measures 3x faster per node
830def MatT.matmul.sq(+d: Nat, +a: MatT(d, d), +b: MatT(d, d)) -> MatT(d, d):
831  match d:
832    case 0n:
833      MatB{x} = a
834      MatB{y} = b
835      MatB{Blk.matmul(x, y)}
836    case 1n+p:
837      MatQ{a0, a1, a2, a3} = a
838      MatQ{b0, b1, b2, b3} = b
839      p0 q0 p1 q1 p2 q2 p3 q3 = MatT.matmul.sq(p, a0, b0) MatT.matmul.sq(p, a1, b2) MatT.matmul.sq(p, a0, b1) MatT.matmul.sq(p, a1, b3) MatT.matmul.sq(p, a2, b0) MatT.matmul.sq(p, a3, b2) MatT.matmul.sq(p, a2, b1) MatT.matmul.sq(p, a3, b3)
840      MatT.add4(p, p, p0, q0, p1, q1, p2, q2, p3, q3)
841
842# a * b for any shapes: MatT(r, k) times MatT(k, c) is MatT(r, c). A node
843# forks its whole burst of products at once (eight at a quadrant), then
844# joins them with a fork of adds. Both operands are shared through +,
845# never copied. Three depths ride every call, which is what the square
846# path below avoids.
847def MatT.matmul.any(+r: Nat, +k: Nat, +c: Nat, +a: MatT(r, k), +b: MatT(k, c)) -> MatT(r, c):
848  match r k c:
849    case 0n 0n 0n:
850      # the leaf goes through the square path on purpose: compiled
851      # without MatT.matmul.sq reachable, this recursion fails at runtime
852      # from depth 2 on (a bend 2.0.16 bug, see README); with it, it runs
853      MatT.matmul.sq(0n, a, b)
854    case 1n+p 0n 0n:
855      MatR{at, ab} = a
856      x y = MatT.matmul.any(p, 0n, 0n, at, b) MatT.matmul.any(p, 0n, 0n, ab, b)
857      MatR{x, y}
858    case 0n 0n 1n+s:
859      MatC{bl, br} = b
860      x y = MatT.matmul.any(0n, 0n, s, a, bl) MatT.matmul.any(0n, 0n, s, a, br)
861      MatC{x, y}
862    case 0n 1n+q 0n:
863      MatC{al, ar} = a
864      MatR{bt, bb} = b
865      x y = MatT.matmul.any(0n, q, 0n, al, bt) MatT.matmul.any(0n, q, 0n, ar, bb)
866      MatT.add(0n, 0n, x, y)
867    case 1n+p 1n+q 0n:
868      MatQ{a0, a1, a2, a3} = a
869      MatR{bt, bb} = b
870      p0 q0 p1 q1 = MatT.matmul.any(p, q, 0n, a0, bt) MatT.matmul.any(p, q, 0n, a1, bb) MatT.matmul.any(p, q, 0n, a2, bt) MatT.matmul.any(p, q, 0n, a3, bb)
871      MatT.add2r(p, p0, q0, p1, q1)
872    case 0n 1n+q 1n+s:
873      MatC{al, ar} = a
874      MatQ{b0, b1, b2, b3} = b
875      p0 q0 p1 q1 = MatT.matmul.any(0n, q, s, al, b0) MatT.matmul.any(0n, q, s, ar, b2) MatT.matmul.any(0n, q, s, al, b1) MatT.matmul.any(0n, q, s, ar, b3)
876      MatT.add2c(s, p0, q0, p1, q1)
877    case 1n+p 0n 1n+s:
878      MatR{at, ab} = a
879      MatC{bl, br} = b
880      w x y z = MatT.matmul.any(p, 0n, s, at, bl) MatT.matmul.any(p, 0n, s, at, br) MatT.matmul.any(p, 0n, s, ab, bl) MatT.matmul.any(p, 0n, s, ab, br)
881      MatQ{w, x, y, z}
882    case 1n+p 1n+q 1n+s:
883      MatQ{a0, a1, a2, a3} = a
884      MatQ{b0, b1, b2, b3} = b
885      p0 q0 p1 q1 p2 q2 p3 q3 = MatT.matmul.any(p, q, s, a0, b0) MatT.matmul.any(p, q, s, a1, b2) MatT.matmul.any(p, q, s, a0, b1) MatT.matmul.any(p, q, s, a1, b3) MatT.matmul.any(p, q, s, a2, b0) MatT.matmul.any(p, q, s, a3, b2) MatT.matmul.any(p, q, s, a2, b1) MatT.matmul.any(p, q, s, a3, b3)
886      MatT.add4(p, s, p0, q0, p1, q1, p2, q2, p3, q3)
887
888# a * b: MatT(r, k) times MatT(k, c) is MatT(r, c). For two square
889# operands of one depth, MatT.matmul.sq is the same recursion with one
890# depth riding instead of three, and 3x faster per node; a runtime
891# dispatch to it is not possible without a cast, and any def that returns
892# or rebuilds a MatT parameter makes the compiler reference-count every
893# MatT node in the program (see AGENTS.md), so the choice is the caller's.
894def MatT.matmul(+r: Nat, +k: Nat, +c: Nat, +a: MatT(r, k), +b: MatT(k, c)) -> MatT(r, c):
895  MatT.matmul.any(r, k, c, a, b)
896
897def VecT.add2(+p: Nat, a0: VecT(p), b0: VecT(p), a1: VecT(p), b1: VecT(p)) -> VecT(1n+p):
898  x y = VecT.add(p, a0, b0) VecT.add(p, a1, b1)
899  VecN{x, y}
900
901# m * v: MatT(r, c) times VecT(c) is VecT(r)
902def MatT.matvec(+r: Nat, +c: Nat, +m: MatT(r, c), +v: VecT(c)) -> VecT(r):
903  match r c:
904    case 0n 0n:
905      MatB{b} = m
906      VecL{x} = v
907      VecL{Blk.matvec(b, x)}
908    case 1n+p 0n:
909      MatR{t, b} = m
910      x y = MatT.matvec(p, 0n, t, v) MatT.matvec(p, 0n, b, v)
911      VecN{x, y}
912    case 0n 1n+q:
913      MatC{l, r2} = m
914      VecN{vl, vr} = v
915      x y = MatT.matvec(0n, q, l, vl) MatT.matvec(0n, q, r2, vr)
916      VecT.add(0n, x, y)
917    case 1n+p 1n+q:
918      MatQ{a, b, c2, d} = m
919      VecN{vl, vr} = v
920      p0 q0 p1 q1 = MatT.matvec(p, q, a, vl) MatT.matvec(p, q, b, vr) MatT.matvec(p, q, c2, vl) MatT.matvec(p, q, d, vr)
921      VecT.add2(p, p0, q0, p1, q1)
922
923# the sum of every cell
924def MatT.sum(+r: Nat, +c: Nat, m: MatT(r, c)) -> F32:
925  match r c:
926    case 0n 0n:
927      MatB{b} = m
928      Blk.sum(b)
929    case 1n+p 0n:
930      MatR{t, b} = m
931      x y = MatT.sum(p, 0n, t) MatT.sum(p, 0n, b)
932      (x + y : F32)
933    case 0n 1n+q:
934      MatC{l, r2} = m
935      x y = MatT.sum(0n, q, l) MatT.sum(0n, q, r2)
936      (x + y : F32)
937    case 1n+p 1n+q:
938      MatQ{a, b, c2, d} = m
939      w x y z = MatT.sum(p, q, a) MatT.sum(p, q, b) MatT.sum(p, q, c2) MatT.sum(p, q, d)
940      ((w + x) + (y + z) : F32)
941
942# each row's sum: VecT(r)
943def MatT.sum_rows(+r: Nat, +c: Nat, m: MatT(r, c)) -> VecT(r):
944  match r c:
945    case 0n 0n:
946      MatB{b} = m
947      VecL{Blk.sum_rows(b)}
948    case 1n+p 0n:
949      MatR{t, b} = m
950      x y = MatT.sum_rows(p, 0n, t) MatT.sum_rows(p, 0n, b)
951      VecN{x, y}
952    case 0n 1n+q:
953      MatC{l, r2} = m
954      x y = MatT.sum_rows(0n, q, l) MatT.sum_rows(0n, q, r2)
955      VecT.add(0n, x, y)
956    case 1n+p 1n+q:
957      MatQ{a, b, c2, d} = m
958      p0 q0 p1 q1 = MatT.sum_rows(p, q, a) MatT.sum_rows(p, q, b) MatT.sum_rows(p, q, c2) MatT.sum_rows(p, q, d)
959      VecT.add2(p, p0, q0, p1, q1)
960
961# each column's sum: VecT(c)
962def MatT.sum_cols(+r: Nat, +c: Nat, m: MatT(r, c)) -> VecT(c):
963  match r c:
964    case 0n 0n:
965      MatB{b} = m
966      VecL{Blk.sum_cols(b)}
967    case 1n+p 0n:
968      MatR{t, b} = m
969      x y = MatT.sum_cols(p, 0n, t) MatT.sum_cols(p, 0n, b)
970      VecT.add(0n, x, y)
971    case 0n 1n+q:
972      MatC{l, r2} = m
973      x y = MatT.sum_cols(0n, q, l) MatT.sum_cols(0n, q, r2)
974      VecN{x, y}
975    case 1n+p 1n+q:
976      MatQ{a, b, c2, d} = m
977      p0 q0 p1 q1 = MatT.sum_cols(p, q, a) MatT.sum_cols(p, q, c2) MatT.sum_cols(p, q, b) MatT.sum_cols(p, q, d)
978      VecT.add2(q, p0, q0, p1, q1)
979
980def VecT.max2(+p: Nat, a0: VecT(p), b0: VecT(p), a1: VecT(p), b1: VecT(p)) -> VecT(1n+p):
981  x y = VecT.zip(~F32.max.of, p, a0, b0) VecT.zip(~F32.max.of, p, a1, b1)
982  VecN{x, y}
983
984# each row's largest element: VecT(r)
985def MatT.max_rows(+r: Nat, +c: Nat, m: MatT(r, c)) -> VecT(r):
986  match r c:
987    case 0n 0n:
988      MatB{b} = m
989      VecL{Blk.max_rows(b)}
990    case 1n+p 0n:
991      MatR{t, b} = m
992      x y = MatT.max_rows(p, 0n, t) MatT.max_rows(p, 0n, b)
993      VecN{x, y}
994    case 0n 1n+q:
995      MatC{l, r2} = m
996      x y = MatT.max_rows(0n, q, l) MatT.max_rows(0n, q, r2)
997      VecT.zip(~F32.max.of, 0n, x, y)
998    case 1n+p 1n+q:
999      MatQ{a, b, c2, d} = m
1000      p0 q0 p1 q1 = MatT.max_rows(p, q, a) MatT.max_rows(p, q, b) MatT.max_rows(p, q, c2) MatT.max_rows(p, q, d)
1001      VecT.max2(p, p0, q0, p1, q1)
1002
1003# v added to every row: v is a VecT(c)
1004def MatT.add_row(+r: Nat, +c: Nat, m: MatT(r, c), +v: VecT(c)) -> MatT(r, c):
1005  match r c:
1006    case 0n 0n:
1007      MatB{b} = m
1008      VecL{x} = v
1009      MatB{Blk.add_row(b, x)}
1010    case 1n+p 0n:
1011      MatR{t, b} = m
1012      x y = MatT.add_row(p, 0n, t, v) MatT.add_row(p, 0n, b, v)
1013      MatR{x, y}
1014    case 0n 1n+q:
1015      MatC{l, r2} = m
1016      VecN{vl, vr} = v
1017      x y = MatT.add_row(0n, q, l, vl) MatT.add_row(0n, q, r2, vr)
1018      MatC{x, y}
1019    case 1n+p 1n+q:
1020      MatQ{a, b, c2, d} = m
1021      VecN{vl, vr} = v
1022      w x y z = MatT.add_row(p, q, a, vl) MatT.add_row(p, q, b, vr) MatT.add_row(p, q, c2, vl) MatT.add_row(p, q, d, vr)
1023      MatQ{w, x, y, z}
1024
1025# row i minus v_i: v is a VecT(r), one scalar per row
1026def MatT.sub_col(+r: Nat, +c: Nat, m: MatT(r, c), +v: VecT(r)) -> MatT(r, c):
1027  match r c:
1028    case 0n 0n:
1029      MatB{b} = m
1030      VecL{x} = v
1031      MatB{Blk.sub_col(b, x)}
1032    case 1n+p 0n:
1033      MatR{t, b} = m
1034      VecN{vt, vb} = v
1035      x y = MatT.sub_col(p, 0n, t, vt) MatT.sub_col(p, 0n, b, vb)
1036      MatR{x, y}
1037    case 0n 1n+q:
1038      MatC{l, r2} = m
1039      x y = MatT.sub_col(0n, q, l, v) MatT.sub_col(0n, q, r2, v)
1040      MatC{x, y}
1041    case 1n+p 1n+q:
1042      MatQ{a, b, c2, d} = m
1043      VecN{vt, vb} = v
1044      w x y z = MatT.sub_col(p, q, a, vt) MatT.sub_col(p, q, b, vt) MatT.sub_col(p, q, c2, vb) MatT.sub_col(p, q, d, vb)
1045      MatQ{w, x, y, z}
1046
1047# row i times v_i: v is a VecT(r), one scalar per row
1048def MatT.scale_rows(+r: Nat, +c: Nat, m: MatT(r, c), +v: VecT(r)) -> MatT(r, c):
1049  match r c:
1050    case 0n 0n:
1051      MatB{b} = m
1052      VecL{x} = v
1053      MatB{Blk.scale_rows(b, x)}
1054    case 1n+p 0n:
1055      MatR{t, b} = m
1056      VecN{vt, vb} = v
1057      x y = MatT.scale_rows(p, 0n, t, vt) MatT.scale_rows(p, 0n, b, vb)
1058      MatR{x, y}
1059    case 0n 1n+q:
1060      MatC{l, r2} = m
1061      x y = MatT.scale_rows(0n, q, l, v) MatT.scale_rows(0n, q, r2, v)
1062      MatC{x, y}
1063    case 1n+p 1n+q:
1064      MatQ{a, b, c2, d} = m
1065      VecN{vt, vb} = v
1066      w x y z = MatT.scale_rows(p, q, a, vt) MatT.scale_rows(p, q, b, vt) MatT.scale_rows(p, q, c2, vb) MatT.scale_rows(p, q, d, vb)
1067      MatQ{w, x, y, z}
1068
1069def MatT.softmax_rows.fin(+r: Nat, +c: Nat, +e: MatT(r, c)) -> MatT(r, c):
1070  MatT.scale_rows(r, c, e, VecT.map(~(s => F32.div(1.0, s)), r, MatT.sum_rows(r, c, e)))
1071
1072# softmax along each row: exp(row - max row) / sum
1073def MatT.softmax_rows(+r: Nat, +c: Nat, +m: MatT(r, c)) -> MatT(r, c):
1074  MatT.softmax_rows.fin(r, c, MatT.exp(r, c, MatT.sub_col(r, c, m, MatT.max_rows(r, c, m))))
1075
1076# Conversions
1077# -----------
1078
1079# the first w of every row
1080def MatT.rows_take(+w: Nat, rows: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
1081  match rows:
1082    case []:
1083      []
1084    case row <> rest:
1085      List.take(&2, F32, row, w) <> MatT.rows_take(w, rest)
1086
1087# every row past its first w
1088def MatT.rows_drop(+w: Nat, rows: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
1089  match rows:
1090    case []:
1091      []
1092    case row <> rest:
1093      List.drop(&2, F32, row, w) <> MatT.rows_drop(w, rest)
1094
1095# the rows of xs side by side with the rows of ys
1096def MatT.rows_cat(xs: List<&2, List<&2, F32>>, ys: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
1097  match xs ys:
1098    case x <> xt y <> yt:
1099      List.append(&2, F32, x, y) <> MatT.rows_cat(xt, yt)
1100    case _ _:
1101      []
1102
1103# a leaf from up to four rows; missing rows and cells pad with 0.0
1104def MatT.from_rows.leaf(rows: List<&2, List<&2, F32>>) -> Blk:
1105  match rows:
1106    case r0 <> r1 <> r2 <> r3 <> rest:
1107      Blk.from_rows(r0, r1, r2, r3)
1108    case r0 <> r1 <> r2 <> []:
1109      Blk.from_rows(r0, r1, r2, [])
1110    case r0 <> r1 <> []:
1111      Blk.from_rows(r0, r1, [], [])
1112    case r0 <> []:
1113      Blk.from_rows(r0, [], [], [])
1114    case []:
1115      Blk.zeros()
1116
1117# a MatT(r, c) from its rows, row-major; short or missing rows pad with 0.0
1118def MatT.from_rows(+r: Nat, +c: Nat, +rows: List<&2, List<&2, F32>>) -> MatT(r, c):
1119  match r c:
1120    case 0n 0n:
1121      MatB{MatT.from_rows.leaf(rows)}
1122    case 1n+p 0n:
1123      +h = VecT.size(p)
1124      t b = MatT.from_rows(p, 0n, List.take(&2, List<&2, F32>, rows, h)) MatT.from_rows(p, 0n, List.drop(&2, List<&2, F32>, rows, h))
1125      MatR{t, b}
1126    case 0n 1n+q:
1127      +w = VecT.size(q)
1128      l r2 = MatT.from_rows(0n, q, MatT.rows_take(w, rows)) MatT.from_rows(0n, q, MatT.rows_drop(w, rows))
1129      MatC{l, r2}
1130    case 1n+p 1n+q:
1131      +h = VecT.size(p)
1132      +w = VecT.size(q)
1133      +top = List.take(&2, List<&2, F32>, rows, h)
1134      +bot = List.drop(&2, List<&2, F32>, rows, h)
1135      a b c2 d = MatT.from_rows(p, q, MatT.rows_take(w, top)) MatT.from_rows(p, q, MatT.rows_drop(w, top)) MatT.from_rows(p, q, MatT.rows_take(w, bot)) MatT.from_rows(p, q, MatT.rows_drop(w, bot))
1136      MatQ{a, b, c2, d}
1137
1138# the rows of a MatT(r, c), row-major
1139def MatT.to_rows(+r: Nat, +c: Nat, m: MatT(r, c)) -> List<&2, List<&2, F32>>:
1140  match r c:
1141    case 0n 0n:
1142      MatB{b} = m
1143      Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = b
1144      [[a00, a01, a02, a03], [a10, a11, a12, a13], [a20, a21, a22, a23], [a30, a31, a32, a33]]
1145    case 1n+p 0n:
1146      MatR{t, b} = m
1147      List.append(&2, List<&2, F32>, MatT.to_rows(p, 0n, t), MatT.to_rows(p, 0n, b))
1148    case 0n 1n+q:
1149      MatC{l, r2} = m
1150      MatT.rows_cat(MatT.to_rows(0n, q, l), MatT.to_rows(0n, q, r2))
1151    case 1n+p 1n+q:
1152      MatQ{a, b, c2, d} = m
1153      List.append(&2, List<&2, F32>, MatT.rows_cat(MatT.to_rows(p, q, a), MatT.to_rows(p, q, b)), MatT.rows_cat(MatT.to_rows(p, q, c2), MatT.to_rows(p, q, d)))
1154
1155# xs cut into n rows of w
1156def MatT.chunks(+n: Nat, +w: Nat, +xs: List<&2, F32>) -> List<&2, List<&2, F32>>:
1157  match n:
1158    case 0n:
1159      []
1160    case 1n+p:
1161      List.take(&2, F32, xs, w) <> MatT.chunks(p, w, List.drop(&2, F32, xs, w))
1162
1163# a MatT(r, c) from 4 * 2^r times 4 * 2^c floats in row-major order
1164def MatT.from_list(+r: Nat, +c: Nat, xs: List<&2, F32>) -> MatT(r, c):
1165  +r2 = r
1166  MatT.from_rows(r2, c, MatT.chunks(VecT.size(r2), VecT.size(c), xs))
1167
1168# the cells in row-major order
1169def MatT.to_list(+r: Nat, +c: Nat, m: MatT(r, c)) -> List<&2, F32>:
1170  List.concat(&2, F32, MatT.to_rows(r, c, m))
1171
1172def MatT.show.row(row: List<&2, F32>) -> String:
1173  List.show(~&2, ~F32, ~F32.show.of, row)
1174
1175def MatT.show(+r: Nat, +c: Nat, m: MatT(r, c)) -> String:
1176  List.show(~&2, ~List<&2, F32>, ~MatT.show.row, MatT.to_rows(r, c, m))
1177
1178# Sizes
1179# =====
1180#
1181# The public types. A Vec<m> holds exactly m floats and a Mat<rows, cols>
1182# exactly rows x cols, in a tree of the smallest depth that fits, padded
1183# with 0.0. This is what every dense library does underneath (BLAS packs
1184# edge tiles into padded buffers; the tensor-core paths want multiples of
1185# 8 or 16; the LLM stacks pad sequence lengths), here made explicit and
1186# kept by every op: an op that could give a padding cell a value other
1187# than 0.0 (exp, shift, map, eye) masks its result, and a reduction that
1188# padding would fool (max, min, argmax, softmax) stops at the size.
1189# Sums, dots and products need nothing: a padding cell adds 0.
1190
1191# the smallest n with 4 * 2^n >= m; the fuel is m itself, log2(m) steps
1192def VecT.depth.go(fuel: Nat, +n: Nat, +cap: Nat, +m: Nat, lt: Bool) -> Nat:
1193  match fuel lt:
1194    case _ False{}:
1195      n
1196    case 0n _:
1197      n
1198    case 1n+f True{}:
1199      +cap2 = Nat.add(cap, cap)
1200      VecT.depth.go(f, 1n+n, cap2, m, Nat.is_lt(cap2, m))
1201
1202def VecT.depth(+m: Nat) -> Nat:
1203  VecT.depth.go(m, 0n, 4n, m, Nat.is_lt(4n, m))
1204
1205# Masked tree ops
1206# ---------------
1207# i0 (and j0) is the index of the subtree's first element; m (rows,
1208# cols) the size. A subtree that starts past the size is computed and
1209# discarded, which costs at most the padding's share of the work.
1210
1211# elements at index >= m set to 0.0
1212def VecT.mask(+n: Nat, v: VecT(n), +i0: Nat, +m: Nat) -> VecT(n):
1213  match n:
1214    case 0n:
1215      VecL{x} = v
1216      VecL{Vec4.mask(Nat.sub(m, i0), x)}
1217    case 1n+p:
1218      VecN{l, r} = v
1219      a b = VecT.mask(p, l, i0, m) VecT.mask(p, r, Nat.add(i0, VecT.size(p)), m)
1220      VecN{a, b}
1221
1222# the largest element among the first m
1223def VecT.max_upto(+n: Nat, v: VecT(n), +i0: Nat, +m: Nat) -> F32:
1224  match n:
1225    case 0n:
1226      VecL{x} = v
1227      Vec4.max_upto(Nat.sub(m, i0), x)
1228    case 1n+p:
1229      VecN{l, r} = v
1230      +h = Nat.add(i0, VecT.size(p))
1231      a b = VecT.max_upto(p, l, i0, m) VecT.max_upto(p, r, h, m)
1232      +a2 = a
1233      Bool.pick(F32, Nat.is_lt(h, m), F32.max(a2, b), a2)
1234
1235def VecT.min_upto(+n: Nat, v: VecT(n), +i0: Nat, +m: Nat) -> F32:
1236  match n:
1237    case 0n:
1238      VecL{x} = v
1239      Vec4.min_upto(Nat.sub(m, i0), x)
1240    case 1n+p:
1241      VecN{l, r} = v
1242      +h = Nat.add(i0, VecT.size(p))
1243      a b = VecT.min_upto(p, l, i0, m) VecT.min_upto(p, r, h, m)
1244      +a2 = a
1245      Bool.pick(F32, Nat.is_lt(h, m), F32.min(a2, b), a2)
1246
1247# the right half's winner counts only when the right half starts before m
1248def VecT.argmax_upto.join(lt: Bool, +off: U32, a: U32 & F32, b: U32 & F32) -> U32 & F32:
1249  match lt:
1250    case True{}:
1251      VecT.argmax.two(off, a, b)
1252    case False{}:
1253      a
1254
1255# the index and value of the largest among the first m; the first wins ties
1256def VecT.argmax_upto(+n: Nat, v: VecT(n), +i0: Nat, +m: Nat) -> U32 & F32:
1257  match n:
1258    case 0n:
1259      VecL{x} = v
1260      Vec4.argmax_upto(Nat.sub(m, i0), x)
1261    case 1n+p:
1262      VecN{l, r} = v
1263      +h = Nat.add(i0, VecT.size(p))
1264      a b = VecT.argmax_upto(p, l, i0, m) VecT.argmax_upto(p, r, h, m)
1265      VecT.argmax_upto.join(Nat.is_lt(h, m), U32.from_nat(VecT.size(p)), a, b)
1266
1267# rows at index >= rows and columns at index >= cols set to 0.0
1268def MatT.mask(+r: Nat, +c: Nat, m: MatT(r, c), +i0: Nat, +j0: Nat, +rows: Nat, +cols: Nat) -> MatT(r, c):
1269  match r c:
1270    case 0n 0n:
1271      MatB{b} = m
1272      MatB{Blk.mask(Nat.sub(rows, i0), Nat.sub(cols, j0), b)}
1273    case 1n+p 0n:
1274      MatR{t, b} = m
1275      x y = MatT.mask(p, 0n, t, i0, j0, rows, cols) MatT.mask(p, 0n, b, Nat.add(i0, VecT.size(p)), j0, rows, cols)
1276      MatR{x, y}
1277    case 0n 1n+q:
1278      MatC{l, rt} = m
1279      x y = MatT.mask(0n, q, l, i0, j0, rows, cols) MatT.mask(0n, q, rt, i0, Nat.add(j0, VecT.size(q)), rows, cols)
1280      MatC{x, y}
1281    case 1n+p 1n+q:
1282      MatQ{a, b, c2, d} = m
1283      +i1 = Nat.add(i0, VecT.size(p))
1284      +j1 = Nat.add(j0, VecT.size(q))
1285      w x y z = MatT.mask(p, q, a, i0, j0, rows, cols) MatT.mask(p, q, b, i0, j1, rows, cols) MatT.mask(p, q, c2, i1, j0, rows, cols) MatT.mask(p, q, d, i1, j1, rows, cols)
1286      MatQ{w, x, y, z}
1287
1288# each row's largest among its first cols columns; a padding row answers 0.0
1289def MatT.max_rows_upto(+r: Nat, +c: Nat, m: MatT(r, c), +j0: Nat, +cols: Nat) -> VecT(r):
1290  match r c:
1291    case 0n 0n:
1292      MatB{b} = m
1293      VecL{Blk.max_rows_upto(Nat.sub(cols, j0), b)}
1294    case 1n+p 0n:
1295      MatR{t, b} = m
1296      x y = MatT.max_rows_upto(p, 0n, t, j0, cols) MatT.max_rows_upto(p, 0n, b, j0, cols)
1297      VecN{x, y}
1298    case 0n 1n+q:
1299      MatC{l, rt} = m
1300      +h = Nat.add(j0, VecT.size(q))
1301      x y = MatT.max_rows_upto(0n, q, l, j0, cols) MatT.max_rows_upto(0n, q, rt, h, cols)
1302      +x2 = x
1303      Bool.pick(VecT(0n), Nat.is_lt(h, cols), VecT.zip(~F32.max.of, 0n, x2, y), x2)
1304    case 1n+p 1n+q:
1305      MatQ{a, b, c2, d} = m
1306      +h = Nat.add(j0, VecT.size(q))
1307      p0 q0 p1 q1 = MatT.max_rows_upto(p, q, a, j0, cols) MatT.max_rows_upto(p, q, b, h, cols) MatT.max_rows_upto(p, q, c2, j0, cols) MatT.max_rows_upto(p, q, d, h, cols)
1308      +p2 = p0
1309      +p3 = p1
1310      +lt = Nat.is_lt(h, cols)
1311      VecN{Bool.pick(VecT(p), lt, VecT.zip(~F32.max.of, p, p2, q0), p2), Bool.pick(VecT(p), lt, VecT.zip(~F32.max.of, p, p3, q1), p3)}
1312
1313# Vec<m>
1314# ------
1315
1316type Vec<-m: Nat> is Data:
1317  VecS{tree: VecT(VecT.depth(m))}
1318
1319def Vec.zeros(+m: Nat) -> Vec<m>:
1320  VecS{VecT.zeros(VecT.depth(m))}
1321
1322def Vec.full(+m: Nat, +x: F32) -> Vec<m>:
1323  VecS{VecT.mask(VecT.depth(m), VecT.full(VecT.depth(m), x), 0n, m)}
1324
1325def Vec.ones(+m: Nat) -> Vec<m>:
1326  Vec.full(m, 1.0)
1327
1328# 0, 1, .., m - 1
1329def Vec.arange(+m: Nat) -> Vec<m>:
1330  VecS{VecT.mask(VecT.depth(m), VecT.arange(VecT.depth(m)), 0n, m)}
1331
1332# the first m of xs, padded with 0.0
1333def Vec.from_list(+m: Nat, xs: List<&2, F32>) -> Vec<m>:
1334  VecS{VecT.mask(VecT.depth(m), VecT.from_list(VecT.depth(m), List.take(&2, F32, xs, m)), 0n, m)}
1335
1336# read borrowed: a bias is shared by every pass of a model
1337def Vec.to_list(+m: Nat, +v: Vec<m>) -> List<&2, F32>:
1338  match v:
1339    case VecS{t}:
1340      List.take(&2, F32, VecT.to_list(VecT.depth(m), t), m)
1341
1342def Vec.map(~f: F32 -> F32, +m: Nat, v: Vec<m>) -> Vec<m>:
1343  VecS{t} = v
1344  VecS{VecT.mask(VecT.depth(m), VecT.map(~f, VecT.depth(m), t), 0n, m)}
1345
1346def Vec.zip(~f: F32 -> F32 -> F32, +m: Nat, a: Vec<m>, b: Vec<m>) -> Vec<m>:
1347  VecS{x} = a
1348  VecS{y} = b
1349  VecS{VecT.mask(VecT.depth(m), VecT.zip(~f, VecT.depth(m), x, y), 0n, m)}
1350
1351def Vec.add(+m: Nat, a: Vec<m>, b: Vec<m>) -> Vec<m>:
1352  VecS{x} = a
1353  VecS{y} = b
1354  VecS{VecT.add(VecT.depth(m), x, y)}
1355
1356def Vec.sub(+m: Nat, a: Vec<m>, b: Vec<m>) -> Vec<m>:
1357  VecS{x} = a
1358  VecS{y} = b
1359  VecS{VecT.sub(VecT.depth(m), x, y)}
1360
1361def Vec.mul(+m: Nat, a: Vec<m>, b: Vec<m>) -> Vec<m>:
1362  VecS{x} = a
1363  VecS{y} = b
1364  VecS{VecT.mul(VecT.depth(m), x, y)}
1365
1366def Vec.exp(+m: Nat, v: Vec<m>) -> Vec<m>:
1367  Vec.map(~F32.exp, m, v)
1368
1369def Vec.scale(+m: Nat, +k: F32, v: Vec<m>) -> Vec<m>:
1370  VecS{t} = v
1371  VecS{VecT.scale(VecT.depth(m), k, t)}
1372
1373def Vec.shift(+m: Nat, +k: F32, v: Vec<m>) -> Vec<m>:
1374  VecS{t} = v
1375  VecS{VecT.mask(VecT.depth(m), VecT.shift(VecT.depth(m), k, t), 0n, m)}
1376
1377def Vec.sum(+m: Nat, v: Vec<m>) -> F32:
1378  VecS{t} = v
1379  VecT.sum(VecT.depth(m), t)
1380
1381def Vec.dot(+m: Nat, a: Vec<m>, b: Vec<m>) -> F32:
1382  VecS{x} = a
1383  VecS{y} = b
1384  VecT.dot(VecT.depth(m), x, y)
1385
1386def Vec.max(+m: Nat, v: Vec<m>) -> F32:
1387  VecS{t} = v
1388  VecT.max_upto(VecT.depth(m), t, 0n, m)
1389
1390def Vec.min(+m: Nat, v: Vec<m>) -> F32:
1391  VecS{t} = v
1392  VecT.min_upto(VecT.depth(m), t, 0n, m)
1393
1394def Vec.argmax(+m: Nat, v: Vec<m>) -> U32 & F32:
1395  VecS{t} = v
1396  VecT.argmax_upto(VecT.depth(m), t, 0n, m)
1397
1398def Vec.mask(+m: Nat, v: Vec<m>) -> Vec<m>:
1399  VecS{t} = v
1400  VecS{VecT.mask(VecT.depth(m), t, 0n, m)}
1401
1402def Vec.softmax.fin(+m: Nat, +e: Vec<m>) -> Vec<m>:
1403  Vec.mask(m, Vec.scale(m, F32.div(1.0, Vec.sum(m, e)), e))
1404
1405# exp(v - max v) / sum, over the m elements
1406def Vec.softmax(+m: Nat, +v: Vec<m>) -> Vec<m>:
1407  Vec.softmax.fin(m, Vec.exp(m, Vec.shift(m, (0.0 - Vec.max(m, v) : F32), v)))
1408
1409def Vec.show(+m: Nat, v: Vec<m>) -> String:
1410  List.show(~&2, ~F32, ~F32.show.of, Vec.to_list(m, v))
1411
1412# Mat<rows, cols>
1413# ---------------
1414
1415type Mat<-rows: Nat, -cols: Nat> is Data:
1416  MatS{tree: MatT(VecT.depth(rows), VecT.depth(cols))}
1417
1418def Mat.zeros(+rows: Nat, +cols: Nat) -> Mat<rows, cols>:
1419  MatS{MatT.zeros(VecT.depth(rows), VecT.depth(cols))}
1420
1421def Mat.full(+rows: Nat, +cols: Nat, +x: F32) -> Mat<rows, cols>:
1422  MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.full(VecT.depth(rows), VecT.depth(cols), x), 0n, 0n, rows, cols)}
1423
1424def Mat.ones(+rows: Nat, +cols: Nat) -> Mat<rows, cols>:
1425  Mat.full(rows, cols, 1.0)
1426
1427def Mat.eye(+n: Nat) -> Mat<n, n>:
1428  MatS{MatT.mask(VecT.depth(n), VecT.depth(n), MatT.eye(VecT.depth(n)), 0n, 0n, n, n)}
1429
1430# a matrix from its rows, row-major; short or missing rows pad with 0.0
1431def Mat.from_rows(+rows: Nat, +cols: Nat, rs: List<&2, List<&2, F32>>) -> Mat<rows, cols>:
1432  MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.from_rows(VecT.depth(rows), VecT.depth(cols), MatT.rows_take(cols, List.take(&2, List<&2, F32>, rs, rows))), 0n, 0n, rows, cols)}
1433
1434# a matrix from rows * cols floats in row-major order
1435def Mat.from_list(+rows: Nat, +cols: Nat, xs: List<&2, F32>) -> Mat<rows, cols>:
1436  MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.from_rows(VecT.depth(rows), VecT.depth(cols), MatT.chunks(rows, cols, xs)), 0n, 0n, rows, cols)}
1437
1438def Mat.to_rows(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> List<&2, List<&2, F32>>:
1439  MatS{t} = m
1440  MatT.rows_take(cols, List.take(&2, List<&2, F32>, MatT.to_rows(VecT.depth(rows), VecT.depth(cols), t), rows))
1441
1442def Mat.to_list(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> List<&2, F32>:
1443  List.concat(&2, F32, Mat.to_rows(rows, cols, m))
1444
1445def Mat.map(~f: F32 -> F32, +rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Mat<rows, cols>:
1446  MatS{t} = m
1447  MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.map(~f, VecT.depth(rows), VecT.depth(cols), t), 0n, 0n, rows, cols)}
1448
1449def Mat.zip(~f: F32 -> F32 -> F32, +rows: Nat, +cols: Nat, m: Mat<rows, cols>, n: Mat<rows, cols>) -> Mat<rows, cols>:
1450  MatS{x} = m
1451  MatS{y} = n
1452  MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.zip(~f, VecT.depth(rows), VecT.depth(cols), x, y), 0n, 0n, rows, cols)}
1453
1454def Mat.add(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, n: Mat<rows, cols>) -> Mat<rows, cols>:
1455  MatS{x} = m
1456  MatS{y} = n
1457  MatS{MatT.add(VecT.depth(rows), VecT.depth(cols), x, y)}
1458
1459def Mat.sub(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, n: Mat<rows, cols>) -> Mat<rows, cols>:
1460  MatS{x} = m
1461  MatS{y} = n
1462  MatS{MatT.sub(VecT.depth(rows), VecT.depth(cols), x, y)}
1463
1464# elementwise (Hadamard) product; the matrix product is Mat.matmul
1465def Mat.mul(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, n: Mat<rows, cols>) -> Mat<rows, cols>:
1466  MatS{x} = m
1467  MatS{y} = n
1468  MatS{MatT.mul(VecT.depth(rows), VecT.depth(cols), x, y)}
1469
1470def Mat.exp(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Mat<rows, cols>:
1471  Mat.map(~F32.exp, rows, cols, m)
1472
1473def Mat.scale(+rows: Nat, +cols: Nat, +k: F32, m: Mat<rows, cols>) -> Mat<rows, cols>:
1474  MatS{t} = m
1475  MatS{MatT.scale(VecT.depth(rows), VecT.depth(cols), k, t)}
1476
1477def Mat.transpose(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Mat<cols, rows>:
1478  MatS{t} = m
1479  MatS{MatT.mask(VecT.depth(cols), VecT.depth(rows), MatT.transpose(VecT.depth(rows), VecT.depth(cols), t), 0n, 0n, cols, rows)}
1480
1481# a * b: Mat<rows, k> times Mat<k, cols> is Mat<rows, cols>
1482def Mat.matmul(+rows: Nat, +k: Nat, +cols: Nat, a: Mat<rows, k>, b: Mat<k, cols>) -> Mat<rows, cols>:
1483  MatS{x} = a
1484  MatS{y} = b
1485  MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.matmul(VecT.depth(rows), VecT.depth(k), VecT.depth(cols), x, y), 0n, 0n, rows, cols)}
1486
1487# a * b for two Mat<n, n>: the one-depth recursion, 3x faster per node
1488def Mat.matmul.sq(+n: Nat, a: Mat<n, n>, b: Mat<n, n>) -> Mat<n, n>:
1489  MatS{x} = a
1490  MatS{y} = b
1491  MatS{MatT.mask(VecT.depth(n), VecT.depth(n), MatT.matmul.sq(VecT.depth(n), x, y), 0n, 0n, n, n)}
1492
1493# m * v: Mat<rows, cols> times Vec<cols> is Vec<rows>
1494def Mat.matvec(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, v: Vec<cols>) -> Vec<rows>:
1495  MatS{t} = m
1496  VecS{x} = v
1497  VecS{VecT.mask(VecT.depth(rows), MatT.matvec(VecT.depth(rows), VecT.depth(cols), t, x), 0n, rows)}
1498
1499def Mat.sum(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> F32:
1500  MatS{t} = m
1501  MatT.sum(VecT.depth(rows), VecT.depth(cols), t)
1502
1503def Mat.sum_rows(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Vec<rows>:
1504  MatS{t} = m
1505  VecS{VecT.mask(VecT.depth(rows), MatT.sum_rows(VecT.depth(rows), VecT.depth(cols), t), 0n, rows)}
1506
1507def Mat.sum_cols(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Vec<cols>:
1508  MatS{t} = m
1509  VecS{VecT.mask(VecT.depth(cols), MatT.sum_cols(VecT.depth(rows), VecT.depth(cols), t), 0n, cols)}
1510
1511def Mat.max_rows(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Vec<rows>:
1512  MatS{t} = m
1513  VecS{VecT.mask(VecT.depth(rows), MatT.max_rows_upto(VecT.depth(rows), VecT.depth(cols), t, 0n, cols), 0n, rows)}
1514
1515# v added to every row. Masked: a padding row plus v is v, not 0.0.
1516def Mat.add_row(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, v: Vec<cols>) -> Mat<rows, cols>:
1517  MatS{t} = m
1518  VecS{x} = v
1519  MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.add_row(VecT.depth(rows), VecT.depth(cols), t, x), 0n, 0n, rows, cols)}
1520
1521# row i minus v_i. Masked: a padding column minus v_i is -v_i, not 0.0.
1522def Mat.sub_col(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, v: Vec<rows>) -> Mat<rows, cols>:
1523  MatS{t} = m
1524  VecS{x} = v
1525  MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.sub_col(VecT.depth(rows), VecT.depth(cols), t, x), 0n, 0n, rows, cols)}
1526
1527# row i times v_i. Masked, so its law needs no fact about x * 0.0.
1528def Mat.scale_rows(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, v: Vec<rows>) -> Mat<rows, cols>:
1529  MatS{t} = m
1530  VecS{x} = v
1531  MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.scale_rows(VecT.depth(rows), VecT.depth(cols), t, x), 0n, 0n, rows, cols)}
1532
1533def Mat.mask(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Mat<rows, cols>:
1534  MatS{t} = m
1535  MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), t, 0n, 0n, rows, cols)}
1536
1537def Mat.softmax_rows.fin(+rows: Nat, +cols: Nat, +e: Mat<rows, cols>) -> Mat<rows, cols>:
1538  Mat.mask(rows, cols, Mat.scale_rows(rows, cols, e, Vec.map(~(s => F32.div(1.0, s)), rows, Mat.sum_rows(rows, cols, e))))
1539
1540# softmax along each row over its cols columns
1541def Mat.softmax_rows(+rows: Nat, +cols: Nat, +m: Mat<rows, cols>) -> Mat<rows, cols>:
1542  Mat.softmax_rows.fin(rows, cols, Mat.exp(rows, cols, Mat.sub_col(rows, cols, m, Mat.max_rows(rows, cols, m))))
1543
1544def Mat.show(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> String:
1545  List.show(~&2, ~List<&2, F32>, ~MatT.show.row, Mat.to_rows(rows, cols, m))
1546
1547# a Mat<rows, cols> around a tree, masked: what the IO products below
1548# answer, and what their padding law is about
1549def Mat.from_tree(+rows: Nat, +cols: Nat, t: MatT(VecT.depth(rows), VecT.depth(cols))) -> Mat<rows, cols>:
1550  MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), t, 0n, 0n, rows, cols)}
1551
1552# Padding
1553# =======
1554#
1555# The predicates the padding laws are stated with. Vec4.Padded(k, v) says
1556# the elements of v at index >= k are 0.0, as a nested pair of equalities
1557# (Unit when k covers all four); VecT.Padded and MatT.Padded lift that
1558# over the trees with the same index bookkeeping as the masks, and
1559# Vec.Padded / Mat.Padded state it for the sized types. A proof of one is
1560# a tree of {==}s; the laws in LAWS.bend say every op answers one.
1561
1562def Vec4.Padded(k: Nat, v: Vec4) -> Type:
1563  match k v:
1564    case 0n Vec4{x0, x1, x2, x3}:
1565      {x0 == 0.0 : F32} & {x1 == 0.0 : F32} & {x2 == 0.0 : F32} & {x3 == 0.0 : F32}
1566    case 1n Vec4{x0, x1, x2, x3}:
1567      {x1 == 0.0 : F32} & {x2 == 0.0 : F32} & {x3 == 0.0 : F32}
1568    case 2n Vec4{x0, x1, x2, x3}:
1569      {x2 == 0.0 : F32} & {x3 == 0.0 : F32}
1570    case 3n Vec4{x0, x1, x2, x3}:
1571      {x3 == 0.0 : F32}
1572    case _ _:
1573      Unit
1574
1575def VecT.Padded(+n: Nat, v: VecT(n), +i0: Nat, +m: Nat) -> Type:
1576  match n:
1577    case 0n:
1578      VecL{x} = v
1579      Vec4.Padded(Nat.sub(m, i0), x)
1580    case 1n+p:
1581      VecN{l, r} = v
1582      VecT.Padded(p, l, i0, m) & VecT.Padded(p, r, Nat.add(i0, VecT.size(p)), m)
1583
1584# row i of a block is padded past kj when i < ki, and entirely when not
1585def Blk.kcol(ki: Nat, i: Nat, +kj: Nat) -> Nat:
1586  Bool.pick(Nat, Nat.is_lt(i, ki), kj, 0n)
1587
1588def Blk.Padded(+ki: Nat, +kj: Nat, b: Blk) -> Type:
1589  Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = b
1590  Vec4.Padded(Blk.kcol(ki, 0n, kj), Vec4{a00, a01, a02, a03}) & Vec4.Padded(Blk.kcol(ki, 1n, kj), Vec4{a10, a11, a12, a13}) & Vec4.Padded(Blk.kcol(ki, 2n, kj), Vec4{a20, a21, a22, a23}) & Vec4.Padded(Blk.kcol(ki, 3n, kj), Vec4{a30, a31, a32, a33})
1591
1592def MatT.Padded(+r: Nat, +c: Nat, m: MatT(r, c), +i0: Nat, +j0: Nat, +rows: Nat, +cols: Nat) -> Type:
1593  match r c:
1594    case 0n 0n:
1595      MatB{b} = m
1596      Blk.Padded(Nat.sub(rows, i0), Nat.sub(cols, j0), b)
1597    case 1n+p 0n:
1598      MatR{t, b} = m
1599      MatT.Padded(p, 0n, t, i0, j0, rows, cols) & MatT.Padded(p, 0n, b, Nat.add(i0, VecT.size(p)), j0, rows, cols)
1600    case 0n 1n+q:
1601      MatC{l, rt} = m
1602      MatT.Padded(0n, q, l, i0, j0, rows, cols) & MatT.Padded(0n, q, rt, i0, Nat.add(j0, VecT.size(q)), rows, cols)
1603    case 1n+p 1n+q:
1604      MatQ{a, b, c2, d} = m
1605      (MatT.Padded(p, q, a, i0, j0, rows, cols) & MatT.Padded(p, q, b, i0, Nat.add(j0, VecT.size(q)), rows, cols)) & (MatT.Padded(p, q, c2, Nat.add(i0, VecT.size(p)), j0, rows, cols) & MatT.Padded(p, q, d, Nat.add(i0, VecT.size(p)), Nat.add(j0, VecT.size(q)), rows, cols))
1606
1607def Vec.Padded(+m: Nat, v: Vec<m>) -> Type:
1608  VecS{t} = v
1609  VecT.Padded(VecT.depth(m), t, 0n, m)
1610
1611def Mat.Padded(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Type:
1612  MatS{t} = m
1613  MatT.Padded(VecT.depth(rows), VecT.depth(cols), t, 0n, 0n, rows, cols)
1614
1615# Flat
1616# ====
1617#
1618# Flat<rows, cols> is rows x cols floats in one row-major Array<F32>,
1619# the layout BLAS reads, so a product hands its operands to the platform
1620# sgemm with no conversion and a model's forward pass never leaves flat
1621# memory. The array has the smallest power-of-two slot count that holds
1622# rows * cols; nothing reads past rows * cols. An Array has one owner, so
1623# a Flat is affine (a Type): the elementwise ops consume their operand
1624# and answer a new one, the reductions consume it, a product consumes
1625# its array operands, and Flat.clone copies. Weights live outside the
1626# heap as a Weights (Data, shared, never consumed) and biases as a Vec
1627# (Data too), so a forward pass is one do block. Every loop here is one
1628# flat tail loop, the reads as stage chains. There are no laws about
1629# Flat beyond its shape: BLAS is a black box and proofs about array
1630# cells are out of reach; the tests check it against numpy.
1631
1632# Effects
1633# -------
1634# C = op(A) op(B) on row-major arrays. m, n, k are the result rows, the
1635# result columns and the inner size; ta and tb (0 or 1) transpose an
1636# operand as stored. Answers a fresh array of the smallest power-of-two
1637# size holding m * n. Blas.sgemm is the platform CPU BLAS (Accelerate,
1638# OpenBLAS), opened at runtime. The _w forms take the right operand as
1639# a registry weight (see Weights) instead of an array. The GPU build
1640# (bend_tensors_gpu.bend) adds Blas.sgemm_gpu: Metal Performance Shaders
1641# on macOS, cuBLAS on Linux.
1642def Blas.sgemm(m: U32, n: U32, k: U32, ta: U32, tb: U32, a: Array<F32>, b: Array<F32>) -> IO(Result<&1, &1, U32 & String, Array<F32>>):
1643  import "./effs/sgemm.c"
1644  import "./effs/sgemm.js"
1645
1646def Blas.sgemm_w(m: U32, n: U32, k: U32, ta: U32, tb: U32, a: Array<F32>, id: U32) -> IO(Result<&1, &1, U32 & String, Array<F32>>):
1647  import "./effs/sgemm_w.c"
1648  import "./effs/sgemm_w.js"
1649
1650# n float32 (little-endian, as numpy writes them) from a file, into the
1651# weight registry (answers the id) or into a fresh array
1652def Weights.load.eff(path: String, n: U32) -> IO(Result<&1, &1, U32 & String, U32>):
1653  import "./effs/weights_load.c"
1654  import "./effs/weights_load.js"
1655
1656def Weights.of_array.eff(a: Array<F32>, n: U32) -> IO(Result<&1, &1, U32 & String, U32>):
1657  import "./effs/weights_of_array.c"
1658  import "./effs/weights_of_array.js"
1659
1660def Flat.load.eff(path: String, n: U32) -> IO(Result<&1, &1, U32 & String, Array<F32>>):
1661  import "./effs/flat_load.c"
1662  import "./effs/flat_load.js"
1663
1664# Weights
1665# -------
1666# A Weights<rows, cols> is rows x cols floats that live outside the
1667# Bend heap, in the BLAS effects' memory, named by an id. It is Data:
1668# share it freely, use it in every pass, it is never copied or consumed.
1669# The GPU effects keep a device copy beside it after the first use.
1670
1671type Weights<-rows: Nat, -cols: Nat> is Data:
1672  W{id: U32}
1673
1674def Weights.wrap(-rows: Nat, -cols: Nat, id: U32) -> IO(Weights<rows, cols>):
1675  IO.pure(Weights<rows, cols>, W{id})
1676
1677# rows * cols float32 from a file, as numpy's tofile writes a C-ordered array
1678def Weights.load(+rows: Nat, +cols: Nat, path: String) -> IO(Weights<rows, cols>):
1679  IO.bind(U32, Weights<rows, cols>, IO.try(U32, Weights.load.eff(path, (U32.from_nat(rows) * U32.from_nat(cols) : U32))), id => Weights.wrap(rows, cols, id))
1680
1681# Flat
1682# ----
1683
1684type Flat<-rows: Nat, -cols: Nat> is Type:
1685  FlatS{a: Array<F32>}
1686
1687# the smallest d with 2^d >= n; the fuel bounds the doubling steps
1688def Flat.depth.go(fuel: Nat, +d: Nat, +cap: Nat, +n: Nat, lt: Bool) -> Nat:
1689  match fuel lt:
1690    case _ False{}:
1691      d
1692    case 0n _:
1693      d
1694    case 1n+f True{}:
1695      +cap2 = Nat.add(cap, cap)
1696      Flat.depth.go(f, 1n+d, cap2, n, Nat.is_lt(cap2, n))
1697
1698def Flat.depth(+n: Nat) -> Nat:
1699  Flat.depth.go(48n, 0n, 1n, n, Nat.is_lt(1n, n))
1700
1701def Flat.count(+rows: Nat, +cols: Nat) -> Nat:
1702  Nat.mul(rows, cols)
1703
1704# Construction
1705# ------------
1706
1707def Flat.full(+rows: Nat, +cols: Nat, +x: F32) -> Flat<rows, cols>:
1708  FlatS{Array.new(F32, Flat.depth(Flat.count(rows, cols)), x)}
1709
1710def Flat.zeros(+rows: Nat, +cols: Nat) -> Flat<rows, cols>:
1711  Flat.full(rows, cols, 0.0)
1712
1713def Flat.ones(+rows: Nat, +cols: Nat) -> Flat<rows, cols>:
1714  Flat.full(rows, cols, 1.0)
1715
1716def Flat.fill.go(xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
1717  match xs:
1718    case []:
1719      a
1720    case x <> rest:
1721      Flat.fill.go(rest, (i + 1 : U32), Array.set(F32, a, i, x))
1722
1723# the first rows * cols of xs, row-major, the rest 0.0
1724def Flat.from_list(+rows: Nat, +cols: Nat, xs: List<&2, F32>) -> Flat<rows, cols>:
1725  FlatS{Flat.fill.go(List.take(&2, F32, xs, Flat.count(rows, cols)), 0, Array.new(F32, Flat.depth(Flat.count(rows, cols)), 0.0))}
1726
1727def Flat.wrap(-rows: Nat, -cols: Nat, a: Array<F32>) -> IO(Flat<rows, cols>):
1728  IO.pure(Flat<rows, cols>, FlatS{a})
1729
1730# rows * cols float32 from a file, row-major
1731def Flat.from_file(+rows: Nat, +cols: Nat, path: String) -> IO(Flat<rows, cols>):
1732  IO.bind(Array<F32>, Flat<rows, cols>, IO.try(Array<F32>, Flat.load.eff(path, (U32.from_nat(rows) * U32.from_nat(cols) : U32))), a => Flat.wrap(rows, cols, a))
1733
1734# a Flat's floats copied into the weight registry; the Flat is consumed
1735def Weights.of_flat(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> IO(Weights<rows, cols>):
1736  FlatS{a} = m
1737  IO.bind(U32, Weights<rows, cols>, IO.try(U32, Weights.of_array.eff(a, (U32.from_nat(rows) * U32.from_nat(cols) : U32))), id => Weights.wrap(rows, cols, id))
1738
1739# the cells in row-major order; i counts down from the last
1740def Flat.to_list.go(n: Nat, +i: U32, r: Array<F32> & F32, acc: List<&2, F32>) -> List<&2, F32>:
1741  match n:
1742    case 0n:
1743      (a, x) = r
1744      x <> acc
1745    case 1n+p:
1746      (a, x) = r
1747      Flat.to_list.go(p, (i - 1 : U32), Array.get(F32, a, (i - 1 : U32)), x <> acc)
1748
1749def Flat.to_list.of(n: Nat, +last: U32, a: Array<F32>) -> List<&2, F32>:
1750  match n:
1751    case 0n:
1752      []
1753    case 1n+p:
1754      Flat.to_list.go(p, last, Array.get(F32, a, last), [])
1755
1756def Flat.to_list(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> List<&2, F32>:
1757  FlatS{a} = m
1758  Flat.to_list.of(Flat.count(rows, cols), (U32.from_nat(Flat.count(rows, cols)) - 1 : U32), a)
1759
1760def Vec.wrap(+m: Nat, a: Array<F32>) -> IO(Vec<m>):
1761  IO.pure(Vec<m>, Vec.from_list(m, Flat.to_list.of(m, (U32.from_nat(m) - 1 : U32), a)))
1762
1763# m float32 from a file, as a Vec (Data: share it, it is never consumed)
1764def Vec.from_file(+m: Nat, path: String) -> IO(Vec<m>):
1765  IO.bind(Array<F32>, Vec<m>, IO.try(Array<F32>, Flat.load.eff(path, U32.from_nat(m))), a => Vec.wrap(m, a))
1766
1767def Flat.clone.fin(-rows: Nat, -cols: Nat, r: Array<F32> & Array<F32>) -> Flat<rows, cols> & Flat<rows, cols>:
1768  (a, b) = r
1769  (FlatS{a}, FlatS{b})
1770
1771# two copies; open the pair in a def of your own (a do block cannot)
1772def Flat.clone(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols> & Flat<rows, cols>:
1773  FlatS{a} = m
1774  Flat.clone.fin(rows, cols, Array.clone(F32, a))
1775
1776def Flat.clone.fst(+rows: Nat, +cols: Nat, r: Flat<rows, cols> & Flat<rows, cols>) -> Flat<rows, cols>:
1777  (a, b) = r
1778  a
1779
1780def Flat.clone.snd(+rows: Nat, +cols: Nat, r: Flat<rows, cols> & Flat<rows, cols>) -> Flat<rows, cols>:
1781  (a, b) = r
1782  b
1783
1784def Flat.to_rows(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> List<&2, List<&2, F32>>:
1785  MatT.chunks(rows, cols, Flat.to_list(rows, cols, m))
1786
1787def Flat.show(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> String:
1788  List.show(~&2, ~List<&2, F32>, ~MatT.show.row, Flat.to_rows(rows, cols, m))
1789
1790# Elementwise
1791# -----------
1792
1793def Flat.map.step(~f: F32 -> F32, +i: U32, r: Array<F32> & F32) -> Array<F32>:
1794  (a, x) = r
1795  Array.set(F32, a, i, f(x))
1796
1797def Flat.map.go(~f: F32 -> F32, n: Nat, +i: U32, a: Array<F32>) -> Array<F32>:
1798  match n:
1799    case 0n:
1800      a
1801    case 1n+p:
1802      Flat.map.go(~f, p, (i + 1 : U32), Flat.map.step(~f, i, Array.get(F32, a, i)))
1803
1804def Flat.map(~f: F32 -> F32, +rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>:
1805  FlatS{a} = m
1806  FlatS{Flat.map.go(~f, Flat.count(rows, cols), 0, a)}
1807
1808def Flat.exp(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>:
1809  Flat.map(~F32.exp, rows, cols, m)
1810
1811def F32.relu(x: F32) -> F32:
1812  F32.max(0.0, x)
1813
1814def Flat.relu(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>:
1815  Flat.map(~F32.relu, rows, cols, m)
1816
1817# gelu, the tanh form: 0.5 x (1 + tanh(sqrt(2 / pi) (x + 0.044715 x^3)))
1818def F32.gelu(x0: F32) -> F32:
1819  +x = x0
1820  (0.5 * x * (1.0 + F32.tanh((0.7978845608 * (x + 0.044715 * x * x * x : F32) : F32))) : F32)
1821
1822def Flat.gelu(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>:
1823  Flat.map(~F32.gelu, rows, cols, m)
1824
1825def Flat.scale.step(+k: F32, +i: U32, r: Array<F32> & F32) -> Array<F32>:
1826  (a, x) = r
1827  Array.set(F32, a, i, (k * x : F32))
1828
1829def Flat.scale.go(n: Nat, +k: F32, +i: U32, a: Array<F32>) -> Array<F32>:
1830  match n:
1831    case 0n:
1832      a
1833    case 1n+p:
1834      Flat.scale.go(p, k, (i + 1 : U32), Flat.scale.step(k, i, Array.get(F32, a, i)))
1835
1836def Flat.scale(+rows: Nat, +cols: Nat, +k: F32, m: Flat<rows, cols>) -> Flat<rows, cols>:
1837  FlatS{a} = m
1838  FlatS{Flat.scale.go(Flat.count(rows, cols), k, 0, a)}
1839
1840def Flat.zip.step2(~f: F32 -> F32 -> F32, +i: U32, x: F32, a: Array<F32>, r: Array<F32> & F32) -> Array<F32> & Array<F32>:
1841  (b, y) = r
1842  (Array.set(F32, a, i, f(x, y)), b)
1843
1844def Flat.zip.step(~f: F32 -> F32 -> F32, +i: U32, b: Array<F32>, r: Array<F32> & F32) -> Array<F32> & Array<F32>:
1845  (a, x) = r
1846  Flat.zip.step2(~f, i, x, a, Array.get(F32, b, i))
1847
1848def Flat.zip.go(~f: F32 -> F32 -> F32, n: Nat, +i: U32, ab: Array<F32> & Array<F32>) -> Array<F32> & Array<F32>:
1849  match n:
1850    case 0n:
1851      ab
1852    case 1n+p:
1853      (a, b) = ab
1854      Flat.zip.go(~f, p, (i + 1 : U32), Flat.zip.step(~f, i, b, Array.get(F32, a, i)))
1855
1856def Flat.zip.fin(-rows: Nat, -cols: Nat, ab: Array<F32> & Array<F32>) -> Flat<rows, cols>:
1857  (a, b) = ab
1858  FlatS{a}
1859
1860# f over the pairs of cells; both operands are consumed
1861def Flat.zip(~f: F32 -> F32 -> F32, +rows: Nat, +cols: Nat, m: Flat<rows, cols>, n: Flat<rows, cols>) -> Flat<rows, cols>:
1862  FlatS{a} = m
1863  FlatS{b} = n
1864  Flat.zip.fin(rows, cols, Flat.zip.go(~f, Flat.count(rows, cols), 0, (a, b)))
1865
1866def Flat.add(+rows: Nat, +cols: Nat, m: Flat<rows, cols>, n: Flat<rows, cols>) -> Flat<rows, cols>:
1867  Flat.zip(~F32.add, rows, cols, m, n)
1868
1869def Flat.sub(+rows: Nat, +cols: Nat, m: Flat<rows, cols>, n: Flat<rows, cols>) -> Flat<rows, cols>:
1870  Flat.zip(~F32.sub, rows, cols, m, n)
1871
1872# elementwise product; the matrix product is Flat.matmul
1873def Flat.mul(+rows: Nat, +cols: Nat, m: Flat<rows, cols>, n: Flat<rows, cols>) -> Flat<rows, cols>:
1874  Flat.zip(~F32.mul, rows, cols, m, n)
1875
1876# cells above the diagonal (j > i) set to -3.4e38, the most negative
1877# finite float, which a softmax turns into 0 with no inf - inf
1878def Flat.mask_causal.step(+idx: U32, +i: U32, +j: U32, r: Array<F32> & F32) -> Array<F32>:
1879  (a, x) = r
1880  Array.set(F32, a, idx, Bool.pick(F32, U32.is_gt(j, i), (0.0 - 3.4028235e38 : F32), x))
1881
1882def Flat.mask_causal.go(n: Nat, +idx: U32, +colsu: U32, a: Array<F32>) -> Array<F32>:
1883  match n:
1884    case 0n:
1885      a
1886    case 1n+p:
1887      Flat.mask_causal.go(p, (idx + 1 : U32), colsu, Flat.mask_causal.step(idx, U32.div(idx, colsu), U32.mod(idx, colsu), Array.get(F32, a, idx)))
1888
1889def Flat.mask_causal(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>:
1890  FlatS{a} = m
1891  FlatS{Flat.mask_causal.go(Flat.count(rows, cols), 0, U32.from_nat(cols), a)}
1892
1893# Rows
1894# ----
1895# The row ops run an inner loop over one row at a time; the row's start
1896# is base, its length cols.
1897
1898def Flat.add_row.step(+idx: U32, +y: F32, r: Array<F32> & F32) -> Array<F32>:
1899  (a, x) = r
1900  Array.set(F32, a, idx, (x + y : F32))
1901
1902# the bias as a list, walked from its head for every row
1903def Flat.add_row.row(n: Nat, +idx: U32, xs: List<&2, F32>, a: Array<F32>) -> Array<F32>:
1904  match n xs:
1905    case 0n _:
1906      a
1907    case 1n+p Nil{}:
1908      a
1909    case 1n+p Con{y, rest}:
1910      Flat.add_row.row(p, (idx + 1 : U32), rest, Flat.add_row.step(idx, y, Array.get(F32, a, idx)))
1911
1912def Flat.add_row.go(n: Nat, +cols: Nat, +colsu: U32, +base: U32, +xs: List<&2, F32>, a: Array<F32>) -> Array<F32>:
1913  match n:
1914    case 0n:
1915      a
1916    case 1n+p:
1917      Flat.add_row.go(p, cols, colsu, (base + colsu : U32), xs, Flat.add_row.row(cols, base, xs, a))
1918
1919# v added to every row; v is a Vec, Data, so it is shared, not consumed
1920def Flat.add_row(+rows: Nat, +cols: Nat, m: Flat<rows, cols>, +v: Vec<cols>) -> Flat<rows, cols>:
1921  FlatS{a} = m
1922  FlatS{Flat.add_row.go(rows, cols, U32.from_nat(cols), 0, Vec.to_list(cols, v), a)}
1923
1924# the largest cell of a row segment
1925def Flat.rowmax.go(n: Nat, +i: U32, r: Array<F32> & F32, +best: F32) -> Array<F32> & F32:
1926  match n:
1927    case 0n:
1928      (a, x) = r
1929      (a, F32.max(best, x))
1930    case 1n+p:
1931      (a, x) = r
1932      Flat.rowmax.go(p, (i + 1 : U32), Array.get(F32, a, (i + 1 : U32)), F32.max(best, x))
1933
1934# exp(x - mx) written back, the row's sum answered
1935def Flat.rowexp.step(+i: U32, +mx: F32, r: Array<F32> & F32, +acc: F32) -> Array<F32> & F32:
1936  (a, x) = r
1937  +y = F32.exp((x - mx : F32))
1938  (Array.set(F32, a, i, y), (acc + y : F32))
1939
1940def Flat.rowexp.go(n: Nat, +i: U32, +mx: F32, r: Array<F32> & F32) -> Array<F32> & F32:
1941  match n:
1942    case 0n:
1943      r
1944    case 1n+p:
1945      (a, acc) = r
1946      Flat.rowexp.go(p, (i + 1 : U32), mx, Flat.rowexp.step(i, mx, Array.get(F32, a, i), acc))
1947
1948def Flat.rowscale.go(n: Nat, +i: U32, +k: F32, a: Array<F32>) -> Array<F32>:
1949  match n:
1950    case 0n:
1951      a
1952    case 1n+p:
1953      Flat.rowscale.go(p, (i + 1 : U32), k, Flat.scale.step(k, i, Array.get(F32, a, i)))
1954
1955# one row of softmax: max, exp and sum, scale
1956def Flat.softmax.row3(+n: Nat, +base: U32, r: Array<F32> & F32) -> Array<F32>:
1957  (a, s) = r
1958  Flat.rowscale.go(n, base, F32.div(1.0, s), a)
1959
1960def Flat.softmax.row2(+n: Nat, +base: U32, r: Array<F32> & F32) -> Array<F32>:
1961  (a, mx) = r
1962  Flat.softmax.row3(n, base, Flat.rowexp.go(n, base, mx, (a, 0.0)))
1963
1964def Flat.softmax.row(+n: Nat, +base: U32, a: Array<F32>) -> Array<F32>:
1965  match n:
1966    case 0n:
1967      a
1968    case 1n+p:
1969      Flat.softmax.row2(1n+p, base, Flat.rowmax.go(p, base, Array.get(F32, a, base), (0.0 - 3.4028235e38 : F32)))
1970
1971def Flat.softmax.go(n: Nat, +cols: Nat, +colsu: U32, +base: U32, a: Array<F32>) -> Array<F32>:
1972  match n:
1973    case 0n:
1974      a
1975    case 1n+p:
1976      Flat.softmax.go(p, cols, colsu, (base + colsu : U32), Flat.softmax.row(cols, base, a))
1977
1978# softmax along each row
1979def Flat.softmax_rows(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>:
1980  FlatS{a} = m
1981  FlatS{Flat.softmax.go(rows, cols, U32.from_nat(cols), 0, a)}
1982
1983# layer norm over each row: (x - mean) / sqrt(var + eps), var the mean
1984# of the squared deviations; a pass for the sum, one for the squares,
1985# one to write
1986def F32.sqdev(+x: F32, +mean: F32) -> F32:
1987  ((x - mean) * (x - mean) : F32)
1988
1989def Flat.rowsq.go(n: Nat, +i: U32, +mean: F32, r: Array<F32> & F32, +acc: F32) -> Array<F32> & F32:
1990  match n:
1991    case 0n:
1992      (a, x) = r
1993      (a, (acc + F32.sqdev(x, mean) : F32))
1994    case 1n+p:
1995      (a, x) = r
1996      Flat.rowsq.go(p, (i + 1 : U32), mean, Array.get(F32, a, (i + 1 : U32)), (acc + F32.sqdev(x, mean) : F32))
1997
1998def Flat.rowsum.go(n: Nat, +i: U32, r: Array<F32> & F32, +acc: F32) -> Array<F32> & F32:
1999  match n:
2000    case 0n:
2001      (a, x) = r
2002      (a, (acc + x : F32))
2003    case 1n+p:
2004      (a, x) = r
2005      Flat.rowsum.go(p, (i + 1 : U32), Array.get(F32, a, (i + 1 : U32)), (acc + x : F32))
2006
2007def Flat.rownorm.step(+i: U32, +mean: F32, +inv: F32, r: Array<F32> & F32) -> Array<F32>:
2008  (a, x) = r
2009  Array.set(F32, a, i, ((x - mean) * inv : F32))
2010
2011def Flat.rownorm.go(n: Nat, +i: U32, +mean: F32, +inv: F32, a: Array<F32>) -> Array<F32>:
2012  match n:
2013    case 0n:
2014      a
2015    case 1n+p:
2016      Flat.rownorm.go(p, (i + 1 : U32), mean, inv, Flat.rownorm.step(i, mean, inv, Array.get(F32, a, i)))
2017
2018def Flat.layer_norm.row3(+n: Nat, +base: U32, +mean: F32, +eps: F32, +cnt: F32, r: Array<F32> & F32) -> Array<F32>:
2019  (a, sq) = r
2020  Flat.rownorm.go(n, base, mean, F32.div(1.0, F32.sqrt((F32.div(sq, cnt) + eps : F32))), a)
2021
2022def Flat.layer_norm.row2(+n: Nat, +p: Nat, +base: U32, +eps: F32, +cnt: F32, r: Array<F32> & F32) -> Array<F32>:
2023  (a, s) = r
2024  +mean = F32.div(s, cnt)
2025  Flat.layer_norm.row3(n, base, mean, eps, cnt, Flat.rowsq.go(p, base, mean, Array.get(F32, a, base), 0.0))
2026
2027def Flat.layer_norm.row(+n: Nat, +base: U32, +eps: F32, a: Array<F32>) -> Array<F32>:
2028  match n:
2029    case 0n:
2030      a
2031    case 1n+p:
2032      Flat.layer_norm.row2(1n+p, p, base, eps, F32.from_nat(1n+p), Flat.rowsum.go(p, base, Array.get(F32, a, base), 0.0))
2033
2034def Flat.layer_norm.go(n: Nat, +cols: Nat, +colsu: U32, +base: U32, +eps: F32, a: Array<F32>) -> Array<F32>:
2035  match n:
2036    case 0n:
2037      a
2038    case 1n+p:
2039      Flat.layer_norm.go(p, cols, colsu, (base + colsu : U32), eps, Flat.layer_norm.row(cols, base, eps, a))
2040
2041def Flat.layer_norm_rows(+rows: Nat, +cols: Nat, +eps: F32, m: Flat<rows, cols>) -> Flat<rows, cols>:
2042  FlatS{a} = m
2043  FlatS{Flat.layer_norm.go(rows, cols, U32.from_nat(cols), 0, eps, a)}
2044
2045# the index of the largest cell of a row segment; the first wins ties
2046def Flat.rowarg.pick(c: Bool, +i: U32, +x: F32, +j: U32, +y: F32) -> U32 & F32:
2047  match c:
2048    case True{}:
2049      (i, x)
2050    case False{}:
2051      (j, y)
2052
2053# the better of the best so far and cell j
2054def Flat.rowarg.best(+bj: U32, +bx: F32, +j: U32, +x: F32) -> U32 & F32:
2055  Flat.rowarg.pick((bx >= x : F32), bj, bx, j, x)
2056
2057def Flat.rowarg.go(n: Nat, +i: U32, +j: U32, r: Array<F32> & F32, best: U32 & F32) -> Array<F32> & (U32 & F32):
2058  match n:
2059    case 0n:
2060      (a, x) = r
2061      (bj, bx) = best
2062      (a, Flat.rowarg.best(bj, bx, j, x))
2063    case 1n+p:
2064      (a, x) = r
2065      (bj, bx) = best
2066      Flat.rowarg.go(p, (i + 1 : U32), (j + 1 : U32), Array.get(F32, a, (i + 1 : U32)), Flat.rowarg.best(bj, bx, j, x))
2067
2068def Flat.argmax.row(+n: Nat, +base: U32, a: Array<F32>) -> Array<F32> & (U32 & F32):
2069  match n:
2070    case 0n:
2071      (a, (0, 0.0))
2072    case 1n+p:
2073      Flat.rowarg.go(p, base, 0, Array.get(F32, a, base), (0, (0.0 - 3.4028235e38 : F32)))
2074
2075# rows from the last down to the first, so the list comes out in row
2076# order; r is the row just done
2077def Flat.argmax.go(n: Nat, +cols: Nat, +colsu: U32, +base: U32, r: Array<F32> & (U32 & F32), acc: List<&2, U32>) -> List<&2, U32>:
2078  match n:
2079    case 0n:
2080      (a, (j, x)) = r
2081      j <> acc
2082    case 1n+p:
2083      (a, (j, x)) = r
2084      Flat.argmax.go(p, cols, colsu, (base - colsu : U32), Flat.argmax.row(cols, (base - colsu : U32), a), j <> acc)
2085
2086def Flat.argmax.of(+rows: Nat, +cols: Nat, a: Array<F32>) -> List<&2, U32>:
2087  match rows:
2088    case 0n:
2089      []
2090    case 1n+p:
2091      +colsu = U32.from_nat(cols)
2092      +last = (U32.from_nat(p) * colsu : U32)
2093      Flat.argmax.go(p, cols, colsu, last, Flat.argmax.row(cols, last, a), [])
2094
2095# each row's argmax, in row order; the matrix is consumed
2096def Flat.argmax_rows(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> List<&2, U32>:
2097  FlatS{a} = m
2098  Flat.argmax.of(rows, cols, a)
2099
2100# Reductions
2101# ----------
2102# Each consumes the matrix; clone it to keep it.
2103
2104def Flat.sum.go(n: Nat, +i: U32, r: Array<F32> & F32, +acc: F32) -> F32:
2105  match n:
2106    case 0n:
2107      (a, x) = r
2108      (acc + x : F32)
2109    case 1n+p:
2110      (a, x) = r
2111      Flat.sum.go(p, (i + 1 : U32), Array.get(F32, a, (i + 1 : U32)), (acc + x : F32))
2112
2113def Flat.sum.of(n: Nat, a: Array<F32>) -> F32:
2114  match n:
2115    case 0n:
2116      0.0
2117    case 1n+p:
2118      Flat.sum.go(p, 0, Array.get(F32, a, 0), 0.0)
2119
2120def Flat.sum(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> F32:
2121  FlatS{a} = m
2122  Flat.sum.of(Flat.count(rows, cols), a)
2123
2124def Flat.mean(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> F32:
2125  F32.div(Flat.sum(rows, cols, m), F32.from_nat(Flat.count(rows, cols)))
2126
2127def Flat.max.go(n: Nat, +i: U32, r: Array<F32> & F32, +best: F32) -> F32:
2128  match n:
2129    case 0n:
2130      (a, x) = r
2131      F32.max(best, x)
2132    case 1n+p:
2133      (a, x) = r
2134      Flat.max.go(p, (i + 1 : U32), Array.get(F32, a, (i + 1 : U32)), F32.max(best, x))
2135
2136def Flat.max.of(n: Nat, a: Array<F32>) -> F32:
2137  match n:
2138    case 0n:
2139      0.0
2140    case 1n+p:
2141      Flat.max.go(p, 0, Array.get(F32, a, 0), (0.0 - 3.4028235e38 : F32))
2142
2143def Flat.max(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> F32:
2144  FlatS{a} = m
2145  Flat.max.of(Flat.count(rows, cols), a)
2146
2147# Transpose
2148# ---------
2149
2150def Flat.transpose.step2(+dst: U32, x: F32, a: Array<F32>, b: Array<F32>) -> Array<F32> & Array<F32>:
2151  (a, Array.set(F32, b, dst, x))
2152
2153def Flat.transpose.step(+dst: U32, b: Array<F32>, r: Array<F32> & F32) -> Array<F32> & Array<F32>:
2154  (a, x) = r
2155  Flat.transpose.step2(dst, x, a, b)
2156
2157# cell (i, j) of a at idx = i * cols + j goes to j * rows + i in b
2158def Flat.transpose.go(n: Nat, +idx: U32, +rowsu: U32, +colsu: U32, ab: Array<F32> & Array<F32>) -> Array<F32> & Array<F32>:
2159  match n:
2160    case 0n:
2161      ab
2162    case 1n+p:
2163      (a, b) = ab
2164      +i = U32.div(idx, colsu)
2165      +j = U32.mod(idx, colsu)
2166      Flat.transpose.go(p, (idx + 1 : U32), rowsu, colsu, Flat.transpose.step((j * rowsu + i : U32), b, Array.get(F32, a, idx)))
2167
2168def Flat.transpose.fin(-rows: Nat, -cols: Nat, ab: Array<F32> & Array<F32>) -> Flat<cols, rows>:
2169  (a, b) = ab
2170  FlatS{b}
2171
2172def Flat.transpose(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<cols, rows>:
2173  FlatS{a} = m
2174  Flat.transpose.fin(rows, cols, Flat.transpose.go(Flat.count(rows, cols), 0, U32.from_nat(rows), U32.from_nat(cols), (a, Array.new(F32, Flat.depth(Flat.count(rows, cols)), 0.0))))
2175
2176# Products
2177# --------
2178# Flat<rows, k> times Flat<k, cols> is Flat<rows, cols>. The _t forms
2179# take the right operand transposed as stored, a Flat<cols, k> (q times
2180# k^T, with no transpose copy); the _w forms take it as a Weights, never
2181# consumed. Array operands are consumed: clone what you keep, or load
2182# it as a Weights. The GPU build adds the _gpu forms of all four.
2183
2184def Flat.matmul.fin(-rows: Nat, -cols: Nat, c: Array<F32>) -> IO(Flat<rows, cols>):
2185  IO.pure(Flat<rows, cols>, FlatS{c})
2186
2187def Flat.matmul.run(+rows: Nat, +k: Nat, +cols: Nat, ta: U32, tb: U32, x: Array<F32>, y: Array<F32>) -> IO(Flat<rows, cols>):
2188  IO.bind(Array<F32>, Flat<rows, cols>, IO.try(Array<F32>, Blas.sgemm(U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), ta, tb, x, y)), c => Flat.matmul.fin(rows, cols, c))
2189
2190def Flat.matmul(+rows: Nat, +k: Nat, +cols: Nat, a: Flat<rows, k>, b: Flat<k, cols>) -> IO(Flat<rows, cols>):
2191  FlatS{x} = a
2192  FlatS{y} = b
2193  Flat.matmul.run(rows, k, cols, 0, 0, x, y)
2194
2195def Flat.matmul_t(+rows: Nat, +k: Nat, +cols: Nat, a: Flat<rows, k>, b: Flat<cols, k>) -> IO(Flat<rows, cols>):
2196  FlatS{x} = a
2197  FlatS{y} = b
2198  Flat.matmul.run(rows, k, cols, 0, 1, x, y)
2199
2200def Flat.matmul_w.run(+rows: Nat, +k: Nat, +cols: Nat, ta: U32, tb: U32, x: Array<F32>, id: U32) -> IO(Flat<rows, cols>):
2201  IO.bind(Array<F32>, Flat<rows, cols>, IO.try(Array<F32>, Blas.sgemm_w(U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), ta, tb, x, id)), c => Flat.matmul.fin(rows, cols, c))
2202
2203def Flat.matmul_w(+rows: Nat, +k: Nat, +cols: Nat, a: Flat<rows, k>, w: Weights<k, cols>) -> IO(Flat<rows, cols>):
2204  FlatS{x} = a
2205  W{id} = w
2206  Flat.matmul_w.run(rows, k, cols, 0, 0, x, id)
2207
2208def Flat.matmul_wt(+rows: Nat, +k: Nat, +cols: Nat, a: Flat<rows, k>, w: Weights<cols, k>) -> IO(Flat<rows, cols>):
2209  FlatS{x} = a
2210  W{id} = w
2211  Flat.matmul_w.run(rows, k, cols, 0, 1, x, id)
2212
2213# Bridges
2214# -------
2215
2216def Flat.to_mat(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Mat<rows, cols>:
2217  Mat.from_list(rows, cols, Flat.to_list(rows, cols, m))
2218
2219def Mat.to_flat(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Flat<rows, cols>:
2220  Flat.from_list(rows, cols, Mat.to_list(rows, cols, m))