1 # bend-ml-tensor-fork: bend-ml-tensor-array's products, with the parallel blocks sharing their operands. 2 # 3 # import bend-ml-tensor-array@0.1.12.0/main.bend as TA 4 # import bend-ml-tensor-fork@0.1.0.0/main.bend as TF 5 # TF.matmul(n, k, m, par, a, b) # the same arguments, result type and bits as TA.Mat.matmul 6 # 7 # The six products (matmul, matmul_nt, matmul_tn and their accumulating forms) take tensor-array's Mat values and 8 # return its result types. A product that splits (tensor-array's split_depth, two or more rows) packs A and B once into 9 # cells of 8, as the sequential path does, and every node of the fork tree shares the two packed operands and C through 10 # Base's Array.fork: two handles to one block, no copy. Each leaf runs tensor-array's cells (t8_rows, t8_cols) straight 11 # into its own rows (or columns) of C. tensor-array's parallel path instead clones an operand at every node and hands 12 # 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 13 # row, runs tensor-array's sequential code. The bits are those of tensor-array at depth 0, at every depth and thread 14 # count (tests/bits.bend). 15 # 16 # Not proved, and marked so: Array.fork and Array.join are @unsafe in Base (two handles to one array break the single 17 # owner that the proofs rely on), so the defs below that call them are @unsafe, and `bend main.bend` names every def 18 # that relies on them. Nothing checks that a leaf never writes A or B, or that two leaves' blocks of C are disjoint; 19 # the blocks are the halves of the fork tree, as in tensor-array's own split. tests/bits.bend checks both by sampling. 20 # All of it is Bend: no foreign code. 21 import Base 22 import bend-ml-tensor-array@0.1.12.0/main.bend as TA 23 24 def tg_c(g: TA.TG) -> Array: 25 match g: 26 case TA.TG{a, b, c}: 27 c 28 29 # rows i0 .. i0+cnt-1 of C, every column 30 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: 31 tg_c(TA.t8_rows(ac, cnt, m, tk, tku, rot, i0, mu, TA.TG{a, b, c})) 32 33 # columns j0 .. j0+cnt-1 of C, rows i.. (n more) 34 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: 35 match n: 36 case 0n: 37 g 38 case 1n+p: 39 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)) 40 41 # 2^d blocks of rows of C, halved at each node; each pair holds two handles to one block 42 @unsafe 43 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: 44 match d pa pb pc: 45 case 0n Tuple{a, a2} Tuple{b, b2} Tuple{c, c2}: 46 rblock(ac, cnt, m, tk, tku, rot, i0, mu, a, b, c) 47 case 1n++q Tuple{a, a2} Tuple{b, b2} Tuple{c, c2}: 48 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)) 49 Array.join(F32, l, r) 50 51 # 2^d blocks of columns of C 52 @unsafe 53 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: 54 match d pa pb pc: 55 case 0n Tuple{a, a2} Tuple{b, b2} Tuple{c, c2}: 56 tg_c(cbrows(ac, n, cnt, tk, tku, rot, 0, j0, mu, TA.TG{a, b, c})) 57 case 1n++q Tuple{a, a2} Tuple{b, b2} Tuple{c, c2}: 58 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)) 59 Array.join(F32, l, r) 60 61 # by rows when C has at least as many rows as columns, as tensor-array splits 62 @unsafe 63 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: 64 match byrows: 65 case True{}: 66 rpar(ac, d, n, m, tk, tku, rot, 0, U32.from_nat(m), pa, pb, pc) 67 case False{}: 68 cpar(ac, d, m, n, tk, tku, rot, 0, U32.from_nat(m), pa, pb, pc) 69 70 # the packed copies are dropped after the tree; A and B (unpacked) go back to the caller 71 @unsafe 72 def t8_go(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array, p: TA.PKK) -> TA.G: 73 match p: 74 case TA.PKK{a, b, ta, tb}: 75 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))} 76 77 # 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 78 # (arow, sa, sb, bcol); A and B come back intact 79 @unsafe 80 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: 81 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(TA.pk_pack(n, Nat.div(Nat.add(k, 7n), 8n), k, arow, sa, a), TA.pk_pack(m, Nat.div(Nat.add(k, 7n), 8n), k, bcol, sb, b))) 82 83 # two or more rows: the shared tree when d > 0, tensor-array's sequential cells at d = 0 84 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: 85 match d: 86 case 0n: 87 TA.gemm_t8_ac(ac, 0n, n, m, k, arow, sa, sb, bcol, a, b, c) 88 case 1n+q: 89 gemm_fork(ac, 1n+q, n, m, k, arow, sa, sb, bcol, a, b, c) 90 91 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: 92 match one: 93 case True{}: 94 TA.gemm_ac(ac, n, m, k, arow, sa, sb, bcol, a, b, c) 95 case False{}: 96 gemm_rows(ac, d, n, m, k, arow, sa, sb, bcol, a, b, c) 97 98 # TA.gemm_any's arguments and result; one row runs the scalar loop at every depth, as in tensor-array 0.1.12.0 99 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: 100 gemm_any2(ac, Nat.is_le(n, 1n), d, n, m, k, arow, sa, sb, bcol, a, b, c) 101 102 # ---- the six products: tensor-array's signatures; par goes through TA.split_depth, as there ---- 103 104 # C = A · B (n x k) · (k x m) 105 def matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat) -> TA.MMul: 106 match a b: 107 case TA.Mat{x} TA.Mat{y}: 108 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, TA.cap_depth(Nat.mul(n, m)), 0.0))) 109 110 # C + A · B, into C 111 def matmul_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat, c: TA.Mat) -> TA.MMul: 112 match a b c: 113 case TA.Mat{x} TA.Mat{y} TA.Mat{z}: 114 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)) 115 116 # C = A · Bᵀ (n x k) · (m x k)ᵀ: weights with one row per output 117 def matmul_nt(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat) -> TA.MMulNT: 118 match a b: 119 case TA.Mat{x} TA.Mat{y}: 120 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, TA.cap_depth(Nat.mul(n, m)), 0.0))) 121 122 # C + A · Bᵀ, into C 123 def matmul_nt_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat, c: TA.Mat) -> TA.MMulNT: 124 match a b c: 125 case TA.Mat{x} TA.Mat{y} TA.Mat{z}: 126 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)) 127 128 # C = Aᵀ · B (k x n)ᵀ · (k x m): a weight gradient, Xᵀ · dY 129 def matmul_tn(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat) -> TA.MMulTN: 130 match a b: 131 case TA.Mat{x} TA.Mat{y}: 132 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, TA.cap_depth(Nat.mul(n, m)), 0.0))) 133 134 # C + Aᵀ · B, into C: dW += Xᵀ · dY 135 def matmul_tn_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat, c: TA.Mat) -> TA.MMulTN: 136 match a b c: 137 case TA.Mat{x} TA.Mat{y} TA.Mat{z}: 138 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))