main.bend
bend-ml-tensor-fork@0.1.0.0 · 138 lines · 0 definitions
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.
21import Base
22import bend-ml-tensor-array@0.1.12.0/main.bend as TA
23
24def tg_c(g: TA.TG) -> Array<F32>:
25 match g:
26 case TA.TG{a, b, c}:
27 c
28
29# rows i0 .. i0+cnt-1 of C, every column
30def rblock(+ac: Bool, cnt: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i0: U32, +mu: U32, a: Array<TA.T8>, b: Array<TA.T8>, c: Array<F32>) -> Array<F32>:
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)
34def 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
43def rpar(+ac: Bool, d: Nat, +cnt: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, +i0: U32, +mu: U32, pa: Array<TA.T8> & Array<TA.T8>, pb: Array<TA.T8> & Array<TA.T8>, pc: Array<F32> & Array<F32>) -> Array<F32>:
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
53def cpar(+ac: Bool, d: Nat, +cnt: Nat, +n: Nat, +tk: Nat, +tku: U32, +rot: Nat, +j0: U32, +mu: U32, pa: Array<TA.T8> & Array<TA.T8>, pb: Array<TA.T8> & Array<TA.T8>, pc: Array<F32> & Array<F32>) -> Array<F32>:
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
63def t8_tree(byrows: Bool, +ac: Bool, d: Nat, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, pa: Array<TA.T8> & Array<TA.T8>, pb: Array<TA.T8> & Array<TA.T8>, pc: Array<F32> & Array<F32>) -> Array<F32>:
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
72def t8_go(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array<F32>, 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
80def gemm_fork(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> 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
84def gemm_rows(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> 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
91def gemm_any2(+ac: Bool, one: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> 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
99def gemm_any(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> 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)
105def matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat<n, k>, b: TA.Mat<k, m>) -> TA.MMul<n, k, m>:
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
111def matmul_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat<n, k>, b: TA.Mat<k, m>, c: TA.Mat<n, m>) -> TA.MMul<n, k, m>:
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
117def matmul_nt(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat<n, k>, b: TA.Mat<m, k>) -> TA.MMulNT<n, k, m>:
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
123def matmul_nt_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat<n, k>, b: TA.Mat<m, k>, c: TA.Mat<n, m>) -> TA.MMulNT<n, k, m>:
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
129def matmul_tn(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat<k, n>, b: TA.Mat<k, m>) -> TA.MMulTN<n, k, m>:
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
135def matmul_tn_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat<k, n>, b: TA.Mat<k, m>, c: TA.Mat<n, m>) -> TA.MMulTN<n, k, m>:
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))