bend_tensors.bend
bend-tensors@0.0.0.1 · 2,220 lines · 0 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))