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.8.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 def dotg(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, acc: F32, ra: Array & F32, rb: Array & F32) -> R: 43 dotr(k, ia, ib, sa, sb, acc, 0.0, 0.0, 0.0, ra, rb) 44 45 # row i of C, columns j.. ; r is the product of the current column (already computed) 46 47 def gcols(mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, +cidx: U32, c: Array, r: R) -> G: 48 match mleft r: 49 case 0n R{a, b, +x}: 50 G{a, b, Array.set(F32, c, cidx, x)} 51 case 1n+p R{a, b, +x}: 52 gcols(p, kn, ia0, (ib0 + bcol : U32), sa, sb, bcol, (cidx + 1 : U32), Array.set(F32, c, cidx, x), dotg(kn, ia0, (ib0 + bcol : U32), sa, sb, 0.0, Array.get(F32, a, ia0), Array.get(F32, b, (ib0 + bcol : U32)))) 53 54 def grows(nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +mu: U32, g: G) -> G: 55 match nleft g: 56 case 0n G{a, b, c}: 57 G{a, b, c} 58 case 1n+p G{a, b, c}: 59 grows(p, mm1, kn, (i + 1 : U32), arow, sa, sb, bcol, mu, gcols(mm1, kn, (i * arow : U32), 0, sa, sb, bcol, (i * mu : U32), c, dotg(kn, (i * arow : U32), 0, sa, sb, 0.0, Array.get(F32, a, (i * arow : U32)), Array.get(F32, b, 0)))) 60 61 # C (n x m) = A · B with k steps; c is the output array (with capacity n*m) 62 63 def gemm(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 64 grows(n, Nat.sub(m, 1n), k, 0, arow, sa, sb, bcol, U32.from_nat(m), G{a, b, c}) 65 66 67 # --------------------------------------------------------------------- 68 # parallel gemm: 2^d blocks of rows of C, each with its own copy of A and B, 69 # returning lists that are written into C at the end 70 # --------------------------------------------------------------------- 71 72 type S3 is Type: 73 S3{a: Array, b: Array, out: List<&2, F32>} 74 75 def lcols(mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, out: List<&2, F32>, r: R) -> S3: 76 match mleft r: 77 case 0n R{a, b, +x}: 78 S3{a, b, x <> out} 79 case 1n+p R{a, b, +x}: 80 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)))) 81 82 def lrows(nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, st: S3) -> S3: 83 match nleft st: 84 case 0n S3{a, b, out}: 85 S3{a, b, out} 86 case 1n+p S3{a, b, out}: 87 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)))) 88 89 def lblock_fin(st: S3) -> List<&2, F32>: 90 match st: 91 case S3{a, b, out}: 92 List.reverse(&2, F32, out) 93 94 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>: 95 lblock_fin(lrows(cnt, mm1, kn, i, arow, sa, sb, bcol, S3{a, b, Nil{}})) 96 97 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>: 98 match d ra rb: 99 case 0n Tuple{a, a2} Tuple{b, b2}: 100 lblock(cnt, mm1, kn, i, arow, sa, sb, bcol, a, b) 101 case 1n++q Tuple{a, a2} Tuple{b, b2}: 102 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)) 103 List.append(&2, F32, x, y) 104 105 def write_l(xs: List<&2, F32>, +idx: U32, c: Array) -> Array: 106 match xs: 107 case Nil{}: 108 c 109 case Con{h, t}: 110 write_l(t, (idx + 1 : U32), Array.set(F32, c, idx, h)) 111 112 # c[idx..] += list 113 114 def gp3(a: Array, b: Array, c: Array, xs: List<&2, F32>) -> G: 115 G{a, b, write_l(xs, 0, c)} 116 117 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: 118 match pa pb: 119 case Tuple{a1, a2} Tuple{b1, b2}: 120 gp3(a1, b1, c, lpar(d, n, mm1, kn, 0, arow, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2))) 121 122 # C (n x m) = A · B, with 2^d blocks of rows in parallel; returns A, B (intact) and C 123 124 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: 125 gp2(d, n, Nat.sub(m, 1n), k, arow, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b)) 126 127 128 # --------------------------------------------------------------------- 129 # parallel gemm by COLUMNS, for n = 1 (matrix · vector): each block computes a 130 # range of columns of C, with its own copy of the operands 131 # --------------------------------------------------------------------- 132 133 def cblock_fin(st: S3) -> List<&2, F32>: 134 match st: 135 case S3{a, b, out}: 136 List.reverse(&2, F32, out) 137 138 # cnt columns starting at j0 (cnt >= 1); row 0 of A is at ia0 139 140 def cblock_go(cnt1: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array) -> List<&2, F32>: 141 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))))) 142 143 def cblock(cnt: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array) -> List<&2, F32>: 144 match cnt: 145 case 0n: 146 Nil{} 147 case 1n+p: 148 cblock_go(p, kn, ia0, j0, sa, sb, bcol, a, b) 149 150 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>: 151 match d ra rb: 152 case 0n Tuple{a, a2} Tuple{b, b2}: 153 cblock(cnt, kn, ia0, j0, sa, sb, bcol, a, b) 154 case 1n++q Tuple{a, a2} Tuple{b, b2}: 155 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)) 156 List.append(&2, F32, x, y) 157 158 def cp2(d: Nat, +m: Nat, +kn: Nat, +sa: U32, +sb: U32, +bcol: U32, c: Array, pa: Array & Array, pb: Array & Array) -> G: 159 match pa pb: 160 case Tuple{a1, a2} Tuple{b1, b2}: 161 gp3(a1, b1, c, cpar(d, m, kn, 0, 0, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2))) 162 163 # C (1 x m) = A (1 x k) · B, with 2^d blocks of columns in parallel 164 165 def gemm_par_cols(d: Nat, +m: Nat, +k: Nat, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 166 cp2(d, m, k, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b)) 167 168 # --------------------------------------------------------------------- 169 # element-wise operations (in place, by index) 170 # --------------------------------------------------------------------- 171 172 # --------------------------------------------------------------------- 173 # gemm: C[i][j] = sum_p A[i*arow + p*sa] * B[p*sb + j*bcol] 174 # (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) 175 # --------------------------------------------------------------------- 176 177 # acc += A[ia + t*sa] * B[ib + t*sb] for k steps; ra, rb already carry the first pair 178 179 type P2 is Type: 180 P2{a: Array, b: Array} 181 182 type AL is Type: 183 AL{c: Array, l: List<&2, F32>} 184 185 # list of m elements starting at idx (reads with Array.get, returning the array) 186 187 def read_l(mleft: Nat, +idx: U32, acc: List<&2, F32>, rc: Array & F32) -> AL: 188 match mleft rc: 189 case 0n Tuple{c, +v}: 190 AL{c, List.reverse(&2, F32, v <> acc)} 191 case 1n+p Tuple{c, +v}: 192 read_l(p, (idx + 1 : U32), v <> acc, Array.get(F32, c, (idx + 1 : U32))) 193 194 # writes the list starting at idx 195 196 def addl(xs: List<&2, F32>, +idx: U32, rc: Array & F32) -> Array: 197 match xs rc: 198 case Nil{} Tuple{c, v}: 199 c 200 case Con{+h, t} Tuple{c, +v}: 201 addl(t, (idx + 1 : U32), Array.get(F32, Array.set(F32, c, idx, (v + h : F32)), (idx + 1 : U32))) 202 203 # adds a bias (a list of m) to each of the n rows of c (n x m) 204 205 def bias_rows(nleft: Nat, +i: U32, +mu: U32, +bl: List<&2, F32>, c: Array) -> Array: 206 match nleft: 207 case 0n: 208 c 209 case 1n+p: 210 bias_rows(p, (i + 1 : U32), mu, bl, addl(bl, (i * mu : U32), Array.get(F32, c, (i * mu : U32)))) 211 212 # in-place relu on the first n elements 213 214 def relu_l(nleft: Nat, +t: U32, rc: Array & F32) -> Array: 215 match nleft rc: 216 case 0n Tuple{c, v}: 217 c 218 case 1n+p Tuple{c, +v}: 219 relu_l(p, (t + 1 : U32), Array.get(F32, Array.set(F32, c, t, F32.max(v, 0.0)), (t + 1 : U32))) 220 221 def pos2(b: Bool) -> F32: 222 match b: 223 case True{}: 224 1.0 225 case False{}: 226 0.0 227 228 def pos(+y: F32) -> F32: 229 pos2(F32.is_gt(y, 0.0)) 230 231 # d[t] *= (h[t] > 0): the relu gradient; returns both arrays 232 233 def mask_l(nleft: Nat, +t: U32, rd: Array & F32, rh: Array & F32) -> P2: 234 match nleft rd rh: 235 case 0n Tuple{d, x} Tuple{h, y}: 236 P2{d, h} 237 case 1n+p Tuple{d, +x} Tuple{h, +y}: 238 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))) 239 240 # w[t] -= lr * dw[t] 241 242 def maxl2(xs: List<&2, F32>, cur: F32) -> F32: 243 match xs: 244 case Nil{}: 245 cur 246 case Con{+h, t}: 247 maxl2(t, F32.max(cur, h)) 248 249 def maxl(xs: List<&2, F32>) -> F32: 250 match xs: 251 case Nil{}: 252 0.0 253 case Con{+h, t}: 254 maxl2(t, h) 255 256 def expl(xs: List<&2, F32>, +m: F32) -> List<&2, F32>: 257 match xs: 258 case Nil{}: 259 Nil{} 260 case Con{+h, t}: 261 F32.exp((h - m : F32)) <> expl(t, m) 262 263 def suml(xs: List<&2, F32>, acc: F32) -> F32: 264 match xs: 265 case Nil{}: 266 acc 267 case Con{+h, t}: 268 suml(t, (acc + h : F32)) 269 270 def divl(xs: List<&2, F32>, +s: F32) -> List<&2, F32>: 271 match xs: 272 case Nil{}: 273 Nil{} 274 case Con{+h, t}: 275 (h / s : F32) <> divl(t, s) 276 277 def softmax2(+es: List<&2, F32>) -> List<&2, F32>: 278 divl(es, suml(es, 0.0)) 279 280 def softmax(+xs: List<&2, F32>) -> List<&2, F32>: 281 softmax2(expl(xs, maxl(xs))) 282 283 # p[label] 284 285 def sgd_l(nleft: Nat, +t: U32, +lr: F32, rw: Array & F32, rd: Array & F32) -> P2: 286 match nleft rw rd: 287 case 0n Tuple{w, x} Tuple{d, y}: 288 P2{w, d} 289 case 1n+p Tuple{w, +x} Tuple{d, +y}: 290 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))) 291 292 # --------------------------------------------------------------------- 293 # softmax + cross-entropy per row (lists of 10 elements: negligible cost) 294 # --------------------------------------------------------------------- 295 296 def pick(ps: List<&2, F32>, n: Nat) -> F32: 297 match ps n: 298 case Con{+h, t} 0n: 299 h 300 case Con{h, t} 1n+q: 301 pick(t, q) 302 case Nil{} _: 303 0.0 304 305 # (p - onehot(label)) / n 306 307 def hit(b: Bool) -> F32: 308 match b: 309 case True{}: 310 1.0 311 case False{}: 312 0.0 313 314 def grad_l(ps: List<&2, F32>, +label: Nat, +j: Nat, +nf: F32) -> List<&2, F32>: 315 match ps: 316 case Nil{}: 317 Nil{} 318 case Con{+h, t}: 319 (((h - hit(Nat.is_eq(label, j)) : F32) / nf) : F32) <> grad_l(t, label, 1n+j, nf) 320 321 type L2 is Type: 322 L2{loss: F32, c: Array} 323 324 def ce_row2(+ps: List<&2, F32>, +label: Nat, +nf: F32, +idx: U32, loss: F32, c: Array) -> L2: 325 L2{(loss - F32.log(F32.max(pick(ps, label), 0.000000000001)) : F32), write_l(grad_l(ps, label, 0n, nf), idx, c)} 326 327 def ce_row(al: AL, +label: Nat, +nf: F32, +idx: U32, loss: F32) -> L2: 328 match al: 329 case AL{c, l}: 330 ce_row2(softmax(l), label, nf, idx, loss, c) 331 332 # 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 333 334 type Mat<-r: Nat, -c: Nat> is Type: 335 Mat{d: Array} 336 337 def ce_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +nf: F32, +i: U32, r: L2) -> L2: 338 match labels r: 339 case Nil{} L2{loss, c}: 340 L2{loss, c} 341 case Con{+y, t} L2{loss, c}: 342 ce_rows(t, mm1, mu, nf, (i + 1 : U32), ce_row(read_l(mm1, (i * mu : U32), Nil{}, Array.get(F32, c, (i * mu : U32))), y, nf, (i * mu : U32), loss)) 343 344 # ---- helper: hits per row ---- 345 346 type CR is Type: 347 CR{c: Array, n: Nat} 348 349 def gt_head(xs: List<&2, F32>, +v: F32) -> Bool: 350 match xs: 351 case Nil{}: 352 False{} 353 case Con{+h, t}: 354 F32.is_gt(h, v) 355 356 def argmax2(xs: List<&2, F32>, g: Bool, +i: Nat, +bi: Nat, +bv: F32) -> Nat: 357 match xs g: 358 case Nil{} _: 359 bi 360 case Con{+h, +t} True{}: 361 argmax2(t, gt_head(t, h), 1n+i, i, h) 362 case Con{h, +t} False{}: 363 argmax2(t, gt_head(t, bv), 1n+i, bi, bv) 364 365 def argmax(xs: List<&2, F32>) -> Nat: 366 match xs: 367 case Nil{}: 368 0n 369 case Con{+h, +t}: 370 argmax2(t, gt_head(t, h), 1n, 0n, h) 371 372 def one(b: Bool) -> Nat: 373 match b: 374 case True{}: 375 1n 376 case False{}: 377 0n 378 379 def count_rows3(+y: Nat, n: Nat, al: AL) -> CR: 380 match al: 381 case AL{c, l}: 382 CR{c, Nat.add(n, one(Nat.is_eq(argmax(l), y)))} 383 384 def count_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +i: U32, r: CR) -> CR: 385 match labels r: 386 case Nil{} CR{c, n}: 387 CR{c, n} 388 case Con{+y, t} CR{c, n}: 389 count_rows(t, mm1, mu, (i + 1 : U32), count_rows3(y, n, read_l(mm1, (i * mu : U32), Nil{}, Array.get(F32, c, (i * mu : U32))))) 390 391 392 # ===================================================================== 393 # Typed API: Mat (the dimensions enter the type) 394 # ===================================================================== 395 396 # ===================================================================== 397 # Array capacity: a PROVED depth for the number of slots 398 # ===================================================================== 399 # 400 # cap_depth(n) is the d such that an Array with 2^d slots holds n elements. It is computed by 401 # structural recursion on n (not with U32.log2) so that the law cap_ok below can prove 402 # `n <= 2^cap_depth(n)`. What stays trusted is only that Array.new(T, d, v) allocates 2^d slots. 403 404 # a type chosen by a Bool: rewriting through it refutes True == False 405 def BD(b: Bool, t: Type, f: Type) -> Type: 406 match b: 407 case True{}: 408 t 409 case False{}: 410 f 411 412 # a <= b, by direct recursion 413 def leb(a: Nat, b: Nat) -> Bool: 414 match a b: 415 case 0n _: 416 True{} 417 case 1n+p 0n: 418 False{} 419 case 1n+p 1n+q: 420 leb(p, q) 421 422 def pow2(d: Nat) -> Nat: 423 match d: 424 case 0n: 425 1n 426 case 1n++q: 427 Nat.add(pow2(q), pow2(q)) 428 429 # smallest power of two >= n, by structural recursion on n: returns (depth, 2^depth) 430 def cap_pick(ok: Bool, d: Nat, +p: Nat) -> Nat & Nat: 431 match ok: 432 case True{}: 433 (d, p) 434 case False{}: 435 (1n+d, Nat.add(p, p)) 436 437 def cap_step(r: Nat & Nat, +n: Nat) -> Nat & Nat: 438 match r: 439 case Tuple{+d, +p}: 440 cap_pick(leb(n, p), d, p) 441 442 def cap_pair(n: Nat) -> Nat & Nat: 443 match n: 444 case 0n: 445 (0n, 1n) 446 case 1n++m: 447 cap_step(cap_pair(m), 1n+m) 448 449 def fst_nat(r: Nat & Nat) -> Nat: 450 match r: 451 case Tuple{d, p}: 452 d 453 454 def snd_nat(r: Nat & Nat) -> Nat: 455 match r: 456 case Tuple{d, p}: 457 p 458 459 # the depth d such that an Array with 2^d slots holds n elements 460 def cap_depth(n: Nat) -> Nat: 461 fst_nat(cap_pair(n)) 462 463 # ---- lemmas ---- 464 465 # a <= b -> a <= 1 + b 466 def leb_succ(a: Nat, b: Nat, e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, 1n+b) : Bool}: 467 match a b: 468 case 0n _: 469 {==} 470 case 1n+p 0n: 471 %e : BD(_, Unit, {True{} == leb(1n+p, 1n) : Bool}) 472 Unit{} 473 case 1n+p 1n+q: 474 leb_succ(p, q, e) 475 476 # a <= b -> a <= c + b 477 def leb_up_l(+a: Nat, +b: Nat, c: Nat, +e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, Nat.add(c, b)) : Bool}: 478 match c: 479 case 0n: 480 e 481 case 1n++c2: 482 leb_succ(a, Nat.add(c2, b), leb_up_l(a, b, c2, e)) 483 484 # a <= b -> a <= b + c 485 def leb_up(a: Nat, b: Nat, c: Nat, e: {True{} == leb(a, b) : Bool}) -> {True{} == leb(a, Nat.add(b, c)) : Bool}: 486 match a b: 487 case 0n _: 488 {==} 489 case 1n+p 0n: 490 %e : BD(_, Unit, {True{} == leb(1n+p, Nat.add(0n, c)) : Bool}) 491 Unit{} 492 case 1n+p 1n+q: 493 leb_up(p, q, c, e) 494 495 # a <= b and c <= d -> a + c <= b + d 496 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}: 497 match a b: 498 case 0n _: 499 leb_up_l(c, d, b, e2) 500 case 1n+p 0n: 501 %e1 : BD(_, Unit, {True{} == leb(Nat.add(1n+p, c), Nat.add(0n, d)) : Bool}) 502 Unit{} 503 case 1n+p 1n+q: 504 leb_add(p, q, c, d, e1, e2) 505 506 # 1 <= 2^d 507 def pow2_pos(d: Nat) -> {True{} == leb(1n, pow2(d)) : Bool}: 508 match d: 509 case 0n: 510 {==} 511 case 1n++q: 512 leb_up(1n, pow2(q), pow2(q), pow2_pos(q)) 513 514 # 1 <= p when p = 2^d 515 def pos_p(d: Nat, p: Nat, i1: {p == pow2(d) : Nat}) -> {True{} == leb(1n, p) : Bool}: 516 %Equal.sym(Nat, p, pow2(d), i1) : {True{} == leb(1n, _) : Bool} 517 pow2_pos(d) 518 519 # p + p == 2^(1+d) when p == 2^d 520 def step_eq(d: Nat, p: Nat, i1: {p == pow2(d) : Nat}) -> {Nat.add(p, p) == pow2(1n+d) : Nat}: 521 %i1 : {Nat.add(p, p) == Nat.add(_, _) : Nat} 522 {==} 523 524 # one step: given the invariant for m, the invariant for 1+m holds whichever way the comparison goes 525 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}: 526 match ok: 527 case True{}: 528 (i1, e) 529 case False{}: 530 (step_eq(d, p, i1), leb_add(1n, p, m, p, pos_p(d, p, i1), i2)) 531 532 533 # the invariant for cap_step on a pair, opened 534 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}: 535 match r: 536 case Tuple{+d, +p}: 537 step_ok(leb(1n+m, p), d, p, m, {==}, i1, i2) 538 539 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}: 540 match ih: 541 case Tuple{i1, i2}: 542 step_pair_ok(m, cap_pair(m), i1, i2) 543 544 # the invariant of cap_pair: the second component is 2^(first) and is >= n 545 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}: 546 match n: 547 case 0n: 548 ({==}, {==}) 549 case 1n++m: 550 cap_succ(m, cap_inv(m)) 551 552 def cap_final(n: Nat, pr: {snd_nat(cap_pair(n)) == pow2(fst_nat(cap_pair(n))) : Nat} & {True{} == leb(n, snd_nat(cap_pair(n))) : Bool}) -> {True{} == leb(n, pow2(cap_depth(n))) : Bool}: 553 match pr: 554 case Tuple{i1, i2}: 555 %i1 : {True{} == leb(n, _) : Bool} 556 i2 557 558 # LAW (cap_ok): an Array with 2^cap_depth(n) slots always has room for n elements. 559 law cap_ok: 560 for +n: Nat 561 {True{} == leb(n, pow2(cap_depth(n))) : Bool} 562 563 def cap_ok(n): 564 cap_final(n, cap_inv(n)) 565 566 # r x c of zeros, with the exact capacity (smallest power of 2) 567 568 def Mat.zeros(+r: Nat, +c: Nat) -> Mat: 569 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0)} 570 571 def Mat.fill(+r: Nat, +c: Nat, +v: F32) -> Mat: 572 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), v)} 573 574 def fill_list(xs: List<&2, F32>, +i: U32, a: Array) -> Array: 575 match xs: 576 case Nil{}: 577 a 578 case Con{h, t}: 579 fill_list(t, (i + 1 : U32), Array.set(F32, a, i, h)) 580 581 def len_is(xs: List<&2, F32>, n: Nat) -> Bool: 582 match xs n: 583 case Nil{} 0n: 584 True{} 585 case Con{h, t} 1n+p: 586 len_is(t, p) 587 case _ _: 588 False{} 589 590 def Mat.of_list2(-r: Nat, -c: Nat, +rr: Nat, +cc: Nat, ok: Bool, xs: List<&2, F32>) -> Maybe<&1, Mat>: 591 match ok: 592 case True{}: 593 Some{Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(rr, cc)), 0.0))}} 594 case False{}: 595 None{} 596 597 # Mat from r*c numbers in row order; None if the size does not match 598 599 def Mat.of_list(+r: Nat, +c: Nat, +xs: List<&2, F32>) -> Maybe<&1, Mat>: 600 Mat.of_list2(r, c, r, c, len_is(xs, Nat.mul(r, c)), xs) 601 602 # reads the first n numbers, returning the matrix 603 604 # Mat without checking the list size (a short list leaves the rest at zero, a long one is cut 605 # by the capacity): only for when the size is known by construction. Prefer Mat.of_list. 606 607 def Mat.from_list(+r: Nat, +c: Nat, xs: List<&2, F32>) -> Mat: 608 Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0))} 609 610 # writes the numbers of xs into m starting at flat index i (row-major); numbers past the 611 # capacity are dropped. For loading a large matrix block by block without building one big list. 612 def Mat.fill_at(-r: Nat, -c: Nat, +i: U32, xs: List<&2, F32>, m: Mat) -> Mat: 613 match m: 614 case Mat{d}: 615 Mat{fill_list(xs, i, d)} 616 617 type ML<-r: Nat, -c: Nat> is Type: 618 ML{m: Mat, l: List<&2, F32>} 619 620 def Mat.to_list2(-r: Nat, -c: Nat, al: AL) -> ML: 621 match al: 622 case AL{d, l}: 623 ML{Mat{d}, l} 624 625 # two independent copies 626 627 def Mat.to_list(+r: Nat, +c: Nat, m: Mat) -> ML: 628 match m: 629 case Mat{d}: 630 Mat.to_list2(r, c, read_l(Nat.sub(Nat.mul(r, c), 1n), 0, Nil{}, Array.get(F32, d, 0))) 631 632 type MM2<-r: Nat, -c: Nat> is Type: 633 MM2{a: Mat, b: Mat} 634 635 def Mat.clone2(-r: Nat, -c: Nat, p: Array & Array) -> MM2: 636 match p: 637 case Tuple{x, y}: 638 MM2{Mat{x}, Mat{y}} 639 640 # ---- products: C (n x m) = A · B, with 2^par blocks of rows in parallel ---- 641 642 # C = A(n x k) · B(k x m) 643 644 def Mat.clone(-r: Nat, -c: Nat, m: Mat) -> MM2: 645 match m: 646 case Mat{d}: 647 Mat.clone2(r, c, Array.clone(F32, d)) 648 649 type MMul<-n: Nat, -k: Nat, -m: Nat> is Type: 650 MMul{a: Mat, b: Mat, c: Mat} 651 652 # C = A(n x k) · Bᵀ, with B stored m x k 653 654 type MMulNT<-n: Nat, -k: Nat, -m: Nat> is Type: 655 MMulNT{a: Mat, b: Mat, c: Mat} 656 657 # C = Aᵀ · B, with A stored k x n 658 659 type MMulTN<-n: Nat, -k: Nat, -m: Nat> is Type: 660 MMulTN{a: Mat, b: Mat, c: Mat} 661 662 # --------------------------------------------------------------------- 663 # 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 664 # 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 665 # 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, 666 # in order, and the four sums are combined as dotr combines them after k steps (its rotation depends on k mod 4). The 667 # padding adds +0 to a sum that starts at +0, which changes nothing, so the results are the same bits as the scalar 668 # loop's. Products with one row (matrix · vector) keep the scalar loop: packing B would cost as much as the product. 669 # --------------------------------------------------------------------- 670 671 type T8 is Data: 672 T8{v0: F32, v1: F32, v2: F32, v3: F32, v4: F32, v5: F32, v6: F32, v7: F32} 673 674 def t8_zero() -> T8: 675 T8{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0} 676 677 def pick0(b: Bool, +v: F32) -> F32: 678 match b: 679 case True{}: 680 v 681 case False{}: 682 0.0 683 684 type PK is Type: 685 PK{a: Array, t: Array} 686 687 # 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 688 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: 689 match r: 690 case Tuple{a, +v}: 691 PK{a, Array.set(T8, t, o, T8{v0, v1, v2, v3, v4, v5, v6, pick0(U32.is_lt((q0 + 7 : U32), lim), v)})} 692 693 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: 694 match r: 695 case Tuple{a, +v}: 696 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))) 697 698 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: 699 match r: 700 case Tuple{a, +v}: 701 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))) 702 703 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: 704 match r: 705 case Tuple{a, +v}: 706 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))) 707 708 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: 709 match r: 710 case Tuple{a, +v}: 711 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))) 712 713 def pk_g2(+v0: F32, +v1: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array, r: Array & F32) -> PK: 714 match r: 715 case Tuple{a, +v}: 716 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))) 717 718 def pk_g1(+v0: F32, +q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array, r: Array & F32) -> PK: 719 match r: 720 case Tuple{a, +v}: 721 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))) 722 723 def pk_g0(+q0: U32, +lim: U32, +idx: U32, +st: U32, +o: U32, t: Array, r: Array & F32) -> PK: 724 match r: 725 case Tuple{a, +v}: 726 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))) 727 728 # the cells of one line, cell o.. of t; q0 counts the line's numbers before the cell 729 def pk_line(tl: Nat, +q0: U32, +lim: U32, +base: U32, +st: U32, +o: U32, r: PK) -> PK: 730 match tl r: 731 case 0n r: 732 r 733 case 1n+q PK{a, t}: 734 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))) 735 736 # lines count lines; line i starts at i * lstep (in a), its cells at i * tku (in t) 737 def pk_lines(cnt: Nat, +tk: Nat, +tku: U32, +ku: U32, +i: U32, +lstep: U32, +st: U32, r: PK) -> PK: 738 match cnt: 739 case 0n: 740 r 741 case 1n+p: 742 pk_lines(p, tk, tku, ku, (i + 1 : U32), lstep, st, pk_line(tk, 0, ku, (i * lstep : U32), st, (i * tku : U32), r)) 743 744 def pk_pack(+cnt: Nat, +tk: Nat, +k: Nat, +lstep: U32, +st: U32, a: Array) -> PK: 745 pk_lines(cnt, tk, U32.from_nat(tk), U32.from_nat(k), 0, lstep, st, PK{a, Array.new(T8, cap_depth(Nat.mul(cnt, tk)), t8_zero())}) 746 747 type TD is Type: 748 TD{a: Array, b: Array, s0: F32, s1: F32, s2: F32, s3: F32} 749 750 # product q of a cell goes to sum q mod 4, the first four before the last four, as dotr adds them 751 def t8_step(s0: F32, s1: F32, s2: F32, s3: F32, ra: Array & T8, rb: Array & T8) -> TD: 752 match ra rb: 753 case Tuple{a2, T8{x0, x1, x2, x3, x4, x5, x6, x7}} Tuple{b2, T8{y0, y1, y2, y3, y4, y5, y6, y7}}: 754 TD{a2, b2, ((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)} 755 756 def t8_dot(tk: Nat, +ia: U32, +ib: U32, d: TD) -> TD: 757 match tk d: 758 case 0n d: 759 d 760 case 1n+p TD{a, b, s0, s1, s2, s3}: 761 t8_dot(p, (ia + 1 : U32), (ib + 1 : U32), t8_step(s0, s1, s2, s3, Array.get(T8, a, ia), Array.get(T8, b, ib))) 762 763 # dotr's final sum after k products: the sums sit rotated by k mod 4 when its loop ends 764 def t8_comb(+rot: Nat, +s0: F32, +s1: F32, +s2: F32, +s3: F32) -> F32: 765 match rot: 766 case 0n: 767 ((s0 + s1 : F32) + (s2 + s3 : F32) : F32) 768 case 1n: 769 ((s1 + s2 : F32) + (s3 + s0 : F32) : F32) 770 case 2n: 771 ((s2 + s3 : F32) + (s0 + s1 : F32) : F32) 772 case _: 773 ((s3 + s0 : F32) + (s1 + s2 : F32) : F32) 774 775 type TR is Type: 776 TR{a: Array, b: Array, x: F32} 777 778 def t8_fin(+rot: Nat, d: TD) -> TR: 779 match d: 780 case TD{a, b, +s0, +s1, +s2, +s3}: 781 TR{a, b, t8_comb(rot, s0, s1, s2, s3)} 782 783 def t8_cell(+tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, a: Array, b: Array) -> TR: 784 t8_fin(rot, t8_dot(tk, (i * tku : U32), (j * tku : U32), TD{a, b, 0.0, 0.0, 0.0, 0.0})) 785 786 type TG is Type: 787 TG{a: Array, b: Array, c: Array} 788 789 def t8_set(+idx: U32, c: Array, r: TR) -> TG: 790 match r: 791 case TR{a, b, x}: 792 TG{a, b, Array.set(F32, c, idx, x)} 793 794 def t8_cols(mleft: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, +cidx: U32, g: TG) -> TG: 795 match mleft g: 796 case 0n g: 797 g 798 case 1n+p TG{a, b, c}: 799 t8_cols(p, tk, tku, rot, i, (j + 1 : U32), (cidx + 1 : U32), t8_set(cidx, c, t8_cell(tk, tku, rot, i, j, a, b))) 800 801 def t8_rows(nleft: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +mu: U32, g: TG) -> TG: 802 match nleft: 803 case 0n: 804 g 805 case 1n+p: 806 t8_rows(p, m, tk, tku, rot, (i + 1 : U32), mu, t8_cols(m, tk, tku, rot, i, 0, (i * mu : U32), g)) 807 808 # rows i0 .. i0+cnt-1 of C as a list, in order (the parallel version's leaves) 809 type TL is Type: 810 TL{a: Array, b: Array, out: List<&2, F32>} 811 812 # cells j.. of row i, m - j of them; r is cell j, already computed 813 def t8_lcols(mleft: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j: U32, out: List<&2, F32>, r: TR) -> TL: 814 match mleft r: 815 case 0n TR{a, b, +x}: 816 TL{a, b, x <> out} 817 case 1n+p TR{a, b, +x}: 818 t8_lcols(p, tk, tku, rot, i, (j + 1 : U32), x <> out, t8_cell(tk, tku, rot, i, (j + 1 : U32), a, b)) 819 820 def t8_lrows(nleft: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, l: TL) -> TL: 821 match nleft l: 822 case 0n l: 823 l 824 case 1n+p TL{a, b, out}: 825 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))) 826 827 def t8_lfin(l: TL) -> List<&2, F32>: 828 match l: 829 case TL{a, b, out}: 830 List.reverse(&2, F32, out) 831 832 def t8_lpar(d: Nat, +cnt: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, ra: Array & Array, rb: Array & Array) -> List<&2, F32>: 833 match d ra rb: 834 case 0n Tuple{a, a2} Tuple{b, b2}: 835 t8_lfin(t8_lrows(cnt, m, tk, tku, rot, i, TL{a, b, Nil{}})) 836 case 1n++q Tuple{a, a2} Tuple{b, b2}: 837 x y = t8_lpar(q, Nat.div(cnt, 2n), m, tk, tku, rot, i, Array.clone(T8, a), Array.clone(T8, b)) t8_lpar(q, Nat.sub(cnt, Nat.div(cnt, 2n)), m, tk, tku, rot, (i + U32.from_nat(Nat.div(cnt, 2n)) : U32), Array.clone(T8, a2), Array.clone(T8, b2)) 838 List.append(&2, F32, x, y) 839 840 type PKK is Type: 841 PKK{a: Array, b: Array, ta: Array, tb: Array} 842 843 def t8_pk2(pa: PK, pb: PK) -> PKK: 844 match pa pb: 845 case PK{a, ta} PK{b, tb}: 846 PKK{a, b, ta, tb} 847 848 def t8_g(a: Array, b: Array, g: TG) -> G: 849 match g: 850 case TG{ta, tb, c}: 851 G{a, b, c} 852 853 def t8_out_seq(+n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array, p: PKK) -> G: 854 match p: 855 case PKK{a, b, ta, tb}: 856 t8_g(a, b, t8_rows(n, m, tk, tku, rot, 0, U32.from_nat(m), TG{ta, tb, c})) 857 858 def t8_out_par(d: Nat, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array, p: PKK) -> G: 859 match p: 860 case PKK{a, b, ta, tb}: 861 G{a, b, write_l(t8_lpar(d, n, m, tk, tku, rot, 0, Array.clone(T8, ta), Array.clone(T8, tb)), 0, c)} 862 863 def gemm_t8b(d: Nat, +n: Nat, +m: Nat, +k: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array, p: PKK) -> G: 864 match d: 865 case 0n: 866 t8_out_seq(n, m, tk, tku, rot, c, p) 867 case 1n+q: 868 t8_out_par(1n+q, n, m, tk, tku, rot, c, p) 869 870 # C (n x m) = A · B in cells of 8, sequential (d = 0) or with 2^d blocks of rows in parallel; A and B come back intact 871 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: 872 gemm_t8b(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), c, t8_pk2(pk_pack(n, Nat.div(Nat.add(k, 7n), 8n), k, arow, sa, a), pk_pack(m, Nat.div(Nat.add(k, 7n), 8n), k, bcol, sb, b))) 873 874 # one row (matrix · vector): the scalar loop, by columns in parallel when d > 0; two or more rows: cells of 8 875 def gemm_any3(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: 876 match one d: 877 case True{} 0n: 878 gemm(n, m, k, arow, sa, sb, bcol, a, b, c) 879 case True{} 1n+q: 880 gemm_par_cols(1n+q, m, k, sa, sb, bcol, a, b, c) 881 case False{} d: 882 gemm_t8(d, n, m, k, arow, sa, sb, bcol, a, b, c) 883 884 def gemm_any2(d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 885 gemm_any3(Nat.is_le(n, 1n), d, n, m, k, arow, sa, sb, bcol, a, b, c) 886 887 def gemm_any(d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 888 gemm_any2(d, n, m, k, arow, sa, sb, bcol, a, b, c) 889 890 def mmul_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMul: 891 match g: 892 case G{a, b, c}: 893 MMul{Mat{a}, Mat{b}, Mat{c}} 894 895 def mmul_nt_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulNT: 896 match g: 897 case G{a, b, c}: 898 MMulNT{Mat{a}, Mat{b}, Mat{c}} 899 900 def mmul_tn_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulTN: 901 match g: 902 case G{a, b, c}: 903 MMulTN{Mat{a}, Mat{b}, Mat{c}} 904 905 def mm_a(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array, b: Array) -> MMul: 906 mmul_fin(n, k, m, gemm_any(d, n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0))) 907 908 # C = A · B (n x k) · (k x m); par = log2 of the number of parallel blocks (0 = sequential) 909 910 def Mat.matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat, b: Mat) -> MMul: 911 match a b: 912 case Mat{x} Mat{y}: 913 mm_a(n, k, m, par, x, y) 914 915 def mm_nt(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array, b: Array) -> MMulNT: 916 mmul_nt_fin(n, k, m, gemm_any(d, n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0))) 917 918 # C = A · Bᵀ (n x k) · (m x k)ᵀ: for weights stored with one row per output 919 920 def Mat.matmul_nt(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat, b: Mat) -> MMulNT: 921 match a b: 922 case Mat{x} Mat{y}: 923 mm_nt(n, k, m, par, x, y) 924 925 def mm_tn(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array, b: Array) -> MMulTN: 926 mmul_tn_fin(n, k, m, gemm_any(d, n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0))) 927 928 # C = Aᵀ · B (k x n)ᵀ · (k x m): weight gradient, X^T · dY 929 930 def Mat.matmul_tn(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat, b: Mat) -> MMulTN: 931 match a b: 932 case Mat{x} Mat{y}: 933 mm_tn(n, k, m, par, x, y) 934 935 # ---- element-wise and training operations ---- 936 937 # adds the bias (one row, m numbers) to each row of C (n x m) 938 939 type MBias<-n: Nat, -m: Nat> is Type: 940 MBias{c: Mat, b: Mat<1n, m>} 941 942 def add_row2(-n: Nat, -m: Nat, +nn: Nat, +mm: Nat, cd: Array, al: AL) -> MBias: 943 match al: 944 case AL{bd, bl}: 945 MBias{Mat{bias_rows(nn, 0, U32.from_nat(mm), bl, cd)}, Mat{bd}} 946 947 def Mat.add_row(+n: Nat, +m: Nat, c: Mat, b: Mat<1n, m>) -> MBias: 948 match c b: 949 case Mat{cd} Mat{bd}: 950 add_row2(n, m, n, m, cd, read_l(Nat.sub(m, 1n), 0, Nil{}, Array.get(F32, bd, 0))) 951 952 # pair of matrices of the same shape (the result of an operation that reads two operands) 953 954 type MPair<-r: Nat, -c: Nat> is Type: 955 MPair{a: Mat, b: Mat} 956 957 def mpair_fin(-r: Nat, -c: Nat, p: P2) -> MPair: 958 match p: 959 case P2{a, b}: 960 MPair{Mat{a}, Mat{b}} 961 962 def relu_fin(-r: Nat, -c: Nat, d: Array) -> Mat: 963 Mat{d} 964 965 def Mat.relu(+r: Nat, +c: Nat, a: Mat) -> Mat: 966 match a: 967 case Mat{d}: 968 relu_fin(r, c, relu_l(Nat.mul(r, c), 0, Array.get(F32, d, 0))) 969 970 # (dy * (h > 0), h): the gradient that passes through the relu 971 972 def Mat.relu_bwd(+r: Nat, +c: Nat, dy: Mat, h: Mat) -> MPair: 973 match dy h: 974 case Mat{d} Mat{x}: 975 mpair_fin(r, c, mask_l(Nat.mul(r, c), 0, Array.get(F32, d, 0), Array.get(F32, x, 0))) 976 977 # (w - lr * dw, dw): one SGD step 978 979 def Mat.sgd(+r: Nat, +c: Nat, +lr: F32, w: Mat, dw: Mat) -> MPair: 980 match w dw: 981 case Mat{x} Mat{y}: 982 mpair_fin(r, c, sgd_l(Nat.mul(r, c), 0, lr, Array.get(F32, x, 0), Array.get(F32, y, 0))) 983 984 # (a + b, b): element-wise sum (a - (-1)*b is exactly a + b) 985 986 def Mat.add(+r: Nat, +c: Nat, a: Mat, b: Mat) -> MPair: 987 match a b: 988 case Mat{x} Mat{y}: 989 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))) 990 991 # sum of each column (n x m -> 1 x m), returning the matrix 992 993 type MColSum<-n: Nat, -m: Nat> is Type: 994 MColSum{a: Mat, s: Mat<1n, m>} 995 996 def colsum_fin(-n: Nat, -m: Nat, g: G) -> MColSum: 997 match g: 998 case G{ones, a, s}: 999 MColSum{Mat{a}, Mat{s}} 1000 1001 def Mat.col_sums(+n: Nat, +m: Nat, a: Mat) -> MColSum: 1002 match a: 1003 case Mat{d}: 1004 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))) 1005 1006 # mean cross-entropy and its gradient on the logits: (softmax - one-hot) / n. 1007 # labels has one entry per row (n); the gradient replaces the logits. 1008 1009 type MCE<-n: Nat, -c: Nat> is Type: 1010 MCE{loss: F32, grad: Mat} 1011 1012 def ce_fin(-n: Nat, -c: Nat, +nf: F32, r: L2) -> MCE: 1013 match r: 1014 case L2{loss, d}: 1015 MCE{(loss / nf : F32), Mat{d}} 1016 1017 def Mat.softmax_ce(+n: Nat, +c: Nat, z: Mat, labels: List<&2, Nat>) -> MCE: 1018 match z: 1019 case Mat{d}: 1020 ce_fin(n, c, F32.from_nat(n), ce_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), F32.from_nat(n), 0, L2{0.0, d})) 1021 1022 # argmax hits per row against labels 1023 1024 type MHits<-n: Nat, -c: Nat> is Type: 1025 MHits{z: Mat, hits: Nat} 1026 1027 def hits_fin(-n: Nat, -c: Nat, r: CR) -> MHits: 1028 match r: 1029 case CR{d, k}: 1030 MHits{Mat{d}, k} 1031 1032 def Mat.count_correct(+n: Nat, +c: Nat, z: Mat, labels: List<&2, Nat>) -> MHits: 1033 match z: 1034 case Mat{d}: 1035 hits_fin(n, c, count_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), 0, CR{d, 0n})) 1036 1037 # a row as a list (returning the matrix) and the other way around 1038 1039 type MRow<-r: Nat, -c: Nat> is Type: 1040 MRow{m: Mat, row: List<&2, F32>} 1041 1042 def row_fin(-r: Nat, -c: Nat, al: AL) -> MRow: 1043 match al: 1044 case AL{d, l}: 1045 MRow{Mat{d}, l} 1046 1047 def Mat.read_row(+r: Nat, +c: Nat, m: Mat, +i: U32) -> MRow: 1048 match m: 1049 case Mat{d}: 1050 row_fin(r, c, read_l(Nat.sub(c, 1n), (i * U32.from_nat(c) : U32), Nil{}, Array.get(F32, d, (i * U32.from_nat(c) : U32)))) 1051 1052 def Mat.write_row(-r: Nat, +c: Nat, m: Mat, +i: U32, xs: List<&2, F32>) -> Mat: 1053 match m: 1054 case Mat{d}: 1055 Mat{write_l(xs, (i * U32.from_nat(c) : U32), d)} 1056 1057 # ---- variants that check the number of labels ---- 1058 1059 # does the list of labels have exactly n entries (one per row)? 1060 def labels_len_is(xs: List<&2, Nat>, n: Nat) -> Bool: 1061 match xs n: 1062 case Nil{} 0n: 1063 True{} 1064 case Con{h, t} 1n+p: 1065 labels_len_is(t, p) 1066 case _ _: 1067 False{} 1068 1069 def ce_checked(ok: Bool, +n: Nat, +c: Nat, z: Mat, labels: List<&2, Nat>) -> Maybe<&1, MCE>: 1070 match ok: 1071 case True{}: 1072 Some{Mat.softmax_ce(n, c, z, labels)} 1073 case False{}: 1074 None{} 1075 1076 # like Mat.softmax_ce, but returns None unless labels has exactly n entries 1077 def Mat.softmax_ce_checked(+n: Nat, +c: Nat, z: Mat, +labels: List<&2, Nat>) -> Maybe<&1, MCE>: 1078 ce_checked(labels_len_is(labels, n), n, c, z, labels) 1079 1080 def hits_checked(ok: Bool, +n: Nat, +c: Nat, z: Mat, labels: List<&2, Nat>) -> Maybe<&1, MHits>: 1081 match ok: 1082 case True{}: 1083 Some{Mat.count_correct(n, c, z, labels)} 1084 case False{}: 1085 None{} 1086 1087 # like Mat.count_correct, but returns None unless labels has exactly n entries 1088 def Mat.count_correct_checked(+n: Nat, +c: Nat, z: Mat, +labels: List<&2, Nat>) -> Maybe<&1, MHits>: 1089 hits_checked(labels_len_is(labels, n), n, c, z, labels) 1090 1091 # --------------------------------------------------------------------- 1092 # Bands: a matrix cut into bands of rows, for matrix · vector in parallel 1093 # without copying the weights. 1094 # 1095 # Mat.matmul_nt with n = 1 can only run in parallel by cloning the whole weight 1096 # matrix for every task, and in matrix · vector each weight is read once, so the 1097 # copies cost more than the products (NOTES.md, exp. 9). Bands keeps each band in 1098 # its own Array: a task takes its band and nothing is copied except the input vector. 1099 # 1100 # A node with r rows has children with half(r) and r - half(r) rows. The split is 1101 # part of the type, so the row count of every band follows from r, and the law 1102 # half_cover says that the two halves always cover the r rows exactly. 1103 # --------------------------------------------------------------------- 1104 1105 # n / 2 rounded down. A tail loop with an accumulator: a non-tail recursion here would put a 1106 # continuation on every node of the band tree and keep the bands from running in parallel. 1107 def half.go(n: Nat, acc: Nat) -> Nat: 1108 match n: 1109 case 0n: 1110 acc 1111 case 1n+p: 1112 match p: 1113 case 0n: 1114 acc 1115 case 1n+q: 1116 half.go(q, 1n+acc) 1117 1118 def half(n: Nat) -> Nat: 1119 half.go(n, 0n) 1120 1121 # a <= a 1122 def leb_refl(+a: Nat) -> {True{} == leb(a, a) : Bool}: 1123 match a: 1124 case 0n: 1125 {==} 1126 case 1n+p: 1127 leb_refl(p) 1128 1129 # 1 + (a + b) == a + (1 + b) 1130 def add_succ_r(+a: Nat, -b: Nat) -> {1n+Nat.add(a, b) == Nat.add(a, 1n+b) : Nat}: 1131 match a: 1132 case 0n: 1133 {==} 1134 case 1n+p: 1135 %add_succ_r(p, b) : {2n+Nat.add(p, b) == 1n+_ : Nat} 1136 {==} 1137 1138 # a + 0 == a 1139 def add_zero_r(+a: Nat) -> {Nat.add(a, 0n) == a : Nat}: 1140 match a: 1141 case 0n: 1142 {==} 1143 case 1n+p: 1144 %add_zero_r(p) : {1n+Nat.add(p, 0n) == 1n+_ : Nat} 1145 {==} 1146 1147 # half.go(n, acc) <= n + acc 1148 def half_go_le(+n: Nat, +acc: Nat) -> {True{} == leb(half.go(n, acc), Nat.add(n, acc)) : Bool}: 1149 match n: 1150 case 0n: 1151 leb_refl(acc) 1152 case 1n+p: 1153 match p: 1154 case 0n: 1155 leb_succ(acc, acc, leb_refl(acc)) 1156 case 1n+q: 1157 %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} 1158 leb_succ(half.go(q, 1n+acc), Nat.add(q, 1n+acc), half_go_le(q, 1n+acc)) 1159 1160 # half(n) <= n 1161 def half_le(+n: Nat) -> {True{} == leb(half(n), n) : Bool}: 1162 %add_zero_r(n) : {True{} == leb(half(n), _) : Bool} 1163 half_go_le(n, 0n) 1164 1165 # a <= r -> r == a + (r - a) 1166 def add_sub(+a: Nat, +r: Nat, e: {True{} == leb(a, r) : Bool}) -> {r == Nat.add(a, Nat.sub(r, a)) : Nat}: 1167 match a r: 1168 case 0n 0n: 1169 {==} 1170 case 0n 1n+q: 1171 {==} 1172 case 1n+p 0n: 1173 %e : BD(_, Unit, {0n == Nat.add(1n+p, Nat.sub(0n, 1n+p)) : Nat}) 1174 Unit{} 1175 case 1n+p 1n+q: 1176 %add_sub(p, q, e) : {1n+q == 1n+_ : Nat} 1177 {==} 1178 1179 # LAW (half_cover): the two bands of a node, half(r) rows and r - half(r) rows, 1180 # add up to exactly r rows: no row is lost and none is counted twice. 1181 law half_cover: 1182 for +n: Nat 1183 {n == Nat.add(half(n), Nat.sub(n, half(n))) : Nat} 1184 1185 def half_cover(n): 1186 add_sub(half(n), n, half_le(n)) 1187 1188 type Bands<-r: Nat, -c: Nat> is Type: 1189 BLeaf{m: Mat} 1190 BNode{x: Bands, y: Bands} 1191 1192 # r x c of zeros in 2^d bands (d = 0: a single band, the same as a Mat) 1193 def Bands.zeros(+d: Nat, +r: Nat, +c: Nat) -> Bands: 1194 match d: 1195 case 0n: 1196 BLeaf{Mat.zeros(r, c)} 1197 case 1n+q: 1198 BNode{Bands.zeros(q, half(r), c), Bands.zeros(q, Nat.sub(r, half(r)), c)} 1199 1200 # the row counts of the two bands follow the row count of the node: if rr == r at run 1201 # time, the same holds for both halves (the bands carry r only in their type) 1202 def eq_half(-a: Nat, -b: Nat, e: {a == b : Nat}) -> {half(a) == half(b) : Nat}: 1203 %e : {half(a) == half(_) : Nat} 1204 {==} 1205 1206 def eq_rest(-a: Nat, -b: Nat, e: {a == b : Nat}) -> {Nat.sub(a, half(a)) == Nat.sub(b, half(b)) : Nat}: 1207 %e : {Nat.sub(a, half(a)) == Nat.sub(_, half(_)) : Nat} 1208 {==} 1209 1210 # does a block of n numbers starting at flat index i end before the second band of a 1211 # node with r rows? does it start in the second band? 1212 def bfill_l(+r: Nat, +c: Nat, +i: Nat, +n: Nat) -> Bool: 1213 Nat.is_le(Nat.add(i, n), Nat.mul(half(r), c)) 1214 1215 def bfill_r(+r: Nat, +c: Nat, +i: Nat) -> Bool: 1216 Nat.is_le(Nat.mul(half(r), c), i) 1217 1218 # writes the first n numbers of xs at index i (and nothing past them) 1219 def fill_n(n: Nat, xs: List<&2, F32>, +i: U32, a: Array) -> Array: 1220 match n xs: 1221 case 1n+p Con{h, t}: 1222 fill_n(p, t, (i + 1 : U32), Array.set(F32, a, i, h)) 1223 case _ _: 1224 a 1225 1226 def fill_band(-r: Nat, -c: Nat, n: Nat, xs: List<&2, F32>, +i: U32, m: Mat) -> Mat: 1227 match m: 1228 case Mat{d}: 1229 Mat{fill_n(n, xs, i, d)} 1230 1231 # writes the first n numbers of xs at flat index i; l and rt say which bands they touch 1232 # (bfill_l, bfill_r); rr is the row count r at run time. A block that crosses into the second 1233 # band is not copied: the first band writes only its part, the second gets the rest by List.drop. 1234 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: 1235 match b l rt: 1236 case BLeaf{m} _ _: 1237 BLeaf{fill_band(r, c, n, xs, U32.from_nat(i), m)} 1238 case BNode{x, y} True{} _: 1239 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} 1240 case BNode{x, y} False{} True{}: 1241 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)} 1242 case BNode{x, y} False{} False{}: 1243 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)))} 1244 1245 def bfill_top(ok: Bool, +r: Nat, +c: Nat, b: Bands, +i: Nat, +n: Nat, +xs: List<&2, F32>) -> Bands: 1246 match ok: 1247 case True{}: 1248 bfill(c, r, b, bfill_l(r, c, i, n), bfill_r(r, c, i), r, {==}, i, n, xs) 1249 case False{}: 1250 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) 1251 1252 # writes the n numbers of xs into b starting at flat index i (row-major, as Mat.fill_at); 1253 # numbers past r*c are dropped. For loading a large matrix block by block. 1254 def Bands.fill_at(+r: Nat, +c: Nat, +i: Nat, +n: Nat, +xs: List<&2, F32>, b: Bands) -> Bands: 1255 bfill_top(Nat.is_le(Nat.add(i, n), Nat.mul(r, c)), r, c, b, i, n, xs) 1256 1257 # ---- y = B · x for B with r rows of c numbers (a weight matrix stored out x in) ---- 1258 1259 # B (intact) and y = B · x, typed like Mat.matmul_nt: x is 1 x c, y is 1 x r 1260 type BV<-r: Nat, -c: Nat> is Type: 1261 BV{b: Bands, y: Mat<1n, r>} 1262 1263 # the bands and y as a list in row order (the result of Bands.matvec_l) 1264 type BL<-r: Nat, -c: Nat> is Type: 1265 BL{b: Bands, y: List<&2, F32>} 1266 1267 def bv_g2(-r: Nat, -c: Nat, w: Array, al: AL) -> BL: 1268 match al: 1269 case AL{d, l}: 1270 BL{BLeaf{Mat{w}}, l} 1271 1272 # the band's 1 + m rows came out in an Array (one number per row); read them back as a list 1273 def bv_g(-r: Nat, -c: Nat, +m: Nat, g: G) -> BL: 1274 match g: 1275 case G{a, b, y}: 1276 bv_g2(r, c, b, read_l(m, 0, Nil{}, Array.get(F32, y, 0))) 1277 1278 # the rows of one band, one dot product each, with the gemm of Mat.matmul_nt (n = 1), so the 1279 # results are identical. Writing the results into an Array and reading them back is faster than 1280 # building the list row by row (NOTES.md, exp. 14). 1281 def bv_gemm(-r: Nat, +c: Nat, rows: Nat, +rr: Nat, x: Array, w: Array) -> BL: 1282 match rows: 1283 case 0n: 1284 BL{BLeaf{Mat{w}}, Nil{}} 1285 case 1n+p: 1286 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))) 1287 1288 def bv_mat(-r: Nat, +c: Nat, +rr: Nat, x: Array, m: Mat) -> BL: 1289 match m: 1290 case Mat{w}: 1291 bv_gemm(r, c, rr, rr, x, w) 1292 1293 def bv_join(-r: Nat, -c: Nat, p: BL, q: BL) -> BL: 1294 match p q: 1295 case BL{bx, yx} BL{by, yy}: 1296 BL{BNode{bx, by}, List.append(&2, F32, yx, yy)} 1297 1298 # the two bands of a node run in parallel, each with its own copy of x (c numbers); 1299 # 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 1300 # parameter at the end of the list turned the parallel let into two sequential calls (NOTES.md, exp. 11; fixed in 1301 # 2.0.37 by bendlang/bend#1377). The order is kept, and harmless. 1302 def bmv(+c: Nat, -r: Nat, b: Bands, +rr: Nat, -e: {rr == r : Nat}, xp: Array & Array) -> BL: 1303 match b xp: 1304 case BLeaf{m} Tuple{x, x2}: 1305 bv_mat(r, c, rr, x, m) 1306 case BNode{bx, by} Tuple{x1, x2}: 1307 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)) 1308 bv_join(r, c, p, q) 1309 1310 def bv_fin(+r: Nat, -c: Nat, bl: BL) -> BV: 1311 match bl: 1312 case BL{b, y}: 1313 BV{b, Mat.from_list(1n, r, y)} 1314 1315 # 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). 1316 # The same product as Mat.matmul_nt(1n, c, r, ..), and the same numbers, bit for bit. 1317 def Bands.matvec(+r: Nat, +c: Nat, b: Bands, x: Mat<1n, c>) -> BV: 1318 match x: 1319 case Mat{xa}: 1320 bv_fin(r, c, bmv(c, r, b, r, {==}, Array.clone(F32, xa))) 1321 1322 # the same product with x and y as lists (x: c numbers, y: r numbers in row order), without the 1323 # Mat conversions of Bands.matvec: for callers that already hold the vector as a list 1324 def Bands.matvec_l(+r: Nat, +c: Nat, b: Bands, +x: List<&2, F32>) -> BL: 1325 bmv(c, r, b, r, {==}, Array.clone(F32, fill_list(x, 0, Array.new(F32, cap_depth(c), 0.0)))) 1326 1327 # ---- reading rows and the whole matrix ---- 1328 1329 type BRow<-r: Nat, -c: Nat> is Type: 1330 BRow{b: Bands, row: List<&2, F32>} 1331 1332 def brow_fin(-r: Nat, -c: Nat, al: AL) -> BRow: 1333 match al: 1334 case AL{d, l}: 1335 BRow{BLeaf{Mat{d}}, l} 1336 1337 # n numbers starting at flat index i of one band (a band may have no rows) 1338 def brow_leaf2(-r: Nat, -c: Nat, n: Nat, +i: U32, d: Array) -> BRow: 1339 match n: 1340 case 0n: 1341 BRow{BLeaf{Mat{d}}, Nil{}} 1342 case 1n+p: 1343 brow_fin(r, c, read_l(p, i, Nil{}, Array.get(F32, d, i))) 1344 1345 def brow_leaf(-r: Nat, -c: Nat, m: Mat, +i: U32, n: Nat) -> BRow: 1346 match m: 1347 case Mat{d}: 1348 brow_leaf2(r, c, n, i, d) 1349 1350 def brow_l(-r: Nat, -c: Nat, p: BRow, y: Bands) -> BRow: 1351 match p: 1352 case BRow{bx, row}: 1353 BRow{BNode{bx, y}, row} 1354 1355 def brow_r(-r: Nat, -c: Nat, x: Bands, q: BRow) -> BRow: 1356 match q: 1357 case BRow{by, row}: 1358 BRow{BNode{x, by}, row} 1359 1360 # row i; left says whether i is in the first band of the node (i < half(rr)) 1361 def brow(+c: Nat, -r: Nat, b: Bands, left: Bool, +rr: Nat, -e: {rr == r : Nat}, +i: Nat) -> BRow: 1362 match b left: 1363 case BLeaf{m} _: 1364 brow_leaf(r, c, m, U32.from_nat(Nat.mul(i, c)), c) 1365 case BNode{x, y} True{}: 1366 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) 1367 case BNode{x, y} False{}: 1368 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)))) 1369 1370 # row i (c numbers), for example an embedding lookup 1371 def Bands.read_row(+r: Nat, +c: Nat, b: Bands, +i: Nat) -> BRow: 1372 brow(c, r, b, Nat.is_lt(i, half(r)), r, {==}, i) 1373 1374 def bl_join(-r: Nat, -c: Nat, p: BRow, q: BRow) -> BRow: 1375 match p q: 1376 case BRow{bx, xs} BRow{by, ys}: 1377 BRow{BNode{bx, by}, List.append(&2, F32, xs, ys)} 1378 1379 def blist(+c: Nat, -r: Nat, b: Bands, +rr: Nat, -e: {rr == r : Nat}) -> BRow: 1380 match b: 1381 case BLeaf{m}: 1382 brow_leaf(r, c, m, 0, Nat.mul(rr, c)) 1383 case BNode{x, y}: 1384 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))) 1385 1386 # all r*c numbers in row order (in BRow.row) 1387 def Bands.to_list(+r: Nat, +c: Nat, b: Bands) -> BRow: 1388 blist(c, r, b, r, {==}) 1389 1390 # ---- the list path returns exactly r numbers ---- 1391 # 1392 # The numbers do not matter here, only how many there are: each lemma below counts the list a 1393 # function builds, following the function's own recursion. 1394 1395 def bl_len(-r: Nat, -c: Nat, bl: BL) -> Nat: 1396 match bl: 1397 case BL{b, y}: 1398 List.length(&2, F32, y) 1399 1400 # s == a + (1 + b) -> s == 1 + (a + b) 1401 def succ_out(-s: Nat, +a: Nat, +b: Nat, e: {s == Nat.add(a, 1n+b) : Nat}) -> {s == 1n+Nat.add(a, b) : Nat}: 1402 %Equal.sym(Nat, 1n+Nat.add(a, b), Nat.add(a, 1n+b), add_succ_r(a, b)) : {s == _ : Nat} 1403 e 1404 1405 # s == 1 + (p + (1 + l)) -> s == 2 + (p + l) 1406 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}: 1407 %Equal.sym(Nat, 1n+Nat.add(p, l), Nat.add(p, 1n+l), add_succ_r(p, l)) : {s == 1n+_ : Nat} 1408 e 1409 1410 # s == n + 0 -> s == n 1411 def zero_out(-s: Nat, +n: Nat, e: {s == Nat.add(n, 0n) : Nat}) -> {s == n : Nat}: 1412 %add_zero_r(n) : {s == _ : Nat} 1413 e 1414 1415 # reversing onto acc adds the lengths 1416 def rev_len(+xs: List<&2, F32>, +acc: List<&2, F32>) -> {List.length(&2, F32, List.reverse.go(&2, F32, xs, acc)) == Nat.add(List.length(&2, F32, xs), List.length(&2, F32, acc)) : Nat}: 1417 match xs: 1418 case Nil{}: 1419 {==} 1420 case Con{h, t}: 1421 succ_out(List.length(&2, F32, List.reverse.go(&2, F32, t, h <> acc)), List.length(&2, F32, t), List.length(&2, F32, acc), rev_len(t, h <> acc)) 1422 1423 # appending adds the lengths 1424 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}: 1425 match xs: 1426 case Nil{}: 1427 {==} 1428 case Con{h, t}: 1429 %app_len(t, ys) : {1n+List.length(&2, F32, List.append(&2, F32, t, ys)) == 1n+_ : Nat} 1430 {==} 1431 1432 # a leaf's list: read_l reads mleft + 1 numbers onto acc and reverses them. By induction on 1433 # mleft: with 0 left it reads one number (rev_len counts the reversed list); with 1 + p left it 1434 # reads one and recurses, and len_step moves the 1 out of the sum. 1435 law leaf_len: 1436 for -r: Nat 1437 for -c: Nat 1438 for w: Array 1439 for +mleft: Nat 1440 for +idx: U32 1441 for +acc: List<&2, F32> 1442 for rc: Array & F32 1443 {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} 1444 1445 def leaf_len(r, c, w, mleft, idx, acc, rc): 1446 match mleft rc: 1447 case 0n Tuple{d, +v}: 1448 zero_out(List.length(&2, F32, List.reverse.go(&2, F32, v <> acc, Nil{})), 1n+List.length(&2, F32, acc), rev_len(v <> acc, Nil{})) 1449 case 1n+p Tuple{d, +v}: 1450 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)))) 1451 1452 # whatever the gemm wrote, reading 1 + p rows back gives 1 + p numbers (G has one constructor, so 1453 # matching g opens it, and leaf_len counts the list) 1454 def g_len(-r: Nat, -c: Nat, +p: Nat, g: G) -> {bl_len(r, c, bv_g(r, c, p, g)) == 1n+p : Nat}: 1455 match g: 1456 case G{a, b, y}: 1457 zero_out(bl_len(r, c, bv_g2(r, c, b, read_l(p, 0, Nil{}, Array.get(F32, y, 0)))), 1n+p, leaf_len(r, c, b, p, 0, Nil{}, Array.get(F32, y, 0))) 1458 1459 # LAW (band_len): the product of one band with `rows` rows gives exactly `rows` numbers. 1460 # 0 rows give the empty list; 1 + p rows go through the gemm, and g_len counts what is read back. 1461 law band_len: 1462 for -r: Nat 1463 for +c: Nat 1464 for +rows: Nat 1465 for x: Array 1466 for w: Array 1467 {bl_len(r, c, bv_gemm(r, c, rows, rows, x, w)) == rows : Nat} 1468 1469 def band_len(r, c, rows, x, w): 1470 match rows: 1471 case 0n: 1472 {==} 1473 case 1n+p: 1474 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))) 1475 1476 # joining two halves of h1 and h2 numbers gives h1 + h2 numbers. (The statement for the whole tree, 1477 # "Bands.matvec_l gives r numbers", would apply this to the results of the two recursive calls and 1478 # to the induction hypotheses about the same calls; that uses each band's Array twice, which Bend's 1479 # affine rules refuse. It stays open, covered by tests: NOTES.md, v3.1.) 1480 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}: 1481 match p q: 1482 case BL{bx, +yx} BL{by, +yy}: 1483 %hp : {List.length(&2, F32, List.append(&2, F32, yx, yy)) == Nat.add(_, h2) : Nat} 1484 %hq : {List.length(&2, F32, List.append(&2, F32, yx, yy)) == Nat.add(List.length(&2, F32, yx), _) : Nat} 1485 app_len(yx, yy) 1486