main.bend
bend-ml-tensor-fork@0.1.4.0 · 309 lines · 0 definitions
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.16.0/main.bend as TA
5# import bend-ml-tensor-fork@0.1.4.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.
23import Base
24import bend-ml-tensor-array@0.1.16.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
33type B2<-A: Type, -B: Type> is Type:
34 B2{a: Array<A>, b: Array<B>}
35
36type B3<-A: Type, -B: Type, -C: Type> is Type:
37 B3{a: Array<A>, b: Array<B>, c: Array<C>}
38
39type B4<-A: Type, -B: Type, -C: Type, -D: Type> is Type:
40 B4{a: Array<A>, b: Array<B>, c: Array<C>, d: Array<D>}
41
42@unsafe
43def blocks1(~E: Data, ~A: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> Array<A> -> Array<A>, +e: E, d: Nat, +i0: U32, +cnt: Nat, pa: Array<A> & Array<A>) -> Array<A>:
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
52def join2(-A: Type, -B: Type, l: B2<A, B>, r: B2<A, B>) -> B2<A, B>:
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
58def blocks2(~E: Data, ~A: Type, ~B: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> B2<A, B> -> B2<A, B>, +e: E, d: Nat, +i0: U32, +cnt: Nat, pa: Array<A> & Array<A>, pb: Array<B> & Array<B>) -> B2<A, B>:
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
67def join3(-A: Type, -B: Type, -C: Type, l: B3<A, B, C>, r: B3<A, B, C>) -> B3<A, B, C>:
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
73def blocks3(~E: Data, ~A: Type, ~B: Type, ~C: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> B3<A, B, C> -> B3<A, B, C>, +e: E, d: Nat, +i0: U32, +cnt: Nat, pa: Array<A> & Array<A>, pb: Array<B> & Array<B>, pc: Array<C> & Array<C>) -> B3<A, B, C>:
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
82def join4(-A: Type, -B: Type, -C: Type, -D: Type, l: B4<A, B, C, D>, r: B4<A, B, C, D>) -> B4<A, B, C, D>:
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
88def blocks4(~E: Data, ~A: Type, ~B: Type, ~C: Type, ~D: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> B4<A, B, C, D> -> B4<A, B, C, D>, +e: E, d: Nat, +i0: U32, +cnt: Nat, pa: Array<A> & Array<A>, pb: Array<B> & Array<B>, pc: Array<C> & Array<C>, pd: Array<D> & Array<D>) -> B4<A, B, C, D>:
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
98def par1(~E: Data, ~A: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> Array<A> -> Array<A>, +e: E, d: Nat, +cnt: Nat, a: Array<A>) -> Array<A>:
99 blocks1(~E, ~A, ~f, e, d, 0, cnt, Array.fork(A, a))
100
101@unsafe
102def par2(~E: Data, ~A: Type, ~B: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> B2<A, B> -> B2<A, B>, +e: E, d: Nat, +cnt: Nat, a: Array<A>, b: Array<B>) -> B2<A, B>:
103 blocks2(~E, ~A, ~B, ~f, e, d, 0, cnt, Array.fork(A, a), Array.fork(B, b))
104
105@unsafe
106def par3(~E: Data, ~A: Type, ~B: Type, ~C: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> B3<A, B, C> -> B3<A, B, C>, +e: E, d: Nat, +cnt: Nat, a: Array<A>, b: Array<B>, c: Array<C>) -> B3<A, B, C>:
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
110def par4(~E: Data, ~A: Type, ~B: Type, ~C: Type, ~D: Type, ~f: @+e: E -> @+i0: U32 -> @+cnt: Nat -> B4<A, B, C, D> -> B4<A, B, C, D>, +e: E, d: Nat, +cnt: Nat, a: Array<A>, b: Array<B>, c: Array<C>, x: Array<D>) -> B4<A, B, C, D>:
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)
115def 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
124def 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
129type PE is Data:
130 PE{tk: Nat, tku: U32, ku: U32, lstep: U32, st: U32}
131
132def pk_b2(r: TA.PK) -> B2<F32, TA.T8>:
133 match r:
134 case TA.PK{a, t}:
135 B2{a, t}
136
137# lines i0 .. i0+cnt-1 into their cells
138def pk_leaf(+e: PE, +i0: U32, +cnt: Nat, r: B2<F32, TA.T8>) -> B2<F32, TA.T8>:
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
143def pk_of(r: B2<F32, TA.T8>) -> 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
149def pk_d(big: Bool, d: Nat) -> Nat:
150 match big:
151 case True{}:
152 d
153 case False{}:
154 0n
155
156@unsafe
157def pk_par(d: Nat, +cnt: Nat, +tk: Nat, +k: Nat, +lstep: U32, +st: U32, a: Array<F32>) -> 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
160def tg_c(g: TA.TG) -> Array<F32>:
161 match g:
162 case TA.TG{a, b, c}:
163 c
164
165# rows i0 .. i0+cnt-1 of C, every column
166def 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>:
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)
170def 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
179def 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>:
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
189def 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>:
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
199def 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>:
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
208def t8_go(+ac: Bool, d: Nat, +n: Nat, +m: Nat, +tk: Nat, +tku: U32, +rot: Nat, c: Array<F32>, 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
216def 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:
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
220def 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:
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
227def 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:
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
235def 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:
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)
241def matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat<n, k>, b: TA.Mat<k, m>) -> TA.MMul<n, k, m>:
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
247def 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>:
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
253def 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>:
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
259def 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>:
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
265def 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>:
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
271def 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>:
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
281def mp_s3(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +w: Nat, +w32: U32, +m32: U32, a: Array<F32>, as: Array<F32>, c: Array<F32>, 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
286def mp_s2(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +w: Nat, +w32: U32, +m32: U32, +o: U32, c: Array<F32>, b: Array<F32>, 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
291def mp_step(+ac: Bool, +par: Nat, +n: Nat, +m: Nat, +k32: U32, +bi: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> 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
295def 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
301def mm_precise2(long: Bool, +n: Nat, +k: Nat, +m: Nat, +par: Nat, a: TA.Mat<n, k>, b: TA.Mat<k, m>) -> TA.MMul<n, k, m>:
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
308def matmul_precise(+n: Nat, +k: Nat, +m: Nat, +par: Nat, a: TA.Mat<n, k>, b: TA.Mat<k, m>) -> TA.MMul<n, k, m>:
309 mm_precise2(Nat.is_lt(1024n, k), n, k, m, par, a, b)