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	
15	import Base
16	
17	# Base's F32.max and F32.show take + parameters, which a template slot
18	# does not accept; these take plain ones.
19	def F32.max.of(a: F32, b: F32) -> F32:
20	  F32.max(a, b)
21	
22	def 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	
34	type 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	
40	type Vec4 is Data:
41	  Vec4{x0: F32, x1: F32, x2: F32, x3: F32}
42	
43	# Vec4
44	# ----
45	
46	def Vec4.full(+x: F32) -> Vec4:
47	  Vec4{x, x, x, x}
48	
49	def Vec4.zeros() -> Vec4:
50	  Vec4{0.0, 0.0, 0.0, 0.0}
51	
52	# x, x + 1, x + 2, x + 3
53	def Vec4.arange(+x: F32) -> Vec4:
54	  Vec4{x, (x + 1.0 : F32), (x + 2.0 : F32), (x + 3.0 : F32)}
55	
56	def 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	
60	def 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	
65	def 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	
70	def 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	
75	def 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	
80	def 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
86	def Vec4.sum(v: Vec4) -> F32:
87	  Vec4{x0, x1, x2, x3} = v
88	  ((x0 + x1) + (x2 + x3) : F32)
89	
90	def 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	
95	def Vec4.max(v: Vec4) -> F32:
96	  Vec4{x0, x1, x2, x3} = v
97	  F32.max(F32.max(x0, x1), F32.max(x2, x3))
98	
99	def 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
104	def 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	
111	def 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	
118	def 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	
122	def 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
127	def 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	
140	def 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	
147	def 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	
153	def Blk.zeros() -> Blk:
154	  Blk.full(0.0)
155	
156	def 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	
162	def 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	
169	def 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	
177	def 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	
185	def 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	
193	def 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	
201	def 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	
208	def 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)
216	def 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
225	def 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)
231	def 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
236	def 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
240	def 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	
244	def 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
249	def 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
258	def 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)
267	def 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
276	def 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
285	def 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
290	def 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
301	def 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	
304	def 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
324	def 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
338	def 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
350	def 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
362	def 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
374	def 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
400	def 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
426	def 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
430	def 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	
450	law VecT:
451	  for n: Nat
452	  Data
453	
454	type VecT.Leaf is Data:
455	  VecL{v: Vec4}
456	
457	type VecT.Node<-p: Nat> is Data:
458	  VecN{l: VecT(p), r: VecT(p)}
459	
460	def 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
468	def 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	
476	def 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	
484	def VecT.zeros(+n: Nat) -> VecT(n):
485	  VecT.full(n, 0.0)
486	
487	def VecT.ones(+n: Nat) -> VecT(n):
488	  VecT.full(n, 1.0)
489	
490	# x, x + 1, x + 2, ..
491	def 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, ..
501	def VecT.arange(+n: Nat) -> VecT(n):
502	  VecT.arange.from(n, 0.0)
503	
504	def 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	
514	def 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	
526	def VecT.add(+n: Nat, a: VecT(n), b: VecT(n)) -> VecT(n):
527	  VecT.zip(~F32.add, n, a, b)
528	
529	def VecT.sub(+n: Nat, a: VecT(n), b: VecT(n)) -> VecT(n):
530	  VecT.zip(~F32.sub, n, a, b)
531	
532	# elementwise product
533	def VecT.mul(+n: Nat, a: VecT(n), b: VecT(n)) -> VecT(n):
534	  VecT.zip(~F32.mul, n, a, b)
535	
536	def VecT.exp(+n: Nat, v: VecT(n)) -> VecT(n):
537	  VecT.map(~F32.exp, n, v)
538	
539	def 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
550	def 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
562	def 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	
572	def 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	
584	def 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	
594	def 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
605	def 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
611	def 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_*)
626	def 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	
633	def 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
637	def 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	
646	def 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
650	def 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	
653	def 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	
668	law MatT:
669	  for r: Nat
670	  for c: Nat
671	  Data
672	
673	type MatT.Leaf is Data:
674	  MatB{b: Blk}
675	
676	type MatT.Rows<-p: Nat, -c: Nat> is Data:
677	  MatR{top: MatT(p, c), bot: MatT(p, c)}
678	
679	type MatT.Cols<-r: Nat, -q: Nat> is Data:
680	  MatC{left: MatT(r, q), right: MatT(r, q)}
681	
682	type 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	
685	def 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	
696	def 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	
710	def MatT.zeros(+r: Nat, +c: Nat) -> MatT(r, c):
711	  MatT.full(r, c, 0.0)
712	
713	def MatT.ones(+r: Nat, +c: Nat) -> MatT(r, c):
714	  MatT.full(r, c, 1.0)
715	
716	# the identity of MatT(n, n)
717	def 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	
725	def 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	
743	def 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	
765	def 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	
768	def 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
772	def 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	
775	def MatT.exp(+r: Nat, +c: Nat, m: MatT(r, c)) -> MatT(r, c):
776	  MatT.map(~F32.exp, r, c, m)
777	
778	def 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	
796	def 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
815	def 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	
820	def 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	
824	def 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
830	def 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.
847	def 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.
894	def 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	
897	def 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)
902	def 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
924	def 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)
943	def 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)
962	def 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	
980	def 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)
985	def 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)
1004	def 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
1026	def 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
1048	def 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	
1069	def 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
1073	def 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
1080	def 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
1088	def 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
1096	def 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
1104	def 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
1118	def 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
1139	def 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
1156	def 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
1164	def 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
1169	def 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	
1172	def MatT.show.row(row: List<&2, F32>) -> String:
1173	  List.show(~&2, ~F32, ~F32.show.of, row)
1174	
1175	def 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
1192	def 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	
1202	def 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
1212	def 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
1223	def 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	
1235	def 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
1248	def 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
1256	def 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
1268	def 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
1289	def 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	
1316	type Vec<-m: Nat> is Data:
1317	  VecS{tree: VecT(VecT.depth(m))}
1318	
1319	def Vec.zeros(+m: Nat) -> Vec<m>:
1320	  VecS{VecT.zeros(VecT.depth(m))}
1321	
1322	def Vec.full(+m: Nat, +x: F32) -> Vec<m>:
1323	  VecS{VecT.mask(VecT.depth(m), VecT.full(VecT.depth(m), x), 0n, m)}
1324	
1325	def Vec.ones(+m: Nat) -> Vec<m>:
1326	  Vec.full(m, 1.0)
1327	
1328	# 0, 1, .., m - 1
1329	def 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
1333	def 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
1337	def 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	
1342	def 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	
1346	def 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	
1351	def 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	
1356	def 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	
1361	def 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	
1366	def Vec.exp(+m: Nat, v: Vec<m>) -> Vec<m>:
1367	  Vec.map(~F32.exp, m, v)
1368	
1369	def 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	
1373	def 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	
1377	def Vec.sum(+m: Nat, v: Vec<m>) -> F32:
1378	  VecS{t} = v
1379	  VecT.sum(VecT.depth(m), t)
1380	
1381	def 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	
1386	def Vec.max(+m: Nat, v: Vec<m>) -> F32:
1387	  VecS{t} = v
1388	  VecT.max_upto(VecT.depth(m), t, 0n, m)
1389	
1390	def Vec.min(+m: Nat, v: Vec<m>) -> F32:
1391	  VecS{t} = v
1392	  VecT.min_upto(VecT.depth(m), t, 0n, m)
1393	
1394	def Vec.argmax(+m: Nat, v: Vec<m>) -> U32 & F32:
1395	  VecS{t} = v
1396	  VecT.argmax_upto(VecT.depth(m), t, 0n, m)
1397	
1398	def Vec.mask(+m: Nat, v: Vec<m>) -> Vec<m>:
1399	  VecS{t} = v
1400	  VecS{VecT.mask(VecT.depth(m), t, 0n, m)}
1401	
1402	def 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
1406	def 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	
1409	def 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	
1415	type Mat<-rows: Nat, -cols: Nat> is Data:
1416	  MatS{tree: MatT(VecT.depth(rows), VecT.depth(cols))}
1417	
1418	def Mat.zeros(+rows: Nat, +cols: Nat) -> Mat<rows, cols>:
1419	  MatS{MatT.zeros(VecT.depth(rows), VecT.depth(cols))}
1420	
1421	def 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	
1424	def Mat.ones(+rows: Nat, +cols: Nat) -> Mat<rows, cols>:
1425	  Mat.full(rows, cols, 1.0)
1426	
1427	def 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
1431	def 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
1435	def 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	
1438	def 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	
1442	def 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	
1445	def 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	
1449	def 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	
1454	def 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	
1459	def 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
1465	def 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	
1470	def Mat.exp(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Mat<rows, cols>:
1471	  Mat.map(~F32.exp, rows, cols, m)
1472	
1473	def 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	
1477	def 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>
1482	def 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
1488	def 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>
1494	def 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	
1499	def 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	
1503	def 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	
1507	def 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	
1511	def 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.
1516	def 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.
1522	def 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.
1528	def 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	
1533	def 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	
1537	def 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
1541	def 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	
1544	def 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
1549	def 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	
1562	def 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	
1575	def 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
1585	def Blk.kcol(ki: Nat, i: Nat, +kj: Nat) -> Nat:
1586	  Bool.pick(Nat, Nat.is_lt(i, ki), kj, 0n)
1587	
1588	def 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	
1592	def 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	
1607	def Vec.Padded(+m: Nat, v: Vec<m>) -> Type:
1608	  VecS{t} = v
1609	  VecT.Padded(VecT.depth(m), t, 0n, m)
1610	
1611	def 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.
1642	def 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	
1646	def 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
1652	def 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	
1656	def 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	
1660	def 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	
1671	type Weights<-rows: Nat, -cols: Nat> is Data:
1672	  W{id: U32}
1673	
1674	def 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
1678	def 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	
1684	type 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
1688	def 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	
1698	def Flat.depth(+n: Nat) -> Nat:
1699	  Flat.depth.go(48n, 0n, 1n, n, Nat.is_lt(1n, n))
1700	
1701	def Flat.count(+rows: Nat, +cols: Nat) -> Nat:
1702	  Nat.mul(rows, cols)
1703	
1704	# Construction
1705	# ------------
1706	
1707	def Flat.full(+rows: Nat, +cols: Nat, +x: F32) -> Flat<rows, cols>:
1708	  FlatS{Array.new(F32, Flat.depth(Flat.count(rows, cols)), x)}
1709	
1710	def Flat.zeros(+rows: Nat, +cols: Nat) -> Flat<rows, cols>:
1711	  Flat.full(rows, cols, 0.0)
1712	
1713	def Flat.ones(+rows: Nat, +cols: Nat) -> Flat<rows, cols>:
1714	  Flat.full(rows, cols, 1.0)
1715	
1716	def 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
1724	def 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	
1727	def 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
1731	def 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
1735	def 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
1740	def 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	
1749	def 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	
1756	def 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	
1760	def 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)
1764	def 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	
1767	def 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)
1772	def 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	
1776	def Flat.clone.fst(+rows: Nat, +cols: Nat, r: Flat<rows, cols> & Flat<rows, cols>) -> Flat<rows, cols>:
1777	  (a, b) = r
1778	  a
1779	
1780	def Flat.clone.snd(+rows: Nat, +cols: Nat, r: Flat<rows, cols> & Flat<rows, cols>) -> Flat<rows, cols>:
1781	  (a, b) = r
1782	  b
1783	
1784	def 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	
1787	def 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	
1793	def 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	
1797	def 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	
1804	def 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	
1808	def Flat.exp(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>:
1809	  Flat.map(~F32.exp, rows, cols, m)
1810	
1811	def F32.relu(x: F32) -> F32:
1812	  F32.max(0.0, x)
1813	
1814	def 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)))
1818	def 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	
1822	def Flat.gelu(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>:
1823	  Flat.map(~F32.gelu, rows, cols, m)
1824	
1825	def 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	
1829	def 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	
1836	def 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	
1840	def 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	
1844	def 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	
1848	def 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	
1856	def 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
1861	def 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	
1866	def 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	
1869	def 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
1873	def 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
1878	def 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	
1882	def 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	
1889	def 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	
1898	def 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
1903	def 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	
1912	def 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
1920	def 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
1925	def 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
1935	def 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	
1940	def 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	
1948	def 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
1956	def 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	
1960	def 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	
1964	def 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	
1971	def 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
1979	def 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
1986	def F32.sqdev(+x: F32, +mean: F32) -> F32:
1987	  ((x - mean) * (x - mean) : F32)
1988	
1989	def 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	
1998	def Flat.rowsum.go(n: Nat, +i: U32, r: Array<F32> & F32, +acc: F32) -> Array<F32> & F32:
1999	  match n:
2000	    case 0n:
Lines 1-2000 of 2220. More: https://jend.leeeet.dev/read.txt?path=bend-tensors@0.0.0.1/bend_tensors.bend&start_line=2001