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