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 (bend-ml 23 # docs/upstream/gpu-math-overloads.md); 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 35 import Base 36 import bend-ml-tensor-array@0.1.14.0/main.bend as TA 37 import bend-ml-tensor-gpu@0.1.2.1/main.bend as TG 38 import bendtorch@1.0.0.0/main.bend as T 39 40 # ---- Linear: y = x·Wᵀ + b ---- 41 42 def lin_fwd(+n: Nat, +i: Nat, +o: Nat, p: T.Lin, x: TA.Mat) -> T.Fwd, TA.Mat, TA.Mat>: 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 47 def lin_b2(+n: Nat, +i: Nat, +o: Nat, b: TA.Mat<1n, o>, x: TA.Mat, g: T.Lin, r: TA.MMul) -> T.Bwd, TA.Mat>: 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 52 def lin_bwd(+n: Nat, +i: Nat, +o: Nat, p: T.Lin, x: TA.Mat, dy: TA.Mat, g: T.Lin) -> T.Bwd, TA.Mat>: 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 64 type EN is Data: 65 EN{n: U32} 66 67 def nn(u: Unit) -> Array: 68 TG.none(u) 69 70 # out[i] = f(x[i]) for the leaf's cells 71 def ew1_loop(~f: @+a: F32 -> F32, k: Nat, +i: U32, +l: U32, o: Array, r: Array & F32) -> Array: 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 76 def ew1(~f: @+a: F32 -> F32, +t: U32, +l: U32, +e: EN, x: Array, o: Array) -> Array: 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 81 def ew2_loop(~f: @+a: F32 -> @+b: F32 -> F32, k: Nat, +i: U32, +l: U32, o: Array, ra: Array & F32, rb: Array & F32) -> Array: 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 86 def ew2(~f: @+a: F32 -> @+b: F32 -> F32, +t: U32, +l: U32, +e: EN, x: Array, y: Array, o: Array) -> Array: 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) 91 def ewi_loop(~f: @+a: F32 -> @+b: F32 -> F32, k: Nat, +i: U32, +l: U32, ro: Array & F32, ry: Array & F32) -> Array: 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 96 def ewi(~f: @+a: F32 -> @+b: F32 -> F32, +t: U32, +l: U32, +e: EN, o: Array, y: Array) -> Array: 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 100 def f32_id(+x: F32) -> F32: 101 x 102 103 def lf_copy(+t: U32, +l: U32, +e: EN, x1: Array, x2: Array, x3: Array, x4: Array, o: Array) -> Array: 104 ew1(~f32_id, t, l, e, x1, o) 105 106 def lf_silu(+t: U32, +l: U32, +e: EN, x1: Array, x2: Array, x3: Array, x4: Array, o: Array) -> Array: 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) 110 def lf_silud(+t: U32, +l: U32, +e: EN, x1: Array, x2: Array, x3: Array, x4: Array, o: Array) -> Array: 111 ewi(~T.silu_d, t, l, e, o, x1) 112 113 # o = a, written in place, x1 = b: a + b 114 def lf_add(+t: U32, +l: U32, +e: EN, x1: Array, x2: Array, x3: Array, x4: Array, o: Array) -> Array: 115 ewi(~T.f32_add, t, l, e, o, x1) 116 117 def lf_mul(+t: U32, +l: U32, +e: EN, x1: Array, x2: Array, x3: Array, x4: Array, o: Array) -> Array: 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 121 def mulb_loop(k: Nat, +i: U32, +l: U32, rd: Array & F32, ra: Array & F32, rb: Array & 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 128 def lf_mulb(+t: U32, +l: U32, +e: EN, x1: Array, x2: Array, x3: Array, x4: Array, 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 132 def en(+cnt: Nat) -> EN: 133 EN{U32.from_nat(cnt)} 134 135 def r1_a(r: TG.R1) -> Array: 136 match r: 137 case TG.R1{x1, x2, x3, x4, a}: a 138 139 # ---- element-wise layers ---- 140 141 def dup_fin(-r: Nat, -c: Nat, q: TG.R1) -> T.Two, TA.Mat>: 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 146 def Mat.dup(+r: Nat, +c: Nat, m: TA.Mat) -> T.Two, TA.Mat>: 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) 151 def Mat.sum2(+r: Nat, +c: Nat, a: TA.Mat, b: TA.Mat) -> TA.Mat: 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 155 def silu_f2(-n: Nat, -c: Nat, q: TG.R1) -> T.Fwd, TA.Mat>: 156 match q: 157 case TG.R1{x1, x2, x3, x4, a}: T.Fwd{Unit{}, TA.Mat{a}, TA.Mat{x1}} 158 159 def silu_f(+n: Nat, +c: Nat, x: TA.Mat) -> T.Fwd, TA.Mat>: 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 163 def silu_b(+n: Nat, +c: Nat, s: TA.Mat, dy: TA.Mat) -> T.Bwd>: 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 168 def SiLU(+n: Nat, +c: Nat) -> T.Layer, TA.Mat, TA.Mat>: 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 171 def mul2_f2(-n: Nat, -c: Nat, q: TG.R1) -> T.Fwd, T.Two, TA.Mat>>: 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 175 def mul2_f(+n: Nat, +c: Nat, x: T.Two, TA.Mat>) -> T.Fwd, T.Two, TA.Mat>>: 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 179 def mul2_b2(-n: Nat, -c: Nat, q: TG.R2) -> T.Bwd, TA.Mat>>: 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 183 def mul2_b(+n: Nat, +c: Nat, s: T.Two, TA.Mat>, dy: TA.Mat) -> T.Bwd, TA.Mat>>: 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 188 def Mul2(+n: Nat, +c: Nat) -> T.Layer, TA.Mat>, TA.Mat, T.Two, TA.Mat>>: 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 191 def add2_f(+n: Nat, +c: Nat, x: T.Two, TA.Mat>) -> T.Fwd, Unit>: 192 match x: 193 case T.Two{a, b}: T.Fwd{Unit{}, Mat.sum2(n, c, a, b), Unit{}} 194 195 def Add2(+n: Nat, +c: Nat) -> T.Layer, TA.Mat>, TA.Mat, 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 199 def Residual(~n: Nat, ~c: Nat, ~P: Type, ~S: Type, ~f: T.Layer, TA.Mat, S>) -> T.Layer, Unit>, TA.Mat, TA.Mat, T.Two, Unit>>: 200 T.Seq(~T.Two, ~Unit, ~TA.Mat, ~T.Two, TA.Mat>, ~TA.Mat, ~T.Two, ~Unit, ~T.Fork(~Unit, ~P, ~TA.Mat, ~TA.Mat, ~TA.Mat, ~Unit, ~S, ~(m => Mat.dup(n, c, m)), ~(a => b => Mat.sum2(n, c, a, b)), ~T.Id(~TA.Mat), ~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 204 type ESGD is Data: 205 ESGD{n: U32, lr: F32} 206 207 def sgd_loop(k: Nat, +i: U32, +l: U32, +lr: F32, rw: Array & F32, rg: Array & 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 213 def lf_sgd(+t: U32, +l: U32, +e: ESGD, x1: Array, x2: Array, x3: Array, x4: Array, 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 217 type EMOM is Data: 218 EMOM{n: U32, lr: F32, mu: F32} 219 220 def mom_loop(k: Nat, +i: U32, +l: U32, +lr: F32, +mu: F32, rw: Array & F32, rg: Array & F32, rm: Array & 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 228 def lf_mom(+t: U32, +l: U32, +e: EMOM, x1: Array, x2: Array, x3: Array, x4: Array, 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 232 type 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 235 def 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, rg: Array & F32, rm: Array & F32, rv: Array & 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 246 def lf_adam(+t: U32, +l: U32, +e: EADAM, x1: Array, x2: Array, x3: Array, x4: Array, 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 250 def sgd_fin(-r: Nat, -c: Nat, q: TG.R2) -> T.MatStep: 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 254 def mom_fin(-r: Nat, -c: Nat, q: TG.R3) -> T.MatStep: 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 258 def opt_fin(-r: Nat, -c: Nat, q: TG.R4) -> T.MatStep: 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 263 def opt_mat(+r: Nat, +c: Nat, +o: T.Opt, w: TA.Mat, g: TA.Mat, st: T.OptState>) -> T.MatStep: 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 case _ w g st: T.MatStep{w, g, st} 272 273 def mat_step(+r: Nat, +c: Nat, +op: T.Opt, p: TA.Mat, g: TA.Mat, st: T.OptState>) -> T.Step>: 274 T.mat_step_fin(r, c, opt_mat(r, c, op, p, g, st)) 275 276 def lin_step2(+i: Nat, +o: Nat, +op: T.Opt, p: T.Lin, g: T.Lin, st: T.Two>, T.OptState>>) -> T.Step>: 277 match p g st: 278 case T.Lin{w, b} T.Lin{gw, gb} T.Two{sw, sb}: 279 T.lin_s2(i, o, opt_mat(o, i, op, w, gw, sw), opt_mat(1n, o, op, b, gb, sb)) 280 281 def lin_step(+i: Nat, +o: Nat, +op: T.Opt, p: T.Lin, g: T.Lin, st: T.OptState>) -> T.Step>: 282 lin_step2(i, o, op, p, g, T.OptState.split(TA.Mat, TA.Mat<1n, o>, T.OptState.map(T.Lin, T.Two, TA.Mat<1n, o>>, st, x => T.lin_pair(i, o, x), x => T.lin_pair(i, o, x)))) 283 284 # ---- the linear layers ---- 285 286 def LinearF(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer, TA.Mat, TA.Mat, TA.Mat>: 287 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))} 288 289 # LinearFirst: no dX (the first layer's input needs no gradient) 290 def linf_bwd(+n: Nat, +i: Nat, +o: Nat, p: T.Lin, x: TA.Mat, dy: TA.Mat, g: T.Lin) -> T.Bwd, TA.Mat>: 291 match p g: 292 case T.Lin{w, b} T.Lin{gw, gb}: 293 T.linf_b3(n, i, o, w, b, gb, TG.matmul_tn_acc(o, n, i, 0n, dy, x, gw)) 294 295 def LinearFirstF(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer, TA.Mat, TA.Mat, TA.Mat>: 296 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))} 297 298 # LinearNB: y = x·Wᵀ, no bias (Llama's projections) 299 def lnb_b2(+n: Nat, +i: Nat, +o: Nat, x: TA.Mat, g: TA.Mat, r: TA.MMul) -> T.Bwd, TA.Mat>: 300 match r: 301 case TA.MMul{dy, w, dx}: 302 T.lnb_b3(n, i, o, w, dx, TG.matmul_tn_acc(o, n, i, 0n, dy, x, g)) 303 304 def LinearNBF(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer, TA.Mat, TA.Mat, TA.Mat>: 305 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))} 306 307 # ---- RMSNorm: a leaf a row with main.bend's own row code (rmsf_row) ---- 308 309 type ERF is Data: 310 ERF{n: U32, d: Nat, d32: U32, df: F32, eps: F32} 311 312 def r3_y(s: T.R3) -> Array: 313 match s: 314 case T.R3{x, g, y}: y 315 316 def rmsf_if(ok: Bool, +t: U32, +d: Nat, +d32: U32, +df: F32, +eps: F32, x: Array, g: Array, y: Array) -> Array: 317 match ok: 318 case True{}: r3_y(T.rmsf_row(d, (t * d32 : U32), df, eps, T.R3{x, g, y})) 319 case False{}: y 320 321 # x1 = x, x2 = the gain g; y a row a leaf 322 def lf_rmsf(+t: U32, +l: U32, +e: ERF, x1: Array, x2: Array, x3: Array, x4: Array, y: Array) -> Array: 323 match e: 324 case ERF{+n, +d, +d32, +df, +eps}: rmsf_if(U32.is_lt(t, n), t, d, d32, df, eps, x1, x2, y) 325 326 def erf(+n: Nat, +d: Nat, +eps: F32) -> ERF: 327 ERF{U32.from_nat(n), d, U32.from_nat(d), F32.from_nat(d), eps} 328 329 def rms_f2(-n: Nat, -d: Nat, q: TG.R1) -> T.Fwd, TA.Mat, TA.Mat>: 330 match q: 331 case TG.R1{x1, x2, x3, x4, y}: T.Fwd{TA.Mat{x2}, TA.Mat{y}, TA.Mat{x1}} 332 333 def rms_f(+n: Nat, +d: Nat, +eps: F32, p: TA.Mat<1n, d>, x: TA.Mat) -> T.Fwd, TA.Mat, TA.Mat>: 334 match p x: 335 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)))) 336 337 # 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 338 # 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 339 # row's term into dg 340 341 def gdx_w(k: Nat, +j: U32, +off: U32, +r: F32, +m: F32, dx: Array, rx: Array & F32, rd: Array & F32, rg: Array & F32) -> Array: 342 match k rx rd rg: 343 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 344 +xh = (a * r : F32) 345 Array.set(F32, dx, (off + j : U32), (((b * c : F32) - (xh * m : F32) : F32) * r : F32)) 346 case 1n+p Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 347 +xh = (a * r : F32) 348 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))) 349 350 def gdx_3(+d: Nat, +off: U32, +df: F32, +r: F32, +t: U32, dx: Array, rr: Array, q: T.R3S) -> TG.O2: 351 match q: 352 case T.R3S{x, dy, g, +s}: 353 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)} 354 355 def gdx_2(+d: Nat, +off: U32, +df: F32, +eps: F32, +t: U32, dy: Array, g: Array, dx: Array, rr: Array, q: T.AS) -> TG.O2: 356 match q: 357 case T.AS{x, +s}: 358 +r = T.rms_inv(s, df, eps) 359 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))) 360 361 def gdx_if(ok: Bool, +t: U32, +d: Nat, +d32: U32, +df: F32, +eps: F32, x: Array, dy: Array, g: Array, dx: Array, rr: Array) -> TG.O2: 362 match ok: 363 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)))) 364 case False{}: TG.O2{dx, rr} 365 366 # x1 = x, x2 = dy, x3 = g; outputs dx and 1/rms of each row 367 def lf_gdx(+t: U32, +l: U32, +e: ERF, x1: Array, x2: Array, x3: Array, x4: Array, o: TG.O2) -> TG.O2: 368 match e o: 369 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) 370 371 # acc + dy[i]·(x[i]·r[row]) down a column, i from j by d 372 def gdg_loop(k: Nat, +i: U32, +d: U32, +row: U32, acc: F32, rx: Array & F32, rd: Array & F32, rr: Array & F32) -> F32: 373 match k rx rd rr: 374 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{r, +c}: (acc + (b * (a * c : F32) : F32) : F32) 375 case 1n+p Tuple{x, +a} Tuple{dy, +b} Tuple{r, +c}: 376 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))) 377 378 def gdg_put(+j: U32, +s: F32, r: Array & F32) -> Array: 379 match r: 380 case Tuple{ga, +v}: Array.set(F32, ga, j, (v + s : F32)) 381 382 type EDG is Data: 383 EDG{d: U32, n: Nat} 384 385 def gdg_if(ok: Bool, +j: U32, +n: Nat, +d: U32, x: Array, dy: Array, rr: Array, ga: Array) -> Array: 386 match ok: 387 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)) 388 case False{}: ga 389 390 # x1 = x, x2 = dy, x3 = 1/rms; the gradient of the gain, in place: a leaf a column 391 def lf_gdg(+t: U32, +l: U32, +e: EDG, x1: Array, x2: Array, x3: Array, x4: Array, ga: Array) -> Array: 392 match e: 393 case EDG{+d, +n}: gdg_if(U32.is_lt(t, d), t, n, d, x1, x2, x3, ga) 394 395 def rms_b3(-n: Nat, -d: Nat, g: Array, dx: Array, q: TG.R1) -> T.Bwd, TA.Mat>: 396 match q: 397 case TG.R1{x1, x2, x3, x4, ga}: T.Bwd{TA.Mat{g}, TA.Mat{dx}, TA.Mat{ga}} 398 399 def rms_b2(+n: Nat, +d: Nat, ga: Array, q: TG.R2) -> T.Bwd, TA.Mat>: 400 match q: 401 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)) 402 403 def rms_b(+n: Nat, +d: Nat, +eps: F32, p: TA.Mat<1n, d>, x: TA.Mat, dy: TA.Mat, gacc: TA.Mat<1n, d>) -> T.Bwd, TA.Mat>: 404 match p x dy gacc: 405 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))) 406 407 # RMSNorm over the last dimension, its gain trained (main.bend's RMSNorm) 408 def RMSNorm(+n: Nat, +d: Nat, +eps: F32) -> T.Layer, TA.Mat, TA.Mat, TA.Mat>: 409 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)} 410 411 # ---- RoPE: the tables on the host (main.bend's rope_tab, so the host's cos and sin), each pair rotated in place by 412 # main.bend's rope_one, the pairs dealt to the leaves round-robin ---- 413 414 type ERO is Data: 415 ERO{n: U32, fwd: Bool, hh: U32, half: U32, w: U32, t: U32} 416 417 def rope_loop(k: Nat, +e: U32, +l: U32, +fwd: Bool, +hh: U32, +half: U32, +w: U32, +t: U32, st: T.RS3) -> T.RS3: 418 match k: 419 case 0n: st 420 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)) 421 422 def rs3_x(st: T.RS3) -> Array: 423 match st: 424 case T.RS3{x, c, s}: x 425 426 # x1 = cos, x2 = sin; x rotated in place 427 def lf_rope(+t: U32, +l: U32, +e: ERO, x1: Array, x2: Array, x3: Array, x4: Array, x: Array) -> Array: 428 match e: 429 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})) 430 431 def rope_go2(+n: Nat, +w: Nat, +fwd: Bool, +t: Nat, +h: Nat, +half: Nat, x: Array, tb: T.CS) -> TA.Mat: 432 match tb: 433 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))} 434 435 def rope_go(+n: Nat, +w: Nat, +fwd: Bool, +t: Nat, +h: Nat, +half: Nat, +base: F32, m: TA.Mat) -> TA.Mat: 436 match m: 437 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))})) 438 439 # RoPE on n = b·t rows of h heads of 2·half (main.bend's RoPE) 440 def 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, TA.Mat, Unit>: 441 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} 442 443 type S4 is Type: 444 S4{a: Array, b: Array, c: Array, d: Array} 445 446 # ---- attention in cell bangs and row bangs. A leaf handles cells or rows t, t + l, .. of one bang. A cell is one number 447 # 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 448 # 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 449 # 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 450 # main.bend's row code (psmc, g_set2), a tail loop over the row's i + 1 positions. Every number is computed by 451 # main.bend's expression in main.bend's order, so the bits are main.bend's (but CUDA's exp in the softmax). 452 # forward: 0. cells: kᵀ 1. cells: P's scores (q·k scaled, j <= i) 2. rows: the softmax over the i + 1 453 # 3. cells: o = P·V over the query's i + 1 positions 454 # backward: 0. cells: kᵀ and vᵀ 1. cells: P's scores again and dP = dO·V into G 2. rows: P's softmax, then 455 # G = (P * (dP - s)) * sc with s = dP·P over the row, as gr3 3. cells: dQ = G·K, dK = Σ Gᵀ·Q and 456 # dV = Σ Pᵀ·dO over the group's heads, dK written over kᵀ and dV over V (neither read any more) 457 # 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 458 # 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 459 # sizes on the RTX 4050; this one takes 17 ms, about half of it filling the new arrays (one lane, Array.new). ---- 460 461 # 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 462 # of a head's P 463 type EAC is Data: 464 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} 465 466 # 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 467 # 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 468 # started at +0.0, which leaves it as it is; so the sums below skip them and give the same bits (for finite numbers: 469 # 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) 470 def toff(+bh: U32, +i: U32, +tr: U32) -> U32: 471 ((bh * tr : U32) + U32.div((i * (i + 1 : U32) : U32), 2) : U32) 472 473 def eac_n(+c: EAC) -> U32: 474 match c: 475 case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}: n 476 477 type S3 is Type: 478 S3{a: Array, b: Array, c: Array} 479 480 def s3_c(s: S3) -> Array: 481 match s: 482 case S3{a, b, c}: c 483 484 def sc_fin(y: Array, s: S3) -> TG.O2: 485 match s: 486 case S3{q, k, p}: TG.O2{p, y} 487 488 # 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 489 # i + 1 positions, · V's (or K's) column of the head's key/value head 490 def pv_val(+e: U32, +c: EAC, a: Array, b: Array) -> T.D2: 491 match c: 492 case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}: 493 +row = U32.div(e, w) 494 +col = U32.mod(e, w) 495 +bi = U32.div(row, t32) 496 +hi = U32.div(col, hd32) 497 +i = U32.mod(row, t32) 498 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) 499 500 # 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 501 # rows (sb apart), from row i on: dot0's sum, but with the step in a growing by one a row 502 def tdot(k: Nat, +ia: U32, +ib: U32, +i: U32, +sb: U32, acc: F32, ra: Array & F32, rb: Array & F32) -> T.D2: 503 match k ra rb: 504 case 0n Tuple{a, x} Tuple{b, y}: T.D2{a, b, acc} 505 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))) 506 507 # 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 508 # apart a head), rows j..t-1 509 def 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: 510 match k r: 511 case 0n T.D2{a, b, +s}: T.D2{a, b, (acc + s : F32)} 512 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)))) 513 514 # 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 515 # heads: the head's column j of a (rows j..t-1) · the head's column of b 516 def kv_val(+e: U32, +c: EAC, a: Array, b: Array) -> T.D2: 517 match c: 518 case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}: 519 +row = U32.div(e, kvw) 520 +col = U32.mod(e, kvw) 521 +bi = U32.div(row, t32) 522 +jj = U32.div(col, hd32) 523 +j = U32.mod(row, t32) 524 +ia = (toff(((bi * h : U32) + (jj * g : U32) : U32), j, tr) + j : U32) 525 +ib = ((((bi * t32 : U32) + j : U32) * w : U32) + (((jj * g : U32) * hd32 : U32) + U32.mod(col, hd32) : U32) : U32) 526 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))) 527 528 # K and V transposed for the score cells (kᵀ: b·kv·hd·t, by sequence, key/value head, column c, position j), so 529 # that neighbouring lanes, which take neighbouring keys j, read neighbouring cells. dot0 reads kᵀ with stride t: the same 530 # numbers in the same order as reading k's row, so the same bits. 531 def tk_src(+e: U32, +t32: U32, +kv: U32, +hd32: U32, +kvw: U32) -> U32: 532 +r = U32.div(e, t32) 533 +bj = U32.div(r, hd32) 534 +row = ((U32.div(bj, kv) * t32 : U32) + U32.mod(e, t32) : U32) 535 ((row * kvw : U32) + ((U32.mod(bj, kv) * hd32 : U32) + U32.mod(r, hd32) : U32) : U32) 536 537 def tk_put(+e: U32, o: Array, r: Array & F32) -> TG.O2: 538 match r: 539 case Tuple{a, +x}: TG.O2{a, Array.set(F32, o, e, x)} 540 541 def tk_loop(k: Nat, +e: U32, +l: U32, +t32: U32, +kv: U32, +hd32: U32, +kvw: U32, s: TG.O2) -> TG.O2: 542 match k s: 543 case 0n r: r 544 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)))) 545 546 def o2_b(s: TG.O2) -> Array: 547 match s: 548 case TG.O2{a, b}: b 549 550 # x1 = k; kᵀ written (made here) 551 def lf_tk(+t: U32, +l: U32, +c: EAC, x1: Array, x2: Array, x3: Array, x4: Array, o: Array) -> Array: 552 match c: 553 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})) 554 555 # both: cells [0, nk) are kᵀ's from k, the next nk vᵀ's from v (S4{k, v, kᵀ, vᵀ}) 556 def tkv_k(v: Array, vt: Array, r: TG.O2) -> S4: 557 match r: 558 case TG.O2{k, kt}: S4{k, v, kt, vt} 559 560 def tkv_v(k: Array, kt: Array, r: TG.O2) -> S4: 561 match r: 562 case TG.O2{v, vt}: S4{k, v, kt, vt} 563 564 def tkv_sel(lo: Bool, +e: U32, +nk: U32, +t32: U32, +kv: U32, +hd32: U32, +kvw: U32, s: S4) -> S4: 565 match lo s: 566 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)))) 567 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)))) 568 569 def tkv_loop(k: Nat, +e: U32, +l: U32, +nk: U32, +t32: U32, +kv: U32, +hd32: U32, +kvw: U32, s: S4) -> S4: 570 match k: 571 case 0n: s 572 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)) 573 574 def tkv_o(dq: Array, s: S4) -> TG.O3: 575 match s: 576 case S4{k, v, kt, vt}: TG.O3{kt, vt, dq} 577 578 # x1 = k, x2 = v; outputs kᵀ and vᵀ (written), dQ (made here, written by backward 3) 579 def lf_tkv(+t: U32, +l: U32, +c: EAC, x1: Array, x2: Array, x3: Array, x4: Array, o: TG.O3) -> TG.O3: 580 match c o: 581 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})) 582 583 # 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 584 # 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); 585 # in the middle row of an odd t the second half repeats the first and is skipped (ok false) 586 type IJ is Data: 587 IJ{ok: Bool, bh: U32, i: U32, j: U32} 588 589 def ij_of(+e: U32, +t32: U32) -> IJ: 590 +np = U32.div((t32 + 1 : U32), 2) 591 +c = U32.mod(e, (t32 + 1 : U32)) 592 +pr = U32.div(e, (t32 + 1 : U32)) 593 +p = U32.mod(pr, np) 594 +lo = U32.is_le(c, p) 595 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))} 596 597 # 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}) 598 def sc_put(+e: U32, +sc: F32, p: Array, r: T.D2) -> S3: 599 match r: 600 case T.D2{q, k, +s}: S3{q, k, Array.set(F32, p, e, (s * sc : F32))} 601 602 def sc_if(ok: Bool, +e: U32, +qa: U32, +ka: U32, +hd: Nat, +t32: U32, +sc: F32, s: S3) -> S3: 603 match ok s: 604 case True{} S3{q, k, p}: sc_put(e, sc, p, T.dot0(hd, qa, ka, 1, t32, q, k)) 605 case False{} S3{q, k, p}: S3{q, k, p} 606 607 def sc_ij(+a: IJ, +t32: U32, +h: U32, +g: U32, +hd: Nat, +hd32: U32, +w: U32, +sc: F32, +tr: U32, s: S3) -> S3: 608 match a: 609 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) 610 611 def sc_cell(+e: U32, +c: EAC, s: S3) -> S3: 612 match c: 613 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) 614 615 def sc_loop(k: Nat, +e: U32, +l: U32, +c: EAC, s: S3) -> S3: 616 match k: 617 case 0n: s 618 case 1n+q: sc_loop(q, (e + l : U32), l, c, sc_cell(e, c, s)) 619 620 # x1 = q, x2 = k; outputs P (written) and o (made here, written by forward 3) 621 def lf_sc(+t: U32, +l: U32, +c: EAC, x1: Array, x2: Array, x3: Array, x4: Array, o: TG.O2) -> TG.O2: 622 match o: 623 case TG.O2{p, y}: sc_fin(y, sc_loop(TG.share(t, eac_n(c), l), t, l, c, S3{x1, x2, p})) 624 625 # forward 2: P's row r (query i = r % t) by main.bend's softmax over its i + 1 positions 626 def sm_loop(k: Nat, +r: U32, +l: U32, +t32: U32, +tr: U32, p: Array) -> Array: 627 match k: 628 case 0n: p 629 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))) 630 631 def lf_sm(+t: U32, +l: U32, +c: EAC, x1: Array, x2: Array, x3: Array, x4: Array, p: Array) -> Array: 632 match c: 633 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) 634 635 # forward 3 (and backward 3's dQ): o[e] from pv_val (S3{a, b, o}) 636 def put_c(+e: U32, o: Array, r: T.D2) -> S3: 637 match r: 638 case T.D2{a, b, +s}: S3{a, b, Array.set(F32, o, e, s)} 639 640 def pv_loop(k: Nat, +e: U32, +l: U32, +c: EAC, s: S3) -> S3: 641 match k s: 642 case 0n r: r 643 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))) 644 645 # x1 = P, x2 = v; o written in place 646 def lf_pv(+t: U32, +l: U32, +c: EAC, x1: Array, x2: Array, x3: Array, x4: Array, o: Array) -> Array: 647 s3_c(pv_loop(TG.share(t, eac_n(c), l), t, l, c, S3{x1, x2, o})) 648 649 # 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}) 650 type S6 is Type: 651 S6{q: Array, k: Array, v: Array, d: Array, p: Array, g: Array} 652 653 def b1_g(+e: U32, q: Array, k: Array, p: Array, g: Array, r: T.D2) -> S6: 654 match r: 655 case T.D2{d, v, +s}: S6{q, k, v, d, p, Array.set(F32, g, e, s)} 656 657 def b1_p(+e: U32, +sc: F32, +qa: U32, +ka: U32, +hd: Nat, +t32: U32, v: Array, d: Array, p: Array, g: Array, r: T.D2) -> S6: 658 match r: 659 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)) 660 661 def b1_if(ok: Bool, +e: U32, +qa: U32, +ka: U32, +hd: Nat, +t32: U32, +sc: F32, s: S6) -> S6: 662 match ok s: 663 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)) 664 case False{} S6{q, k, v, d, p, g}: S6{q, k, v, d, p, g} 665 666 def b1_ij(+a: IJ, +t32: U32, +h: U32, +g: U32, +hd: Nat, +hd32: U32, +w: U32, +sc: F32, +tr: U32, s: S6) -> S6: 667 match a: 668 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) 669 670 def b1_cell(+e: U32, +c: EAC, s: S6) -> S6: 671 match c: 672 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) 673 674 def b1_loop(k: Nat, +e: U32, +l: U32, +c: EAC, s: S6) -> S6: 675 match k: 676 case 0n: s 677 case 1n+q: b1_loop(q, (e + l : U32), l, c, b1_cell(e, c, s)) 678 679 def b1_o(s: S6) -> TG.O2: 680 match s: 681 case S6{q, k, v, d, p, g}: TG.O2{p, g} 682 683 # x1 = q, x2 = kᵀ, x3 = vᵀ, x4 = dO; outputs P and G (made here, written) 684 def lf_b1(+t: U32, +l: U32, +c: EAC, x1: Array, x2: Array, x3: Array, x4: Array, o: TG.O2) -> TG.O2: 685 match o: 686 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})) 687 688 # 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, 689 # main.bend's gr3 and g_set2 690 def gr_2(+n: Nat, +off: U32, +sc: F32, g: Array, r: T.PC) -> T.A2: 691 match r: 692 case T.PC{p, c}: T.g_set2(n, sc, off, T.dot0(n, off, off, 1, 1, g, p)) 693 694 def gr_loop(k: Nat, +r: U32, +l: U32, +t32: U32, +tr: U32, +sc: F32, s: T.A2) -> T.A2: 695 match k s: 696 case 0n r: r 697 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))) 698 699 def a2_o2(s: T.A2) -> TG.O2: 700 match s: 701 case T.A2{p, g}: TG.O2{p, g} 702 703 def lf_gr(+t: U32, +l: U32, +c: EAC, x1: Array, x2: Array, x3: Array, x4: Array, o: TG.O2) -> TG.O2: 704 match c o: 705 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})) 706 707 # 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}) 708 type S8 is Type: 709 S8{q: Array, k: Array, d: Array, g: Array, p: Array, dq: Array, dk: Array, v: Array} 710 711 def b3_q(+e: U32, q: Array, d: Array, p: Array, dq: Array, dk: Array, v: Array, r: T.D2) -> S8: 712 match r: 713 case T.D2{g, k, +s}: S8{q, k, d, g, p, Array.set(F32, dq, e, s), dk, v} 714 715 def b3_k(+e: U32, k: Array, d: Array, p: Array, dq: Array, dk: Array, v: Array, r: T.D2) -> S8: 716 match r: 717 case T.D2{g, q, +s}: S8{q, k, d, g, p, dq, Array.set(F32, dk, e, s), v} 718 719 def b3_v(+e: U32, q: Array, k: Array, g: Array, dq: Array, dk: Array, v: Array, r: T.D2) -> S8: 720 match r: 721 case T.D2{p, d, +s}: S8{q, k, d, g, p, dq, dk, Array.set(F32, v, e, s)} 722 723 def b3_sel(a: Bool, b: Bool, +e: U32, +nw: U32, +nk: U32, +c: EAC, s: S8) -> S8: 724 match a b s: 725 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)) 726 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)) 727 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)) 728 729 def b3_cell(+e: U32, +c: EAC, s: S8) -> S8: 730 match c: 731 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) 732 733 def b3_loop(k: Nat, +e: U32, +l: U32, +c: EAC, s: S8) -> S8: 734 match k: 735 case 0n: s 736 case 1n+q: b3_loop(q, (e + l : U32), l, c, b3_cell(e, c, s)) 737 738 def b3_o(s: S8) -> TG.O4: 739 match s: 740 case S8{q, k, d, g, p, dq, dk, v}: TG.O4{p, dq, dk, v} 741 742 # x1 = q, x2 = k, x3 = dO, x4 = G; P read from the first output (unchanged), dQ, dK and dV (over V) written 743 def lf_b3(+t: U32, +l: U32, +c: EAC, x1: Array, x2: Array, x3: Array, x4: Array, o: TG.O4) -> TG.O4: 744 match o: 745 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})) 746 747 # the sizes of a bang of cnt cells (or rows) 748 def eac(+cnt: Nat, +n: Nat, +w: Nat, +kvw: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32) -> EAC: 749 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)} 750 751 # 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) 752 def pcells(+b: Nat, +t: Nat, +kv: Nat, +g: Nat) -> Nat: 753 Nat.mul(Nat.mul(b, Nat.mul(kv, g)), Nat.div(Nat.mul(t, 1n+t), 2n)) 754 755 def tcells(+b: Nat, +t: Nat, +kv: Nat, +g: Nat) -> Nat: 756 Nat.mul(Nat.mul(Nat.mul(b, Nat.mul(kv, g)), Nat.div(1n+t, 2n)), 1n+t) 757 758 def prows(+b: Nat, +t: Nat, +kv: Nat, +g: Nat) -> Nat: 759 Nat.mul(Nat.mul(b, Nat.mul(kv, g)), t) 760 761 def af3(-n: Nat, -w: Nat, -kvw: Nat, q: Array, k: Array, r: TG.R1) -> T.Fwd, T.AttnS>: 762 match r: 763 case TG.R1{p, v, x3, x4, o}: T.Fwd{Unit{}, TA.Mat{o}, T.AttnS{q, k, v}} 764 765 def af2(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, q: Array, k: Array, v: Array, o: Array, r: TG.R1) -> T.Fwd, T.AttnS>: 766 match r: 767 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)) 768 769 def af1(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, k: Array, v: Array, r: TG.R2) -> T.Fwd, T.AttnS>: 770 match r: 771 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)) 772 773 def af0(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, q: Array, v: Array, r: TG.R1) -> T.Fwd, T.AttnS>: 774 match r: 775 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)))) 776 777 def attn_f(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, x: T.Two, T.Two, TA.Mat>>) -> T.Fwd, T.AttnS>: 778 match x: 779 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)))) 780 781 def ab4(-n: Nat, -w: Nat, -kvw: Nat, r: TG.R4) -> T.Bwd, T.Two, TA.Mat>>>: 782 match r: 783 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{}} 784 785 # backward 3 writes dK over kᵀ and dV over v (neither read any more) 786 def ab3(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, q: Array, k: Array, v: Array, d: Array, dq: Array, kt: Array, r: TG.R2) -> T.Bwd, T.Two, TA.Mat>>>: 787 match r: 788 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)) 789 790 def ab2(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, k: Array, v: Array, dq: Array, r: TG.R2) -> T.Bwd, T.Two, TA.Mat>>>: 791 match r: 792 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)) 793 794 def ab1(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, q: Array, d: Array, r: TG.R3) -> T.Bwd, T.Two, TA.Mat>>>: 795 match r: 796 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)))) 797 798 def 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) -> T.Bwd, T.Two, TA.Mat>>>: 799 match s dy: 800 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)))) 801 802 # causal attention with grouped heads over b sequences of t positions (main.bend's Attention) 803 def 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, T.Two, TA.Mat>>, TA.Mat, T.AttnS>: 804 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} 805 806 # ---- MSE on the device. main.bend's MSE writes each gradient cell 2 (y - t) / (n·c) over y and sums the squares with 807 # tree_sum, whose passes add neighbouring pairs; its first p passes carry no odd cell when 2^p divides the count, so a 808 # leaf that sums a block of 2^p squares as a balanced tree of pairs gets exactly the number those passes leave at the 809 # block's place. The leaves do that on the device (y written in place), and the host finishes tree_sum over the 810 # count / 2^p block sums: the same bits as main.bend's MSE, with only the block sums read back. ---- 811 812 type EMS is Data: 813 EMS{nb: U32, p: Nat, bs: U32, kk: F32} 814 815 type MS is Type: 816 MS{y: Array, t: Array, s: F32} 817 818 def ms_add(+s: F32, r: MS) -> MS: 819 match r: 820 case MS{y, t, +s2}: MS{y, t, (s + s2 : F32)} 821 822 def ms_then(r: MS, k: @+s: F32 -> Array -> Array -> MS) -> MS: 823 match r: 824 case MS{y, t, +s}: k(s, y, t) 825 826 def ms_cell(+i: U32, +kk: F32, ry: Array & F32, rt: Array & F32) -> MS: 827 match ry rt: 828 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)} 829 830 # the squares of the block of 2^p cells from i summed as a balanced tree of pairs (h = 2^(p-1)), gradients written 831 def ms_blk(p: Nat, +i: U32, +h: U32, +kk: F32, y: Array, t: Array) -> MS: 832 match p: 833 case 0n: ms_cell(i, kk, Array.get(F32, y, i), Array.get(F32, t, i)) 834 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))) 835 836 def ms_put(+t: U32, ps: Array, r: MS) -> TG.O2: 837 match r: 838 case MS{y, tt, +s}: TG.O2{y, Array.set(F32, ps, t, s)} 839 840 def ms_if(ok: Bool, +t: U32, +p: Nat, +bs: U32, +kk: F32, tg: Array, y: Array, ps: Array) -> TG.O2: 841 match ok: 842 case True{}: ms_put(t, ps, ms_blk(p, (t * bs : U32), U32.shr(bs), kk, y, tg)) 843 case False{}: TG.O2{y, ps} 844 845 # x1 = the target; y written in place with the gradient, ps[t] = leaf t's block sum (one leaf a block) 846 def lf_mse(+t: U32, +l: U32, +e: EMS, x1: Array, x2: Array, x3: Array, x4: Array, o: TG.O2) -> TG.O2: 847 match e o: 848 case EMS{+nb, +p, +bs, +kk} TG.O2{y, ps}: ms_if(U32.is_lt(t, nb), t, p, bs, kk, x1, y, ps) 849 850 # the largest p <= k with 2^p dividing m (m > 0) 851 def ms_p(k: Nat, +m: U32) -> Nat: 852 match k: 853 case 0n: 0n 854 case 1n+q: Bool.pick(Nat, U32.is_eq(U32.mod(m, 2), 0), 1n+ms_p(q, U32.div(m, 2)), 0n) 855 856 def mse_fin(-n: Nat, -c: Nat, +nb: U32, +nc: F32, q: TG.R2) -> T.LossOut>: 857 match q: 858 case TG.R2{x1, x2, x3, x4, TG.O2{y, ps}}: T.LossOut{(T.tree_sum(nb, ps) / nc : F32), TA.Mat{y}} 859 860 def mse_go(+n: Nat, +c: Nat, +m: U32, +p: Nat, y: Array, t: Array) -> T.LossOut>: 861 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))))) 862 863 # mean squared error over every number (main.bend's MSE, the same bits) 864 def MSE(+n: Nat, +c: Nat, y: TA.Mat, t: TA.Mat) -> T.LossOut>: 865 match y t: 866 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) 867 868 # ---- 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 869 # 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 ---- 870 871 def u2f(x: U32) -> F32: 872 match x: 873 case U32{w}: F32{w} 874 875 def f2u(x: F32) -> U32: 876 match x: 877 case F32{w}: U32{w} 878 879 type IA is Type: 880 IA{a: Array, k: Nat} 881 882 def ia_go(left: Nat, ids: List<&2, Nat>, +i: U32, k: Nat, a: Array) -> IA: 883 match left ids: 884 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)))) 885 case _ _: IA{a, k} 886 887 def ids_arr(+n: Nat, ids: List<&2, Nat>) -> IA: 888 ia_go(n, ids, 0, 0n, T.zeros_a(n)) 889 890 # ---- Embedding: a leaf a column j of the table (d of them), each looping over the ids in order: the forward pass 891 # 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 892 # after the other, so every cell of the gradient gets its rows added in main.bend's order, with no two leaves writing 893 # one cell. Neighbouring lanes read neighbouring cells. ---- 894 895 type EEM is Data: 896 EEM{k: Nat, d: U32} 897 898 type ES3 is Type: 899 ES3{ids: Array, a: Array, o: Array} 900 901 def ef_3(+i: U32, +j: U32, +d: U32, ids: Array, o: Array, r: Array & F32) -> ES3: 902 match r: 903 case Tuple{e, +x}: ES3{ids, e, Array.set(F32, o, ((i * d : U32) + j : U32), x)} 904 905 def ef_2(+i: U32, +j: U32, +d: U32, e: Array, o: Array, r: Array & F32) -> ES3: 906 match r: 907 case Tuple{ids, +v}: ef_3(i, j, d, ids, o, Array.get(F32, e, ((f2u(v) * d : U32) + j : U32))) 908 909 # o[i, j] = e[ids[i], j] 910 def ef_cell(+i: U32, +j: U32, +d: U32, st: ES3) -> ES3: 911 match st: 912 case ES3{ids, e, o}: ef_2(i, j, d, e, o, Array.get(F32, ids, i)) 913 914 def ef_loop(k: Nat, +i: U32, +j: U32, +d: U32, st: ES3) -> ES3: 915 match k: 916 case 0n: st 917 case 1n+q: ef_loop(q, (i + 1 : U32), j, d, ef_cell(i, j, d, st)) 918 919 def es3_o(st: ES3) -> Array: 920 match st: 921 case ES3{ids, a, o}: o 922 923 def ef_if(ok: Bool, +t: U32, +k: Nat, +d: U32, e: Array, ids: Array, o: Array) -> Array: 924 match ok: 925 case True{}: es3_o(ef_loop(k, 0, t, d, ES3{ids, e, o})) 926 case False{}: o 927 928 # x1 = the table, x2 = the ids; o (made on the device, zeros) gets the rows 929 def lf_embf(+t: U32, +l: U32, +e: EEM, x1: Array, x2: Array, x3: Array, x4: Array, o: Array) -> Array: 930 match e: 931 case EEM{+k, +d}: ef_if(U32.is_lt(t, d), t, k, d, x1, x2, o) 932 933 def eb_4(+gi: U32, ids: Array, g: Array, +a: F32, r: Array & F32) -> ES3: 934 match r: 935 case Tuple{dy, +b}: ES3{ids, dy, Array.set(F32, g, gi, (a + b : F32))} 936 937 def eb_3(+i: U32, +j: U32, +d: U32, +gi: U32, ids: Array, dy: Array, r: Array & F32) -> ES3: 938 match r: 939 case Tuple{g, +a}: eb_4(gi, ids, g, a, Array.get(F32, dy, ((i * d : U32) + j : U32))) 940 941 def eb_2(+i: U32, +j: U32, +d: U32, dy: Array, g: Array, r: Array & F32) -> ES3: 942 match r: 943 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))) 944 945 # g[ids[i], j] = g[ids[i], j] + dy[i, j] 946 def eb_cell(+i: U32, +j: U32, +d: U32, st: ES3) -> ES3: 947 match st: 948 case ES3{ids, dy, g}: eb_2(i, j, d, dy, g, Array.get(F32, ids, i)) 949 950 def eb_loop(k: Nat, +i: U32, +j: U32, +d: U32, st: ES3) -> ES3: 951 match k: 952 case 0n: st 953 case 1n+q: eb_loop(q, (i + 1 : U32), j, d, eb_cell(i, j, d, st)) 954 955 def eb_if(ok: Bool, +t: U32, +k: Nat, +d: U32, dy: Array, ids: Array, g: Array) -> Array: 956 match ok: 957 case True{}: es3_o(eb_loop(k, 0, t, d, ES3{ids, dy, g})) 958 case False{}: g 959 960 # x1 = dy, x2 = the ids; the gradient g, in place 961 def lf_embb(+t: U32, +l: U32, +e: EEM, x1: Array, x2: Array, x3: Array, x4: Array, g: Array) -> Array: 962 match e: 963 case EEM{+k, +d}: eb_if(U32.is_lt(t, d), t, k, d, x1, x2, g) 964 965 def emb_f3(-n: Nat, -v: Nat, -d: Nat, +ids: List<&2, Nat>, q: TG.R1) -> T.Fwd, TA.Mat, List<&2, Nat>>: 966 match q: 967 case TG.R1{x1, x2, x3, x4, o}: T.Fwd{TA.Mat{x1}, TA.Mat{o}, ids} 968 969 def emb_f2(+n: Nat, +v: Nat, +d: Nat, +ids: List<&2, Nat>, ea: Array, r: IA) -> T.Fwd, TA.Mat, List<&2, Nat>>: 970 match r: 971 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)))) 972 973 def emb_f(+n: Nat, +v: Nat, +d: Nat, +ids: List<&2, Nat>, e: TA.Mat) -> T.Fwd, TA.Mat, List<&2, Nat>>: 974 match e: 975 case TA.Mat{ea}: emb_f2(n, v, d, ids, ea, ids_arr(n, ids)) 976 977 def emb_b3(-v: Nat, -d: Nat, e: TA.Mat, q: TG.R1) -> T.Bwd, List<&2, Nat>>: 978 match q: 979 case TG.R1{x1, x2, x3, x4, g}: T.Bwd{e, Nil{}, TA.Mat{g}} 980 981 def emb_b2(+v: Nat, +d: Nat, e: TA.Mat, da: Array, ga: Array, r: IA) -> T.Bwd, List<&2, Nat>>: 982 match r: 983 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)) 984 985 def emb_bwd(+n: Nat, +v: Nat, +d: Nat, e: TA.Mat, ids: List<&2, Nat>, dy: TA.Mat, g: TA.Mat) -> T.Bwd, List<&2, Nat>>: 986 match dy g: 987 case TA.Mat{da} TA.Mat{ga}: emb_b2(v, d, e, da, ga, ids_arr(n, ids)) 988 989 # 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 990 def Embedding(+n: Nat, +v: Nat, +d: Nat) -> T.Layer, List<&2, Nat>, TA.Mat, List<&2, Nat>>: 991 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))} 992 993 # ---- cross-entropy on the device, the same numbers as main.bend's CrossEntropy (but CUDA's exp on the GPU), in four 994 # bangs; only 3 numbers a row come back to the host, and the gradient stays on the device: 995 # 1. a leaf a row: its max m by main.bend's pmax0 (the same order of comparisons) and x[y] - m; 996 # 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 997 # 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 998 # these pairs (no odd cell is carried while 2^p divides the count), so a block's sum is the number they leave there; 999 # 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; 1000 # 4. a leaf for 2^14 cells at a time: ((e / s) - [j == y]) / n, pdiv's and ce_g's two expressions. 1001 # 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 1002 # they are, as in main.bend; labels past the n rows are not read (main.bend would write past the matrix). ---- 1003 1004 type TS is Type: 1005 TS{a: Array, s: F32} 1006 1007 def ts_leaf(r: Array & F32) -> TS: 1008 match r: 1009 case Tuple{a, +x}: TS{a, x} 1010 1011 def ts_plus(+s: F32, r: TS) -> TS: 1012 match r: 1013 case TS{a, +s2}: TS{a, (s + s2 : F32)} 1014 1015 def ts_then(r: TS, k: @+s: F32 -> Array -> TS) -> TS: 1016 match r: 1017 case TS{a, +s}: k(s, a) 1018 1019 # 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 = 1020 # 2^(l-1) cells) and its right child (the rest) when there is a rest (big = h < len), and is its left child otherwise 1021 # (tree_sum's odd last cell carried up a pass) 1022 def tnode(l: Nat, big: Bool, +lo: U32, +len: U32, +h: U32, a: Array) -> TS: 1023 match l big: 1024 case 0n _: ts_leaf(Array.get(F32, a, lo)) 1025 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))) 1026 case 1n+q False{}: tnode(q, U32.is_lt(U32.shr(h), len), lo, len, U32.shr(h), a) 1027 1028 # the passes tree_sum makes over m cells: the smallest l with 2^l >= m (f is fuel) 1029 def ts_lv(f: Nat, +m: U32, +p: U32) -> Nat: 1030 match f: 1031 case 0n: 0n 1032 case 1n+q: Bool.pick(Nat, U32.is_lt(p, m), 1n+ts_lv(q, m, (p * 2 : U32)), 0n) 1033 1034 # 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 1035 # last class 1036 type LS is Type: 1037 LS{a: Array, l: F32} 1038 1039 def ce_t2(ok: Bool, +loss: F32, +xm: F32, r: Array & F32) -> LS: 1040 match ok r: 1041 case True{} Tuple{a, +s}: LS{a, (loss - (xm - F32.log(s) : F32) : F32)} 1042 case False{} Tuple{a, s}: LS{a, (loss - F32.log(F32.max(0.0, 0.000000000001)) : F32)} 1043 1044 def ce_g4(+n: Nat, -c: Nat, z: Array, r: LS) -> T.LossOut>: 1045 match r: 1046 case LS{a, +l}: T.LossOut{(l / F32.from_nat(n) : F32), TA.Mat{z}} 1047 1048 # 1. a leaf a row: st[3t] = x[y] - m, st[3t + 1] = m (x[y] read only when y is a class) 1049 type EC1 is Data: 1050 EC1{k: U32, c1: Nat, c32: U32} 1051 1052 def c1_3(+t: U32, +m: F32, st: Array, r: Array & F32) -> Array: 1053 match r: 1054 case Tuple{z, +xy}: Array.set(F32, Array.set(F32, st, (t * 3 : U32), (xy - m : F32)), ((t * 3 : U32) + 1 : U32), m) 1055 1056 def c1_2(ok: Bool, +t: U32, +off: U32, +y: U32, st: Array, r: T.AS) -> Array: 1057 match ok r: 1058 case True{} T.AS{z, +m}: c1_3(t, m, st, Array.get(F32, z, (off + y : U32))) 1059 case False{} T.AS{z, +m}: c1_3(t, m, st, Tuple{z, 0.0}) 1060 1061 def c1_1(+t: U32, +off: U32, +c1: Nat, +c32: U32, z: Array, st: Array, r: Array & F32) -> Array: 1062 match r: 1063 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))) 1064 1065 def c1_if(ok: Bool, +t: U32, +c1: Nat, +c32: U32, lab: Array, z: Array, st: Array) -> Array: 1066 match ok: 1067 case True{}: c1_1(t, (t * c32 : U32), c1, c32, z, st, Array.get(F32, lab, t)) 1068 case False{}: st 1069 1070 def lf_c1(+t: U32, +l: U32, +e: EC1, x1: Array, x2: Array, x3: Array, x4: Array, st: Array) -> Array: 1071 match e: 1072 case EC1{+k, +c1, +c32}: c1_if(U32.is_lt(t, k), t, c1, c32, x1, x2, st) 1073 1074 # 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)) 1075 def tsx_leaf(+lo: U32, +m: F32, r: Array & F32) -> TS: 1076 match r: 1077 case Tuple{a, +x}: 1078 +e = F32.exp((x - m : F32)) 1079 TS{Array.set(F32, a, lo, e), e} 1080 1081 def tsx(l: Nat, +lo: U32, +h: U32, +m: F32, a: Array) -> TS: 1082 match l: 1083 case 0n: tsx_leaf(lo, m, Array.get(F32, a, lo)) 1084 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))) 1085 1086 # the same over 8 cells in straight-line code (no closure a node): e0..e7 written, then ((e0 + e1) + (e2 + e3)) + 1087 # ((e4 + e5) + (e6 + e7)), the balanced tree's three levels 1088 def x8_7(+lo: U32, +m: F32, +e0: F32, +e1: F32, +e2: F32, +e3: F32, +e4: F32, +e5: F32, +e6: F32, r: Array & F32) -> TS: 1089 match r: 1090 case Tuple{a, +x}: 1091 +e7 = F32.exp((x - m : F32)) 1092 TS{Array.set(F32, a, (lo + 7 : U32), e7), ((((e0 + e1 : F32) + (e2 + e3 : F32) : F32) + ((e4 + e5 : F32) + (e6 + e7 : F32) : F32)) : F32)} 1093 1094 def x8_6(+lo: U32, +m: F32, +e0: F32, +e1: F32, +e2: F32, +e3: F32, +e4: F32, +e5: F32, r: Array & F32) -> TS: 1095 match r: 1096 case Tuple{a, +x}: 1097 +e = F32.exp((x - m : F32)) 1098 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))) 1099 1100 def x8_5(+lo: U32, +m: F32, +e0: F32, +e1: F32, +e2: F32, +e3: F32, +e4: F32, r: Array & F32) -> TS: 1101 match r: 1102 case Tuple{a, +x}: 1103 +e = F32.exp((x - m : F32)) 1104 x8_6(lo, m, e0, e1, e2, e3, e4, e, Array.get(F32, Array.set(F32, a, (lo + 5 : U32), e), (lo + 6 : U32))) 1105 1106 def x8_4(+lo: U32, +m: F32, +e0: F32, +e1: F32, +e2: F32, +e3: F32, r: Array & F32) -> TS: 1107 match r: 1108 case Tuple{a, +x}: 1109 +e = F32.exp((x - m : F32)) 1110 x8_5(lo, m, e0, e1, e2, e3, e, Array.get(F32, Array.set(F32, a, (lo + 4 : U32), e), (lo + 5 : U32))) 1111 1112 def x8_3(+lo: U32, +m: F32, +e0: F32, +e1: F32, +e2: F32, r: Array & F32) -> TS: 1113 match r: 1114 case Tuple{a, +x}: 1115 +e = F32.exp((x - m : F32)) 1116 x8_4(lo, m, e0, e1, e2, e, Array.get(F32, Array.set(F32, a, (lo + 3 : U32), e), (lo + 4 : U32))) 1117 1118 def x8_2(+lo: U32, +m: F32, +e0: F32, +e1: F32, r: Array & F32) -> TS: 1119 match r: 1120 case Tuple{a, +x}: 1121 +e = F32.exp((x - m : F32)) 1122 x8_3(lo, m, e0, e1, e, Array.get(F32, Array.set(F32, a, (lo + 2 : U32), e), (lo + 3 : U32))) 1123 1124 def x8_1(+lo: U32, +m: F32, +e0: F32, r: Array & F32) -> TS: 1125 match r: 1126 case Tuple{a, +x}: 1127 +e = F32.exp((x - m : F32)) 1128 x8_2(lo, m, e0, e, Array.get(F32, Array.set(F32, a, (lo + 1 : U32), e), (lo + 2 : U32))) 1129 1130 def x8_0(+lo: U32, +m: F32, r: Array & F32) -> TS: 1131 match r: 1132 case Tuple{a, +x}: 1133 +e = F32.exp((x - m : F32)) 1134 x8_1(lo, m, e, Array.get(F32, Array.set(F32, a, lo, e), (lo + 1 : U32))) 1135 1136 # tsx over 2^(l+3) cells, its last three levels in blocks of 8 1137 def tsx8(l: Nat, +lo: U32, +h: U32, +m: F32, a: Array) -> TS: 1138 match l: 1139 case 0n: x8_0(lo, m, Array.get(F32, a, lo)) 1140 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))) 1141 1142 # 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 1143 type EC2 is Data: 1144 EC2{nb: U32, per: U32, p: Nat, bsz: U32, c32: U32} 1145 1146 type C3S is Type: 1147 C3S{a: Array, b: Array, c: Array} 1148 1149 def c2_put(+b: U32, st: Array, bs: Array, r: TS) -> C3S: 1150 match r: 1151 case TS{z, +s}: C3S{st, z, Array.set(F32, bs, b, s)} 1152 1153 def c2_m(b8: Bool, +b: U32, +lo: U32, +p: Nat, +h: U32, z: Array, bs: Array, r: Array & F32) -> C3S: 1154 match b8 r: 1155 case True{} Tuple{st, +m}: c2_put(b, st, bs, tsx8(Nat.sub(p, 3n), lo, h, m, z)) 1156 case False{} Tuple{st, +m}: c2_put(b, st, bs, tsx(p, lo, h, m, z)) 1157 1158 def c2_cell(+b: U32, +per: U32, +p: Nat, +bsz: U32, +c32: U32, s: C3S) -> C3S: 1159 match s: 1160 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))) 1161 1162 def c2_loop(k: Nat, +b: U32, +l: U32, +per: U32, +p: Nat, +bsz: U32, +c32: U32, s: C3S) -> C3S: 1163 match k: 1164 case 0n: s 1165 case 1n+q: c2_loop(q, (b + l : U32), l, per, p, bsz, c32, c2_cell(b, per, p, bsz, c32, s)) 1166 1167 def c2_o(s: C3S) -> TG.O2: 1168 match s: 1169 case C3S{st, z, bs}: TG.O2{z, bs} 1170 1171 def lf_c2(+t: U32, +l: U32, +e: EC2, x1: Array, x2: Array, x3: Array, x4: Array, o: TG.O2) -> TG.O2: 1172 match e o: 1173 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})) 1174 1175 # 3. a leaf a row: st[3t + 2] = s, tree_sum's remaining passes over the row's per block sums 1176 type EC3 is Data: 1177 EC3{k: U32, per: U32, tl: Nat, th: U32} 1178 1179 def c3_put(+t: U32, st: Array, r: TS) -> Array: 1180 match r: 1181 case TS{bs, +s}: Array.set(F32, st, ((t * 3 : U32) + 2 : U32), s) 1182 1183 def c3_if(ok: Bool, +t: U32, +per: U32, +tl: Nat, +th: U32, bs: Array, st: Array) -> Array: 1184 match ok: 1185 case True{}: c3_put(t, st, tnode(tl, U32.is_lt(th, per), (t * per : U32), per, th, bs)) 1186 case False{}: st 1187 1188 def lf_c3(+t: U32, +l: U32, +e: EC3, x1: Array, x2: Array, x3: Array, x4: Array, st: Array) -> Array: 1189 match e: 1190 case EC3{+k, +per, +tl, +th}: c3_if(U32.is_lt(t, k), t, per, tl, th, x1, st) 1191 1192 # 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 1193 type EC4 is Data: 1194 EC4{n: U32, c32: U32, nf: F32} 1195 1196 def c4_4(+i: U32, +j: U32, +y: U32, +s: F32, +nf: F32, lab: Array, st: Array, r: Array & F32) -> C3S: 1197 match r: 1198 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))} 1199 1200 def c4_3(+i: U32, +j: U32, +y: U32, +nf: F32, lab: Array, z: Array, r: Array & F32) -> C3S: 1201 match r: 1202 case Tuple{st, +s}: c4_4(i, j, y, s, nf, lab, st, Array.get(F32, z, i)) 1203 1204 def c4_2(+i: U32, +row: U32, +j: U32, +nf: F32, st: Array, z: Array, r: Array & F32) -> C3S: 1205 match r: 1206 case Tuple{lab, +yv}: c4_3(i, j, f2u(yv), nf, lab, z, Array.get(F32, st, ((row * 3 : U32) + 2 : U32))) 1207 1208 def c4_cell(+i: U32, +c32: U32, +nf: F32, s: C3S) -> C3S: 1209 match s: 1210 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))) 1211 1212 def c4_loop(k: Nat, +i: U32, +l: U32, +c32: U32, +nf: F32, s: C3S) -> C3S: 1213 match k: 1214 case 0n: s 1215 case 1n+q: c4_loop(q, (i + l : U32), l, c32, nf, c4_cell(i, c32, nf, s)) 1216 1217 def c3s_c(s: C3S) -> Array: 1218 match s: 1219 case C3S{a, b, c}: c 1220 1221 def lf_c4(+t: U32, +l: U32, +e: EC4, x1: Array, x2: Array, x3: Array, x4: Array, z: Array) -> Array: 1222 match e: 1223 case EC4{+n, +c32, +nf}: c3s_c(c4_loop(TG.share(t, n, l), t, l, c32, nf, C3S{x1, x2, z})) 1224 1225 # the host's part over st, row by row 1226 def cx_t1(ok: Bool, +i: U32, +loss: F32, r: Array & F32) -> LS: 1227 match r: 1228 case Tuple{a, +xm}: ce_t2(ok, loss, xm, Array.get(F32, a, ((i * 3 : U32) + 2 : U32))) 1229 1230 def cx_ls(left: Nat, ls: List<&2, Nat>, +i: U32, +c: Nat, r: LS) -> LS: 1231 match left ls r: 1232 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)))) 1233 case _ _ r: r 1234 1235 def cx_5(+n: Nat, +c: Nat, +labels: List<&2, Nat>, +k: Nat, q: TG.R1) -> T.LossOut>: 1236 match q: 1237 case TG.R1{x1, x2, x3, x4, z}: ce_g4(n, c, z, cx_ls(k, labels, 0, c, LS{x2, 0.0})) 1238 1239 def cx_4(+n: Nat, +c: Nat, +labels: List<&2, Nat>, +k: Nat, la: Array, z: Array, q: TG.R1) -> T.LossOut>: 1240 match q: 1241 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)) 1242 1243 def cx_3(+n: Nat, +c: Nat, +labels: List<&2, Nat>, +k: Nat, +per: U32, la: Array, q: TG.R2) -> T.LossOut>: 1244 match q: 1245 case TG.R2{st, x2, x3, x4, TG.O2{z, bs}}: cx_4(n, c, labels, k, la, z, TG.run1(~EC3, ~lf_c3, TA.cap_depth(k), EC3{U32.from_nat(k), per, ts_lv(32n, per, 1), U32.shr(U32.pow(2, ts_lv(32n, per, 1)))}, bs, nn(Unit{}), nn(Unit{}), nn(Unit{}), st)) 1246 1247 def cx_2(+n: Nat, +c: Nat, +labels: List<&2, Nat>, +k: Nat, +p: Nat, q: TG.R1) -> T.LossOut>: 1248 match q: 1249 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)))))) 1250 1251 def cx_1(+n: Nat, +c: Nat, +labels: List<&2, Nat>, z: Array, r: IA) -> T.LossOut>: 1252 match r: 1253 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)))) 1254 1255 # cross-entropy over n rows of c classes against the labels (main.bend's CrossEntropy; its bits but CUDA's exp) 1256 def CrossEntropy(+n: Nat, +c: Nat, z: TA.Mat, +labels: List<&2, Nat>) -> T.LossOut>: 1257 match z: 1258 case TA.Mat{d}: cx_1(n, c, labels, d, ids_arr(n, labels))