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.19.0/main.bend as TA 5 # import bend-ml-tensor-fork@0.1.6.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.19.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) 310 311 # ---- softmax_ce with the rows in blocks (TA.Mat.softmax_ce's bits) ---- 312 # tensor-array walks the rows one after another (cea_rows), adding each row's loss term to the running loss. Here the 313 # rows run in 2^d blocks through par3, sharing the logits (each block writes only its own rows), the labels (read) and 314 # an array t of one loss term per row (each block writes only its own cells). The loss is then summed from t in row 315 # order, from +0.0, with the same subtraction as cea_rows (loss - log(max(p, 1e-12))), so the bits are the same at 316 # every depth and thread count. The rows are tensor-array's own loops (cea_max, cea_exp, cea_p, cea_grad). 317 # Labels past the n rows (a list longer than the matrix) run after the blocks through TA.cea_rows, in order, as in 318 # tensor-array: those rows reach past the matrix and may wrap around onto its first rows, so they cannot share a 319 # block. A matrix with no columns, or par = 0, runs TA.Mat.softmax_ce itself. 320 321 type SE is Data: 322 SE{mm1: Nat, mu: U32, nf: F32} 323 324 # a row's loss term log(max(p, 1e-12)) and its gradient: cea_row3..cea_row0 without the running loss 325 def sr3(+mm1: Nat, +o: U32, +label: Nat, +s: F32, +nf: F32, q: TA.AF) -> TA.L2: 326 match q: 327 case TA.AF{c, +p}: 328 TA.L2{F32.log(F32.max(p, 0.000000000001)), TA.cea_grad(mm1, o, 0n, label, s, nf, Array.get(F32, c, o))} 329 330 def sr2(+mm1: Nat, +o: U32, +label: Nat, +nf: F32, q: TA.AF) -> TA.L2: 331 match q: 332 case TA.AF{c, +s}: 333 sr3(mm1, o, label, s, nf, TA.cea_p(Nat.is_le(label, mm1), o, label, s, c)) 334 335 def sr1(+mm1: Nat, +o: U32, +label: Nat, +nf: F32, q: TA.AF) -> TA.L2: 336 match q: 337 case TA.AF{c, +m}: 338 sr2(mm1, o, label, nf, TA.cea_exp(mm1, o, m, 0.0, Array.get(F32, c, o))) 339 340 def sr0(+mm1: Nat, +o: U32, +label: Nat, +nf: F32, r: Array & F32) -> TA.L2: 341 match r: 342 case Tuple{c, +x0}: 343 sr1(mm1, o, label, nf, TA.cea_max(mm1, o, x0, Array.get(F32, c, o))) 344 345 def ce_put(+i: U32, lab: Array, t: Array, q: TA.L2) -> B3: 346 match q: 347 case TA.L2{+x, z}: 348 B3{z, lab, Array.set(F32, t, i, x)} 349 350 def ce_row(+mm1: Nat, +mu: U32, +nf: F32, +i: U32, z: Array, t: Array, r: Array & U32) -> B3: 351 match r: 352 case Tuple{lab, +y}: 353 ce_put(i, lab, t, sr0(mm1, (i * mu : U32), U32.to_nat(y), nf, Array.get(F32, z, (i * mu : U32)))) 354 355 def ce_step(+mm1: Nat, +mu: U32, +nf: F32, +i: U32, g: B3) -> B3: 356 match g: 357 case B3{z, lab, t}: 358 ce_row(mm1, mu, nf, i, z, t, Array.get(U32, lab, i)) 359 360 def ce_loop(k: Nat, +mm1: Nat, +mu: U32, +nf: F32, +i: U32, g: B3) -> B3: 361 match k: 362 case 0n: 363 g 364 case 1n+p: 365 ce_loop(p, mm1, mu, nf, (i + 1 : U32), ce_step(mm1, mu, nf, i, g)) 366 367 # rows i0 .. i0+cnt-1: their gradients in the logits, their loss terms in t 368 def ce_leaf(+e: SE, +i0: U32, +cnt: Nat, g: B3) -> B3: 369 match e: 370 case SE{+mm1, +mu, +nf}: 371 ce_loop(cnt, mm1, mu, nf, i0, g) 372 373 # the loss over k terms, in row order from +0.0, as cea_rows adds them 374 def ce_sum(k: Nat, +i: U32, +loss: F32, r: Array & F32) -> F32: 375 match k r: 376 case 0n Tuple{t, x}: 377 loss 378 case 1n+p Tuple{t, +x}: 379 ce_sum(p, (i + 1 : U32), (loss - x : F32), Array.get(F32, t, (i + 1 : U32))) 380 381 type CL is Type: 382 CL{a: Array, k: Nat, rest: List<&2, Nat>} 383 384 # a label past the row gives p = 0 and no one-hot cell, as c does: kept as c so that it fits a U32 385 def lab_fit(inside: Bool, +y: Nat, +c: Nat) -> Nat: 386 match inside: 387 case True{}: 388 y 389 case False{}: 390 c 391 392 # up to n labels into a, counted in k; the rest stays a list 393 def lab_take(n: Nat, labels: List<&2, Nat>, +mm1: Nat, +c: Nat, +i: U32, +k: Nat, a: Array) -> CL: 394 match n labels: 395 case 0n ls: 396 CL{a, k, ls} 397 case 1n+p Nil{}: 398 CL{a, k, Nil{}} 399 case 1n+p Con{+y, t}: 400 lab_take(p, t, mm1, c, (i + 1 : U32), 1n+k, Array.set(U32, a, i, U32.from_nat(lab_fit(Nat.is_le(y, mm1), y, c)))) 401 402 # after the blocks: the loss from t, then any labels past the n rows, in order, then tensor-array's mean 403 def ce_fin2(-n: Nat, -c: Nat, +nf: F32, +mm1: Nat, +mu: U32, +k: Nat, rest: List<&2, Nat>, g: B3) -> TA.MCE: 404 match g: 405 case B3{z, lab, t}: 406 TA.ce_fin(n, c, nf, TA.cea_rows(rest, mm1, mu, nf, U32.from_nat(k), TA.L2{ce_sum(k, 0, 0.0, Array.get(F32, t, 0)), z})) 407 408 @unsafe 409 def ce_par3(-n: Nat, -c: Nat, +nf: F32, +mm1: Nat, +mu: U32, +d: Nat, z: Array, lt: CL) -> TA.MCE: 410 match lt: 411 case CL{lab, +k, rest}: 412 ce_fin2(n, c, nf, mm1, mu, k, rest, par3(~SE, ~F32, ~U32, ~F32, ~ce_leaf, SE{mm1, mu, nf}, Nat.min(d, depth(k)), k, z, lab, Array.new(F32, depth(k), 0.0))) 413 414 def softmax_ce2(seq: Bool, +n: Nat, +c: Nat, +d: Nat, z: TA.Mat, labels: List<&2, Nat>) -> TA.MCE: 415 match seq z: 416 case True{} z: 417 TA.Mat.softmax_ce(n, c, z, labels) 418 case False{} TA.Mat{x}: 419 ce_par3(n, c, F32.from_nat(n), Nat.sub(c, 1n), U32.from_nat(c), d, x, lab_take(n, labels, Nat.sub(c, 1n), c, 0, 0n, Array.new(U32, depth(n), 0))) 420 421 # mean cross-entropy and its gradient, TA.Mat.softmax_ce's arguments and bits, with the rows in up to 2^par blocks 422 # (par = 0: tensor-array's own loop) 423 def softmax_ce(+n: Nat, +c: Nat, +par: Nat, z: TA.Mat, labels: List<&2, Nat>) -> TA.MCE: 424 softmax_ce2(Bool.or(Nat.is_eq(par, 0n), Nat.is_eq(c, 0n)), n, c, par, z, labels)