main.bend

bend-ml-tensor-gpu@0.1.5.1 · 783 lines · 147 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.5.1/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 `+RTS --gpu off -RTS`, 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# Counter-based random numbers (fill_uniform, fill_gauss, uniform_mat, gauss_mat and their _at forms; the last section):
21# n uniforms or Gaussians of a counter made in one bang, each a pure function of (counter, index), with the same bits on
22# every lane.
23#
24# Not proved, and marked so: Array.fork and Array.join are @unsafe in Base, so the defs below that call them, or take
25# a shared array, are @unsafe, and `bend main.bend` names every def that relies on them. Nothing checks that a leaf
26# writes only its own tile of C, or never writes A or B; the tiles are disjoint by construction (leaf t writes tile t).
27# All of it is Bend: no foreign code. Written by bend-ml's scripts/gen_tensor_gpu.py (4 x 4 tiles).
28# bend-lint: disable-file bend-lint-rules/performance/gpu_fit -- the bangs carry the arrays to the device by design; that is the package not-expected-because: products timed on the GPU against 16 CPU threads (8.4 against 12.2 ms, 11-13 against 55-56 ms, NOTES 2026-10-18); a caller whose kernel does not win runs the same bang on the CPU threads
29import Base
30import bend-ml-tensor-array@0.1.19.0/main.bend as TA
31
32# 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
33# (oa, ob: 0 except for matmul_precise's blocks)
34type E is Data:
35  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}
36
37# a product that counts only when its step is inside k (the last turn's steps past k add +0.0)
38def pk(ok: Bool, +x: F32) -> F32:
39  match ok:
40    case True{}: x
41    case False{}: 0.0
42
43# the values of one turn: A's 4 rows and B's 4 columns at four consecutive steps, read through the arrays in order
44type GR is Type:
45  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}
46
47def 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:
48  match rv:
49    case Tuple{b, +y3_3}:
50      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}
51def 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:
52  match rv:
53    case Tuple{b, +y3_2}:
54      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)))
55def 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:
56  match rv:
57    case Tuple{b, +y3_1}:
58      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)))
59def 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:
60  match rv:
61    case Tuple{b, +y3_0}:
62      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)))
63def 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:
64  match rv:
65    case Tuple{b, +y2_3}:
66      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)))
67def 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:
68  match rv:
69    case Tuple{b, +y2_2}:
70      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)))
71def 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:
72  match rv:
73    case Tuple{b, +y2_1}:
74      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)))
75def 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:
76  match rv:
77    case Tuple{b, +y2_0}:
78      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)))
79def 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:
80  match rv:
81    case Tuple{b, +y1_3}:
82      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)))
83def 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:
84  match rv:
85    case Tuple{b, +y1_2}:
86      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)))
87def 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:
88  match rv:
89    case Tuple{b, +y1_1}:
90      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)))
91def 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:
92  match rv:
93    case Tuple{b, +y1_0}:
94      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)))
95def 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:
96  match rv:
97    case Tuple{b, +y0_3}:
98      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)))
99def 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:
100  match rv:
101    case Tuple{b, +y0_2}:
102      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)))
103def 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:
104  match rv:
105    case Tuple{b, +y0_1}:
106      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)))
107def 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:
108  match rv:
109    case Tuple{b, +y0_0}:
110      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)))
111def 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:
112  match rv:
113    case Tuple{a, +x3_3}:
114      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)))
115def 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:
116  match rv:
117    case Tuple{a, +x3_2}:
118      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)
119def 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:
120  match rv:
121    case Tuple{a, +x3_1}:
122      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)
123def 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:
124  match rv:
125    case Tuple{a, +x3_0}:
126      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)
127def 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:
128  match rv:
129    case Tuple{a, +x2_3}:
130      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)
131def 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:
132  match rv:
133    case Tuple{a, +x2_2}:
134      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)
135def 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:
136  match rv:
137    case Tuple{a, +x2_1}:
138      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)
139def 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:
140  match rv:
141    case Tuple{a, +x2_0}:
142      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)
143def 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:
144  match rv:
145    case Tuple{a, +x1_3}:
146      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)
147def 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:
148  match rv:
149    case Tuple{a, +x1_2}:
150      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)
151def 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:
152  match rv:
153    case Tuple{a, +x1_1}:
154      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)
155def 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:
156  match rv:
157    case Tuple{a, +x1_0}:
158      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)
159def 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:
160  match rv:
161    case Tuple{a, +x0_3}:
162      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)
163def 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:
164  match rv:
165    case Tuple{a, +x0_2}:
166      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)
167def 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:
168  match rv:
169    case Tuple{a, +x0_1}:
170      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)
171def gr_0(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, rv: Array<F32> & F32, b: Array<F32>) -> GR:
172  match rv:
173    case Tuple{a, +x0_0}:
174      gr_1(ia, ib, arow, sa, sb, bcol, x0_0, Array.get(F32, a, (ia + ((sa * 0 : U32) + (arow * 1 : U32) : U32) : U32)), b)
175# the turn whose A values start at ia and B values at ib
176def turn(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>) -> GR:
177  gr_0(ia, ib, arow, sa, sb, bcol, Array.get(F32, a, ia), b)
178
179# the tile's cells after the four sums are combined, row by row
180type T is Type:
181  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}
182
183# kg more full turns; g holds the current turn's values. The last call (kg = 0) adds the turn that holds the steps
184# past the last full turn (rem = k mod 4 of them), counting only those, and combines the sums.
185@unsafe
186def 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:
187  match kg g:
188    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}:
189      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))}
190    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}:
191      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))
192
193# writes cell (i, j) of C when it is inside n x m
194def wr(ok: Bool, +ac: Bool, +idx: U32, +x: F32, c: Array<F32>) -> Array<F32>:
195  match ok:
196    case True{}: TA.cwr(ac, idx, x, c)
197    case False{}: c
198
199# the tile at rows i0.., columns j0.. of C, from the loop's result; A and B are dropped (the caller holds them)
200@unsafe
201def put(+ac: Bool, +n: U32, +m: U32, +i0: U32, +j0: U32, c: Array<F32>, t: T) -> Array<F32>:
202  match t:
203    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}:
204      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))))))))))))))))
205
206@unsafe
207def tile(+e: E, +i0: U32, +j0: U32, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> Array<F32>:
208  match e:
209    case E{+n, +m, +k, +tcols, +tiles, +arow, +sa, +sb, +bcol, +kg, +rot, +rem, +ac, +oa, +ob}:
210      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)))
211
212# leaf t computes tile t (row-major over the tiles); leaves past the last tile write nothing
213@unsafe
214def leaf_if(ok: Bool, +e: E, +t: U32, +tcols: U32, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> Array<F32>:
215  match ok:
216    case True{}: tile(e, (U32.div(t, tcols) * 4 : U32), (U32.mod(t, tcols) * 4 : U32), a, b, c)
217    case False{}: c
218
219@unsafe
220def leaf(+e: E, +t: U32, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> Array<F32>:
221  match e:
222    case E{+n, +m, +k, +tcols, +tiles, +arow, +sa, +sb, +bcol, +kg, +rot, +rem, +ac, +oa, +ob}:
223      leaf_if(U32.is_lt(t, tiles), e, t, tcols, a, b, c)
224
225@unsafe
226def tree(+e: E, d: Nat, +t: U32, +a: Array<F32>, +b: Array<F32>, +c: Array<F32>) -> Array<F32>:
227  match d:
228    case 0n: leaf(e, t, a, b, c)
229    case 1n++q:
230      l r = tree(e, q, (t * 2 : U32), a, b, c) tree(e, q, ((t * 2 : U32) + 1 : U32), a, b, c)
231      Array.join(F32, l, r)
232
233# C made inside the bang: the root of the ! call allocates it on the device. Made on the host instead, a new C often
234# reuses a block whose pages sit on the device, so zeroing it on the host pulls them back over PCIe and the bang moves
235# them again (a 256 x 1536 product: 9.2 ms against 6.0, measured 2026-10-18). The leaves write every cell of C.
236@unsafe
237def tree_new(+e: E, d: Nat, +a: Array<F32>, +b: Array<F32>, +cd: Nat) -> Array<F32>:
238  tree(e, d, 0, a, b, Array.new(F32, cd, 0.0))
239
240@unsafe
241def bang_new(+e: E, d: Nat, +a: Array<F32>, +b: Array<F32>, +cd: Nat) -> Array<F32>:
242  tree_new!(e, d, a, b, cd)
243
244# one bang over the tiles; A and B are shared with every leaf and come back unchanged
245@unsafe
246def bang(+e: E, d: Nat, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> Array<F32>:
247  tree!(e, d, 0, a, b, c)
248
249type ABC is Type:
250  ABC{a: Array<F32>, b: Array<F32>, c: Array<F32>}
251
252@unsafe
253def run_e(+e: E, +tiles: U32, +a: Array<F32>, +b: Array<F32>, c: Array<F32>) -> ABC:
254  ABC{a, b, bang(e, TA.cap_depth(U32.to_nat(tiles)), a, b, c)}
255
256# the product's sizes and strides
257def e_ofo(+ac: Bool, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +oa: U32, +ob: U32) -> E:
258  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, oa, ob}
259
260def e_of(+ac: Bool, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32) -> E:
261  e_ofo(ac, n, m, k, arow, sa, sb, bcol, 0, 0)
262
263def tiles(+n: Nat, +m: Nat) -> Nat:
264  Nat.mul(Nat.div(Nat.add(n, 3n), 4n), Nat.div(Nat.add(m, 3n), 4n))
265
266# C (n x m) = A · B with k steps into a new C made on the device
267@unsafe
268def gemm_new(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +a: Array<F32>, +b: Array<F32>) -> ABC:
269  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)))}
270
271# 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]
272@unsafe
273def 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:
274  run_e(e_of(ac, n, m, k, arow, sa, sb, bcol), U32.from_nat(tiles(n, m)), a, b, c)
275
276def mm_fin(-n: Nat, -k: Nat, -m: Nat, g: ABC) -> TA.MMul<n, k, m>:
277  match g:
278    case ABC{a, b, c}: TA.MMul{TA.Mat{a}, TA.Mat{b}, TA.Mat{c}}
279
280def mm_nt_fin(-n: Nat, -k: Nat, -m: Nat, g: ABC) -> TA.MMulNT<n, k, m>:
281  match g:
282    case ABC{a, b, c}: TA.MMulNT{TA.Mat{a}, TA.Mat{b}, TA.Mat{c}}
283
284def mm_tn_fin(-n: Nat, -k: Nat, -m: Nat, g: ABC) -> TA.MMulTN<n, k, m>:
285  match g:
286    case ABC{a, b, c}: TA.MMulTN{TA.Mat{a}, TA.Mat{b}, TA.Mat{c}}
287
288# C = A · B   (n x k) · (k x m), tensor-array's TA.Mat.matmul bits
289@unsafe
290def matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat<n, k>, b: TA.Mat<k, m>) -> TA.MMul<n, k, m>:
291  match a b:
292    case TA.Mat{x} TA.Mat{y}:
293      mm_fin(n, k, m, gemm_new(n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, x, y))
294
295# C + A · B, into C (n x m)
296@unsafe
297def 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>:
298  match a b c:
299    case TA.Mat{x} TA.Mat{y} TA.Mat{z}:
300      mm_fin(n, k, m, gemm(True{}, n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, x, y, z))
301
302# C = A · Bᵀ   (n x k) · (m x k)ᵀ: for weights stored with one row per output
303@unsafe
304def 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>:
305  match a b:
306    case TA.Mat{x} TA.Mat{y}:
307      mm_nt_fin(n, k, m, gemm_new(n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), x, y))
308
309# C + A · Bᵀ, into C (n x m)
310@unsafe
311def 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>:
312  match a b c:
313    case TA.Mat{x} TA.Mat{y} TA.Mat{z}:
314      mm_nt_fin(n, k, m, gemm(True{}, n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), x, y, z))
315
316# C = Aᵀ · B   (k x n)ᵀ · (k x m): weight gradient, Xᵀ · dY
317@unsafe
318def 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>:
319  match a b:
320    case TA.Mat{x} TA.Mat{y}:
321      mm_tn_fin(n, k, m, gemm_new(n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, x, y))
322
323# C + Aᵀ · B, into C (n x m): dW += Xᵀ · dY
324@unsafe
325def 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>:
326  match a b c:
327    case TA.Mat{x} TA.Mat{y} TA.Mat{z}:
328      mm_tn_fin(n, k, m, gemm(True{}, n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, x, y, z))
329
330# =====================================================================================================================
331# 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
332# 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
333# the outputs. The leaf code is a template (~f), written by the caller in plain Bend: it takes every array as an ordinary
334# argument (it reads them with Array.get and may write the outputs), so it needs no @unsafe; the sharing is here, where
335# every node of the tree holds a handle to each array (Array.fork in Base, upstream's one-block pattern) and joins the
336# outputs (Array.join). runK takes the outputs from the caller and writes them in place (an optimizer's weights and
337# moments); runK_new makes them inside the bang with 2^c1.. cells, since an array made on the host pulls device pages
338# back (see bang_new above). Both return the inputs unchanged with the outputs (RK). An unused input is any array, such
339# as none(Unit{}).
340#
341# Not checked, as for the products: that a leaf writes only its own cells of the outputs (two leaves writing one cell
342# race) and never an input. Leaves past the work must return their outputs unchanged; share(t, n, l) is how many of the
343# cells t, t + l, t + 2l, .. lie below n (leaf t's cells when n cells are dealt round-robin, so neighbouring lanes read
344# neighbouring cells), and depth(n) is a tree for n cells: 2^14 leaves once there are that many, else the fewest that
345# cover them. The leaves run in a run-time order inside each group of 128 (bendlang/bend's CUDA lowering), so a leaf's
346# result must not depend on another's. Every number a leaf computes is the caller's own code, so the bits are the
347# 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
348# on the device takes the float overload; the known limit of bendlang/bend#1305).
349# =====================================================================================================================
350
351def depth2(big: Bool, +n: Nat) -> Nat:
352  match big:
353    case True{}: 14n
354    case False{}: TA.cap_depth(n)
355
356# leaves for n cells: 2^14 when there are that many, else the fewest powers of two that cover them
357def depth(+n: Nat) -> Nat:
358  depth2(Nat.is_le(16384n, n), n)
359
360# how many of the cells t, t + l, t + 2l, .. lie below n
361def share(+t: U32, +n: U32, +l: U32) -> Nat:
362  U32.to_nat(Bool.pick(U32, U32.is_lt(t, n), U32.div(((n - t : U32) + (l - 1 : U32) : U32), l), 0))
363
364# an array of one cell, for an unused input
365def none(u: Unit) -> Array<F32>:
366  Array.new(F32, 0n, 0.0)
367
368type O2 is Type:
369  O2{a: Array<F32>, b: Array<F32>}
370
371type O3 is Type:
372  O3{a: Array<F32>, b: Array<F32>, c: Array<F32>}
373
374type O4 is Type:
375  O4{a: Array<F32>, b: Array<F32>, c: Array<F32>, d: Array<F32>}
376
377type R1 is Type:
378  R1{x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, a: Array<F32>}
379
380type R2 is Type:
381  R2{x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: O2}
382
383type R3 is Type:
384  R3{x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: O3}
385
386type R4 is Type:
387  R4{x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: O4}
388
389@unsafe
390def join2(x: O2, y: O2) -> O2:
391  match x y:
392    case O2{a, b} O2{a2, b2}: O2{Array.join(F32, a, a2), Array.join(F32, b, b2)}
393
394@unsafe
395def join3(x: O3, y: O3) -> O3:
396  match x y:
397    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)}
398
399@unsafe
400def join4(x: O4, y: O4) -> O4:
401  match x y:
402    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)}
403
404@unsafe
405def 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>:
406  match d:
407    case 0n: f(t, l, e, x1, x2, x3, x4, oa)
408    case 1n+q:
409      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)
410      Array.join(F32, x, y)
411
412@unsafe
413def 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>:
414  tree1(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa)
415
416@unsafe
417def 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>:
418  root1(~E, ~f, d, e, x1, x2, x3, x4, Array.new(F32, ca, 0.0))
419
420@unsafe
421def 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:
422  R1{x1, x2, x3, x4, root1!(~E, ~f, d, e, x1, x2, x3, x4, oa)}
423
424@unsafe
425def 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:
426  R1{x1, x2, x3, x4, root1_new!(~E, ~f, d, e, x1, x2, x3, x4, ca)}
427
428# one bang over 2^d leaves with 1 output written in place
429@unsafe
430def 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:
431  go1(~E, ~f, d, e, x1, x2, x3, x4, oa)
432
433# one bang over 2^d leaves with 1 new output of 2^c.. cells, made on the device
434@unsafe
435def 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:
436  go1_new(~E, ~f, d, e, x1, x2, x3, x4, ca)
437
438@unsafe
439def 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:
440  match d:
441    case 0n: f(t, l, e, x1, x2, x3, x4, O2{oa, ob})
442    case 1n+q:
443      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)
444      join2(x, y)
445
446@unsafe
447def 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:
448  tree2(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa, ob)
449
450@unsafe
451def 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:
452  root2(~E, ~f, d, e, x1, x2, x3, x4, Array.new(F32, ca, 0.0), Array.new(F32, cb, 0.0))
453
454@unsafe
455def 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:
456  R2{x1, x2, x3, x4, root2!(~E, ~f, d, e, x1, x2, x3, x4, oa, ob)}
457
458@unsafe
459def 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:
460  R2{x1, x2, x3, x4, root2_new!(~E, ~f, d, e, x1, x2, x3, x4, ca, cb)}
461
462# one bang over 2^d leaves with 2 outputs written in place
463@unsafe
464def 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:
465  go2(~E, ~f, d, e, x1, x2, x3, x4, oa, ob)
466
467# one bang over 2^d leaves with 2 new outputs of 2^c.. cells, made on the device
468@unsafe
469def 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:
470  go2_new(~E, ~f, d, e, x1, x2, x3, x4, ca, cb)
471
472@unsafe
473def 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:
474  match d:
475    case 0n: f(t, l, e, x1, x2, x3, x4, O3{oa, ob, oc})
476    case 1n+q:
477      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)
478      join3(x, y)
479
480@unsafe
481def 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:
482  tree3(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa, ob, oc)
483
484@unsafe
485def 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:
486  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))
487
488@unsafe
489def 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:
490  R3{x1, x2, x3, x4, root3!(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc)}
491
492@unsafe
493def 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:
494  R3{x1, x2, x3, x4, root3_new!(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc)}
495
496# one bang over 2^d leaves with 3 outputs written in place
497@unsafe
498def 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:
499  go3(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc)
500
501# one bang over 2^d leaves with 3 new outputs of 2^c.. cells, made on the device
502@unsafe
503def 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:
504  go3_new(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc)
505
506@unsafe
507def 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:
508  match d:
509    case 0n: f(t, l, e, x1, x2, x3, x4, O4{oa, ob, oc, od})
510    case 1n+q:
511      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)
512      join4(x, y)
513
514@unsafe
515def 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:
516  tree4(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa, ob, oc, od)
517
518@unsafe
519def 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:
520  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))
521
522@unsafe
523def 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:
524  R4{x1, x2, x3, x4, root4!(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc, od)}
525
526@unsafe
527def 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:
528  R4{x1, x2, x3, x4, root4_new!(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc, cd)}
529
530# one bang over 2^d leaves with 4 outputs written in place
531@unsafe
532def 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:
533  go4(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc, od)
534
535# one bang over 2^d leaves with 4 new outputs of 2^c.. cells, made on the device
536@unsafe
537def 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:
538  go4_new(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc, cd)
539
540# ---- tensor-array 0.1.14.0's precise product on the GPU (opt-in) ----
541# 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
542# copy): the first block's bang makes C on the device, each next one adds its cells into C as the accumulating kernel
543# does (once, as tensor-array's cwr), so every cell is ((x0 + x1) + x2) .., the same bits as tensor-array's
544# (tests/precise.bend). Only plain (not accumulating) products, as tensor-array's matmul_precise.
545
546def blk_len(+bi: U32, +k: U32) -> Nat:
547  U32.to_nat(Bool.pick(U32, U32.is_lt((k - (bi * 1024 : U32) : U32), 1024), (k - (bi * 1024 : U32) : U32), 1024))
548
549# the blocks after the first: nb more, from block bi, each added into C
550@unsafe
551def 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>:
552  match nb:
553    case 0n: c
554    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))
555
556@unsafe
557def precise_c(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +a: Array<F32>, +b: Array<F32>) -> Array<F32>:
558  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))))
559
560@unsafe
561def mm_precise3(+n: Nat, +k: Nat, +m: Nat, +a: Array<F32>, +b: Array<F32>) -> TA.MMul<n, k, m>:
562  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)})
563
564@unsafe
565def 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>:
566  match long a b:
567    case True{} TA.Mat{x} TA.Mat{y}: mm_precise3(n, k, m, x, y)
568    case False{} a b: matmul(n, k, m, par, a, b)
569
570# C = A · B with k summed in blocks of 1024: TA.Mat.matmul_precise's bits (opt-in)
571@unsafe
572def 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>:
573  mm_precise2(Nat.is_lt(1024n, k), n, k, m, par, a, b)
574
575# =====================================================================================================================
576# Counter-based random fills (prototype). fill_uniform and fill_gauss make n numbers of a counter s in one bang, each
577# number a pure function of (s, its index), so a lane computes its cells without reading anything; uniform_mat and
578# gauss_mat give them as a tensor-array Mat. The _at forms (fill_uniform_at, fill_gauss_at, uniform_mat_at, gauss_mat_at)
579# make the cells i0 .. i0 + n - 1 of the same counter, so a batch can be made in slices: their cell k is the cell
580# function at i0 + k, and the Gaussians' pairs stay global (2p, 2p + 1), so an odd i0 starts with the sine half of the
581# pair (i0 - 1, i0). The plain forms are the case i0 = 0. Indexes are U32 and wrap: i0 + k is taken mod 2^32. The
582# functions behind them (rng_uniform, rng_gauss2, rng_gauss) are plain
583# Bend that a caller's own leaf (runK) can call too, to make noise where it is used (x_t = a * x_0 + b * noise in one
584# bang) instead of in an array of its own.
585#
586# The counter hash and the uniform are mlbend-generative 0.1.0.0's (its mix32 over s * 0x9E3779B9 + i + 0x1234567A,
587# the low 24 bits plus one over 2^24, a uniform in (0, 1]), so rng_uniform gives that package's words. A Gaussian is
588# Box-Muller over the pair (2p, 2p + 1): cell 2p is r cos(2 pi u_b), cell 2p + 1 is r sin(2 pi u_b), with
589# r = sqrt(-2 ln u_a), u_a the uniform of index 2p and u_b of 2p + 1, so each pair costs one ln, one sqrt and one sine
590# and cosine. ln, sin and cos are written here with + - * / and the exponent's bits, not F32.log/sin/cos: CUDA's float
591# functions give other bits than the host's (bendlang/bend#1305), while + - * / and sqrt are
592# IEEE's on every lane and Bend fuses no multiply-add, so these numbers have the same bits on the GPU, on the CPU
593# threads and on the JS lane. They are within a few ulp of the libm values (tests/fill.bend measures it), not equal.
594# =====================================================================================================================
595
596def rng_word(x: F32) -> U32:
597  match x:
598    case F32{w}: U32{w}
599
600def rng_f32(x: U32) -> F32:
601  match x:
602    case U32{w}: F32{w}
603
604# mlbend-generative's mix32: xor-shift 16, multiply, xor-shift 13, multiply, xor-shift 16, over U32 that wraps
605def rng_mix(+z: U32) -> U32:
606  +a = U32.xor(z, U32.shrn(z, 16n))
607  +b = (a * 2246822519 : U32)
608  +c = U32.xor(b, U32.shrn(b, 13n))
609  +e = (c * 3266489917 : U32)
610  U32.xor(e, U32.shrn(e, 16n))
611
612# the hash of index i of counter s
613def rng_hash(+s: U32, +i: U32) -> U32:
614  rng_mix(((s * 2654435769 : U32) + (i + 305419898 : U32) : U32))
615
616# the uniform of index i of counter s, in (0, 1]: the hash's low 24 bits plus one, over 2^24 (exact)
617def rng_uniform(+s: U32, +i: U32) -> F32:
618  (U32.to_f32((U32.and(rng_hash(s, i), 16777215) + 1 : U32)) / 16777216.0 : F32)
619
620# ln m for m in [sqrt(1/2), sqrt(2)], from s = (m - 1) / (m + 1): ln m = 2 (s + s^3/3 + s^5/5 + s^7/7 + s^9/9)
621def rng_ln_m(+ex: F32, +m: F32) -> F32:
622  +s = ((m - 1.0 : F32) / (m + 1.0 : F32) : F32)
623  +z = (s * s : F32)
624  +c3 = (0.14285715 + (z * 0.11111111 : F32) : F32)
625  +c2 = (0.2 + (z * c3 : F32) : F32)
626  +p = (z * (0.33333334 + (z * c2 : F32) : F32) : F32)
627  +t = (s + s : F32)
628  ((ex * 0.693145751953125 : F32) + ((t + (t * p : F32) : F32) + (ex * 0.0000014286067653 : F32) : F32) : F32)
629
630# ln x for a positive normal x: x = m * 2^ex with m in [sqrt(1/2), sqrt(2)) read from the bits (a mantissa above
631# sqrt(2)'s, 0x3504F3, is halved and the exponent raised by one); ex * ln 2 in two parts, the first exact
632def rng_ln_w(+w: U32) -> F32:
633  +big = U32.is_lt(3474675, U32.and(w, 8388607))
634  +mb = U32.or(U32.and(w, 8388607), Bool.pick(U32, big, 1056964608, 1065353216))
635  rng_ln_m((U32.to_f32(U32.shrn(w, 23n)) - Bool.pick(F32, big, 126.0, 127.0) : F32), rng_f32(mb))
636
637def rng_ln(x: F32) -> F32:
638  rng_ln_w(rng_word(x))
639
640# sin x and cos x for x in [0, pi/4], Taylor to x^9 and x^10 (the next terms are below 2^-28)
641def rng_sin(+x: F32) -> F32:
642  +z = (x * x : F32)
643  +c4 = (0.0001984127 - (z * 0.0000027557319 : F32) : F32)
644  +c3 = (0.008333334 - (z * c4 : F32) : F32)
645  +c2 = (0.16666667 - (z * c3 : F32) : F32)
646  (x - ((x * z : F32) * c2 : F32) : F32)
647
648def rng_cos(+x: F32) -> F32:
649  +z = (x * x : F32)
650  +c5 = (0.0000248015873 - (z * 0.00000027557319 : F32) : F32)
651  +c4 = (0.0013888889 - (z * c5 : F32) : F32)
652  +c3 = (0.041666668 - (z * c4 : F32) : F32)
653  +c2 = (0.5 - (z * c3 : F32) : F32)
654  (1.0 - (z * c2 : F32) : F32)
655
656# the result of rng_sc: sin and cos of one angle
657type SC is Data:
658  SC{s: F32, c: F32}
659
660# sin and cos of a quarter turn's first eighth (x) or second (pi/2 - x, the two swapped), then turned by q quarters
661def rng_sc_q(+q: U32, +sx: F32, +cx: F32) -> SC:
662  SC{Bool.pick(F32, U32.is_eq(q, 0), sx, Bool.pick(F32, U32.is_eq(q, 1), cx, Bool.pick(F32, U32.is_eq(q, 2), F32.neg(sx), F32.neg(cx)))), Bool.pick(F32, U32.is_eq(q, 0), cx, Bool.pick(F32, U32.is_eq(q, 1), F32.neg(sx), Bool.pick(F32, U32.is_eq(q, 2), F32.neg(cx), sx)))}
663
664def rng_sc_x(+q: U32, +hi: Bool, +x: F32) -> SC:
665  +a = rng_sin(x)
666  +b = rng_cos(x)
667  rng_sc_q(q, Bool.pick(F32, hi, b, a), Bool.pick(F32, hi, a, b))
668
669# sin and cos of 2 pi v / 2^24 (v < 2^24): the quarter q is v's top 2 bits, w the rest; past the eighth (2^21) the angle
670# is taken from the quarter's end (2^22 - w, exact), so the polynomial only sees [0, pi/4]
671def rng_sc_w(+q: U32, +w: U32) -> SC:
672  +hi = U32.is_lt(2097152, w)
673  rng_sc_x(q, hi, (U32.to_f32(Bool.pick(U32, hi, (4194304 - w : U32), w)) * 0.00000037450702 : F32))
674
675def rng_sc(+v: U32) -> SC:
676  rng_sc_w(U32.shrn(v, 22n), U32.and(v, 4194303))
677
678def rng_pair(+r: F32, a: SC) -> F32 & F32:
679  match a:
680    case SC{+s, +c}: Tuple{(r * c : F32), (r * s : F32)}
681
682# the Gaussians of cells 2p and 2p + 1 of counter s (Box-Muller). The angle's 24 bits are u_b's (its hash plus one,
683# mod 2^24), so the angle is 2 pi u_b exactly before rounding
684def rng_gauss2(+s: U32, +p: U32) -> F32 & F32:
685  +i = (p + p : U32)
686  rng_pair(F32.sqrt((0.0 - (2.0 * rng_ln(rng_uniform(s, i)) : F32) : F32)), rng_sc(U32.and((rng_hash(s, (i + 1 : U32)) + 1 : U32), 16777215)))
687
688def rng_pick(+odd: Bool, g: F32 & F32) -> F32:
689  match g:
690    case Tuple{+a, +b}: Bool.pick(F32, odd, b, a)
691
692# the Gaussian of cell i alone (fill_gauss's cell i)
693def rng_gauss(+s: U32, +i: U32) -> F32:
694  rng_pick(U32.is_eq(U32.and(i, 1), 1), rng_gauss2(s, U32.shrn(i, 1n)))
695
696# the counter, the first index and the number of cells, the same for every leaf
697type FE is Data:
698  FE{s: U32, i0: U32, n: U32}
699
700def fu_go(k: Nat, +i: U32, +l: U32, +s: U32, +i0: U32, o: Array<F32>) -> Array<F32>:
701  match k:
702    case 0n: o
703    case 1n+q: fu_go(q, (i + l : U32), l, s, i0, Array.set(F32, o, i, rng_uniform(s, (i0 + i : U32))))
704
705# leaf t writes the cells t, t + l, .. below n, cell k the uniform of index i0 + k
706def fu_leaf(+t: U32, +l: U32, +e: FE, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: Array<F32>) -> Array<F32>:
707  match e:
708    case FE{+s, +i0, +n}: fu_go(share(t, n, l), t, l, s, i0, o)
709
710# writes cell i when it is below n (the cell before the first of an odd i0, which wraps to 2^32 - 1, and the cell after
711# the last are not written)
712def fg_if(ok: Bool, +i: U32, +x: F32, o: Array<F32>) -> Array<F32>:
713  match ok:
714    case True{}: Array.set(F32, o, i, x)
715    case False{}: o
716
717def fg_put(+i: U32, +n: U32, g: F32 & F32, o: Array<F32>) -> Array<F32>:
718  match g:
719    case Tuple{+a, +b}: fg_if(U32.is_lt((i + 1 : U32), n), (i + 1 : U32), b, fg_if(U32.is_lt(i, n), i, a, o))
720
721# pair j of the slice is the global pair q + j (q = i0 / 2), cells 2j - h and 2j - h + 1 of the output (h = i0 mod 2)
722def fg_go(k: Nat, +j: U32, +l: U32, +s: U32, +q: U32, +h: U32, +n: U32, o: Array<F32>) -> Array<F32>:
723  match k:
724    case 0n: o
725    case 1n+r: fg_go(r, (j + l : U32), l, s, q, h, n, fg_put(((j + j : U32) - h : U32), n, rng_gauss2(s, (q + j : U32)), o))
726
727# the number of global pairs that cells i0 .. i0 + n - 1 touch: (i0 mod 2 + n + 1) / 2
728def fg_pairs(+i0: U32, +n: U32) -> U32:
729  U32.shrn(((U32.and(i0, 1) + n : U32) + 1 : U32), 1n)
730
731# leaf t writes the pairs t, t + l, .. of the slice's pairs, each into the cells of it that lie in the slice
732def fg_leaf(+t: U32, +l: U32, +e: FE, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: Array<F32>) -> Array<F32>:
733  match e:
734    case FE{+s, +i0, +n}: fg_go(share(t, fg_pairs(i0, n), l), t, l, s, U32.shrn(i0, 1n), U32.and(i0, 1), n, o)
735
736def r1_out(r: R1) -> Array<F32>:
737  match r:
738    case R1{x1, x2, x3, x4, a}: a
739
740# n uniforms of counter s from index i0, cell k = rng_uniform(s, i0 + k) (mod 2^32), in a new array of 2^c cells made
741# on the device (one bang)
742@unsafe
743def fill_uniform_at(+s: U32, +i0: U32, +n: Nat, +c: Nat) -> Array<F32>:
744  r1_out(run1_new(~FE, ~fu_leaf, depth(n), FE{s, i0, U32.from_nat(n)}, none(Unit{}), none(Unit{}), none(Unit{}), none(Unit{}), c))
745
746# n Gaussians of counter s from index i0, cell k = rng_gauss(s, i0 + k) (mod 2^32), in a new array of 2^c cells made
747# on the device (one bang, a lane a global pair of cells; an odd i0 starts with the sine half of the pair (i0 - 1, i0))
748@unsafe
749def fill_gauss_at(+s: U32, +i0: U32, +n: Nat, +c: Nat) -> Array<F32>:
750  r1_out(run1_new(~FE, ~fg_leaf, depth(U32.to_nat(fg_pairs(i0, U32.from_nat(n)))), FE{s, i0, U32.from_nat(n)}, none(Unit{}), none(Unit{}), none(Unit{}), none(Unit{}), c))
751
752# n uniforms of counter s, cell i = rng_uniform(s, i), in a new array of 2^c cells made on the device (one bang)
753@unsafe
754def fill_uniform(+s: U32, +n: Nat, +c: Nat) -> Array<F32>:
755  fill_uniform_at(s, 0, n, c)
756
757# n Gaussians of counter s, cell i = rng_gauss(s, i), in a new array of 2^c cells made on the device (one bang, a lane a
758# pair of cells)
759@unsafe
760def fill_gauss(+s: U32, +n: Nat, +c: Nat) -> Array<F32>:
761  fill_gauss_at(s, 0, n, c)
762
763# an r x m Mat of uniforms of counter s, row-major (cell i * m + j = rng_uniform(s, i * m + j))
764@unsafe
765def uniform_mat(+s: U32, +r: Nat, +m: Nat) -> TA.Mat<r, m>:
766  TA.Mat{fill_uniform(s, Nat.mul(r, m), TA.cap_depth(Nat.mul(r, m)))}
767
768# an r x m Mat of Gaussians of counter s, row-major (cell i * m + j = rng_gauss(s, i * m + j))
769@unsafe
770def gauss_mat(+s: U32, +r: Nat, +m: Nat) -> TA.Mat<r, m>:
771  TA.Mat{fill_gauss(s, Nat.mul(r, m), TA.cap_depth(Nat.mul(r, m)))}
772
773# an r x m Mat of uniforms of counter s from index i0 (cell i * m + j = rng_uniform(s, i0 + i * m + j)); rows b .. b + r - 1
774# of an R x m uniform_mat are uniform_mat_at(s, b * m, r, m)
775@unsafe
776def uniform_mat_at(+s: U32, +i0: U32, +r: Nat, +m: Nat) -> TA.Mat<r, m>:
777  TA.Mat{fill_uniform_at(s, i0, Nat.mul(r, m), TA.cap_depth(Nat.mul(r, m)))}
778
779# an r x m Mat of Gaussians of counter s from index i0 (cell i * m + j = rng_gauss(s, i0 + i * m + j)); rows
780# b .. b + r - 1 of an R x m gauss_mat are gauss_mat_at(s, b * m, r, m)
781@unsafe
782def gauss_mat_at(+s: U32, +i0: U32, +r: Nat, +m: Nat) -> TA.Mat<r, m>:
783  TA.Mat{fill_gauss_at(s, i0, Nat.mul(r, m), TA.cap_depth(Nat.mul(r, m)))}