1 # bend-ml-tensor-gpu: tensor-array's products on the GPU, with tensor-array's bits. 2 # 3 # import bend-ml-tensor-array@0.1.13.0/main.bend as TA 4 # import bend-ml-tensor-gpu@0.1.1.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 `--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). 24 import Base 25 import bend-ml-tensor-array@0.1.13.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 28 type 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) 32 def 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 38 type GR is Type: 39 GR{a: Array, b: Array, 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 41 def 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, rv: Array & 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} 45 def 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, rv: Array & 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))) 49 def 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, rv: Array & 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))) 53 def 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, rv: Array & 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))) 57 def 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, rv: Array & 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))) 61 def 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, rv: Array & 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))) 65 def 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, rv: Array & 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))) 69 def 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, rv: Array & 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))) 73 def 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, rv: Array & 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))) 77 def 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, rv: Array & 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))) 81 def 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, rv: Array & 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))) 85 def 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, rv: Array & 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))) 89 def 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, rv: Array & 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))) 93 def 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, rv: Array & 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))) 97 def 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, rv: Array & 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))) 101 def 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, rv: Array & 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))) 105 def 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, b: Array) -> 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))) 109 def 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, b: Array) -> 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) 113 def 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, b: Array) -> 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) 117 def 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, b: Array) -> 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) 121 def 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, b: Array) -> 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) 125 def 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, b: Array) -> 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) 129 def 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, b: Array) -> 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) 133 def 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, b: Array) -> 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) 137 def 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, b: Array) -> 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) 141 def 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, b: Array) -> 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) 145 def 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, b: Array) -> 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) 149 def 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, b: Array) -> 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) 153 def 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, b: Array) -> 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) 157 def gr_2(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, +x0_1: F32, rv: Array & F32, b: Array) -> 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) 161 def gr_1(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +x0_0: F32, rv: Array & F32, b: Array) -> 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) 165 def gr_0(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, rv: Array & F32, b: Array) -> 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 170 def turn(+ia: U32, +ib: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array) -> 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 174 type T is Type: 175 T{a: Array, b: Array, 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 180 def 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_l 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), 186 187 # writes cell (i, j) of C when it is inside n x m 188 def wr(ok: Bool, +ac: Bool, +idx: U32, +x: F32, c: Array) -> Array: 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 195 def put(+ac: Bool, +n: U32, +m: U32, +i0: U32, +j0: U32, c: Array, t: T) -> Array: 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 : U 199 200 @unsafe 201 def tile(+e: E, +i0: U32, +j0: U32, +a: Array, +b: Array, c: Array) -> Array: 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 208 def leaf_if(ok: Bool, +e: E, +t: U32, +tcols: U32, +a: Array, +b: Array, c: Array) -> Array: 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 214 def leaf(+e: E, +t: U32, +a: Array, +b: Array, c: Array) -> Array: 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 220 def tree(+e: E, d: Nat, +t: U32, +a: Array, +b: Array, +c: Array) -> Array: 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 231 def tree_new(+e: E, d: Nat, +a: Array, +b: Array, +cd: Nat) -> Array: 232 tree(e, d, 0, a, b, Array.new(F32, cd, 0.0)) 233 234 @unsafe 235 def bang_new(+e: E, d: Nat, +a: Array, +b: Array, +cd: Nat) -> Array: 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 240 def bang(+e: E, d: Nat, +a: Array, +b: Array, c: Array) -> Array: 241 tree!(e, d, 0, a, b, c) 242 243 type ABC is Type: 244 ABC{a: Array, b: Array, c: Array} 245 246 @unsafe 247 def run_e(+e: E, +tiles: U32, +a: Array, +b: Array, c: Array) -> 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 251 def 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 254 def 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 259 def gemm_new(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +a: Array, +b: Array) -> 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 264 def gemm(+ac: Bool, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array, b: Array, c: Array) -> ABC: 265 run_e(e_of(ac, n, m, k, arow, sa, sb, bcol), U32.from_nat(tiles(n, m)), a, b, c) 266 267 def mm_fin(-n: Nat, -k: Nat, -m: Nat, g: ABC) -> TA.MMul: 268 match g: 269 case ABC{a, b, c}: TA.MMul{TA.Mat{a}, TA.Mat{b}, TA.Mat{c}} 270 271 def mm_nt_fin(-n: Nat, -k: Nat, -m: Nat, g: ABC) -> TA.MMulNT: 272 match g: 273 case ABC{a, b, c}: TA.MMulNT{TA.Mat{a}, TA.Mat{b}, TA.Mat{c}} 274 275 def mm_tn_fin(-n: Nat, -k: Nat, -m: Nat, g: ABC) -> TA.MMulTN: 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 281 def matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat) -> TA.MMul: 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 288 def matmul_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat, c: TA.Mat) -> TA.MMul: 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 295 def matmul_nt(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat) -> TA.MMulNT: 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 302 def matmul_nt_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat, c: TA.Mat) -> TA.MMulNT: 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 309 def matmul_tn(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat) -> TA.MMulTN: 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 316 def matmul_tn_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat, b: TA.Mat, c: TA.Mat) -> TA.MMulTN: 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 342 def 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 348 def 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 352 def 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 356 def none(u: Unit) -> Array: 357 Array.new(F32, 0n, 0.0) 358 359 type O2 is Type: 360 O2{a: Array, b: Array} 361 362 type O3 is Type: 363 O3{a: Array, b: Array, c: Array} 364 365 type O4 is Type: 366 O4{a: Array, b: Array, c: Array, d: Array} 367 368 type R1 is Type: 369 R1{x1: Array, x2: Array, x3: Array, x4: Array, a: Array} 370 371 type R2 is Type: 372 R2{x1: Array, x2: Array, x3: Array, x4: Array, o: O2} 373 374 type R3 is Type: 375 R3{x1: Array, x2: Array, x3: Array, x4: Array, o: O3} 376 377 type R4 is Type: 378 R4{x1: Array, x2: Array, x3: Array, x4: Array, o: O4} 379 380 @unsafe 381 def 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 386 def 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 391 def 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 396 def tree1(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> Array -> Array, +d: Nat, +t: U32, +l: U32, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, +oa: Array) -> Array: 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 404 def root1(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> Array -> Array, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, oa: Array) -> Array: 405 tree1(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa) 406 407 @unsafe 408 def root1_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> Array -> Array, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, +ca: Nat) -> Array: 409 root1(~E, ~f, d, e, x1, x2, x3, x4, Array.new(F32, ca, 0.0)) 410 411 @unsafe 412 def go1(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> Array -> Array, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, oa: Array) -> R1: 413 R1{x1, x2, x3, x4, root1!(~E, ~f, d, e, x1, x2, x3, x4, oa)} 414 415 @unsafe 416 def go1_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> Array -> Array, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, +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 421 def run1(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> Array -> Array, +d: Nat, +e: E, x1: Array, x2: Array, x3: Array, x4: Array, oa: Array) -> 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 426 def run1_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> Array -> Array, +d: Nat, +e: E, x1: Array, x2: Array, x3: Array, x4: Array, +ca: Nat) -> R1: 427 go1_new(~E, ~f, d, e, x1, x2, x3, x4, ca) 428 429 @unsafe 430 def tree2(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O2 -> O2, +d: Nat, +t: U32, +l: U32, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, +oa: Array, +ob: Array) -> 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 438 def root2(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O2 -> O2, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, oa: Array, ob: Array) -> O2: 439 tree2(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa, ob) 440 441 @unsafe 442 def root2_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O2 -> O2, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, +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 446 def go2(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O2 -> O2, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, oa: Array, ob: Array) -> R2: 447 R2{x1, x2, x3, x4, root2!(~E, ~f, d, e, x1, x2, x3, x4, oa, ob)} 448 449 @unsafe 450 def go2_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O2 -> O2, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, +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 455 def run2(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O2 -> O2, +d: Nat, +e: E, x1: Array, x2: Array, x3: Array, x4: Array, oa: Array, ob: Array) -> 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 460 def run2_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O2 -> O2, +d: Nat, +e: E, x1: Array, x2: Array, x3: Array, x4: Array, +ca: Nat, +cb: Nat) -> R2: 461 go2_new(~E, ~f, d, e, x1, x2, x3, x4, ca, cb) 462 463 @unsafe 464 def tree3(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O3 -> O3, +d: Nat, +t: U32, +l: U32, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, +oa: Array, +ob: Array, +oc: Array) -> 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 472 def root3(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O3 -> O3, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, oa: Array, ob: Array, oc: Array) -> O3: 473 tree3(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa, ob, oc) 474 475 @unsafe 476 def root3_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O3 -> O3, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, +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 480 def go3(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O3 -> O3, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, oa: Array, ob: Array, oc: Array) -> R3: 481 R3{x1, x2, x3, x4, root3!(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc)} 482 483 @unsafe 484 def go3_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O3 -> O3, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, +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 489 def run3(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O3 -> O3, +d: Nat, +e: E, x1: Array, x2: Array, x3: Array, x4: Array, oa: Array, ob: Array, oc: Array) -> 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 494 def run3_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O3 -> O3, +d: Nat, +e: E, x1: Array, x2: Array, x3: Array, x4: Array, +ca: Nat, +cb: Nat, +cc: Nat) -> R3: 495 go3_new(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc) 496 497 @unsafe 498 def tree4(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O4 -> O4, +d: Nat, +t: U32, +l: U32, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, +oa: Array, +ob: Array, +oc: Array, +od: Array) -> 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 506 def root4(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O4 -> O4, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, oa: Array, ob: Array, oc: Array, od: Array) -> O4: 507 tree4(~E, ~f, d, 0, U32.pow(2, d), e, x1, x2, x3, x4, oa, ob, oc, od) 508 509 @unsafe 510 def root4_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O4 -> O4, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, +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 514 def go4(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O4 -> O4, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, oa: Array, ob: Array, oc: Array, od: Array) -> R4: 515 R4{x1, x2, x3, x4, root4!(~E, ~f, d, e, x1, x2, x3, x4, oa, ob, oc, od)} 516 517 @unsafe 518 def go4_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O4 -> O4, +d: Nat, +e: E, +x1: Array, +x2: Array, +x3: Array, +x4: Array, +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 523 def run4(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O4 -> O4, +d: Nat, +e: E, x1: Array, x2: Array, x3: Array, x4: Array, oa: Array, ob: Array, oc: Array, od: Array) -> 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 528 def run4_new(~E: Data, ~f: @+t: U32 -> @+l: U32 -> @+e: E -> Array -> Array -> Array -> Array -> O4 -> O4, +d: Nat, +e: E, x1: Array, x2: Array, x3: Array, x4: Array, +ca: Nat, +cb: Nat, +cc: Nat, +cd: Nat) -> R4: 529 go4_new(~E, ~f, d, e, x1, x2, x3, x4, ca, cb, cc, cd)