main.bend
bend-ml-tensor-gpu@0.1.3.0 · 632 lines · 0 definitions
1# bend-ml-tensor-gpu: tensor-array's products on the GPU, with tensor-array's bits.
2#
3# import ../tensor-array/main.bend as TA
4# import bend-ml-tensor-gpu@0.1.3.0/main.bend as TG
5# TG.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 return
8# its result types; par is ignored (the GPU decides the split). Each product is one `!` call (a bang): a balanced fork
9# tree whose leaves each compute a 4 x 2 tile of C, every node sharing A, B and C through Base's Array.fork (two
10# handles to one block, no copy), as upstream's one-block pattern does (bendlang/bend tests/run/stencil3d.bend). On a
11# machine without a GPU, or with `--gpu off`, the same bang runs on the CPU threads.
12#
13# The bits: a cell of C is computed in tensor-array's order. Product p of a dot goes to partial sum p mod 8, in
14# increasing p, each sum starting at +0.0, and the eight sums are combined by TA.t8_comb with rot = k mod 8, as
15# tensor-array's dotr and t8_dot do; an accumulating product adds the cell to C once with TA.cwr. The loop takes eight
16# steps of k a turn; the steps past k in the last turn add +0.0, which leaves a sum unchanged (a sum that starts at +0.0
17# is never -0.0). Bend builds the device code with --fmad=false, so no multiply-add is fused. tests/bits.bend checks the
18# bits against tensor-array on many shapes, on the GPU and the CPU.
19#
20# Not proved, and marked so: Array.fork and Array.join are @unsafe in Base, so the defs below that call them, or take
21# a shared array, are @unsafe, and `bend main.bend` names every def that relies on them. Nothing checks that a leaf
22# writes only its own tile of C, or never writes A or B; the tiles are disjoint by construction (leaf t writes tile t).
23# All of it is Bend: no foreign code. Written by bend-ml's scripts/gen_tensor_gpu.py (4 x 2 tiles).
24import Base
25import bend-ml-tensor-array@0.1.15.0/main.bend as TA
26
27# the sizes and strides of one product: C[i][j] = sum_p A[oa + i*arow + p*sa] * B[ob + p*sb + j*bcol], n x m, k steps
28# (oa, ob: 0 except for matmul_precise's blocks)
29type E is Data:
30 E{n: U32, m: U32, k: U32, tcols: U32, tiles: U32, arow: U32, sa: U32, sb: U32, bcol: U32, kg: Nat, rot: Nat, rem: U32, ac: Bool, oa: U32, ob: U32}
31
32# a product that counts only when its step is inside k (the last turn's steps past k add +0.0)
33def pk(ok: Bool, +x: F32) -> F32:
34 match ok:
35 case True{}: x
36 case False{}: 0.0
37
38# the values of one turn: A's 4 rows and B's 2 columns at 8 consecutive steps, read through the arrays in order
39type GR is Type:
40 GR{a: Array<F32>, b: Array<F32>, x0_0: F32, x0_1: F32, x0_2: F32, x0_3: F32, x1_0: F32, x1_1: F32, x1_2: F32, x1_3: F32, x2_0: F32, x2_1: F32, x2_2: F32, x2_3: F32, x3_0: F32, x3_1: F32, x3_2: F32, x3_3: F32, x4_0: F32, x4_1: F32, x4_2: F32, x4_3: F32, x5_0: F32, x5_1: F32, x5_2: F32, x5_3: F32, x6_0: F32, x6_1: F32, x6_2: F32, x6_3: F32, x7_0: F32, x7_1: F32, x7_2: F32, x7_3: F32, y0_0: F32, y0_1: F32, y1_0: F32, y1_1: F32, y2_0: F32, y2_1: F32, y3_0: F32, y3_1: F32, y4_0: F32, y4_1: F32, y5_0: F32, y5_1: F32, y6_0: F32, y6_1: F32, y7_0: F32, y7_1: F32}
41
42def gr_47(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, +y1_0: F32, +y1_1: F32, +y2_0: F32, +y2_1: F32, +y3_0: F32, +y3_1: F32, +y4_0: F32, +y4_1: F32, +y5_0: F32, +y5_1: F32, +y6_0: F32, +y6_1: F32, +y7_0: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
43 match rv:
44 case Tuple{b, +y7_1}:
45 GR{a, b, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, y1_1, y2_0, y2_1, y3_0, y3_1, y4_0, y4_1, y5_0, y5_1, y6_0, y6_1, y7_0, y7_1}
46def gr_46(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, +y1_0: F32, +y1_1: F32, +y2_0: F32, +y2_1: F32, +y3_0: F32, +y3_1: F32, +y4_0: F32, +y4_1: F32, +y5_0: F32, +y5_1: F32, +y6_0: F32, +y6_1: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
47 match rv:
48 case Tuple{b, +y7_0}:
49 gr_47(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, y1_1, y2_0, y2_1, y3_0, y3_1, y4_0, y4_1, y5_0, y5_1, y6_0, y6_1, y7_0, a, Array.get(F32, b, (ib + ((sb * 7 : U32) + (bcol * 1 : U32) : U32) : U32)))
50def gr_45(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, +y1_0: F32, +y1_1: F32, +y2_0: F32, +y2_1: F32, +y3_0: F32, +y3_1: F32, +y4_0: F32, +y4_1: F32, +y5_0: F32, +y5_1: F32, +y6_0: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
51 match rv:
52 case Tuple{b, +y6_1}:
53 gr_46(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, y1_1, y2_0, y2_1, y3_0, y3_1, y4_0, y4_1, y5_0, y5_1, y6_0, y6_1, a, Array.get(F32, b, (ib + ((sb * 7 : U32) + (bcol * 0 : U32) : U32) : U32)))
54def gr_44(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, +y1_0: F32, +y1_1: F32, +y2_0: F32, +y2_1: F32, +y3_0: F32, +y3_1: F32, +y4_0: F32, +y4_1: F32, +y5_0: F32, +y5_1: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
55 match rv:
56 case Tuple{b, +y6_0}:
57 gr_45(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, y1_1, y2_0, y2_1, y3_0, y3_1, y4_0, y4_1, y5_0, y5_1, y6_0, a, Array.get(F32, b, (ib + ((sb * 6 : U32) + (bcol * 1 : U32) : U32) : U32)))
58def gr_43(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, +y1_0: F32, +y1_1: F32, +y2_0: F32, +y2_1: F32, +y3_0: F32, +y3_1: F32, +y4_0: F32, +y4_1: F32, +y5_0: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
59 match rv:
60 case Tuple{b, +y5_1}:
61 gr_44(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, y1_1, y2_0, y2_1, y3_0, y3_1, y4_0, y4_1, y5_0, y5_1, a, Array.get(F32, b, (ib + ((sb * 6 : U32) + (bcol * 0 : U32) : U32) : U32)))
62def gr_42(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, +y1_0: F32, +y1_1: F32, +y2_0: F32, +y2_1: F32, +y3_0: F32, +y3_1: F32, +y4_0: F32, +y4_1: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
63 match rv:
64 case Tuple{b, +y5_0}:
65 gr_43(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, y1_1, y2_0, y2_1, y3_0, y3_1, y4_0, y4_1, y5_0, a, Array.get(F32, b, (ib + ((sb * 5 : U32) + (bcol * 1 : U32) : U32) : U32)))
66def gr_41(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, +y1_0: F32, +y1_1: F32, +y2_0: F32, +y2_1: F32, +y3_0: F32, +y3_1: F32, +y4_0: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
67 match rv:
68 case Tuple{b, +y4_1}:
69 gr_42(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, y1_1, y2_0, y2_1, y3_0, y3_1, y4_0, y4_1, a, Array.get(F32, b, (ib + ((sb * 5 : U32) + (bcol * 0 : U32) : U32) : U32)))
70def gr_40(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, +y1_0: F32, +y1_1: F32, +y2_0: F32, +y2_1: F32, +y3_0: F32, +y3_1: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
71 match rv:
72 case Tuple{b, +y4_0}:
73 gr_41(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, y1_1, y2_0, y2_1, y3_0, y3_1, y4_0, a, Array.get(F32, b, (ib + ((sb * 4 : U32) + (bcol * 1 : U32) : U32) : U32)))
74def gr_39(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, +y1_0: F32, +y1_1: F32, +y2_0: F32, +y2_1: F32, +y3_0: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
75 match rv:
76 case Tuple{b, +y3_1}:
77 gr_40(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, y1_1, y2_0, y2_1, y3_0, y3_1, a, Array.get(F32, b, (ib + ((sb * 4 : U32) + (bcol * 0 : U32) : U32) : U32)))
78def gr_38(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, +y1_0: F32, +y1_1: F32, +y2_0: F32, +y2_1: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
79 match rv:
80 case Tuple{b, +y3_0}:
81 gr_39(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, y1_1, y2_0, y2_1, y3_0, a, Array.get(F32, b, (ib + ((sb * 3 : U32) + (bcol * 1 : U32) : U32) : U32)))
82def gr_37(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, +y1_0: F32, +y1_1: F32, +y2_0: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
83 match rv:
84 case Tuple{b, +y2_1}:
85 gr_38(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, y1_1, y2_0, y2_1, a, Array.get(F32, b, (ib + ((sb * 3 : U32) + (bcol * 0 : U32) : U32) : U32)))
86def gr_36(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, +y1_0: F32, +y1_1: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
87 match rv:
88 case Tuple{b, +y2_0}:
89 gr_37(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, y1_1, y2_0, a, Array.get(F32, b, (ib + ((sb * 2 : U32) + (bcol * 1 : U32) : U32) : U32)))
90def gr_35(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, +y1_0: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
91 match rv:
92 case Tuple{b, +y1_1}:
93 gr_36(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, y1_1, a, Array.get(F32, b, (ib + ((sb * 2 : U32) + (bcol * 0 : U32) : U32) : U32)))
94def gr_34(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, +y0_1: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
95 match rv:
96 case Tuple{b, +y1_0}:
97 gr_35(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, y1_0, a, Array.get(F32, b, (ib + ((sb * 1 : U32) + (bcol * 1 : U32) : U32) : U32)))
98def gr_33(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, +y0_0: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
99 match rv:
100 case Tuple{b, +y0_1}:
101 gr_34(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, y0_1, a, Array.get(F32, b, (ib + ((sb * 1 : U32) + (bcol * 0 : U32) : U32) : U32)))
102def gr_32(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, +x7_3: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
103 match rv:
104 case Tuple{b, +y0_0}:
105 gr_33(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, y0_0, a, Array.get(F32, b, (ib + ((sb * 0 : U32) + (bcol * 1 : U32) : U32) : U32)))
106def gr_31(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, +x7_2: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
107 match rv:
108 case Tuple{a, +x7_3}:
109 gr_32(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, x7_3, a, Array.get(F32, b, (ib + ((sb * 0 : U32) + (bcol * 0 : U32) : U32) : U32)))
110def gr_30(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, +x7_1: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
111 match rv:
112 case Tuple{a, +x7_2}:
113 gr_31(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, x7_2, Array.get(F32, a, (ia + ((sa * 7 : U32) + (arow * 3 : U32) : U32) : U32)), b)
114def gr_29(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, +x7_0: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
115 match rv:
116 case Tuple{a, +x7_1}:
117 gr_30(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, x7_1, Array.get(F32, a, (ia + ((sa * 7 : U32) + (arow * 2 : U32) : U32) : U32)), b)
118def gr_28(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, +x6_3: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
119 match rv:
120 case Tuple{a, +x7_0}:
121 gr_29(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, x7_0, Array.get(F32, a, (ia + ((sa * 7 : U32) + (arow * 1 : U32) : U32) : U32)), b)
122def gr_27(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, +x6_2: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
123 match rv:
124 case Tuple{a, +x6_3}:
125 gr_28(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, x6_3, Array.get(F32, a, (ia + ((sa * 7 : U32) + (arow * 0 : U32) : U32) : U32)), b)
126def gr_26(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, +x6_1: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
127 match rv:
128 case Tuple{a, +x6_2}:
129 gr_27(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, x6_2, Array.get(F32, a, (ia + ((sa * 6 : U32) + (arow * 3 : U32) : U32) : U32)), b)
130def gr_25(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, +x6_0: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
131 match rv:
132 case Tuple{a, +x6_1}:
133 gr_26(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, x6_1, Array.get(F32, a, (ia + ((sa * 6 : U32) + (arow * 2 : U32) : U32) : U32)), b)
134def gr_24(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, +x5_3: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
135 match rv:
136 case Tuple{a, +x6_0}:
137 gr_25(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, x6_0, Array.get(F32, a, (ia + ((sa * 6 : U32) + (arow * 1 : U32) : U32) : U32)), b)
138def gr_23(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, +x5_2: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
139 match rv:
140 case Tuple{a, +x5_3}:
141 gr_24(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, x5_3, Array.get(F32, a, (ia + ((sa * 6 : U32) + (arow * 0 : U32) : U32) : U32)), b)
142def gr_22(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, +x5_1: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
143 match rv:
144 case Tuple{a, +x5_2}:
145 gr_23(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, x5_2, Array.get(F32, a, (ia + ((sa * 5 : U32) + (arow * 3 : U32) : U32) : U32)), b)
146def gr_21(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, +x5_0: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
147 match rv:
148 case Tuple{a, +x5_1}:
149 gr_22(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, x5_1, Array.get(F32, a, (ia + ((sa * 5 : U32) + (arow * 2 : U32) : U32) : U32)), b)
150def gr_20(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, +x4_3: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
151 match rv:
152 case Tuple{a, +x5_0}:
153 gr_21(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, x5_0, Array.get(F32, a, (ia + ((sa * 5 : U32) + (arow * 1 : U32) : U32) : U32)), b)
154def gr_19(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, +x4_2: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
155 match rv:
156 case Tuple{a, +x4_3}:
157 gr_20(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, x4_3, Array.get(F32, a, (ia + ((sa * 5 : U32) + (arow * 0 : U32) : U32) : U32)), b)
158def gr_18(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, +x4_1: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
159 match rv:
160 case Tuple{a, +x4_2}:
161 gr_19(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, x4_2, Array.get(F32, a, (ia + ((sa * 4 : U32) + (arow * 3 : U32) : U32) : U32)), b)
162def gr_17(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, +x4_0: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
163 match rv:
164 case Tuple{a, +x4_1}:
165 gr_18(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, x4_1, Array.get(F32, a, (ia + ((sa * 4 : U32) + (arow * 2 : U32) : U32) : U32)), b)
166def gr_16(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, +x3_3: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
167 match rv:
168 case Tuple{a, +x4_0}:
169 gr_17(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, x4_0, Array.get(F32, a, (ia + ((sa * 4 : U32) + (arow * 1 : U32) : U32) : U32)), b)
170def gr_15(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, +x3_2: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
171 match rv:
172 case Tuple{a, +x3_3}:
173 gr_16(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, x3_3, Array.get(F32, a, (ia + ((sa * 4 : U32) + (arow * 0 : U32) : U32) : U32)), b)
174def gr_14(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, +x3_1: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
175 match rv:
176 case Tuple{a, +x3_2}:
177 gr_15(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, x3_2, Array.get(F32, a, (ia + ((sa * 3 : U32) + (arow * 3 : U32) : U32) : U32)), b)
178def gr_13(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, +x3_0: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
179 match rv:
180 case Tuple{a, +x3_1}:
181 gr_14(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, x3_1, Array.get(F32, a, (ia + ((sa * 3 : U32) + (arow * 2 : U32) : U32) : U32)), b)
182def gr_12(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, +x2_3: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
183 match rv:
184 case Tuple{a, +x3_0}:
185 gr_13(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, x3_0, Array.get(F32, a, (ia + ((sa * 3 : U32) + (arow * 1 : U32) : U32) : U32)), b)
186def gr_11(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, +x2_2: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
187 match rv:
188 case Tuple{a, +x2_3}:
189 gr_12(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, x2_3, Array.get(F32, a, (ia + ((sa * 3 : U32) + (arow * 0 : U32) : U32) : U32)), b)
190def gr_10(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, +x2_1: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
191 match rv:
192 case Tuple{a, +x2_2}:
193 gr_11(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, x2_2, Array.get(F32, a, (ia + ((sa * 2 : U32) + (arow * 3 : U32) : U32) : U32)), b)
194def gr_9(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, +x2_0: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
195 match rv:
196 case Tuple{a, +x2_1}:
197 gr_10(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, x2_1, Array.get(F32, a, (ia + ((sa * 2 : U32) + (arow * 2 : U32) : U32) : U32)), b)
198def gr_8(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, +x1_3: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
199 match rv:
200 case Tuple{a, +x2_0}:
201 gr_9(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, x2_0, Array.get(F32, a, (ia + ((sa * 2 : U32) + (arow * 1 : U32) : U32) : U32)), b)
202def gr_7(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, +x1_2: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
203 match rv:
204 case Tuple{a, +x1_3}:
205 gr_8(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, x1_3, Array.get(F32, a, (ia + ((sa * 2 : U32) + (arow * 0 : U32) : U32) : U32)), b)
206def gr_6(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, +x1_1: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
207 match rv:
208 case Tuple{a, +x1_2}:
209 gr_7(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, x1_2, Array.get(F32, a, (ia + ((sa * 1 : U32) + (arow * 3 : U32) : U32) : U32)), b)
210def gr_5(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, +x1_0: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
211 match rv:
212 case Tuple{a, +x1_1}:
213 gr_6(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, x1_1, Array.get(F32, a, (ia + ((sa * 1 : U32) + (arow * 2 : U32) : U32) : U32)), b)
214def gr_4(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, +x0_3: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
215 match rv:
216 case Tuple{a, +x1_0}:
217 gr_5(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, x1_0, Array.get(F32, a, (ia + ((sa * 1 : U32) + (arow * 1 : U32) : U32) : U32)), b)
218def gr_3(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, +x0_2: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
219 match rv:
220 case Tuple{a, +x0_3}:
221 gr_4(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, x0_3, Array.get(F32, a, (ia + ((sa * 1 : U32) + (arow * 0 : U32) : U32) : U32)), b)
222def gr_2(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
223 match rv:
224 case Tuple{a, +x0_2}:
225 gr_3(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, x0_2, Array.get(F32, a, (ia + ((sa * 0 : U32) + (arow * 3 : U32) : U32) : U32)), b)
226def gr_1(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
227 match rv:
228 case Tuple{a, +x0_1}:
229 gr_2(ia, ib, arow, sa, sb, bcol, x0_0, x0_1, Array.get(F32, a, (ia + ((sa * 0 : U32) + (arow * 2 : U32) : U32) : U32)), b)
230def gr_0(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
231 match rv:
232 case Tuple{a, +x0_0}:
233 gr_1(ia, ib, arow, sa, sb, bcol, x0_0, Array.get(F32, a, (ia + ((sa * 0 : U32) + (arow * 1 : U32) : U32) : U32)), b)
234# the turn whose A values start at ia and B values at ib
235def turn(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>) -> GR:
236 gr_0(ia, ib, arow, sa, sb, bcol, Array.get(F32, a, ia), b)
237
238# the tile's cells after the 8 sums are combined, row by row
239type T is Type:
240 T{a: Array<F32>, b: Array<F32>, c0_0: F32, c0_1: F32, c1_0: F32, c1_1: F32, c2_0: F32, c2_1: F32, c3_0: F32, c3_1: F32}
241
242# kg more full turns; g holds the current turn's values. The last call (kg = 0) adds the turn that holds the steps
243# past the last full turn (rem = k mod 8 of them), counting only those, and combines the sums.
244@unsafe
245def mm(kg: Nat, +ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +rem: U32, +rot: Nat, +s0_0_0: F32, +s1_0_0: F32, +s2_0_0: F32, +s3_0_0: F32, +s4_0_0: F32, +s5_0_0: F32, +s6_0_0: F32, +s7_0_0: F32, +s0_0_1: F32, +s1_0_1: F32, +s2_0_1: F32, +s3_0_1: F32, +s4_0_1: F32, +s5_0_1: F32, +s6_0_1: F32, +s7_0_1: F32, +s0_1_0: F32, +s1_1_0: F32, +s2_1_0: F32, +s3_1_0: F32, +s4_1_0: F32, +s5_1_0: F32, +s6_1_0: F32, +s7_1_0: F32, +s0_1_1: F32, +s1_1_1: F32, +s2_1_1: F32, +s3_1_1: F32, +s4_1_1: F32, +s5_1_1: F32, +s6_1_1: F32, +s7_1_1: F32, +s0_2_0: F32, +s1_2_0: F32, +s2_2_0: F32, +s3_2_0: F32, +s4_2_0: F32, +s5_2_0: F32, +s6_2_0: F32, +s7_2_0: F32, +s0_2_1: F32, +s1_2_1: F32, +s2_2_1: F32, +s3_2_1: F32, +s4_2_1: F32, +s5_2_1: F32, +s6_2_1: F32, +s7_2_1: F32, +s0_3_0: F32, +s1_3_0: F32, +s2_3_0: F32, +s3_3_0: F32, +s4_3_0: F32, +s5_3_0: F32, +s6_3_0: F32, +s7_3_0: F32, +s0_3_1: F32, +s1_3_1: F32, +s2_3_1: F32, +s3_3_1: F32, +s4_3_1: F32, +s5_3_1: F32, +s6_3_1: F32, +s7_3_1: F32, g: GR) -> T:
246 match kg g:
247 case 0n GR{a, b, +x0_0, +x0_1, +x0_2, +x0_3, +x1_0, +x1_1, +x1_2, +x1_3, +x2_0, +x2_1, +x2_2, +x2_3, +x3_0, +x3_1, +x3_2, +x3_3, +x4_0, +x4_1, +x4_2, +x4_3, +x5_0, +x5_1, +x5_2, +x5_3, +x6_0, +x6_1, +x6_2, +x6_3, +x7_0, +x7_1, +x7_2, +x7_3, +y0_0, +y0_1, +y1_0, +y1_1, +y2_0, +y2_1, +y3_0, +y3_1, +y4_0, +y4_1, +y5_0, +y5_1, +y6_0, +y6_1, +y7_0, +y7_1}:
248 T{a, b, TA.t8_comb(rot, (s0_0_0 + pk(U32.is_lt(0, rem), (x0_0 * y0_0 : F32)) : F32), (s1_0_0 + pk(U32.is_lt(1, rem), (x1_0 * y1_0 : F32)) : F32), (s2_0_0 + pk(U32.is_lt(2, rem), (x2_0 * y2_0 : F32)) : F32), (s3_0_0 + pk(U32.is_lt(3, rem), (x3_0 * y3_0 : F32)) : F32), (s4_0_0 + pk(U32.is_lt(4, rem), (x4_0 * y4_0 : F32)) : F32), (s5_0_0 + pk(U32.is_lt(5, rem), (x5_0 * y5_0 : F32)) : F32), (s6_0_0 + pk(U32.is_lt(6, rem), (x6_0 * y6_0 : F32)) : F32), (s7_0_0 + pk(U32.is_lt(7, rem), (x7_0 * y7_0 : F32)) : F32)), TA.t8_comb(rot, (s0_0_1 + pk(U32.is_lt(0, rem), (x0_0 * y0_1 : F32)) : F32), (s1_0_1 + pk(U32.is_lt(1, rem), (x1_0 * y1_1 : F32)) : F32), (s2_0_1 + pk(U32.is_lt(2, rem), (x2_0 * y2_1 : F32)) : F32), (s3_0_1 + pk(U32.is_lt(3, rem), (x3_0 * y3_1 : F32)) : F32), (s4_0_1 + pk(U32.is_lt(4, rem), (x4_0 * y4_1 : F32)) : F32), (s5_0_1 + pk(U32.is_lt(5, rem), (x5_0 * y5_1 : F32)) : F32), (s6_0_1 + pk(U32.is_lt(6, rem), (x6_0 * y6_1 : F32)) : F32), (s7_0_1 + pk(U32.is_lt(7, rem), (x7_0 * y7_1 : F32)) : F32)), TA.t8_comb(rot, (s0_1_0 + pk(U32.is_lt(0, rem), (x0_1 * y0_0 : F32)) : F32), (s1_1_0 + pk(U32.is_lt(1, rem), (x1_1 * y1_0 : F32)) : F32), (s2_1_0 + pk(U32.is_lt(2, rem), (x2_1 * y2_0 : F32)) : F32), (s3_1_0 + pk(U32.is_lt(3, rem), (x3_1 * y3_0 : F32)) : F32), (s4_1_0 + pk(U32.is_lt(4, rem), (x4_1 * y4_0 : F32)) : F32), (s5_1_0 + pk(U32.is_lt(5, rem), (x5_1 * y5_0 : F32)) : F32), (s6_1_0 + pk(U32.is_lt(6, rem), (x6_1 * y6_0 : F32)) : F32), (s7_1_0 + pk(U32.is_lt(7, rem), (x7_1 * y7_0 : F32)) : F32)), TA.t8_comb(rot, (s0_1_1 + pk(U32.is_lt(0, rem), (x0_1 * y0_1 : F32)) : F32), (s1_1_1 + pk(U32.is_lt(1, rem), (x1_1 * y1_1 : F32)) : F32), (s2_1_1 + pk(U32.is_lt(2, rem), (x2_1 * y2_1 : F32)) : F32), (s3_1_1 + pk(U32.is_lt(3, rem), (x3_1 * y3_1 : F32)) : F32), (s4_1_1 + pk(U32.is_lt(4, rem), (x4_1 * y4_1 : F32)) : F32), (s5_1_1 + pk(U32.is_lt(5, rem), (x5_1 * y5_1 : F32)) : F32), (s6_1_1 + pk(U32.is_lt(6, rem), (x6_1 * y6_1 : F32)) : F32), (s7_1_1 + pk(U32.is_lt(7, rem), (x7_1 * y7_1 : F32)) : F32)), TA.t8_comb(rot, (s0_2_0 + pk(U32.is_lt(0, rem), (x0_2 * y0_0 : F32)) : F32), (s1_2_0 + pk(U32.is_lt(1, rem), (x1_2 * y1_0 : F32)) : F32), (s2_2_0 + pk(U32.is_lt(2, rem), (x2_2 * y2_0 : F32)) : F32), (s3_2_0 + pk(U32.is_lt(3, rem), (x3_2 * y3_0 : F32)) : F32), (s4_2_0 + pk(U32.is_lt(4, rem), (x4_2 * y4_0 : F32)) : F32), (s5_2_0 + pk(U32.is_lt(5, rem), (x5_2 * y5_0 : F32)) : F32), (s6_2_0 + pk(U32.is_lt(6, rem), (x6_2 * y6_0 : F32)) : F32), (s7_2_0 + pk(U32.is_lt(7, rem), (x7_2 * y7_0 : F32)) : F32)), TA.t8_comb(rot, (s0_2_1 + pk(U32.is_lt(0, rem), (x0_2 * y0_1 : F32)) : F32), (s1_2_1 + pk(U32.is_lt(1, rem), (x1_2 * y1_1 : F32)) : F32), (s2_2_1 + pk(U32.is_lt(2, rem), (x2_2 * y2_1 : F32)) : F32), (s3_2_1 + pk(U32.is_lt(3, rem), (x3_2 * y3_1 : F32)) : F32), (s4_2_1 + pk(U32.is_lt(4, rem), (x4_2 * y4_1 : F32)) : F32), (s5_2_1 + pk(U32.is_lt(5, rem), (x5_2 * y5_1 : F32)) : F32), (s6_2_1 + pk(U32.is_lt(6, rem), (x6_2 * y6_1 : F32)) : F32), (s7_2_1 + pk(U32.is_lt(7, rem), (x7_2 * y7_1 : F32)) : F32)), TA.t8_comb(rot, (s0_3_0 + pk(U32.is_lt(0, rem), (x0_3 * y0_0 : F32)) : F32), (s1_3_0 + pk(U32.is_lt(1, rem), (x1_3 * y1_0 : F32)) : F32), (s2_3_0 + pk(U32.is_lt(2, rem), (x2_3 * y2_0 : F32)) : F32), (s3_3_0 + pk(U32.is_lt(3, rem), (x3_3 * y3_0 : F32)) : F32), (s4_3_0 + pk(U32.is_lt(4, rem), (x4_3 * y4_0 : F32)) : F32), (s5_3_0 + pk(U32.is_lt(5, rem), (x5_3 * y5_0 : F32)) : F32), (s6_3_0 + pk(U32.is_lt(6, rem), (x6_3 * y6_0 : F32)) : F32), (s7_3_0 + pk(U32.is_lt(7, rem), (x7_3 * y7_0 : F32)) : F32)), TA.t8_comb(rot, (s0_3_1 + pk(U32.is_lt(0, rem), (x0_3 * y0_1 : F32)) : F32), (s1_3_1 + pk(U32.is_lt(1, rem), (x1_3 * y1_1 : F32)) : F32), (s2_3_1 + pk(U32.is_lt(2, rem), (x2_3 * y2_1 : F32)) : F32), (s3_3_1 + pk(U32.is_lt(3, rem), (x3_3 * y3_1 : F32)) : F32), (s4_3_1 + pk(U32.is_lt(4, rem), (x4_3 * y4_1 : F32)) : F32), (s5_3_1 + pk(U32.is_lt(5, rem), (x5_3 * y5_1 : F32)) : F32), (s6_3_1 + pk(U32.is_lt(6, rem), (x6_3 * y6_1 : F32)) : F32), (s7_3_1 + pk(U32.is_lt(7, rem), (x7_3 * y7_1 : F32)) : F32))}
249 case 1n+q GR{a, b, +x0_0, +x0_1, +x0_2, +x0_3, +x1_0, +x1_1, +x1_2, +x1_3, +x2_0, +x2_1, +x2_2, +x2_3, +x3_0, +x3_1, +x3_2, +x3_3, +x4_0, +x4_1, +x4_2, +x4_3, +x5_0, +x5_1, +x5_2, +x5_3, +x6_0, +x6_1, +x6_2, +x6_3, +x7_0, +x7_1, +x7_2, +x7_3, +y0_0, +y0_1, +y1_0, +y1_1, +y2_0, +y2_1, +y3_0, +y3_1, +y4_0, +y4_1, +y5_0, +y5_1, +y6_0, +y6_1, +y7_0, +y7_1}:
250 mm(q, (ia + (sa * 8 : U32) : U32), (ib + (sb * 8 : U32) : U32), arow, sa, sb, bcol, rem, rot, (s0_0_0 + (x0_0 * y0_0 : F32) : F32), (s1_0_0 + (x1_0 * y1_0 : F32) : F32), (s2_0_0 + (x2_0 * y2_0 : F32) : F32), (s3_0_0 + (x3_0 * y3_0 : F32) : F32), (s4_0_0 + (x4_0 * y4_0 : F32) : F32), (s5_0_0 + (x5_0 * y5_0 : F32) : F32), (s6_0_0 + (x6_0 * y6_0 : F32) : F32), (s7_0_0 + (x7_0 * y7_0 : F32) : F32), (s0_0_1 + (x0_0 * y0_1 : F32) : F32), (s1_0_1 + (x1_0 * y1_1 : F32) : F32), (s2_0_1 + (x2_0 * y2_1 : F32) : F32), (s3_0_1 + (x3_0 * y3_1 : F32) : F32), (s4_0_1 + (x4_0 * y4_1 : F32) : F32), (s5_0_1 + (x5_0 * y5_1 : F32) : F32), (s6_0_1 + (x6_0 * y6_1 : F32) : F32), (s7_0_1 + (x7_0 * y7_1 : F32) : F32), (s0_1_0 + (x0_1 * y0_0 : F32) : F32), (s1_1_0 + (x1_1 * y1_0 : F32) : F32), (s2_1_0 + (x2_1 * y2_0 : F32) : F32), (s3_1_0 + (x3_1 * y3_0 : F32) : F32), (s4_1_0 + (x4_1 * y4_0 : F32) : F32), (s5_1_0 + (x5_1 * y5_0 : F32) : F32), (s6_1_0 + (x6_1 * y6_0 : F32) : F32), (s7_1_0 + (x7_1 * y7_0 : F32) : F32), (s0_1_1 + (x0_1 * y0_1 : F32) : F32), (s1_1_1 + (x1_1 * y1_1 : F32) : F32), (s2_1_1 + (x2_1 * y2_1 : F32) : F32), (s3_1_1 + (x3_1 * y3_1 : F32) : F32), (s4_1_1 + (x4_1 * y4_1 : F32) : F32), (s5_1_1 + (x5_1 * y5_1 : F32) : F32), (s6_1_1 + (x6_1 * y6_1 : F32) : F32), (s7_1_1 + (x7_1 * y7_1 : F32) : F32), (s0_2_0 + (x0_2 * y0_0 : F32) : F32), (s1_2_0 + (x1_2 * y1_0 : F32) : F32), (s2_2_0 + (x2_2 * y2_0 : F32) : F32), (s3_2_0 + (x3_2 * y3_0 : F32) : F32), (s4_2_0 + (x4_2 * y4_0 : F32) : F32), (s5_2_0 + (x5_2 * y5_0 : F32) : F32), (s6_2_0 + (x6_2 * y6_0 : F32) : F32), (s7_2_0 + (x7_2 * y7_0 : F32) : F32), (s0_2_1 + (x0_2 * y0_1 : F32) : F32), (s1_2_1 + (x1_2 * y1_1 : F32) : F32), (s2_2_1 + (x2_2 * y2_1 : F32) : F32), (s3_2_1 + (x3_2 * y3_1 : F32) : F32), (s4_2_1 + (x4_2 * y4_1 : F32) : F32), (s5_2_1 + (x5_2 * y5_1 : F32) : F32), (s6_2_1 + (x6_2 * y6_1 : F32) : F32), (s7_2_1 + (x7_2 * y7_1 : F32) : F32), (s0_3_0 + (x0_3 * y0_0 : F32) : F32), (s1_3_0 + (x1_3 * y1_0 : F32) : F32), (s2_3_0 + (x2_3 * y2_0 : F32) : F32), (s3_3_0 + (x3_3 * y3_0 : F32) : F32), (s4_3_0 + (x4_3 * y4_0 : F32) : F32), (s5_3_0 + (x5_3 * y5_0 : F32) : F32), (s6_3_0 + (x6_3 * y6_0 : F32) : F32), (s7_3_0 + (x7_3 * y7_0 : F32) : F32), (s0_3_1 + (x0_3 * y0_1 : F32) : F32), (s1_3_1 + (x1_3 * y1_1 : F32) : F32), (s2_3_1 + (x2_3 * y2_1 : F32) : F32), (s3_3_1 + (x3_3 * y3_1 : F32) : F32), (s4_3_1 + (x4_3 * y4_1 : F32) : F32), (s5_3_1 + (x5_3 * y5_1 : F32) : F32), (s6_3_1 + (x6_3 * y6_1 : F32) : F32), (s7_3_1 + (x7_3 * y7_1 : F32) : F32), turn((ia + (sa * 8 : U32) : U32), (ib + (sb * 8 : U32) : U32), arow, sa, sb, bcol, a, b))
251
252# writes cell (i, j) of C when it is inside n x m
253def wr(ok: Bool, +ac: Bool, +idx: U32, +x: F32, c: Array<F32>) -> Array<F32>:
254 match ok:
255 case True{}: TA.cwr(ac, idx, x, c)
256 case False{}: c
257
258# the tile at rows i0.., columns j0.. of C, from the loop's result; A and B are dropped (the caller holds them)
259@unsafe
260def put(+ac: Bool, +n: U32, +m: U32, +i0: U32, +j0: U32, c: Array<F32>, t: T) -> Array<F32>:
261 match t:
262 case T{a, b, +c0_0, +c0_1, +c1_0, +c1_1, +c2_0, +c2_1, +c3_0, +c3_1}:
263 wr(Bool.and(U32.is_lt((i0 + 3 : U32), n), U32.is_lt((j0 + 1 : U32), m)), ac, (((i0 + 3 : U32) * m : U32) + (j0 + 1 : U32) : U32), c3_1, wr(Bool.and(U32.is_lt((i0 + 3 : U32), n), U32.is_lt((j0 + 0 : U32), m)), ac, (((i0 + 3 : U32) * m : U32) + (j0 + 0 : U32) : U32), c3_0, wr(Bool.and(U32.is_lt((i0 + 2 : U32), n), U32.is_lt((j0 + 1 : U32), m)), ac, (((i0 + 2 : U32) * m : U32) + (j0 + 1 : U32) : U32), c2_1, wr(Bool.and(U32.is_lt((i0 + 2 : U32), n), U32.is_lt((j0 + 0 : U32), m)), ac, (((i0 + 2 : U32) * m : U32) + (j0 + 0 : U32) : U32), c2_0, wr(Bool.and(U32.is_lt((i0 + 1 : U32), n), U32.is_lt((j0 + 1 : U32), m)), ac, (((i0 + 1 : U32) * m : U32) + (j0 + 1 : U32) : U32), c1_1, wr(Bool.and(U32.is_lt((i0 + 1 : U32), n), U32.is_lt((j0 + 0 : U32), m)), ac, (((i0 + 1 : U32) * m : U32) + (j0 + 0 : U32) : U32), c1_0, wr(Bool.and(U32.is_lt((i0 + 0 : U32), n), U32.is_lt((j0 + 1 : U32), m)), ac, (((i0 + 0 : U32) * m : U32) + (j0 + 1 : U32) : U32), c0_1, wr(Bool.and(U32.is_lt((i0 + 0 : U32), n), U32.is_lt((j0 + 0 : U32), m)), ac, (((i0 + 0 : U32) * m : U32) + (j0 + 0 : U32) : U32), c0_0, c))))))))
264
265@unsafe
266def tile(+e: E, +i0: U32, +j0: U32, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> Array<F32>:
267 match e:
268 case E{+n, +m, +k, +tcols, +tiles, +arow, +sa, +sb, +bcol, +kg, +rot, +rem, +ac, +oa, +ob}:
269 put(ac, n, m, i0, j0, c, mm(kg, (oa + (i0 * arow : U32) : U32), (ob + (j0 * bcol : U32) : U32), arow, sa, sb, bcol, rem, rot, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, turn((oa + (i0 * arow : U32) : U32), (ob + (j0 * bcol : U32) : U32), arow, sa, sb, bcol, a, b)))
270
271# leaf t computes tile t (row-major over the tiles); leaves past the last tile write nothing
272@unsafe
273def leaf_if(ok: Bool, +e: E, +t: U32, +tcols: U32, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> Array<F32>:
274 match ok:
275 case True{}: tile(e, (U32.div(t, tcols) * 4 : U32), (U32.mod(t, tcols) * 2 : U32), a, b, c)
276 case False{}: c
277
278@unsafe
279def leaf(+e: E, +t: U32, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> Array<F32>:
280 match e:
281 case E{+n, +m, +k, +tcols, +tiles, +arow, +sa, +sb, +bcol, +kg, +rot, +rem, +ac, +oa, +ob}:
282 leaf_if(U32.is_lt(t, tiles), e, t, tcols, a, b, c)
283
284@unsafe
285def tree(+e: E, d: Nat, +t: U32, +a: Array<F32>, +b: Array<F32>, +c: Array<F32>) -> Array<F32>:
286 match d:
287 case 0n: leaf(e, t, a, b, c)
288 case 1n++q:
289 l r = tree(e, q, (t * 2 : U32), a, b, c) tree(e, q, ((t * 2 : U32) + 1 : U32), a, b, c)
290 Array.join(F32, l, r)
291
292# C made inside the bang: the root of the ! call allocates it on the device. Made on the host instead, a new C often
293# reuses a block whose pages sit on the device, so zeroing it on the host pulls them back over PCIe and the bang moves
294# them again (a 256 x 1536 product: 9.2 ms against 6.0, measured 2026-10-18). The leaves write every cell of C.
295@unsafe
296def tree_new(+e: E, d: Nat, +a: Array<F32>, +b: Array<F32>, +cd: Nat) -> Array<F32>:
297 tree(e, d, 0, a, b, Array.new(F32, cd, 0.0))
298
299@unsafe
300def bang_new(+e: E, d: Nat, +a: Array<F32>, +b: Array<F32>, +cd: Nat) -> Array<F32>:
301 tree_new!(e, d, a, b, cd)
302
303# one bang over the tiles; A and B are shared with every leaf and come back unchanged
304@unsafe
305def bang(+e: E, d: Nat, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> Array<F32>:
306 tree!(e, d, 0, a, b, c)
307
308type ABC is Type:
309 ABC{a: Array<F32>, b: Array<F32>, c: Array<F32>}
310
311@unsafe
312def run_e(+e: E, +tiles: U32, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> ABC:
313 ABC{a, b, bang(e, TA.cap_depth(U32.to_nat(tiles)), a, b, c)}
314
315# the product's sizes and strides
316def e_ofo(+ac: Bool, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +oa: U32, +ob: U32) -> E:
317 E{U32.from_nat(n), U32.from_nat(m), U32.from_nat(k), U32.from_nat(Nat.div(Nat.add(m, 1n), 2n)), U32.from_nat(Nat.mul(Nat.div(Nat.add(n, 3n), 4n), Nat.div(Nat.add(m, 1n), 2n))), arow, sa, sb, bcol, Nat.div(k, 8n), Nat.mod(k, 8n), U32.from_nat(Nat.mod(k, 8n)), ac, oa, ob}
318
319def e_of(+ac: Bool, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32) -> E:
320 e_ofo(ac, n, m, k, arow, sa, sb, bcol, 0, 0)
321
322def tiles(+n: Nat, +m: Nat) -> Nat:
323 Nat.mul(Nat.div(Nat.add(n, 3n), 4n), Nat.div(Nat.add(m, 1n), 2n))
324
325# C (n x m) = A · B with k steps into a new C made on the device
326@unsafe
327def gemm_new(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +a: Array<F32>, +b: Array<F32>) -> ABC:
328 ABC{a, b, bang_new(e_of(False{}, n, m, k, arow, sa, sb, bcol), TA.cap_depth(tiles(n, m)), a, b, TA.cap_depth(Nat.mul(n, m)))}
329
330# C (n x m) = A · B with k steps, or C + A · B when ac, C[i][j] = sum_p A[i*arow + p*sa] * B[p*sb + j*bcol]
331@unsafe
332def gemm(+ac: Bool, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> ABC:
333 run_e(e_of(ac, n, m, k, arow, sa, sb, bcol), U32.from_nat(tiles(n, m)), a, b, c)
334
335def mm_fin(-n: Nat, -k: Nat, -m: Nat, g: ABC) -> TA.MMul<n, k, m>:
336 match g:
337 case ABC{a, b, c}: TA.MMul{TA.Mat{a}, TA.Mat{b}, TA.Mat{c}}
338
339def mm_nt_fin(-n: Nat, -k: Nat, -m: Nat, g: ABC) -> TA.MMulNT<n, k, m>:
340 match g:
341 case ABC{a, b, c}: TA.MMulNT{TA.Mat{a}, TA.Mat{b}, TA.Mat{c}}
342
343def mm_tn_fin(-n: Nat, -k: Nat, -m: Nat, g: ABC) -> TA.MMulTN<n, k, m>:
344 match g:
345 case ABC{a, b, c}: TA.MMulTN{TA.Mat{a}, TA.Mat{b}, TA.Mat{c}}
346
347# C = A · B (n x k) · (k x m), tensor-array's TA.Mat.matmul bits
348@unsafe
349def matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat<n, k>, b: TA.Mat<k, m>) -> TA.MMul<n, k, m>:
350 match a b:
351 case TA.Mat{x} TA.Mat{y}:
352 mm_fin(n, k, m, gemm_new(n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, x, y))
353
354# C + A · B, into C (n x m)
355@unsafe
356def 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>:
357 match a b c:
358 case TA.Mat{x} TA.Mat{y} TA.Mat{z}:
359 mm_fin(n, k, m, gemm(True{}, n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, x, y, z))
360
361# C = A · Bᵀ (n x k) · (m x k)ᵀ: for weights stored with one row per output
362@unsafe
363def 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>:
364 match a b:
365 case TA.Mat{x} TA.Mat{y}:
366 mm_nt_fin(n, k, m, gemm_new(n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), x, y))
367
368# C + A · Bᵀ, into C (n x m)
369@unsafe
370def 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>:
371 match a b c:
372 case TA.Mat{x} TA.Mat{y} TA.Mat{z}:
373 mm_nt_fin(n, k, m, gemm(True{}, n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), x, y, z))
374
375# C = Aᵀ · B (k x n)ᵀ · (k x m): weight gradient, Xᵀ · dY
376@unsafe
377def 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>:
378 match a b:
379 case TA.Mat{x} TA.Mat{y}:
380 mm_tn_fin(n, k, m, gemm_new(n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, x, y))
381
382# C + Aᵀ · B, into C (n x m): dW += Xᵀ · dY
383@unsafe
384def 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>:
385 match a b c:
386 case TA.Mat{x} TA.Mat{y} TA.Mat{z}:
387 mm_tn_fin(n, k, m, gemm(True{}, n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, x, y, z))
388
389# =====================================================================================================================
390# Bangs over any leaf code (0.1.1.0). runK and runK_new (K = 1..4 outputs) run one bang over 2^d leaves: leaf t of
391# l = 2^d gets t, l, a record e of scalars (the same for every leaf), four input arrays and the K outputs, and returns
392# the outputs. The leaf code is a template (~f), written by the caller in plain Bend: it takes every array as an ordinary
393# argument (it reads them with Array.get and may write the outputs), so it needs no @unsafe; the sharing is here, where
394# every node of the tree holds a handle to each array (Array.fork in Base, upstream's one-block pattern) and joins the
395# outputs (Array.join). runK takes the outputs from the caller and writes them in place (an optimizer's weights and
396# moments); runK_new makes them inside the bang with 2^c1.. cells, since an array made on the host pulls device pages
397# back (see bang_new above). Both return the inputs unchanged with the outputs (RK). An unused input is any array, such
398# as none(Unit{}).
399#
400# Not checked, as for the products: that a leaf writes only its own cells of the outputs (two leaves writing one cell
401# race) and never an input. Leaves past the work must return their outputs unchanged; share(t, n, l) is how many of the
402# cells t, t + l, t + 2l, .. lie below n (leaf t's cells when n cells are dealt round-robin, so neighbouring lanes read
403# neighbouring cells), and depth(n) is a tree for n cells: 2^14 leaves once there are that many, else the fewest that
404# cover them. The leaves run in a run-time order inside each group of 128 (bendlang/bend's CUDA lowering), so a leaf's
405# result must not depend on another's. Every number a leaf computes is the caller's own code, so the bits are the
406# caller's: the same on the GPU as on the CPU threads, except where CUDA's math functions differ from the host's (an exp
407# on the device takes the float overload; bend-ml's docs/upstream/gpu-math-overloads.md).
408# =====================================================================================================================
409
410def depth2(big: Bool, +n: Nat) -> Nat:
411 match big:
412 case True{}: 14n
413 case False{}: TA.cap_depth(n)
414
415# leaves for n cells: 2^14 when there are that many, else the fewest powers of two that cover them
416def depth(+n: Nat) -> Nat:
417 depth2(Nat.is_le(16384n, n), n)
418
419# how many of the cells t, t + l, t + 2l, .. lie below n
420def share(+t: U32, +n: U32, +l: U32) -> Nat:
421 U32.to_nat(Bool.pick(U32, U32.is_lt(t, n), U32.div(((n - t : U32) + (l - 1 : U32) : U32), l), 0))
422
423# an array of one cell, for an unused input
424def none(u: Unit) -> Array<F32>:
425 Array.new(F32, 0n, 0.0)
426
427type O2 is Type:
428 O2{a: Array<F32>, b: Array<F32>}
429
430type O3 is Type:
431 O3{a: Array<F32>, b: Array<F32>, c: Array<F32>}
432
433type O4 is Type:
434 O4{a: Array<F32>, b: Array<F32>, c: Array<F32>, d: Array<F32>}
435
436type R1 is Type:
437 R1{x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, a: Array<F32>}
438
439type R2 is Type:
440 R2{x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: O2}
441
442type R3 is Type:
443 R3{x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: O3}
444
445type R4 is Type:
446 R4{x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: O4}
447
448@unsafe
449def join2(x: O2, y: O2) -> O2:
450 match x y:
451 case O2{a, b} O2{a2, b2}: O2{Array.join(F32, a, a2), Array.join(F32, b, b2)}
452
453@unsafe
454def join3(x: O3, y: O3) -> O3:
455 match x y:
456 case O3{a, b, c} O3{a2, b2, c2}: O3{Array.join(F32, a, a2), Array.join(F32, b, b2), Array.join(F32, c, c2)}
457
458@unsafe
459def join4(x: O4, y: O4) -> O4:
460 match x y:
461 case O4{a, b, c, d} O4{a2, b2, c2, d2}: O4{Array.join(F32, a, a2), Array.join(F32, b, b2), Array.join(F32, c, c2), Array.join(F32, d, d2)}
462
463@unsafe
464def tree1(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32>, +d: Nat, +t: U32, +l: U32, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, +oa: Array<F32>) -> Array<F32>:
465 match d:
466 case 0n: f(t, l, e, x1, x2, x3, x4, oa)
467 case 1n+q:
468 x y = tree1(~E, ~f, q, (t * 2 : U32), l, e, x1, x2, x3, x4, oa) tree1(~E, ~f, q, ((t * 2 : U32) + 1 : U32), l, e, x1, x2, x3, x4, oa)
469 Array.join(F32, x, y)
470
471@unsafe
472def root1(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32>, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, oa: Array<F32>) -> Array<F32>:
473 tree1(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa)
474
475@unsafe
476def root1_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32>, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, +ca: Nat) -> Array<F32>:
477 root1(~E, ~f, d, e, x1, x2, x3, x4, Array.new(F32, ca, 0.0))
478
479@unsafe
480def go1(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32>, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, oa: Array<F32>) -> R1:
481 R1{x1, x2, x3, x4, root1!(~E, ~f, d, e, x1, x2, x3, x4, oa)}
482
483@unsafe
484def go1_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32>, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, +ca: Nat) -> R1:
485 R1{x1, x2, x3, x4, root1_new!(~E, ~f, d, e, x1, x2, x3, x4, ca)}
486
487# one bang over 2^d leaves with 1 output written in place
488@unsafe
489def run1(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32>, +d: Nat, +e: E, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, oa: Array<F32>) -> R1:
490 go1(~E, ~f, d, e, x1, x2, x3, x4, oa)
491
492# one bang over 2^d leaves with 1 new output of 2^c.. cells, made on the device
493@unsafe
494def run1_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32>, +d: Nat, +e: E, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, +ca: Nat) -> R1:
495 go1_new(~E, ~f, d, e, x1, x2, x3, x4, ca)
496
497@unsafe
498def tree2(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O2 -> O2, +d: Nat, +t: U32, +l: U32, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, +oa: Array<F32>, +ob: Array<F32>) -> O2:
499 match d:
500 case 0n: f(t, l, e, x1, x2, x3, x4, O2{oa, ob})
501 case 1n+q:
502 x y = tree2(~E, ~f, q, (t * 2 : U32), l, e, x1, x2, x3, x4, oa, ob) tree2(~E, ~f, q, ((t * 2 : U32) + 1 : U32), l, e, x1, x2, x3, x4, oa, ob)
503 join2(x, y)
504
505@unsafe
506def root2(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O2 -> O2, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, oa: Array<F32>, ob: Array<F32>) -> O2:
507 tree2(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa, ob)
508
509@unsafe
510def root2_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O2 -> O2, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, +ca: Nat, +cb: Nat) -> O2:
511 root2(~E, ~f, d, e, x1, x2, x3, x4, Array.new(F32, ca, 0.0), Array.new(F32, cb, 0.0))
512
513@unsafe
514def go2(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O2 -> O2, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, oa: Array<F32>, ob: Array<F32>) -> R2:
515 R2{x1, x2, x3, x4, root2!(~E, ~f, d, e, x1, x2, x3, x4, oa, ob)}
516
517@unsafe
518def go2_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O2 -> O2, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, +ca: Nat, +cb: Nat) -> R2:
519 R2{x1, x2, x3, x4, root2_new!(~E, ~f, d, e, x1, x2, x3, x4, ca, cb)}
520
521# one bang over 2^d leaves with 2 outputs written in place
522@unsafe
523def run2(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O2 -> O2, +d: Nat, +e: E, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, oa: Array<F32>, ob: Array<F32>) -> R2:
524 go2(~E, ~f, d, e, x1, x2, x3, x4, oa, ob)
525
526# one bang over 2^d leaves with 2 new outputs of 2^c.. cells, made on the device
527@unsafe
528def run2_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O2 -> O2, +d: Nat, +e: E, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, +ca: Nat, +cb: Nat) -> R2:
529 go2_new(~E, ~f, d, e, x1, x2, x3, x4, ca, cb)
530
531@unsafe
532def tree3(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O3 -> O3, +d: Nat, +t: U32, +l: U32, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, +oa: Array<F32>, +ob: Array<F32>, +oc: Array<F32>) -> O3:
533 match d:
534 case 0n: f(t, l, e, x1, x2, x3, x4, O3{oa, ob, oc})
535 case 1n+q:
536 x y = tree3(~E, ~f, q, (t * 2 : U32), l, e, x1, x2, x3, x4, oa, ob, oc) tree3(~E, ~f, q, ((t * 2 : U32) + 1 : U32), l, e, x1, x2, x3, x4, oa, ob, oc)
537 join3(x, y)
538
539@unsafe
540def root3(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O3 -> O3, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, oa: Array<F32>, ob: Array<F32>, oc: Array<F32>) -> O3:
541 tree3(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa, ob, oc)
542
543@unsafe
544def root3_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O3 -> O3, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, +ca: Nat, +cb: Nat, +cc: Nat) -> O3:
545 root3(~E, ~f, d, e, x1, x2, x3, x4, Array.new(F32, ca, 0.0), Array.new(F32, cb, 0.0), Array.new(F32, cc, 0.0))
546
547@unsafe
548def go3(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O3 -> O3, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, oa: Array<F32>, ob: Array<F32>, oc: Array<F32>) -> R3:
549 R3{x1, x2, x3, x4, root3!(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc)}
550
551@unsafe
552def go3_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O3 -> O3, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, +ca: Nat, +cb: Nat, +cc: Nat) -> R3:
553 R3{x1, x2, x3, x4, root3_new!(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc)}
554
555# one bang over 2^d leaves with 3 outputs written in place
556@unsafe
557def run3(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O3 -> O3, +d: Nat, +e: E, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, oa: Array<F32>, ob: Array<F32>, oc: Array<F32>) -> R3:
558 go3(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc)
559
560# one bang over 2^d leaves with 3 new outputs of 2^c.. cells, made on the device
561@unsafe
562def run3_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O3 -> O3, +d: Nat, +e: E, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, +ca: Nat, +cb: Nat, +cc: Nat) -> R3:
563 go3_new(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc)
564
565@unsafe
566def tree4(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O4 -> O4, +d: Nat, +t: U32, +l: U32, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, +oa: Array<F32>, +ob: Array<F32>, +oc: Array<F32>, +od: Array<F32>) -> O4:
567 match d:
568 case 0n: f(t, l, e, x1, x2, x3, x4, O4{oa, ob, oc, od})
569 case 1n+q:
570 x y = tree4(~E, ~f, q, (t * 2 : U32), l, e, x1, x2, x3, x4, oa, ob, oc, od) tree4(~E, ~f, q, ((t * 2 : U32) + 1 : U32), l, e, x1, x2, x3, x4, oa, ob, oc, od)
571 join4(x, y)
572
573@unsafe
574def root4(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O4 -> O4, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, oa: Array<F32>, ob: Array<F32>, oc: Array<F32>, od: Array<F32>) -> O4:
575 tree4(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa, ob, oc, od)
576
577@unsafe
578def root4_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O4 -> O4, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, +ca: Nat, +cb: Nat, +cc: Nat, +cd: Nat) -> O4:
579 root4(~E, ~f, d, e, x1, x2, x3, x4, Array.new(F32, ca, 0.0), Array.new(F32, cb, 0.0), Array.new(F32, cc, 0.0), Array.new(F32, cd, 0.0))
580
581@unsafe
582def go4(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O4 -> O4, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, oa: Array<F32>, ob: Array<F32>, oc: Array<F32>, od: Array<F32>) -> R4:
583 R4{x1, x2, x3, x4, root4!(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc, od)}
584
585@unsafe
586def go4_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O4 -> O4, +d: Nat, +e: E, +x1: Array<F32>, +x2: Array<F32>, +x3: Array<F32>, +x4: Array<F32>, +ca: Nat, +cb: Nat, +cc: Nat, +cd: Nat) -> R4:
587 R4{x1, x2, x3, x4, root4_new!(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc, cd)}
588
589# one bang over 2^d leaves with 4 outputs written in place
590@unsafe
591def run4(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O4 -> O4, +d: Nat, +e: E, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, oa: Array<F32>, ob: Array<F32>, oc: Array<F32>, od: Array<F32>) -> R4:
592 go4(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc, od)
593
594# one bang over 2^d leaves with 4 new outputs of 2^c.. cells, made on the device
595@unsafe
596def run4_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array<F32> -> Array<F32> -> Array<F32> -> Array<F32> -> O4 -> O4, +d: Nat, +e: E, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, +ca: Nat, +cb: Nat, +cc: Nat, +cd: Nat) -> R4:
597 go4_new(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc, cd)
598
599# ---- tensor-array 0.1.14.0's precise product on the GPU (opt-in) ----
600# TA.Mat.matmul_precise's sums, a bang a block of 1024 steps that starts at the block's place in A and B (oa, ob; no
601# copy): the first block's bang makes C on the device, each next one adds its cells into C as the accumulating kernel
602# does (once, as tensor-array's cwr), so every cell is ((x0 + x1) + x2) .., the same bits as tensor-array's
603# (tests/precise.bend). Only plain (not accumulating) products, as tensor-array's matmul_precise.
604
605def blk_len(+bi: U32, +k: U32) -> Nat:
606 U32.to_nat(Bool.pick(U32, U32.is_lt((k - (bi * 1024 : U32) : U32), 1024), (k - (bi * 1024 : U32) : U32), 1024))
607
608# the blocks after the first: nb more, from block bi, each added into C
609@unsafe
610def blocks_acc(nb: Nat, +bi: U32, +n: Nat, +m: Nat, +k: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> Array<F32>:
611 match nb:
612 case 0n: c
613 case 1n++q: blocks_acc(q, (bi + 1 : U32), n, m, k, arow, sa, sb, bcol, a, b, bang(e_ofo(True{}, n, m, blk_len(bi, k), arow, sa, sb, bcol, ((bi * 1024 : U32) * sa : U32), ((bi * 1024 : U32) * sb : U32)), TA.cap_depth(tiles(n, m)), a, b, c))
614
615@unsafe
616def precise_c(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +a: Array<F32>, +b: Array<F32>) -> Array<F32>:
617 blocks_acc(Nat.sub(U32.to_nat(U32.div((U32.from_nat(k) + 1023 : U32), 1024)), 1n), 1, n, m, U32.from_nat(k), arow, sa, sb, bcol, a, b, bang_new(e_ofo(False{}, n, m, 1024n, arow, sa, sb, bcol, 0, 0), TA.cap_depth(tiles(n, m)), a, b, TA.cap_depth(Nat.mul(n, m))))
618
619@unsafe
620def mm_precise3(+n: Nat, +k: Nat, +m: Nat, +a: Array<F32>, +b: Array<F32>) -> TA.MMul<n, k, m>:
621 mm_fin(n, k, m, ABC{a, b, precise_c(n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, a, b)})
622
623@unsafe
624def 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>:
625 match long a b:
626 case True{} TA.Mat{x} TA.Mat{y}: mm_precise3(n, k, m, x, y)
627 case False{} a b: matmul(n, k, m, par, a, b)
628
629# C = A · B with k summed in blocks of 1024: TA.Mat.matmul_precise's bits (opt-in)
630@unsafe
631def 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>:
632 mm_precise2(Nat.is_lt(1024n, k), n, k, m, par, a, b)