1 # bend-ml-tensor-fork: bend-ml-tensor-array's products, with the parallel blocks sharing their operands; and par1..par4, 2 # any kernel over blocks of a range sharing up to four arrays. 3 # 4 # import bend-ml-tensor-array@0.1.18.0/main.bend as TA 5 # import bend-ml-tensor-fork@0.1.5.0/main.bend as TF 6 # TF.matmul(n, k, m, par, a, b) # the same arguments, result type and bits as TA.Mat.matmul 7 # 8 # The six products (matmul, matmul_nt, matmul_tn and their accumulating forms) take tensor-array's Mat values and 9 # return its result types. A product that splits (tensor-array's split_depth, two or more rows) packs A and B once into 10 # cells of 8, as the sequential path does (in blocks of lines when an operand holds 16384 cells or more), and every 11 # node of the fork tree shares the two packed operands and C through Base's Array.fork: two handles to one block, no 12 # copy. Each leaf runs tensor-array's cells (t8_rows, t8_cols) straight 13 # into its own rows (or columns) of C. tensor-array's parallel path instead clones an operand at every node and hands 14 # each block's cells back as a list that is then written into C. A product that does not split, and a product with one 15 # row, runs tensor-array's sequential code. The bits are those of tensor-array at depth 0, at every depth and thread 16 # count (tests/bits.bend). 17 # 18 # Not proved, and marked so: Array.fork and Array.join are @unsafe in Base (two handles to one array break the single 19 # owner that the proofs rely on), so the defs below that call them are @unsafe, and `bend main.bend` names every def 20 # that relies on them. Nothing checks that a leaf never writes A or B, or that two leaves' blocks of C are disjoint; 21 # the blocks are the halves of the fork tree, as in tensor-array's own split. tests/bits.bend checks both by sampling. 22 # All of it is Bend: no foreign code. 23 import Base 24 import bend-ml-tensor-array@0.1.18.0/main.bend as TA 25 26 # ---- blocks: any kernel over ranges of shared arrays ---- 27 # parN runs a kernel f over 2^d blocks of the range 0 .. cnt-1 in a tree of parallel lets; every block gets a handle to 28 # each of the N arrays (no copy) and the start and length of its range, plus e, a reusable record of whatever else the 29 # kernel needs (sizes, rates). The caller's kernel must write only inside its block's range, and a block's result may not 30 # depend on another's writes; then the arrays hold the same bits at every depth and thread count as f over the whole 31 # range at once. Nothing checks this (Array.fork and Array.join are @unsafe). 32 33 type B2<-A: Type, -B: Type> is Type: 34 B2{a: Array, b: Array} 35 36 type B3<-A: Type, -B: Type, -C: Type> is Type: 37 B3{a: Array, b: Array, c: Array} 38 39 type B4<-A: Type, -B: Type, -C: Type, -D: Type> is Type: 40 B4{a: Array, b: Array, c: Array, d: Array} 41 42 @unsafe 43 def blocks1(~E: Data, ~A: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> Array -> Array, +e: E, d: Nat, +i0: U32, +cnt: Nat, pa: Array & Array) -> Array: 44 match d pa: 45 case 0n Tuple{a, a2}: 46 f(e, i0, cnt, a) 47 case 1n++q Tuple{a, a2}: 48 l r = blocks1(~E, ~A, ~f, e, q, i0, Nat.div(cnt, 2n), Array.fork(A, a)) blocks1(~E, ~A, ~f, e, q, (i0 + U32.from_nat(Nat.div(cnt, 2n)) : U32), Nat.sub(cnt, Nat.div(cnt, 2n)), Array.fork(A, a2)) 49 Array.join(A, l, r) 50 51 @unsafe 52 def join2(-A: Type, -B: Type, l: B2, r: B2) -> B2: 53 match l r: 54 case B2{a, b} B2{a2, b2}: 55 B2{Array.join(A, a, a2), Array.join(B, b, b2)} 56 57 @unsafe 58 def blocks2(~E: Data, ~A: Type, ~B: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> B2 -> B2, +e: E, d: Nat, +i0: U32, +cnt: Nat, pa: Array & Array, pb: Array & Array) -> B2: 59 match d pa pb: 60 case 0n Tuple{a, a2} Tuple{b, b2}: 61 f(e, i0, cnt, B2{a, b}) 62 case 1n++q Tuple{a, a2} Tuple{b, b2}: 63 l r = blocks2(~E, ~A, ~B, ~f, e, q, i0, Nat.div(cnt, 2n), Array.fork(A, a), Array.fork(B, b)) blocks2(~E, ~A, ~B, ~f, e, q, (i0 + U32.from_nat(Nat.div(cnt, 2n)) : U32), Nat.sub(cnt, Nat.div(cnt, 2n)), Array.fork(A, a2), Array.fork(B, b2)) 64 join2(A, B, l, r) 65 66 @unsafe 67 def join3(-A: Type, -B: Type, -C: Type, l: B3, r: B3) -> B3: 68 match l r: 69 case B3{a, b, c} B3{a2, b2, c2}: 70 B3{Array.join(A, a, a2), Array.join(B, b, b2), Array.join(C, c, c2)} 71 72 @unsafe 73 def blocks3(~E: Data, ~A: Type, ~B: Type, ~C: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> B3 -> B3, +e: E, d: Nat, +i0: U32, +cnt: Nat, pa: Array & Array, pb: Array & Array, pc: Array & Array) -> B3: 74 match d pa pb pc: 75 case 0n Tuple{a, a2} Tuple{b, b2} Tuple{c, c2}: 76 f(e, i0, cnt, B3{a, b, c}) 77 case 1n++q Tuple{a, a2} Tuple{b, b2} Tuple{c, c2}: 78 l r = blocks3(~E, ~A, ~B, ~C, ~f, e, q, i0, Nat.div(cnt, 2n), Array.fork(A, a), Array.fork(B, b), Array.fork(C, c)) blocks3(~E, ~A, ~B, ~C, ~f, e, q, (i0 + U32.from_nat(Nat.div(cnt, 2n)) : U32), Nat.sub(cnt, Nat.div(cnt, 2n)), Array.fork(A, a2), Array.fork(B, b2), Array.fork(C, c2)) 79 join3(A, B, C, l, r) 80 81 @unsafe 82 def join4(-A: Type, -B: Type, -C: Type, -D: Type, l: B4, r: B4) -> B4: 83 match l r: 84 case B4{a, b, c, d} B4{a2, b2, c2, d2}: 85 B4{Array.join(A, a, a2), Array.join(B, b, b2), Array.join(C, c, c2), Array.join(D, d, d2)} 86 87 @unsafe 88 def blocks4(~E: Data, ~A: Type, ~B: Type, ~C: Type, ~D: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> B4 -> B4, +e: E, d: Nat, +i0: U32, +cnt: Nat, pa: Array & Array, pb: Array & Array, pc: Array & Array, pd: Array & Array) -> B4: 89 match d pa pb pc pd: 90 case 0n Tuple{a, a2} Tuple{b, b2} Tuple{c, c2} Tuple{x, x2}: 91 f(e, i0, cnt, B4{a, b, c, x}) 92 case 1n++q Tuple{a, a2} Tuple{b, b2} Tuple{c, c2} Tuple{x, x2}: 93 l r = blocks4(~E, ~A, ~B, ~C, ~D, ~f, e, q, i0, Nat.div(cnt, 2n), Array.fork(A, a), Array.fork(B, b), Array.fork(C, c), Array.fork(D, x)) blocks4(~E, ~A, ~B, ~C, ~D, ~f, e, q, (i0 + U32.from_nat(Nat.div(cnt, 2n)) : U32), Nat.sub(cnt, Nat.div(cnt, 2n)), Array.fork(A, a2), Array.fork(B, b2), Array.fork(C, c2), Array.fork(D, x2)) 94 join4(A, B, C, D, l, r) 95 96 # f over 2^d blocks of 0 .. cnt-1, sharing the arrays 97 @unsafe 98 def par1(~E: Data, ~A: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> Array -> Array, +e: E, d: Nat, +cnt: Nat, a: Array) -> Array: 99 blocks1(~E, ~A, ~f, e, d, 0, cnt, Array.fork(A, a)) 100 101 @unsafe 102 def par2(~E: Data, ~A: Type, ~B: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> B2 -> B2, +e: E, d: Nat, +cnt: Nat, a: Array, b: Array) -> B2: 103 blocks2(~E, ~A, ~B, ~f, e, d, 0, cnt, Array.fork(A, a), Array.fork(B, b)) 104 105 @unsafe 106 def par3(~E: Data, ~A: Type, ~B: Type, ~C: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> B3 -> B3, +e: E, d: Nat, +cnt: Nat, a: Array, b: Array, c: Array) -> B3: 107 blocks3(~E, ~A, ~B, ~C, ~f, e, d, 0, cnt, Array.fork(A, a), Array.fork(B, b), Array.fork(C, c)) 108 109 @unsafe 110 def par4(~E: Data, ~A: Type, ~B: Type, ~C: Type, ~D: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> B4 -> B4, +e: E, d: Nat, +cnt: Nat, a: Array, b: Array, c: Array, x: Array) -> B4: 111 blocks4(~E, ~A, ~B, ~C, ~D, ~f, e, d, 0, cnt, Array.fork(A, a), Array.fork(B, b), Array.fork(C, c), Array.fork(D, x)) 112 113 # the smallest d with 2^d >= n (TA.cap_depth's answer) by doubling a U32: log2(n) steps, where cap_depth compares Nats 114 # (a share of every product's time at 512 x 1536) 115 def depth_go(f: Nat, ok: Bool, +n: U32, +p: U32, d: Nat) -> Nat: 116 match f ok: 117 case _ True{}: 118 d 119 case 0n False{}: 120 d 121 case 1n+g False{}: 122 depth_go(g, U32.is_le(n, (p + p : U32)), n, (p + p : U32), 1n+d) 123 124 def depth(+n: Nat) -> Nat: 125 depth_go(32n, U32.is_le(U32.from_nat(n), 1), U32.from_nat(n), 1, 0n) 126 127 # ---- packing in parallel: tensor-array's pk_pack, its lines in 2^d blocks ---- 128 129 type PE is Data: 130 PE{tk: Nat, tku: U32, ku: U32, lstep: U32, st: U32} 131 132 def pk_b2(r: TA.PK) -> B2: 133 match r: 134 case TA.PK{a, t}: 135 B2{a, t} 136 137 # lines i0 .. i0+cnt-1 into their cells 138 def pk_leaf(+e: PE, +i0: U32, +cnt: Nat, r: B2) -> B2: 139 match e r: 140 case PE{+tk, +tku, +ku, +lstep, +st} B2{a, t}: 141 pk_b2(TA.pk_lines(cnt, tk, tku, ku, i0, i0, lstep, st, TA.PK{a, t})) 142 143 def pk_of(r: B2) -> TA.PK: 144 match r: 145 case B2{a, t}: 146 TA.PK{a, t} 147 148 # one block below 16384 cells (MNIST's 100 x 784 and 128 x 784 pack faster in one), else 2^d 149 def pk_d(big: Bool, d: Nat) -> Nat: 150 match big: 151 case True{}: 152 d 153 case False{}: 154 0n 155 156 @unsafe 157 def pk_par(d: Nat, +cnt: Nat, +tk: Nat, +k: Nat, +lstep: U32, +st: U32, a: Array) -> TA.PK: 158 pk_of(par2(~PE, ~F32, ~TA.T8, ~pk_leaf, PE{tk, U32.from_nat(tk), U32.from_nat(k), lstep, st}, pk_d(Nat.is_le(16384n, Nat.mul(cnt, tk)), d), cnt, a, Array.new(TA.T8, depth(Nat.mul(cnt, tk)), TA.t8_zero()))) 159 160 def tg_c(g: TA.TG) -> Array: 161 match g: 162 case TA.TG{a, b, c}: 163 c 164 165 # rows i0 .. i0+cnt-1 of C, every column 166 def rblock(+ac: Bool, cnt: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i0: U32, +mu: U32, a: Array, b: Array, c: Array) -> Array: 167 tg_c(TA.t8_rows(ac, cnt, m, tk, tku, rot, i0, mu, TA.TG{a, b, c})) 168 169 # columns j0 .. j0+cnt-1 of C, rows i.. (n more) 170 def cbrows(+ac: Bool, n: Nat, +cnt: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i: U32, +j0: U32, +mu: U32, g: TA.TG) -> TA.TG: 171 match n: 172 case 0n: 173 g 174 case 1n+p: 175 cbrows(ac, p, cnt, tk, tku, rot, (i + 1 : U32), j0, mu, TA.t8_cols(ac, cnt, tk, tku, rot, i, j0, ((i * mu : U32) + j0 : U32), g)) 176 177 # 2^d blocks of rows of C, halved at each node; each pair holds two handles to one block 178 @unsafe 179 def rpar(+ac: Bool, d: Nat, +cnt: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i0: U32, +mu: U32, pa: Array & Array, pb: Array & Array, pc: Array & Array) -> Array: 180 match d pa pb pc: 181 case 0n Tuple{a, a2} Tuple{b, b2} Tuple{c, c2}: 182 rblock(ac, cnt, m, tk, tku, rot, i0, mu, a, b, c) 183 case 1n++q Tuple{a, a2} Tuple{b, b2} Tuple{c, c2}: 184 l r = rpar(ac, q, Nat.div(cnt, 2n), m, tk, tku, rot, i0, mu, Array.fork(TA.T8, a), Array.fork(TA.T8, b), Array.fork(F32, c)) rpar(ac, q, Nat.sub(cnt, Nat.div(cnt, 2n)), m, tk, tku, rot, (i0 + U32.from_nat(Nat.div(cnt, 2n)) : U32), mu, Array.fork(TA.T8, a2), Array.fork(TA.T8, b2), Array.fork(F32, c2)) 185 Array.join(F32, l, r) 186 187 # 2^d blocks of columns of C 188 @unsafe 189 def cpar(+ac: Bool, d: Nat, +cnt: Nat, +n: Nat, +tk: Nat, +tku: U32, +rot: Nat, +j0: U32, +mu: U32, pa: Array & Array, pb: Array & Array, pc: Array & Array) -> Array: 190 match d pa pb pc: 191 case 0n Tuple{a, a2} Tuple{b, b2} Tuple{c, c2}: 192 tg_c(cbrows(ac, n, cnt, tk, tku, rot, 0, j0, mu, TA.TG{a, b, c})) 193 case 1n++q Tuple{a, a2} Tuple{b, b2} Tuple{c, c2}: 194 l r = cpar(ac, q, Nat.div(cnt, 2n), n, tk, tku, rot, j0, mu, Array.fork(TA.T8, a), Array.fork(TA.T8, b), Array.fork(F32, c)) cpar(ac, q, Nat.sub(cnt, Nat.div(cnt, 2n)), n, tk, tku, rot, (j0 + U32.from_nat(Nat.div(cnt, 2n)) : U32), mu, Array.fork(TA.T8, a2), Array.fork(TA.T8, b2), Array.fork(F32, c2)) 195 Array.join(F32, l, r) 196 197 # by rows when C has at least as many rows as columns, as tensor-array splits 198 @unsafe 199 def t8_tree(byrows: Bool, +ac: Bool, d: Nat, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, pa: Array & Array, pb: Array & Array, pc: Array & Array) -> Array: 200 match byrows: 201 case True{}: 202 rpar(ac, d, n, m, tk, tku, rot, 0, U32.from_nat(m), pa, pb, pc) 203 case False{}: 204 cpar(ac, d, m, n, tk, tku, rot, 0, U32.from_nat(m), pa, pb, pc) 205 206 # the packed copies are dropped after the tree; A and B (unpacked) go back to the caller 207 @unsafe 208 def t8_go(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array, p: TA.PKK) -> TA.G: 209 match p: 210 case TA.PKK{a, b, ta, tb}: 211 TA.G{a, b, t8_tree(Nat.is_le(m, n), ac, d, n, m, tk, tku, rot, Array.fork(TA.T8, ta), Array.fork(TA.T8, tb), Array.fork(F32, c))} 212 213 # C (n x m) = A · B (or C + A · B when ac), n >= 2, with 2^d blocks sharing A, B and C; in any layout of TA.gemm_any 214 # (arow, sa, sb, bcol); A and B come back intact 215 @unsafe 216 def gemm_fork(+ac: Bool, +d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> TA.G: 217 t8_go(ac, d, n, m, Nat.div(Nat.add(k, 7n), 8n), U32.from_nat(Nat.div(Nat.add(k, 7n), 8n)), Nat.mod(k, 4n), c, TA.t8_pk2(pk_par(d, n, Nat.div(Nat.add(k, 7n), 8n), k, arow, sa, a), pk_par(d, m, Nat.div(Nat.add(k, 7n), 8n), k, bcol, sb, b))) 218 219 # two or more rows: the shared tree when d > 0, tensor-array's sequential cells at d = 0 220 def gemm_rows(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> TA.G: 221 match d: 222 case 0n: 223 TA.gemm_t8_ac(ac, 0n, n, m, k, arow, sa, sb, bcol, a, b, c) 224 case 1n+q: 225 gemm_fork(ac, 1n+q, n, m, k, arow, sa, sb, bcol, a, b, c) 226 227 def gemm_any2(+ac: Bool, one: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> TA.G: 228 match one: 229 case True{}: 230 TA.gemm_ac(ac, n, m, k, arow, sa, sb, bcol, a, b, c) 231 case False{}: 232 gemm_rows(ac, d, n, m, k, arow, sa, sb, bcol, a, b, c) 233 234 # TA.gemm_any's arguments and result; one row runs the scalar loop at every depth, as in tensor-array 0.1.12.0 235 def gemm_any(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> TA.G: 236 gemm_any2(ac, Nat.is_le(n, 1n), d, n, m, k, arow, sa, sb, bcol, a, b, c) 237 238 # ---- the six products: tensor-array's signatures; par goes through TA.split_depth, as there ---- 239 240 # C = A · B (n x k) · (k x m) 241 def matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat) -> TA.MMul: 242 match a b: 243 case TA.Mat{x} TA.Mat{y}: 244 TA.mmul_fin(n, k, m, gemm_any(False{}, TA.split_depth(par, n, m, k), n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, x, y, Array.new(F32, depth(Nat.mul(n, m)), 0.0))) 245 246 # C + A · B, into C 247 def matmul_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat, c: TA.Mat) -> TA.MMul: 248 match a b c: 249 case TA.Mat{x} TA.Mat{y} TA.Mat{z}: 250 TA.mmul_fin(n, k, m, gemm_any(True{}, TA.split_depth(par, n, m, k), n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, x, y, z)) 251 252 # C = A · Bᵀ (n x k) · (m x k)ᵀ: weights with one row per output 253 def matmul_nt(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat) -> TA.MMulNT: 254 match a b: 255 case TA.Mat{x} TA.Mat{y}: 256 TA.mmul_nt_fin(n, k, m, gemm_any(False{}, TA.split_depth(par, n, m, k), n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), x, y, Array.new(F32, depth(Nat.mul(n, m)), 0.0))) 257 258 # C + A · Bᵀ, into C 259 def matmul_nt_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat, c: TA.Mat) -> TA.MMulNT: 260 match a b c: 261 case TA.Mat{x} TA.Mat{y} TA.Mat{z}: 262 TA.mmul_nt_fin(n, k, m, gemm_any(True{}, TA.split_depth(par, n, m, k), n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), x, y, z)) 263 264 # C = Aᵀ · B (k x n)ᵀ · (k x m): a weight gradient, Xᵀ · dY 265 def matmul_tn(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat) -> TA.MMulTN: 266 match a b: 267 case TA.Mat{x} TA.Mat{y}: 268 TA.mmul_tn_fin(n, k, m, gemm_any(False{}, TA.split_depth(par, n, m, k), n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, x, y, Array.new(F32, depth(Nat.mul(n, m)), 0.0))) 269 270 # C + Aᵀ · B, into C: dW += Xᵀ · dY 271 def matmul_tn_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat, c: TA.Mat) -> TA.MMulTN: 272 match a b c: 273 case TA.Mat{x} TA.Mat{y} TA.Mat{z}: 274 TA.mmul_tn_fin(n, k, m, gemm_any(True{}, TA.split_depth(par, n, m, k), n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, x, y, z)) 275 276 # ---- tensor-array 0.1.14.0's precise product with these blocks (opt-in) ---- 277 # TA.Mat.matmul_precise's composition, each block's product computed by gemm_any above: k in blocks of 1024, the first 278 # block's product written into C, each next one's added to it, so the same bits as tensor-array's (tests/bits.bend). 279 280 @unsafe 281 def mp_s3(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +w: Nat, +w32: U32, +m32: U32, a: Array, as: Array, c: Array, sb: TA.SL) -> TA.G: 282 match sb: 283 case TA.SL{b, bs}: TA.mp_s4(a, b, gemm_any(ac, TA.split_depth(par, n, m, w), n, m, w, w32, 1, m32, 1, as, bs, c)) 284 285 @unsafe 286 def mp_s2(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +w: Nat, +w32: U32, +m32: U32, +o: U32, c: Array, b: Array, sa: TA.SL) -> TA.G: 287 match sa: 288 case TA.SL{a, as}: mp_s3(ac, par, n, m, w, w32, m32, a, as, c, TA.sl_row(Nat.mul(w, m), (o * m32 : U32), 0, Array.new(F32, depth(Nat.mul(w, m)), 0.0), Array.get(F32, b, (o * m32 : U32)))) 289 290 @unsafe 291 def mp_step(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +k32: U32, +bi: U32, a: Array, b: Array, c: Array) -> TA.G: 292 mp_s2(ac, par, n, m, TA.blk_len(bi, k32), U32.from_nat(TA.blk_len(bi, k32)), U32.from_nat(m), (bi * 1024 : U32), c, b, TA.sl_rows(n, 0, TA.blk_len(bi, k32), U32.from_nat(TA.blk_len(bi, k32)), k32, (bi * 1024 : U32), TA.SL{a, Array.new(F32, depth(Nat.mul(n, TA.blk_len(bi, k32))), 0.0)})) 293 294 @unsafe 295 def mp_go(nb: Nat, +bi: U32, +par: Nat, +n: Nat, +m: Nat, +k32: U32, g: TA.G) -> TA.G: 296 match nb g: 297 case 0n g: g 298 case 1n++q TA.G{a, b, c}: mp_go(q, (bi + 1 : U32), par, n, m, k32, mp_step(U32.is_lt(0, bi), par, n, m, k32, bi, a, b, c)) 299 300 @unsafe 301 def mm_precise2(long: Bool, +n: Nat, +k: Nat, +m: Nat, +par: Nat, a: TA.Mat, b: TA.Mat) -> TA.MMul: 302 match long a b: 303 case True{} TA.Mat{x} TA.Mat{y}: TA.mmul_fin(n, k, m, mp_go(U32.to_nat(U32.div((U32.from_nat(k) + 1023 : U32), 1024)), 0, par, n, m, U32.from_nat(k), TA.G{x, y, Array.new(F32, depth(Nat.mul(n, m)), 0.0)})) 304 case False{} a b: matmul(n, k, m, par, a, b) 305 306 # C = A · B with k summed in blocks of 1024: TA.Mat.matmul_precise's bits (opt-in) 307 @unsafe 308 def matmul_precise(+n: Nat, +k: Nat, +m: Nat, +par: Nat, a: TA.Mat, b: TA.Mat) -> TA.MMul: 309 mm_precise2(Nat.is_lt(1024n, k), n, k, m, par, a, b)