main.bend

bendtorch-gpu@1.3.0.0 · 1,627 lines · 324 definitions

1# BendTorch's layers with their work on the GPU (bend-ml-tensor-gpu): the same layers, types and bits as
2# torch/main.bend's and torch/fast.bend's.
3#
4#   import ../torch/main.bend as T
5#   import ../torch/gpu.bend as F
6#   F.LinearFirstF(100n, 784n, 128n, 4n)     # this layer's products and optimizer step run as bangs
7#
8# Every layer here runs its work as bangs (one `!` call each), which run on the GPU in a binary built with CUDA and on
9# the CPU threads otherwise:
10# - LinearF, LinearFirstF and LinearNBF: their three products (x·Wᵀ, dY·W, dYᵀ·X added into dW) from tensor-gpu, and
11#   their optimizer step.
12# - SiLU, Mul2, Add2, Residual, Mat.dup and Mat.sum2: their element-wise loops.
13# - RMSNorm, RoPE and Attention (main.bend's names; fast.bend's are main.bend's, on the CPU).
14# - MSE, CrossEntropy and Embedding (its rows, its gradient and its step), so that a training step's work, the loss
15#   included, stays on the device; the host reads back a few numbers a row for the loss.
16# The names are fast.bend's, so a program switches by its import line alone (the Makefile builds the GPU demos so).
17#
18# Each bang is one of tensor-gpu's run1..run4 (outputs in place) or run1_new..run4_new (outputs made on the device):
19# a tree of leaves, every node sharing the arrays. The code of a leaf is below, in plain Bend, and computes each number
20# with main.bend's own expression or row code (rmsf_row, ssq, rmsb_dot, rope_one, dot0, psmc, g_set2, adam_m..), so the bits
21# are main.bend's, except where an exponential is taken on the device (SiLU, attention's softmax, cross-entropy): CUDA
22# takes the float overload of exp there and its last bit can differ from the host's (the known
23# limit of bendlang/bend#1305); 1 ulp is accepted there, my decision of 2026-10-18 (tests: demo/gpu_test.bend
24# and demo/gpu_layers_test.bend). Sizes, depths and capacities are
25# computed on the host and passed in. The arrays stay on the device from one bang to the next as long as the host only
26# passes their handles; a bang that follows the host's parallel lets in one evaluation runs on the CPU threads.
27# A sum whose order fixes the bits (a row's softmax, a row's squares) is a row bang: a leaf a row, a tail loop along
28# it. Everything else is a cell bang: a leaf takes cells t, t + l, .., each one number (a dot product, or one
29# expression of a row's sums), so a bang has as many independent pieces of work as the output has numbers.
30#
31# Outside the proofs, on purpose, like fast.bend: tensor-gpu shares arrays between the leaves of a bang through Base's
32# @unsafe Array.fork, so `bend` names the defs below that reach it, and no file with laws imports this one (`make laws`
33# checks both). This file writes no @unsafe of its own. All of it is Bend.
34
35import Base
36import bend-ml-tensor-array@0.1.18.0/main.bend as TA
37import bend-ml-tensor-gpu@0.1.5.0/main.bend as TG
38import bendtorch@1.3.0.0/main.bend as T
39
40# ---- Linear: y = x·Wᵀ + b ----
41
42def lin_fwd(+n: Nat, +i: Nat, +o: Nat, p: T.Lin<i, o>, x: TA.Mat<n, i>) -> T.Fwd<T.Lin<i, o>, TA.Mat<n, o>, TA.Mat<n, i>>:
43  match p:
44    case T.Lin{w, b}:
45      T.lin_f2(n, i, o, b, TG.matmul_nt(n, i, o, 0n, x, w))
46
47def lin_b2(+n: Nat, +i: Nat, +o: Nat, b: TA.Mat<1n, o>, x: TA.Mat<n, i>, g: T.Lin<i, o>, r: TA.MMul<n, o, i>) -> T.Bwd<T.Lin<i, o>, TA.Mat<n, i>>:
48  match g r:
49    case T.Lin{gw, gb} TA.MMul{dy, w, dx}:
50      T.lin_b3(n, i, o, w, b, dx, gb, TG.matmul_tn_acc(o, n, i, 0n, dy, x, gw))
51
52def lin_bwd(+n: Nat, +i: Nat, +o: Nat, p: T.Lin<i, o>, x: TA.Mat<n, i>, dy: TA.Mat<n, o>, g: T.Lin<i, o>) -> T.Bwd<T.Lin<i, o>, TA.Mat<n, i>>:
53  match p:
54    case T.Lin{w, b}:
55      lin_b2(n, i, o, b, x, g, TG.matmul_precise(n, o, i, 0n, dy, w))
56
57# =====================================================================================================
58# The leaves' shared pieces. A leaf gets its index t among l leaves, a record of scalars and four input arrays, and
59# returns its outputs; an input it does not use is tensor-gpu's none(Unit{}). Element-wise leaves take the cells t,
60# t + l, .. (TG.share), so neighbouring lanes read neighbouring cells; row and column leaves take one row or column
61# each and leave their outputs unchanged past the last one.
62# =====================================================================================================
63
64type EN is Data:
65  EN{n: U32}
66
67def nn(u: Unit) -> Array<F32>:
68  TG.none(u)
69
70# out[i] = f(x[i]) for the leaf's cells
71def ew1_loop(~f: @+a: F32 -> F32, k: Nat, +i: U32, +l: U32, o: Array<F32>, r: Array<F32> & F32) -> Array<F32>:
72  match k r:
73    case 0n Tuple{a, +x}: o
74    case 1n+q Tuple{a, +x}: ew1_loop(~f, q, (i + l : U32), l, Array.set(F32, o, i, f(x)), Array.get(F32, a, (i + l : U32)))
75
76def ew1(~f: @+a: F32 -> F32, +t: U32, +l: U32, +e: EN, x: Array<F32>, o: Array<F32>) -> Array<F32>:
77  match e:
78    case EN{+n}: ew1_loop(~f, TG.share(t, n, l), t, l, o, Array.get(F32, x, t))
79
80# out[i] = f(x[i], y[i]) for the leaf's cells
81def ew2_loop(~f: @+a: F32 -> @+b: F32 -> F32, k: Nat, +i: U32, +l: U32, o: Array<F32>, ra: Array<F32> & F32, rb: Array<F32> & F32) -> Array<F32>:
82  match k ra rb:
83    case 0n Tuple{a, +x} Tuple{b, +y}: o
84    case 1n+q Tuple{a, +x} Tuple{b, +y}: ew2_loop(~f, q, (i + l : U32), l, Array.set(F32, o, i, f(x, y)), Array.get(F32, a, (i + l : U32)), Array.get(F32, b, (i + l : U32)))
85
86def ew2(~f: @+a: F32 -> @+b: F32 -> F32, +t: U32, +l: U32, +e: EN, x: Array<F32>, y: Array<F32>, o: Array<F32>) -> Array<F32>:
87  match e:
88    case EN{+n}: ew2_loop(~f, TG.share(t, n, l), t, l, o, Array.get(F32, x, t), Array.get(F32, y, t))
89
90# o[i] = f(o[i], y[i]) for the leaf's cells, in place (o's old value is the first argument)
91def ewi_loop(~f: @+a: F32 -> @+b: F32 -> F32, k: Nat, +i: U32, +l: U32, ro: Array<F32> & F32, ry: Array<F32> & F32) -> Array<F32>:
92  match k ro ry:
93    case 0n Tuple{o, +x} Tuple{y, +v}: o
94    case 1n+q Tuple{o, +x} Tuple{y, +v}: ewi_loop(~f, q, (i + l : U32), l, Array.get(F32, Array.set(F32, o, i, f(x, v)), (i + l : U32)), Array.get(F32, y, (i + l : U32)))
95
96def ewi(~f: @+a: F32 -> @+b: F32 -> F32, +t: U32, +l: U32, +e: EN, o: Array<F32>, y: Array<F32>) -> Array<F32>:
97  match e:
98    case EN{+n}: ewi_loop(~f, TG.share(t, n, l), t, l, Array.get(F32, o, t), Array.get(F32, y, t))
99
100def f32_id(+x: F32) -> F32:
101  x
102
103def lf_copy(+t: U32, +l: U32, +e: EN, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: Array<F32>) -> Array<F32>:
104  ew1(~f32_id, t, l, e, x1, o)
105
106def lf_silu(+t: U32, +l: U32, +e: EN, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: Array<F32>) -> Array<F32>:
107  ew1(~T.silu_v, t, l, e, x1, o)
108
109# o = dy, written in place, x1 = x: main.bend's silu_d(dy, x)
110def lf_silud(+t: U32, +l: U32, +e: EN, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: Array<F32>) -> Array<F32>:
111  ewi(~T.silu_d, t, l, e, o, x1)
112
113# o = a, written in place, x1 = b: a + b
114def lf_add(+t: U32, +l: U32, +e: EN, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: Array<F32>) -> Array<F32>:
115  ewi(~T.f32_add, t, l, e, o, x1)
116
117def lf_mul(+t: U32, +l: U32, +e: EN, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: Array<F32>) -> Array<F32>:
118  ew2(~T.f32_mul, t, l, e, x1, x2, o)
119
120# x1 = dy; the saved a and b, written in place: a becomes db = dy * a and b becomes da = dy * b, in one bang
121def mulb_loop(k: Nat, +i: U32, +l: U32, rd: Array<F32> & F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> TG.O2:
122  match k rd ra rb:
123    case 0n Tuple{d, +x} Tuple{a, +y} Tuple{b, +z}: TG.O2{a, b}
124    case 1n+q Tuple{d, +x} Tuple{a, +y} Tuple{b, +z}:
125      +u = (i + l : U32)
126      mulb_loop(q, u, l, Array.get(F32, d, u), Array.get(F32, Array.set(F32, a, i, (x * y : F32)), u), Array.get(F32, Array.set(F32, b, i, (x * z : F32)), u))
127
128def lf_mulb(+t: U32, +l: U32, +e: EN, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
129  match e o:
130    case EN{+n} TG.O2{a, b}: mulb_loop(TG.share(t, n, l), t, l, Array.get(F32, x1, t), Array.get(F32, a, t), Array.get(F32, b, t))
131
132def en(+cnt: Nat) -> EN:
133  EN{U32.from_nat(cnt)}
134
135def r1_a(r: TG.R1) -> Array<F32>:
136  match r:
137    case TG.R1{x1, x2, x3, x4, a}: a
138
139# ---- element-wise layers ----
140
141def dup_fin(-r: Nat, -c: Nat, q: TG.R1) -> T.Two<TA.Mat<r, c>, TA.Mat<r, c>>:
142  match q:
143    case TG.R1{x1, x2, x3, x4, a}: T.Two{TA.Mat{x1}, TA.Mat{a}}
144
145# two copies of a matrix, the second made on the device
146def Mat.dup(+r: Nat, +c: Nat, m: TA.Mat<r, c>) -> T.Two<TA.Mat<r, c>, TA.Mat<r, c>>:
147  match m:
148    case TA.Mat{a}: dup_fin(r, c, TG.run1_new(~EN, ~lf_copy, TG.depth(Nat.mul(r, c)), en(Nat.mul(r, c)), a, nn(Unit{}), nn(Unit{}), nn(Unit{}), TA.cap_depth(Nat.mul(r, c))))
149
150# a + b, as main.bend's Mat.sum2 adds each cell, written over a (both are consumed)
151def Mat.sum2(+r: Nat, +c: Nat, a: TA.Mat<r, c>, b: TA.Mat<r, c>) -> TA.Mat<r, c>:
152  match a b:
153    case TA.Mat{x} TA.Mat{y}: TA.Mat{r1_a(TG.run1(~EN, ~lf_add, TG.depth(Nat.mul(r, c)), en(Nat.mul(r, c)), y, nn(Unit{}), nn(Unit{}), nn(Unit{}), x))}
154
155def silu_f2(-n: Nat, -c: Nat, q: TG.R1) -> T.Fwd<Unit, TA.Mat<n, c>, TA.Mat<n, c>>:
156  match q:
157    case TG.R1{x1, x2, x3, x4, a}: T.Fwd{Unit{}, TA.Mat{a}, TA.Mat{x1}}
158
159def silu_f(+n: Nat, +c: Nat, x: TA.Mat<n, c>) -> T.Fwd<Unit, TA.Mat<n, c>, TA.Mat<n, c>>:
160  match x:
161    case TA.Mat{a}: silu_f2(n, c, TG.run1_new(~EN, ~lf_silu, TG.depth(Nat.mul(n, c)), en(Nat.mul(n, c)), a, nn(Unit{}), nn(Unit{}), nn(Unit{}), TA.cap_depth(Nat.mul(n, c))))
162
163def silu_b(+n: Nat, +c: Nat, s: TA.Mat<n, c>, dy: TA.Mat<n, c>) -> T.Bwd<Unit, TA.Mat<n, c>>:
164  match s dy:
165    case TA.Mat{x} TA.Mat{d}: T.Bwd{Unit{}, TA.Mat{r1_a(TG.run1(~EN, ~lf_silud, TG.depth(Nat.mul(n, c)), en(Nat.mul(n, c)), x, nn(Unit{}), nn(Unit{}), nn(Unit{}), d))}, Unit{}}
166
167# SiLU: y = silu(x), the input saved; dx = silu'(x) * dy
168def SiLU(+n: Nat, +c: Nat) -> T.Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
169  T.Layer{p => x => silu_f(n, c, x), p => s => dy => g => silu_b(n, c, s, dy), op => p => g => st => T.no_step(Unit, op, p, g, st), u => Unit{}, p => T.PL{p, Nil{}}, p => xs => off => T.PA{p, xs, off}, p => xs => off => T.PA{p, xs, off}, u => 0}
170
171def mul2_f2(-n: Nat, -c: Nat, q: TG.R1) -> T.Fwd<Unit, TA.Mat<n, c>, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
172  match q:
173    case TG.R1{x1, x2, x3, x4, y}: T.Fwd{Unit{}, TA.Mat{y}, T.Two{TA.Mat{x1}, TA.Mat{x2}}}
174
175def mul2_f(+n: Nat, +c: Nat, x: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> T.Fwd<Unit, TA.Mat<n, c>, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
176  match x:
177    case T.Two{TA.Mat{a}, TA.Mat{b}}: mul2_f2(n, c, TG.run1_new(~EN, ~lf_mul, TG.depth(Nat.mul(n, c)), en(Nat.mul(n, c)), a, b, nn(Unit{}), nn(Unit{}), TA.cap_depth(Nat.mul(n, c))))
178
179def mul2_b2(-n: Nat, -c: Nat, q: TG.R2) -> T.Bwd<Unit, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
180  match q:
181    case TG.R2{x1, x2, x3, x4, TG.O2{db, da}}: T.Bwd{Unit{}, T.Two{TA.Mat{da}, TA.Mat{db}}, Unit{}}
182
183def mul2_b(+n: Nat, +c: Nat, s: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>, dy: TA.Mat<n, c>) -> T.Bwd<Unit, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
184  match s dy:
185    case T.Two{TA.Mat{a}, TA.Mat{b}} TA.Mat{d}: mul2_b2(n, c, TG.run2(~EN, ~lf_mulb, TG.depth(Nat.mul(n, c)), en(Nat.mul(n, c)), d, nn(Unit{}), nn(Unit{}), nn(Unit{}), a, b))
186
187# Mul2: y = a * b, the inputs saved; da = dy * b, db = dy * a
188def Mul2(+n: Nat, +c: Nat) -> T.Layer<Unit, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>, TA.Mat<n, c>, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
189  T.Layer{p => x => mul2_f(n, c, x), p => s => dy => g => mul2_b(n, c, s, dy), op => p => g => st => T.no_step(Unit, op, p, g, st), u => Unit{}, p => T.PL{p, Nil{}}, p => xs => off => T.PA{p, xs, off}, p => xs => off => T.PA{p, xs, off}, u => 0}
190
191def add2_f(+n: Nat, +c: Nat, x: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> T.Fwd<Unit, TA.Mat<n, c>, Unit>:
192  match x:
193    case T.Two{a, b}: T.Fwd{Unit{}, Mat.sum2(n, c, a, b), Unit{}}
194
195def Add2(+n: Nat, +c: Nat) -> T.Layer<Unit, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>, TA.Mat<n, c>, Unit>:
196  T.Layer{p => x => add2_f(n, c, x), p => s => dy => g => T.Bwd{p, Mat.dup(n, c, dy), g}, op => p => g => st => T.no_step(Unit, op, p, g, st), u => Unit{}, p => T.PL{p, Nil{}}, p => xs => off => T.PA{p, xs, off}, p => xs => off => T.PA{p, xs, off}, u => 0}
197
198# T.Residual with these copies and sums
199def Residual(~n: Nat, ~c: Nat, ~P: Type, ~S: Type, ~f: T.Layer<P, TA.Mat<n, c>, TA.Mat<n, c>, S>) -> T.Layer<T.Two<T.Two<Unit, P>, Unit>, TA.Mat<n, c>, TA.Mat<n, c>, T.Two<T.Two<Unit, S>, Unit>>:
200  T.Seq(~T.Two<Unit, P>, ~Unit, ~TA.Mat<n, c>, ~T.Two<TA.Mat<n, c>, TA.Mat<n, c>>, ~TA.Mat<n, c>, ~T.Two<Unit, S>, ~Unit, ~T.Fork(~Unit, ~P, ~TA.Mat<n, c>, ~TA.Mat<n, c>, ~TA.Mat<n, c>, ~Unit, ~S, ~(m => Mat.dup(n, c, m)), ~(a => b => Mat.sum2(n, c, a, b)), ~T.Id(~TA.Mat<n, c>), ~f), ~Add2(n, c))
201
202# ---- the optimizer step in place, each cell as main.bend's sgd_l, mom_l and adam_l compute it ----
203
204type ESGD is Data:
205  ESGD{n: U32, lr: F32}
206
207def sgd_loop(k: Nat, +i: U32, +l: U32, +lr: F32, rw: Array<F32> & F32, rg: Array<F32> & F32) -> TG.O2:
208  match k rw rg:
209    case 0n Tuple{w, +x} Tuple{g, +y}: TG.O2{w, g}
210    case 1n+q Tuple{w, +x} Tuple{g, +y}:
211      sgd_loop(q, (i + l : U32), l, lr, Array.get(F32, Array.set(F32, w, i, (x + (F32.neg(lr) * y : F32) : F32)), (i + l : U32)), Array.get(F32, Array.set(F32, g, i, 0.0), (i + l : U32)))
212
213def lf_sgd(+t: U32, +l: U32, +e: ESGD, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
214  match e o:
215    case ESGD{+n, +lr} TG.O2{w, g}: sgd_loop(TG.share(t, n, l), t, l, lr, Array.get(F32, w, t), Array.get(F32, g, t))
216
217type EMOM is Data:
218  EMOM{n: U32, lr: F32, mu: F32}
219
220def mom_loop(k: Nat, +i: U32, +l: U32, +lr: F32, +mu: F32, rw: Array<F32> & F32, rg: Array<F32> & F32, rm: Array<F32> & F32) -> TG.O3:
221  match k rw rg rm:
222    case 0n Tuple{w, +x} Tuple{g, +y} Tuple{m, +b}: TG.O3{w, g, m}
223    case 1n+q Tuple{w, +x} Tuple{g, +y} Tuple{m, +b}:
224      +nb = ((b * mu : F32) + y : F32)
225      +u = (i + l : U32)
226      mom_loop(q, u, l, lr, mu, Array.get(F32, Array.set(F32, w, i, (x + (F32.neg(lr) * nb : F32) : F32)), u), Array.get(F32, Array.set(F32, g, i, 0.0), u), Array.get(F32, Array.set(F32, m, i, nb), u))
227
228def lf_mom(+t: U32, +l: U32, +e: EMOM, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O3) -> TG.O3:
229  match e o:
230    case EMOM{+n, +lr, +mu} TG.O3{w, g, m}: mom_loop(TG.share(t, n, l), t, l, lr, mu, Array.get(F32, w, t), Array.get(F32, g, t), Array.get(F32, m, t))
231
232type EADAM is Data:
233  EADAM{n: U32, lr: F32, c1: F32, b2: F32, c2: F32, eps: F32, wd: F32, dec: Bool, step: F32, bc2: F32}
234
235def adam_loop(k: Nat, +i: U32, +l: U32, +lr: F32, +c1: F32, +b2: F32, +c2: F32, +eps: F32, +wd: F32, +dec: Bool, +step: F32, +bc2: F32, rw: Array<F32> & F32, rg: Array<F32> & F32, rm: Array<F32> & F32, rv: Array<F32> & F32) -> TG.O4:
236  match k rw rg rm rv:
237    case 0n Tuple{w, +x} Tuple{g, +y} Tuple{m, +a} Tuple{v, +b}: TG.O4{w, g, m, v}
238    case 1n+q Tuple{w, +x} Tuple{g, +y} Tuple{m, +a} Tuple{v, +b}:
239      +gg = T.adam_grad(dec, wd, x, y)
240      +nm = T.adam_m(c1, a, gg)
241      +nv = T.adam_v(b2, c2, b, gg)
242      +nw = T.adam_w(step, bc2, eps, T.adam_decay(dec, lr, wd, x), nm, nv)
243      +u = (i + l : U32)
244      adam_loop(q, u, l, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, Array.set(F32, w, i, nw), u), Array.get(F32, Array.set(F32, g, i, 0.0), u), Array.get(F32, Array.set(F32, m, i, nm), u), Array.get(F32, Array.set(F32, v, i, nv), u))
245
246def lf_adam(+t: U32, +l: U32, +e: EADAM, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O4) -> TG.O4:
247  match e o:
248    case EADAM{+n, +lr, +c1, +b2, +c2, +eps, +wd, +dec, +step, +bc2} TG.O4{w, g, m, v}: adam_loop(TG.share(t, n, l), t, l, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, w, t), Array.get(F32, g, t), Array.get(F32, m, t), Array.get(F32, v, t))
249
250def sgd_fin(-r: Nat, -c: Nat, q: TG.R2) -> T.MatStep<r, c>:
251  match q:
252    case TG.R2{x1, x2, x3, x4, TG.O2{w, g}}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.NoMom{}}
253
254def mom_fin(-r: Nat, -c: Nat, q: TG.R3) -> T.MatStep<r, c>:
255  match q:
256    case TG.R3{x1, x2, x3, x4, TG.O3{w, g, m}}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.OneMom{TA.Mat{m}}}
257
258def opt_fin(-r: Nat, -c: Nat, q: TG.R4) -> T.MatStep<r, c>:
259  match q:
260    case TG.R4{x1, x2, x3, x4, TG.O4{w, g, m, v}}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.TwoMom{TA.Mat{m}, TA.Mat{v}}}
261
262# one matrix's step, in place: SGD writes two arrays, momentum three, Adam four
263def opt_mat(+r: Nat, +c: Nat, +o: T.Opt, w: TA.Mat<r, c>, g: TA.Mat<r, c>, st: T.OptState<TA.Mat<r, c>>) -> T.MatStep<r, c>:
264  match o w g st:
265    case T.SGD{+lr} TA.Mat{wa} TA.Mat{ga} _:
266      sgd_fin(r, c, TG.run2(~ESGD, ~lf_sgd, TG.depth(Nat.mul(r, c)), ESGD{U32.from_nat(Nat.mul(r, c)), lr}, nn(Unit{}), nn(Unit{}), nn(Unit{}), nn(Unit{}), wa, ga))
267    case T.Momentum{+lr, +mu} TA.Mat{wa} TA.Mat{ga} T.OneMom{TA.Mat{ma}}:
268      mom_fin(r, c, TG.run3(~EMOM, ~lf_mom, TG.depth(Nat.mul(r, c)), EMOM{U32.from_nat(Nat.mul(r, c)), lr, mu}, nn(Unit{}), nn(Unit{}), nn(Unit{}), nn(Unit{}), wa, ga, ma))
269    case T.Adam{+lr, +b1, +b2, +eps, +wd, +dec, +step, +bc2, +c1, +c2} TA.Mat{wa} TA.Mat{ga} T.TwoMom{TA.Mat{ma}, TA.Mat{va}}:
270      opt_fin(r, c, TG.run4(~EADAM, ~lf_adam, TG.depth(Nat.mul(r, c)), EADAM{U32.from_nat(Nat.mul(r, c)), lr, c1, b2, c2, eps, wd, dec, step, bc2}, nn(Unit{}), nn(Unit{}), nn(Unit{}), nn(Unit{}), wa, ga, ma, va))
271    # packed BF16 moments (CAdamWBF16): main.bend's step, in one block
272    case o w g T.Packed{t}: T.Opt.mat(r, c, o, w, g, T.Packed{t})
273    case _ w g st: T.MatStep{w, g, st}
274
275def mat_step(+r: Nat, +c: Nat, +op: T.Opt, p: TA.Mat<r, c>, g: TA.Mat<r, c>, st: T.OptState<TA.Mat<r, c>>) -> T.Step<TA.Mat<r, c>>:
276  T.mat_step_fin(r, c, opt_mat(r, c, op, p, g, st))
277
278def lin_step2(+i: Nat, +o: Nat, +op: T.Opt, p: T.Lin<i, o>, g: T.Lin<i, o>, st: T.Two<T.OptState<TA.Mat<o, i>>, T.OptState<TA.Mat<1n, o>>>) -> T.Step<T.Lin<i, o>>:
279  match p g st:
280    case T.Lin{w, b} T.Lin{gw, gb} T.Two{sw, sb}:
281      T.lin_s2(i, o, opt_mat(o, i, op, w, gw, sw), opt_mat(1n, o, op, b, gb, sb))
282
283def lin_step(+i: Nat, +o: Nat, +op: T.Opt, p: T.Lin<i, o>, g: T.Lin<i, o>, st: T.OptState<T.Lin<i, o>>) -> T.Step<T.Lin<i, o>>:
284  lin_step2(i, o, op, p, g, T.OptState.split(TA.Mat<o, i>, TA.Mat<1n, o>, T.OptState.map(T.Lin<i, o>, T.Two<TA.Mat<o, i>, TA.Mat<1n, o>>, st, x => T.lin_pair(i, o, x), x => T.lin_pair(i, o, x))))
285
286# ---- the linear layers ----
287
288def LinearF(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer<T.Lin<i, o>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
289  T.Layer{p => x => lin_fwd(n, i, o, p, x), p => s => dy => g => lin_bwd(n, i, o, p, s, dy, g), op => p => g => st => lin_step(i, o, op, p, g, st), u => T.lin_zeros(i, o), p => T.lin_dump(i, o, p), p => xs => off => T.lin_load(i, o, p, xs, off), p => xs => off => T.lin_store(i, o, p, xs, off), u => U32.from_nat(Nat.mul(o, 1n+i))}
290
291# LinearFirst: no dX (the first layer's input needs no gradient)
292def linf_bwd(+n: Nat, +i: Nat, +o: Nat, p: T.Lin<i, o>, x: TA.Mat<n, i>, dy: TA.Mat<n, o>, g: T.Lin<i, o>) -> T.Bwd<T.Lin<i, o>, TA.Mat<n, i>>:
293  match p g:
294    case T.Lin{w, b} T.Lin{gw, gb}:
295      T.linf_b3(n, i, o, w, b, gb, TG.matmul_tn_acc(o, n, i, 0n, dy, x, gw))
296
297def LinearFirstF(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer<T.Lin<i, o>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
298  T.Layer{p => x => lin_fwd(n, i, o, p, x), p => s => dy => g => linf_bwd(n, i, o, p, s, dy, g), op => p => g => st => lin_step(i, o, op, p, g, st), u => T.lin_zeros(i, o), p => T.lin_dump(i, o, p), p => xs => off => T.lin_load(i, o, p, xs, off), p => xs => off => T.lin_store(i, o, p, xs, off), u => U32.from_nat(Nat.mul(o, 1n+i))}
299
300# LinearNB: y = x·Wᵀ, no bias (Llama's projections)
301def lnb_b2(+n: Nat, +i: Nat, +o: Nat, x: TA.Mat<n, i>, g: TA.Mat<o, i>, r: TA.MMul<n, o, i>) -> T.Bwd<TA.Mat<o, i>, TA.Mat<n, i>>:
302  match r:
303    case TA.MMul{dy, w, dx}:
304      T.lnb_b3(n, i, o, w, dx, TG.matmul_tn_acc(o, n, i, 0n, dy, x, g))
305
306def LinearNBF(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer<TA.Mat<o, i>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
307  T.Layer{w => x => T.lnb_f2(n, i, o, TG.matmul_nt(n, i, o, 0n, x, w)), w => s => dy => g => lnb_b2(n, i, o, s, g, TG.matmul_precise(n, o, i, 0n, dy, w)), op => p => g => st => mat_step(o, i, op, p, g, st), u => TA.Mat.zeros(o, i), p => T.Mat.dump(o, i, p), p => xs => off => T.Mat.restore(o, i, p, xs, off), p => xs => off => T.Mat.store(o, i, p, xs, off), u => U32.from_nat(Nat.mul(o, i))}
308
309# ---- RMSNorm: a leaf a row with main.bend's own row code (rmsf_row) ----
310
311type ERF is Data:
312  ERF{n: U32, d: Nat, d32: U32, df: F32, eps: F32}
313
314def r3_y(s: T.R3) -> Array<F32>:
315  match s:
316    case T.R3{x, g, y}: y
317
318def rmsf_if(ok: Bool, +t: U32, +d: Nat, +d32: U32, +df: F32, +eps: F32, x: Array<F32>, g: Array<F32>, y: Array<F32>) -> Array<F32>:
319  match ok:
320    case True{}: r3_y(T.rmsf_row(d, (t * d32 : U32), df, eps, T.R3{x, g, y}))
321    case False{}: y
322
323# x1 = x, x2 = the gain g; y a row a leaf
324def lf_rmsf(+t: U32, +l: U32, +e: ERF, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, y: Array<F32>) -> Array<F32>:
325  match e:
326    case ERF{+n, +d, +d32, +df, +eps}: rmsf_if(U32.is_lt(t, n), t, d, d32, df, eps, x1, x2, y)
327
328def erf(+n: Nat, +d: Nat, +eps: F32) -> ERF:
329  ERF{U32.from_nat(n), d, U32.from_nat(d), F32.from_nat(d), eps}
330
331def rms_f2(-n: Nat, -d: Nat, q: TG.R1) -> T.Fwd<TA.Mat<1n, d>, TA.Mat<n, d>, TA.Mat<n, d>>:
332  match q:
333    case TG.R1{x1, x2, x3, x4, y}: T.Fwd{TA.Mat{x2}, TA.Mat{y}, TA.Mat{x1}}
334
335def rms_f(+n: Nat, +d: Nat, +eps: F32, p: TA.Mat<1n, d>, x: TA.Mat<n, d>) -> T.Fwd<TA.Mat<1n, d>, TA.Mat<n, d>, TA.Mat<n, d>>:
336  match p x:
337    case TA.Mat{g} TA.Mat{a}: rms_f2(n, d, TG.run1_new(~ERF, ~lf_rmsf, TA.cap_depth(n), erf(n, d, eps), a, g, nn(Unit{}), nn(Unit{}), TA.cap_depth(Nat.mul(n, d))))
338
339# the backward pass in two bangs: a leaf a row writes dx and the row's 1/rms (main.bend's rmsb_row without dg), then a
340# leaf a column sums dg down the rows in order from +0.0 and adds it to the gradient it is given, as main.bend adds each
341# row's term into dg
342
343def gdx_w(k: Nat, +j: U32, +off: U32, +r: F32, +m: F32, dx: Array<F32>, rx: Array<F32> & F32, rd: Array<F32> & F32, rg: Array<F32> & F32) -> Array<F32>:
344  match k rx rd rg:
345    case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
346      +xh = (a * r : F32)
347      Array.set(F32, dx, (off + j : U32), (((b * c : F32) - (xh * m : F32) : F32) * r : F32))
348    case 1n+p Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
349      +xh = (a * r : F32)
350      gdx_w(p, (j + 1 : U32), off, r, m, Array.set(F32, dx, (off + j : U32), (((b * c : F32) - (xh * m : F32) : F32) * r : F32)), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32)))
351
352def gdx_3(+d: Nat, +off: U32, +df: F32, +r: F32, +t: U32, dx: Array<F32>, rr: Array<F32>, q: T.R3S) -> TG.O2:
353  match q:
354    case T.R3S{x, dy, g, +s}:
355      TG.O2{gdx_w(Nat.sub(d, 1n), 0, off, r, (s / df : F32), dx, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0)), Array.set(F32, rr, t, r)}
356
357def gdx_2(+d: Nat, +off: U32, +df: F32, +eps: F32, +t: U32, dy: Array<F32>, g: Array<F32>, dx: Array<F32>, rr: Array<F32>, q: T.AS) -> TG.O2:
358  match q:
359    case T.AS{x, +s}:
360      +r = T.rms_inv(s, df, eps)
361      gdx_3(d, off, df, r, t, dx, rr, T.rmsb_dot(Nat.sub(d, 1n), 0, off, r, 0.0, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0)))
362
363def gdx_if(ok: Bool, +t: U32, +d: Nat, +d32: U32, +df: F32, +eps: F32, x: Array<F32>, dy: Array<F32>, g: Array<F32>, dx: Array<F32>, rr: Array<F32>) -> TG.O2:
364  match ok:
365    case True{}: gdx_2(d, (t * d32 : U32), df, eps, t, dy, g, dx, rr, T.ssq(Nat.sub(d, 1n), (t * d32 : U32), 0.0, Array.get(F32, x, (t * d32 : U32))))
366    case False{}: TG.O2{dx, rr}
367
368# x1 = x, x2 = dy, x3 = g; outputs dx and 1/rms of each row
369def lf_gdx(+t: U32, +l: U32, +e: ERF, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
370  match e o:
371    case ERF{+n, +d, +d32, +df, +eps} TG.O2{dx, rr}: gdx_if(U32.is_lt(t, n), t, d, d32, df, eps, x1, x2, x3, dx, rr)
372
373# acc + dy[i]·(x[i]·r[row]) down a column, i from j by d
374def gdg_loop(k: Nat, +i: U32, +d: U32, +row: U32, acc: F32, rx: Array<F32> & F32, rd: Array<F32> & F32, rr: Array<F32> & F32) -> F32:
375  match k rx rd rr:
376    case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{r, +c}: (acc + (b * (a * c : F32) : F32) : F32)
377    case 1n+p Tuple{x, +a} Tuple{dy, +b} Tuple{r, +c}:
378      gdg_loop(p, (i + d : U32), d, (row + 1 : U32), (acc + (b * (a * c : F32) : F32) : F32), Array.get(F32, x, (i + d : U32)), Array.get(F32, dy, (i + d : U32)), Array.get(F32, r, (row + 1 : U32)))
379
380def gdg_put(+j: U32, +s: F32, r: Array<F32> & F32) -> Array<F32>:
381  match r:
382    case Tuple{ga, +v}: Array.set(F32, ga, j, (v + s : F32))
383
384type EDG is Data:
385  EDG{d: U32, n: Nat}
386
387def gdg_if(ok: Bool, +j: U32, +n: Nat, +d: U32, x: Array<F32>, dy: Array<F32>, rr: Array<F32>, ga: Array<F32>) -> Array<F32>:
388  match ok:
389    case True{}: gdg_put(j, gdg_loop(Nat.sub(n, 1n), j, d, 0, 0.0, Array.get(F32, x, j), Array.get(F32, dy, j), Array.get(F32, rr, 0)), Array.get(F32, ga, j))
390    case False{}: ga
391
392# x1 = x, x2 = dy, x3 = 1/rms; the gradient of the gain, in place: a leaf a column
393def lf_gdg(+t: U32, +l: U32, +e: EDG, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, ga: Array<F32>) -> Array<F32>:
394  match e:
395    case EDG{+d, +n}: gdg_if(U32.is_lt(t, d), t, n, d, x1, x2, x3, ga)
396
397def rms_b3(-n: Nat, -d: Nat, g: Array<F32>, dx: Array<F32>, q: TG.R1) -> T.Bwd<TA.Mat<1n, d>, TA.Mat<n, d>>:
398  match q:
399    case TG.R1{x1, x2, x3, x4, ga}: T.Bwd{TA.Mat{g}, TA.Mat{dx}, TA.Mat{ga}}
400
401def rms_b2(+n: Nat, +d: Nat, ga: Array<F32>, q: TG.R2) -> T.Bwd<TA.Mat<1n, d>, TA.Mat<n, d>>:
402  match q:
403    case TG.R2{x, dy, g, x4, TG.O2{dx, rr}}: rms_b3(n, d, g, dx, TG.run1(~EDG, ~lf_gdg, TA.cap_depth(d), EDG{U32.from_nat(d), n}, x, dy, rr, x4, ga))
404
405def rms_b(+n: Nat, +d: Nat, +eps: F32, p: TA.Mat<1n, d>, x: TA.Mat<n, d>, dy: TA.Mat<n, d>, gacc: TA.Mat<1n, d>) -> T.Bwd<TA.Mat<1n, d>, TA.Mat<n, d>>:
406  match p x dy gacc:
407    case TA.Mat{g} TA.Mat{xa} TA.Mat{da} TA.Mat{ga}: rms_b2(n, d, ga, TG.run2_new(~ERF, ~lf_gdx, TA.cap_depth(n), erf(n, d, eps), xa, da, g, nn(Unit{}), TA.cap_depth(Nat.mul(n, d)), TA.cap_depth(n)))
408
409# RMSNorm over the last dimension, its gain trained (main.bend's RMSNorm)
410def RMSNorm(+n: Nat, +d: Nat, +eps: F32) -> T.Layer<TA.Mat<1n, d>, TA.Mat<n, d>, TA.Mat<n, d>, TA.Mat<n, d>>:
411  T.Layer{p => x => rms_f(n, d, eps, p, x), p => s => dy => g => rms_b(n, d, eps, p, s, dy, g), op => p => g => st => mat_step(1n, d, op, p, g, st), u => TA.Mat.zeros(1n, d), p => T.Mat.dump(1n, d, p), p => xs => off => T.Mat.restore(1n, d, p, xs, off), p => xs => off => T.Mat.store(1n, d, p, xs, off), u => U32.from_nat(d)}
412
413# ---- GroupNorm: a leaf an (image, group) block with main.bend's own block code (gnf_grp, gn_mr, gnb_sums) ----
414# Block t is image t / ng's group t % ng: its first cell is (t / ng)·hw·c and its first channel (t % ng)·cg, the
415# offsets main.bend's gnf_images and gnf_groups reach by steps of hw·c and cg.
416
417type EGN is Data:
418  EGN{nb: U32, ng: U32, hw: Nat, hwc: U32, c: U32, cg: Nat, cg32: U32, nf: F32, eps: F32}
419
420def egn(+n: Nat, +hw: Nat, +ng: Nat, +cg: Nat, +eps: F32) -> EGN:
421  EGN{U32.from_nat(Nat.mul(n, ng)), U32.from_nat(ng), hw, U32.from_nat(Nat.mul(hw, Nat.mul(ng, cg))), U32.from_nat(Nat.mul(ng, cg)), cg, U32.from_nat(cg), F32.from_nat(Nat.mul(hw, cg)), eps}
422
423def r4_y(s: T.R4) -> Array<F32>:
424  match s:
425    case T.R4{x, g, b, y}: y
426
427def gnf_if(ok: Bool, +t: U32, +ng: U32, +hw: Nat, +hwc: U32, +c: U32, +cg: Nat, +cg32: U32, +nf: F32, +eps: F32, x: Array<F32>, g: Array<F32>, b: Array<F32>, y: Array<F32>) -> Array<F32>:
428  match ok:
429    case True{}: r4_y(T.gnf_grp(hw, (U32.div(t, ng) * hwc : U32), (U32.mod(t, ng) * cg32 : U32), c, cg, nf, eps, T.R4{x, g, b, y}))
430    case False{}: y
431
432# x1 = x, x2 = the gain g, x3 = the bias b; y a block a leaf
433def lf_gnf(+t: U32, +l: U32, +e: EGN, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, y: Array<F32>) -> Array<F32>:
434  match e:
435    case EGN{+nb, +ng, +hw, +hwc, +c, +cg, +cg32, +nf, +eps}: gnf_if(U32.is_lt(t, nb), t, ng, hw, hwc, c, cg, cg32, nf, eps, x1, x2, x3, y)
436
437def gn_f2(-n: Nat, -d: Nat, q: TG.R1) -> T.Fwd<T.LN<d>, TA.Mat<n, d>, TA.Mat<n, d>>:
438  match q:
439    case TG.R1{x1, x2, x3, x4, y}: T.Fwd{T.LN{TA.Mat{x2}, TA.Mat{x3}}, TA.Mat{y}, TA.Mat{x1}}
440
441def gn_f(+n: Nat, +hw: Nat, +ng: Nat, +cg: Nat, +eps: F32, p: T.LN<Nat.mul(ng, cg)>, x: TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>) -> T.Fwd<T.LN<Nat.mul(ng, cg)>, TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>, TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>>:
442  match p x:
443    case T.LN{TA.Mat{g}, TA.Mat{b}} TA.Mat{a}: gn_f2(Nat.mul(n, hw), Nat.mul(ng, cg), TG.run1_new(~EGN, ~lf_gnf, TA.cap_depth(Nat.mul(n, ng)), egn(n, hw, ng, cg, eps), a, g, b, nn(Unit{}), TA.cap_depth(Nat.mul(Nat.mul(n, hw), Nat.mul(ng, cg)))))
444
445# the backward pass in two bangs: a leaf a block writes its cells' dx (main.bend's lnb_w without dg and db) and the
446# block's mean and 1/sqrt(var + eps) at 2t and 2t + 1; then a leaf a channel sums dy·xh and dy down its rows in order
447# from +0.0, the order main.bend adds them in, and adds the sums to the gradients it is given
448
449type GB is Type:
450  GB{x: Array<F32>, dy: Array<F32>, g: Array<F32>, dx: Array<F32>}
451
452def gb_dx(s: GB) -> Array<F32>:
453  match s:
454    case GB{x, dy, g, dx}: dx
455
456# dx[off + j] = (((dy·g - ma) - xh·mb)·r) along one row's segment
457def gnb_seg(k: Nat, +j: U32, +off: U32, +mu: F32, +r: F32, +ma: F32, +mb: F32, dx: Array<F32>, rx: Array<F32> & F32, rd: Array<F32> & F32, rg: Array<F32> & F32) -> GB:
458  match k rx rd rg:
459    case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
460      +xh = ((a - mu : F32) * r : F32)
461      GB{x, dy, g, Array.set(F32, dx, (off + j : U32), ((((b * c : F32) - ma : F32) - (xh * mb : F32) : F32) * r : F32))}
462    case 1n+q Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
463      +xh = ((a - mu : F32) * r : F32)
464      gnb_seg(q, (j + 1 : U32), off, mu, r, ma, mb, Array.set(F32, dx, (off + j : U32), ((((b * c : F32) - ma : F32) - (xh * mb : F32) : F32) * r : F32)), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32)))
465
466def gnb_segs(left: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +mu: F32, +r: F32, +ma: F32, +mb: F32, s: GB) -> GB:
467  match left s:
468    case 0n GB{x, dy, g, dx}:
469      gnb_seg(Nat.sub(cg, 1n), j, off, mu, r, ma, mb, dx, Array.get(F32, x, (off + j : U32)), Array.get(F32, dy, (off + j : U32)), Array.get(F32, g, j))
470    case 1n+l GB{x, dy, g, dx}:
471      gnb_segs(l, (off + c : U32), j, c, cg, mu, r, ma, mb, gnb_seg(Nat.sub(cg, 1n), j, off, mu, r, ma, mb, dx, Array.get(F32, x, (off + j : U32)), Array.get(F32, dy, (off + j : U32)), Array.get(F32, g, j)))
472
473def gnb_blk3(+hw: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +nf: F32, +mu: F32, +r: F32, +t: U32, dx: Array<F32>, mr: Array<F32>, q: T.R3SS) -> TG.O2:
474  match q:
475    case T.R3SS{x, dy, g, +sa, +sb}:
476      TG.O2{gb_dx(gnb_segs(Nat.sub(hw, 1n), off, j, c, cg, mu, r, (sa / nf : F32), (sb / nf : F32), GB{x, dy, g, dx})), Array.set(F32, Array.set(F32, mr, (t * 2 : U32), mu), ((t * 2 : U32) + 1 : U32), r)}
477
478def gnb_blk2(+hw: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +nf: F32, +t: U32, dy: Array<F32>, g: Array<F32>, dx: Array<F32>, mr: Array<F32>, m: T.ASM) -> TG.O2:
479  match m:
480    case T.ASM{x, +mu, +r}:
481      gnb_blk3(hw, off, j, c, cg, nf, mu, r, t, dx, mr, T.gnb_sums(Nat.sub(hw, 1n), off, j, c, cg, mu, r, T.lnb_sums(Nat.sub(cg, 1n), j, off, mu, r, 0.0, 0.0, Array.get(F32, x, (off + j : U32)), Array.get(F32, dy, (off + j : U32)), Array.get(F32, g, j))))
482
483def gnb_blk(+hw: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +nf: F32, +eps: F32, +t: U32, x: Array<F32>, dy: Array<F32>, g: Array<F32>, dx: Array<F32>, mr: Array<F32>) -> TG.O2:
484  gnb_blk2(hw, off, j, c, cg, nf, t, dy, g, dx, mr, T.gn_mr(hw, (off + j : U32), c, cg, nf, eps, x))
485
486def gnb_if(ok: Bool, +t: U32, +ng: U32, +hw: Nat, +hwc: U32, +c: U32, +cg: Nat, +cg32: U32, +nf: F32, +eps: F32, x: Array<F32>, dy: Array<F32>, g: Array<F32>, dx: Array<F32>, mr: Array<F32>) -> TG.O2:
487  match ok:
488    case True{}: gnb_blk(hw, (U32.div(t, ng) * hwc : U32), (U32.mod(t, ng) * cg32 : U32), c, cg, nf, eps, t, x, dy, g, dx, mr)
489    case False{}: TG.O2{dx, mr}
490
491# x1 = x, x2 = dy, x3 = g; outputs dx and each block's mean and 1/sqrt(var + eps)
492def lf_gnb(+t: U32, +l: U32, +e: EGN, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
493  match e o:
494    case EGN{+nb, +ng, +hw, +hwc, +c, +cg, +cg32, +nf, +eps} TG.O2{dx, mr}: gnb_if(U32.is_lt(t, nb), t, ng, hw, hwc, c, cg, cg32, nf, eps, x1, x2, x3, dx, mr)
495
496type GS is Type:
497  GS{x: Array<F32>, dy: Array<F32>, mr: Array<F32>, sg: F32, sb: F32}
498
499# sg + dy·xh and sb + dy down one image's rows of a channel, i its cell, by c
500def gdw_rows(k: Nat, +i: U32, +c: U32, +mu: F32, +r: F32, sg: F32, sb: F32, mr: Array<F32>, rx: Array<F32> & F32, rd: Array<F32> & F32) -> GS:
501  match k rx rd:
502    case 0n Tuple{x, +a} Tuple{dy, +b}:
503      GS{x, dy, mr, (sg + (b * ((a - mu : F32) * r : F32) : F32) : F32), (sb + b : F32)}
504    case 1n+q Tuple{x, +a} Tuple{dy, +b}:
505      gdw_rows(q, (i + c : U32), c, mu, r, (sg + (b * ((a - mu : F32) * r : F32) : F32) : F32), (sb + b : F32), mr, Array.get(F32, x, (i + c : U32)), Array.get(F32, dy, (i + c : U32)))
506
507def gdw_blk2(+hw: Nat, +i: U32, +c: U32, +mu: F32, sg: F32, sb: F32, x: Array<F32>, dy: Array<F32>, rr: Array<F32> & F32) -> GS:
508  match rr:
509    case Tuple{mr, +r}: gdw_rows(Nat.sub(hw, 1n), i, c, mu, r, sg, sb, mr, Array.get(F32, x, i), Array.get(F32, dy, i))
510
511def gdw_blk1(+hw: Nat, +i: U32, +bl: U32, +c: U32, sg: F32, sb: F32, x: Array<F32>, dy: Array<F32>, rm: Array<F32> & F32) -> GS:
512  match rm:
513    case Tuple{mr, +mu}: gdw_blk2(hw, i, c, mu, sg, sb, x, dy, Array.get(F32, mr, ((bl * 2 : U32) + 1 : U32)))
514
515# block bl's mean and 1/sqrt(var + eps), then its rows
516def gdw_blk(+hw: Nat, +i: U32, +bl: U32, +c: U32, s: GS) -> GS:
517  match s:
518    case GS{x, dy, mr, sg, sb}: gdw_blk1(hw, i, bl, c, sg, sb, x, dy, Array.get(F32, mr, (bl * 2 : U32)))
519
520# a channel's images in order: i steps by hw·c to the next image's first row, bl by ng to its block of the same group
521def gdw_imgs(left: Nat, +i: U32, +bl: U32, +hw: Nat, +hwc: U32, +c: U32, +ng: U32, s: GS) -> GS:
522  match left:
523    case 0n: gdw_blk(hw, i, bl, c, s)
524    case 1n+q: gdw_imgs(q, (i + hwc : U32), (bl + ng : U32), hw, hwc, c, ng, gdw_blk(hw, i, bl, c, s))
525
526def gdw_put(+j: U32, ga: Array<F32>, gb: Array<F32>, s: GS) -> TG.O2:
527  match s:
528    case GS{x, dy, mr, +sg, +sb}: TG.O2{gdg_put(j, sg, Array.get(F32, ga, j)), gdg_put(j, sb, Array.get(F32, gb, j))}
529
530type EGW is Data:
531  EGW{c: U32, n: Nat, hw: Nat, hwc: U32, ng: U32, cg: U32}
532
533def gdw_if(ok: Bool, +t: U32, +n: Nat, +hw: Nat, +hwc: U32, +c: U32, +ng: U32, +cg: U32, x: Array<F32>, dy: Array<F32>, mr: Array<F32>, ga: Array<F32>, gb: Array<F32>) -> TG.O2:
534  match ok:
535    case True{}: gdw_put(t, ga, gb, gdw_imgs(Nat.sub(n, 1n), t, U32.div(t, cg), hw, hwc, c, ng, GS{x, dy, mr, 0.0, 0.0}))
536    case False{}: TG.O2{ga, gb}
537
538# x1 = x, x2 = dy, x3 = the blocks' means and 1/sqrt(var + eps); the gain's and the bias's gradients, in place: a leaf a
539# channel
540def lf_gdw(+t: U32, +l: U32, +e: EGW, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
541  match e o:
542    case EGW{+c, +n, +hw, +hwc, +ng, +cg} TG.O2{ga, gb}: gdw_if(U32.is_lt(t, c), t, n, hw, hwc, c, ng, cg, x1, x2, x3, ga, gb)
543
544def gn_b3(-m: Nat, -d: Nat, g: Array<F32>, b: TA.Mat<1n, d>, dx: Array<F32>, q: TG.R2) -> T.Bwd<T.LN<d>, TA.Mat<m, d>>:
545  match q:
546    case TG.R2{x1, x2, x3, x4, TG.O2{ga, gb}}: T.Bwd{T.LN{TA.Mat{g}, b}, TA.Mat{dx}, T.LN{TA.Mat{ga}, TA.Mat{gb}}}
547
548def gn_b2(+n: Nat, +hw: Nat, +ng: Nat, +cg: Nat, b: TA.Mat<1n, Nat.mul(ng, cg)>, ga: Array<F32>, gb: Array<F32>, q: TG.R2) -> T.Bwd<T.LN<Nat.mul(ng, cg)>, TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>>:
549  match q:
550    case TG.R2{x, dy, g, x4, TG.O2{dx, mr}}: gn_b3(Nat.mul(n, hw), Nat.mul(ng, cg), g, b, dx, TG.run2(~EGW, ~lf_gdw, TA.cap_depth(Nat.mul(ng, cg)), EGW{U32.from_nat(Nat.mul(ng, cg)), n, hw, U32.from_nat(Nat.mul(hw, Nat.mul(ng, cg))), U32.from_nat(ng), U32.from_nat(cg)}, x, dy, mr, x4, ga, gb))
551
552def gn_b(+n: Nat, +hw: Nat, +ng: Nat, +cg: Nat, +eps: F32, p: T.LN<Nat.mul(ng, cg)>, x: TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>, dy: TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>, gacc: T.LN<Nat.mul(ng, cg)>) -> T.Bwd<T.LN<Nat.mul(ng, cg)>, TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>>:
553  match p x dy gacc:
554    case T.LN{TA.Mat{g}, b} TA.Mat{xa} TA.Mat{da} T.LN{TA.Mat{ga}, TA.Mat{gb}}: gn_b2(n, hw, ng, cg, b, ga, gb, TG.run2_new(~EGN, ~lf_gnb, TA.cap_depth(Nat.mul(n, ng)), egn(n, hw, ng, cg, eps), xa, da, g, nn(Unit{}), TA.cap_depth(Nat.mul(Nat.mul(n, hw), Nat.mul(ng, cg))), TA.cap_depth(Nat.mul(2n, Nat.mul(n, ng)))))
555
556# the gain and the bias take this file's optimizer step, as Linear's two matrices do
557def ln_step2(+d: Nat, +op: T.Opt, p: T.LN<d>, g: T.LN<d>, st: T.Two<T.OptState<TA.Mat<1n, d>>, T.OptState<TA.Mat<1n, d>>>) -> T.Step<T.LN<d>>:
558  match p g st:
559    case T.LN{w, b} T.LN{gw, gb} T.Two{sw, sb}:
560      T.ln_s2(d, opt_mat(1n, d, op, w, gw, sw), opt_mat(1n, d, op, b, gb, sb))
561
562def ln_step(+d: Nat, +op: T.Opt, p: T.LN<d>, g: T.LN<d>, st: T.OptState<T.LN<d>>) -> T.Step<T.LN<d>>:
563  ln_step2(d, op, p, g, T.OptState.split(TA.Mat<1n, d>, TA.Mat<1n, d>, T.OptState.map(T.LN<d>, T.Two<TA.Mat<1n, d>, TA.Mat<1n, d>>, st, x => T.ln_pair(d, x), x => T.ln_pair(d, x))))
564
565# GroupNorm over n images of hw pixels and ng groups of cg channels, channels last (main.bend's GroupNorm)
566def GroupNorm(+n: Nat, +hw: Nat, +ng: Nat, +cg: Nat, +eps: F32) -> T.Layer<T.LN<Nat.mul(ng, cg)>, TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>, TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>, TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>>:
567  T.Layer{p => x => gn_f(n, hw, ng, cg, eps, p, x), p => s => dy => g => gn_b(n, hw, ng, cg, eps, p, s, dy, g), op => p => g => st => ln_step(Nat.mul(ng, cg), op, p, g, st), u => T.LN{TA.Mat.zeros(1n, Nat.mul(ng, cg)), TA.Mat.zeros(1n, Nat.mul(ng, cg))}, p => T.ln_dump(Nat.mul(ng, cg), p), p => xs => off => T.ln_load(Nat.mul(ng, cg), p, xs, off), p => xs => off => T.ln_store(Nat.mul(ng, cg), p, xs, off), u => U32.from_nat(Nat.mul(2n, Nat.mul(ng, cg)))}
568
569# ---- RoPE: the tables on the host (main.bend's rope_tab, so the host's cos and sin), each pair rotated in place by
570# main.bend's rope_one, the pairs dealt to the leaves round-robin ----
571
572type ERO is Data:
573  ERO{n: U32, fwd: Bool, hh: U32, half: U32, w: U32, t: U32}
574
575def rope_loop(k: Nat, +e: U32, +l: U32, +fwd: Bool, +hh: U32, +half: U32, +w: U32, +t: U32, st: T.RS3) -> T.RS3:
576  match k:
577    case 0n: st
578    case 1n+q: rope_loop(q, (e + l : U32), l, fwd, hh, half, w, t, T.rope_one(e, fwd, hh, half, w, t, st))
579
580def rs3_x(st: T.RS3) -> Array<F32>:
581  match st:
582    case T.RS3{x, c, s}: x
583
584# x1 = cos, x2 = sin; x rotated in place
585def lf_rope(+t: U32, +l: U32, +e: ERO, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, x: Array<F32>) -> Array<F32>:
586  match e:
587    case ERO{+n, +fwd, +hh, +half, +w, +tt}: rs3_x(rope_loop(TG.share(t, n, l), t, l, fwd, hh, half, w, tt, T.RS3{x, x1, x2}))
588
589def rope_go2(+n: Nat, +w: Nat, +fwd: Bool, +t: Nat, +h: Nat, +half: Nat, x: Array<F32>, tb: T.CS) -> TA.Mat<n, w>:
590  match tb:
591    case T.CS{cs, sn}: TA.Mat{r1_a(TG.run1(~ERO, ~lf_rope, TG.depth(Nat.mul(n, Nat.mul(h, half))), ERO{U32.from_nat(Nat.mul(n, Nat.mul(h, half))), fwd, U32.from_nat(Nat.mul(h, half)), U32.from_nat(half), U32.from_nat(w), U32.from_nat(t)}, cs, sn, nn(Unit{}), nn(Unit{}), x))}
592
593def rope_go(+n: Nat, +w: Nat, +fwd: Bool, +t: Nat, +h: Nat, +half: Nat, +base: F32, m: TA.Mat<n, w>) -> TA.Mat<n, w>:
594  match m:
595    case TA.Mat{x}: rope_go2(n, w, fwd, t, h, half, x, T.rope_tab(Nat.sub(Nat.mul(t, half), 1n), 0, U32.from_nat(half), base, T.CS{T.zeros_a(Nat.mul(t, half)), T.zeros_a(Nat.mul(t, half))}))
596
597# RoPE on n = b·t rows of h heads of 2·half (main.bend's RoPE)
598def RoPE(+n: Nat, +w: Nat, +t: Nat, +h: Nat, +half: Nat, +base: F32, -e: {w == Nat.mul(h, Nat.add(half, half)) : Nat}) -> T.Layer<Unit, TA.Mat<n, w>, TA.Mat<n, w>, Unit>:
599  T.Layer{p => x => T.Fwd{p, rope_go(n, w, True{}, t, h, half, base, x), Unit{}}, p => s => dy => g => T.Bwd{p, rope_go(n, w, False{}, t, h, half, base, dy), g}, op => p => g => st => T.no_step(Unit, op, p, g, st), u => Unit{}, p => T.PL{p, Nil{}}, p => xs => off => T.PA{p, xs, off}, p => xs => off => T.PA{p, xs, off}, u => 0}
600
601type S4 is Type:
602  S4{a: Array<F32>, b: Array<F32>, c: Array<F32>, d: Array<F32>}
603
604# ---- attention in cell bangs and row bangs. A leaf handles cells or rows t, t + l, .. of one bang. A cell is one number
605# of an output: a score, a dP, a number of o, dQ, dK or dV, each one dot product with main.bend's own code (dot0, or
606# tdot, dot0's sum down a column of the packed triangle below; tsum_h adds the group's heads as main.bend's pacc). A
607# row is one row of P (and of G), whose softmax and its backward pass are sums along the row in order: a leaf a row with
608# main.bend's row code (psmc, g_set2), a tail loop over the row's i + 1 positions. Every number is computed by
609# main.bend's expression in main.bend's order, so the bits are main.bend's (but CUDA's exp in the softmax).
610#   forward:  0. cells: kᵀ  1. cells: P's scores (q·k scaled, j <= i)  2. rows: the softmax over the i + 1
611#             3. cells: o = P·V over the query's i + 1 positions
612#   backward: 0. cells: kᵀ and vᵀ  1. cells: P's scores again and dP = dO·V into G  2. rows: P's softmax, then
613#             G = (P * (dP - s)) * sc with s = dP·P over the row, as gr3  3. cells: dQ = G·K, dK = Σ Gᵀ·Q and
614#             dV = Σ Pᵀ·dO over the group's heads, dK written over kᵀ and dV over V (neither read any more)
615# The old form (a leaf per row doing the scores, the softmax and the row's hd outputs, ~8,000 steps in a row at
616# SmolLM2's t = 128, and dK and dV a leaf per key position) took 56 ms a layer's forward and backward at SmolLM2-135M's
617# sizes on the RTX 4050; this one takes 17 ms, about half of it filling the new arrays (one lane, Array.new). ----
618
619# the sizes; n is how many cells (or rows) the bang has, nw and nk the cells of dQ and of dK, tr = t(t + 1)/2 the cells
620# of a head's P
621type EAC is Data:
622  EAC{n: U32, t: Nat, t32: U32, h: U32, g: U32, gn: Nat, hd: Nat, hd32: U32, w: U32, kvw: U32, sc: F32, nw: U32, nk: U32, tr: U32}
623
624# where row i of head bh starts in P and G, which keep only the lower triangle (j <= i), row after row: main.bend's P
625# and G are t·t with zeros above the diagonal, and a product by one of those zeros adds +0.0 or -0.0 to a sum that
626# started at +0.0, which leaves it as it is; so the sums below skip them and give the same bits (for finite numbers:
627# an infinity or a NaN in a masked key's row would reach main.bend's sums as 0·inf = NaN, and does not reach these)
628def toff(+bh: U32, +i: U32, +tr: U32) -> U32:
629  ((bh * tr : U32) + U32.div((i * (i + 1 : U32) : U32), 2) : U32)
630
631def eac_n(+c: EAC) -> U32:
632  match c:
633    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}: n
634
635type S3 is Type:
636  S3{a: Array<F32>, b: Array<F32>, c: Array<F32>}
637
638def s3_c(s: S3) -> Array<F32>:
639  match s:
640    case S3{a, b, c}: c
641
642def sc_fin(y: Array<F32>, s: S3) -> TG.O2:
643  match s:
644    case S3{q, k, p}: TG.O2{p, y}
645
646# the cell e = row·w + col of o (a = P, b = V) or of dQ (a = G, b = K): row i of P (or G) for the query's head, its
647# i + 1 positions, · V's (or K's) column of the head's key/value head
648def pv_val(+e: U32, +c: EAC, a: Array<F32>, b: Array<F32>) -> T.D2:
649  match c:
650    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}:
651      +row = U32.div(e, w)
652      +col = U32.mod(e, w)
653      +bi = U32.div(row, t32)
654      +hi = U32.div(col, hd32)
655      +i = U32.mod(row, t32)
656      T.dot0(U32.to_nat((i + 1 : U32)), toff(((bi * h : U32) + hi : U32), i, tr), ((((bi * t32 : U32) * kvw : U32) + (U32.div(hi, g) * hd32 : U32) : U32) + U32.mod(col, hd32) : U32), 1, kvw, a, b)
657
658# acc + a[ia]·b[ib] + .. down column j of a head's packed triangle (row i + 1's cell is i + 1 after row i's) and b's
659# rows (sb apart), from row i on: dot0's sum, but with the step in a growing by one a row
660def tdot(k: Nat, +ia: U32, +ib: U32, +i: U32, +sb: U32, acc: F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> T.D2:
661  match k ra rb:
662    case 0n Tuple{a, x} Tuple{b, y}: T.D2{a, b, acc}
663    case 1n+p Tuple{a, +x} Tuple{b, +y}: tdot(p, (ia + (i + 1 : U32) : U32), (ib + sb : U32), (i + 1 : U32), sb, (acc + (x * y : F32) : F32), Array.get(F32, a, (ia + (i + 1 : U32) : U32)), Array.get(F32, b, (ib + sb : U32)))
664
665# acc + the group's heads' sums, in head order (main.bend's pacc): a's column j (tr apart a head) · b's column (hd
666# apart a head), rows j..t-1
667def tsum_h(k: Nat, +ia: U32, +ib: U32, +j: U32, +tn: Nat, +w: U32, +sa: U32, +sb: U32, acc: F32, r: T.D2) -> T.D2:
668  match k r:
669    case 0n T.D2{a, b, +s}: T.D2{a, b, (acc + s : F32)}
670    case 1n+q T.D2{a, b, +s}: tsum_h(q, (ia + sa : U32), (ib + sb : U32), j, tn, w, sa, sb, (acc + s : F32), tdot(tn, (ia + sa : U32), (ib + sb : U32), j, w, 0.0, Array.get(F32, a, (ia + sa : U32)), Array.get(F32, b, (ib + sb : U32))))
671
672# the cell e = row·kvw + col of dK (a = G, b = Q) or of dV (a = P, b = dO): key j = row % t, over the group's g query
673# heads: the head's column j of a (rows j..t-1) · the head's column of b
674def kv_val(+e: U32, +c: EAC, a: Array<F32>, b: Array<F32>) -> T.D2:
675  match c:
676    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}:
677      +row = U32.div(e, kvw)
678      +col = U32.mod(e, kvw)
679      +bi = U32.div(row, t32)
680      +jj = U32.div(col, hd32)
681      +j = U32.mod(row, t32)
682      +ia = (toff(((bi * h : U32) + (jj * g : U32) : U32), j, tr) + j : U32)
683      +ib = ((((bi * t32 : U32) + j : U32) * w : U32) + (((jj * g : U32) * hd32 : U32) + U32.mod(col, hd32) : U32) : U32)
684      tsum_h(Nat.sub(gn, 1n), ia, ib, j, U32.to_nat((t32 - j : U32)), w, tr, hd32, 0.0, tdot(U32.to_nat((t32 - j : U32)), ia, ib, j, w, 0.0, Array.get(F32, a, ia), Array.get(F32, b, ib)))
685
686# K and V transposed for the score cells (kᵀ: b·kv·hd·t, by sequence, key/value head, column c, position j), so
687# that neighbouring lanes, which take neighbouring keys j, read neighbouring cells. dot0 reads kᵀ with stride t: the same
688# numbers in the same order as reading k's row, so the same bits.
689def tk_src(+e: U32, +t32: U32, +kv: U32, +hd32: U32, +kvw: U32) -> U32:
690  +r = U32.div(e, t32)
691  +bj = U32.div(r, hd32)
692  +row = ((U32.div(bj, kv) * t32 : U32) + U32.mod(e, t32) : U32)
693  ((row * kvw : U32) + ((U32.mod(bj, kv) * hd32 : U32) + U32.mod(r, hd32) : U32) : U32)
694
695def tk_put(+e: U32, o: Array<F32>, r: Array<F32> & F32) -> TG.O2:
696  match r:
697    case Tuple{a, +x}: TG.O2{a, Array.set(F32, o, e, x)}
698
699def tk_loop(k: Nat, +e: U32, +l: U32, +t32: U32, +kv: U32, +hd32: U32, +kvw: U32, s: TG.O2) -> TG.O2:
700  match k s:
701    case 0n r: r
702    case 1n+q TG.O2{a, o}: tk_loop(q, (e + l : U32), l, t32, kv, hd32, kvw, tk_put(e, o, Array.get(F32, a, tk_src(e, t32, kv, hd32, kvw))))
703
704def o2_b(s: TG.O2) -> Array<F32>:
705  match s:
706    case TG.O2{a, b}: b
707
708# x1 = k; kᵀ written (made here)
709def lf_tk(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: Array<F32>) -> Array<F32>:
710  match c:
711    case EAC{+n, +tt, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}: o2_b(tk_loop(TG.share(t, n, l), t, l, t32, U32.div(h, g), hd32, kvw, TG.O2{x1, o}))
712
713# both: cells [0, nk) are kᵀ's from k, the next nk vᵀ's from v (S4{k, v, kᵀ, vᵀ})
714def tkv_k(v: Array<F32>, vt: Array<F32>, r: TG.O2) -> S4:
715  match r:
716    case TG.O2{k, kt}: S4{k, v, kt, vt}
717
718def tkv_v(k: Array<F32>, kt: Array<F32>, r: TG.O2) -> S4:
719  match r:
720    case TG.O2{v, vt}: S4{k, v, kt, vt}
721
722def tkv_sel(lo: Bool, +e: U32, +nk: U32, +t32: U32, +kv: U32, +hd32: U32, +kvw: U32, s: S4) -> S4:
723  match lo s:
724    case True{} S4{k, v, kt, vt}: tkv_k(v, vt, tk_put(e, kt, Array.get(F32, k, tk_src(e, t32, kv, hd32, kvw))))
725    case False{} S4{k, v, kt, vt}: tkv_v(k, kt, tk_put((e - nk : U32), vt, Array.get(F32, v, tk_src((e - nk : U32), t32, kv, hd32, kvw))))
726
727def tkv_loop(k: Nat, +e: U32, +l: U32, +nk: U32, +t32: U32, +kv: U32, +hd32: U32, +kvw: U32, s: S4) -> S4:
728  match k:
729    case 0n: s
730    case 1n+q: tkv_loop(q, (e + l : U32), l, nk, t32, kv, hd32, kvw, tkv_sel(U32.is_lt(e, nk), e, nk, t32, kv, hd32, kvw, s))
731
732def tkv_o(dq: Array<F32>, s: S4) -> TG.O3:
733  match s:
734    case S4{k, v, kt, vt}: TG.O3{kt, vt, dq}
735
736# x1 = k, x2 = v; outputs kᵀ and vᵀ (written), dQ (made here, written by backward 3)
737def lf_tkv(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O3) -> TG.O3:
738  match c o:
739    case EAC{+n, +tt, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr} TG.O3{kt, vt, dq}: tkv_o(dq, tkv_loop(TG.share(t, n, l), t, l, nk, t32, U32.div(h, g), hd32, kvw, S4{x1, x2, kt, vt}))
740
741# the score cells are P's lower triangle (j <= i), dealt in pairs of rows i and t - 1 - i (t + 1 cells a pair), so
742# every lane gets real work: cell e is pair p of head bh, place c; c <= p is (p, c), the rest (t - 1 - p, c - p - 1);
743# in the middle row of an odd t the second half repeats the first and is skipped (ok false)
744type IJ is Data:
745  IJ{ok: Bool, bh: U32, i: U32, j: U32}
746
747def ij_of(+e: U32, +t32: U32) -> IJ:
748  +np = U32.div((t32 + 1 : U32), 2)
749  +c = U32.mod(e, (t32 + 1 : U32))
750  +pr = U32.div(e, (t32 + 1 : U32))
751  +p = U32.mod(pr, np)
752  +lo = U32.is_le(c, p)
753  IJ{Bool.or(lo, U32.is_ne(((t32 - 1 : U32) - p : U32), p)), U32.div(pr, np), Bool.pick(U32, lo, p, ((t32 - 1 : U32) - p : U32)), Bool.pick(U32, lo, c, ((c - p : U32) - 1 : U32))}
754
755# forward 1: P's cell (i, j) = (q's row i · k's row j) * sc, as qrow's dot and pw's product (S3{q, kᵀ, p})
756def sc_put(+e: U32, +sc: F32, p: Array<F32>, r: T.D2) -> S3:
757  match r:
758    case T.D2{q, k, +s}: S3{q, k, Array.set(F32, p, e, (s * sc : F32))}
759
760def sc_if(ok: Bool, +e: U32, +qa: U32, +ka: U32, +hd: Nat, +t32: U32, +sc: F32, s: S3) -> S3:
761  match ok s:
762    case True{} S3{q, k, p}: sc_put(e, sc, p, T.dot0(hd, qa, ka, 1, t32, q, k))
763    case False{} S3{q, k, p}: S3{q, k, p}
764
765def sc_ij(+a: IJ, +t32: U32, +h: U32, +g: U32, +hd: Nat, +hd32: U32, +w: U32, +sc: F32, +tr: U32, s: S3) -> S3:
766  match a:
767    case IJ{+ok, +bh, +i, +j}: sc_if(ok, (toff(bh, i, tr) + j : U32), ((((U32.div(bh, h) * t32 : U32) + i : U32) * w : U32) + (U32.mod(bh, h) * hd32 : U32) : U32), (((U32.div(bh, g) * hd32 : U32) * t32 : U32) + j : U32), hd, t32, sc, s)
768
769def sc_cell(+e: U32, +c: EAC, s: S3) -> S3:
770  match c:
771    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}: sc_ij(ij_of(e, t32), t32, h, g, hd, hd32, w, sc, tr, s)
772
773def sc_loop(k: Nat, +e: U32, +l: U32, +c: EAC, s: S3) -> S3:
774  match k:
775    case 0n: s
776    case 1n+q: sc_loop(q, (e + l : U32), l, c, sc_cell(e, c, s))
777
778# x1 = q, x2 = k; outputs P (written) and o (made here, written by forward 3)
779def lf_sc(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
780  match o:
781    case TG.O2{p, y}: sc_fin(y, sc_loop(TG.share(t, eac_n(c), l), t, l, c, S3{x1, x2, p}))
782
783# forward 2: P's row r (query i = r % t) by main.bend's softmax over its i + 1 positions
784def sm_loop(k: Nat, +r: U32, +l: U32, +t32: U32, +tr: U32, p: Array<F32>) -> Array<F32>:
785  match k:
786    case 0n: p
787    case 1n+q: sm_loop(q, (r + l : U32), l, t32, tr, T.pc_a(T.psmc(U32.to_nat((U32.mod(r, t32) + 1 : U32)), toff(U32.div(r, t32), U32.mod(r, t32), tr), p)))
788
789def lf_sm(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, p: Array<F32>) -> Array<F32>:
790  match c:
791    case EAC{+n, +tt, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}: sm_loop(TG.share(t, n, l), t, l, t32, tr, p)
792
793# forward 3 (and backward 3's dQ): o[e] from pv_val (S3{a, b, o})
794def put_c(+e: U32, o: Array<F32>, r: T.D2) -> S3:
795  match r:
796    case T.D2{a, b, +s}: S3{a, b, Array.set(F32, o, e, s)}
797
798def pv_loop(k: Nat, +e: U32, +l: U32, +c: EAC, s: S3) -> S3:
799  match k s:
800    case 0n r: r
801    case 1n+q S3{a, b, o}: pv_loop(q, (e + l : U32), l, c, put_c(e, o, pv_val(e, c, a, b)))
802
803# x1 = P, x2 = v; o written in place
804def lf_pv(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: Array<F32>) -> Array<F32>:
805  s3_c(pv_loop(TG.share(t, eac_n(c), l), t, l, c, S3{x1, x2, o}))
806
807# backward 1: P's score cell as forward 1, and G's cell dP = dO's row · v's row, as drow's dot (S6{q, k, v, d, p, g})
808type S6 is Type:
809  S6{q: Array<F32>, k: Array<F32>, v: Array<F32>, d: Array<F32>, p: Array<F32>, g: Array<F32>}
810
811def b1_g(+e: U32, q: Array<F32>, k: Array<F32>, p: Array<F32>, g: Array<F32>, r: T.D2) -> S6:
812  match r:
813    case T.D2{d, v, +s}: S6{q, k, v, d, p, Array.set(F32, g, e, s)}
814
815def b1_p(+e: U32, +sc: F32, +qa: U32, +ka: U32, +hd: Nat, +t32: U32, v: Array<F32>, d: Array<F32>, p: Array<F32>, g: Array<F32>, r: T.D2) -> S6:
816  match r:
817    case T.D2{q, k, +s}: b1_g(e, q, k, Array.set(F32, p, e, (s * sc : F32)), g, T.dot0(hd, qa, ka, 1, t32, d, v))
818
819def b1_if(ok: Bool, +e: U32, +qa: U32, +ka: U32, +hd: Nat, +t32: U32, +sc: F32, s: S6) -> S6:
820  match ok s:
821    case True{} S6{q, k, v, d, p, g}: b1_p(e, sc, qa, ka, hd, t32, v, d, p, g, T.dot0(hd, qa, ka, 1, t32, q, k))
822    case False{} S6{q, k, v, d, p, g}: S6{q, k, v, d, p, g}
823
824def b1_ij(+a: IJ, +t32: U32, +h: U32, +g: U32, +hd: Nat, +hd32: U32, +w: U32, +sc: F32, +tr: U32, s: S6) -> S6:
825  match a:
826    case IJ{+ok, +bh, +i, +j}: b1_if(ok, (toff(bh, i, tr) + j : U32), ((((U32.div(bh, h) * t32 : U32) + i : U32) * w : U32) + (U32.mod(bh, h) * hd32 : U32) : U32), (((U32.div(bh, g) * hd32 : U32) * t32 : U32) + j : U32), hd, t32, sc, s)
827
828def b1_cell(+e: U32, +c: EAC, s: S6) -> S6:
829  match c:
830    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}: b1_ij(ij_of(e, t32), t32, h, g, hd, hd32, w, sc, tr, s)
831
832def b1_loop(k: Nat, +e: U32, +l: U32, +c: EAC, s: S6) -> S6:
833  match k:
834    case 0n: s
835    case 1n+q: b1_loop(q, (e + l : U32), l, c, b1_cell(e, c, s))
836
837def b1_o(s: S6) -> TG.O2:
838  match s:
839    case S6{q, k, v, d, p, g}: TG.O2{p, g}
840
841# x1 = q, x2 = kᵀ, x3 = vᵀ, x4 = dO; outputs P and G (made here, written)
842def lf_b1(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
843  match o:
844    case TG.O2{p, g}: b1_o(b1_loop(TG.share(t, eac_n(c), l), t, l, c, S6{x1, x2, x3, x4, p, g}))
845
846# backward 2: row r of P by the softmax, then G's row: s = dP·P over the i + 1 positions and G = (P * (dP - s)) * sc,
847# main.bend's gr3 and g_set2
848def gr_2(+n: Nat, +off: U32, +sc: F32, g: Array<F32>, r: T.PC) -> T.A2:
849  match r:
850    case T.PC{p, c}: T.g_set2(n, sc, off, T.dot0(n, off, off, 1, 1, g, p))
851
852def gr_loop(k: Nat, +r: U32, +l: U32, +t32: U32, +tr: U32, +sc: F32, s: T.A2) -> T.A2:
853  match k s:
854    case 0n r: r
855    case 1n+q T.A2{p, g}: gr_loop(q, (r + l : U32), l, t32, tr, sc, gr_2(U32.to_nat((U32.mod(r, t32) + 1 : U32)), toff(U32.div(r, t32), U32.mod(r, t32), tr), sc, g, T.psmc(U32.to_nat((U32.mod(r, t32) + 1 : U32)), toff(U32.div(r, t32), U32.mod(r, t32), tr), p)))
856
857def a2_o2(s: T.A2) -> TG.O2:
858  match s:
859    case T.A2{p, g}: TG.O2{p, g}
860
861def lf_gr(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
862  match c o:
863    case EAC{+n, +tt, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr} TG.O2{p, gb}: a2_o2(gr_loop(TG.share(t, n, l), t, l, t32, tr, sc, T.A2{p, gb}))
864
865# backward 3: cells [0, nw) are dQ's, [nw, nw + nk) dK's, the next nk dV's (S8{q, k, d, g, p, dq, dk, v})
866type S8 is Type:
867  S8{q: Array<F32>, k: Array<F32>, d: Array<F32>, g: Array<F32>, p: Array<F32>, dq: Array<F32>, dk: Array<F32>, v: Array<F32>}
868
869def b3_q(+e: U32, q: Array<F32>, d: Array<F32>, p: Array<F32>, dq: Array<F32>, dk: Array<F32>, v: Array<F32>, r: T.D2) -> S8:
870  match r:
871    case T.D2{g, k, +s}: S8{q, k, d, g, p, Array.set(F32, dq, e, s), dk, v}
872
873def b3_k(+e: U32, k: Array<F32>, d: Array<F32>, p: Array<F32>, dq: Array<F32>, dk: Array<F32>, v: Array<F32>, r: T.D2) -> S8:
874  match r:
875    case T.D2{g, q, +s}: S8{q, k, d, g, p, dq, Array.set(F32, dk, e, s), v}
876
877def b3_v(+e: U32, q: Array<F32>, k: Array<F32>, g: Array<F32>, dq: Array<F32>, dk: Array<F32>, v: Array<F32>, r: T.D2) -> S8:
878  match r:
879    case T.D2{p, d, +s}: S8{q, k, d, g, p, dq, dk, Array.set(F32, v, e, s)}
880
881def b3_sel(a: Bool, b: Bool, +e: U32, +nw: U32, +nk: U32, +c: EAC, s: S8) -> S8:
882  match a b s:
883    case True{} _ S8{q, k, d, g, p, dq, dk, v}: b3_q(e, q, d, p, dq, dk, v, pv_val(e, c, g, k))
884    case False{} True{} S8{q, k, d, g, p, dq, dk, v}: b3_k((e - nw : U32), k, d, p, dq, dk, v, kv_val((e - nw : U32), c, g, q))
885    case False{} False{} S8{q, k, d, g, p, dq, dk, v}: b3_v(((e - nw : U32) - nk : U32), q, k, g, dq, dk, v, kv_val(((e - nw : U32) - nk : U32), c, p, d))
886
887def b3_cell(+e: U32, +c: EAC, s: S8) -> S8:
888  match c:
889    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}: b3_sel(U32.is_lt(e, nw), U32.is_lt(e, (nw + nk : U32)), e, nw, nk, c, s)
890
891def b3_loop(k: Nat, +e: U32, +l: U32, +c: EAC, s: S8) -> S8:
892  match k:
893    case 0n: s
894    case 1n+q: b3_loop(q, (e + l : U32), l, c, b3_cell(e, c, s))
895
896def b3_o(s: S8) -> TG.O4:
897  match s:
898    case S8{q, k, d, g, p, dq, dk, v}: TG.O4{p, dq, dk, v}
899
900# x1 = q, x2 = k, x3 = dO, x4 = G; P read from the first output (unchanged), dQ, dK and dV (over V) written
901def lf_b3(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O4) -> TG.O4:
902  match o:
903    case TG.O4{p, dq, dk, v}: b3_o(b3_loop(TG.share(t, eac_n(c), l), t, l, c, S8{x1, x2, x3, x4, p, dq, dk, v}))
904
905# the sizes of a bang of cnt cells (or rows)
906def eac(+cnt: Nat, +n: Nat, +w: Nat, +kvw: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32) -> EAC:
907  EAC{U32.from_nat(cnt), t, U32.from_nat(t), U32.from_nat(Nat.mul(kv, g)), U32.from_nat(g), g, hd, U32.from_nat(hd), U32.from_nat(w), U32.from_nat(kvw), sc, U32.from_nat(Nat.mul(n, w)), U32.from_nat(Nat.mul(n, kvw)), U32.div((U32.from_nat(t) * (U32.from_nat(t) + 1 : U32) : U32), 2)}
908
909# P's cells (b·h·t(t + 1)/2), the score cells dealt in pairs of rows (b·h·ceil(t/2)·(t + 1)) and P's rows (b·h·t)
910def pcells(+b: Nat, +t: Nat, +kv: Nat, +g: Nat) -> Nat:
911  Nat.mul(Nat.mul(b, Nat.mul(kv, g)), Nat.div(Nat.mul(t, 1n+t), 2n))
912
913def tcells(+b: Nat, +t: Nat, +kv: Nat, +g: Nat) -> Nat:
914  Nat.mul(Nat.mul(Nat.mul(b, Nat.mul(kv, g)), Nat.div(1n+t, 2n)), 1n+t)
915
916def prows(+b: Nat, +t: Nat, +kv: Nat, +g: Nat) -> Nat:
917  Nat.mul(Nat.mul(b, Nat.mul(kv, g)), t)
918
919def af3(-n: Nat, -w: Nat, -kvw: Nat, q: Array<F32>, k: Array<F32>, r: TG.R1) -> T.Fwd<Unit, TA.Mat<n, w>, T.AttnS>:
920  match r:
921    case TG.R1{p, v, x3, x4, o}: T.Fwd{Unit{}, TA.Mat{o}, T.AttnS{q, k, v}}
922
923def af2(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, q: Array<F32>, k: Array<F32>, v: Array<F32>, o: Array<F32>, r: TG.R1) -> T.Fwd<Unit, TA.Mat<n, w>, T.AttnS>:
924  match r:
925    case TG.R1{x1, x2, x3, x4, p}: af3(n, w, kvw, q, k, TG.run1(~EAC, ~lf_pv, TG.depth(Nat.mul(n, w)), eac(Nat.mul(n, w), n, w, kvw, t, kv, g, hd, sc), p, v, nn(Unit{}), nn(Unit{}), o))
926
927def af1(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, k: Array<F32>, v: Array<F32>, r: TG.R2) -> T.Fwd<Unit, TA.Mat<n, w>, T.AttnS>:
928  match r:
929    case TG.R2{q, kt, x3, x4, TG.O2{p, o}}: af2(n, w, kvw, b, t, kv, g, hd, sc, q, k, v, o, TG.run1(~EAC, ~lf_sm, TG.depth(prows(b, t, kv, g)), eac(prows(b, t, kv, g), n, w, kvw, t, kv, g, hd, sc), nn(Unit{}), nn(Unit{}), nn(Unit{}), nn(Unit{}), p))
930
931def af0(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, q: Array<F32>, v: Array<F32>, r: TG.R1) -> T.Fwd<Unit, TA.Mat<n, w>, T.AttnS>:
932  match r:
933    case TG.R1{k, x2, x3, x4, kt}: af1(n, w, kvw, b, t, kv, g, hd, sc, k, v, TG.run2_new(~EAC, ~lf_sc, TG.depth(tcells(b, t, kv, g)), eac(tcells(b, t, kv, g), n, w, kvw, t, kv, g, hd, sc), q, kt, nn(Unit{}), nn(Unit{}), TA.cap_depth(pcells(b, t, kv, g)), TA.cap_depth(Nat.mul(n, w))))
934
935def attn_f(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, x: T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>) -> T.Fwd<Unit, TA.Mat<n, w>, T.AttnS>:
936  match x:
937    case T.Two{TA.Mat{q}, T.Two{TA.Mat{k}, TA.Mat{v}}}: af0(n, w, kvw, b, t, kv, g, hd, sc, q, v, TG.run1_new(~EAC, ~lf_tk, TG.depth(Nat.mul(n, kvw)), eac(Nat.mul(n, kvw), n, w, kvw, t, kv, g, hd, sc), k, nn(Unit{}), nn(Unit{}), nn(Unit{}), TA.cap_depth(Nat.mul(n, kvw))))
938
939def ab4(-n: Nat, -w: Nat, -kvw: Nat, r: TG.R4) -> T.Bwd<Unit, T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
940  match r:
941    case TG.R4{x1, x2, x3, x4, TG.O4{p, dq, dk, dv}}: T.Bwd{Unit{}, T.Two{TA.Mat{dq}, T.Two{TA.Mat{dk}, TA.Mat{dv}}}, Unit{}}
942
943# backward 3 writes dK over kᵀ and dV over v (neither read any more)
944def ab3(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, q: Array<F32>, k: Array<F32>, v: Array<F32>, d: Array<F32>, dq: Array<F32>, kt: Array<F32>, r: TG.R2) -> T.Bwd<Unit, T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
945  match r:
946    case TG.R2{x1, x2, x3, x4, TG.O2{p, gb}}: ab4(n, w, kvw, TG.run4(~EAC, ~lf_b3, TG.depth(Nat.add(Nat.mul(n, w), Nat.mul(2n, Nat.mul(n, kvw)))), eac(Nat.add(Nat.mul(n, w), Nat.mul(2n, Nat.mul(n, kvw))), n, w, kvw, t, kv, g, hd, sc), q, k, d, gb, p, dq, kt, v))
947
948def ab2(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, k: Array<F32>, v: Array<F32>, dq: Array<F32>, r: TG.R2) -> T.Bwd<Unit, T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
949  match r:
950    case TG.R2{q, kt, vt, d, TG.O2{p, gb}}: ab3(n, w, kvw, b, t, kv, g, hd, sc, q, k, v, d, dq, kt, TG.run2(~EAC, ~lf_gr, TG.depth(prows(b, t, kv, g)), eac(prows(b, t, kv, g), n, w, kvw, t, kv, g, hd, sc), nn(Unit{}), nn(Unit{}), nn(Unit{}), nn(Unit{}), p, gb))
951
952def ab1(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, q: Array<F32>, d: Array<F32>, r: TG.R3) -> T.Bwd<Unit, T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
953  match r:
954    case TG.R3{k, v, x3, x4, TG.O3{kt, vt, dq}}: ab2(n, w, kvw, b, t, kv, g, hd, sc, k, v, dq, TG.run2_new(~EAC, ~lf_b1, TG.depth(tcells(b, t, kv, g)), eac(tcells(b, t, kv, g), n, w, kvw, t, kv, g, hd, sc), q, kt, vt, d, TA.cap_depth(pcells(b, t, kv, g)), TA.cap_depth(pcells(b, t, kv, g))))
955
956def attn_b(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, s: T.AttnS, dy: TA.Mat<n, w>) -> T.Bwd<Unit, T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
957  match s dy:
958    case T.AttnS{q, k, v} TA.Mat{d}: ab1(n, w, kvw, b, t, kv, g, hd, sc, q, d, TG.run3_new(~EAC, ~lf_tkv, TG.depth(Nat.mul(2n, Nat.mul(n, kvw))), eac(Nat.mul(2n, Nat.mul(n, kvw)), n, w, kvw, t, kv, g, hd, sc), k, v, nn(Unit{}), nn(Unit{}), TA.cap_depth(Nat.mul(n, kvw)), TA.cap_depth(Nat.mul(n, kvw)), TA.cap_depth(Nat.mul(n, w))))
959
960# ---- FullAttention on the device: the same leaves with a dense P. main.bend's full row is t wide, so P here is t x t
961# per head (offset bh·t² + i·t) and every sum runs over all t keys; the causal lane packs the lower triangle instead.
962# The cell map is direct, e -> (bh, i, j), because there is no triangle to pair up.
963
964def toff_f(+bh: U32, +i: U32, +t32: U32, +tr: U32) -> U32:
965  ((bh * tr : U32) + i * t32 : U32)
966
967def pcf(+b: Nat, +t: Nat, +kv: Nat, +g: Nat) -> Nat:
968  Nat.mul(Nat.mul(Nat.mul(b, Nat.mul(kv, g)), t), t)
969
970def eacf(+cnt: Nat, +n: Nat, +w: Nat, +kvw: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32) -> EAC:
971  EAC{U32.from_nat(cnt), t, U32.from_nat(t), U32.from_nat(Nat.mul(kv, g)), U32.from_nat(g), g, hd, U32.from_nat(hd), U32.from_nat(w), U32.from_nat(kvw), sc, U32.from_nat(Nat.mul(n, w)), U32.from_nat(Nat.mul(n, kvw)), U32.mul(U32.from_nat(t), U32.from_nat(t))}
972
973def ij_of_f(+e: U32, +t32: U32) -> IJ:
974  +tt = (t32 * t32 : U32)
975  +r = U32.mod(e, tt)
976  IJ{True{}, U32.div(e, tt), U32.div(r, t32), U32.mod(r, t32)}
977
978def sc_ij_f(+a: IJ, +t32: U32, +h: U32, +g: U32, +hd: Nat, +hd32: U32, +w: U32, +sc: F32, +tr: U32, s: S3) -> S3:
979  match a:
980    case IJ{+ok, +bh, +i, +j}: sc_if(ok, (toff_f(bh, i, t32, tr) + j : U32), ((((U32.div(bh, h) * t32 : U32) + i : U32) * w : U32) + (U32.mod(bh, h) * hd32 : U32) : U32), (((U32.div(bh, g) * hd32 : U32) * t32 : U32) + j : U32), hd, t32, sc, s)
981
982def sc_cell_f(+e: U32, +c: EAC, s: S3) -> S3:
983  match c:
984    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}: sc_ij_f(ij_of_f(e, t32), t32, h, g, hd, hd32, w, sc, tr, s)
985
986def sc_loop_f(k: Nat, +e: U32, +l: U32, +c: EAC, s: S3) -> S3:
987  match k:
988    case 0n: s
989    case 1n+q: sc_loop_f(q, (e + l : U32), l, c, sc_cell_f(e, c, s))
990
991def lf_sc_f(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
992  match o:
993    case TG.O2{p, y}: sc_fin(y, sc_loop_f(TG.share(t, eac_n(c), l), t, l, c, S3{x1, x2, p}))
994
995# softmax over all t of row i, at the dense row start
996def sm_loop_f(k: Nat, +r: U32, +l: U32, +t32: U32, +tr: U32, p: Array<F32>) -> Array<F32>:
997  match k:
998    case 0n: p
999    case 1n+q: sm_loop_f(q, (r + l : U32), l, t32, tr, T.pc_a(T.psmc(U32.to_nat(t32), toff_f(U32.div(r, t32), U32.mod(r, t32), t32, tr), p)))
1000
1001def lf_sm_f(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, p: Array<F32>) -> Array<F32>:
1002  match c:
1003    case EAC{+n, +tt, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}: sm_loop_f(TG.share(t, n, l), t, l, t32, tr, p)
1004
1005def pv_val_f(+e: U32, +c: EAC, a: Array<F32>, b: Array<F32>) -> T.D2:
1006  match c:
1007    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}:
1008      +row = U32.div(e, w)
1009      +col = U32.mod(e, w)
1010      +bi = U32.div(row, t32)
1011      +hi = U32.div(col, hd32)
1012      +i = U32.mod(row, t32)
1013      T.dot0(t, toff_f(((bi * h : U32) + hi : U32), i, t32, tr), ((((bi * t32 : U32) * kvw : U32) + (U32.div(hi, g) * hd32 : U32) : U32) + U32.mod(col, hd32) : U32), 1, kvw, a, b)
1014
1015def pv_loop_f(k: Nat, +e: U32, +l: U32, +c: EAC, s: S3) -> S3:
1016  match k s:
1017    case 0n r: r
1018    case 1n+q S3{a, b, o}: pv_loop_f(q, (e + l : U32), l, c, put_c(e, o, pv_val_f(e, c, a, b)))
1019
1020def lf_pv_f(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: Array<F32>) -> Array<F32>:
1021  s3_c(pv_loop_f(TG.share(t, eac_n(c), l), t, l, c, S3{x1, x2, o}))
1022
1023def af3_f(-n: Nat, -w: Nat, -kvw: Nat, q: Array<F32>, k: Array<F32>, r: TG.R1) -> T.Fwd<Unit, TA.Mat<n, w>, T.AttnS>:
1024  match r:
1025    case TG.R1{p, v, x3, x4, o}: T.Fwd{Unit{}, TA.Mat{o}, T.AttnS{q, k, v}}
1026
1027def af2_f(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, q: Array<F32>, k: Array<F32>, v: Array<F32>, o: Array<F32>, r: TG.R1) -> T.Fwd<Unit, TA.Mat<n, w>, T.AttnS>:
1028  match r:
1029    case TG.R1{x1, x2, x3, x4, p}: af3_f(n, w, kvw, q, k, TG.run1(~EAC, ~lf_pv_f, TG.depth(Nat.mul(n, w)), eacf(Nat.mul(n, w), n, w, kvw, t, kv, g, hd, sc), p, v, nn(Unit{}), nn(Unit{}), o))
1030
1031def af1_f(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, k: Array<F32>, v: Array<F32>, r: TG.R2) -> T.Fwd<Unit, TA.Mat<n, w>, T.AttnS>:
1032  match r:
1033    case TG.R2{q, kt, x3, x4, TG.O2{p, o}}: af2_f(n, w, kvw, b, t, kv, g, hd, sc, q, k, v, o, TG.run1(~EAC, ~lf_sm_f, TG.depth(prows(b, t, kv, g)), eacf(prows(b, t, kv, g), n, w, kvw, t, kv, g, hd, sc), nn(Unit{}), nn(Unit{}), nn(Unit{}), nn(Unit{}), p))
1034
1035def af0_f(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, q: Array<F32>, v: Array<F32>, r: TG.R1) -> T.Fwd<Unit, TA.Mat<n, w>, T.AttnS>:
1036  match r:
1037    case TG.R1{k, x2, x3, x4, kt}: af1_f(n, w, kvw, b, t, kv, g, hd, sc, k, v, TG.run2_new(~EAC, ~lf_sc_f, TG.depth(pcf(b, t, kv, g)), eacf(pcf(b, t, kv, g), n, w, kvw, t, kv, g, hd, sc), q, kt, nn(Unit{}), nn(Unit{}), TA.cap_depth(pcf(b, t, kv, g)), TA.cap_depth(Nat.mul(n, w))))
1038
1039def attn_f_f(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, x: T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>) -> T.Fwd<Unit, TA.Mat<n, w>, T.AttnS>:
1040  match x:
1041    case T.Two{TA.Mat{q}, T.Two{TA.Mat{k}, TA.Mat{v}}}: af0_f(n, w, kvw, b, t, kv, g, hd, sc, q, v, TG.run1_new(~EAC, ~lf_tk, TG.depth(Nat.mul(n, kvw)), eacf(Nat.mul(n, kvw), n, w, kvw, t, kv, g, hd, sc), k, nn(Unit{}), nn(Unit{}), nn(Unit{}), TA.cap_depth(Nat.mul(n, kvw))))
1042
1043# ---- the full lane's backward. Every sum runs over all t keys, so the packed-triangle walks (tdot, tsum_h) become a
1044# stride-t dot on the dense row, and the row functions of backward 2 use the whole row. b1_if, b1_g, b1_p, b3_q, b3_k,
1045# b3_v, b3_o, gr_2 and the S6/S8 records are the causal lane's own code; only the cell map and the row bounds differ. ----
1046
1047# backward 1, dense: the score cell as forward 1, and G's dP = dO's row · v's row, at the dense offsets
1048def b1_ij_f(+a: IJ, +t32: U32, +h: U32, +g: U32, +hd: Nat, +hd32: U32, +w: U32, +sc: F32, +tr: U32, s: S6) -> S6:
1049  match a:
1050    case IJ{+ok, +bh, +i, +j}: b1_if(ok, (toff_f(bh, i, t32, tr) + j : U32), ((((U32.div(bh, h) * t32 : U32) + i : U32) * w : U32) + (U32.mod(bh, h) * hd32 : U32) : U32), (((U32.div(bh, g) * hd32 : U32) * t32 : U32) + j : U32), hd, t32, sc, s)
1051
1052def b1_cell_f(+e: U32, +c: EAC, s: S6) -> S6:
1053  match c:
1054    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}: b1_ij_f(ij_of_f(e, t32), t32, h, g, hd, hd32, w, sc, tr, s)
1055
1056def b1_loop_f(k: Nat, +e: U32, +l: U32, +c: EAC, s: S6) -> S6:
1057  match k:
1058    case 0n: s
1059    case 1n+q: b1_loop_f(q, (e + l : U32), l, c, b1_cell_f(e, c, s))
1060
1061# x1 = q, x2 = kᵀ, x3 = vᵀ, x4 = dO; outputs P and G-dP (made here, written)
1062def lf_b1_f(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
1063  match o:
1064    case TG.O2{p, g}: b1_o(b1_loop_f(TG.share(t, eac_n(c), l), t, l, c, S6{x1, x2, x3, x4, p, g}))
1065
1066# backward 2, dense: row r's softmax over all t, then G's row: s = dP·P over all t and G = (P * (dP - s)) * sc.
1067# gr_2's own arithmetic: it takes the row's count and offset, so the dense lane only changes those two numbers.
1068def gr_2_f(+n: Nat, +off: U32, +sc: F32, g: Array<F32>, r: T.PC) -> T.A2:
1069  gr_2(n, off, sc, g, r)
1070
1071def gr_loop_f(k: Nat, +r: U32, +l: U32, +t32: U32, +tr: U32, +sc: F32, s: T.A2) -> T.A2:
1072  match k s:
1073    case 0n r: r
1074    case 1n+q T.A2{p, g}: gr_loop_f(q, (r + l : U32), l, t32, tr, sc, gr_2_f(U32.to_nat(t32), toff_f(U32.div(r, t32), U32.mod(r, t32), t32, tr), sc, g, T.psmc(U32.to_nat(t32), toff_f(U32.div(r, t32), U32.mod(r, t32), t32, tr), p)))
1075
1076def lf_gr_f(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
1077  match c o:
1078    case EAC{+n, +tt, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr} TG.O2{p, gb}: a2_o2(gr_loop_f(TG.share(t, n, l), t, l, t32, tr, sc, T.A2{p, gb}))
1079
1080# acc + a's cells ia, ia + sa, .. times b's ib, ib + sb, .. (k of them): dot0's sum with the strides given, as tdot
1081def ddot(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, acc: F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> T.D2:
1082  match k ra rb:
1083    case 0n Tuple{a, x} Tuple{b, y}: T.D2{a, b, acc}
1084    case 1n+p Tuple{a, +x} Tuple{b, +y}: ddot(p, (ia + sa : U32), (ib + sb : U32), sa, sb, (acc + (x * y : F32) : F32), Array.get(F32, a, (ia + sa : U32)), Array.get(F32, b, (ib + sb : U32)))
1085
1086# acc + the group's heads' sums, in head order (main.bend's pacc): a's column j from row 0 (rows t32 apart) · b's
1087# column (rows w apart), all tn rows; the dense form of tsum_h and tdot
1088def dsum_h(k: Nat, +ia: U32, +ib: U32, +tn: Nat, +t32: U32, +w: U32, +sa: U32, +sb: U32, acc: F32, r: T.D2) -> T.D2:
1089  match k r:
1090    case 0n T.D2{a, b, +s}: T.D2{a, b, (acc + s : F32)}
1091    case 1n+q T.D2{a, b, +s}: dsum_h(q, (ia + sa : U32), (ib + sb : U32), tn, t32, w, sa, sb, (acc + s : F32), ddot(tn, (ia + sa : U32), (ib + sb : U32), t32, w, 0.0, Array.get(F32, a, (ia + sa : U32)), Array.get(F32, b, (ib + sb : U32))))
1092
1093# the cell e = row·kvw + col of dK (a = G, b = Q) or of dV (a = P, b = dO), dense: key j = row % t, over the group's
1094# g query heads, each the head's column j of a from row 0 (t apart a row) · the head's column of b (w apart a row)
1095def kv_val_f(+e: U32, +c: EAC, a: Array<F32>, b: Array<F32>) -> T.D2:
1096  match c:
1097    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}:
1098      +row = U32.div(e, kvw)
1099      +col = U32.mod(e, kvw)
1100      +bi = U32.div(row, t32)
1101      +jj = U32.div(col, hd32)
1102      +j = U32.mod(row, t32)
1103      +ia = (toff_f(((bi * h : U32) + (jj * g : U32) : U32), 0, t32, tr) + j : U32)
1104      +ib = (((bi * t32 : U32) * w : U32) + (((jj * g : U32) * hd32 : U32) + U32.mod(col, hd32) : U32) : U32)
1105      dsum_h(Nat.sub(gn, 1n), ia, ib, t, t32, w, tr, hd32, 0.0, ddot(t, ia, ib, t32, w, 0.0, Array.get(F32, a, ia), Array.get(F32, b, ib)))
1106
1107# backward 3, dense: cells [0, nw) are dQ's, [nw, nw + nk) dK's, the next nk dV's (S8 as the causal lane's)
1108def b3_sel_f(a: Bool, b: Bool, +e: U32, +nw: U32, +nk: U32, +c: EAC, s: S8) -> S8:
1109  match a b s:
1110    case True{} _ S8{q, k, d, g, p, dq, dk, v}: b3_q(e, q, d, p, dq, dk, v, pv_val_f(e, c, g, k))
1111    case False{} True{} S8{q, k, d, g, p, dq, dk, v}: b3_k((e - nw : U32), k, d, p, dq, dk, v, kv_val_f((e - nw : U32), c, g, q))
1112    case False{} False{} S8{q, k, d, g, p, dq, dk, v}: b3_v(((e - nw : U32) - nk : U32), q, k, g, dq, dk, v, kv_val_f(((e - nw : U32) - nk : U32), c, p, d))
1113
1114def b3_cell_f(+e: U32, +c: EAC, s: S8) -> S8:
1115  match c:
1116    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}: b3_sel_f(U32.is_lt(e, nw), U32.is_lt(e, (nw + nk : U32)), e, nw, nk, c, s)
1117
1118def b3_loop_f(k: Nat, +e: U32, +l: U32, +c: EAC, s: S8) -> S8:
1119  match k:
1120    case 0n: s
1121    case 1n+q: b3_loop_f(q, (e + l : U32), l, c, b3_cell_f(e, c, s))
1122
1123# x1 = q, x2 = k, x3 = dO, x4 = G; P read from the first output (unchanged), dQ, dK and dV (over V) written
1124def lf_b3_f(+t: U32, +l: U32, +c: EAC, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O4) -> TG.O4:
1125  match o:
1126    case TG.O4{p, dq, dk, v}: b3_o(b3_loop_f(TG.share(t, eac_n(c), l), t, l, c, S8{x1, x2, x3, x4, p, dq, dk, v}))
1127
1128# backward 3's bang: dense cells, dK over kᵀ and dV over V (neither read any more)
1129def ab3_f(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, q: Array<F32>, k: Array<F32>, v: Array<F32>, d: Array<F32>, dq: Array<F32>, kt: Array<F32>, r: TG.R2) -> T.Bwd<Unit, T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
1130  match r:
1131    case TG.R2{x1, x2, x3, x4, TG.O2{p, gb}}: ab4(n, w, kvw, TG.run4(~EAC, ~lf_b3_f, TG.depth(Nat.add(Nat.mul(n, w), Nat.mul(2n, Nat.mul(n, kvw)))), eacf(Nat.add(Nat.mul(n, w), Nat.mul(2n, Nat.mul(n, kvw))), n, w, kvw, t, kv, g, hd, sc), q, k, d, gb, p, dq, kt, v))
1132
1133# backward 2's bang: the softmax over all t, then G, in place
1134def ab2_f(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, k: Array<F32>, v: Array<F32>, dq: Array<F32>, r: TG.R2) -> T.Bwd<Unit, T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
1135  match r:
1136    case TG.R2{q, kt, vt, d, TG.O2{p, gb}}: ab3_f(n, w, kvw, b, t, kv, g, hd, sc, q, k, v, d, dq, kt, TG.run2(~EAC, ~lf_gr_f, TG.depth(prows(b, t, kv, g)), eacf(prows(b, t, kv, g), n, w, kvw, t, kv, g, hd, sc), nn(Unit{}), nn(Unit{}), nn(Unit{}), nn(Unit{}), p, gb))
1137
1138# backward 1's bang: dense P and G-dP (made here), kᵀ and vᵀ read
1139def ab1_f(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, q: Array<F32>, d: Array<F32>, r: TG.R3) -> T.Bwd<Unit, T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
1140  match r:
1141    case TG.R3{k, v, x3, x4, TG.O3{kt, vt, dq}}: ab2_f(n, w, kvw, b, t, kv, g, hd, sc, k, v, dq, TG.run2_new(~EAC, ~lf_b1_f, TG.depth(pcf(b, t, kv, g)), eacf(pcf(b, t, kv, g), n, w, kvw, t, kv, g, hd, sc), q, kt, vt, d, TA.cap_depth(pcf(b, t, kv, g)), TA.cap_depth(pcf(b, t, kv, g))))
1142
1143# x1 = q, x2 = k, x3 = v, x4 = dO; kᵀ and vᵀ made here, then dQ, dK and dV
1144def attn_b_f(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, s: T.AttnS, dy: TA.Mat<n, w>) -> T.Bwd<Unit, T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
1145  match s dy:
1146    case T.AttnS{q, k, v} TA.Mat{d}: ab1_f(n, w, kvw, b, t, kv, g, hd, sc, q, d, TG.run3_new(~EAC, ~lf_tkv, TG.depth(Nat.mul(2n, Nat.mul(n, kvw))), eacf(Nat.mul(2n, Nat.mul(n, kvw)), n, w, kvw, t, kv, g, hd, sc), k, v, nn(Unit{}), nn(Unit{}), TA.cap_depth(Nat.mul(n, kvw)), TA.cap_depth(Nat.mul(n, kvw)), TA.cap_depth(Nat.mul(n, w))))
1147
1148# causal attention with grouped heads over b sequences of t positions (main.bend's Attention)
1149def Attention(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, -e1: {n == Nat.mul(b, t) : Nat}, -e2: {w == Nat.mul(Nat.mul(kv, g), hd) : Nat}, -e3: {kvw == Nat.mul(kv, hd) : Nat}) -> T.Layer<Unit, T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>, TA.Mat<n, w>, T.AttnS>:
1150  T.Layer{p => x => attn_f(n, w, kvw, b, t, kv, g, hd, sc, x), p => s => dy => g2 => attn_b(n, w, kvw, b, t, kv, g, hd, sc, s, dy), op => p => g2 => st => T.no_step(Unit, op, p, g2, st), u => Unit{}, p => T.PL{p, Nil{}}, p => xs => off => T.PA{p, xs, off}, p => xs => off => T.PA{p, xs, off}, u => 0}
1151
1152# the same layer with every key visible (main.bend's FullAttention on the device); Attention is untouched, so the
1153# causal lane's bits stay as they are
1154def FullAttention(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, -e1: {n == Nat.mul(b, t) : Nat}, -e2: {w == Nat.mul(Nat.mul(kv, g), hd) : Nat}, -e3: {kvw == Nat.mul(kv, hd) : Nat}) -> T.Layer<Unit, T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>, TA.Mat<n, w>, T.AttnS>:
1155  T.Layer{p => x => attn_f_f(n, w, kvw, b, t, kv, g, hd, sc, x), p => s => dy => g2 => attn_b_f(n, w, kvw, b, t, kv, g, hd, sc, s, dy), op => p => g2 => st => T.no_step(Unit, op, p, g2, st), u => Unit{}, p => T.PL{p, Nil{}}, p => xs => off => T.PA{p, xs, off}, p => xs => off => T.PA{p, xs, off}, u => 0}
1156
1157# ---- MSE on the device. main.bend's MSE writes each gradient cell 2 (y - t) / (n·c) over y and sums the squares with
1158# tree_sum, whose passes add neighbouring pairs; its first p passes carry no odd cell when 2^p divides the count, so a
1159# leaf that sums a block of 2^p squares as a balanced tree of pairs gets exactly the number those passes leave at the
1160# block's place. The leaves do that on the device (y written in place), and the host finishes tree_sum over the
1161# count / 2^p block sums: the same bits as main.bend's MSE, with only the block sums read back. ----
1162
1163type EMS is Data:
1164  EMS{nb: U32, p: Nat, bs: U32, kk: F32}
1165
1166type MS is Type:
1167  MS{y: Array<F32>, t: Array<F32>, s: F32}
1168
1169def ms_add(+s: F32, r: MS) -> MS:
1170  match r:
1171    case MS{y, t, +s2}: MS{y, t, (s + s2 : F32)}
1172
1173def ms_then(r: MS, k: @+s: F32 -> Array<F32> -> Array<F32> -> MS) -> MS:
1174  match r:
1175    case MS{y, t, +s}: k(s, y, t)
1176
1177def ms_cell(+i: U32, +kk: F32, ry: Array<F32> & F32, rt: Array<F32> & F32) -> MS:
1178  match ry rt:
1179    case Tuple{y, +a} Tuple{t, +b}: MS{Array.set(F32, y, i, ((a - b : F32) * kk : F32)), t, ((a - b : F32) * (a - b : F32) : F32)}
1180
1181# the squares of the block of 2^p cells from i summed as a balanced tree of pairs (h = 2^(p-1)), gradients written
1182def ms_blk(p: Nat, +i: U32, +h: U32, +kk: F32, y: Array<F32>, t: Array<F32>) -> MS:
1183  match p:
1184    case 0n: ms_cell(i, kk, Array.get(F32, y, i), Array.get(F32, t, i))
1185    case 1n++q: ms_then(ms_blk(q, i, U32.shr(h), kk, y, t), +s => y2 => t2 => ms_add(s, ms_blk(q, (i + h : U32), U32.shr(h), kk, y2, t2)))
1186
1187def ms_put(+t: U32, ps: Array<F32>, r: MS) -> TG.O2:
1188  match r:
1189    case MS{y, tt, +s}: TG.O2{y, Array.set(F32, ps, t, s)}
1190
1191def ms_if(ok: Bool, +t: U32, +p: Nat, +bs: U32, +kk: F32, tg: Array<F32>, y: Array<F32>, ps: Array<F32>) -> TG.O2:
1192  match ok:
1193    case True{}: ms_put(t, ps, ms_blk(p, (t * bs : U32), U32.shr(bs), kk, y, tg))
1194    case False{}: TG.O2{y, ps}
1195
1196# x1 = the target; y written in place with the gradient, ps[t] = leaf t's block sum (one leaf a block)
1197def lf_mse(+t: U32, +l: U32, +e: EMS, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
1198  match e o:
1199    case EMS{+nb, +p, +bs, +kk} TG.O2{y, ps}: ms_if(U32.is_lt(t, nb), t, p, bs, kk, x1, y, ps)
1200
1201# the largest p <= k with 2^p dividing m (m > 0)
1202def ms_p(k: Nat, +m: U32) -> Nat:
1203  match k:
1204    case 0n: 0n
1205    # bend-lint: disable-next bend-lint-rules/performance/eager_branch -- a loop of at most 32 U32 steps on the host, once per bang launch not-expected-because: a launch costs about 1 ms, the whole walk well under a microsecond
1206    case 1n+q: Bool.pick(Nat, U32.is_eq(U32.mod(m, 2), 0), 1n+ms_p(q, U32.div(m, 2)), 0n)
1207
1208def mse_fin(-n: Nat, -c: Nat, +nb: U32, +nc: F32, q: TG.R2) -> T.LossOut<TA.Mat<n, c>>:
1209  match q:
1210    case TG.R2{x1, x2, x3, x4, TG.O2{y, ps}}: T.LossOut{(T.tree_sum(nb, ps) / nc : F32), TA.Mat{y}}
1211
1212def mse_go(+n: Nat, +c: Nat, +m: U32, +p: Nat, y: Array<F32>, t: Array<F32>) -> T.LossOut<TA.Mat<n, c>>:
1213  mse_fin(n, c, U32.shrn(m, p), F32.from_nat(Nat.mul(n, c)), TG.run2(~EMS, ~lf_mse, TA.cap_depth(U32.to_nat(U32.shrn(m, p))), EMS{U32.shrn(m, p), p, U32.pow(2, p), (2.0 / F32.from_nat(Nat.mul(n, c)) : F32)}, t, nn(Unit{}), nn(Unit{}), nn(Unit{}), y, T.zeros_a(U32.to_nat(U32.shrn(m, p)))))
1214
1215# mean squared error over every number (main.bend's MSE, the same bits)
1216def MSE(+n: Nat, +c: Nat, y: TA.Mat<n, c>, t: TA.Mat<n, c>) -> T.LossOut<TA.Mat<n, c>>:
1217  match y t:
1218    case TA.Mat{ya} TA.Mat{ta}: mse_go(n, c, U32.from_nat(Nat.mul(n, c)), ms_p(6n, U32.from_nat(Nat.mul(n, c))), ya, ta)
1219
1220# ---- ids and labels on the device: a list of n Nats as an array of n cells, each holding the id's U32 word as an F32's
1221# bits (a bit copy, not a conversion, so every U32 id is exact); the count is how many of the list's first n there are ----
1222
1223def u2f(x: U32) -> F32:
1224  match x:
1225    case U32{w}: F32{w}
1226
1227def f2u(x: F32) -> U32:
1228  match x:
1229    case F32{w}: U32{w}
1230
1231type IA is Type:
1232  IA{a: Array<F32>, k: Nat}
1233
1234def ia_go(left: Nat, ids: List<&2, Nat>, +i: U32, k: Nat, a: Array<F32>) -> IA:
1235  match left ids:
1236    case 1n+q Con{+id, t}: ia_go(q, t, (i + 1 : U32), 1n+k, Array.set(F32, a, i, u2f(U32.from_nat(id))))
1237    case _ _: IA{a, k}
1238
1239def ids_arr(+n: Nat, ids: List<&2, Nat>) -> IA:
1240  ia_go(n, ids, 0, 0n, T.zeros_a(n))
1241
1242# ---- Embedding: a leaf a column j of the table (d of them), each looping over the ids in order: the forward pass
1243# copies e[id_i, j] into row i, the backward adds dy[i, j] into g[id_i, j] as main.bend's emb_add does (g + dy), one id
1244# after the other, so every cell of the gradient gets its rows added in main.bend's order, with no two leaves writing
1245# one cell. Neighbouring lanes read neighbouring cells. ----
1246
1247type EEM is Data:
1248  EEM{k: Nat, d: U32}
1249
1250type ES3 is Type:
1251  ES3{ids: Array<F32>, a: Array<F32>, o: Array<F32>}
1252
1253def ef_3(+i: U32, +j: U32, +d: U32, ids: Array<F32>, o: Array<F32>, r: Array<F32> & F32) -> ES3:
1254  match r:
1255    case Tuple{e, +x}: ES3{ids, e, Array.set(F32, o, ((i * d : U32) + j : U32), x)}
1256
1257def ef_2(+i: U32, +j: U32, +d: U32, e: Array<F32>, o: Array<F32>, r: Array<F32> & F32) -> ES3:
1258  match r:
1259    case Tuple{ids, +v}: ef_3(i, j, d, ids, o, Array.get(F32, e, ((f2u(v) * d : U32) + j : U32)))
1260
1261# o[i, j] = e[ids[i], j]
1262def ef_cell(+i: U32, +j: U32, +d: U32, st: ES3) -> ES3:
1263  match st:
1264    case ES3{ids, e, o}: ef_2(i, j, d, e, o, Array.get(F32, ids, i))
1265
1266def ef_loop(k: Nat, +i: U32, +j: U32, +d: U32, st: ES3) -> ES3:
1267  match k:
1268    case 0n: st
1269    case 1n+q: ef_loop(q, (i + 1 : U32), j, d, ef_cell(i, j, d, st))
1270
1271def es3_o(st: ES3) -> Array<F32>:
1272  match st:
1273    case ES3{ids, a, o}: o
1274
1275def ef_if(ok: Bool, +t: U32, +k: Nat, +d: U32, e: Array<F32>, ids: Array<F32>, o: Array<F32>) -> Array<F32>:
1276  match ok:
1277    case True{}: es3_o(ef_loop(k, 0, t, d, ES3{ids, e, o}))
1278    case False{}: o
1279
1280# x1 = the table, x2 = the ids; o (made on the device, zeros) gets the rows
1281def lf_embf(+t: U32, +l: U32, +e: EEM, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: Array<F32>) -> Array<F32>:
1282  match e:
1283    case EEM{+k, +d}: ef_if(U32.is_lt(t, d), t, k, d, x1, x2, o)
1284
1285def eb_4(+gi: U32, ids: Array<F32>, g: Array<F32>, +a: F32, r: Array<F32> & F32) -> ES3:
1286  match r:
1287    case Tuple{dy, +b}: ES3{ids, dy, Array.set(F32, g, gi, (a + b : F32))}
1288
1289def eb_3(+i: U32, +j: U32, +d: U32, +gi: U32, ids: Array<F32>, dy: Array<F32>, r: Array<F32> & F32) -> ES3:
1290  match r:
1291    case Tuple{g, +a}: eb_4(gi, ids, g, a, Array.get(F32, dy, ((i * d : U32) + j : U32)))
1292
1293def eb_2(+i: U32, +j: U32, +d: U32, dy: Array<F32>, g: Array<F32>, r: Array<F32> & F32) -> ES3:
1294  match r:
1295    case Tuple{ids, +v}: eb_3(i, j, d, ((f2u(v) * d : U32) + j : U32), ids, dy, Array.get(F32, g, ((f2u(v) * d : U32) + j : U32)))
1296
1297# g[ids[i], j] = g[ids[i], j] + dy[i, j]
1298def eb_cell(+i: U32, +j: U32, +d: U32, st: ES3) -> ES3:
1299  match st:
1300    case ES3{ids, dy, g}: eb_2(i, j, d, dy, g, Array.get(F32, ids, i))
1301
1302def eb_loop(k: Nat, +i: U32, +j: U32, +d: U32, st: ES3) -> ES3:
1303  match k:
1304    case 0n: st
1305    case 1n+q: eb_loop(q, (i + 1 : U32), j, d, eb_cell(i, j, d, st))
1306
1307def eb_if(ok: Bool, +t: U32, +k: Nat, +d: U32, dy: Array<F32>, ids: Array<F32>, g: Array<F32>) -> Array<F32>:
1308  match ok:
1309    case True{}: es3_o(eb_loop(k, 0, t, d, ES3{ids, dy, g}))
1310    case False{}: g
1311
1312# x1 = dy, x2 = the ids; the gradient g, in place
1313def lf_embb(+t: U32, +l: U32, +e: EEM, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, g: Array<F32>) -> Array<F32>:
1314  match e:
1315    case EEM{+k, +d}: eb_if(U32.is_lt(t, d), t, k, d, x1, x2, g)
1316
1317def emb_f3(-n: Nat, -v: Nat, -d: Nat, +ids: List<&2, Nat>, q: TG.R1) -> T.Fwd<TA.Mat<v, d>, TA.Mat<n, d>, List<&2, Nat>>:
1318  match q:
1319    case TG.R1{x1, x2, x3, x4, o}: T.Fwd{TA.Mat{x1}, TA.Mat{o}, ids}
1320
1321def emb_f2(+n: Nat, +v: Nat, +d: Nat, +ids: List<&2, Nat>, ea: Array<F32>, r: IA) -> T.Fwd<TA.Mat<v, d>, TA.Mat<n, d>, List<&2, Nat>>:
1322  match r:
1323    case IA{ia, +k}: emb_f3(n, v, d, ids, TG.run1_new(~EEM, ~lf_embf, TA.cap_depth(d), EEM{k, U32.from_nat(d)}, ea, ia, nn(Unit{}), nn(Unit{}), TA.cap_depth(Nat.mul(n, d))))
1324
1325def emb_f(+n: Nat, +v: Nat, +d: Nat, +ids: List<&2, Nat>, e: TA.Mat<v, d>) -> T.Fwd<TA.Mat<v, d>, TA.Mat<n, d>, List<&2, Nat>>:
1326  match e:
1327    case TA.Mat{ea}: emb_f2(n, v, d, ids, ea, ids_arr(n, ids))
1328
1329def emb_b3(-v: Nat, -d: Nat, e: TA.Mat<v, d>, q: TG.R1) -> T.Bwd<TA.Mat<v, d>, List<&2, Nat>>:
1330  match q:
1331    case TG.R1{x1, x2, x3, x4, g}: T.Bwd{e, Nil{}, TA.Mat{g}}
1332
1333def emb_b2(+v: Nat, +d: Nat, e: TA.Mat<v, d>, da: Array<F32>, ga: Array<F32>, r: IA) -> T.Bwd<TA.Mat<v, d>, List<&2, Nat>>:
1334  match r:
1335    case IA{ia, +k}: emb_b3(v, d, e, TG.run1(~EEM, ~lf_embb, TA.cap_depth(d), EEM{k, U32.from_nat(d)}, da, ia, nn(Unit{}), nn(Unit{}), ga))
1336
1337def emb_bwd(+n: Nat, +v: Nat, +d: Nat, e: TA.Mat<v, d>, ids: List<&2, Nat>, dy: TA.Mat<n, d>, g: TA.Mat<v, d>) -> T.Bwd<TA.Mat<v, d>, List<&2, Nat>>:
1338  match dy g:
1339    case TA.Mat{da} TA.Mat{ga}: emb_b2(v, d, e, da, ga, ids_arr(n, ids))
1340
1341# a table of v rows of d numbers, looked up by n ids (main.bend's Embedding, the same bits); its step on the device
1342def Embedding(+n: Nat, +v: Nat, +d: Nat) -> T.Layer<TA.Mat<v, d>, List<&2, Nat>, TA.Mat<n, d>, List<&2, Nat>>:
1343  T.Layer{p => x => emb_f(n, v, d, x, p), p => s => dy => g => emb_bwd(n, v, d, p, s, dy, g), op => p => g => st => mat_step(v, d, op, p, g, st), u => TA.Mat.zeros(v, d), p => T.Mat.dump(v, d, p), p => xs => off => T.Mat.restore(v, d, p, xs, off), p => xs => off => T.Mat.store(v, d, p, xs, off), u => U32.from_nat(Nat.mul(v, d))}
1344
1345# ---- cross-entropy on the device, the same numbers as main.bend's CrossEntropy (but CUDA's exp on the GPU), in four
1346# bangs; only 3 numbers a row come back to the host, and the gradient stays on the device:
1347# 1. a leaf a row: its max m by main.bend's pmax0 (the same order of comparisons) and x[y] - m;
1348# 2. a leaf for each block of 2^p cells of a row (2^p dividing c, p <= 9): e = exp(x - m) written over each cell, as
1349#    pexp writes it, and the block's sum as a balanced tree of pairs. tree_sum's first p passes over a row add exactly
1350#    these pairs (no odd cell is carried while 2^p divides the count), so a block's sum is the number they leave there;
1351# 3. a leaf a row: tree_sum's remaining passes over its c / 2^p block sums (tnode below), which is s, cesm's sum;
1352# 4. a leaf for 2^14 cells at a time: ((e / s) - [j == y]) / n, pdiv's and ce_g's two expressions.
1353# The host takes log(s) and adds the rows' terms in order from 0.0, as ce_rows does. Rows past the labels are left as
1354# they are, as in main.bend; labels past the n rows are not read (main.bend would write past the matrix). ----
1355
1356type TS is Type:
1357  TS{a: Array<F32>, s: F32}
1358
1359def ts_leaf(r: Array<F32> & F32) -> TS:
1360  match r:
1361    case Tuple{a, +x}: TS{a, x}
1362
1363def ts_plus(+s: F32, r: TS) -> TS:
1364  match r:
1365    case TS{a, +s2}: TS{a, (s + s2 : F32)}
1366
1367def ts_then(r: TS, k: @+s: F32 -> Array<F32> -> TS) -> TS:
1368  match r:
1369    case TS{a, +s}: k(s, a)
1370
1371# tree_sum of the len cells from lo, read in place: a node l passes above the cells sums its left child (the first h =
1372# 2^(l-1) cells) and its right child (the rest) when there is a rest (big = h < len), and is its left child otherwise
1373# (tree_sum's odd last cell carried up a pass)
1374def tnode(l: Nat, big: Bool, +lo: U32, +len: U32, +h: U32, a: Array<F32>) -> TS:
1375  match l big:
1376    case 0n _: ts_leaf(Array.get(F32, a, lo))
1377    case 1n++q True{}: ts_then(tnode(q, U32.is_lt(U32.shr(h), h), lo, h, U32.shr(h), a), +s => a2 => ts_plus(s, tnode(q, U32.is_lt(U32.shr(h), (len - h : U32)), (lo + h : U32), (len - h : U32), U32.shr(h), a2)))
1378    case 1n+q False{}: tnode(q, U32.is_lt(U32.shr(h), len), lo, len, U32.shr(h), a)
1379
1380# the passes tree_sum makes over m cells: the smallest l with 2^l >= m (f is fuel)
1381def ts_lv(f: Nat, +m: U32, +p: U32) -> Nat:
1382  match f:
1383    case 0n: 0n
1384    # bend-lint: disable-next bend-lint-rules/performance/eager_branch -- a loop of at most 32 U32 steps on the host, once per bang launch not-expected-because: a launch costs about 1 ms, the whole walk well under a microsecond
1385    case 1n+q: Bool.pick(Nat, U32.is_lt(p, m), 1n+ts_lv(q, m, (p * 2 : U32)), 0n)
1386
1387# the host's part, row by row in order: loss - ((x[y] - m) - log(s)), or main.bend's clamped term for a label past the
1388# last class
1389type LS is Type:
1390  LS{a: Array<F32>, l: F32}
1391
1392def ce_t2(ok: Bool, +loss: F32, +xm: F32, r: Array<F32> & F32) -> LS:
1393  match ok r:
1394    case True{} Tuple{a, +s}: LS{a, (loss - (xm - F32.log(s) : F32) : F32)}
1395    case False{} Tuple{a, s}: LS{a, (loss - F32.log(F32.max(0.0, 0.000000000001)) : F32)}
1396
1397def ce_g4(+n: Nat, -c: Nat, z: Array<F32>, r: LS) -> T.LossOut<TA.Mat<n, c>>:
1398  match r:
1399    case LS{a, +l}: T.LossOut{(l / F32.from_nat(n) : F32), TA.Mat{z}}
1400
1401# 1. a leaf a row: st[3t] = x[y] - m, st[3t + 1] = m (x[y] read only when y is a class)
1402type EC1 is Data:
1403  EC1{k: U32, c1: Nat, c32: U32}
1404
1405def c1_3(+t: U32, +m: F32, st: Array<F32>, r: Array<F32> & F32) -> Array<F32>:
1406  match r:
1407    case Tuple{z, +xy}: Array.set(F32, Array.set(F32, st, (t * 3 : U32), (xy - m : F32)), ((t * 3 : U32) + 1 : U32), m)
1408
1409def c1_2(ok: Bool, +t: U32, +off: U32, +y: U32, st: Array<F32>, r: T.AS) -> Array<F32>:
1410  match ok r:
1411    case True{} T.AS{z, +m}: c1_3(t, m, st, Array.get(F32, z, (off + y : U32)))
1412    case False{} T.AS{z, +m}: c1_3(t, m, st, Tuple{z, 0.0})
1413
1414def c1_1(+t: U32, +off: U32, +c1: Nat, +c32: U32, z: Array<F32>, st: Array<F32>, r: Array<F32> & F32) -> Array<F32>:
1415  match r:
1416    case Tuple{la, +yv}: c1_2(U32.is_lt(f2u(yv), c32), t, off, f2u(yv), st, T.pmax0(c1, off, Array.get(F32, z, off)))
1417
1418def c1_if(ok: Bool, +t: U32, +c1: Nat, +c32: U32, lab: Array<F32>, z: Array<F32>, st: Array<F32>) -> Array<F32>:
1419  match ok:
1420    case True{}: c1_1(t, (t * c32 : U32), c1, c32, z, st, Array.get(F32, lab, t))
1421    case False{}: st
1422
1423def lf_c1(+t: U32, +l: U32, +e: EC1, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, st: Array<F32>) -> Array<F32>:
1424  match e:
1425    case EC1{+k, +c1, +c32}: c1_if(U32.is_lt(t, k), t, c1, c32, x1, x2, st)
1426
1427# 2. the 2^l cells from lo: e = exp(x - m) written over each, and their sum as a balanced tree of pairs (h = 2^(l-1))
1428def tsx_leaf(+lo: U32, +m: F32, r: Array<F32> & F32) -> TS:
1429  match r:
1430    case Tuple{a, +x}:
1431      +e = F32.exp((x - m : F32))
1432      TS{Array.set(F32, a, lo, e), e}
1433
1434def tsx(l: Nat, +lo: U32, +h: U32, +m: F32, a: Array<F32>) -> TS:
1435  match l:
1436    case 0n: tsx_leaf(lo, m, Array.get(F32, a, lo))
1437    case 1n++q: ts_then(tsx(q, lo, U32.shr(h), m, a), +s => a2 => ts_plus(s, tsx(q, (lo + h : U32), U32.shr(h), m, a2)))
1438
1439# the same over 8 cells in straight-line code (no closure a node): e0..e7 written, then ((e0 + e1) + (e2 + e3)) +
1440# ((e4 + e5) + (e6 + e7)), the balanced tree's three levels
1441def x8_7(+lo: U32, +m: F32, +e0: F32, +e1: F32, +e2: F32, +e3: F32, +e4: F32, +e5: F32, +e6: F32, r: Array<F32> & F32) -> TS:
1442  match r:
1443    case Tuple{a, +x}:
1444      +e7 = F32.exp((x - m : F32))
1445      TS{Array.set(F32, a, (lo + 7 : U32), e7), ((((e0 + e1 : F32) + (e2 + e3 : F32) : F32) + ((e4 + e5 : F32) + (e6 + e7 : F32) : F32)) : F32)}
1446
1447def x8_6(+lo: U32, +m: F32, +e0: F32, +e1: F32, +e2: F32, +e3: F32, +e4: F32, +e5: F32, r: Array<F32> & F32) -> TS:
1448  match r:
1449    case Tuple{a, +x}:
1450      +e = F32.exp((x - m : F32))
1451      x8_7(lo, m, e0, e1, e2, e3, e4, e5, e, Array.get(F32, Array.set(F32, a, (lo + 6 : U32), e), (lo + 7 : U32)))
1452
1453def x8_5(+lo: U32, +m: F32, +e0: F32, +e1: F32, +e2: F32, +e3: F32, +e4: F32, r: Array<F32> & F32) -> TS:
1454  match r:
1455    case Tuple{a, +x}:
1456      +e = F32.exp((x - m : F32))
1457      x8_6(lo, m, e0, e1, e2, e3, e4, e, Array.get(F32, Array.set(F32, a, (lo + 5 : U32), e), (lo + 6 : U32)))
1458
1459def x8_4(+lo: U32, +m: F32, +e0: F32, +e1: F32, +e2: F32, +e3: F32, r: Array<F32> & F32) -> TS:
1460  match r:
1461    case Tuple{a, +x}:
1462      +e = F32.exp((x - m : F32))
1463      x8_5(lo, m, e0, e1, e2, e3, e, Array.get(F32, Array.set(F32, a, (lo + 4 : U32), e), (lo + 5 : U32)))
1464
1465def x8_3(+lo: U32, +m: F32, +e0: F32, +e1: F32, +e2: F32, r: Array<F32> & F32) -> TS:
1466  match r:
1467    case Tuple{a, +x}:
1468      +e = F32.exp((x - m : F32))
1469      x8_4(lo, m, e0, e1, e2, e, Array.get(F32, Array.set(F32, a, (lo + 3 : U32), e), (lo + 4 : U32)))
1470
1471def x8_2(+lo: U32, +m: F32, +e0: F32, +e1: F32, r: Array<F32> & F32) -> TS:
1472  match r:
1473    case Tuple{a, +x}:
1474      +e = F32.exp((x - m : F32))
1475      x8_3(lo, m, e0, e1, e, Array.get(F32, Array.set(F32, a, (lo + 2 : U32), e), (lo + 3 : U32)))
1476
1477def x8_1(+lo: U32, +m: F32, +e0: F32, r: Array<F32> & F32) -> TS:
1478  match r:
1479    case Tuple{a, +x}:
1480      +e = F32.exp((x - m : F32))
1481      x8_2(lo, m, e0, e, Array.get(F32, Array.set(F32, a, (lo + 1 : U32), e), (lo + 2 : U32)))
1482
1483def x8_0(+lo: U32, +m: F32, r: Array<F32> & F32) -> TS:
1484  match r:
1485    case Tuple{a, +x}:
1486      +e = F32.exp((x - m : F32))
1487      x8_1(lo, m, e, Array.get(F32, Array.set(F32, a, lo, e), (lo + 1 : U32)))
1488
1489# tsx over 2^(l+3) cells, its last three levels in blocks of 8
1490def tsx8(l: Nat, +lo: U32, +h: U32, +m: F32, a: Array<F32>) -> TS:
1491  match l:
1492    case 0n: x8_0(lo, m, Array.get(F32, a, lo))
1493    case 1n++q: ts_then(tsx8(q, lo, U32.shr(h), m, a), +s => a2 => ts_plus(s, tsx8(q, (lo + h : U32), U32.shr(h), m, a2)))
1494
1495# a leaf takes the blocks b, b + l, .. of the nb = rows * per; block b is row b / per, cells (b % per) * 2^p on
1496type EC2 is Data:
1497  EC2{nb: U32, per: U32, p: Nat, bsz: U32, c32: U32}
1498
1499type C3S is Type:
1500  C3S{a: Array<F32>, b: Array<F32>, c: Array<F32>}
1501
1502def c2_put(+b: U32, st: Array<F32>, bs: Array<F32>, r: TS) -> C3S:
1503  match r:
1504    case TS{z, +s}: C3S{st, z, Array.set(F32, bs, b, s)}
1505
1506def c2_m(b8: Bool, +b: U32, +lo: U32, +p: Nat, +h: U32, z: Array<F32>, bs: Array<F32>, r: Array<F32> & F32) -> C3S:
1507  match b8 r:
1508    case True{} Tuple{st, +m}: c2_put(b, st, bs, tsx8(Nat.sub(p, 3n), lo, h, m, z))
1509    case False{} Tuple{st, +m}: c2_put(b, st, bs, tsx(p, lo, h, m, z))
1510
1511def c2_cell(+b: U32, +per: U32, +p: Nat, +bsz: U32, +c32: U32, s: C3S) -> C3S:
1512  match s:
1513    case C3S{st, z, bs}: c2_m(Nat.is_le(3n, p), b, ((U32.div(b, per) * c32 : U32) + (U32.mod(b, per) * bsz : U32) : U32), p, U32.shr(bsz), z, bs, Array.get(F32, st, ((U32.div(b, per) * 3 : U32) + 1 : U32)))
1514
1515def c2_loop(k: Nat, +b: U32, +l: U32, +per: U32, +p: Nat, +bsz: U32, +c32: U32, s: C3S) -> C3S:
1516  match k:
1517    case 0n: s
1518    case 1n+q: c2_loop(q, (b + l : U32), l, per, p, bsz, c32, c2_cell(b, per, p, bsz, c32, s))
1519
1520def c2_o(s: C3S) -> TG.O2:
1521  match s:
1522    case C3S{st, z, bs}: TG.O2{z, bs}
1523
1524def lf_c2(+t: U32, +l: U32, +e: EC2, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
1525  match e o:
1526    case EC2{+nb, +per, +p, +bsz, +c32} TG.O2{z, bs}: c2_o(c2_loop(TG.share(t, nb, l), t, l, per, p, bsz, c32, C3S{x1, z, bs}))
1527
1528# 3. a leaf a row: st[3t + 2] = s, tree_sum's remaining passes over the row's per block sums
1529type EC3 is Data:
1530  EC3{k: U32, per: U32, tl: Nat, th: U32}
1531
1532def c3_put(+t: U32, st: Array<F32>, r: TS) -> Array<F32>:
1533  match r:
1534    case TS{bs, +s}: Array.set(F32, st, ((t * 3 : U32) + 2 : U32), s)
1535
1536def c3_if(ok: Bool, +t: U32, +per: U32, +tl: Nat, +th: U32, bs: Array<F32>, st: Array<F32>) -> Array<F32>:
1537  match ok:
1538    case True{}: c3_put(t, st, tnode(tl, U32.is_lt(th, per), (t * per : U32), per, th, bs))
1539    case False{}: st
1540
1541def lf_c3(+t: U32, +l: U32, +e: EC3, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, st: Array<F32>) -> Array<F32>:
1542  match e:
1543    case EC3{+k, +per, +tl, +th}: c3_if(U32.is_lt(t, k), t, per, tl, th, x1, st)
1544
1545# 4. element-wise over the first k rows: z[i] = ((e / s) - [j == y]) / nf, j = i % c, s and y of row i / c
1546type EC4 is Data:
1547  EC4{n: U32, c32: U32, nf: F32}
1548
1549def c4_4(+i: U32, +j: U32, +y: U32, +s: F32, +nf: F32, lab: Array<F32>, st: Array<F32>, r: Array<F32> & F32) -> C3S:
1550  match r:
1551    case Tuple{z, +x}: C3S{lab, st, Array.set(F32, z, i, (((x / s : F32) - T.hit(U32.is_eq(j, y)) : F32) / nf : F32))}
1552
1553def c4_3(+i: U32, +j: U32, +y: U32, +nf: F32, lab: Array<F32>, z: Array<F32>, r: Array<F32> & F32) -> C3S:
1554  match r:
1555    case Tuple{st, +s}: c4_4(i, j, y, s, nf, lab, st, Array.get(F32, z, i))
1556
1557def c4_2(+i: U32, +row: U32, +j: U32, +nf: F32, st: Array<F32>, z: Array<F32>, r: Array<F32> & F32) -> C3S:
1558  match r:
1559    case Tuple{lab, +yv}: c4_3(i, j, f2u(yv), nf, lab, z, Array.get(F32, st, ((row * 3 : U32) + 2 : U32)))
1560
1561def c4_cell(+i: U32, +c32: U32, +nf: F32, s: C3S) -> C3S:
1562  match s:
1563    case C3S{lab, st, z}: c4_2(i, U32.div(i, c32), (i - (U32.div(i, c32) * c32 : U32) : U32), nf, st, z, Array.get(F32, lab, U32.div(i, c32)))
1564
1565def c4_loop(k: Nat, +i: U32, +l: U32, +c32: U32, +nf: F32, s: C3S) -> C3S:
1566  match k:
1567    case 0n: s
1568    case 1n+q: c4_loop(q, (i + l : U32), l, c32, nf, c4_cell(i, c32, nf, s))
1569
1570def c3s_c(s: C3S) -> Array<F32>:
1571  match s:
1572    case C3S{a, b, c}: c
1573
1574def lf_c4(+t: U32, +l: U32, +e: EC4, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, z: Array<F32>) -> Array<F32>:
1575  match e:
1576    case EC4{+n, +c32, +nf}: c3s_c(c4_loop(TG.share(t, n, l), t, l, c32, nf, C3S{x1, x2, z}))
1577
1578# the host's part over st, row by row
1579def cx_t1(ok: Bool, +i: U32, +loss: F32, r: Array<F32> & F32) -> LS:
1580  match r:
1581    case Tuple{a, +xm}: ce_t2(ok, loss, xm, Array.get(F32, a, ((i * 3 : U32) + 2 : U32)))
1582
1583def cx_ls(left: Nat, ls: List<&2, Nat>, +i: U32, +c: Nat, r: LS) -> LS:
1584  match left ls r:
1585    case 1n+q Con{+y, t} LS{a, +loss}: cx_ls(q, t, (i + 1 : U32), c, cx_t1(Nat.is_le(1n+y, c), i, loss, Array.get(F32, a, (i * 3 : U32))))
1586    case _ _ r: r
1587
1588def cx_5(+n: Nat, +c: Nat, +labels: List<&2, Nat>, +k: Nat, q: TG.R1) -> T.LossOut<TA.Mat<n, c>>:
1589  match q:
1590    case TG.R1{x1, x2, x3, x4, z}: ce_g4(n, c, z, cx_ls(k, labels, 0, c, LS{x2, 0.0}))
1591
1592def cx_4(+n: Nat, +c: Nat, +labels: List<&2, Nat>, +k: Nat, la: Array<F32>, z: Array<F32>, q: TG.R1) -> T.LossOut<TA.Mat<n, c>>:
1593  match q:
1594    case TG.R1{x1, x2, x3, x4, st}: cx_5(n, c, labels, k, TG.run1(~EC4, ~lf_c4, TG.depth(Nat.mul(k, c)), EC4{U32.from_nat(Nat.mul(k, c)), U32.from_nat(c), F32.from_nat(n)}, la, st, nn(Unit{}), nn(Unit{}), z))
1595
1596def cx_3(+n: Nat, +c: Nat, +labels: List<&2, Nat>, +k: Nat, +per: U32, la: Array<F32>, q: TG.R2) -> T.LossOut<TA.Mat<n, c>>:
1597  match q:
1598    case TG.R2{st, x2, x3, x4, TG.O2{z, bs}}:
1599      +lv = ts_lv(32n, per, 1)
1600      cx_4(n, c, labels, k, la, z, TG.run1(~EC3, ~lf_c3, TA.cap_depth(k), EC3{U32.from_nat(k), per, lv, U32.shr(U32.pow(2, lv))}, bs, nn(Unit{}), nn(Unit{}), nn(Unit{}), st))
1601
1602def cx_2(+n: Nat, +c: Nat, +labels: List<&2, Nat>, +k: Nat, +p: Nat, q: TG.R1) -> T.LossOut<TA.Mat<n, c>>:
1603  match q:
1604    case TG.R1{la, z, x3, x4, st}: cx_3(n, c, labels, k, U32.shrn(U32.from_nat(c), p), la, TG.run2(~EC2, ~lf_c2, TG.depth(Nat.mul(k, U32.to_nat(U32.shrn(U32.from_nat(c), p)))), EC2{(U32.from_nat(k) * U32.shrn(U32.from_nat(c), p) : U32), U32.shrn(U32.from_nat(c), p), p, U32.pow(2, p), U32.from_nat(c)}, st, nn(Unit{}), nn(Unit{}), nn(Unit{}), z, T.zeros_a(Nat.mul(k, U32.to_nat(U32.shrn(U32.from_nat(c), p))))))
1605
1606def cx_1(+n: Nat, +c: Nat, +labels: List<&2, Nat>, z: Array<F32>, r: IA) -> T.LossOut<TA.Mat<n, c>>:
1607  match r:
1608    case IA{la, +k}: cx_2(n, c, labels, k, ms_p(9n, U32.from_nat(c)), TG.run1(~EC1, ~lf_c1, TA.cap_depth(n), EC1{U32.from_nat(k), Nat.sub(c, 1n), U32.from_nat(c)}, la, z, nn(Unit{}), nn(Unit{}), T.zeros_a(Nat.mul(3n, n))))
1609
1610# cross-entropy over n rows of c classes against the labels (main.bend's CrossEntropy; its bits but CUDA's exp)
1611def CrossEntropy(+n: Nat, +c: Nat, z: TA.Mat<n, c>, +labels: List<&2, Nat>) -> T.LossOut<TA.Mat<n, c>>:
1612  match z:
1613    case TA.Mat{d}: cx_1(n, c, labels, d, ids_arr(n, labels))
1614
1615# ---- main.bend's names for the same layers (#13) ----
1616#
1617# A program written against torch/main.bend switches to this lane by its import line alone:
1618# `import ../torch/gpu.bend as F` and F.Linear / F.LinearFirst / F.LinearNB, no renames.
1619
1620def Linear(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer<T.Lin<i, o>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
1621  LinearF(n, i, o, par)
1622
1623def LinearFirst(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer<T.Lin<i, o>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
1624  LinearFirstF(n, i, o, par)
1625
1626def LinearNB(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer<TA.Mat<o, i>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
1627  LinearNBF(n, i, o, par)