main.bend

bend-ml-tensor-fork@0.1.6.0 · 424 lines · 63 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.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.
23import Base
24import 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
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)
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
321type 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
325def 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
330def 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
335def 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
340def sr0(+mm1: Nat, +o: U32, +label: Nat, +nf: F32, r: Array<F32> & 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
345def ce_put(+i: U32, lab: Array<U32>, t: Array<F32>, q: TA.L2) -> B3<F32, U32, F32>:
346  match q:
347    case TA.L2{+x, z}:
348      B3{z, lab, Array.set(F32, t, i, x)}
349
350def ce_row(+mm1: Nat, +mu: U32, +nf: F32, +i: U32, z: Array<F32>, t: Array<F32>, r: Array<U32> & U32) -> B3<F32, U32, F32>:
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
355def ce_step(+mm1: Nat, +mu: U32, +nf: F32, +i: U32, g: B3<F32, U32, F32>) -> B3<F32, U32, F32>:
356  match g:
357    case B3{z, lab, t}:
358      ce_row(mm1, mu, nf, i, z, t, Array.get(U32, lab, i))
359
360def ce_loop(k: Nat, +mm1: Nat, +mu: U32, +nf: F32, +i: U32, g: B3<F32, U32, F32>) -> B3<F32, U32, F32>:
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
368def ce_leaf(+e: SE, +i0: U32, +cnt: Nat, g: B3<F32, U32, F32>) -> B3<F32, U32, F32>:
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
374def ce_sum(k: Nat, +i: U32, +loss: F32, r: Array<F32> & 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
381type CL is Type:
382  CL{a: Array<U32>, 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
385def 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
393def lab_take(n: Nat, labels: List<&2, Nat>, +mm1: Nat, +c: Nat, +i: U32, +k: Nat, a: Array<U32>) -> 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
403def ce_fin2(-n: Nat, -c: Nat, +nf: F32, +mm1: Nat, +mu: U32, +k: Nat, rest: List<&2, Nat>, g: B3<F32, U32, F32>) -> TA.MCE<n, c>:
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
409def ce_par3(-n: Nat, -c: Nat, +nf: F32, +mm1: Nat, +mu: U32, +d: Nat, z: Array<F32>, lt: CL) -> TA.MCE<n, c>:
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
414def softmax_ce2(seq: Bool, +n: Nat, +c: Nat, +d: Nat, z: TA.Mat<n, c>, labels: List<&2, Nat>) -> TA.MCE<n, c>:
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)
423def softmax_ce(+n: Nat, +c: Nat, +par: Nat, z: TA.Mat<n, c>, labels: List<&2, Nat>) -> TA.MCE<n, c>:
424  softmax_ce2(Bool.or(Nat.is_eq(par, 0n), Nat.is_eq(c, 0n)), n, c, par, z, labels)