1 # bend-ml-tensor-array: tensors over a flat Array, with the shape in the type. 2 # 3 # import bend-ml-tensor-array@0.1.19.0/main.bend as TA 4 # 5 # Mat stores r*c numbers in an Array, row by row (index i*c + j). The 6 # dimensions are erased type parameters: the same guarantee as bend-ml-tensor 7 # (a product with wrong dimensions does not compile), but ~50x faster than lists: 8 # Array.get/set by index, and products in parallel blocks of rows. 9 # 10 # An Array is affine (a single owner): every operation that reads an operand returns it 11 # together with the result (as Array.get does), in records MM, AR, ... Use Mat.clone 12 # when you need two copies. 13 # 14 # What is NOT proved: the invariant "the Array has capacity >= r*c" holds because 15 # the constructors (Mat.zeros, Mat.of_list) allocate the right size; and F32 is 16 # validated by tests against PyTorch (reference/test_tensor_array.py), not by proof. 17 18 import Base 19 20 # --------------------------------------------------------------------- 21 # helper types 22 # --------------------------------------------------------------------- 23 24 type R is Type: 25 R{a: Array, b: Array, x: F32} 26 27 type G is Type: 28 G{a: Array, b: Array, c: Array} 29 30 # A dot product with four partial sums. Each step adds its product to a0 and rotates the sums 31 # (a1, a2, a3, a0), so each sum gets every fourth product and the processor can work on the next 32 # addition before the last one has finished; at the end the sums are added in pairs. The order of 33 # the additions is fixed, so the result is the same on every run and every thread count, but it is 34 # not the order of a single running sum (NOTES.md, exp. 17). 35 def dotr(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, a0: F32, a1: F32, a2: F32, a3: F32, ra: Array & F32, rb: Array & F32) -> R: 36 match k ra rb: 37 case 0n Tuple{a, x} Tuple{b, y}: 38 R{a, b, ((a0 + a1 : F32) + (a2 + a3 : F32) : F32)} 39 case 1n+p Tuple{a, +x} Tuple{b, +y}: 40 dotr(p, (ia + sa : U32), (ib + sb : U32), sa, sb, a1, a2, a3, (a0 + (x * y : F32) : F32), Array.get(F32, a, (ia + sa : U32)), Array.get(F32, b, (ib + sb : U32))) 41 42 # dotr's state between two steps: the four partial sums and the two pending reads 43 type D4 is Type: 44 D4{a0: F32, a1: F32, a2: F32, a3: F32, ra: Array & F32, rb: Array & F32} 45 46 # one step of dotr: adds the pending product to a0, rotates the sums, reads the next pair at ia, ib 47 def dstep(+ia: U32, +ib: U32, s: D4) -> D4: 48 match s: 49 case D4{a0, a1, a2, a3, ra, rb}: 50 match ra rb: 51 case Tuple{a, +x} Tuple{b, +y}: 52 D4{a1, a2, a3, (a0 + (x * y : F32) : F32), Array.get(F32, a, ia), Array.get(F32, b, ib)} 53 54 def dfin(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, s: D4) -> R: 55 match s: 56 case D4{a0, a1, a2, a3, ra, rb}: 57 dotr(k, ia, ib, sa, sb, a0, a1, a2, a3, ra, rb) 58 59 # dotr with its loop unrolled by four: q rounds of four dstep, then the last rem (< 4) steps by dotr itself. 60 # Every product goes to the same sum in the same order as in dotr (four rotations bring the names back), so the 61 # result has dotr's bits (tests/dot4.bend). The compiled loop runs ~16% fewer instructions; one thread on a 62 # 50,257 x 768 matrix · vector goes from 26.7 to 20.0 ms, smaller ones and 16 threads stay within noise (NOTES.md, 63 # 2026-11-01). Eight steps a round were no faster. 64 def dot4(q: Nat, +rem: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, s: D4) -> R: 65 match q: 66 case 0n: 67 dfin(rem, ia, ib, sa, sb, s) 68 case 1n+p: 69 dot4(p, rem, (ia + (4 * sa : U32) : U32), (ib + (4 * sb : U32) : U32), sa, sb, dstep((ia + (4 * sa : U32) : U32), (ib + (4 * sb : U32) : U32), dstep((ia + (3 * sa : U32) : U32), (ib + (3 * sb : U32) : U32), dstep((ia + (2 * sa : U32) : U32), (ib + (2 * sb : U32) : U32), dstep((ia + sa : U32), (ib + sb : U32), s))))) 70 71 def dotg(+k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, acc: F32, ra: Array & F32, rb: Array & F32) -> R: 72 dot4(Nat.div(k, 4n), Nat.sub(k, Nat.mul(4n, Nat.div(k, 4n))), ia, ib, sa, sb, D4{acc, 0.0, 0.0, 0.0, ra, rb}) 73 74 # writes one cell of C: c[idx] = x, or c[idx] = c[idx] + x when ac (an accumulating product, 0.1.10.0). x is computed 75 # exactly as without ac and added once, so C + A·B has the bits of adding a separately computed A·B to C. 76 77 def cwr_add(+idx: U32, +x: F32, r: Array & F32) -> Array: 78 match r: 79 case Tuple{c, +v}: 80 Array.set(F32, c, idx, (v + x : F32)) 81 82 def cwr(ac: Bool, +idx: U32, +x: F32, c: Array) -> Array: 83 match ac: 84 case True{}: 85 cwr_add(idx, x, Array.get(F32, c, idx)) 86 case False{}: 87 Array.set(F32, c, idx, x) 88 89 # row i of C, columns j.. ; r is the product of the current column (already computed) 90 91 def gcols(+ac: Bool, mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, +cidx: U32, c: Array, r: R) -> G: 92 match mleft r: 93 case 0n R{a, b, +x}: 94 G{a, b, cwr(ac, cidx, x, c)} 95 case 1n+p R{a, b, +x}: 96 gcols(ac, p, kn, ia0, (ib0 + bcol : U32), sa, sb, bcol, (cidx + 1 : U32), cwr(ac, cidx, x, c), dotg(kn, ia0, (ib0 + bcol : U32), sa, sb, 0.0, Array.get(F32, a, ia0), Array.get(F32, b, (ib0 + bcol : U32)))) 97 98 def grows(+ac: Bool, nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +mu: U32, g: G) -> G: 99 match nleft g: 100 case 0n G{a, b, c}: 101 G{a, b, c} 102 case 1n+p G{a, b, c}: 103 grows(ac, p, mm1, kn, (i + 1 : U32), arow, sa, sb, bcol, mu, gcols(ac, mm1, kn, (i * arow : U32), 0, sa, sb, bcol, (i * mu : U32), c, dotg(kn, (i * arow : U32), 0, sa, sb, 0.0, Array.get(F32, a, (i * arow : U32)), Array.get(F32, b, 0)))) 104 105 # C (n x m) = A · B with k steps, or C + A · B when ac; c is the output array (with capacity n*m) 106 107 def gemm_ac(+ac: Bool, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 108 grows(ac, n, Nat.sub(m, 1n), k, 0, arow, sa, sb, bcol, U32.from_nat(m), G{a, b, c}) 109 110 def gemm(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 111 gemm_ac(False{}, n, m, k, arow, sa, sb, bcol, a, b, c) 112 113 114 # --------------------------------------------------------------------- 115 # parallel gemm: 2^d blocks of rows of C, each with its own copy of A and B, 116 # returning lists that are written into C at the end 117 # --------------------------------------------------------------------- 118 119 type S3 is Type: 120 S3{a: Array, b: Array, out: List<&2, F32>} 121 122 def lcols(mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, out: List<&2, F32>, r: R) -> S3: 123 match mleft r: 124 case 0n R{a, b, +x}: 125 S3{a, b, x <> out} 126 case 1n+p R{a, b, +x}: 127 lcols(p, kn, ia0, (ib0 + bcol : U32), sa, sb, bcol, x <> out, dotg(kn, ia0, (ib0 + bcol : U32), sa, sb, 0.0, Array.get(F32, a, ia0), Array.get(F32, b, (ib0 + bcol : U32)))) 128 129 def lrows(nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, st: S3) -> S3: 130 match nleft st: 131 case 0n S3{a, b, out}: 132 S3{a, b, out} 133 case 1n+p S3{a, b, out}: 134 lrows(p, mm1, kn, (i + 1 : U32), arow, sa, sb, bcol, lcols(mm1, kn, (i * arow : U32), 0, sa, sb, bcol, out, dotg(kn, (i * arow : U32), 0, sa, sb, 0.0, Array.get(F32, a, (i * arow : U32)), Array.get(F32, b, 0)))) 135 136 def lblock_fin(st: S3) -> List<&2, F32>: 137 match st: 138 case S3{a, b, out}: 139 List.reverse(&2, F32, out) 140 141 def lblock(cnt: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array) -> List<&2, F32>: 142 lblock_fin(lrows(cnt, mm1, kn, i, arow, sa, sb, bcol, S3{a, b, Nil{}})) 143 144 def lpar(d: Nat, +cnt: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, ra: Array & Array, rb: Array & Array) -> List<&2, F32>: 145 match d ra rb: 146 case 0n Tuple{a, a2} Tuple{b, b2}: 147 lblock(cnt, mm1, kn, i, arow, sa, sb, bcol, a, b) 148 case 1n++q Tuple{a, a2} Tuple{b, b2}: 149 x y = lpar(q, Nat.div(cnt, 2n), mm1, kn, i, arow, sa, sb, bcol, Array.clone(F32, a), Array.clone(F32, b)) lpar(q, Nat.sub(cnt, Nat.div(cnt, 2n)), mm1, kn, (i + U32.from_nat(Nat.div(cnt, 2n)) : U32), arow, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2)) 150 List.append(&2, F32, x, y) 151 152 def write_l(xs: List<&2, F32>, +idx: U32, c: Array) -> Array: 153 match xs: 154 case Nil{}: 155 c 156 case Con{h, t}: 157 write_l(t, (idx + 1 : U32), Array.set(F32, c, idx, h)) 158 159 # c[idx..] = list, or c[idx..] += list when ac 160 161 def write_la(+ac: Bool, xs: List<&2, F32>, +idx: U32, c: Array) -> Array: 162 match xs: 163 case Nil{}: 164 c 165 case Con{+h, t}: 166 write_la(ac, t, (idx + 1 : U32), cwr(ac, idx, h, c)) 167 168 def gp3(+ac: Bool, a: Array, b: Array, c: Array, xs: List<&2, F32>) -> G: 169 G{a, b, write_la(ac, xs, 0, c)} 170 171 def gp2(d: Nat, +n: Nat, +mm1: Nat, +kn: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, c: Array, pa: Array & Array, pb: Array & Array) -> G: 172 match pa pb: 173 case Tuple{a1, a2} Tuple{b1, b2}: 174 gp3(False{}, a1, b1, c, lpar(d, n, mm1, kn, 0, arow, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2))) 175 176 # C (n x m) = A · B, with 2^d blocks of rows in parallel; returns A, B (intact) and C 177 178 def gemm_par(d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 179 gp2(d, n, Nat.sub(m, 1n), k, arow, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b)) 180 181 182 # --------------------------------------------------------------------- 183 # parallel gemm by COLUMNS, for n = 1 (matrix · vector): each block computes a 184 # range of columns of C, with its own copy of the operands 185 # --------------------------------------------------------------------- 186 187 def cblock_fin(st: S3) -> List<&2, F32>: 188 match st: 189 case S3{a, b, out}: 190 List.reverse(&2, F32, out) 191 192 # cnt columns starting at j0 (cnt >= 1); row 0 of A is at ia0 193 194 def cblock_go(cnt1: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array) -> List<&2, F32>: 195 cblock_fin(lcols(cnt1, kn, ia0, (j0 * bcol : U32), sa, sb, bcol, Nil{}, dotg(kn, ia0, (j0 * bcol : U32), sa, sb, 0.0, Array.get(F32, a, ia0), Array.get(F32, b, (j0 * bcol : U32))))) 196 197 def cblock(cnt: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array) -> List<&2, F32>: 198 match cnt: 199 case 0n: 200 Nil{} 201 case 1n+p: 202 cblock_go(p, kn, ia0, j0, sa, sb, bcol, a, b) 203 204 def cpar(d: Nat, +cnt: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, ra: Array & Array, rb: Array & Array) -> List<&2, F32>: 205 match d ra rb: 206 case 0n Tuple{a, a2} Tuple{b, b2}: 207 cblock(cnt, kn, ia0, j0, sa, sb, bcol, a, b) 208 case 1n++q Tuple{a, a2} Tuple{b, b2}: 209 x y = cpar(q, Nat.div(cnt, 2n), kn, ia0, j0, sa, sb, bcol, Array.clone(F32, a), Array.clone(F32, b)) cpar(q, Nat.sub(cnt, Nat.div(cnt, 2n)), kn, ia0, (j0 + U32.from_nat(Nat.div(cnt, 2n)) : U32), sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2)) 210 List.append(&2, F32, x, y) 211 212 def cp2(+ac: Bool, d: Nat, +m: Nat, +kn: Nat, +sa: U32, +sb: U32, +bcol: U32, c: Array, pa: Array & Array, pb: Array & Array) -> G: 213 match pa pb: 214 case Tuple{a1, a2} Tuple{b1, b2}: 215 gp3(ac, a1, b1, c, cpar(d, m, kn, 0, 0, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2))) 216 217 # C (1 x m) = A (1 x k) · B, with 2^d blocks of columns in parallel. Since 0.1.12.0 no product calls it (one row runs 218 # the scalar loop, see gemm_any3): it clones both operands at every node. Kept for callers of earlier versions. 219 220 def gemm_par_cols(+ac: Bool, d: Nat, +m: Nat, +k: Nat, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 221 cp2(ac, d, m, k, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b)) 222 223 # --------------------------------------------------------------------- 224 # element-wise operations (in place, by index) 225 # --------------------------------------------------------------------- 226 227 # --------------------------------------------------------------------- 228 # gemm: C[i][j] = sum_p A[i*arow + p*sa] * B[p*sb + j*bcol] 229 # (nn: arow=K, sa=1, sb=M, bcol=1 | nt: B stored m x k: sb=1, bcol=K | tn: A stored k x n: arow=1, sa=N) 230 # --------------------------------------------------------------------- 231 232 # acc += A[ia + t*sa] * B[ib + t*sb] for k steps; ra, rb already carry the first pair 233 234 type P2 is Type: 235 P2{a: Array, b: Array} 236 237 type AL is Type: 238 AL{c: Array, l: List<&2, F32>} 239 240 # list of mleft + 1 elements ending at idx (reads with Array.get, returning the array). It reads from 241 # idx down and conses each number onto acc, so the list comes out in index order with no reverse 242 # (NOTES.md, 2026-11-01): callers pass the last index and its read. 243 244 def read_l(mleft: Nat, +idx: U32, acc: List<&2, F32>, rc: Array & F32) -> AL: 245 match mleft rc: 246 case 0n Tuple{c, +v}: 247 AL{c, v <> acc} 248 case 1n+p Tuple{c, +v}: 249 read_l(p, (idx - 1 : U32), v <> acc, Array.get(F32, c, (idx - 1 : U32))) 250 251 # writes the list starting at idx 252 253 def addl(xs: List<&2, F32>, +idx: U32, rc: Array & F32) -> Array: 254 match xs rc: 255 case Nil{} Tuple{c, v}: 256 c 257 case Con{+h, t} Tuple{c, +v}: 258 addl(t, (idx + 1 : U32), Array.get(F32, Array.set(F32, c, idx, (v + h : F32)), (idx + 1 : U32))) 259 260 # adds a bias (a list of m) to each of the n rows of c (n x m) 261 262 def bias_rows(nleft: Nat, +i: U32, +mu: U32, +bl: List<&2, F32>, c: Array) -> Array: 263 match nleft: 264 case 0n: 265 c 266 case 1n+p: 267 bias_rows(p, (i + 1 : U32), mu, bl, addl(bl, (i * mu : U32), Array.get(F32, c, (i * mu : U32)))) 268 269 # in-place relu on the first n elements 270 271 def relu_l(nleft: Nat, +t: U32, rc: Array & F32) -> Array: 272 match nleft rc: 273 case 0n Tuple{c, v}: 274 c 275 case 1n+p Tuple{c, +v}: 276 relu_l(p, (t + 1 : U32), Array.get(F32, Array.set(F32, c, t, F32.max(v, 0.0)), (t + 1 : U32))) 277 278 def pos2(b: Bool) -> F32: 279 match b: 280 case True{}: 281 1.0 282 case False{}: 283 0.0 284 285 def pos(+y: F32) -> F32: 286 pos2(F32.is_gt(y, 0.0)) 287 288 # d[t] *= (h[t] > 0): the relu gradient; returns both arrays 289 290 def mask_l(nleft: Nat, +t: U32, rd: Array & F32, rh: Array & F32) -> P2: 291 match nleft rd rh: 292 case 0n Tuple{d, x} Tuple{h, y}: 293 P2{d, h} 294 case 1n+p Tuple{d, +x} Tuple{h, +y}: 295 mask_l(p, (t + 1 : U32), Array.get(F32, Array.set(F32, d, t, (x * pos(y) : F32)), (t + 1 : U32)), Array.get(F32, h, (t + 1 : U32))) 296 297 # w[t] -= lr * dw[t] 298 299 def maxl2(xs: List<&2, F32>, cur: F32) -> F32: 300 match xs: 301 case Nil{}: 302 cur 303 case Con{+h, t}: 304 maxl2(t, F32.max(cur, h)) 305 306 def maxl(xs: List<&2, F32>) -> F32: 307 match xs: 308 case Nil{}: 309 0.0 310 case Con{+h, t}: 311 maxl2(t, h) 312 313 def expl(xs: List<&2, F32>, +m: F32) -> List<&2, F32>: 314 match xs: 315 case Nil{}: 316 Nil{} 317 case Con{+h, t}: 318 F32.exp((h - m : F32)) <> expl(t, m) 319 320 def suml(xs: List<&2, F32>, acc: F32) -> F32: 321 match xs: 322 case Nil{}: 323 acc 324 case Con{+h, t}: 325 suml(t, (acc + h : F32)) 326 327 def divl(xs: List<&2, F32>, +s: F32) -> List<&2, F32>: 328 match xs: 329 case Nil{}: 330 Nil{} 331 case Con{+h, t}: 332 (h / s : F32) <> divl(t, s) 333 334 def softmax2(+es: List<&2, F32>) -> List<&2, F32>: 335 divl(es, suml(es, 0.0)) 336 337 def softmax(+xs: List<&2, F32>) -> List<&2, F32>: 338 softmax2(expl(xs, maxl(xs))) 339 340 # p[label] 341 342 def sgd_l(nleft: Nat, +t: U32, +lr: F32, rw: Array & F32, rd: Array & F32) -> P2: 343 match nleft rw rd: 344 case 0n Tuple{w, x} Tuple{d, y}: 345 P2{w, d} 346 case 1n+p Tuple{w, +x} Tuple{d, +y}: 347 sgd_l(p, (t + 1 : U32), lr, Array.get(F32, Array.set(F32, w, t, (x - (lr * y : F32) : F32)), (t + 1 : U32)), Array.get(F32, d, (t + 1 : U32))) 348 349 # --------------------------------------------------------------------- 350 # softmax + cross-entropy per row (lists of 10 elements: negligible cost) 351 # --------------------------------------------------------------------- 352 353 def pick(ps: List<&2, F32>, n: Nat) -> F32: 354 match ps n: 355 case Con{+h, t} 0n: 356 h 357 case Con{h, t} 1n+q: 358 pick(t, q) 359 case Nil{} _: 360 0.0 361 362 # (p - onehot(label)) / n 363 364 def hit(b: Bool) -> F32: 365 match b: 366 case True{}: 367 1.0 368 case False{}: 369 0.0 370 371 def grad_l(ps: List<&2, F32>, +label: Nat, +j: Nat, +nf: F32) -> List<&2, F32>: 372 match ps: 373 case Nil{}: 374 Nil{} 375 case Con{+h, t}: 376 (((h - hit(Nat.is_eq(label, j)) : F32) / nf) : F32) <> grad_l(t, label, 1n+j, nf) 377 378 type L2 is Type: 379 L2{loss: F32, c: Array} 380 381 def ce_row2(+ps: List<&2, F32>, +label: Nat, +nf: F32, +idx: U32, loss: F32, c: Array) -> L2: 382 L2{(loss - F32.log(F32.max(pick(ps, label), 0.000000000001)) : F32), write_l(grad_l(ps, label, 0n, nf), idx, c)} 383 384 def ce_row(al: AL, +label: Nat, +nf: F32, +idx: U32, loss: F32) -> L2: 385 match al: 386 case AL{c, l}: 387 ce_row2(softmax(l), label, nf, idx, loss, c) 388 389 # walks the rows of z2 (n x m): r carries the accumulated loss and the array; at the end the array has become the gradient 390 391 type Mat<-r: Nat, -c: Nat> is Type: 392 Mat{d: Array} 393 394 def ce_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +nf: F32, +i: U32, r: L2) -> L2: 395 match labels r: 396 case Nil{} L2{loss, c}: 397 L2{loss, c} 398 case Con{+y, t} L2{loss, c}: 399 ce_rows(t, mm1, mu, nf, (i + 1 : U32), ce_row(read_l(mm1, ((i * mu : U32) + U32.from_nat(mm1) : U32), Nil{}, Array.get(F32, c, ((i * mu : U32) + U32.from_nat(mm1) : U32))), y, nf, (i * mu : U32), loss)) 400 401 # ---- the same rows over the array itself: no lists ---- 402 # The list version above reads each row into a list and walks it five times. This one walks the row in place with tail 403 # loops, with the same operations in the same order (the max from the first cell, exp(x - max) summed from +0.0 left to 404 # right, p = e / s, the loss term read before the row is overwritten, (p - onehot) / n written back), so the same bits; 405 # MetalBend measured 3.4x fewer cycles on 64 rows of 49,152 (2026-10-25). The list version stays as the reference 406 # tests/ce_array.bend compares it with. 407 408 type AF is Type: 409 AF{c: Array, x: F32} 410 411 # the max of k + 1 cells from idx, from m (the first cell's value, already read in r) 412 def cea_max(k: Nat, +idx: U32, +m: F32, r: Array & F32) -> AF: 413 match k r: 414 case 0n Tuple{c, +x}: 415 AF{c, F32.max(m, x)} 416 case 1n+p Tuple{c, +x}: 417 cea_max(p, (idx + 1 : U32), F32.max(m, x), Array.get(F32, c, (idx + 1 : U32))) 418 419 # cells idx..idx + k become exp(cell - m); returns their sum from s, added left to right 420 def cea_exp(k: Nat, +idx: U32, +m: F32, +s: F32, r: Array & F32) -> AF: 421 match k r: 422 case 0n Tuple{c, +x}: 423 +e = F32.exp((x - m : F32)) 424 AF{Array.set(F32, c, idx, e), (s + e : F32)} 425 case 1n+p Tuple{c, +x}: 426 +e = F32.exp((x - m : F32)) 427 cea_exp(p, (idx + 1 : U32), m, (s + e : F32), Array.get(F32, Array.set(F32, c, idx, e), (idx + 1 : U32))) 428 429 # cell idx (column j) becomes ((e / s) - onehot) / nf, for k + 1 cells 430 def cea_grad(k: Nat, +idx: U32, +j: Nat, +label: Nat, +s: F32, +nf: F32, r: Array & F32) -> Array: 431 match k r: 432 case 0n Tuple{c, +x}: 433 Array.set(F32, c, idx, ((((x / s : F32) - hit(Nat.is_eq(label, j)) : F32) / nf) : F32)) 434 case 1n+p Tuple{c, +x}: 435 cea_grad(p, (idx + 1 : U32), 1n+j, label, s, nf, Array.get(F32, Array.set(F32, c, idx, ((((x / s : F32) - hit(Nat.is_eq(label, j)) : F32) / nf) : F32)), (idx + 1 : U32))) 436 437 def cea_p2(+s: F32, r: Array & F32) -> AF: 438 match r: 439 case Tuple{c, +e}: 440 AF{c, (e / s : F32)} 441 442 # p[label] = e[label] / s, or 0.0 for a label past the row (pick's Nil case) 443 def cea_p(inside: Bool, +o: U32, +label: Nat, +s: F32, c: Array) -> AF: 444 match inside: 445 case True{}: 446 cea_p2(s, Array.get(F32, c, (o + U32.from_nat(label) : U32))) 447 case False{}: 448 AF{c, 0.0} 449 450 def cea_row3(+mm1: Nat, +o: U32, +label: Nat, +s: F32, +nf: F32, +loss: F32, q: AF) -> L2: 451 match q: 452 case AF{c, +p}: 453 L2{(loss - F32.log(F32.max(p, 0.000000000001)) : F32), cea_grad(mm1, o, 0n, label, s, nf, Array.get(F32, c, o))} 454 455 def cea_row2(+mm1: Nat, +o: U32, +label: Nat, +nf: F32, +loss: F32, q: AF) -> L2: 456 match q: 457 case AF{c, +s}: 458 cea_row3(mm1, o, label, s, nf, loss, cea_p(Nat.is_le(label, mm1), o, label, s, c)) 459 460 def cea_row1(+mm1: Nat, +o: U32, +label: Nat, +nf: F32, +loss: F32, q: AF) -> L2: 461 match q: 462 case AF{c, +m}: 463 cea_row2(mm1, o, label, nf, loss, cea_exp(mm1, o, m, 0.0, Array.get(F32, c, o))) 464 465 # a row's max starts from its first cell (maxl), then folds over the k + 1 cells from it 466 def cea_row0(+mm1: Nat, +o: U32, +label: Nat, +nf: F32, +loss: F32, r: Array & F32) -> L2: 467 match r: 468 case Tuple{c, +x0}: 469 cea_row1(mm1, o, label, nf, loss, cea_max(mm1, o, x0, Array.get(F32, c, o))) 470 471 # one row per label, in order, as ce_rows 472 def cea_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +nf: F32, +i: U32, r: L2) -> L2: 473 match labels r: 474 case Nil{} L2{loss, c}: 475 L2{loss, c} 476 case Con{+y, t} L2{+loss, c}: 477 cea_rows(t, mm1, mu, nf, (i + 1 : U32), cea_row0(mm1, (i * mu : U32), y, nf, loss, Array.get(F32, c, (i * mu : U32)))) 478 479 # ---- helper: hits per row ---- 480 481 type CR is Type: 482 CR{c: Array, n: Nat} 483 484 def gt_head(xs: List<&2, F32>, +v: F32) -> Bool: 485 match xs: 486 case Nil{}: 487 False{} 488 case Con{+h, t}: 489 F32.is_gt(h, v) 490 491 def argmax2(xs: List<&2, F32>, g: Bool, +i: Nat, +bi: Nat, +bv: F32) -> Nat: 492 match xs g: 493 case Nil{} _: 494 bi 495 case Con{+h, +t} True{}: 496 argmax2(t, gt_head(t, h), 1n+i, i, h) 497 case Con{h, +t} False{}: 498 argmax2(t, gt_head(t, bv), 1n+i, bi, bv) 499 500 def argmax(xs: List<&2, F32>) -> Nat: 501 match xs: 502 case Nil{}: 503 0n 504 case Con{+h, +t}: 505 argmax2(t, gt_head(t, h), 1n, 0n, h) 506 507 def one(b: Bool) -> Nat: 508 match b: 509 case True{}: 510 1n 511 case False{}: 512 0n 513 514 def count_rows3(+y: Nat, n: Nat, al: AL) -> CR: 515 match al: 516 case AL{c, l}: 517 CR{c, Nat.add(n, one(Nat.is_eq(argmax(l), y)))} 518 519 def count_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +i: U32, r: CR) -> CR: 520 match labels r: 521 case Nil{} CR{c, n}: 522 CR{c, n} 523 case Con{+y, t} CR{c, n}: 524 count_rows(t, mm1, mu, (i + 1 : U32), count_rows3(y, n, read_l(mm1, ((i * mu : U32) + U32.from_nat(mm1) : U32), Nil{}, Array.get(F32, c, ((i * mu : U32) + U32.from_nat(mm1) : U32))))) 525 526 527 # ===================================================================== 528 # Typed API: Mat (the dimensions enter the type) 529 # ===================================================================== 530 531 # ===================================================================== 532 # Array capacity: a PROVED depth for the number of slots 533 # ===================================================================== 534 # 535 # cap_depth(n) is the d such that an Array with 2^d slots holds n elements. It is computed by 536 # structural recursion on n (not with U32.log2) so that the law cap_ok below can prove 537 # `n <= 2^cap_depth(n)`. What stays trusted is only that Array.new(T, d, v) allocates 2^d slots. 538 539 # a type chosen by a Bool: rewriting through it refutes True == False 540 def BD(b: Bool, t: Type, f: Type) -> Type: 541 match b: 542 case True{}: 543 t 544 case False{}: 545 f 546 547 # a <= b, by direct recursion 548 def leb(a: Nat, b: Nat) -> Bool: 549 match a b: 550 case 0n _: 551 True{} 552 case 1n+p 0n: 553 False{} 554 case 1n+p 1n+q: 555 leb(p, q) 556 557 def pow2(d: Nat) -> Nat: 558 match d: 559 case 0n: 560 1n 561 case 1n++q: 562 # bend-lint: disable-next bend-lint-rules/performance/repeated_call -- pow2 appears only in proofs' types not-expected-because: the package never runs it (NOTES 2026-10-17) 563 Nat.add(pow2(q), pow2(q)) 564 565 # smallest power of two >= n, by structural recursion on n: returns (depth, 2^depth) 566 def cap_pick(ok: Bool, d: Nat, +p: Nat) -> Nat & Nat: 567 match ok: 568 case True{}: 569 (d, p) 570 case False{}: 571 (1n+d, Nat.add(p, p)) 572 573 def cap_step(r: Nat & Nat, +n: Nat) -> Nat & Nat: 574 match r: 575 case Tuple{+d, +p}: 576 cap_pick(leb(n, p), d, p) 577 578 def cap_pair(n: Nat) -> Nat & Nat: 579 match n: 580 case 0n: 581 (0n, 1n) 582 case 1n++m: 583 cap_step(cap_pair(m), 1n+m) 584 585 def fst_nat(r: Nat & Nat) -> Nat: 586 match r: 587 case Tuple{d, p}: 588 d 589 590 def snd_nat(r: Nat & Nat) -> Nat: 591 match r: 592 case Tuple{d, p}: 593 p 594 595 # leb is Base's Nat.is_le, case for case. leb walks its arguments one step at a time in the compiled program; Base's 596 # comparison is one native comparison there (0.1.13.0). 597 def le_leb(a: Nat, b: Nat) -> {Nat.is_le(a, b) == leb(a, b) : Bool}: 598 match a b: 599 case 0n 0n: 600 {==} 601 case 0n 1n+q: 602 {==} 603 case 1n+p 0n: 604 {==} 605 case 1n+p 1n+q: 606 le_leb(p, q) 607 608 def leb_of(+a: Nat, +b: Nat, e: {True{} == Nat.is_le(a, b) : Bool}) -> {True{} == leb(a, b) : Bool}: 609 %le_leb(a, b) : {True{} == _ : Bool} 610 e 611 612 # the same pair by doubling (0.1.10.0): p doubles until n <= p, log2(n) steps in a tail loop, where cap_pair takes n 613 # steps of a recursion that is not a tail call, each with a comparison of up to n steps. Its fuel is n, which always 614 # suffices (2^n >= n); if it ran out, the answer would be cap_pair's, so the law below needs no arithmetic about how 615 # many doublings there are. ok is Nat.is_le(n, p), computed by the caller (match only inspects a parameter); with leb 616 # it was n steps in all, 152 ms for n = 786,432 on one GPU lane (0.1.13.0). 617 def cap_dbl(f: Nat, ok: Bool, +n: Nat, +d: Nat, +p: Nat) -> Nat & Nat: 618 match f ok: 619 case _ True{}: 620 (d, p) 621 case 0n False{}: 622 cap_pair(n) 623 case 1n+g False{}: 624 cap_dbl(g, Nat.is_le(n, Nat.add(p, p)), n, 1n+d, Nat.add(p, p)) 625 626 # the depth d such that an Array with 2^d slots holds n elements (the smallest such d) 627 def cap_depth(+n: Nat) -> Nat: 628 fst_nat(cap_dbl(n, Nat.is_le(n, 1n), n, 0n, 1n)) 629 630 # ---- lemmas ---- 631 632 # a <= b -> a <= 1 + b 633 def leb_succ(a: Nat, b: Nat, e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, 1n+b) : Bool}: 634 match a b: 635 case 0n _: 636 {==} 637 case 1n+p 0n: 638 %e : BD(_, Unit, {True{} == leb(1n+p, 1n) : Bool}) 639 Unit{} 640 case 1n+p 1n+q: 641 leb_succ(p, q, e) 642 643 # a <= b -> a <= c + b 644 def leb_up_l(+a: Nat, +b: Nat, c: Nat, +e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, Nat.add(c, b)) : Bool}: 645 match c: 646 case 0n: 647 e 648 case 1n++c2: 649 leb_succ(a, Nat.add(c2, b), leb_up_l(a, b, c2, e)) 650 651 # a <= b -> a <= b + c 652 def leb_up(a: Nat, b: Nat, c: Nat, e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, Nat.add(b, c)) : Bool}: 653 match a b: 654 case 0n _: 655 {==} 656 case 1n+p 0n: 657 %e : BD(_, Unit, {True{} == leb(1n+p, Nat.add(0n, c)) : Bool}) 658 Unit{} 659 case 1n+p 1n+q: 660 leb_up(p, q, c, e) 661 662 # a <= b and c <= d -> a + c <= b + d 663 def leb_add(a: Nat, b: Nat, c: Nat, d: Nat, e1: {True{} == leb(a, b) : Bool}, e2: {True{} == leb(c, d) : Bool}) -> {True{} == leb(Nat.add(a, c), Nat.add(b, d)) : Bool}: 664 match a b: 665 case 0n _: 666 leb_up_l(c, d, b, e2) 667 case 1n+p 0n: 668 %e1 : BD(_, Unit, {True{} == leb(Nat.add(1n+p, c), Nat.add(0n, d)) : Bool}) 669 Unit{} 670 case 1n+p 1n+q: 671 leb_add(p, q, c, d, e1, e2) 672 673 # 1 <= 2^d 674 def pow2_pos(d: Nat) -> {True{} == leb(1n, pow2(d)) : Bool}: 675 match d: 676 case 0n: 677 {==} 678 case 1n++q: 679 leb_up(1n, pow2(q), pow2(q), pow2_pos(q)) 680 681 # 1 <= p when p = 2^d 682 def pos_p(d: Nat, p: Nat, i1: {p == pow2(d) : Nat}) -> {True{} == leb(1n, p) : Bool}: 683 %Equal.sym(Nat, p, pow2(d), i1) : {True{} == leb(1n, _) : Bool} 684 pow2_pos(d) 685 686 # p + p == 2^(1+d) when p == 2^d 687 def step_eq(d: Nat, p: Nat, i1: {p == pow2(d) : Nat}) -> {Nat.add(p, p) == pow2(1n+d) : Nat}: 688 %i1 : {Nat.add(p, p) == Nat.add(_, _) : Nat} 689 {==} 690 691 # one step: given the invariant for m, the invariant for 1+m holds whichever way the comparison goes 692 def step_ok(ok: Bool, +d: Nat, +p: Nat, +m: Nat, +e: {ok == leb(1n+m, p) : Bool}, +i1: {p == pow2(d) : Nat}, +i2: {True{} == leb(m, p) : Bool}) -> {snd_nat(cap_pick(ok, d, p)) == pow2(fst_nat(cap_pick(ok, d, p))) : Nat} & {True{} == leb(1n+m, snd_nat(cap_pick(ok, d, p))) : Bool}: 693 match ok: 694 case True{}: 695 (i1, e) 696 case False{}: 697 (step_eq(d, p, i1), leb_add(1n, p, m, p, pos_p(d, p, i1), i2)) 698 699 700 # the invariant for cap_step on a pair, opened 701 def step_pair_ok(+m: Nat, r: Nat & Nat, i1: {snd_nat(r) == pow2(fst_nat(r)) : Nat}, i2: {True{} == leb(m, snd_nat(r)) : Bool}) -> {snd_nat(cap_step(r, 1n+m)) == pow2(fst_nat(cap_step(r, 1n+m))) : Nat} & {True{} == leb(1n+m, snd_nat(cap_step(r, 1n+m))) : Bool}: 702 match r: 703 case Tuple{+d, +p}: 704 step_ok(leb(1n+m, p), d, p, m, {==}, i1, i2) 705 706 def cap_succ(+m: Nat, ih: {snd_nat(cap_pair(m)) == pow2(fst_nat(cap_pair(m))) : Nat} & {True{} == leb(m, snd_nat(cap_pair(m))) : Bool}) -> {snd_nat(cap_pair(1n+m)) == pow2(fst_nat(cap_pair(1n+m))) : Nat} & {True{} == leb(1n+m, snd_nat(cap_pair(1n+m))) : Bool}: 707 match ih: 708 case Tuple{i1, i2}: 709 step_pair_ok(m, cap_pair(m), i1, i2) 710 711 # the invariant of cap_pair: the second component is 2^(first) and is >= n 712 def cap_inv(n: Nat) -> {snd_nat(cap_pair(n)) == pow2(fst_nat(cap_pair(n))) : Nat} & {True{} == leb(n, snd_nat(cap_pair(n))) : Bool}: 713 match n: 714 case 0n: 715 ({==}, {==}) 716 case 1n++m: 717 cap_succ(m, cap_inv(m)) 718 719 # the invariant of cap_dbl: whichever way it ends (n <= p found, or the fuel out and cap_pair's answer), the second 720 # component is 2^(first) and is >= n 721 def dbl_inv(f: Nat, ok: Bool, +n: Nat, +d: Nat, +p: Nat, e: {ok == Nat.is_le(n, p) : Bool}, i1: {p == pow2(d) : Nat}) -> {snd_nat(cap_dbl(f, ok, n, d, p)) == pow2(fst_nat(cap_dbl(f, ok, n, d, p))) : Nat} & {True{} == leb(n, snd_nat(cap_dbl(f, ok, n, d, p))) : Bool}: 722 match f ok: 723 case _ True{}: 724 (i1, leb_of(n, p, e)) 725 case 0n False{}: 726 cap_inv(n) 727 case 1n+g False{}: 728 dbl_inv(g, Nat.is_le(n, Nat.add(p, p)), n, 1n+d, Nat.add(p, p), {==}, step_eq(d, p, i1)) 729 730 def dbl_final(n: Nat, pr: {snd_nat(cap_dbl(n, Nat.is_le(n, 1n), n, 0n, 1n)) == pow2(fst_nat(cap_dbl(n, Nat.is_le(n, 1n), n, 0n, 1n))) : Nat} & {True{} == leb(n, snd_nat(cap_dbl(n, Nat.is_le(n, 1n), n, 0n, 1n))) : Bool}) -> {True{} == leb(n, pow2(cap_depth(n))) : Bool}: 731 match pr: 732 case Tuple{i1, i2}: 733 %i1 : {True{} == leb(n, _) : Bool} 734 i2 735 736 # LAW (cap_ok): an Array with 2^cap_depth(n) slots always has room for n elements. 737 law cap_ok: 738 for +n: Nat 739 {True{} == leb(n, pow2(cap_depth(n))) : Bool} 740 741 def cap_ok(n): 742 dbl_final(n, dbl_inv(n, Nat.is_le(n, 1n), n, 0n, 1n, {==}, {==})) 743 744 # r x c of zeros, with the exact capacity (smallest power of 2) 745 746 def Mat.zeros(+r: Nat, +c: Nat) -> Mat: 747 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0)} 748 749 def Mat.fill(+r: Nat, +c: Nat, +v: F32) -> Mat: 750 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), v)} 751 752 def fill_list(xs: List<&2, F32>, +i: U32, a: Array) -> Array: 753 match xs: 754 case Nil{}: 755 a 756 case Con{h, t}: 757 fill_list(t, (i + 1 : U32), Array.set(F32, a, i, h)) 758 759 def len_is(xs: List<&2, F32>, n: Nat) -> Bool: 760 match xs n: 761 case Nil{} 0n: 762 True{} 763 case Con{h, t} 1n+p: 764 len_is(t, p) 765 case _ _: 766 False{} 767 768 def Mat.of_list2(-r: Nat, -c: Nat, +rr: Nat, +cc: Nat, ok: Bool, xs: List<&2, F32>) -> Maybe<&1, Mat>: 769 match ok: 770 case True{}: 771 Some{Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(rr, cc)), 0.0))}} 772 case False{}: 773 None{} 774 775 # Mat from r*c numbers in row order; None if the size does not match 776 777 def Mat.of_list(+r: Nat, +c: Nat, +xs: List<&2, F32>) -> Maybe<&1, Mat>: 778 Mat.of_list2(r, c, r, c, len_is(xs, Nat.mul(r, c)), xs) 779 780 # reads the first n numbers, returning the matrix 781 782 # Mat without checking the list size (a short list leaves the rest at zero, a long one is cut 783 # by the capacity): only for when the size is known by construction. Prefer Mat.of_list. 784 785 def Mat.from_list(+r: Nat, +c: Nat, xs: List<&2, F32>) -> Mat: 786 Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0))} 787 788 # writes the numbers of xs into m starting at flat index i (row-major); numbers past the 789 # capacity are dropped. For loading a large matrix block by block without building one big list. 790 def Mat.fill_at(-r: Nat, -c: Nat, +i: U32, xs: List<&2, F32>, m: Mat) -> Mat: 791 match m: 792 case Mat{d}: 793 Mat{fill_list(xs, i, d)} 794 795 type ML<-r: Nat, -c: Nat> is Type: 796 ML{m: Mat, l: List<&2, F32>} 797 798 def Mat.to_list2(-r: Nat, -c: Nat, al: AL) -> ML: 799 match al: 800 case AL{d, l}: 801 ML{Mat{d}, l} 802 803 def ML.len(-r: Nat, -c: Nat, ml: ML) -> Nat: 804 match ml: 805 case ML{m, l}: 806 List.length(&2, F32, l) 807 808 # n numbers from index 0, in row order; n = 0 gives the empty list (a Mat with no numbers) 809 def Mat.to_list_count(-r: Nat, -c: Nat, d: Array, +n: Nat) -> ML: 810 match n: 811 case 0n: 812 ML{Mat{d}, Nil{}} 813 case 1n+p: 814 Mat.to_list2(r, c, read_l(p, (U32.from_nat(n) - 1 : U32), Nil{}, Array.get(F32, d, (U32.from_nat(n) - 1 : U32)))) 815 816 # two independent copies 817 818 def Mat.to_list(+r: Nat, +c: Nat, m: Mat) -> ML: 819 match m: 820 case Mat{d}: 821 Mat.to_list_count(r, c, d, Nat.mul(r, c)) 822 823 type MM2<-r: Nat, -c: Nat> is Type: 824 MM2{a: Mat, b: Mat} 825 826 def Mat.clone2(-r: Nat, -c: Nat, p: Array & Array) -> MM2: 827 match p: 828 case Tuple{x, y}: 829 MM2{Mat{x}, Mat{y}} 830 831 # ---- products: C (n x m) = A · B, with 2^par blocks of rows in parallel ---- 832 833 # C = A(n x k) · B(k x m) 834 835 def Mat.clone(-r: Nat, -c: Nat, m: Mat) -> MM2: 836 match m: 837 case Mat{d}: 838 Mat.clone2(r, c, Array.clone(F32, d)) 839 840 type MMul<-n: Nat, -k: Nat, -m: Nat> is Type: 841 MMul{a: Mat, b: Mat, c: Mat} 842 843 # C = A(n x k) · Bᵀ, with B stored m x k 844 845 type MMulNT<-n: Nat, -k: Nat, -m: Nat> is Type: 846 MMulNT{a: Mat, b: Mat, c: Mat} 847 848 # C = Aᵀ · B, with A stored k x n 849 850 type MMulTN<-n: Nat, -k: Nat, -m: Nat> is Type: 851 MMulTN{a: Mat, b: Mat, c: Mat} 852 853 # --------------------------------------------------------------------- 854 # gemm in tiles of 8 (0.1.8.0): for products with 2 or more rows, A's rows and B's columns are packed once into cells of 855 # 8 F32 (a row of k numbers is ceil(k/8) cells, zeros after the last), so the inner loop reads 8 numbers per array 856 # access instead of one. The order of the additions is that of `dotr`: product p of a dot goes to partial sum p mod 4, 857 # in order, and the four sums are combined as dotr combines them after k steps (its rotation depends on k mod 4). The 858 # padding adds +0 to a sum that starts at +0, which changes nothing, so the results are the same bits as the scalar 859 # loop's. Products with one row (matrix · vector) keep the scalar loop: packing B would cost as much as the product. 860 # --------------------------------------------------------------------- 861 862 type T8 is Data: 863 T8{v0: F32, v1: F32, v2: F32, v3: F32, v4: F32, v5: F32, v6: F32, v7: F32} 864 865 def t8_zero() -> T8: 866 T8{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0} 867 868 def pick0(b: Bool, +v: F32) -> F32: 869 match b: 870 case True{}: 871 v 872 case False{}: 873 0.0 874 875 type PK is Type: 876 PK{a: Array, t: Array} 877 878 # one cell: the 8 numbers of a at idx, idx + st, ..., those at positions q0 + q >= lim (past the end of the line) as 0 879 def pk_g7(+v0: F32, +v1: F32, +v2: F32, +v3: F32, +v4: F32, +v5: F32, +v6: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array, r: Array & F32) -> PK: 880 match r: 881 case Tuple{a, +v}: 882 PK{a, Array.set(T8, t, o, T8{v0, v1, v2, v3, v4, v5, v6, pick0(U32.is_lt((q0 + 7 : U32), lim), v)})} 883 884 def pk_g6(+v0: F32, +v1: F32, +v2: F32, +v3: F32, +v4: F32, +v5: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array, r: Array & F32) -> PK: 885 match r: 886 case Tuple{a, +v}: 887 pk_g7(v0, v1, v2, v3, v4, v5, pick0(U32.is_lt((q0 + 6 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 7 : U32) : U32))) 888 889 def pk_g5(+v0: F32, +v1: F32, +v2: F32, +v3: F32, +v4: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array, r: Array & F32) -> PK: 890 match r: 891 case Tuple{a, +v}: 892 pk_g6(v0, v1, v2, v3, v4, pick0(U32.is_lt((q0 + 5 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 6 : U32) : U32))) 893 894 def pk_g4(+v0: F32, +v1: F32, +v2: F32, +v3: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array, r: Array & F32) -> PK: 895 match r: 896 case Tuple{a, +v}: 897 pk_g5(v0, v1, v2, v3, pick0(U32.is_lt((q0 + 4 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 5 : U32) : U32))) 898 899 def pk_g3(+v0: F32, +v1: F32, +v2: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array, r: Array & F32) -> PK: 900 match r: 901 case Tuple{a, +v}: 902 pk_g4(v0, v1, v2, pick0(U32.is_lt((q0 + 3 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 4 : U32) : U32))) 903 904 def pk_g2(+v0: F32, +v1: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array, r: Array & F32) -> PK: 905 match r: 906 case Tuple{a, +v}: 907 pk_g3(v0, v1, pick0(U32.is_lt((q0 + 2 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 3 : U32) : U32))) 908 909 def pk_g1(+v0: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array, r: Array & F32) -> PK: 910 match r: 911 case Tuple{a, +v}: 912 pk_g2(v0, pick0(U32.is_lt((q0 + 1 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 2 : U32) : U32))) 913 914 def pk_g0(+q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array, r: Array & F32) -> PK: 915 match r: 916 case Tuple{a, +v}: 917 pk_g1(pick0(U32.is_lt((q0 + 0 : U32), lim), v), q0, lim, idx, st, o, t, Array.get(F32, a, (idx + (st * 1 : U32) : U32))) 918 919 # the cells of one line, cell o.. of t; q0 counts the line's numbers before the cell 920 def pk_line(tl: Nat, +q0: U32, +lim: U32, +base: U32, +st: U32, +o: U32, r: PK) -> PK: 921 match tl r: 922 case 0n r: 923 r 924 case 1n+q PK{a, t}: 925 pk_line(q, (q0 + 8 : U32), lim, (base + (st * 8 : U32) : U32), st, (o + 1 : U32), pk_g0(q0, lim, base, st, o, t, Array.get(F32, a, base))) 926 927 # lines count lines; source line i starts at i * lstep (in a), and goes to cell row o (cells o * tku .. in t) 928 def pk_lines(cnt: Nat, +tk: Nat, +tku: U32, +ku: U32, +i: U32, +o: U32, +lstep: U32, +st: U32, r: PK) -> PK: 929 match cnt: 930 case 0n: 931 r 932 case 1n+p: 933 pk_lines(p, tk, tku, ku, (i + 1 : U32), (o + 1 : U32), lstep, st, pk_line(tk, 0, ku, (i * lstep : U32), st, (o * tku : U32), r)) 934 935 def pk_pack(+cnt: Nat, +tk: Nat, +k: Nat, +lstep: U32, +st: U32, a: Array) -> PK: 936 pk_lines(cnt, tk, U32.from_nat(tk), U32.from_nat(k), 0, 0, lstep, st, PK{a, Array.new(T8, cap_depth(Nat.mul(cnt, tk)), t8_zero())}) 937 938 # the dot of two lines of cells: tk cells from ia (in a) and ib (in b); ra, rb hold the arrays and the cells at ia, ib, 939 # already read. Product q of a cell goes to sum q mod 4, the first four before the last four, as dotr adds them. 940 def t8_dot(tk: Nat, +ia: U32, +ib: U32, s0: F32, s1: F32, s2: F32, s3: F32, ra: Array & T8, rb: Array & T8) -> TD: 941 match tk ra rb: 942 case 0n Tuple{a, x} Tuple{b, y}: 943 TD{a, b, s0, s1, s2, s3} 944 case 1n+p Tuple{a, T8{x0, x1, x2, x3, x4, x5, x6, x7}} Tuple{b, T8{y0, y1, y2, y3, y4, y5, y6, y7}}: 945 t8_dot(p, (ia + 1 : U32), (ib + 1 : U32), ((s0 + (x0 * y0 : F32) : F32) + (x4 * y4 : F32) : F32), ((s1 + (x1 * y1 : F32) : F32) + (x5 * y5 : F32) : F32), ((s2 + (x2 * y2 : F32) : F32) + (x6 * y6 : F32) : F32), ((s3 + (x3 * y3 : F32) : F32) + (x7 * y7 : F32) : F32), Array.get(T8, a, (ia + 1 : U32)), Array.get(T8, b, (ib + 1 : U32))) 946 947 type TD is Type: 948 TD{a: Array, b: Array, s0: F32, s1: F32, s2: F32, s3: F32} 949 950 # dotr's final sum after k products: the sums sit rotated by k mod 4 when its loop ends 951 def t8_comb(+rot: Nat, +s0: F32, +s1: F32, +s2: F32, +s3: F32) -> F32: 952 match rot: 953 case 0n: 954 ((s0 + s1 : F32) + (s2 + s3 : F32) : F32) 955 case 1n: 956 ((s1 + s2 : F32) + (s3 + s0 : F32) : F32) 957 case 2n: 958 ((s2 + s3 : F32) + (s0 + s1 : F32) : F32) 959 case _: 960 ((s3 + s0 : F32) + (s1 + s2 : F32) : F32) 961 962 type TR is Type: 963 TR{a: Array, b: Array, x: F32} 964 965 def t8_fin(+rot: Nat, d: TD) -> TR: 966 match d: 967 case TD{a, b, +s0, +s1, +s2, +s3}: 968 TR{a, b, t8_comb(rot, s0, s1, s2, s3)} 969 970 def t8_cell(+tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, a: Array, b: Array) -> TR: 971 t8_fin(rot, t8_dot(tk, (i * tku : U32), (j * tku : U32), 0.0, 0.0, 0.0, 0.0, Array.get(T8, a, (i * tku : U32)), Array.get(T8, b, (j * tku : U32)))) 972 973 type TG is Type: 974 TG{a: Array, b: Array, c: Array} 975 976 def t8_set(+ac: Bool, +idx: U32, c: Array, r: TR) -> TG: 977 match r: 978 case TR{a, b, +x}: 979 TG{a, b, cwr(ac, idx, x, c)} 980 981 def t8_cols(+ac: Bool, mleft: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, +cidx: U32, g: TG) -> TG: 982 match mleft g: 983 case 0n g: 984 g 985 case 1n+p TG{a, b, c}: 986 t8_cols(ac, p, tk, tku, rot, i, (j + 1 : U32), (cidx + 1 : U32), t8_set(ac, cidx, c, t8_cell(tk, tku, rot, i, j, a, b))) 987 988 def t8_rows(+ac: Bool, nleft: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +mu: U32, g: TG) -> TG: 989 match nleft: 990 case 0n: 991 g 992 case 1n+p: 993 t8_rows(ac, p, m, tk, tku, rot, (i + 1 : U32), mu, t8_cols(ac, m, tk, tku, rot, i, 0, (i * mu : U32), g)) 994 995 # rows i0 .. i0+cnt-1 of C as a list, in order (the parallel version's leaves) 996 type TL is Type: 997 TL{a: Array, b: Array, out: List<&2, F32>} 998 999 # cells j.. of row i, m - j of them; r is cell j, already computed 1000 def t8_lcols(mleft: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, out: List<&2, F32>, r: TR) -> TL: 1001 match mleft r: 1002 case 0n TR{a, b, +x}: 1003 TL{a, b, x <> out} 1004 case 1n+p TR{a, b, +x}: 1005 t8_lcols(p, tk, tku, rot, i, (j + 1 : U32), x <> out, t8_cell(tk, tku, rot, i, (j + 1 : U32), a, b)) 1006 1007 def t8_lrows(nleft: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, l: TL) -> TL: 1008 match nleft l: 1009 case 0n l: 1010 l 1011 case 1n+p TL{a, b, out}: 1012 t8_lrows(p, m, tk, tku, rot, (i + 1 : U32), t8_lcols(Nat.sub(m, 1n), tk, tku, rot, i, 0, out, t8_cell(tk, tku, rot, i, 0, a, b))) 1013 1014 def t8_lfin(l: TL) -> List<&2, F32>: 1015 match l: 1016 case TL{a, b, out}: 1017 List.reverse(&2, F32, out) 1018 1019 type PKK is Type: 1020 PKK{a: Array, b: Array, ta: Array, tb: Array} 1021 1022 def t8_pk2(pa: PK, pb: PK) -> PKK: 1023 match pa pb: 1024 case PK{a, ta} PK{b, tb}: 1025 PKK{a, b, ta, tb} 1026 1027 def t8_g(a: Array, b: Array, g: TG) -> G: 1028 match g: 1029 case TG{ta, tb, c}: 1030 G{a, b, c} 1031 1032 def t8_out_seq(+ac: Bool, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array, p: PKK) -> G: 1033 match p: 1034 case PKK{a, b, ta, tb}: 1035 t8_g(a, b, t8_rows(ac, n, m, tk, tku, rot, 0, U32.from_nat(m), TG{ta, tb, c})) 1036 1037 # ---- the parallel path (0.1.9.0): the larger operand is packed into one piece per block, so a block owns its piece; 1038 # the other operand is copied once per node of the fork tree (2^d - 1 copies), instead of both at every node ---- 1039 1040 type PT is Type: 1041 PLeaf{t: Array} 1042 PNode{l: PT, r: PT} 1043 1044 type PKT is Type: 1045 PKT{a: Array, t: PT} 1046 1047 def pkt_leaf(r: PK) -> PKT: 1048 match r: 1049 case PK{a, t}: 1050 PKT{a, PLeaf{t}} 1051 1052 def pkt_then(r: PKT, k: Array -> PT -> PKT) -> PKT: 1053 match r: 1054 case PKT{a, t}: 1055 k(a, t) 1056 1057 # lines i .. i+cnt-1 of a, packed into 2^d pieces (halves of the line count, as the fork tree splits them) 1058 def pk_tree(+d: Nat, +cnt: Nat, +i: U32, +tk: Nat, +tku: U32, +ku: U32, +lstep: U32, +st: U32, a: Array) -> PKT: 1059 match d: 1060 case 0n: 1061 pkt_leaf(pk_lines(cnt, tk, tku, ku, i, 0, lstep, st, PK{a, Array.new(T8, cap_depth(Nat.mul(cnt, tk)), t8_zero())})) 1062 case 1n+q: 1063 pkt_then(pk_tree(q, Nat.div(cnt, 2n), i, tk, tku, ku, lstep, st, a), a2 => l => pkt_then(pk_tree(q, Nat.sub(cnt, Nat.div(cnt, 2n)), (i + U32.from_nat(Nat.div(cnt, 2n)) : U32), tk, tku, ku, lstep, st, a2), a3 => r => PKT{a3, PNode{l, r}})) 1064 1065 def t8_with2(-R: Type, p: Array & Array, k: Array -> Array -> R) -> R: 1066 match p: 1067 case Tuple{x, y}: 1068 k(x, y) 1069 1070 # split by rows of C: the pieces are A's rows, B is shared; the result is C's rows in order 1071 def t8_rpar(t: PT, +cnt: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, b: Array) -> List<&2, F32>: 1072 match t: 1073 case PLeaf{a}: 1074 t8_lfin(t8_lrows(cnt, m, tk, tku, rot, 0, TL{a, b, Nil{}})) 1075 case PNode{l, r}: 1076 t8_with2(List<&2, F32>, Array.clone(T8, b), b1 => b2 => ( 1077 x y = t8_rpar(l, Nat.div(cnt, 2n), m, tk, tku, rot, b1) t8_rpar(r, Nat.sub(cnt, Nat.div(cnt, 2n)), m, tk, tku, rot, b2) 1078 List.append(&2, F32, x, y))) 1079 1080 # a block of C's columns j0 .. j0+w-1 for every row, row by row 1081 type CB is Data: 1082 CB{j0: U32, w: Nat, xs: List<&2, F32>} 1083 1084 # split by columns of C: the pieces are B's columns, A is shared 1085 def t8_cpar(t: PT, +cnt: Nat, +j0: U32, +n: Nat, +tk: Nat, +tku: U32, +rot: Nat, a: Array) -> List<&2, CB>: 1086 match t: 1087 case PLeaf{b}: 1088 [CB{j0, cnt, t8_lfin(t8_lrows(n, cnt, tk, tku, rot, 0, TL{a, b, Nil{}}))}] 1089 case PNode{l, r}: 1090 t8_with2(List<&2, CB>, Array.clone(T8, a), a1 => a2 => ( 1091 x y = t8_cpar(l, Nat.div(cnt, 2n), j0, n, tk, tku, rot, a1) t8_cpar(r, Nat.sub(cnt, Nat.div(cnt, 2n)), (j0 + U32.from_nat(Nat.div(cnt, 2n)) : U32), n, tk, tku, rot, a2) 1092 List.append(&2, CB, x, y))) 1093 1094 type RW is Type: 1095 RW{c: Array, xs: List<&2, F32>} 1096 1097 # w numbers of xs into c from idx; what is left of xs comes back 1098 def cb_row(+ac: Bool, w: Nat, +idx: U32, c: Array, xs: List<&2, F32>) -> RW: 1099 match w xs: 1100 case 1n+p Con{+h, t}: 1101 cb_row(ac, p, (idx + 1 : U32), cwr(ac, idx, h, c), t) 1102 case _ _: 1103 RW{c, xs} 1104 1105 def cb_rows(+ac: Bool, n: Nat, +i: U32, +j0: U32, +w: Nat, +mu: U32, r: RW) -> Array: 1106 match n r: 1107 case 0n RW{c, xs}: 1108 c 1109 case 1n+p RW{c, xs}: 1110 cb_rows(ac, p, (i + 1 : U32), j0, w, mu, cb_row(ac, w, ((i * mu : U32) + j0 : U32), c, xs)) 1111 1112 def cb_write(+ac: Bool, bs: List<&2, CB>, +n: Nat, +mu: U32, c: Array) -> Array: 1113 match bs: 1114 case Nil{}: 1115 c 1116 case Con{CB{+j0, +w, xs}, t}: 1117 cb_write(ac, t, n, mu, cb_rows(ac, n, 0, j0, w, mu, RW{c, xs})) 1118 1119 def t8_pr2(+ac: Bool, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array, ra: PKT, rb: PK) -> G: 1120 match ra rb: 1121 case PKT{a, t} PK{b, tb}: 1122 G{a, b, write_la(ac, t8_rpar(t, n, m, tk, tku, rot, tb), 0, c)} 1123 1124 def t8_pc2(+ac: Bool, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array, ra: PK, rb: PKT) -> G: 1125 match ra rb: 1126 case PK{a, ta} PKT{b, t}: 1127 G{a, b, cb_write(ac, t8_cpar(t, m, 0, n, tk, tku, rot, ta), n, U32.from_nat(m), c)} 1128 1129 # 2^d blocks along the larger of n (rows of C, pieces of A) and m (columns of C, pieces of B) 1130 def t8_par2(+ac: Bool, byrows: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +tk: Nat, +tku: U32, +rot: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 1131 match byrows: 1132 case True{}: 1133 t8_pr2(ac, n, m, tk, tku, rot, c, pk_tree(d, n, 0, tk, tku, U32.from_nat(k), arow, sa, a), pk_pack(m, tk, k, bcol, sb, b)) 1134 case False{}: 1135 t8_pc2(ac, n, m, tk, tku, rot, c, pk_pack(n, tk, k, arow, sa, a), pk_tree(d, m, 0, tk, tku, U32.from_nat(k), bcol, sb, b)) 1136 1137 def gemm_t8b(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +tk: Nat, +tku: U32, +rot: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 1138 match d: 1139 case 0n: 1140 t8_out_seq(ac, n, m, tk, tku, rot, c, t8_pk2(pk_pack(n, tk, k, arow, sa, a), pk_pack(m, tk, k, bcol, sb, b))) 1141 case 1n+q: 1142 t8_par2(ac, Nat.is_le(m, n), 1n+q, n, m, k, tk, tku, rot, arow, sa, sb, bcol, a, b, c) 1143 1144 # C (n x m) = A · B (or C + A · B when ac) in cells of 8, sequential (d = 0) or with 2^d blocks of rows in parallel; A 1145 # and B come back intact 1146 def gemm_t8_ac(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 1147 gemm_t8b(ac, d, n, m, k, Nat.div(Nat.add(k, 7n), 8n), U32.from_nat(Nat.div(Nat.add(k, 7n), 8n)), Nat.mod(k, 4n), arow, sa, sb, bcol, a, b, c) 1148 1149 def gemm_t8(d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 1150 gemm_t8_ac(False{}, d, n, m, k, arow, sa, sb, bcol, a, b, c) 1151 1152 # one row (matrix · vector): the scalar loop at every depth; two or more rows: cells of 8. Up to 0.1.11.0 one row in 1153 # parallel ran gemm_par_cols, which clones both operands at every node of the fork tree: a 50,257 x 768 weight matrix at 1154 # depth 4 ran out of the default heap (5.1 GB) and 8192 x 2048 took 124 ms against 10 ms sequentially. Packing the 1155 # matrix into cells of 8 instead costs more than the product, which reads each weight once (64 ms against 24 at 1156 # 50,257 x 768). Matrix · vector in parallel is Bands.matvec. Measured 2026-10-16 (NOTES.md). 1157 def gemm_any3(+ac: Bool, one: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 1158 match one: 1159 case True{}: 1160 gemm_ac(ac, n, m, k, arow, sa, sb, bcol, a, b, c) 1161 case False{}: 1162 gemm_t8_ac(ac, d, n, m, k, arow, sa, sb, bcol, a, b, c) 1163 1164 def gemm_any2(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 1165 gemm_any3(ac, Nat.is_le(n, 1n), d, n, m, k, arow, sa, sb, bcol, a, b, c) 1166 1167 def gemm_any(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 1168 gemm_any2(ac, d, n, m, k, arow, sa, sb, bcol, a, b, c) 1169 1170 def mmul_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMul: 1171 match g: 1172 case G{a, b, c}: 1173 MMul{Mat{a}, Mat{b}, Mat{c}} 1174 1175 def mmul_nt_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulNT: 1176 match g: 1177 case G{a, b, c}: 1178 MMulNT{Mat{a}, Mat{b}, Mat{c}} 1179 1180 def mmul_tn_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulTN: 1181 match g: 1182 case G{a, b, c}: 1183 MMulTN{Mat{a}, Mat{b}, Mat{c}} 1184 1185 # ---- the split depth of a product (0.1.11.0): par is an upper bound, and a product too small to gain splits 0 times ---- 1186 # 1187 # A product of n x k by k x m splits into 2^par blocks only when it does at least 2^23 multiply-adds and each cell of C 1188 # sums at least 64 of them; otherwise it runs sequentially. Measured on 16 threads (2026-10-16, 18 shapes): below 2^23 the 1189 # split costs more than it saves, and with a short inner dimension (a weight gradient over a few tokens) every split 1190 # was 2-3.6x slower, because each block hands its rows of C back as a list that is then written into C, a cost per cell 1191 # of C that k multiply-adds per cell do not pay for. The result does not depend on the depth: the same bits at every 1192 # depth (tests/t8.bend, tests/acc.bend, tests/split.bend). 1193 def split_pick(big: Bool, d: Nat) -> Nat: 1194 match big: 1195 case True{}: 1196 d 1197 case False{}: 1198 0n 1199 1200 def split_depth(d: Nat, +n: Nat, +m: Nat, +k: Nat) -> Nat: 1201 split_pick(Bool.and(Nat.is_le(64n, k), Nat.is_le(8388608n, Nat.mul(Nat.mul(n, m), k))), d) 1202 1203 def mm_a(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array, b: Array) -> MMul: 1204 mmul_fin(n, k, m, gemm_any(False{}, d, n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0))) 1205 1206 # C = A · B (n x k) · (k x m); par = log2 of the most parallel blocks (0 = sequential); split_depth picks 0 or par 1207 1208 def Mat.matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat, b: Mat) -> MMul: 1209 match a b: 1210 case Mat{x} Mat{y}: 1211 mm_a(n, k, m, split_depth(par, n, m, k), x, y) 1212 1213 # C + A · B, into C (n x m); A and B come back with the result, as Mat.matmul returns them. Each cell of A · B is 1214 # computed as Mat.matmul computes it and added once to C's cell: the same bits as Mat.add(c, Mat.matmul(a, b)'s c), 1215 # without allocating the product or walking it a second time. 1216 1217 def Mat.matmul_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat, b: Mat, c: Mat) -> MMul: 1218 match a b c: 1219 case Mat{x} Mat{y} Mat{z}: 1220 mmul_fin(n, k, m, gemm_any(True{}, split_depth(par, n, m, k), n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, x, y, z)) 1221 1222 def mm_nt(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array, b: Array) -> MMulNT: 1223 mmul_nt_fin(n, k, m, gemm_any(False{}, d, n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0))) 1224 1225 # C = A · Bᵀ (n x k) · (m x k)ᵀ: for weights stored with one row per output 1226 1227 def Mat.matmul_nt(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat, b: Mat) -> MMulNT: 1228 match a b: 1229 case Mat{x} Mat{y}: 1230 mm_nt(n, k, m, split_depth(par, n, m, k), x, y) 1231 1232 # C + A · Bᵀ, into C (n x m) 1233 1234 def Mat.matmul_nt_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat, b: Mat, c: Mat) -> MMulNT: 1235 match a b c: 1236 case Mat{x} Mat{y} Mat{z}: 1237 mmul_nt_fin(n, k, m, gemm_any(True{}, split_depth(par, n, m, k), n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), x, y, z)) 1238 1239 def mm_tn(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array, b: Array) -> MMulTN: 1240 mmul_tn_fin(n, k, m, gemm_any(False{}, d, n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0))) 1241 1242 # C = Aᵀ · B (k x n)ᵀ · (k x m): weight gradient, X^T · dY 1243 1244 def Mat.matmul_tn(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat, b: Mat) -> MMulTN: 1245 match a b: 1246 case Mat{x} Mat{y}: 1247 mm_tn(n, k, m, split_depth(par, n, m, k), x, y) 1248 1249 # C + Aᵀ · B, into C (n x m): a weight gradient added into the gradient it accumulates, dW += Xᵀ · dY 1250 1251 def Mat.matmul_tn_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat, b: Mat, c: Mat) -> MMulTN: 1252 match a b c: 1253 case Mat{x} Mat{y} Mat{z}: 1254 mmul_tn_fin(n, k, m, gemm_any(True{}, split_depth(par, n, m, k), n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, x, y, z)) 1255 1256 # ---- a precise product for long dots (0.1.14.0) ---- 1257 # 1258 # Mat.matmul_precise is Mat.matmul with k summed in blocks of 1024 (the last one 1..1024): block b's columns of A and 1259 # rows of B are copied out, the first block's product written into C, each next one's added to it as matmul_acc adds, 1260 # so every cell is ((x0 + x1) + x2) .., x_b the block's dot with the usual four partial sums. A dot of 49,152 products 1261 # as four running sums of 12,288 lost the small terms (SmolLM2's head: its gradient 1.8x further from float64 than 1262 # PyTorch's float32, 0.03x in blocks; BendTorch NOTES, 2026-10-18). Opt-in: it copies the slices and runs one product 1263 # a block, and the default products keep their speed and bits. Up to 1024 steps it is Mat.matmul itself. 1264 1265 type SL is Type: 1266 SL{a: Array, s: Array} 1267 1268 # d[di..] = src[si..], k numbers; r holds src[si], already read 1269 def sl_row(k: Nat, +si: U32, +di: U32, d: Array, r: Array & F32) -> SL: 1270 match k r: 1271 case 0n Tuple{a, x}: SL{a, d} 1272 case 1n+p Tuple{a, +x}: sl_row(p, (si + 1 : U32), (di + 1 : U32), Array.set(F32, d, di, x), Array.get(F32, a, (si + 1 : U32))) 1273 1274 # rows i.. of a slice: w numbers a row from src[i*st + o], into d[i*w] 1275 def sl_rows(r: Nat, +i: U32, +w: Nat, +w32: U32, +st: U32, +o: U32, sl: SL) -> SL: 1276 match r sl: 1277 case 0n sl: sl 1278 case 1n+p SL{a, d}: sl_rows(p, (i + 1 : U32), w, w32, st, o, sl_row(w, ((i * st : U32) + o : U32), (i * w32 : U32), d, Array.get(F32, a, ((i * st : U32) + o : U32)))) 1279 1280 def blk_len(+bi: U32, +k: U32) -> Nat: 1281 U32.to_nat(Bool.pick(U32, U32.is_lt((k - (bi * 1024 : U32) : U32), 1024), (k - (bi * 1024 : U32) : U32), 1024)) 1282 1283 def mp_s4(a: Array, b: Array, g: G) -> G: 1284 match g: 1285 case G{x, y, c}: G{a, b, c} 1286 1287 def mp_s3(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +w: Nat, +w32: U32, +m32: U32, a: Array, as: Array, c: Array, sb: SL) -> G: 1288 match sb: 1289 case SL{b, bs}: mp_s4(a, b, gemm_any(ac, split_depth(par, n, m, w), n, m, w, w32, 1, m32, 1, as, bs, c)) 1290 1291 def mp_s2(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +w: Nat, +w32: U32, +m32: U32, +o: U32, c: Array, b: Array, sa: SL) -> G: 1292 match sa: 1293 case SL{a, as}: mp_s3(ac, par, n, m, w, w32, m32, a, as, c, sl_row(Nat.mul(w, m), (o * m32 : U32), 0, Array.new(F32, cap_depth(Nat.mul(w, m)), 0.0), Array.get(F32, b, (o * m32 : U32)))) 1294 1295 # block bi: A's columns o..o+w (A is n x k) and B's rows o..o+w (B is k x m), multiplied into C (added when ac) 1296 def mp_step(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +k32: U32, +bi: U32, a: Array, b: Array, c: Array) -> G: 1297 mp_s2(ac, par, n, m, blk_len(bi, k32), U32.from_nat(blk_len(bi, k32)), U32.from_nat(m), (bi * 1024 : U32), c, b, sl_rows(n, 0, blk_len(bi, k32), U32.from_nat(blk_len(bi, k32)), k32, (bi * 1024 : U32), SL{a, Array.new(F32, cap_depth(Nat.mul(n, blk_len(bi, k32))), 0.0)})) 1298 1299 def mp_go(nb: Nat, +bi: U32, +par: Nat, +n: Nat, +m: Nat, +k32: U32, g: G) -> G: 1300 match nb g: 1301 case 0n g: g 1302 case 1n++q G{a, b, c}: mp_go(q, (bi + 1 : U32), par, n, m, k32, mp_step(U32.is_lt(0, bi), par, n, m, k32, bi, a, b, c)) 1303 1304 def mm_precise2(long: Bool, +n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat, b: Mat) -> MMul: 1305 match long a b: 1306 case True{} Mat{x} Mat{y}: mmul_fin(n, k, m, mp_go(U32.to_nat(U32.div((U32.from_nat(k) + 1023 : U32), 1024)), 0, par, n, m, U32.from_nat(k), G{x, y, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0)})) 1307 case False{} a b: Mat.matmul(n, k, m, par, a, b) 1308 1309 # C = A · B (n x k) · (k x m), k summed in blocks of 1024 (above); opt-in 1310 def Mat.matmul_precise(+n: Nat, +k: Nat, +m: Nat, +par: Nat, a: Mat, b: Mat) -> MMul: 1311 mm_precise2(Nat.is_lt(1024n, k), n, k, m, par, a, b) 1312 1313 # ---- element-wise and training operations ---- 1314 1315 # adds the bias (one row, m numbers) to each row of C (n x m) 1316 1317 type MBias<-n: Nat, -m: Nat> is Type: 1318 MBias{c: Mat, b: Mat<1n, m>} 1319 1320 def add_row2(-n: Nat, -m: Nat, +nn: Nat, +mm: Nat, cd: Array, al: AL) -> MBias: 1321 match al: 1322 case AL{bd, bl}: 1323 MBias{Mat{bias_rows(nn, 0, U32.from_nat(mm), bl, cd)}, Mat{bd}} 1324 1325 def Mat.add_row(+n: Nat, +m: Nat, c: Mat, b: Mat<1n, m>) -> MBias: 1326 match c b: 1327 case Mat{cd} Mat{bd}: 1328 add_row2(n, m, n, m, cd, read_l(Nat.sub(m, 1n), U32.from_nat(Nat.sub(m, 1n)), Nil{}, Array.get(F32, bd, U32.from_nat(Nat.sub(m, 1n))))) 1329 1330 # pair of matrices of the same shape (the result of an operation that reads two operands) 1331 1332 type MPair<-r: Nat, -c: Nat> is Type: 1333 MPair{a: Mat, b: Mat} 1334 1335 def mpair_fin(-r: Nat, -c: Nat, p: P2) -> MPair: 1336 match p: 1337 case P2{a, b}: 1338 MPair{Mat{a}, Mat{b}} 1339 1340 def relu_fin(-r: Nat, -c: Nat, d: Array) -> Mat: 1341 Mat{d} 1342 1343 def Mat.relu(+r: Nat, +c: Nat, a: Mat) -> Mat: 1344 match a: 1345 case Mat{d}: 1346 relu_fin(r, c, relu_l(Nat.mul(r, c), 0, Array.get(F32, d, 0))) 1347 1348 # (dy * (h > 0), h): the gradient that passes through the relu 1349 1350 def Mat.relu_bwd(+r: Nat, +c: Nat, dy: Mat, h: Mat) -> MPair: 1351 match dy h: 1352 case Mat{d} Mat{x}: 1353 mpair_fin(r, c, mask_l(Nat.mul(r, c), 0, Array.get(F32, d, 0), Array.get(F32, x, 0))) 1354 1355 # (w - lr * dw, dw): one SGD step 1356 1357 def Mat.sgd(+r: Nat, +c: Nat, +lr: F32, w: Mat, dw: Mat) -> MPair: 1358 match w dw: 1359 case Mat{x} Mat{y}: 1360 mpair_fin(r, c, sgd_l(Nat.mul(r, c), 0, lr, Array.get(F32, x, 0), Array.get(F32, y, 0))) 1361 1362 # (a + b, b): element-wise sum (a - (-1)*b is exactly a + b) 1363 1364 def Mat.add(+r: Nat, +c: Nat, a: Mat, b: Mat) -> MPair: 1365 match a b: 1366 case Mat{x} Mat{y}: 1367 mpair_fin(r, c, sgd_l(Nat.mul(r, c), 0, (0.0 - 1.0 : F32), Array.get(F32, x, 0), Array.get(F32, y, 0))) 1368 1369 # sum of each column (n x m -> 1 x m), returning the matrix 1370 1371 type MColSum<-n: Nat, -m: Nat> is Type: 1372 MColSum{a: Mat, s: Mat<1n, m>} 1373 1374 def colsum_fin(-n: Nat, -m: Nat, g: G) -> MColSum: 1375 match g: 1376 case G{ones, a, s}: 1377 MColSum{Mat{a}, Mat{s}} 1378 1379 def Mat.col_sums(+n: Nat, +m: Nat, a: Mat) -> MColSum: 1380 match a: 1381 case Mat{d}: 1382 colsum_fin(n, m, gemm(1n, m, n, 0, 1, U32.from_nat(m), 1, Array.new(F32, cap_depth(n), 1.0), d, Array.new(F32, cap_depth(m), 0.0))) 1383 1384 # s + the sum of each column, into s (1 x m): a bias gradient added into its accumulator, with the bits of 1385 # Mat.add(s, Mat.col_sums(a)'s s) 1386 1387 def Mat.col_sums_acc(+n: Nat, +m: Nat, a: Mat, s: Mat<1n, m>) -> MColSum: 1388 match a s: 1389 case Mat{d} Mat{z}: 1390 colsum_fin(n, m, gemm_ac(True{}, 1n, m, n, 0, 1, U32.from_nat(m), 1, Array.new(F32, cap_depth(n), 1.0), d, z)) 1391 1392 # mean cross-entropy and its gradient on the logits: (softmax - one-hot) / n. 1393 # labels has one entry per row (n); the gradient replaces the logits. 1394 1395 type MCE<-n: Nat, -c: Nat> is Type: 1396 MCE{loss: F32, grad: Mat} 1397 1398 def ce_fin(-n: Nat, -c: Nat, +nf: F32, r: L2) -> MCE: 1399 match r: 1400 case L2{loss, d}: 1401 MCE{(loss / nf : F32), Mat{d}} 1402 1403 def Mat.softmax_ce(+n: Nat, +c: Nat, z: Mat, labels: List<&2, Nat>) -> MCE: 1404 match z: 1405 case Mat{d}: 1406 ce_fin(n, c, F32.from_nat(n), cea_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), F32.from_nat(n), 0, L2{0.0, d})) 1407 1408 # argmax hits per row against labels 1409 1410 type MHits<-n: Nat, -c: Nat> is Type: 1411 MHits{z: Mat, hits: Nat} 1412 1413 def hits_fin(-n: Nat, -c: Nat, r: CR) -> MHits: 1414 match r: 1415 case CR{d, k}: 1416 MHits{Mat{d}, k} 1417 1418 def Mat.count_correct(+n: Nat, +c: Nat, z: Mat, labels: List<&2, Nat>) -> MHits: 1419 match z: 1420 case Mat{d}: 1421 hits_fin(n, c, count_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), 0, CR{d, 0n})) 1422 1423 # a row as a list (returning the matrix) and the other way around 1424 1425 type MRow<-r: Nat, -c: Nat> is Type: 1426 MRow{m: Mat, row: List<&2, F32>} 1427 1428 def row_fin(-r: Nat, -c: Nat, al: AL) -> MRow: 1429 match al: 1430 case AL{d, l}: 1431 MRow{Mat{d}, l} 1432 1433 def Mat.read_row(+r: Nat, +c: Nat, m: Mat, +i: U32) -> MRow: 1434 match m: 1435 case Mat{d}: 1436 row_fin(r, c, read_l(Nat.sub(c, 1n), ((i * U32.from_nat(c) : U32) + U32.from_nat(Nat.sub(c, 1n)) : U32), Nil{}, Array.get(F32, d, ((i * U32.from_nat(c) : U32) + U32.from_nat(Nat.sub(c, 1n)) : U32)))) 1437 1438 def Mat.write_row(-r: Nat, +c: Nat, m: Mat, +i: U32, xs: List<&2, F32>) -> Mat: 1439 match m: 1440 case Mat{d}: 1441 Mat{write_l(xs, (i * U32.from_nat(c) : U32), d)} 1442 1443 # ---- variants that check the number of labels ---- 1444 1445 # does the list of labels have exactly n entries (one per row)? 1446 def labels_len_is(xs: List<&2, Nat>, n: Nat) -> Bool: 1447 match xs n: 1448 case Nil{} 0n: 1449 True{} 1450 case Con{h, t} 1n+p: 1451 labels_len_is(t, p) 1452 case _ _: 1453 False{} 1454 1455 def ce_checked(ok: Bool, +n: Nat, +c: Nat, z: Mat, labels: List<&2, Nat>) -> Maybe<&1, MCE>: 1456 match ok: 1457 case True{}: 1458 Some{Mat.softmax_ce(n, c, z, labels)} 1459 case False{}: 1460 None{} 1461 1462 # like Mat.softmax_ce, but returns None unless labels has exactly n entries 1463 def Mat.softmax_ce_checked(+n: Nat, +c: Nat, z: Mat, +labels: List<&2, Nat>) -> Maybe<&1, MCE>: 1464 ce_checked(labels_len_is(labels, n), n, c, z, labels) 1465 1466 def hits_checked(ok: Bool, +n: Nat, +c: Nat, z: Mat, labels: List<&2, Nat>) -> Maybe<&1, MHits>: 1467 match ok: 1468 case True{}: 1469 Some{Mat.count_correct(n, c, z, labels)} 1470 case False{}: 1471 None{} 1472 1473 # like Mat.count_correct, but returns None unless labels has exactly n entries 1474 def Mat.count_correct_checked(+n: Nat, +c: Nat, z: Mat, +labels: List<&2, Nat>) -> Maybe<&1, MHits>: 1475 hits_checked(labels_len_is(labels, n), n, c, z, labels) 1476 1477 # --------------------------------------------------------------------- 1478 # Bands: a matrix cut into bands of rows, for matrix · vector in parallel 1479 # without copying the weights. 1480 # 1481 # Mat.matmul_nt with n = 1 can only run in parallel by cloning the whole weight 1482 # matrix for every task, and in matrix · vector each weight is read once, so the 1483 # copies cost more than the products (NOTES.md, exp. 9). Bands keeps each band in 1484 # its own Array: a task takes its band and nothing is copied except the input vector. 1485 # 1486 # A node with r rows has children with half(r) and r - half(r) rows. The split is 1487 # part of the type, so the row count of every band follows from r, and the law 1488 # half_cover says that the two halves always cover the r rows exactly. 1489 # --------------------------------------------------------------------- 1490 1491 # n / 2 rounded down. A tail loop with an accumulator: a non-tail recursion here would put a 1492 # continuation on every node of the band tree and keep the bands from running in parallel. 1493 def half.go(n: Nat, acc: Nat) -> Nat: 1494 match n: 1495 case 0n: 1496 acc 1497 case 1n+p: 1498 match p: 1499 case 0n: 1500 acc 1501 case 1n+q: 1502 half.go(q, 1n+acc) 1503 1504 def half(n: Nat) -> Nat: 1505 half.go(n, 0n) 1506 1507 # a <= a 1508 def leb_refl(+a: Nat) -> {True{} == leb(a, a) : Bool}: 1509 match a: 1510 case 0n: 1511 {==} 1512 case 1n+p: 1513 leb_refl(p) 1514 1515 # 1 + (a + b) == a + (1 + b) 1516 def add_succ_r(+a: Nat, -b: Nat) -> {1n+Nat.add(a, b) == Nat.add(a, 1n+b) : Nat}: 1517 match a: 1518 case 0n: 1519 {==} 1520 case 1n+p: 1521 %add_succ_r(p, b) : {2n+Nat.add(p, b) == 1n+_ : Nat} 1522 {==} 1523 1524 # a + 0 == a 1525 def add_zero_r(+a: Nat) -> {Nat.add(a, 0n) == a : Nat}: 1526 match a: 1527 case 0n: 1528 {==} 1529 case 1n+p: 1530 %add_zero_r(p) : {1n+Nat.add(p, 0n) == 1n+_ : Nat} 1531 {==} 1532 1533 # half.go(n, acc) <= n + acc 1534 def half_go_le(+n: Nat, +acc: Nat) -> {True{} == leb(half.go(n, acc), Nat.add(n, acc)) : Bool}: 1535 match n: 1536 case 0n: 1537 leb_refl(acc) 1538 case 1n+p: 1539 match p: 1540 case 0n: 1541 leb_succ(acc, acc, leb_refl(acc)) 1542 case 1n+q: 1543 %Equal.sym(Nat, 1n+Nat.add(q, acc), Nat.add(q, 1n+acc), add_succ_r(q, acc)) : {True{} == leb(half.go(q, 1n+acc), 1n+_) : Bool} 1544 leb_succ(half.go(q, 1n+acc), Nat.add(q, 1n+acc), half_go_le(q, 1n+acc)) 1545 1546 # half(n) <= n 1547 def half_le(+n: Nat) -> {True{} == leb(half(n), n) : Bool}: 1548 %add_zero_r(n) : {True{} == leb(half(n), _) : Bool} 1549 half_go_le(n, 0n) 1550 1551 # a <= r -> r == a + (r - a) 1552 def add_sub(+a: Nat, +r: Nat, e: {True{} == leb(a, r) : Bool}) -> {r == Nat.add(a, Nat.sub(r, a)) : Nat}: 1553 match a r: 1554 case 0n 0n: 1555 {==} 1556 case 0n 1n+q: 1557 {==} 1558 case 1n+p 0n: 1559 %e : BD(_, Unit, {0n == Nat.add(1n+p, Nat.sub(0n, 1n+p)) : Nat}) 1560 Unit{} 1561 case 1n+p 1n+q: 1562 %add_sub(p, q, e) : {1n+q == 1n+_ : Nat} 1563 {==} 1564 1565 # LAW (half_cover): the two bands of a node, half(r) rows and r - half(r) rows, 1566 # add up to exactly r rows: no row is lost and none is counted twice. 1567 law half_cover: 1568 for +n: Nat 1569 {n == Nat.add(half(n), Nat.sub(n, half(n))) : Nat} 1570 1571 def half_cover(n): 1572 add_sub(half(n), n, half_le(n)) 1573 1574 type Bands<-r: Nat, -c: Nat> is Type: 1575 BLeaf{m: Mat} 1576 BNode{x: Bands, y: Bands} 1577 1578 # r x c of zeros in 2^d bands (d = 0: a single band, the same as a Mat) 1579 def Bands.zeros(+d: Nat, +r: Nat, +c: Nat) -> Bands: 1580 match d: 1581 case 0n: 1582 BLeaf{Mat.zeros(r, c)} 1583 case 1n+q: 1584 # bend-lint: disable-next bend-lint-rules/performance/missed_fork -- bands are built once at load not-expected-because: forking them measured as no change (NOTES 2026-10-16) 1585 BNode{Bands.zeros(q, half(r), c), Bands.zeros(q, Nat.sub(r, half(r)), c)} 1586 1587 # the row counts of the two bands follow the row count of the node: if rr == r at run 1588 # time, the same holds for both halves (the bands carry r only in their type) 1589 def eq_half(-a: Nat, -b: Nat, e: {a == b : Nat}) -> {half(a) == half(b) : Nat}: 1590 %e : {half(a) == half(_) : Nat} 1591 {==} 1592 1593 def eq_rest(-a: Nat, -b: Nat, e: {a == b : Nat}) -> {Nat.sub(a, half(a)) == Nat.sub(b, half(b)) : Nat}: 1594 %e : {Nat.sub(a, half(a)) == Nat.sub(_, half(_)) : Nat} 1595 {==} 1596 1597 # does a block of n numbers starting at flat index i end before the second band of a 1598 # node with r rows? does it start in the second band? 1599 def bfill_l(+r: Nat, +c: Nat, +i: Nat, +n: Nat) -> Bool: 1600 Nat.is_le(Nat.add(i, n), Nat.mul(half(r), c)) 1601 1602 def bfill_r(+r: Nat, +c: Nat, +i: Nat) -> Bool: 1603 Nat.is_le(Nat.mul(half(r), c), i) 1604 1605 # writes the first n numbers of xs at index i (and nothing past them) 1606 def fill_n(n: Nat, xs: List<&2, F32>, +i: U32, a: Array) -> Array: 1607 match n xs: 1608 case 1n+p Con{h, t}: 1609 fill_n(p, t, (i + 1 : U32), Array.set(F32, a, i, h)) 1610 case _ _: 1611 a 1612 1613 def fill_band(-r: Nat, -c: Nat, n: Nat, xs: List<&2, F32>, +i: U32, m: Mat) -> Mat: 1614 match m: 1615 case Mat{d}: 1616 Mat{fill_n(n, xs, i, d)} 1617 1618 # writes the first n numbers of xs at flat index i; l and rt say which bands they touch 1619 # (bfill_l, bfill_r); rr is the row count r at run time. A block that crosses into the second 1620 # band is not copied: the first band writes only its part, the second gets the rest by List.drop. 1621 def bfill(+c: Nat, -r: Nat, b: Bands, l: Bool, rt: Bool, +rr: Nat, -e: {rr == r : Nat}, +i: Nat, +n: Nat, +xs: List<&2, F32>) -> Bands: 1622 match b l rt: 1623 case BLeaf{m} _ _: 1624 BLeaf{fill_band(r, c, n, xs, U32.from_nat(i), m)} 1625 case BNode{x, y} True{} _: 1626 BNode{bfill(c, half(r), x, bfill_l(half(rr), c, i, n), bfill_r(half(rr), c, i), half(rr), eq_half(rr, r, e), i, n, xs), y} 1627 case BNode{x, y} False{} True{}: 1628 BNode{x, bfill(c, Nat.sub(r, half(r)), y, bfill_l(Nat.sub(rr, half(rr)), c, Nat.sub(i, Nat.mul(half(rr), c)), n), bfill_r(Nat.sub(rr, half(rr)), c, Nat.sub(i, Nat.mul(half(rr), c))), Nat.sub(rr, half(rr)), eq_rest(rr, r, e), Nat.sub(i, Nat.mul(half(rr), c)), n, xs)} 1629 case BNode{x, y} False{} False{}: 1630 # bend-lint: disable-next bend-lint-rules/performance/missed_fork -- bands are built once at load not-expected-because: forking them measured as no change (NOTES 2026-10-16) 1631 BNode{bfill(c, half(r), x, bfill_l(half(rr), c, i, Nat.sub(Nat.mul(half(rr), c), i)), bfill_r(half(rr), c, i), half(rr), eq_half(rr, r, e), i, Nat.sub(Nat.mul(half(rr), c), i), xs), bfill(c, Nat.sub(r, half(r)), y, bfill_l(Nat.sub(rr, half(rr)), c, 0n, Nat.sub(n, Nat.sub(Nat.mul(half(rr), c), i))), bfill_r(Nat.sub(rr, half(rr)), c, 0n), Nat.sub(rr, half(rr)), eq_rest(rr, r, e), 0n, Nat.sub(n, Nat.sub(Nat.mul(half(rr), c), i)), List.drop(&2, F32, xs, Nat.sub(Nat.mul(half(rr), c), i)))} 1632 1633 def bfill_top(ok: Bool, +r: Nat, +c: Nat, b: Bands, +i: Nat, +n: Nat, +xs: List<&2, F32>) -> Bands: 1634 match ok: 1635 case True{}: 1636 bfill(c, r, b, bfill_l(r, c, i, n), bfill_r(r, c, i), r, {==}, i, n, xs) 1637 case False{}: 1638 bfill(c, r, b, bfill_l(r, c, i, Nat.sub(Nat.mul(r, c), i)), bfill_r(r, c, i), r, {==}, i, Nat.sub(Nat.mul(r, c), i), xs) 1639 1640 # writes the n numbers of xs into b starting at flat index i (row-major, as Mat.fill_at); 1641 # numbers past r*c are dropped. For loading a large matrix block by block. 1642 def Bands.fill_at(+r: Nat, +c: Nat, +i: Nat, +n: Nat, +xs: List<&2, F32>, b: Bands) -> Bands: 1643 bfill_top(Nat.is_le(Nat.add(i, n), Nat.mul(r, c)), r, c, b, i, n, xs) 1644 1645 # ---- y = B · x for B with r rows of c numbers (a weight matrix stored out x in) ---- 1646 1647 # B (intact) and y = B · x, typed like Mat.matmul_nt: x is 1 x c, y is 1 x r 1648 type BV<-r: Nat, -c: Nat> is Type: 1649 BV{b: Bands, y: Mat<1n, r>} 1650 1651 # the bands and y as a list in row order (the result of Bands.matvec_l) 1652 type BL<-r: Nat, -c: Nat> is Type: 1653 BL{b: Bands, y: List<&2, F32>} 1654 1655 def bv_g2(-r: Nat, -c: Nat, w: Array, al: AL) -> BL: 1656 match al: 1657 case AL{d, l}: 1658 BL{BLeaf{Mat{w}}, l} 1659 1660 # the band's 1 + m rows came out in an Array (one number per row); read them back as a list 1661 def bv_g(-r: Nat, -c: Nat, +m: Nat, g: G) -> BL: 1662 match g: 1663 case G{a, b, y}: 1664 bv_g2(r, c, b, read_l(m, U32.from_nat(m), Nil{}, Array.get(F32, y, U32.from_nat(m)))) 1665 1666 # the rows of one band, one dot product each, with the gemm of Mat.matmul_nt (n = 1), so the 1667 # results are identical. Writing the results into an Array and reading them back is faster than 1668 # building the list row by row (NOTES.md, exp. 14). 1669 def bv_gemm(-r: Nat, +c: Nat, rows: Nat, +rr: Nat, x: Array, w: Array) -> BL: 1670 match rows: 1671 case 0n: 1672 BL{BLeaf{Mat{w}}, Nil{}} 1673 case 1n+p: 1674 bv_g(r, c, p, gemm(1n, rr, c, U32.from_nat(c), 1, 1, U32.from_nat(c), x, w, Array.new(F32, cap_depth(rr), 0.0))) 1675 1676 def bv_mat(-r: Nat, +c: Nat, +rr: Nat, x: Array, m: Mat) -> BL: 1677 match m: 1678 case Mat{w}: 1679 bv_gemm(r, c, rr, rr, x, w) 1680 1681 def bv_join(-r: Nat, -c: Nat, p: BL, q: BL) -> BL: 1682 match p q: 1683 case BL{bx, yx} BL{by, yy}: 1684 BL{BNode{bx, by}, List.append(&2, F32, yx, yy)} 1685 1686 # the two bands of a node run in parallel, each with its own copy of x (c numbers); 1687 # xp is a pair of copies of x (a leaf uses one). xp comes after the erased e: up to Bend 2.0.36 an erased 1688 # parameter at the end of the list turned the parallel let into two sequential calls (NOTES.md, exp. 11; fixed in 1689 # 2.0.37 by bendlang/bend#1377). The order is kept, and harmless. 1690 def bmv(+c: Nat, -r: Nat, b: Bands, +rr: Nat, -e: {rr == r : Nat}, xp: Array & Array) -> BL: 1691 match b xp: 1692 case BLeaf{m} Tuple{x, x2}: 1693 bv_mat(r, c, rr, x, m) 1694 case BNode{bx, by} Tuple{x1, x2}: 1695 p q = bmv(c, half(r), bx, half(rr), eq_half(rr, r, e), Array.clone(F32, x1)) bmv(c, Nat.sub(r, half(r)), by, Nat.sub(rr, half(rr)), eq_rest(rr, r, e), Array.clone(F32, x2)) 1696 bv_join(r, c, p, q) 1697 1698 def bv_fin(+r: Nat, -c: Nat, bl: BL) -> BV: 1699 match bl: 1700 case BL{b, y}: 1701 BV{b, Mat.from_list(1n, r, y)} 1702 1703 # y = B · x for B (r x c, one row per output) and x (1 x c); returns B intact and y (1 x r). 1704 # The same product as Mat.matmul_nt(1n, c, r, ..), and the same numbers, bit for bit. 1705 def Bands.matvec(+r: Nat, +c: Nat, b: Bands, x: Mat<1n, c>) -> BV: 1706 match x: 1707 case Mat{xa}: 1708 bv_fin(r, c, bmv(c, r, b, r, {==}, Array.clone(F32, xa))) 1709 1710 # the same product with x and y as lists (x: c numbers, y: r numbers in row order), without the 1711 # Mat conversions of Bands.matvec: for callers that already hold the vector as a list 1712 def Bands.matvec_l(+r: Nat, +c: Nat, b: Bands, +x: List<&2, F32>) -> BL: 1713 bmv(c, r, b, r, {==}, Array.clone(F32, fill_list(x, 0, Array.new(F32, cap_depth(c), 0.0)))) 1714 1715 # ---- reading rows and the whole matrix ---- 1716 1717 type BRow<-r: Nat, -c: Nat> is Type: 1718 BRow{b: Bands, row: List<&2, F32>} 1719 1720 def brow_fin(-r: Nat, -c: Nat, al: AL) -> BRow: 1721 match al: 1722 case AL{d, l}: 1723 BRow{BLeaf{Mat{d}}, l} 1724 1725 # n numbers starting at flat index i of one band (a band may have no rows) 1726 def brow_leaf2(-r: Nat, -c: Nat, +n: Nat, +i: U32, d: Array) -> BRow: 1727 match n: 1728 case 0n: 1729 BRow{BLeaf{Mat{d}}, Nil{}} 1730 case 1n+p: 1731 brow_fin(r, c, read_l(p, ((i + U32.from_nat(n) : U32) - 1 : U32), Nil{}, Array.get(F32, d, ((i + U32.from_nat(n) : U32) - 1 : U32)))) 1732 1733 def brow_leaf(-r: Nat, -c: Nat, m: Mat, +i: U32, n: Nat) -> BRow: 1734 match m: 1735 case Mat{d}: 1736 brow_leaf2(r, c, n, i, d) 1737 1738 def brow_l(-r: Nat, -c: Nat, p: BRow, y: Bands) -> BRow: 1739 match p: 1740 case BRow{bx, row}: 1741 BRow{BNode{bx, y}, row} 1742 1743 def brow_r(-r: Nat, -c: Nat, x: Bands, q: BRow) -> BRow: 1744 match q: 1745 case BRow{by, row}: 1746 BRow{BNode{x, by}, row} 1747 1748 # row i; left says whether i is in the first band of the node (i < half(rr)) 1749 def brow(+c: Nat, -r: Nat, b: Bands, left: Bool, +rr: Nat, -e: {rr == r : Nat}, +i: Nat) -> BRow: 1750 match b left: 1751 case BLeaf{m} _: 1752 brow_leaf(r, c, m, U32.from_nat(Nat.mul(i, c)), c) 1753 case BNode{x, y} True{}: 1754 brow_l(r, c, brow(c, half(r), x, Nat.is_lt(i, half(half(rr))), half(rr), eq_half(rr, r, e), i), y) 1755 case BNode{x, y} False{}: 1756 brow_r(r, c, x, brow(c, Nat.sub(r, half(r)), y, Nat.is_lt(Nat.sub(i, half(rr)), half(Nat.sub(rr, half(rr)))), Nat.sub(rr, half(rr)), eq_rest(rr, r, e), Nat.sub(i, half(rr)))) 1757 1758 # row i (c numbers), for example an embedding lookup 1759 def Bands.read_row(+r: Nat, +c: Nat, b: Bands, +i: Nat) -> BRow: 1760 brow(c, r, b, Nat.is_lt(i, half(r)), r, {==}, i) 1761 1762 def bl_join(-r: Nat, -c: Nat, p: BRow, q: BRow) -> BRow: 1763 match p q: 1764 case BRow{bx, xs} BRow{by, ys}: 1765 BRow{BNode{bx, by}, List.append(&2, F32, xs, ys)} 1766 1767 def blist(+c: Nat, -r: Nat, b: Bands, +rr: Nat, -e: {rr == r : Nat}) -> BRow: 1768 match b: 1769 case BLeaf{m}: 1770 brow_leaf(r, c, m, 0, Nat.mul(rr, c)) 1771 case BNode{x, y}: 1772 # bend-lint: disable-next bend-lint-rules/performance/missed_fork -- bands are built once at load not-expected-because: forking them measured as no change (NOTES 2026-10-16) 1773 bl_join(r, c, blist(c, half(r), x, half(rr), eq_half(rr, r, e)), blist(c, Nat.sub(r, half(r)), y, Nat.sub(rr, half(rr)), eq_rest(rr, r, e))) 1774 1775 # all r*c numbers in row order (in BRow.row) 1776 def Bands.to_list(+r: Nat, +c: Nat, b: Bands) -> BRow: 1777 blist(c, r, b, r, {==}) 1778 1779 # ---- the list path returns exactly r numbers ---- 1780 # 1781 # The numbers do not matter here, only how many there are: each lemma below counts the list a 1782 # function builds, following the function's own recursion. 1783 1784 def bl_len(-r: Nat, -c: Nat, bl: BL) -> Nat: 1785 match bl: 1786 case BL{b, y}: 1787 List.length(&2, F32, y) 1788 1789 # s == a + (1 + b) -> s == 1 + (a + b) 1790 def succ_out(-s: Nat, +a: Nat, +b: Nat, e: {s == Nat.add(a, 1n+b) : Nat}) -> {s == 1n+Nat.add(a, b) : Nat}: 1791 %Equal.sym(Nat, 1n+Nat.add(a, b), Nat.add(a, 1n+b), add_succ_r(a, b)) : {s == _ : Nat} 1792 e 1793 1794 # s == 1 + (p + (1 + l)) -> s == 2 + (p + l) 1795 def len_step(-s: Nat, +p: Nat, +l: Nat, e: {s == 1n+Nat.add(p, 1n+l) : Nat}) -> {s == 2n+Nat.add(p, l) : Nat}: 1796 %Equal.sym(Nat, 1n+Nat.add(p, l), Nat.add(p, 1n+l), add_succ_r(p, l)) : {s == 1n+_ : Nat} 1797 e 1798 1799 # s == n + 0 -> s == n 1800 def zero_out(-s: Nat, +n: Nat, e: {s == Nat.add(n, 0n) : Nat}) -> {s == n : Nat}: 1801 %add_zero_r(n) : {s == _ : Nat} 1802 e 1803 1804 # appending adds the lengths 1805 def app_len(+xs: List<&2, F32>, +ys: List<&2, F32>) -> {List.length(&2, F32, List.append(&2, F32, xs, ys)) == Nat.add(List.length(&2, F32, xs), List.length(&2, F32, ys)) : Nat}: 1806 match xs: 1807 case Nil{}: 1808 {==} 1809 case Con{h, t}: 1810 %app_len(t, ys) : {1n+List.length(&2, F32, List.append(&2, F32, t, ys)) == 1n+_ : Nat} 1811 {==} 1812 1813 # the same count for the list of a whole Mat: to_list2 wraps read_l's answer, and ML.len counts it 1814 def read_ml_len(-r: Nat, -c: Nat, +mleft: Nat, +idx: U32, +acc: List<&2, F32>, rc: Array & F32) -> {ML.len(r, c, Mat.to_list2(r, c, read_l(mleft, idx, acc, rc))) == Nat.add(1n+mleft, List.length(&2, F32, acc)) : Nat}: 1815 match mleft rc: 1816 case 0n Tuple{d, +v}: 1817 {==} 1818 case 1n+p Tuple{d, +v}: 1819 len_step(ML.len(r, c, Mat.to_list2(r, c, read_l(p, (idx - 1 : U32), v <> acc, Array.get(F32, d, (idx - 1 : U32))))), p, List.length(&2, F32, acc), read_ml_len(r, c, p, (idx - 1 : U32), v <> acc, Array.get(F32, d, (idx - 1 : U32)))) 1820 1821 # a leaf's list: read_l reads mleft + 1 numbers onto acc. By induction on mleft: with 0 left it 1822 # reads one number onto acc (1 + the length of acc, by the definitions); with 1 + p left it reads 1823 # one and recurses, and len_step moves the 1 out of the sum. 1824 law leaf_len: 1825 for -r: Nat 1826 for -c: Nat 1827 for w: Array 1828 for +mleft: Nat 1829 for +idx: U32 1830 for +acc: List<&2, F32> 1831 for rc: Array & F32 1832 {bl_len(r, c, bv_g2(r, c, w, read_l(mleft, idx, acc, rc))) == Nat.add(1n+mleft, List.length(&2, F32, acc)) : Nat} 1833 1834 def leaf_len(r, c, w, mleft, idx, acc, rc): 1835 match mleft rc: 1836 case 0n Tuple{d, +v}: 1837 {==} 1838 case 1n+p Tuple{d, +v}: 1839 len_step(bl_len(r, c, bv_g2(r, c, w, read_l(p, (idx - 1 : U32), v <> acc, Array.get(F32, d, (idx - 1 : U32))))), p, List.length(&2, F32, acc), leaf_len(r, c, w, p, (idx - 1 : U32), v <> acc, Array.get(F32, d, (idx - 1 : U32)))) 1840 1841 # whatever the gemm wrote, reading 1 + p rows back gives 1 + p numbers (G has one constructor, so 1842 # matching g opens it, and leaf_len counts the list) 1843 def g_len(-r: Nat, -c: Nat, +p: Nat, g: G) -> {bl_len(r, c, bv_g(r, c, p, g)) == 1n+p : Nat}: 1844 match g: 1845 case G{a, b, y}: 1846 zero_out(bl_len(r, c, bv_g2(r, c, b, read_l(p, U32.from_nat(p), Nil{}, Array.get(F32, y, U32.from_nat(p))))), 1n+p, leaf_len(r, c, b, p, U32.from_nat(p), Nil{}, Array.get(F32, y, U32.from_nat(p)))) 1847 1848 # LAW (band_len): the product of one band with `rows` rows gives exactly `rows` numbers. 1849 # 0 rows give the empty list; 1 + p rows go through the gemm, and g_len counts what is read back. 1850 law band_len: 1851 for -r: Nat 1852 for +c: Nat 1853 for +rows: Nat 1854 for x: Array 1855 for w: Array 1856 {bl_len(r, c, bv_gemm(r, c, rows, rows, x, w)) == rows : Nat} 1857 1858 def band_len(r, c, rows, x, w): 1859 match rows: 1860 case 0n: 1861 {==} 1862 case 1n+p: 1863 g_len(r, c, p, gemm(1n, 1n+p, c, U32.from_nat(c), 1, 1, U32.from_nat(c), x, w, Array.new(F32, cap_depth(1n+p), 0.0))) 1864 1865 # joining two halves of h1 and h2 numbers gives h1 + h2 numbers. (The statement for the whole tree, 1866 # "Bands.matvec_l gives r numbers", would apply this to the results of the two recursive calls and 1867 # to the induction hypotheses about the same calls; that uses each band's Array twice, which Bend's 1868 # affine rules refuse. It stays open, covered by tests: NOTES.md, v3.1.) 1869 def join_len(-r: Nat, -c: Nat, -h1: Nat, -h2: Nat, p: BL, q: BL, hp: {bl_len(half(r), c, p) == h1 : Nat}, hq: {bl_len(Nat.sub(r, half(r)), c, q) == h2 : Nat}) -> {bl_len(r, c, bv_join(r, c, p, q)) == Nat.add(h1, h2) : Nat}: 1870 match p q: 1871 case BL{bx, +yx} BL{by, +yy}: 1872 %hp : {List.length(&2, F32, List.append(&2, F32, yx, yy)) == Nat.add(_, h2) : Nat} 1873 %hq : {List.length(&2, F32, List.append(&2, F32, yx, yy)) == Nat.add(List.length(&2, F32, yx), _) : Nat} 1874 app_len(yx, yy) 1875 1876 # Mat.to_list_count reads exactly n numbers, and n = 0 reads none. 1877 def to_list_count_len(-r: Nat, -c: Nat, d: Array, +n: Nat) -> {ML.len(r, c, Mat.to_list_count(r, c, d, n)) == n : Nat}: 1878 match n: 1879 case 0n: 1880 {==} 1881 case 1n+p: 1882 zero_out(ML.len(r, c, Mat.to_list_count(r, c, d, 1n+p)), 1n+p, read_ml_len(r, c, p, (U32.from_nat(1n+p) - 1 : U32), Nil{}, Array.get(F32, d, (U32.from_nat(1n+p) - 1 : U32)))) 1883 1884 # LAW (to_list_len): Mat.to_list reads exactly r*c numbers, no more and no fewer. 1885 law to_list_len: 1886 for +r: Nat 1887 for +c: Nat 1888 for m: Mat 1889 {ML.len(r, c, Mat.to_list(r, c, m)) == Nat.mul(r, c) : Nat} 1890 1891 def to_list_len(r, c, m): 1892 match m: 1893 case Mat{d}: 1894 to_list_count_len(r, c, d, Nat.mul(r, c)) 1895 1896 # --------------------------------------------------------------------- 1897 # Base's Array: three laws about Array.set and Array.get 1898 # --------------------------------------------------------------------- 1899 # 1900 # - size_set: Array.set keeps Array.size; 1901 # - get_after_set: what Array.set writes at i, Array.get reads back at i; 1902 # - get_other_set: a write at i leaves alone the cell that a read at k reaches, when the two descents part. 1903 # They are about Base's Array, for any element type, and live here until Base states them 1904 # (bendlang/bend#1934; nuxyel/bend-ml#26). The helpers are named ASet.*. 1905 # 1906 # The first two hold for every i (the two descents wrap i the same way) and for every tree Base can build, 1907 # balanced or not, since a write replaces one leaf and leaves the tree's shape alone. "The size" is the 1908 # second component of the pair Array.size returns, read with Base's Pair.snd (ASet.sz); the value is 1909 # Pair.snd of the pair Array.get returns. 1910 # 1911 # The third law's hypothesis says that the descent Array.set makes for i and the descent Array.get makes 1912 # for k take different branches at some node before either reaches a leaf (ASet.Apart). It is not "i and 1913 # k differ after the mask", because that is false on an unbalanced tree: Array.size follows the leftmost 1914 # path, so in ANode{ANode{ALeaf, ALeaf}, ALeaf} (size 4) the indices 2 and 3 differ but both reach the 1915 # last leaf, and a write at 2 changes what a read at 3 returns (tests/array_laws.bend prints it: 9.0). On a 1916 # perfect tree of 2^d leaves, which Array.new and literals build, the descents part exactly when the masked 1917 # indices differ; that equivalence is not proved here, since it needs lemmas about U32 bits. For a concrete 1918 # tree ASet.Apart computes to Unit, so Unit{} proves it (tests/array_laws.bend). 1919 # 1920 # Why the proofs are shaped this way: an Array is affine, and so is it inside a proof, so a proof may use a 1921 # given array once. The obvious lemmas each need the same subtree again (one to say Array.size of it is 1922 # (it, its size), one for the induction step, one to open the pair the descent returns), so they cannot 1923 # be combined. Instead, one pass over the tree (ASet.facts) proves all six facts the step needs at once; 1924 # each fact is an equation, and equations can be reused (+). A second pass (ASet.copy) hands back the same 1925 # array and its size as a plain value, which the top-level proof needs to pick the first branch. The third 1926 # law uses one pass too (ASet.both): it proves the six facts and four more about the read along k, and on 1927 # the half the write does not enter it runs a smaller pass (ASet.side). 1928 # 1929 # The trailing Unit in ASet.Facts, ASet.Copy, ASet.Side and ASet.Reads works around bendlang/bend#1077: a + 1930 # binder on the last field of a nested pair fails to check, so no fact sits in the last field. 1931 1932 1933 # ASet.sz(T, a): the size of a as a plain U32, the second component of the pair Array.size returns 1934 def ASet.sz(-T: Data, a: Array) -> U32: 1935 Pair.snd(Array, U32, Array.size(T, a)) 1936 1937 # the three results the descents build: the swapped pair, the written array, and the read pair 1938 def ASet.G(-T: Data, a: Array, +n: U32, +j: U32, +v: T, +z: Bool) -> Array & T: 1939 Array.swap.go(T, a, n, j, v, z) 1940 1941 def ASet.NA(-T: Data, a: Array, +n: U32, +j: U32, +v: T, +z: Bool) -> Array: 1942 Array.set.fin(T, ASet.G(T, a, n, j, v, z)) 1943 1944 def ASet.K(-T: Data, a: Array, +n: U32, +j: U32, +v: T, +z: Bool) -> Array & T: 1945 Array.get.go(T, ASet.NA(T, a, n, j, v, z), n, j, z) 1946 1947 # ASet.Facts: what one pass over the tree a proves, for any n, j, v, z (n need not be a's real size): 1948 # 1. Array.size(a) is the pair (a, ASet.sz(a)); 1949 # 2. the pair the write's descent returns is (the written array, the old value); 1950 # 3. Array.size of the written array is the pair (written array, its sz); 1951 # 4. the written array has a's sz; 1952 # 5. reading the written array along the same path gives v; 1953 # 6. the pair the read returns is (its first component, its second). 1954 def ASet.Facts(-T: Data, -a: Array, -n: U32, -j: U32, -v: T, -z: Bool) -> Type: 1955 {(a, ASet.sz(T, a)) == Array.size(T, a) : Array & U32} & {(ASet.NA(T, a, n, j, v, z), Pair.snd(Array, T, ASet.G(T, a, n, j, v, z))) == ASet.G(T, a, n, j, v, z) : Array & T} & {(ASet.NA(T, a, n, j, v, z), ASet.sz(T, ASet.NA(T, a, n, j, v, z))) == Array.size(T, ASet.NA(T, a, n, j, v, z)) : Array & U32} & {ASet.sz(T, ASet.NA(T, a, n, j, v, z)) == ASet.sz(T, a) : U32} & {Pair.snd(Array, T, ASet.K(T, a, n, j, v, z)) == v : T} & {(Pair.fst(Array, T, ASet.K(T, a, n, j, v, z)), Pair.snd(Array, T, ASet.K(T, a, n, j, v, z))) == ASet.K(T, a, n, j, v, z) : Array & T} & Unit 1956 1957 # Array.size hands back the array it measured, next to ASet.sz(T, a) 1958 def ASet.size_pair(-T: Data, a: Array) -> {(a, ASet.sz(T, a)) == Array.size(T, a) : Array & U32}: 1959 match a: 1960 case ALeaf{x}: 1961 {==} 1962 case ANode{xs, ys}: 1963 %ASet.size_pair(T, xs) : {(ANode{xs, ys}, Pair.snd(Array, U32, Array.size.node(T, ys, _))) == Array.size.node(T, ys, _) : Array & U32} 1964 {==} 1965 1966 def ASet.node_t(-T: Data, -xs: Array, -ys: Array, +n: U32, +j: U32, +v: T, f: ASet.Facts(T, xs, U32.shr(n), j, v, U32.is_lt(j, U32.shr(U32.shr(n))))) -> ASet.Facts(T, ANode{xs, ys}, n, j, v, True{}): 1967 match f: 1968 case Tuple{+c1, r1}: 1969 match r1: 1970 case Tuple{+c2, r2}: 1971 match r2: 1972 case Tuple{+c3, r3}: 1973 match r3: 1974 case Tuple{+c4, r4}: 1975 match r4: 1976 case Tuple{+c5, r5}: 1977 match r5: 1978 case Tuple{+c6, u}: 1979 ({(%c1 : {(ANode{xs, ys}, Pair.snd(Array, U32, Array.size.node(T, ys, _))) == Array.size.node(T, ys, _) : Array & U32}; {==}) : {(ANode{xs, ys}, ASet.sz(T, ANode{xs, ys})) == Array.size(T, ANode{xs, ys}) : Array & U32}}, {(%c2 : {(Array.set.fin(T, Array.swap.lo(T, ys, _)), Pair.snd(Array, T, Array.swap.lo(T, ys, _))) == Array.swap.lo(T, ys, _) : Array & T}; {==}) : {(ASet.NA(T, ANode{xs, ys}, n, j, v, True{}), Pair.snd(Array, T, ASet.G(T, ANode{xs, ys}, n, j, v, True{}))) == ASet.G(T, ANode{xs, ys}, n, j, v, True{}) : Array & T}}, {(%c2 : {(Array.set.fin(T, Array.swap.lo(T, ys, _)), ASet.sz(T, Array.set.fin(T, Array.swap.lo(T, ys, _)))) == Array.size(T, Array.set.fin(T, Array.swap.lo(T, ys, _))) : Array & U32}; %c3 : {(ANode{Array.set.fin(T, ASet.G(T, xs, U32.shr(n), j, v, U32.is_lt(j, U32.shr(U32.shr(n))))), ys}, Pair.snd(Array, U32, Array.size.node(T, ys, _))) == Array.size.node(T, ys, _) : Array & U32}; {==}) : {(ASet.NA(T, ANode{xs, ys}, n, j, v, True{}), ASet.sz(T, ASet.NA(T, ANode{xs, ys}, n, j, v, True{}))) == Array.size(T, ASet.NA(T, ANode{xs, ys}, n, j, v, True{})) : Array & U32}}, {(%c2 : {ASet.sz(T, Array.set.fin(T, Array.swap.lo(T, ys, _))) == ASet.sz(T, ANode{xs, ys}) : U32}; %c3 : {Pair.snd(Array, U32, Array.size.node(T, ys, _)) == ASet.sz(T, ANode{xs, ys}) : U32}; %c1 : {U32.shl(ASet.sz(T, Array.set.fin(T, ASet.G(T, xs, U32.shr(n), j, v, U32.is_lt(j, U32.shr(U32.shr(n))))))) == Pair.snd(Array, U32, Array.size.node(T, ys, _)) : U32}; %c4 : {U32.shl(ASet.sz(T, Array.set.fin(T, ASet.G(T, xs, U32.shr(n), j, v, U32.is_lt(j, U32.shr(U32.shr(n))))))) == U32.shl(_) : U32}; {==}) : {ASet.sz(T, ASet.NA(T, ANode{xs, ys}, n, j, v, True{})) == ASet.sz(T, ANode{xs, ys}) : U32}}, {(%c2 : {Pair.snd(Array, T, Array.get.go(T, Array.set.fin(T, Array.swap.lo(T, ys, _)), n, j, True{})) == v : T}; %c6 : {Pair.snd(Array, T, Array.swap.lo(T, ys, _)) == v : T}; c5) : {Pair.snd(Array, T, ASet.K(T, ANode{ 1980 1981 def ASet.node_f(-T: Data, -xs: Array, -ys: Array, +n: U32, +j: U32, +v: T, +s: {(xs, ASet.sz(T, xs)) == Array.size(T, xs) : Array & U32}, f: ASet.Facts(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n))))) -> ASet.Facts(T, ANode{xs, ys}, n, j, v, False{}): 1982 match f: 1983 case Tuple{+c1, r1}: 1984 match r1: 1985 case Tuple{+c2, r2}: 1986 match r2: 1987 case Tuple{+c3, r3}: 1988 match r3: 1989 case Tuple{+c4, r4}: 1990 match r4: 1991 case Tuple{+c5, r5}: 1992 match r5: 1993 case Tuple{+c6, u}: 1994 ({(%s : {(ANode{xs, ys}, Pair.snd(Array, U32, Array.size.node(T, ys, _))) == Array.size.node(T, ys, _) : Array & U32}; {==}) : {(ANode{xs, ys}, ASet.sz(T, ANode{xs, ys})) == Array.size(T, ANode{xs, ys}) : Array & U32}}, {(%c2 : {(Array.set.fin(T, Array.swap.hi(T, xs, _)), Pair.snd(Array, T, Array.swap.hi(T, xs, _))) == Array.swap.hi(T, xs, _) : Array & T}; {==}) : {(ASet.NA(T, ANode{xs, ys}, n, j, v, False{}), Pair.snd(Array, T, ASet.G(T, ANode{xs, ys}, n, j, v, False{}))) == ASet.G(T, ANode{xs, ys}, n, j, v, False{}) : Array & T}}, {(%c2 : {(Array.set.fin(T, Array.swap.hi(T, xs, _)), ASet.sz(T, Array.set.fin(T, Array.swap.hi(T, xs, _)))) == Array.size(T, Array.set.fin(T, Array.swap.hi(T, xs, _))) : Array & U32}; %s : {(ANode{xs, Array.set.fin(T, ASet.G(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n)))))}, Pair.snd(Array, U32, Array.size.node(T, Array.set.fin(T, ASet.G(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n))))), _))) == Array.size.node(T, Array.set.fin(T, ASet.G(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n))))), _) : Array & U32}; {==}) : {(ASet.NA(T, ANode{xs, ys}, n, j, v, False{}), ASet.sz(T, ASet.NA(T, ANode{xs, ys}, n, j, v, False{}))) == Array.size(T, ASet.NA(T, ANode{xs, ys}, n, j, v, False{})) : Array & U32}}, {(%c2 : {ASet.sz(T, Array.set.fin(T, Array.swap.hi(T, xs, _))) == ASet.sz(T, ANode{xs, ys}) : U32}; %s : {Pair.snd(Array, U32, Array.size.node(T, Array.set.fin(T, ASet.G(T, ys, U32.shr(n), U32.sub(j, U32.shr(n)), v, U32.is_lt(U32.sub(j, U32.shr(n)), U32.shr(U32.shr(n))))), _)) == Pair.snd(Array, U32, Array.size.node(T, ys, _)) : U32}; {==}) : {ASet.sz(T, ASet.NA(T, ANode{xs, ys}, n, j, v, False{})) == ASet.sz(T, ANode{xs, ys}) : U32}}, {(%c2 : {Pair.snd(Array, T, Array.get.go(T, Array.set.fin(T, Array.swap.hi(T, xs, _)), n, 1995 1996 def ASet.facts(-T: Data, a: Array, +n: U32, +j: U32, +v: T, +z: Bool) -> ASet.Facts(T, a, n, j, v, z): 1997 match a z: 1998 case ALeaf{x} _: 1999 ({==}, {==}, {==}, {==}, {==}, {==}, Unit{}) 2000 case ANode{xs, ys} True{}: Lines 1-2000 of 2340. More: https://jend.leeeet.dev/read.txt?path=bend-ml-tensor-array@0.1.19.0/main.bend&start_line=2001