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.2.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 # smallest d with 2^d >= n (n >= 2) 31 32 def dotg(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, acc: F32, ra: Array & F32, rb: Array & F32) -> R: 33 match k ra rb: 34 case 0n Tuple{a, x} Tuple{b, y}: 35 R{a, b, acc} 36 case 1n+p Tuple{a, +x} Tuple{b, +y}: 37 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))) 38 39 # row i of C, columns j.. ; r is the product of the current column (already computed) 40 41 def gcols(mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, +cidx: U32, c: Array, r: R) -> G: 42 match mleft r: 43 case 0n R{a, b, +x}: 44 G{a, b, Array.set(F32, c, cidx, x)} 45 case 1n+p R{a, b, +x}: 46 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)))) 47 48 def grows(nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +mu: U32, g: G) -> G: 49 match nleft g: 50 case 0n G{a, b, c}: 51 G{a, b, c} 52 case 1n+p G{a, b, c}: 53 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)))) 54 55 # C (n x m) = A · B with k steps; c is the output array (with capacity n*m) 56 57 def gemm(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 58 grows(n, Nat.sub(m, 1n), k, 0, arow, sa, sb, bcol, U32.from_nat(m), G{a, b, c}) 59 60 61 # --------------------------------------------------------------------- 62 # parallel gemm: 2^d blocks of rows of C, each with its own copy of A and B, 63 # returning lists that are written into C at the end 64 # --------------------------------------------------------------------- 65 66 type S3 is Type: 67 S3{a: Array, b: Array, out: List<&2, F32>} 68 69 def lcols(mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, out: List<&2, F32>, r: R) -> S3: 70 match mleft r: 71 case 0n R{a, b, +x}: 72 S3{a, b, x <> out} 73 case 1n+p R{a, b, +x}: 74 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)))) 75 76 def lrows(nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, st: S3) -> S3: 77 match nleft st: 78 case 0n S3{a, b, out}: 79 S3{a, b, out} 80 case 1n+p S3{a, b, out}: 81 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)))) 82 83 def lblock_fin(st: S3) -> List<&2, F32>: 84 match st: 85 case S3{a, b, out}: 86 List.reverse(&2, F32, out) 87 88 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>: 89 lblock_fin(lrows(cnt, mm1, kn, i, arow, sa, sb, bcol, S3{a, b, Nil{}})) 90 91 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>: 92 match d ra rb: 93 case 0n Tuple{a, a2} Tuple{b, b2}: 94 lblock(cnt, mm1, kn, i, arow, sa, sb, bcol, a, b) 95 case 1n++q Tuple{a, a2} Tuple{b, b2}: 96 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)) 97 List.append(&2, F32, x, y) 98 99 def write_l(xs: List<&2, F32>, +idx: U32, c: Array) -> Array: 100 match xs: 101 case Nil{}: 102 c 103 case Con{h, t}: 104 write_l(t, (idx + 1 : U32), Array.set(F32, c, idx, h)) 105 106 # c[idx..] += list 107 108 def gp3(a: Array, b: Array, c: Array, xs: List<&2, F32>) -> G: 109 G{a, b, write_l(xs, 0, c)} 110 111 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: 112 match pa pb: 113 case Tuple{a1, a2} Tuple{b1, b2}: 114 gp3(a1, b1, c, lpar(d, n, mm1, kn, 0, arow, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2))) 115 116 # C (n x m) = A · B, with 2^d blocks of rows in parallel; returns A, B (intact) and C 117 118 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: 119 gp2(d, n, Nat.sub(m, 1n), k, arow, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b)) 120 121 122 # --------------------------------------------------------------------- 123 # parallel gemm by COLUMNS, for n = 1 (matrix · vector): each block computes a 124 # range of columns of C, with its own copy of the operands 125 # --------------------------------------------------------------------- 126 127 def cblock_fin(st: S3) -> List<&2, F32>: 128 match st: 129 case S3{a, b, out}: 130 List.reverse(&2, F32, out) 131 132 # cnt columns starting at j0 (cnt >= 1); row 0 of A is at ia0 133 134 def cblock_go(cnt1: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array) -> List<&2, F32>: 135 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))))) 136 137 def cblock(cnt: Nat, +kn: Nat, +ia0: U32, +j0: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array) -> List<&2, F32>: 138 match cnt: 139 case 0n: 140 Nil{} 141 case 1n+p: 142 cblock_go(p, kn, ia0, j0, sa, sb, bcol, a, b) 143 144 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>: 145 match d ra rb: 146 case 0n Tuple{a, a2} Tuple{b, b2}: 147 cblock(cnt, kn, ia0, j0, sa, sb, bcol, a, b) 148 case 1n++q Tuple{a, a2} Tuple{b, b2}: 149 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)) 150 List.append(&2, F32, x, y) 151 152 def cp2(d: Nat, +m: Nat, +kn: Nat, +sa: U32, +sb: U32, +bcol: U32, c: Array, pa: Array & Array, pb: Array & Array) -> G: 153 match pa pb: 154 case Tuple{a1, a2} Tuple{b1, b2}: 155 gp3(a1, b1, c, cpar(d, m, kn, 0, 0, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2))) 156 157 # C (1 x m) = A (1 x k) · B, with 2^d blocks of columns in parallel 158 159 def gemm_par_cols(d: Nat, +m: Nat, +k: Nat, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> G: 160 cp2(d, m, k, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b)) 161 162 # --------------------------------------------------------------------- 163 # element-wise operations (in place, by index) 164 # --------------------------------------------------------------------- 165 166 def depth(n: Nat) -> Nat: 167 Nat.add(U32.log2(U32.from_nat(Nat.sub(n, 1n))), 1n) 168 169 # --------------------------------------------------------------------- 170 # gemm: C[i][j] = sum_p A[i*arow + p*sa] * B[p*sb + j*bcol] 171 # (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) 172 # --------------------------------------------------------------------- 173 174 # acc += A[ia + t*sa] * B[ib + t*sb] for k steps; ra, rb already carry the first pair 175 176 type P2 is Type: 177 P2{a: Array, b: Array} 178 179 type AL is Type: 180 AL{c: Array, l: List<&2, F32>} 181 182 # list of m elements starting at idx (reads with Array.get, returning the array) 183 184 def read_l(mleft: Nat, +idx: U32, acc: List<&2, F32>, rc: Array & F32) -> AL: 185 match mleft rc: 186 case 0n Tuple{c, +v}: 187 AL{c, List.reverse(&2, F32, v <> acc)} 188 case 1n+p Tuple{c, +v}: 189 read_l(p, (idx + 1 : U32), v <> acc, Array.get(F32, c, (idx + 1 : U32))) 190 191 # writes the list starting at idx 192 193 def addl(xs: List<&2, F32>, +idx: U32, rc: Array & F32) -> Array: 194 match xs rc: 195 case Nil{} Tuple{c, v}: 196 c 197 case Con{+h, t} Tuple{c, +v}: 198 addl(t, (idx + 1 : U32), Array.get(F32, Array.set(F32, c, idx, (v + h : F32)), (idx + 1 : U32))) 199 200 # adds a bias (a list of m) to each of the n rows of c (n x m) 201 202 def bias_rows(nleft: Nat, +i: U32, +mu: U32, +bl: List<&2, F32>, c: Array) -> Array: 203 match nleft: 204 case 0n: 205 c 206 case 1n+p: 207 bias_rows(p, (i + 1 : U32), mu, bl, addl(bl, (i * mu : U32), Array.get(F32, c, (i * mu : U32)))) 208 209 # in-place relu on the first n elements 210 211 def relu_l(nleft: Nat, +t: U32, rc: Array & F32) -> Array: 212 match nleft rc: 213 case 0n Tuple{c, v}: 214 c 215 case 1n+p Tuple{c, +v}: 216 relu_l(p, (t + 1 : U32), Array.get(F32, Array.set(F32, c, t, F32.max(v, 0.0)), (t + 1 : U32))) 217 218 def pos2(b: Bool) -> F32: 219 match b: 220 case True{}: 221 1.0 222 case False{}: 223 0.0 224 225 def pos(+y: F32) -> F32: 226 pos2(F32.is_gt(y, 0.0)) 227 228 # d[t] *= (h[t] > 0): the relu gradient; returns both arrays 229 230 def mask_l(nleft: Nat, +t: U32, rd: Array & F32, rh: Array & F32) -> P2: 231 match nleft rd rh: 232 case 0n Tuple{d, x} Tuple{h, y}: 233 P2{d, h} 234 case 1n+p Tuple{d, +x} Tuple{h, +y}: 235 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))) 236 237 # w[t] -= lr * dw[t] 238 239 def maxl2(xs: List<&2, F32>, cur: F32) -> F32: 240 match xs: 241 case Nil{}: 242 cur 243 case Con{+h, t}: 244 maxl2(t, F32.max(cur, h)) 245 246 def maxl(xs: List<&2, F32>) -> F32: 247 match xs: 248 case Nil{}: 249 0.0 250 case Con{+h, t}: 251 maxl2(t, h) 252 253 def expl(xs: List<&2, F32>, +m: F32) -> List<&2, F32>: 254 match xs: 255 case Nil{}: 256 Nil{} 257 case Con{+h, t}: 258 F32.exp((h - m : F32)) <> expl(t, m) 259 260 def suml(xs: List<&2, F32>, acc: F32) -> F32: 261 match xs: 262 case Nil{}: 263 acc 264 case Con{+h, t}: 265 suml(t, (acc + h : F32)) 266 267 def divl(xs: List<&2, F32>, +s: F32) -> List<&2, F32>: 268 match xs: 269 case Nil{}: 270 Nil{} 271 case Con{+h, t}: 272 (h / s : F32) <> divl(t, s) 273 274 def softmax2(+es: List<&2, F32>) -> List<&2, F32>: 275 divl(es, suml(es, 0.0)) 276 277 def softmax(+xs: List<&2, F32>) -> List<&2, F32>: 278 softmax2(expl(xs, maxl(xs))) 279 280 # p[label] 281 282 def sgd_l(nleft: Nat, +t: U32, +lr: F32, rw: Array & F32, rd: Array & F32) -> P2: 283 match nleft rw rd: 284 case 0n Tuple{w, x} Tuple{d, y}: 285 P2{w, d} 286 case 1n+p Tuple{w, +x} Tuple{d, +y}: 287 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))) 288 289 # --------------------------------------------------------------------- 290 # softmax + cross-entropy per row (lists of 10 elements: negligible cost) 291 # --------------------------------------------------------------------- 292 293 def pick(ps: List<&2, F32>, n: Nat) -> F32: 294 match ps n: 295 case Con{+h, t} 0n: 296 h 297 case Con{h, t} 1n+q: 298 pick(t, q) 299 case Nil{} _: 300 0.0 301 302 # (p - onehot(label)) / n 303 304 def hit(b: Bool) -> F32: 305 match b: 306 case True{}: 307 1.0 308 case False{}: 309 0.0 310 311 def grad_l(ps: List<&2, F32>, +label: Nat, +j: Nat, +nf: F32) -> List<&2, F32>: 312 match ps: 313 case Nil{}: 314 Nil{} 315 case Con{+h, t}: 316 (((h - hit(Nat.is_eq(label, j)) : F32) / nf) : F32) <> grad_l(t, label, 1n+j, nf) 317 318 type L2 is Type: 319 L2{loss: F32, c: Array} 320 321 def ce_row2(+ps: List<&2, F32>, +label: Nat, +nf: F32, +idx: U32, loss: F32, c: Array) -> L2: 322 L2{(loss - F32.log(F32.max(pick(ps, label), 0.000000000001)) : F32), write_l(grad_l(ps, label, 0n, nf), idx, c)} 323 324 def ce_row(al: AL, +label: Nat, +nf: F32, +idx: U32, loss: F32) -> L2: 325 match al: 326 case AL{c, l}: 327 ce_row2(softmax(l), label, nf, idx, loss, c) 328 329 # 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 330 331 type Mat<-r: Nat, -c: Nat> is Type: 332 Mat{d: Array} 333 334 # smallest d with 2^d >= n (n >= 2); n < 2 gives 0 335 336 def ce_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +nf: F32, +i: U32, r: L2) -> L2: 337 match labels r: 338 case Nil{} L2{loss, c}: 339 L2{loss, c} 340 case Con{+y, t} L2{loss, c}: 341 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)) 342 343 # ---- helper: hits per row ---- 344 345 type CR is Type: 346 CR{c: Array, n: Nat} 347 348 def gt_head(xs: List<&2, F32>, +v: F32) -> Bool: 349 match xs: 350 case Nil{}: 351 False{} 352 case Con{+h, t}: 353 F32.is_gt(h, v) 354 355 def argmax2(xs: List<&2, F32>, g: Bool, +i: Nat, +bi: Nat, +bv: F32) -> Nat: 356 match xs g: 357 case Nil{} _: 358 bi 359 case Con{+h, +t} True{}: 360 argmax2(t, gt_head(t, h), 1n+i, i, h) 361 case Con{h, +t} False{}: 362 argmax2(t, gt_head(t, bv), 1n+i, bi, bv) 363 364 def argmax(xs: List<&2, F32>) -> Nat: 365 match xs: 366 case Nil{}: 367 0n 368 case Con{+h, +t}: 369 argmax2(t, gt_head(t, h), 1n, 0n, h) 370 371 def one(b: Bool) -> Nat: 372 match b: 373 case True{}: 374 1n 375 case False{}: 376 0n 377 378 def count_rows3(+y: Nat, n: Nat, al: AL) -> CR: 379 match al: 380 case AL{c, l}: 381 CR{c, Nat.add(n, one(Nat.is_eq(argmax(l), y)))} 382 383 def count_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +i: U32, r: CR) -> CR: 384 match labels r: 385 case Nil{} CR{c, n}: 386 CR{c, n} 387 case Con{+y, t} CR{c, n}: 388 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))))) 389 390 391 # ===================================================================== 392 # Typed API: Mat (the dimensions enter the type) 393 # ===================================================================== 394 395 def cap_depth2(small: Bool, n: Nat) -> Nat: 396 match small: 397 case True{}: 398 0n 399 case False{}: 400 depth(n) 401 402 def cap_depth(+n: Nat) -> Nat: 403 cap_depth2(Nat.is_lt(n, 2n), n) 404 405 # r x c of zeros, with the exact capacity (smallest power of 2) 406 407 def Mat.zeros(+r: Nat, +c: Nat) -> Mat: 408 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0)} 409 410 def Mat.fill(+r: Nat, +c: Nat, +v: F32) -> Mat: 411 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), v)} 412 413 def fill_list(xs: List<&2, F32>, +i: U32, a: Array) -> Array: 414 match xs: 415 case Nil{}: 416 a 417 case Con{h, t}: 418 fill_list(t, (i + 1 : U32), Array.set(F32, a, i, h)) 419 420 def len_is(xs: List<&2, F32>, n: Nat) -> Bool: 421 match xs n: 422 case Nil{} 0n: 423 True{} 424 case Con{h, t} 1n+p: 425 len_is(t, p) 426 case _ _: 427 False{} 428 429 def Mat.of_list2(-r: Nat, -c: Nat, +rr: Nat, +cc: Nat, ok: Bool, xs: List<&2, F32>) -> Maybe<&1, Mat>: 430 match ok: 431 case True{}: 432 Some{Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(rr, cc)), 0.0))}} 433 case False{}: 434 None{} 435 436 # Mat from r*c numbers in row order; None if the size does not match 437 438 def Mat.of_list(+r: Nat, +c: Nat, +xs: List<&2, F32>) -> Maybe<&1, Mat>: 439 Mat.of_list2(r, c, r, c, len_is(xs, Nat.mul(r, c)), xs) 440 441 # reads the first n numbers, returning the matrix 442 443 # Mat without checking the list size (a short list leaves the rest at zero, a long one is cut 444 # by the capacity): only for when the size is known by construction. Prefer Mat.of_list. 445 446 def Mat.from_list(+r: Nat, +c: Nat, xs: List<&2, F32>) -> Mat: 447 Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0))} 448 449 type ML<-r: Nat, -c: Nat> is Type: 450 ML{m: Mat, l: List<&2, F32>} 451 452 def Mat.to_list2(-r: Nat, -c: Nat, al: AL) -> ML: 453 match al: 454 case AL{d, l}: 455 ML{Mat{d}, l} 456 457 # two independent copies 458 459 def Mat.to_list(+r: Nat, +c: Nat, m: Mat) -> ML: 460 match m: 461 case Mat{d}: 462 Mat.to_list2(r, c, read_l(Nat.sub(Nat.mul(r, c), 1n), 0, Nil{}, Array.get(F32, d, 0))) 463 464 type MM2<-r: Nat, -c: Nat> is Type: 465 MM2{a: Mat, b: Mat} 466 467 def Mat.clone2(-r: Nat, -c: Nat, p: Array & Array) -> MM2: 468 match p: 469 case Tuple{x, y}: 470 MM2{Mat{x}, Mat{y}} 471 472 # ---- products: C (n x m) = A · B, with 2^par blocks of rows in parallel ---- 473 474 # C = A(n x k) · B(k x m) 475 476 def Mat.clone(-r: Nat, -c: Nat, m: Mat) -> MM2: 477 match m: 478 case Mat{d}: 479 Mat.clone2(r, c, Array.clone(F32, d)) 480 481 type MMul<-n: Nat, -k: Nat, -m: Nat> is Type: 482 MMul{a: Mat, b: Mat, c: Mat} 483 484 # C = A(n x k) · Bᵀ, with B stored m x k 485 486 type MMulNT<-n: Nat, -k: Nat, -m: Nat> is Type: 487 MMulNT{a: Mat, b: Mat, c: Mat} 488 489 # C = Aᵀ · B, with A stored k x n 490 491 type MMulTN<-n: Nat, -k: Nat, -m: Nat> is Type: 492 MMulTN{a: Mat, b: Mat, c: Mat} 493 494 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: 495 match cols: 496 case True{}: 497 gemm_par_cols(d, m, k, sa, sb, bcol, a, b, c) 498 case False{}: 499 gemm_par(d, n, m, k, arow, sa, sb, bcol, a, b, c) 500 501 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: 502 match d: 503 case 0n: 504 gemm(n, m, k, arow, sa, sb, bcol, a, b, c) 505 case 1n+q: 506 gemm_any3(Nat.is_eq(n, 1n), 1n+q, n, m, k, arow, sa, sb, bcol, a, b, c) 507 508 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: 509 gemm_any2(d, n, m, k, arow, sa, sb, bcol, a, b, c) 510 511 def mmul_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMul: 512 match g: 513 case G{a, b, c}: 514 MMul{Mat{a}, Mat{b}, Mat{c}} 515 516 def mmul_nt_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulNT: 517 match g: 518 case G{a, b, c}: 519 MMulNT{Mat{a}, Mat{b}, Mat{c}} 520 521 def mmul_tn_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulTN: 522 match g: 523 case G{a, b, c}: 524 MMulTN{Mat{a}, Mat{b}, Mat{c}} 525 526 def mm_a(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array, b: Array) -> MMul: 527 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))) 528 529 # C = A · B (n x k) · (k x m); par = log2 of the number of parallel blocks (0 = sequential) 530 531 def Mat.matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat, b: Mat) -> MMul: 532 match a b: 533 case Mat{x} Mat{y}: 534 mm_a(n, k, m, par, x, y) 535 536 def mm_nt(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array, b: Array) -> MMulNT: 537 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))) 538 539 # C = A · Bᵀ (n x k) · (m x k)ᵀ: for weights stored with one row per output 540 541 def Mat.matmul_nt(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat, b: Mat) -> MMulNT: 542 match a b: 543 case Mat{x} Mat{y}: 544 mm_nt(n, k, m, par, x, y) 545 546 def mm_tn(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array, b: Array) -> MMulTN: 547 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))) 548 549 # C = Aᵀ · B (k x n)ᵀ · (k x m): weight gradient, X^T · dY 550 551 def Mat.matmul_tn(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat, b: Mat) -> MMulTN: 552 match a b: 553 case Mat{x} Mat{y}: 554 mm_tn(n, k, m, par, x, y) 555 556 # ---- element-wise and training operations ---- 557 558 # adds the bias (one row, m numbers) to each row of C (n x m) 559 560 type MBias<-n: Nat, -m: Nat> is Type: 561 MBias{c: Mat, b: Mat<1n, m>} 562 563 def add_row2(-n: Nat, -m: Nat, +nn: Nat, +mm: Nat, cd: Array, al: AL) -> MBias: 564 match al: 565 case AL{bd, bl}: 566 MBias{Mat{bias_rows(nn, 0, U32.from_nat(mm), bl, cd)}, Mat{bd}} 567 568 def Mat.add_row(+n: Nat, +m: Nat, c: Mat, b: Mat<1n, m>) -> MBias: 569 match c b: 570 case Mat{cd} Mat{bd}: 571 add_row2(n, m, n, m, cd, read_l(Nat.sub(m, 1n), 0, Nil{}, Array.get(F32, bd, 0))) 572 573 # pair of matrices of the same shape (the result of an operation that reads two operands) 574 575 type MPair<-r: Nat, -c: Nat> is Type: 576 MPair{a: Mat, b: Mat} 577 578 def mpair_fin(-r: Nat, -c: Nat, p: P2) -> MPair: 579 match p: 580 case P2{a, b}: 581 MPair{Mat{a}, Mat{b}} 582 583 def relu_fin(-r: Nat, -c: Nat, d: Array) -> Mat: 584 Mat{d} 585 586 def Mat.relu(+r: Nat, +c: Nat, a: Mat) -> Mat: 587 match a: 588 case Mat{d}: 589 relu_fin(r, c, relu_l(Nat.mul(r, c), 0, Array.get(F32, d, 0))) 590 591 # (dy * (h > 0), h): the gradient that passes through the relu 592 593 def Mat.relu_bwd(+r: Nat, +c: Nat, dy: Mat, h: Mat) -> MPair: 594 match dy h: 595 case Mat{d} Mat{x}: 596 mpair_fin(r, c, mask_l(Nat.mul(r, c), 0, Array.get(F32, d, 0), Array.get(F32, x, 0))) 597 598 # (w - lr * dw, dw): one SGD step 599 600 def Mat.sgd(+r: Nat, +c: Nat, +lr: F32, w: Mat, dw: Mat) -> MPair: 601 match w dw: 602 case Mat{x} Mat{y}: 603 mpair_fin(r, c, sgd_l(Nat.mul(r, c), 0, lr, Array.get(F32, x, 0), Array.get(F32, y, 0))) 604 605 # (a + b, b): element-wise sum (a - (-1)*b is exactly a + b) 606 607 def Mat.add(+r: Nat, +c: Nat, a: Mat, b: Mat) -> MPair: 608 match a b: 609 case Mat{x} Mat{y}: 610 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))) 611 612 # sum of each column (n x m -> 1 x m), returning the matrix 613 614 type MColSum<-n: Nat, -m: Nat> is Type: 615 MColSum{a: Mat, s: Mat<1n, m>} 616 617 def colsum_fin(-n: Nat, -m: Nat, g: G) -> MColSum: 618 match g: 619 case G{ones, a, s}: 620 MColSum{Mat{a}, Mat{s}} 621 622 def Mat.col_sums(+n: Nat, +m: Nat, a: Mat) -> MColSum: 623 match a: 624 case Mat{d}: 625 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))) 626 627 # mean cross-entropy and its gradient on the logits: (softmax - one-hot) / n. 628 # labels has one entry per row (n); the gradient replaces the logits. 629 630 type MCE<-n: Nat, -c: Nat> is Type: 631 MCE{loss: F32, grad: Mat} 632 633 def ce_fin(-n: Nat, -c: Nat, +nf: F32, r: L2) -> MCE: 634 match r: 635 case L2{loss, d}: 636 MCE{(loss / nf : F32), Mat{d}} 637 638 def Mat.softmax_ce(+n: Nat, +c: Nat, z: Mat, labels: List<&2, Nat>) -> MCE: 639 match z: 640 case Mat{d}: 641 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})) 642 643 # argmax hits per row against labels 644 645 type MHits<-n: Nat, -c: Nat> is Type: 646 MHits{z: Mat, hits: Nat} 647 648 def hits_fin(-n: Nat, -c: Nat, r: CR) -> MHits: 649 match r: 650 case CR{d, k}: 651 MHits{Mat{d}, k} 652 653 def Mat.count_correct(+n: Nat, +c: Nat, z: Mat, labels: List<&2, Nat>) -> MHits: 654 match z: 655 case Mat{d}: 656 hits_fin(n, c, count_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), 0, CR{d, 0n})) 657 658 # a row as a list (returning the matrix) and the other way around 659 660 type MRow<-r: Nat, -c: Nat> is Type: 661 MRow{m: Mat, row: List<&2, F32>} 662 663 def row_fin(-r: Nat, -c: Nat, al: AL) -> MRow: 664 match al: 665 case AL{d, l}: 666 MRow{Mat{d}, l} 667 668 def Mat.read_row(+r: Nat, +c: Nat, m: Mat, +i: U32) -> MRow: 669 match m: 670 case Mat{d}: 671 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)))) 672 673 def Mat.write_row(-r: Nat, +c: Nat, m: Mat, +i: U32, xs: List<&2, F32>) -> Mat: 674 match m: 675 case Mat{d}: 676 Mat{write_l(xs, (i * U32.from_nat(c) : U32), d)}