main.bend
bend-ml-tensor-gpu@0.1.2.0 · 571 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.2.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 4 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 4, in
14# increasing p, each sum starting at +0.0, and the four sums are combined by TA.t8_comb with rot = k mod 4, 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 four
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 4 tiles).
24import Base
25import bend-ml-tensor-array@0.1.14.0/main.bend as TA
26
27# the sizes and strides of one product: C[i][j] = sum_p A[i*arow + p*sa] * B[p*sb + j*bcol], n x m, k steps
28type E is Data:
29 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}
30
31# a product that counts only when its step is inside k (the last turn's steps past k add +0.0)
32def pk(ok: Bool, +x: F32) -> F32:
33 match ok:
34 case True{}: x
35 case False{}: 0.0
36
37# the values of one turn: A's 4 rows and B's 4 columns at four consecutive steps, read through the arrays in order
38type GR is Type:
39 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, y0_0: F32, y0_1: F32, y0_2: F32, y0_3: F32, y1_0: F32, y1_1: F32, y1_2: F32, y1_3: F32, y2_0: F32, y2_1: F32, y2_2: F32, y2_3: F32, y3_0: F32, y3_1: F32, y3_2: F32, y3_3: F32}
40
41def 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, +y0_0: F32, +y0_1: F32, +y0_2: F32, +y0_3: F32, +y1_0: F32, +y1_1: F32, +y1_2: F32, +y1_3: F32, +y2_0: F32, +y2_1: F32, +y2_2: F32, +y2_3: F32, +y3_0: F32, +y3_1: F32, +y3_2: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
42 match rv:
43 case Tuple{b, +y3_3}:
44 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, y0_0, y0_1, y0_2, y0_3, y1_0, y1_1, y1_2, y1_3, y2_0, y2_1, y2_2, y2_3, y3_0, y3_1, y3_2, y3_3}
45def 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, +y0_0: F32, +y0_1: F32, +y0_2: F32, +y0_3: F32, +y1_0: F32, +y1_1: F32, +y1_2: F32, +y1_3: F32, +y2_0: F32, +y2_1: F32, +y2_2: F32, +y2_3: F32, +y3_0: F32, +y3_1: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
46 match rv:
47 case Tuple{b, +y3_2}:
48 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, y0_0, y0_1, y0_2, y0_3, y1_0, y1_1, y1_2, y1_3, y2_0, y2_1, y2_2, y2_3, y3_0, y3_1, y3_2, a, Array.get(F32, b, (ib + ((sb * 3 : U32) + (bcol * 3 : U32) : U32) : U32)))
49def 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, +y0_0: F32, +y0_1: F32, +y0_2: F32, +y0_3: F32, +y1_0: F32, +y1_1: F32, +y1_2: F32, +y1_3: F32, +y2_0: F32, +y2_1: F32, +y2_2: F32, +y2_3: F32, +y3_0: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
50 match rv:
51 case Tuple{b, +y3_1}:
52 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, y0_0, y0_1, y0_2, y0_3, y1_0, y1_1, y1_2, y1_3, y2_0, y2_1, y2_2, y2_3, y3_0, y3_1, a, Array.get(F32, b, (ib + ((sb * 3 : U32) + (bcol * 2 : U32) : U32) : U32)))
53def 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, +y0_0: F32, +y0_1: F32, +y0_2: F32, +y0_3: F32, +y1_0: F32, +y1_1: F32, +y1_2: F32, +y1_3: F32, +y2_0: F32, +y2_1: F32, +y2_2: F32, +y2_3: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
54 match rv:
55 case Tuple{b, +y3_0}:
56 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, y0_0, y0_1, y0_2, y0_3, y1_0, y1_1, y1_2, y1_3, y2_0, y2_1, y2_2, y2_3, y3_0, a, Array.get(F32, b, (ib + ((sb * 3 : U32) + (bcol * 1 : U32) : U32) : U32)))
57def 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, +y0_0: F32, +y0_1: F32, +y0_2: F32, +y0_3: F32, +y1_0: F32, +y1_1: F32, +y1_2: F32, +y1_3: F32, +y2_0: F32, +y2_1: F32, +y2_2: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
58 match rv:
59 case Tuple{b, +y2_3}:
60 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, y0_0, y0_1, y0_2, y0_3, y1_0, y1_1, y1_2, y1_3, y2_0, y2_1, y2_2, y2_3, a, Array.get(F32, b, (ib + ((sb * 3 : U32) + (bcol * 0 : U32) : U32) : U32)))
61def 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, +y0_0: F32, +y0_1: F32, +y0_2: F32, +y0_3: F32, +y1_0: F32, +y1_1: F32, +y1_2: F32, +y1_3: F32, +y2_0: F32, +y2_1: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
62 match rv:
63 case Tuple{b, +y2_2}:
64 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, y0_0, y0_1, y0_2, y0_3, y1_0, y1_1, y1_2, y1_3, y2_0, y2_1, y2_2, a, Array.get(F32, b, (ib + ((sb * 2 : U32) + (bcol * 3 : U32) : U32) : U32)))
65def 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, +y0_0: F32, +y0_1: F32, +y0_2: F32, +y0_3: F32, +y1_0: F32, +y1_1: F32, +y1_2: F32, +y1_3: F32, +y2_0: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
66 match rv:
67 case Tuple{b, +y2_1}:
68 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, y0_0, y0_1, y0_2, y0_3, y1_0, y1_1, y1_2, y1_3, y2_0, y2_1, a, Array.get(F32, b, (ib + ((sb * 2 : U32) + (bcol * 2 : U32) : U32) : U32)))
69def 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, +y0_0: F32, +y0_1: F32, +y0_2: F32, +y0_3: F32, +y1_0: F32, +y1_1: F32, +y1_2: F32, +y1_3: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
70 match rv:
71 case Tuple{b, +y2_0}:
72 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, y0_0, y0_1, y0_2, y0_3, y1_0, y1_1, y1_2, y1_3, y2_0, a, Array.get(F32, b, (ib + ((sb * 2 : U32) + (bcol * 1 : U32) : U32) : U32)))
73def 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, +y0_0: F32, +y0_1: F32, +y0_2: F32, +y0_3: F32, +y1_0: F32, +y1_1: F32, +y1_2: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
74 match rv:
75 case Tuple{b, +y1_3}:
76 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, y0_0, y0_1, y0_2, y0_3, y1_0, y1_1, y1_2, y1_3, a, Array.get(F32, b, (ib + ((sb * 2 : U32) + (bcol * 0 : U32) : U32) : U32)))
77def 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, +y0_0: F32, +y0_1: F32, +y0_2: F32, +y0_3: F32, +y1_0: F32, +y1_1: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
78 match rv:
79 case Tuple{b, +y1_2}:
80 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, y0_0, y0_1, y0_2, y0_3, y1_0, y1_1, y1_2, a, Array.get(F32, b, (ib + ((sb * 1 : U32) + (bcol * 3 : U32) : U32) : U32)))
81def 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, +y0_0: F32, +y0_1: F32, +y0_2: F32, +y0_3: F32, +y1_0: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
82 match rv:
83 case Tuple{b, +y1_1}:
84 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, y0_0, y0_1, y0_2, y0_3, y1_0, y1_1, a, Array.get(F32, b, (ib + ((sb * 1 : U32) + (bcol * 2 : U32) : U32) : U32)))
85def 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, +y0_0: F32, +y0_1: F32, +y0_2: F32, +y0_3: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
86 match rv:
87 case Tuple{b, +y1_0}:
88 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, y0_0, y0_1, y0_2, y0_3, y1_0, a, Array.get(F32, b, (ib + ((sb * 1 : U32) + (bcol * 1 : U32) : U32) : U32)))
89def 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, +y0_0: F32, +y0_1: F32, +y0_2: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
90 match rv:
91 case Tuple{b, +y0_3}:
92 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, y0_0, y0_1, y0_2, y0_3, a, Array.get(F32, b, (ib + ((sb * 1 : U32) + (bcol * 0 : U32) : U32) : U32)))
93def 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, +y0_0: F32, +y0_1: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
94 match rv:
95 case Tuple{b, +y0_2}:
96 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, y0_0, y0_1, y0_2, a, Array.get(F32, b, (ib + ((sb * 0 : U32) + (bcol * 3 : U32) : U32) : U32)))
97def 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, +y0_0: F32, a: Array<F32>, rv: Array<F32> & F32) -> GR:
98 match rv:
99 case Tuple{b, +y0_1}:
100 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, y0_0, y0_1, a, Array.get(F32, b, (ib + ((sb * 0 : U32) + (bcol * 2 : U32) : U32) : U32)))
101def 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, a: Array<F32>, rv: Array<F32> & F32) -> GR:
102 match rv:
103 case Tuple{b, +y0_0}:
104 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, y0_0, a, Array.get(F32, b, (ib + ((sb * 0 : U32) + (bcol * 1 : U32) : U32) : U32)))
105def 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:
106 match rv:
107 case Tuple{a, +x3_3}:
108 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, a, Array.get(F32, b, (ib + ((sb * 0 : U32) + (bcol * 0 : U32) : U32) : U32)))
109def 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:
110 match rv:
111 case Tuple{a, +x3_2}:
112 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)
113def 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:
114 match rv:
115 case Tuple{a, +x3_1}:
116 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)
117def 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:
118 match rv:
119 case Tuple{a, +x3_0}:
120 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)
121def 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:
122 match rv:
123 case Tuple{a, +x2_3}:
124 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)
125def 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:
126 match rv:
127 case Tuple{a, +x2_2}:
128 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)
129def 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:
130 match rv:
131 case Tuple{a, +x2_1}:
132 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)
133def 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:
134 match rv:
135 case Tuple{a, +x2_0}:
136 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)
137def 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:
138 match rv:
139 case Tuple{a, +x1_3}:
140 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)
141def 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:
142 match rv:
143 case Tuple{a, +x1_2}:
144 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)
145def 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:
146 match rv:
147 case Tuple{a, +x1_1}:
148 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)
149def 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:
150 match rv:
151 case Tuple{a, +x1_0}:
152 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)
153def 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:
154 match rv:
155 case Tuple{a, +x0_3}:
156 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)
157def 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:
158 match rv:
159 case Tuple{a, +x0_2}:
160 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)
161def 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:
162 match rv:
163 case Tuple{a, +x0_1}:
164 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)
165def gr_0(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
166 match rv:
167 case Tuple{a, +x0_0}:
168 gr_1(ia, ib, arow, sa, sb, bcol, x0_0, Array.get(F32, a, (ia + ((sa * 0 : U32) + (arow * 1 : U32) : U32) : U32)), b)
169# the turn whose A values start at ia and B values at ib
170def turn(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>) -> GR:
171 gr_0(ia, ib, arow, sa, sb, bcol, Array.get(F32, a, ia), b)
172
173# the tile's cells after the four sums are combined, row by row
174type T is Type:
175 T{a: Array<F32>, b: Array<F32>, c0_0: F32, c0_1: F32, c0_2: F32, c0_3: F32, c1_0: F32, c1_1: F32, c1_2: F32, c1_3: F32, c2_0: F32, c2_1: F32, c2_2: F32, c2_3: F32, c3_0: F32, c3_1: F32, c3_2: F32, c3_3: F32}
176
177# kg more full turns; g holds the current turn's values. The last call (kg = 0) adds the turn that holds the steps
178# past the last full turn (rem = k mod 4 of them), counting only those, and combines the sums.
179@unsafe
180def 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, +s0_0_1: F32, +s1_0_1: F32, +s2_0_1: F32, +s3_0_1: F32, +s0_0_2: F32, +s1_0_2: F32, +s2_0_2: F32, +s3_0_2: F32, +s0_0_3: F32, +s1_0_3: F32, +s2_0_3: F32, +s3_0_3: F32, +s0_1_0: F32, +s1_1_0: F32, +s2_1_0: F32, +s3_1_0: F32, +s0_1_1: F32, +s1_1_1: F32, +s2_1_1: F32, +s3_1_1: F32, +s0_1_2: F32, +s1_1_2: F32, +s2_1_2: F32, +s3_1_2: F32, +s0_1_3: F32, +s1_1_3: F32, +s2_1_3: F32, +s3_1_3: F32, +s0_2_0: F32, +s1_2_0: F32, +s2_2_0: F32, +s3_2_0: F32, +s0_2_1: F32, +s1_2_1: F32, +s2_2_1: F32, +s3_2_1: F32, +s0_2_2: F32, +s1_2_2: F32, +s2_2_2: F32, +s3_2_2: F32, +s0_2_3: F32, +s1_2_3: F32, +s2_2_3: F32, +s3_2_3: F32, +s0_3_0: F32, +s1_3_0: F32, +s2_3_0: F32, +s3_3_0: F32, +s0_3_1: F32, +s1_3_1: F32, +s2_3_1: F32, +s3_3_1: F32, +s0_3_2: F32, +s1_3_2: F32, +s2_3_2: F32, +s3_3_2: F32, +s0_3_3: F32, +s1_3_3: F32, +s2_3_3: F32, +s3_3_3: F32, g: GR) -> T:
181 match kg g:
182 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, +y0_0, +y0_1, +y0_2, +y0_3, +y1_0, +y1_1, +y1_2, +y1_3, +y2_0, +y2_1, +y2_2, +y2_3, +y3_0, +y3_1, +y3_2, +y3_3}:
183 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)), 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)), TA.t8_comb(rot, (s0_0_2 + pk(U32.is_lt(0, rem), (x0_0 * y0_2 : F32)) : F32), (s1_0_2 + pk(U32.is_lt(1, rem), (x1_0 * y1_2 : F32)) : F32), (s2_0_2 + pk(U32.is_lt(2, rem), (x2_0 * y2_2 : F32)) : F32), (s3_0_2 + pk(U32.is_lt(3, rem), (x3_0 * y3_2 : F32)) : F32)), TA.t8_comb(rot, (s0_0_3 + pk(U32.is_lt(0, rem), (x0_0 * y0_3 : F32)) : F32), (s1_0_3 + pk(U32.is_lt(1, rem), (x1_0 * y1_3 : F32)) : F32), (s2_0_3 + pk(U32.is_lt(2, rem), (x2_0 * y2_3 : F32)) : F32), (s3_0_3 + pk(U32.is_lt(3, rem), (x3_0 * y3_3 : 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)), 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)), TA.t8_comb(rot, (s0_1_2 + pk(U32.is_lt(0, rem), (x0_1 * y0_2 : F32)) : F32), (s1_1_2 + pk(U32.is_lt(1, rem), (x1_1 * y1_2 : F32)) : F32), (s2_1_2 + pk(U32.is_lt(2, rem), (x2_1 * y2_2 : F32)) : F32), (s3_1_2 + pk(U32.is_lt(3, rem), (x3_1 * y3_2 : F32)) : F32)), TA.t8_comb(rot, (s0_1_3 + pk(U32.is_lt(0, rem), (x0_1 * y0_3 : F32)) : F32), (s1_1_3 + pk(U32.is_lt(1, rem), (x1_1 * y1_3 : F32)) : F32), (s2_1_3 + pk(U32.is_lt(2, rem), (x2_1 * y2_3 : F32)) : F32), (s3_1_3 + pk(U32.is_lt(3, rem), (x3_1 * y3_3 : 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)), 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)), TA.t8_comb(rot, (s0_2_2 + pk(U32.is_lt(0, rem), (x0_2 * y0_2 : F32)) : F32), (s1_2_2 + pk(U32.is_lt(1, rem), (x1_2 * y1_2 : F32)) : F32), (s2_2_2 + pk(U32.is_lt(2, rem), (x2_2 * y2_2 : F32)) : F32), (s3_2_2 + pk(U32.is_lt(3, rem), (x3_2 * y3_2 : F32)) : F32)), TA.t8_comb(rot, (s0_2_3 + pk(U32.is_lt(0, rem), (x0_2 * y0_3 : F32)) : F32), (s1_2_3 + pk(U32.is_lt(1, rem), (x1_2 * y1_3 : F32)) : F32), (s2_2_3 + pk(U32.is_lt(2, rem), (x2_2 * y2_3 : F32)) : F32), (s3_2_3 + pk(U32.is_lt(3, rem), (x3_2 * y3_3 : 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)), 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)), TA.t8_comb(rot, (s0_3_2 + pk(U32.is_lt(0, rem), (x0_3 * y0_2 : F32)) : F32), (s1_3_2 + pk(U32.is_lt(1, rem), (x1_3 * y1_2 : F32)) : F32), (s2_3_2 + pk(U32.is_lt(2, rem), (x2_3 * y2_2 : F32)) : F32), (s3_3_2 + pk(U32.is_lt(3, rem), (x3_3 * y3_2 : F32)) : F32)), TA.t8_comb(rot, (s0_3_3 + pk(U32.is_lt(0, rem), (x0_3 * y0_3 : F32)) : F32), (s1_3_3 + pk(U32.is_lt(1, rem), (x1_3 * y1_3 : F32)) : F32), (s2_3_3 + pk(U32.is_lt(2, rem), (x2_3 * y2_3 : F32)) : F32), (s3_3_3 + pk(U32.is_lt(3, rem), (x3_3 * y3_3 : F32)) : F32))}
184 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, +y0_0, +y0_1, +y0_2, +y0_3, +y1_0, +y1_1, +y1_2, +y1_3, +y2_0, +y2_1, +y2_2, +y2_3, +y3_0, +y3_1, +y3_2, +y3_3}:
185 mm(q, (ia + (sa * 4 : U32) : U32), (ib + (sb * 4 : 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), (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), (s0_0_2 + (x0_0 * y0_2 : F32) : F32), (s1_0_2 + (x1_0 * y1_2 : F32) : F32), (s2_0_2 + (x2_0 * y2_2 : F32) : F32), (s3_0_2 + (x3_0 * y3_2 : F32) : F32), (s0_0_3 + (x0_0 * y0_3 : F32) : F32), (s1_0_3 + (x1_0 * y1_3 : F32) : F32), (s2_0_3 + (x2_0 * y2_3 : F32) : F32), (s3_0_3 + (x3_0 * y3_3 : 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), (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), (s0_1_2 + (x0_1 * y0_2 : F32) : F32), (s1_1_2 + (x1_1 * y1_2 : F32) : F32), (s2_1_2 + (x2_1 * y2_2 : F32) : F32), (s3_1_2 + (x3_1 * y3_2 : F32) : F32), (s0_1_3 + (x0_1 * y0_3 : F32) : F32), (s1_1_3 + (x1_1 * y1_3 : F32) : F32), (s2_1_3 + (x2_1 * y2_3 : F32) : F32), (s3_1_3 + (x3_1 * y3_3 : 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), (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), (s0_2_2 + (x0_2 * y0_2 : F32) : F32), (s1_2_2 + (x1_2 * y1_2 : F32) : F32), (s2_2_2 + (x2_2 * y2_2 : F32) : F32), (s3_2_2 + (x3_2 * y3_2 : F32) : F32), (s0_2_3 + (x0_2 * y0_3 : F32) : F32), (s1_2_3 + (x1_2 * y1_3 : F32) : F32), (s2_2_3 + (x2_2 * y2_3 : F32) : F32), (s3_2_3 + (x3_2 * y3_3 : 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), (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), (s0_3_2 + (x0_3 * y0_2 : F32) : F32), (s1_3_2 + (x1_3 * y1_2 : F32) : F32), (s2_3_2 + (x2_3 * y2_2 : F32) : F32), (s3_3_2 + (x3_3 * y3_2 : F32) : F32), (s0_3_3 + (x0_3 * y0_3 : F32) : F32), (s1_3_3 + (x1_3 * y1_3 : F32) : F32), (s2_3_3 + (x2_3 * y2_3 : F32) : F32), (s3_3_3 + (x3_3 * y3_3 : F32) : F32), turn((ia + (sa * 4 : U32) : U32), (ib + (sb * 4 : U32) : U32), arow, sa, sb, bcol, a, b))
186
187# writes cell (i, j) of C when it is inside n x m
188def wr(ok: Bool, +ac: Bool, +idx: U32, +x: F32, c: Array<F32>) -> Array<F32>:
189 match ok:
190 case True{}: TA.cwr(ac, idx, x, c)
191 case False{}: c
192
193# the tile at rows i0.., columns j0.. of C, from the loop's result; A and B are dropped (the caller holds them)
194@unsafe
195def put(+ac: Bool, +n: U32, +m: U32, +i0: U32, +j0: U32, c: Array<F32>, t: T) -> Array<F32>:
196 match t:
197 case T{a, b, +c0_0, +c0_1, +c0_2, +c0_3, +c1_0, +c1_1, +c1_2, +c1_3, +c2_0, +c2_1, +c2_2, +c2_3, +c3_0, +c3_1, +c3_2, +c3_3}:
198 wr(Bool.and(U32.is_lt((i0 + 3 : U32), n), U32.is_lt((j0 + 3 : U32), m)), ac, (((i0 + 3 : U32) * m : U32) + (j0 + 3 : U32) : U32), c3_3, wr(Bool.and(U32.is_lt((i0 + 3 : U32), n), U32.is_lt((j0 + 2 : U32), m)), ac, (((i0 + 3 : U32) * m : U32) + (j0 + 2 : U32) : U32), c3_2, 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 + 3 : U32), m)), ac, (((i0 + 2 : U32) * m : U32) + (j0 + 3 : U32) : U32), c2_3, wr(Bool.and(U32.is_lt((i0 + 2 : U32), n), U32.is_lt((j0 + 2 : U32), m)), ac, (((i0 + 2 : U32) * m : U32) + (j0 + 2 : U32) : U32), c2_2, 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 + 3 : U32), m)), ac, (((i0 + 1 : U32) * m : U32) + (j0 + 3 : U32) : U32), c1_3, wr(Bool.and(U32.is_lt((i0 + 1 : U32), n), U32.is_lt((j0 + 2 : U32), m)), ac, (((i0 + 1 : U32) * m : U32) + (j0 + 2 : U32) : U32), c1_2, 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 + 3 : U32), m)), ac, (((i0 + 0 : U32) * m : U32) + (j0 + 3 : U32) : U32), c0_3, wr(Bool.and(U32.is_lt((i0 + 0 : U32), n), U32.is_lt((j0 + 2 : U32), m)), ac, (((i0 + 0 : U32) * m : U32) + (j0 + 2 : U32) : U32), c0_2, 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))))))))))))))))
199
200@unsafe
201def tile(+e: E, +i0: U32, +j0: U32, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> Array<F32>:
202 match e:
203 case E{+n, +m, +k, +tcols, +tiles, +arow, +sa, +sb, +bcol, +kg, +rot, +rem, +ac}:
204 put(ac, n, m, i0, j0, c, mm(kg, (i0 * arow : U32), (j0 * bcol : 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((i0 * arow : U32), (j0 * bcol : U32), arow, sa, sb, bcol, a, b)))
205
206# leaf t computes tile t (row-major over the tiles); leaves past the last tile write nothing
207@unsafe
208def leaf_if(ok: Bool, +e: E, +t: U32, +tcols: U32, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> Array<F32>:
209 match ok:
210 case True{}: tile(e, (U32.div(t, tcols) * 4 : U32), (U32.mod(t, tcols) * 4 : U32), a, b, c)
211 case False{}: c
212
213@unsafe
214def leaf(+e: E, +t: U32, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> Array<F32>:
215 match e:
216 case E{+n, +m, +k, +tcols, +tiles, +arow, +sa, +sb, +bcol, +kg, +rot, +rem, +ac}:
217 leaf_if(U32.is_lt(t, tiles), e, t, tcols, a, b, c)
218
219@unsafe
220def tree(+e: E, d: Nat, +t: U32, +a: Array<F32>, +b: Array<F32>, +c: Array<F32>) -> Array<F32>:
221 match d:
222 case 0n: leaf(e, t, a, b, c)
223 case 1n++q:
224 l r = tree(e, q, (t * 2 : U32), a, b, c) tree(e, q, ((t * 2 : U32) + 1 : U32), a, b, c)
225 Array.join(F32, l, r)
226
227# C made inside the bang: the root of the ! call allocates it on the device. Made on the host instead, a new C often
228# reuses a block whose pages sit on the device, so zeroing it on the host pulls them back over PCIe and the bang moves
229# them again (a 256 x 1536 product: 9.2 ms against 6.0, measured 2026-10-18). The leaves write every cell of C.
230@unsafe
231def tree_new(+e: E, d: Nat, +a: Array<F32>, +b: Array<F32>, +cd: Nat) -> Array<F32>:
232 tree(e, d, 0, a, b, Array.new(F32, cd, 0.0))
233
234@unsafe
235def bang_new(+e: E, d: Nat, +a: Array<F32>, +b: Array<F32>, +cd: Nat) -> Array<F32>:
236 tree_new!(e, d, a, b, cd)
237
238# one bang over the tiles; A and B are shared with every leaf and come back unchanged
239@unsafe
240def bang(+e: E, d: Nat, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> Array<F32>:
241 tree!(e, d, 0, a, b, c)
242
243type ABC is Type:
244 ABC{a: Array<F32>, b: Array<F32>, c: Array<F32>}
245
246@unsafe
247def run_e(+e: E, +tiles: U32, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> ABC:
248 ABC{a, b, bang(e, TA.cap_depth(U32.to_nat(tiles)), a, b, c)}
249
250# the product's sizes and strides
251def e_of(+ac: Bool, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32) -> E:
252 E{U32.from_nat(n), U32.from_nat(m), U32.from_nat(k), U32.from_nat(Nat.div(Nat.add(m, 3n), 4n)), U32.from_nat(Nat.mul(Nat.div(Nat.add(n, 3n), 4n), Nat.div(Nat.add(m, 3n), 4n))), arow, sa, sb, bcol, Nat.div(k, 4n), Nat.mod(k, 4n), U32.from_nat(Nat.mod(k, 4n)), ac}
253
254def tiles(+n: Nat, +m: Nat) -> Nat:
255 Nat.mul(Nat.div(Nat.add(n, 3n), 4n), Nat.div(Nat.add(m, 3n), 4n))
256
257# C (n x m) = A · B with k steps into a new C made on the device
258@unsafe
259def gemm_new(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +a: Array<F32>, +b: Array<F32>) -> ABC:
260 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)))}
261
262# 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]
263@unsafe
264def 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:
265 run_e(e_of(ac, n, m, k, arow, sa, sb, bcol), U32.from_nat(tiles(n, m)), a, b, c)
266
267def mm_fin(-n: Nat, -k: Nat, -m: Nat, g: ABC) -> TA.MMul<n, k, m>:
268 match g:
269 case ABC{a, b, c}: TA.MMul{TA.Mat{a}, TA.Mat{b}, TA.Mat{c}}
270
271def mm_nt_fin(-n: Nat, -k: Nat, -m: Nat, g: ABC) -> TA.MMulNT<n, k, m>:
272 match g:
273 case ABC{a, b, c}: TA.MMulNT{TA.Mat{a}, TA.Mat{b}, TA.Mat{c}}
274
275def mm_tn_fin(-n: Nat, -k: Nat, -m: Nat, g: ABC) -> TA.MMulTN<n, k, m>:
276 match g:
277 case ABC{a, b, c}: TA.MMulTN{TA.Mat{a}, TA.Mat{b}, TA.Mat{c}}
278
279# C = A · B (n x k) · (k x m), tensor-array's TA.Mat.matmul bits
280@unsafe
281def matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat<n, k>, b: TA.Mat<k, m>) -> TA.MMul<n, k, m>:
282 match a b:
283 case TA.Mat{x} TA.Mat{y}:
284 mm_fin(n, k, m, gemm_new(n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, x, y))
285
286# C + A · B, into C (n x m)
287@unsafe
288def 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>:
289 match a b c:
290 case TA.Mat{x} TA.Mat{y} TA.Mat{z}:
291 mm_fin(n, k, m, gemm(True{}, n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, x, y, z))
292
293# C = A · Bᵀ (n x k) · (m x k)ᵀ: for weights stored with one row per output
294@unsafe
295def 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>:
296 match a b:
297 case TA.Mat{x} TA.Mat{y}:
298 mm_nt_fin(n, k, m, gemm_new(n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), x, y))
299
300# C + A · Bᵀ, into C (n x m)
301@unsafe
302def 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>:
303 match a b c:
304 case TA.Mat{x} TA.Mat{y} TA.Mat{z}:
305 mm_nt_fin(n, k, m, gemm(True{}, n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), x, y, z))
306
307# C = Aᵀ · B (k x n)ᵀ · (k x m): weight gradient, Xᵀ · dY
308@unsafe
309def 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>:
310 match a b:
311 case TA.Mat{x} TA.Mat{y}:
312 mm_tn_fin(n, k, m, gemm_new(n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, x, y))
313
314# C + Aᵀ · B, into C (n x m): dW += Xᵀ · dY
315@unsafe
316def 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>:
317 match a b c:
318 case TA.Mat{x} TA.Mat{y} TA.Mat{z}:
319 mm_tn_fin(n, k, m, gemm(True{}, n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, x, y, z))
320
321# =====================================================================================================================
322# 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
323# 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
324# the outputs. The leaf code is a template (~f), written by the caller in plain Bend: it takes every array as an ordinary
325# argument (it reads them with Array.get and may write the outputs), so it needs no @unsafe; the sharing is here, where
326# every node of the tree holds a handle to each array (Array.fork in Base, upstream's one-block pattern) and joins the
327# outputs (Array.join). runK takes the outputs from the caller and writes them in place (an optimizer's weights and
328# moments); runK_new makes them inside the bang with 2^c1.. cells, since an array made on the host pulls device pages
329# back (see bang_new above). Both return the inputs unchanged with the outputs (RK). An unused input is any array, such
330# as none(Unit{}).
331#
332# Not checked, as for the products: that a leaf writes only its own cells of the outputs (two leaves writing one cell
333# race) and never an input. Leaves past the work must return their outputs unchanged; share(t, n, l) is how many of the
334# cells t, t + l, t + 2l, .. lie below n (leaf t's cells when n cells are dealt round-robin, so neighbouring lanes read
335# neighbouring cells), and depth(n) is a tree for n cells: 2^14 leaves once there are that many, else the fewest that
336# cover them. The leaves run in a run-time order inside each group of 128 (bendlang/bend's CUDA lowering), so a leaf's
337# result must not depend on another's. Every number a leaf computes is the caller's own code, so the bits are the
338# 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
339# on the device takes the float overload; bend-ml's docs/upstream/gpu-math-overloads.md).
340# =====================================================================================================================
341
342def depth2(big: Bool, +n: Nat) -> Nat:
343 match big:
344 case True{}: 14n
345 case False{}: TA.cap_depth(n)
346
347# leaves for n cells: 2^14 when there are that many, else the fewest powers of two that cover them
348def depth(+n: Nat) -> Nat:
349 depth2(Nat.is_le(16384n, n), n)
350
351# how many of the cells t, t + l, t + 2l, .. lie below n
352def share(+t: U32, +n: U32, +l: U32) -> Nat:
353 U32.to_nat(Bool.pick(U32, U32.is_lt(t, n), U32.div(((n - t : U32) + (l - 1 : U32) : U32), l), 0))
354
355# an array of one cell, for an unused input
356def none(u: Unit) -> Array<F32>:
357 Array.new(F32, 0n, 0.0)
358
359type O2 is Type:
360 O2{a: Array<F32>, b: Array<F32>}
361
362type O3 is Type:
363 O3{a: Array<F32>, b: Array<F32>, c: Array<F32>}
364
365type O4 is Type:
366 O4{a: Array<F32>, b: Array<F32>, c: Array<F32>, d: Array<F32>}
367
368type R1 is Type:
369 R1{x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, a: Array<F32>}
370
371type R2 is Type:
372 R2{x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: O2}
373
374type R3 is Type:
375 R3{x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: O3}
376
377type R4 is Type:
378 R4{x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: O4}
379
380@unsafe
381def join2(x: O2, y: O2) -> O2:
382 match x y:
383 case O2{a, b} O2{a2, b2}: O2{Array.join(F32, a, a2), Array.join(F32, b, b2)}
384
385@unsafe
386def join3(x: O3, y: O3) -> O3:
387 match x y:
388 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)}
389
390@unsafe
391def join4(x: O4, y: O4) -> O4:
392 match x y:
393 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)}
394
395@unsafe
396def 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>:
397 match d:
398 case 0n: f(t, l, e, x1, x2, x3, x4, oa)
399 case 1n+q:
400 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)
401 Array.join(F32, x, y)
402
403@unsafe
404def 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>:
405 tree1(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa)
406
407@unsafe
408def 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>:
409 root1(~E, ~f, d, e, x1, x2, x3, x4, Array.new(F32, ca, 0.0))
410
411@unsafe
412def 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:
413 R1{x1, x2, x3, x4, root1!(~E, ~f, d, e, x1, x2, x3, x4, oa)}
414
415@unsafe
416def 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:
417 R1{x1, x2, x3, x4, root1_new!(~E, ~f, d, e, x1, x2, x3, x4, ca)}
418
419# one bang over 2^d leaves with 1 output written in place
420@unsafe
421def 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:
422 go1(~E, ~f, d, e, x1, x2, x3, x4, oa)
423
424# one bang over 2^d leaves with 1 new output of 2^c.. cells, made on the device
425@unsafe
426def 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:
427 go1_new(~E, ~f, d, e, x1, x2, x3, x4, ca)
428
429@unsafe
430def 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:
431 match d:
432 case 0n: f(t, l, e, x1, x2, x3, x4, O2{oa, ob})
433 case 1n+q:
434 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)
435 join2(x, y)
436
437@unsafe
438def 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:
439 tree2(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa, ob)
440
441@unsafe
442def 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:
443 root2(~E, ~f, d, e, x1, x2, x3, x4, Array.new(F32, ca, 0.0), Array.new(F32, cb, 0.0))
444
445@unsafe
446def 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:
447 R2{x1, x2, x3, x4, root2!(~E, ~f, d, e, x1, x2, x3, x4, oa, ob)}
448
449@unsafe
450def 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:
451 R2{x1, x2, x3, x4, root2_new!(~E, ~f, d, e, x1, x2, x3, x4, ca, cb)}
452
453# one bang over 2^d leaves with 2 outputs written in place
454@unsafe
455def 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:
456 go2(~E, ~f, d, e, x1, x2, x3, x4, oa, ob)
457
458# one bang over 2^d leaves with 2 new outputs of 2^c.. cells, made on the device
459@unsafe
460def 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:
461 go2_new(~E, ~f, d, e, x1, x2, x3, x4, ca, cb)
462
463@unsafe
464def 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:
465 match d:
466 case 0n: f(t, l, e, x1, x2, x3, x4, O3{oa, ob, oc})
467 case 1n+q:
468 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)
469 join3(x, y)
470
471@unsafe
472def 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:
473 tree3(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa, ob, oc)
474
475@unsafe
476def 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:
477 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))
478
479@unsafe
480def 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:
481 R3{x1, x2, x3, x4, root3!(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc)}
482
483@unsafe
484def 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:
485 R3{x1, x2, x3, x4, root3_new!(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc)}
486
487# one bang over 2^d leaves with 3 outputs written in place
488@unsafe
489def 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:
490 go3(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc)
491
492# one bang over 2^d leaves with 3 new outputs of 2^c.. cells, made on the device
493@unsafe
494def 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:
495 go3_new(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc)
496
497@unsafe
498def 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:
499 match d:
500 case 0n: f(t, l, e, x1, x2, x3, x4, O4{oa, ob, oc, od})
501 case 1n+q:
502 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)
503 join4(x, y)
504
505@unsafe
506def 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:
507 tree4(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa, ob, oc, od)
508
509@unsafe
510def 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:
511 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))
512
513@unsafe
514def 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:
515 R4{x1, x2, x3, x4, root4!(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc, od)}
516
517@unsafe
518def 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:
519 R4{x1, x2, x3, x4, root4_new!(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc, cd)}
520
521# one bang over 2^d leaves with 4 outputs written in place
522@unsafe
523def 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:
524 go4(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc, od)
525
526# one bang over 2^d leaves with 4 new outputs of 2^c.. cells, made on the device
527@unsafe
528def 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:
529 go4_new(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc, cd)
530
531# ---- tensor-array 0.1.14.0's precise product on the GPU (opt-in) ----
532# TA.Mat.matmul_precise's composition, each block one bang: the first block's product made on the device, each next
533# one added into C (the accumulating kernel adds a cell once, as tensor-array's cwr), so the same bits as tensor-array's
534# (tests/bits.bend). The slices of A and B are copied on the host, which pulls their pages back from the device: this
535# is for accuracy, not speed.
536
537def abc_g(a: Array<F32>, b: Array<F32>, r: ABC) -> TA.G:
538 match r:
539 case ABC{x, y, c}: TA.G{a, b, c}
540
541@unsafe
542def mp_s3(+ac: Bool, +n: Nat, +m: Nat, +w: Nat, +w32: U32, +m32: U32, a: Array<F32>, as: Array<F32>, c: Array<F32>, sb: TA.SL) -> TA.G:
543 match ac sb:
544 case True{} TA.SL{b, bs}: abc_g(a, b, gemm(True{}, n, m, w, w32, 1, m32, 1, as, bs, c))
545 case False{} TA.SL{b, bs}: abc_g(a, b, gemm_new(n, m, w, w32, 1, m32, 1, as, bs))
546
547@unsafe
548def mp_s2(+ac: Bool, +n: Nat, +m: Nat, +w: Nat, +w32: U32, +m32: U32, +o: U32, c: Array<F32>, b: Array<F32>, sa: TA.SL) -> TA.G:
549 match sa:
550 case TA.SL{a, as}: mp_s3(ac, n, m, w, w32, m32, a, as, c, TA.sl_row(Nat.mul(w, m), (o * m32 : U32), 0, Array.new(F32, TA.cap_depth(Nat.mul(w, m)), 0.0), Array.get(F32, b, (o * m32 : U32))))
551
552@unsafe
553def mp_step(+ac: Bool, +n: Nat, +m: Nat, +k32: U32, +bi: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> TA.G:
554 mp_s2(ac, 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, TA.cap_depth(Nat.mul(n, TA.blk_len(bi, k32))), 0.0)}))
555
556@unsafe
557def mp_go(nb: Nat, +bi: U32, +n: Nat, +m: Nat, +k32: U32, g: TA.G) -> TA.G:
558 match nb g:
559 case 0n g: g
560 case 1n++q TA.G{a, b, c}: mp_go(q, (bi + 1 : U32), n, m, k32, mp_step(U32.is_lt(0, bi), n, m, k32, bi, a, b, c))
561
562@unsafe
563def 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>:
564 match long a b:
565 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, n, m, U32.from_nat(k), TA.G{x, y, none(Unit{})}))
566 case False{} a b: matmul(n, k, m, par, a, b)
567
568# C = A · B with k summed in blocks of 1024: TA.Mat.matmul_precise's bits (opt-in)
569@unsafe
570def 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>:
571 mm_precise2(Nat.is_lt(1024n, k), n, k, m, par, a, b)