main.bend

bendtorch-fast@1.3.0.0 · 362 lines · 66 definitions

1# BendTorch's layers over bend-ml-tensor-fork: the same layers, types and bits as torch/main.bend's, faster where the
2# work splits into parallel blocks.
3#
4#   import ../torch/main.bend as T
5#   import ../torch/fast.bend as F
6#   F.LinearFirstF(100n, 784n, 128n, 3n)     # instead of T.LinearFirst(100n, 784n, 128n, 3n)
7#
8# LinearF, LinearFirstF and LinearNBF are main.bend's Linear, LinearFirst and LinearNB with their three products
9# (x·Wᵀ, dY·W, dYᵀ·X added into dW) taken from bend-ml-tensor-fork instead of bend-ml-tensor-array, and their optimizer
10# step in blocks. SiLU, ActIn, Mul2, Add2, Residual, Mat.dup and Mat.sum2 are main.bend's with their element-wise loops
11# in blocks (tensor-fork's par1..par4), the optimizer step (mat_step, lin_step) too. Every number is computed by
12# main.bend's own loop or tensor-array's own cells, so the bits are main.bend's at every depth and thread count (tests:
13# demo/fast_test.bend). Small matrices run in one block (ew_depth), as in main.bend.
14#
15# Outside the proofs, on purpose: tensor-fork shares arrays through Base's @unsafe Array.fork, so `bend` names the defs
16# below that reach it, and no file with laws imports this one (`make laws` checks both). All of it is Bend.
17
18import Base
19import bend-ml-tensor-array@0.1.18.0/main.bend as TA
20import bend-ml-tensor-fork@0.1.5.0/main.bend as TF
21import bendtorch@1.3.0.0/main.bend as T
22
23# =====================================================================================================
24# Element-wise work and the optimizer step in parallel blocks (bend-ml-tensor-fork's par1..par4): every number is
25# computed by main.bend's own loop, so the bits are main.bend's at every depth and thread count. Small matrices run in
26# one block (ew_depth).
27# =====================================================================================================
28
29# log2 of the blocks for cnt numbers: one block below 131072 (MNIST's and the tiny Llama's matrices; at 100352 the
30# blocks cost MNIST more than they gave on 16 threads), else one per 16384, at most 64
31def ew_d2(small: Bool, big: Bool, +cnt: Nat) -> Nat:
32  match small big:
33    case True{} _:
34      0n
35    case False{} True{}:
36      6n
37    case False{} False{}:
38      TF.depth(Nat.div(cnt, 16384n))
39
40def ew_depth(+cnt: Nat) -> Nat:
41  ew_d2(Nat.is_lt(cnt, 131072n), Nat.is_le(1048576n, cnt), cnt)
42
43# ---- the optimizer step on a matrix ----
44
45def p2_b2(q: TA.P2) -> TF.B2<F32, F32>:
46  match q:
47    case TA.P2{a, b}:
48      TF.B2{a, b}
49
50def sgd_leaf(+lr: F32, +i0: U32, +cnt: Nat, r: TF.B2<F32, F32>) -> TF.B2<F32, F32>:
51  match r:
52    case TF.B2{w, g}:
53      p2_b2(T.sgd_l(cnt, i0, lr, Array.get(F32, w, i0), Array.get(F32, g, i0)))
54
55def mom_b(q: T.Q3) -> TF.B3<F32, F32, F32>:
56  match q:
57    case T.Q3{w, g, m}: TF.B3{w, g, m}
58
59def mom_leaf(+op: T.Opt, +i0: U32, +cnt: Nat, r: TF.B3<F32, F32, F32>) -> TF.B3<F32, F32, F32>:
60  match op r:
61    case T.Momentum{+lr, +mu} TF.B3{w, g, m}:
62      mom_b(T.mom_l(cnt, i0, lr, mu, Array.get(F32, w, i0), Array.get(F32, g, i0), Array.get(F32, m, i0)))
63    case _ r: r
64
65def q4_b(q: T.Q4) -> TF.B4<F32, F32, F32, F32>:
66  match q:
67    case T.Q4{w, g, m, v}: TF.B4{w, g, m, v}
68
69def opt_leaf(+o: T.Opt, +i0: U32, +cnt: Nat, r: TF.B4<F32, F32, F32, F32>) -> TF.B4<F32, F32, F32, F32>:
70  match o r:
71    case T.Adam{+lr, +b1, +b2, +eps, +wd, +dec, +step, +bc2, +c1, +c2} TF.B4{w, g, m, v}:
72      q4_b(T.adam_l(cnt, i0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, w, i0), Array.get(F32, g, i0), Array.get(F32, m, i0), Array.get(F32, v, i0)))
73    case _ r: r
74
75def sgd_fin(-r: Nat, -c: Nat, q: TF.B2<F32, F32>) -> T.MatStep<r, c>:
76  match q:
77    case TF.B2{w, g}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.NoMom{}}
78
79def mom_fin(-r: Nat, -c: Nat, q: TF.B3<F32, F32, F32>) -> T.MatStep<r, c>:
80  match q:
81    case TF.B3{w, g, m}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.OneMom{TA.Mat{m}}}
82
83def opt_fin(-r: Nat, -c: Nat, q: TF.B4<F32, F32, F32, F32>) -> T.MatStep<r, c>:
84  match q:
85    case TF.B4{w, g, m, v}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.TwoMom{TA.Mat{m}, TA.Mat{v}}}
86
87# ---- AdamW over packed BF16 moments (CAdamWBF16), in blocks of whole words ----
88# Two moments share a word, so the blocks split the words, not the elements: a block of words i0 .. i0 + cnt - 1 runs
89# main.bend's own loop (adamh_l, or adamhv_l when v stays F32) over elements 2·i0 .. 2·(i0 + cnt) - 1, the last block
90# stopping at the matrix's n. No word is written by two blocks, and each element's arithmetic is main.bend's.
91
92type PE is Data:
93  PE{o: T.Opt, n: U32}
94
95# the elements of a block of cnt words from element e0, the last block stopping at n
96def ph_cnt(+n: U32, +e0: U32, +cnt: Nat) -> Nat:
97  +k = U32.from_nat(Nat.add(cnt, cnt))
98  U32.to_nat(T.ph_pick(U32.is_lt(k, (n - e0 : U32)), k, (n - e0 : U32)))
99
100def pq_b(q: T.PQ) -> TF.B4<F32, F32, U32, U32>:
101  match q:
102    case T.PQ{w, g, m, v}: TF.B4{w, g, m, v}
103
104def pqv_b(q: T.PQV) -> TF.B4<F32, F32, U32, F32>:
105  match q:
106    case T.PQV{w, g, m, v}: TF.B4{w, g, m, v}
107
108def ph_go(+o: T.Opt, +n: U32, +i0: U32, +cnt: Nat, r: TF.B4<F32, F32, U32, U32>) -> TF.B4<F32, F32, U32, U32>:
109  match o r:
110    case T.Adam{+lr, +b1, +b2, +eps, +wd, +dec, +step, +bc2, +c1, +c2} TF.B4{w, g, m, v}:
111      +e0 = (i0 + i0 : U32)
112      pq_b(T.adamh_l(ph_cnt(n, e0, cnt), e0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, w, e0), Array.get(F32, g, e0), Array.get(U32, m, i0), Array.get(U32, v, i0)))
113    case _ r: r
114
115def ph_leaf(+e: PE, +i0: U32, +cnt: Nat, r: TF.B4<F32, F32, U32, U32>) -> TF.B4<F32, F32, U32, U32>:
116  match e:
117    case PE{+o, +n}: ph_go(o, n, i0, cnt, r)
118
119def phv_go(+o: T.Opt, +n: U32, +i0: U32, +cnt: Nat, r: TF.B4<F32, F32, U32, F32>) -> TF.B4<F32, F32, U32, F32>:
120  match o r:
121    case T.Adam{+lr, +b1, +b2, +eps, +wd, +dec, +step, +bc2, +c1, +c2} TF.B4{w, g, m, v}:
122      +e0 = (i0 + i0 : U32)
123      pqv_b(T.adamhv_l(ph_cnt(n, e0, cnt), e0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, w, e0), Array.get(F32, g, e0), Array.get(U32, m, i0), Array.get(F32, v, e0)))
124    case _ r: r
125
126def phv_leaf(+e: PE, +i0: U32, +cnt: Nat, r: TF.B4<F32, F32, U32, F32>) -> TF.B4<F32, F32, U32, F32>:
127  match e:
128    case PE{+o, +n}: phv_go(o, n, i0, cnt, r)
129
130def ph_fin(-r: Nat, -c: Nat, +n: U32, q: TF.B4<F32, F32, U32, U32>) -> T.MatStep<r, c>:
131  match q:
132    case TF.B4{w, g, m, v}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.Packed{T.PLeaf{n, m, v}}}
133
134def phv_fin(-r: Nat, -c: Nat, +n: U32, q: TF.B4<F32, F32, U32, F32>) -> T.MatStep<r, c>:
135  match q:
136    case TF.B4{w, g, m, v}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.Packed{T.PLeafV{n, m, v}}}
137
138# the words of r·c elements, ceil(r·c / 2)
139def ph_nw(+r: Nat, +c: Nat) -> Nat:
140  Nat.div(Nat.add(Nat.mul(r, c), 1n), 2n)
141
142def ph_leafs(+r: Nat, +c: Nat, +o: T.Opt, wa: Array<F32>, ga: Array<F32>, t: T.PTree) -> T.MatStep<r, c>:
143  match o t:
144    case T.Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} T.PLeaf{+n, ma, va}:
145      ph_fin(r, c, n, TF.par4(~PE, ~F32, ~F32, ~U32, ~U32, ~ph_leaf, PE{o, n}, ew_depth(Nat.mul(r, c)), ph_nw(r, c), wa, ga, ma, va))
146    case T.Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} T.PLeafV{+n, ma, va}:
147      phv_fin(r, c, n, TF.par4(~PE, ~F32, ~F32, ~U32, ~F32, ~phv_leaf, PE{o, n}, ew_depth(Nat.mul(r, c)), ph_nw(r, c), wa, ga, ma, va))
148    case _ t:
149      T.MatStep{TA.Mat{wa}, TA.Mat{ga}, T.Packed{t}}
150
151def ph_mat(+r: Nat, +c: Nat, +o: T.Opt, wa: Array<F32>, ga: Array<F32>, t: T.PTree) -> T.MatStep<r, c>:
152  match t:
153    case T.PNew{v16}:
154      ph_leafs(r, c, o, wa, ga, T.ph_new(Nat.mul(r, c), v16))
155    case t:
156      ph_leafs(r, c, o, wa, ga, t)
157
158# SGD splits two arrays, momentum three, Adam four. No unused parameter-sized buffers.
159def 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>:
160  match o w g st:
161    case T.SGD{+lr} TA.Mat{wd} TA.Mat{gd} _:
162      sgd_fin(r, c, TF.par2(~F32, ~F32, ~F32, ~sgd_leaf, lr, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), wd, gd))
163    case T.Momentum{lr, mu} TA.Mat{wd} TA.Mat{gd} T.OneMom{TA.Mat{md}}:
164      mom_fin(r, c, TF.par3(~T.Opt, ~F32, ~F32, ~F32, ~mom_leaf, o, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), wd, gd, md))
165    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}}:
166      opt_fin(r, c, TF.par4(~T.Opt, ~F32, ~F32, ~F32, ~F32, ~opt_leaf, o, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), wa, ga, ma, va))
167    # packed BF16 moments (CAdamWBF16): blocks of whole words
168    case T.Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} TA.Mat{wa} TA.Mat{ga} T.Packed{t}:
169      ph_mat(r, c, o, wa, ga, t)
170    case _ w g st: T.MatStep{w, g, st}
171
172def 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>>:
173  T.mat_step_fin(r, c, opt_mat(r, c, op, p, g, st))
174
175def 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>>:
176  match p g st:
177    case T.Lin{w, b} T.Lin{gw, gb} T.Two{sw, sb}:
178      T.lin_s2(i, o, opt_mat(o, i, op, w, gw, sw), opt_mat(1n, o, op, b, gb, sb))
179
180def 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>>:
181  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))))
182
183# ---- element-wise maps, zips and copies ----
184
185def ew1_leaf(~f: @+a: F32 -> F32, +e: Unit, +i0: U32, +cnt: Nat, a: Array<F32>) -> Array<F32>:
186  T.ew1_l(~f, cnt, i0, Array.get(F32, a, i0))
187
188def map(~f: @+a: F32 -> F32, +r: Nat, +c: Nat, m: TA.Mat<r, c>) -> TA.Mat<r, c>:
189  match m:
190    case TA.Mat{d}:
191      TA.Mat{TF.par1(~Unit, ~F32, ~(e => i0 => k => a => ew1_leaf(~f, e, i0, k, a)), Unit{}, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), d)}
192
193
194def ew2_leaf(~f: @+a: F32 -> @+b: F32 -> F32, +e: Unit, +i0: U32, +cnt: Nat, r: TF.B2<F32, F32>) -> TF.B2<F32, F32>:
195  match r:
196    case TF.B2{a, b}:
197      p2_b2(T.ew2_l(~f, cnt, i0, Array.get(F32, a, i0), Array.get(F32, b, i0)))
198
199def zip_fin(-r: Nat, -c: Nat, q: TF.B2<F32, F32>) -> TA.MPair<r, c>:
200  match q:
201    case TF.B2{a, b}:
202      TA.MPair{TA.Mat{a}, TA.Mat{b}}
203
204# a[i] = f(a[i], b[i]); b comes back unchanged
205def zip(~f: @+a: F32 -> @+b: F32 -> F32, +r: Nat, +c: Nat, a: TA.Mat<r, c>, b: TA.Mat<r, c>) -> TA.MPair<r, c>:
206  match a b:
207    case TA.Mat{x} TA.Mat{y}:
208      zip_fin(r, c, TF.par2(~Unit, ~F32, ~F32, ~(e => i0 => k => q => ew2_leaf(~f, e, i0, k, q)), Unit{}, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), x, y))
209
210def a2_b2(q: T.A2) -> TF.B2<F32, F32>:
211  match q:
212    case T.A2{a, b}:
213      TF.B2{a, b}
214
215# numbers i0 .. i0+cnt-1 of a into b
216def copy_leaf(+e: Unit, +i0: U32, +cnt: Nat, r: TF.B2<F32, F32>) -> TF.B2<F32, F32>:
217  match cnt r:
218    case 0n r:
219      r
220    case 1n+k TF.B2{a, b}:
221      a2_b2(T.copy_n(k, i0, i0, b, Array.get(F32, a, i0)))
222
223def dup_fin(-r: Nat, -c: Nat, q: TF.B2<F32, F32>) -> T.Two<TA.Mat<r, c>, TA.Mat<r, c>>:
224  match q:
225    case TF.B2{a, b}:
226      T.Two{TA.Mat{a}, TA.Mat{b}}
227
228# two copies of a matrix (T.Mat.dup)
229def Mat.dup(+r: Nat, +c: Nat, m: TA.Mat<r, c>) -> T.Two<TA.Mat<r, c>, TA.Mat<r, c>>:
230  match m:
231    case TA.Mat{a}:
232      dup_fin(r, c, TF.par2(~Unit, ~F32, ~F32, ~copy_leaf, Unit{}, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), a, T.zeros_a(Nat.mul(r, c))))
233
234def mpair_a(-r: Nat, -c: Nat, q: TA.MPair<r, c>) -> TA.Mat<r, c>:
235  match q:
236    case TA.MPair{a, b}:
237      a
238
239# a + b, dropping b (T.Mat.sum2)
240def Mat.sum2(+r: Nat, +c: Nat, a: TA.Mat<r, c>, b: TA.Mat<r, c>) -> TA.Mat<r, c>:
241  mpair_a(r, c, zip(~T.f32_add, r, c, a, b))
242
243# ---- Linear: y = x·Wᵀ + b ----
244
245def lin_fwd(+n: Nat, +i: Nat, +o: Nat, par: 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>>:
246  match p:
247    case T.Lin{w, b}:
248      T.lin_f2(n, i, o, b, TF.matmul_nt(n, i, o, par, x, w))
249
250def lin_b2(+n: Nat, +i: Nat, +o: Nat, par: 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>>:
251  match g r:
252    case T.Lin{gw, gb} TA.MMul{dy, w, dx}:
253      T.lin_b3(n, i, o, w, b, dx, gb, TF.matmul_tn_acc(o, n, i, par, dy, x, gw))
254
255def lin_bwd(+n: Nat, +i: Nat, +o: Nat, +par: 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>>:
256  match p:
257    case T.Lin{w, b}:
258      lin_b2(n, i, o, par, b, x, g, TF.matmul_precise(n, o, i, par, dy, w))
259
260def 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>>:
261  T.Layer{p => x => lin_fwd(n, i, o, par, p, x), p => s => dy => g => lin_bwd(n, i, o, par, 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))}
262
263# ---- LinearFirst: no dX (the first layer's input needs no gradient) ----
264
265def linf_bwd(+n: Nat, +i: Nat, +o: Nat, +par: 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>>:
266  match p g:
267    case T.Lin{w, b} T.Lin{gw, gb}:
268      T.linf_b3(n, i, o, w, b, gb, TF.matmul_tn_acc(o, n, i, par, dy, x, gw))
269
270def 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>>:
271  T.Layer{p => x => lin_fwd(n, i, o, par, p, x), p => s => dy => g => linf_bwd(n, i, o, par, 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))}
272
273# ---- LinearNB: y = x·Wᵀ, no bias (Llama's projections) ----
274
275def lnb_b2(+n: Nat, +i: Nat, +o: Nat, +par: 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>>:
276  match r:
277    case TA.MMul{dy, w, dx}:
278      T.lnb_b3(n, i, o, w, dx, TF.matmul_tn_acc(o, n, i, par, dy, x, g))
279
280def 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>>:
281  T.Layer{w => x => T.lnb_f2(n, i, o, TF.matmul_nt(n, i, o, par, x, w)), w => s => dy => g => lnb_b2(n, i, o, par, s, g, TF.matmul_precise(n, o, i, par, 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))}
282
283# ---- element-wise layers ----
284
285def acti_f3(~f: @+a: F32 -> F32, +n: Nat, +c: Nat, m: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> T.Fwd<Unit, TA.Mat<n, c>, TA.Mat<n, c>>:
286  match m:
287    case T.Two{x, s}:
288      T.Fwd{Unit{}, map(~f, n, c, x), s}
289
290# T.ActIn: f on a copy, the input saved
291def ActIn(~f: @+a: F32 -> F32, ~d: @+a: F32 -> @+b: F32 -> F32, +n: Nat, +c: Nat) -> T.Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
292  T.Layer{p => x => acti_f3(~f, n, c, Mat.dup(n, c, x)), p => s => dy => g => T.act_b2(n, c, zip(~d, n, c, dy, s)), 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}
293
294def SiLU(+n: Nat, +c: Nat) -> T.Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
295  ActIn(~T.silu_v, ~T.silu_d, n, c)
296
297def 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>:
298  match x:
299    case T.Two{a, b}:
300      T.Fwd{Unit{}, Mat.sum2(n, c, a, b), Unit{}}
301
302def Add2(+n: Nat, +c: Nat) -> T.Layer<Unit, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>, TA.Mat<n, c>, Unit>:
303  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}
304
305def mul2_f2(+n: Nat, +c: Nat, da: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>, db: 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>>>:
306  match da db:
307    case T.Two{a1, a2} T.Two{b1, b2}:
308      T.mul2_f3(n, c, a2, b2, zip(~T.f32_mul, n, c, a1, b1))
309
310def 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>>>:
311  match x:
312    case T.Two{a, b}:
313      mul2_f2(n, c, Mat.dup(n, c, a), Mat.dup(n, c, b))
314
315def mul2_b2(+n: Nat, +c: Nat, s: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>, d: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> T.Bwd<Unit, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
316  match s d:
317    case T.Two{a, b} T.Two{d1, d2}:
318      T.Bwd{Unit{}, T.Two{mpair_a(n, c, zip(~T.f32_mul, n, c, d1, b)), mpair_a(n, c, zip(~T.f32_mul, n, c, d2, a))}, Unit{}}
319
320def 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>>>:
321  T.Layer{p => x => mul2_f(n, c, x), p => s => dy => g => mul2_b2(n, c, s, Mat.dup(n, c, 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}
322
323# T.Residual with these copies and sums (T.ResidualOf, the composition the law residual is about)
324def 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>>:
325  T.ResidualOf(~TA.Mat<n, c>, ~P, ~S, ~(m => Mat.dup(n, c, m)), ~(a => b => Mat.sum2(n, c, a, b)), ~Add2(n, c), ~f)
326
327# main.bend's RMSNorm, RoPE, Attention, MSE and CrossEntropy under their own names (Embedding below, with its step in
328# blocks), so that a program that takes its layers from here can take device versions of the same names from another
329# file by its import line alone (they run as in main.bend)
330def 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>>:
331  T.RMSNorm(n, d, eps)
332
333def 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>:
334  T.RoPE(n, w, t, h, half, base, e)
335
336def 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>:
337  T.Attention(n, w, kvw, b, t, kv, g, hd, sc, e1, e2, e3)
338
339def MSE(+n: Nat, +c: Nat, y: TA.Mat<n, c>, t: TA.Mat<n, c>) -> T.LossOut<TA.Mat<n, c>>:
340  T.MSE(n, c, y, t)
341
342# main.bend's Embedding with its optimizer step in blocks (mat_step): the table is SmolLM2's largest matrix (49152 x 576,
343# tied to the output head), and main.bend's step updates it on one thread
344def Embedding(+n: Nat, +v: Nat, +d: Nat) -> T.Layer<TA.Mat<v, d>, List<&2, Nat>, TA.Mat<n, d>, List<&2, Nat>>:
345  T.Layer{p => x => T.emb_f(n, v, d, x, p), p => s => dy => g => T.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))}
346
347def CrossEntropy(+n: Nat, +c: Nat, z: TA.Mat<n, c>, +labels: List<&2, Nat>) -> T.LossOut<TA.Mat<n, c>>:
348  T.CrossEntropy(n, c, z, labels)
349
350# ---- main.bend's names for the same layers (#13) ----
351#
352# A program written against torch/main.bend switches to this lane by its import line alone:
353# `import ../torch/fast.bend as F` and F.Linear / F.LinearFirst / F.LinearNB, no renames.
354
355def 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>>:
356  LinearF(n, i, o, par)
357
358def 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>>:
359  LinearFirstF(n, i, o, par)
360
361def 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>>:
362  LinearNBF(n, i, o, par)