diff --git a/bend2/base.bend b/bend2/base.bend index d7d118e546..18e7a24e8f 100644 --- a/bend2/base.bend +++ b/bend2/base.bend @@ -2480,6 +2480,889 @@ def Array.get.at(-T: Data, i: U32, an: Array & U32) -> Array & T: def Array.get(-T: Data, a: Array, +i: U32) -> Array & T: Array.get.at(T, i, Array.size(T, a)) +# Array.q4mv(w, x, y, r, n, cols, rows) is the MLX affine 4-bit product +# (bits 4, group 64) of rows r .. r+n of w, a matrix of `rows` rows of `cols` +# weights (cols a multiple of 128): the codes (rows * cols/8 words, weight +# 8k+m of a row in bits 4m..4m+3 of its word k), then the bfloat16 scales +# (rows * cols/128 words, group g in the low half when g is even), then the +# biases laid out as the scales. x holds the cols activations, then at +# cols + g the sum of group g. Each row is summed in a fixed tree that +# SIMD lanes can run: lane j of a group sums nibble j of its 8 words times +# x in word order; the lanes meet as ((0+4)+(2+6))+((1+5)+(3+7)) into qx; +# the group's term scale * qx + bias * x-sum goes onto partial g % 32 (from +# 0.0, groups in order); the 32 partials fold as p[l] += p[l + o] for +# o = 16, 8, 4, 2, 1. p[0] is rounded to bfloat16 (nearest even, a NaN to +# sign | 0x7fff) and written to y at the row's index. The compiler lowers +# it to NEON and Metal SIMD-group loops; the defs below are its meaning. + +def Array.q4mv.f32(+u: U32) -> F32: + match u: + case U32{+w}: + F32{w} + +def Array.q4mv.bits(+x: F32) -> U32: + match x: + case F32{+w}: + U32{w} + +def Array.q4mv.half(even: Bool, +u: U32) -> F32: + match even: + case True{}: + Array.q4mv.f32(((u .&. 65535 : U32) << 16n : U32)) + case False{}: + Array.q4mv.f32((u .&. 4294901760 : U32)) + +def Array.q4mv.narrow(nan: Bool, +u: U32) -> U32: + match nan: + case True{}: + (((u >> 16n) .&. 32768) .|. 32767 : U32) + case False{}: + ((u + 32767 + ((u >> 16n) .&. 1)) >> 16n : U32) + +def Array.q4mv.round(+x: F32) -> F32: + +u = Array.q4mv.bits(x) + Array.q4mv.f32((Array.q4mv.narrow(U32.is_gt((u .&. 2147483647 : U32), 2139095040), u) << 16n : U32)) + +def Array.q4mv.mac(g: Array & F32, w: Array, +c: U32, +sh: U32, +s: F32) -> (Array & Array) & F32: + (x, +a) = g + ((w, x), (s + U32.to_f32(((c >> U32.to_nat(sh)) .&. 15 : U32)) * a : F32)) + +def Array.q4mv.word(g: Array & U32, x: Array, +j: U32, +sh: U32, +s: F32) -> (Array & Array) & F32: + (w, +c) = g + Array.q4mv.mac(Array.get(F32, x, j), w, c, sh, s) + +# Lane sum: nibble sh/4 of words i .. i+n against x[j], x[j + 8], ..., in order. +def Array.q4mv.lane(n: Nat, q: (Array & Array) & F32, +i: U32, +j: U32, +sh: U32) -> (Array & Array) & F32: + match n: + case 0n: + q + case 1n+p: + ((w, x), +s) = q + Array.q4mv.lane(p, Array.q4mv.word(Array.get(U32, w, i), x, j, sh, s), (i + 1 : U32), (j + 8 : U32), sh) + +def Array.q4mv.add(q: (Array & Array) & F32, +a: F32) -> (Array & Array) & F32: + (wx, +b) = q + (wx, (a + b : F32)) + +# Lanes m and m + 4 of the group at code word i, activation j, added. +def Array.q4mv.l2(q: (Array & Array) & F32, +i: U32, +j: U32, +m: U32) -> (Array & Array) & F32: + (wx, +a) = q + Array.q4mv.add(Array.q4mv.lane(8n, (wx, 0.0), i, (j + m + 4 : U32), (m * 4 + 16 : U32)), a) + +def Array.q4mv.p2(wx: Array & Array, +i: U32, +j: U32, +m: U32) -> (Array & Array) & F32: + Array.q4mv.l2(Array.q4mv.lane(8n, (wx, 0.0), i, (j + m : U32), (m * 4 : U32)), i, j, m) + +def Array.q4mv.p4b(q: (Array & Array) & F32, +i: U32, +j: U32, +m: U32) -> (Array & Array) & F32: + (wx, +a) = q + Array.q4mv.add(Array.q4mv.p2(wx, i, j, (m + 2 : U32)), a) + +def Array.q4mv.p4(wx: Array & Array, +i: U32, +j: U32, +m: U32) -> (Array & Array) & F32: + Array.q4mv.p4b(Array.q4mv.p2(wx, i, j, m), i, j, m) + +def Array.q4mv.qxb(q: (Array & Array) & F32, +i: U32, +j: U32) -> (Array & Array) & F32: + (wx, +a) = q + Array.q4mv.add(Array.q4mv.p4(wx, i, j, 1), a) + +def Array.q4mv.gb(g: Array & U32, x: Array, +sw: U32, +qx: F32, +xs: F32, +even: Bool) -> (Array & Array) & F32: + (w, +bw) = g + +sc = Array.q4mv.half(even, sw) + ((w, x), (sc * qx + Array.q4mv.half(even, bw) * xs : F32)) + +def Array.q4mv.gs(g: Array & U32, h: Array & F32, +qx: F32, +even: Bool, +bat: U32) -> (Array & Array) & F32: + (w, +sw) = g + (x, +xs) = h + Array.q4mv.gb(Array.get(U32, w, bat), x, sw, qx, xs, even) + +def Array.q4mv.gfin(q: (Array & Array) & F32, +at: U32, +d: U32, +xg: U32, +even: Bool) -> (Array & Array) & F32: + ((w, x), +qx) = q + Array.q4mv.gs(Array.get(U32, w, at), Array.get(F32, x, xg), qx, even, (at + d : U32)) + +# Group g's term: its lanes meet as ((0+4)+(2+6))+((1+5)+(3+7)) into qx, then +# scale * qx + bias * the group's x-sum. +def Array.q4mv.term(wx: Array & Array, +i: U32, +g: U32, +sa: U32, +d: U32, +cols: U32) -> (Array & Array) & F32: + +k = (i + g * 8 : U32) + +j = (g * 64 : U32) + Array.q4mv.gfin(Array.q4mv.qxb(Array.q4mv.p4(wx, k, j, 0), k, j), (sa + (g >> 1n) : U32), d, (cols + g : U32), + U32.is_zero((g .&. 1 : U32))) + +# Partial l: 0.0 plus the terms of groups l, l + 32, ... in order. +def Array.q4mv.part(n: Nat, q: (Array & Array) & F32, +i: U32, +g: U32, +sa: U32, +d: U32, +cols: U32) -> (Array & Array) & F32: + match n: + case 0n: + q + case 1n+p: + (wx, +s) = q + Array.q4mv.part(p, Array.q4mv.add(Array.q4mv.term(wx, i, g, sa, d, cols), s), i, (g + 32 : U32), sa, d, cols) + +def Array.q4mv.foldb(q: (Array & Array) & F32, +n: Nat, +i: U32, +l: U32, +s: U32, +sa: U32, +d: U32, +cols: U32) -> (Array & Array) & F32: + (wx, +a) = q + Array.q4mv.add(Array.q4mv.fold(n, wx, i, (l + s : U32), (s * 2 : U32), sa, d, cols), a) + +# The 32 partials folded as p[l] += p[l + o] for o = 16, 8, 4, 2, 1: the value +# at lane l, stride s, is the partial l when s = 32, else its two halves' sum. +def Array.q4mv.fold(+n: Nat, wx: Array & Array, +i: U32, +l: U32, +s: U32, +sa: U32, +d: U32, +cols: U32) -> (Array & Array) & F32: + match n: + case 0n: + Array.q4mv.part(U32.to_nat((((cols >> 6n) + 31 - l) >> 5n : U32)), (wx, 0.0), i, l, sa, d, cols) + case 1n+p: + Array.q4mv.foldb(Array.q4mv.fold(p, wx, i, l, (s * 2 : U32), sa, d, cols), p, i, l, s, sa, d, cols) + +def Array.q4mv.put(q: (Array & Array) & F32, y: Array, +i: U32) -> (Array & Array) & Array: + (wx, +acc) = q + (wx, Array.set(F32, y, i, Array.q4mv.round(acc))) + +def Array.q4mv.rows(n: Nat, q: (Array & Array) & Array, +i: U32, +cols: U32, +cw: U32, +s0: U32, +sw: U32, +d: U32) -> (Array & Array) & Array: + match n: + case 0n: + q + case 1n+p: + (wx, y) = q + Array.q4mv.rows(p, Array.q4mv.put(Array.q4mv.fold(5n, wx, (i * cw : U32), 0, 1, (s0 + i * sw : U32), d, cols), y, i), + (i + 1 : U32), cols, cw, s0, sw, d) + +def Array.q4mv(w: Array, x: Array, y: Array, +r: U32, +n: U32, +cols: U32, +rows: U32) -> (Array & Array) & Array: + +cw = (cols >> 3n : U32) + +sw = (cols >> 7n : U32) + Array.q4mv.rows(U32.to_nat(n), ((w, x), y), r, cols, cw, (rows * cw : U32), sw, (rows * sw : U32)) + +# Array.q4go is Array.q4mv, except that a native build with a Metal heap may +# queue the product on the GPU: y then holds it only after the next +# Array.q4wait. Queued products run concurrently, so none may read another's +# y. The GPU reads w and x and writes y until the wait, and a value is freed +# at its last use: wait on x, and use w and y after the wait. +def Array.q4go(w: Array, x: Array, y: Array, +r: U32, +n: U32, +cols: U32, +rows: U32) -> (Array & Array) & Array: + Array.q4mv(w, x, y, r, n, cols, rows) + +# Array.q4wait(x) is x, once every queued Array.q4go has run. +def Array.q4wait(x: Array) -> Array: + x + +# Array.rmsq(x, y, w, o, n, eps, add) is the residual add and RMS norm that +# feeds Array.q4mv, for n a multiple of 64. When add is not 0, x[i] becomes +# x[i] + y[i] rounded to bfloat16 (as Array.q4mv rounds). s sums x[i] * x[i] +# from 0.0 in index order, and inv is 1.0 / sqrt(s / n + eps). o[i] becomes +# w[i] * (x[i] * inv) rounded to bfloat16 (w holds F32 bits), and o[n + g] +# the sum of group g as Array.q4mv reads it: each 32-value half summed from +# 0.0 in order, then the halves added. A native build with a Metal heap +# queues it as Array.q4go, after every product queued before it and before +# every one queued after; use x, y, w and o after the next Array.q4wait. +# The Metal kernel sums in a tree, so its last bits may differ. +def Array.rmsq.two(a: Array, r: Array & F32, +u: F32) -> (Array & Array) & (F32 & F32): + (b, +v) = r + ((a, b), (u, v)) + +def Array.rmsq.get(r: Array & F32, b: Array, +i: U32) -> (Array & Array) & (F32 & F32): + (a, +u) = r + Array.rmsq.two(a, Array.get(F32, b, i), u) + +def Array.rmsq.add1(r: (Array & Array) & (F32 & F32), +i: U32) -> Array & Array: + ((a, b), (u, v)) = r + (Array.set(F32, a, i, Array.q4mv.round((u + v : F32))), b) + +def Array.rmsq.adds(k: Nat, ab: Array & Array, +i: U32) -> Array & Array: + match k: + case 0n: + ab + case 1n+p: + (a, b) = ab + Array.rmsq.adds(p, Array.rmsq.add1(Array.rmsq.get(Array.get(F32, a, i), b, i), i), (i + 1 : U32)) + +def Array.rmsq.sq1(r: Array & F32, +s: F32) -> Array & F32: + (a, +u) = r + (a, (s + u * u : F32)) + +def Array.rmsq.sqs(k: Nat, as: Array & F32, +i: U32) -> Array & F32: + match k: + case 0n: + as + case 1n+p: + (a, +s) = as + Array.rmsq.sqs(p, Array.rmsq.sq1(Array.get(F32, a, i), s), (i + 1 : U32)) + +def Array.rmsq.wtwo(a: Array, r: Array & U32, +u: F32) -> (Array & Array) & (F32 & F32): + (b, +v) = r + ((a, b), (u, Array.q4mv.f32(v))) + +def Array.rmsq.getw(r: Array & F32, b: Array, +i: U32) -> (Array & Array) & (F32 & F32): + (a, +u) = r + Array.rmsq.wtwo(a, Array.get(U32, b, i), u) + +def Array.rmsq.mul1(r: (Array & Array) & (F32 & F32), o: Array, +i: U32, +k: F32) -> (Array & Array) & Array: + (xw, (u, v)) = r + (xw, Array.set(F32, o, i, Array.q4mv.round((v * (u * k : F32) : F32)))) + +def Array.rmsq.muls(n: Nat, q: (Array & Array) & Array, +i: U32, +k: F32) -> (Array & Array) & Array: + match n: + case 0n: + q + case 1n+p: + ((x, w), o) = q + Array.rmsq.muls(p, Array.rmsq.mul1(Array.rmsq.getw(Array.get(F32, x, i), w, i), o, i, k), (i + 1 : U32), k) + +def Array.rmsq.ad1(r: Array & F32, +s: F32) -> Array & F32: + (a, +u) = r + (a, (s + u : F32)) + +def Array.rmsq.half(k: Nat, as: Array & F32, +i: U32) -> Array & F32: + match k: + case 0n: + as + case 1n+p: + (a, +s) = as + Array.rmsq.half(p, Array.rmsq.ad1(Array.get(F32, a, i), s), (i + 1 : U32)) + +def Array.rmsq.gset(r: Array & F32, +sa: F32, +j: U32) -> Array: + (a, +sb) = r + Array.set(F32, a, j, (sa + sb : F32)) + +def Array.rmsq.g1(r: Array & F32, +i: U32, +j: U32) -> Array: + (a, +sa) = r + Array.rmsq.gset(Array.rmsq.half(32n, (a, 0.0), (i + 32 : U32)), sa, j) + +def Array.rmsq.gs(k: Nat, o: Array, +g: U32, +n: U32) -> Array: + match k: + case 0n: + o + case 1n+p: + Array.rmsq.gs(p, Array.rmsq.g1(Array.rmsq.half(32n, (o, 0.0), (g * 64 : U32)), (g * 64 : U32), (n + g : U32)), + (g + 1 : U32), n) + +def Array.rmsq.fin(q: (Array & Array) & Array, y: Array, +n: U32) -> ((Array & Array) & Array) & Array: + ((x, w), o) = q + (((x, y), w), Array.rmsq.gs(U32.to_nat((n >> 6n : U32)), o, 0, n)) + +def Array.rmsq.norm(xs: Array & F32, w: Array, o: Array, y: Array, +n: U32, +eps: F32) -> ((Array & Array) & Array) & Array: + (x, +s) = xs + +k = (1.0 / F32.sqrt((s / U32.to_f32(n) + eps : F32)) : F32) + Array.rmsq.fin(Array.rmsq.muls(U32.to_nat(n), ((x, w), o), 0, k), y, n) + +def Array.rmsq.go(xy: Array & Array, w: Array, o: Array, +n: U32, +eps: F32) -> ((Array & Array) & Array) & Array: + (x, y) = xy + Array.rmsq.norm(Array.rmsq.sqs(U32.to_nat(n), (x, 0.0), 0), w, o, y, n, eps) + +def Array.rmsq.sel(z: Bool, x: Array, y: Array, +n: U32) -> Array & Array: + match z: + case True{}: + (x, y) + case False{}: + Array.rmsq.adds(U32.to_nat(n), (x, y), 0) + +def Array.rmsq(x: Array, y: Array, w: Array, o: Array, +n: U32, +eps: F32, +add: U32) -> ((Array & Array) & Array) & Array: + Array.rmsq.go(Array.rmsq.sel(U32.is_zero(add), x, y, n), w, o, n, eps) + +# Array.swq(a, o, n) is the SwiGLU that feeds Array.q4mv, for n a multiple +# of 64: o[i] becomes silu(a[i]) * a[n + i] rounded to bfloat16, with +# silu(x) = x / (1.0 + exp(-x)), and o[n + g] the sum of group g as +# Array.rmsq writes it. Queued like Array.rmsq; use a and o after the next +# Array.q4wait. Metal's exp and its tree sums may differ in the last bits. +def Array.swq.put(r: Array & F32, o: Array, +g: F32, +i: U32) -> Array & Array: + (a, +u) = r + (a, Array.set(F32, o, i, Array.q4mv.round(((g / (1.0 + F32.exp(F32.neg(g))) : F32) * u : F32)))) + +def Array.swq.at(r: Array & F32, o: Array, +i: U32, +n: U32) -> Array & Array: + (a, +g) = r + Array.swq.put(Array.get(F32, a, (n + i : U32)), o, g, i) + +def Array.swq.go(k: Nat, r: Array & Array, +i: U32, +n: U32) -> Array & Array: + match k: + case 0n: + r + case 1n+p: + (a, o) = r + Array.swq.go(p, Array.swq.at(Array.get(F32, a, i), o, i, n), (i + 1 : U32), n) + +def Array.swq.fin(r: Array & Array, +n: U32) -> Array & Array: + (a, o) = r + (a, Array.rmsq.gs(U32.to_nat((n >> 6n : U32)), o, 0, n)) + +def Array.swq(a: Array, o: Array, +n: U32) -> Array & Array: + Array.swq.fin(Array.swq.go(U32.to_nat(n), (a, o), 0, n), n) + +# Array.gdn(y, w, s, cw, ab, sn, o, c, dk, dv, nvh, din, ko, eps) is one +# token of a gated DeltaNet mixer (Qwen3.5's linear attention), feeding +# Array.q4mv. y holds the projections: qkv (c values: q at 0, k at ko, v at +# 2 * ko), then z (din), beta (nvh) and alpha (nvh). w is the conv window +# (rows of c values, oldest first), s the state (value head h, row j at +# (h * dv + j) * dk), cw the conv taps (4 a channel, oldest first), ab +# a (nvh) then dt (nvh), and sn the gated norm's weights (dv); cw, ab and +# sn hold F32 bits. Value head h reads key head h. Every value below is +# rounded to bfloat16 where it says round, and sums run in index order: +# 1. channel i: v = round(w0 * a + .. + w3 * y[i]) over the window and +# y[i], conv = round(v * round(sigmoid(v))); the window moves on a row. +# 2. each q head of dk: q = round(round(q * inv) * round(1 / dk)) with +# inv = 1 / sqrt(sum q^2 / dk + eps / dk); each k head likewise with +# round(1 / sqrt(dk)). +# 3. head h: g = exp(round(max(x, 0) + log(1 + exp(-|x|))) * a[h]) with +# x = round(alpha[h] + dt[h]), beta = round(sigmoid(beta[h])); row j: +# d = sum (s * g) * k, then s = s * g + ((v_j - d) * beta) * k in place +# and o_j = sum s * q. +# 4. head h: o = round(o), inv = 1 / sqrt(sum o^2 / dk + eps), then +# o = round(round(sn * (o * inv)) * silu(z)). +# 5. o[din + g] holds group g's sum as Array.rmsq writes it. +# Queued like Array.rmsq; use every array after the next Array.q4wait. +# Metal's exp and its tree sums may differ in the last bits. +def Array.gdn.sig(+x: F32) -> F32: + (1.0 / (1.0 + F32.exp(F32.neg(x))) : F32) + +def Array.gdn.uf(r: Array & U32) -> Array & F32: + (a, +u) = r + (a, Array.q4mv.f32(u)) + +def Array.gdn.ui(a: Array, +i: U32) -> Array & F32: + Array.gdn.uf(Array.get(U32, a, i)) + +# Conv of channel i: the window's three values and y[i] against the taps. +def Array.gdn.cv4(r: Array & F32, y: Array, w: Array, co: Array, +d: F32, +acc: F32, +i: U32) -> + ((Array & Array) & Array) & Array: + (cw, +t) = r + +v = Array.q4mv.round((acc + d * t : F32)) + (((y, w), cw), Array.set(F32, co, i, Array.q4mv.round((v * Array.q4mv.round(Array.gdn.sig(v)) : F32)))) + +def Array.gdn.cv3(r: Array & F32, w: Array, cw: Array, co: Array, +acc: F32, +i: U32, +c: U32) -> + ((Array & Array) & Array) & Array: + (y, +d) = r + Array.gdn.cv4(Array.gdn.ui(cw, (i * 4 + 3 : U32)), y, Array.set(F32, w, (i + c + c : U32), d), co, d, acc, i) + +def Array.gdn.cv2(r: Array & F32, y: Array, w: Array, co: Array, +x: F32, +acc: F32, +i: U32, +c: U32) -> + ((Array & Array) & Array) & Array: + (cw, +t) = r + Array.gdn.cv3(Array.get(F32, y, i), Array.set(F32, w, (i + c : U32), x), cw, co, (acc + x * t : F32), i, c) + +def Array.gdn.cvc(r: Array & F32, y: Array, cw: Array, co: Array, +acc: F32, +i: U32, +c: U32) -> + ((Array & Array) & Array) & Array: + (w, +x) = r + Array.gdn.cv2(Array.gdn.ui(cw, (i * 4 + 2 : U32)), y, w, co, x, acc, i, c) + +def Array.gdn.cv1(r: Array & F32, y: Array, w: Array, co: Array, +x: F32, +acc: F32, +i: U32, +c: U32) -> + ((Array & Array) & Array) & Array: + (cw, +t) = r + Array.gdn.cvc(Array.get(F32, Array.set(F32, w, i, x), (i + c + c : U32)), y, cw, co, (acc + x * t : F32), i, c) + +def Array.gdn.cvb(r: Array & F32, y: Array, cw: Array, co: Array, +acc: F32, +i: U32, +c: U32) -> + ((Array & Array) & Array) & Array: + (w, +x) = r + Array.gdn.cv1(Array.gdn.ui(cw, (i * 4 + 1 : U32)), y, w, co, x, acc, i, c) + +def Array.gdn.cv0(r: Array & F32, y: Array, w: Array, co: Array, +a: F32, +i: U32, +c: U32) -> + ((Array & Array) & Array) & Array: + (cw, +t) = r + Array.gdn.cvb(Array.get(F32, w, (i + c : U32)), y, cw, co, (a * t : F32), i, c) + +def Array.gdn.cva(r: Array & F32, y: Array, cw: Array, co: Array, +i: U32, +c: U32) -> + ((Array & Array) & Array) & Array: + (w, +a) = r + Array.gdn.cv0(Array.gdn.ui(cw, (i * 4 : U32)), y, w, co, a, i, c) + +def Array.gdn.conv(k: Nat, q: ((Array & Array) & Array) & Array, +i: U32, +c: U32) -> + ((Array & Array) & Array) & Array: + match k: + case 0n: + q + case 1n+p: + (((y, w), cw), co) = q + Array.gdn.conv(p, Array.gdn.cva(Array.get(F32, w, i), y, cw, co, i, c), (i + 1 : U32), c) + +# a[i ..] = round(a[i ..] * k) over n slots. +def Array.gdn.sc1(r: Array & F32, +i: U32, +k: F32) -> Array: + (a, +x) = r + Array.set(F32, a, i, Array.q4mv.round((x * k : F32))) + +def Array.gdn.scl(n: Nat, a: Array, +i: U32, +k: F32) -> Array: + match n: + case 0n: + a + case 1n+p: + Array.gdn.scl(p, Array.gdn.sc1(Array.get(F32, a, i), i, k), (i + 1 : U32), k) + +def Array.gdn.l2f(r: Array & F32, +i: U32, +n: U32, +eps: F32, +sc: F32) -> Array: + (a, +s) = r + +nf = U32.to_f32(n) + +inv = (1.0 / F32.sqrt((s / nf + eps / nf : F32)) : F32) + Array.gdn.scl(U32.to_nat(n), Array.gdn.scl(U32.to_nat(n), a, i, inv), i, sc) + +# k heads of n values from i, each scaled to unit norm and by sc, in place. +def Array.gdn.l2(k: Nat, a: Array, +i: U32, +n: U32, +eps: F32, +sc: F32) -> Array: + match k: + case 0n: + a + case 1n+p: + Array.gdn.l2(p, Array.gdn.l2f(Array.rmsq.sqs(U32.to_nat(n), (a, 0.0), i), i, n, eps, sc), (i + n : U32), n, eps, sc) + +# The delta rule's state rows: s[ri ..] against k = co[ki ..], q = co[qi ..]. +def Array.gdn.d1(g: Array & F32, s: Array, +a: F32, +eg: F32, +acc: F32) -> (Array & Array) & F32: + (co, +k) = g + ((s, co), (acc + (a * eg : F32) * k : F32)) + +def Array.gdn.d0(g: Array & F32, co: Array, +ki: U32, +eg: F32, +acc: F32) -> (Array & Array) & F32: + (s, +a) = g + Array.gdn.d1(Array.get(F32, co, ki), s, a, eg, acc) + +def Array.gdn.dot(n: Nat, r: (Array & Array) & F32, +ri: U32, +ki: U32, +eg: F32) -> (Array & Array) & F32: + match n: + case 0n: + r + case 1n+p: + ((s, co), +acc) = r + Array.gdn.dot(p, Array.gdn.d0(Array.get(F32, s, ri), co, ki, eg, acc), (ri + 1 : U32), (ki + 1 : U32), eg) + +def Array.gdn.u2(g: Array & F32, s: Array, +nv: F32, +o: F32) -> (Array & Array) & F32: + (co, +q) = g + ((s, co), (o + nv * q : F32)) + +def Array.gdn.u1(g: Array & F32, s: Array, +a: F32, +ri: U32, +qi: U32, +eg: F32, +dj: F32, +o: F32) -> + (Array & Array) & F32: + (co, +k) = g + +nv = (a * eg + dj * k : F32) + Array.gdn.u2(Array.get(F32, co, qi), Array.set(F32, s, ri, nv), nv, o) + +def Array.gdn.u0(g: Array & F32, co: Array, +ri: U32, +ki: U32, +qi: U32, +eg: F32, +dj: F32, +o: F32) -> + (Array & Array) & F32: + (s, +a) = g + Array.gdn.u1(Array.get(F32, co, ki), s, a, ri, qi, eg, dj, o) + +def Array.gdn.upd(n: Nat, r: (Array & Array) & F32, +ri: U32, +ki: U32, +qi: U32, +eg: F32, +dj: F32) -> + (Array & Array) & F32: + match n: + case 0n: + r + case 1n+p: + ((s, co), +o) = r + Array.gdn.upd(p, Array.gdn.u0(Array.get(F32, s, ri), co, ri, ki, qi, eg, dj, o), (ri + 1 : U32), (ki + 1 : U32), + (qi + 1 : U32), eg, dj) + +def Array.gdn.ro(r: (Array & Array) & F32, o: Array, +oi: U32) -> (Array & Array) & Array: + ((s, co), +ov) = r + ((s, co), Array.set(F32, o, oi, ov)) + +def Array.gdn.rv(r: Array & F32, s: Array, o: Array, +d: F32, +ri: U32, +ki: U32, +qi: U32, +oi: U32, + +eg: F32, +beta: F32, +dk: U32) -> (Array & Array) & Array: + (co, +v) = r + Array.gdn.ro(Array.gdn.upd(U32.to_nat(dk), ((s, co), 0.0), ri, ki, qi, eg, ((v - d) * beta : F32)), o, oi) + +def Array.gdn.rd(r: (Array & Array) & F32, o: Array, +vi: U32, +ri: U32, +ki: U32, +qi: U32, +oi: U32, + +eg: F32, +beta: F32, +dk: U32) -> (Array & Array) & Array: + ((s, co), +d) = r + Array.gdn.rv(Array.get(F32, co, vi), s, o, d, ri, ki, qi, oi, eg, beta, dk) + +# Rows j .. of value head h; each reads s, co (q at qi, k at ki, v at vi) +# and writes o[oi + j]. +def Array.gdn.rows(n: Nat, q: (Array & Array) & Array, +j: U32, +ri: U32, +ki: U32, +qi: U32, +vi: U32, + +oi: U32, +eg: F32, +beta: F32, +dk: U32) -> (Array & Array) & Array: + match n: + case 0n: + q + case 1n+p: + (sc, o) = q + Array.gdn.rows(p, Array.gdn.rd(Array.gdn.dot(U32.to_nat(dk), (sc, 0.0), ri, ki, eg), o, (vi + j : U32), ri, ki, qi, + (oi + j : U32), eg, beta, dk), (j + 1 : U32), (ri + dk : U32), ki, qi, vi, oi, eg, beta, dk) + +def Array.gdn.soft(+x: F32) -> F32: + Array.q4mv.round((F32.max(x, 0.0) + F32.log((1.0 + F32.exp(F32.neg(F32.abs(x))) : F32)) : F32)) + +def Array.gdn.hb(r: Array & F32, s: Array, co: Array, o: Array, ab: Array, +eg: F32, +h: U32, + +dk: U32, +dv: U32, +ko: U32) -> (((Array & Array) & Array) & Array) & Array: + (y, +b) = r + ((Array.gdn.rows(U32.to_nat(dv), ((s, co), o), 0, (h * dv * dk : U32), (ko + h * dk : U32), (h * dk : U32), + (ko + ko + h * dv : U32), (h * dv : U32), eg, Array.q4mv.round(Array.gdn.sig(b)), dk), y), ab) + +def Array.gdn.hp(g: Array & F32, ab: Array, s: Array, co: Array, o: Array, +a: F32, +dt: F32, + +h: U32, +dk: U32, +dv: U32, +nvh: U32, +ko: U32, +bi: U32) -> (((Array & Array) & Array) & Array) & Array: + (y, +al) = g + Array.gdn.hb(Array.get(F32, y, (bi + h : U32)), s, co, o, ab, + F32.exp((Array.gdn.soft(Array.q4mv.round((al + dt : F32))) * a : F32)), h, dk, dv, ko) + +def Array.gdn.hd(r: Array & F32, y: Array, s: Array, co: Array, o: Array, +a: F32, +h: U32, + +dk: U32, +dv: U32, +nvh: U32, +ko: U32, +bi: U32) -> (((Array & Array) & Array) & Array) & Array: + (ab, +dt) = r + Array.gdn.hp(Array.get(F32, y, (bi + nvh + h : U32)), ab, s, co, o, a, dt, h, dk, dv, nvh, ko, bi) + +def Array.gdn.ha(r: Array & F32, y: Array, s: Array, co: Array, o: Array, +h: U32, +dk: U32, + +dv: U32, +nvh: U32, +ko: U32, +bi: U32) -> (((Array & Array) & Array) & Array) & Array: + (ab, +a) = r + Array.gdn.hd(Array.gdn.ui(ab, (nvh + h : U32)), y, s, co, o, a, h, dk, dv, nvh, ko, bi) + +# Value heads h ..: decay, beta and the delta rule's rows. +def Array.gdn.heads(n: Nat, q: (((Array & Array) & Array) & Array) & Array, +h: U32, +dk: U32, + +dv: U32, +nvh: U32, +ko: U32, +bi: U32) -> (((Array & Array) & Array) & Array) & Array: + match n: + case 0n: + q + case 1n+p: + ((((s, co), o), y), ab) = q + Array.gdn.heads(p, Array.gdn.ha(Array.gdn.ui(ab, h), y, s, co, o, h, dk, dv, nvh, ko, bi), (h + 1 : U32), dk, dv, + nvh, ko, bi) + +# The gated norm: o[i] = round(round(sn[j] * (o[i] * inv)) * silu(y[zi])). +def Array.gdn.g2(r: Array & F32, o: Array, sn: Array, +x: F32, +w: F32, +i: U32) -> + (Array & Array) & Array: + (y, +z) = r + ((Array.set(F32, o, i, Array.q4mv.round((x * (z / (1.0 + F32.exp(F32.neg(z))) : F32) : F32))), y), sn) + +def Array.gdn.g1(r: Array & F32, o: Array, y: Array, +v: F32, +inv: F32, +i: U32, +zi: U32) -> + (Array & Array) & Array: + (sn, +w) = r + Array.gdn.g2(Array.get(F32, y, zi), o, sn, Array.q4mv.round((w * (v * inv : F32) : F32)), w, i) + +def Array.gdn.g0(r: Array & F32, y: Array, sn: Array, +inv: F32, +i: U32, +j: U32, +zi: U32) -> + (Array & Array) & Array: + (o, +v) = r + Array.gdn.g1(Array.gdn.ui(sn, j), o, y, v, inv, i, zi) + +def Array.gdn.gate(n: Nat, q: (Array & Array) & Array, +inv: F32, +i: U32, +j: U32, +zi: U32) -> + (Array & Array) & Array: + match n: + case 0n: + q + case 1n+p: + ((o, y), sn) = q + Array.gdn.gate(p, Array.gdn.g0(Array.get(F32, o, i), y, sn, inv, i, j, zi), inv, (i + 1 : U32), (j + 1 : U32), + (zi + 1 : U32)) + +def Array.gdn.gnf(r: Array & F32, y: Array, sn: Array, +i: U32, +zi: U32, +n: U32, +nf: F32, +eps: F32) -> + (Array & Array) & Array: + (o, +s) = r + Array.gdn.gate(U32.to_nat(n), ((o, y), sn), (1.0 / F32.sqrt((s / nf + eps : F32)) : F32), i, 0, zi) + +def Array.gdn.gn(k: Nat, q: (Array & Array) & Array, +i: U32, +zi: U32, +n: U32, +nf: F32, +eps: F32) -> + (Array & Array) & Array: + match k: + case 0n: + q + case 1n+p: + ((o, y), sn) = q + Array.gdn.gn(p, Array.gdn.gnf(Array.rmsq.sqs(U32.to_nat(n), (Array.gdn.scl(U32.to_nat(n), o, i, 1.0), 0.0), i), y, + sn, i, zi, n, nf, eps), (i + n : U32), (zi + n : U32), n, nf, eps) + +def Array.gdn.fin(q: (Array & Array) & Array, co: Array, w: Array, s: Array, cw: Array, + ab: Array, +din: U32) -> (((((Array & Array) & Array) & Array) & Array) & Array) & Array: + ((o, y), sn) = q + _ = co + ((((((y, w), s), cw), ab), sn), Array.rmsq.gs(U32.to_nat((din >> 6n : U32)), o, 0, din)) + +def Array.gdn.mix(q: (((Array & Array) & Array) & Array) & Array, w: Array, cw: Array, + sn: Array, +c: U32, +dk: U32, +dv: U32, +nvh: U32, +din: U32, +eps: F32) -> + (((((Array & Array) & Array) & Array) & Array) & Array) & Array: + ((((s, co), o), y), ab) = q + Array.gdn.fin(Array.gdn.gn(U32.to_nat(nvh), ((o, y), sn), 0, c, dv, U32.to_f32(dk), eps), co, w, s, cw, ab, din) + +def Array.gdn.norm(q: ((Array & Array) & Array) & Array, s: Array, ab: Array, sn: Array, + o: Array, +c: U32, +dk: U32, +dv: U32, +nvh: U32, +din: U32, +ko: U32, +eps: F32) -> + (((((Array & Array) & Array) & Array) & Array) & Array) & Array: + (((y, w), cw), co) = q + +nk = U32.to_nat(U32.div(ko, dk)) + +df = U32.to_f32(dk) + co2 = Array.gdn.l2(nk, Array.gdn.l2(nk, co, 0, dk, eps, Array.q4mv.round((1.0 / df : F32))), ko, dk, eps, + Array.q4mv.round((1.0 / F32.sqrt(df) : F32))) + Array.gdn.mix(Array.gdn.heads(U32.to_nat(nvh), ((((s, co2), o), y), ab), 0, dk, dv, nvh, ko, (c + din : U32)), w, cw, sn, + c, dk, dv, nvh, din, eps) + +def Array.gdn.co(r: Array & Array, w: Array, s: Array, cw: Array, ab: Array, sn: Array, + o: Array, +c: U32, +dk: U32, +dv: U32, +nvh: U32, +din: U32, +ko: U32, +eps: F32) -> + (((((Array & Array) & Array) & Array) & Array) & Array) & Array: + (y, co) = r + Array.gdn.norm(Array.gdn.conv(U32.to_nat(c), (((y, w), cw), co), 0, c), s, ab, sn, o, c, dk, dv, nvh, din, ko, eps) + +def Array.gdn(y: Array, w: Array, s: Array, cw: Array, ab: Array, sn: Array, o: Array, + +c: U32, +dk: U32, +dv: U32, +nvh: U32, +din: U32, +ko: U32, +eps: F32) -> + (((((Array & Array) & Array) & Array) & Array) & Array) & Array: + Array.gdn.co(Array.clone(F32, y), w, s, cw, ab, sn, o, c, dk, dv, nvh, din, ko, eps) + +# Array.atq(y, kc, vc, kw, cs, o, pos, nq, nkv, hd, half, eps) is one token of +# gated full attention (Qwen3.5's), feeding Array.q4mv. y holds the +# projections: query head h's q at h * 2 * hd and its gate after it, then +# nkv keys (at nq * 2 * hd) and nkv values, hd values each. kc and vc are +# the key and value caches as F32 bits, position p's head g at +# (p * nkv + g) * hd; kw holds the q norm's weights then the k norm's (hd +# each, F32 bits); cs the rotation's cos then sin (half each). In order: +# 1. each key head, then each query head, in place in y: x = round(w * (x * +# inv)) with inv = 1 / sqrt(sum x^2 / hd + eps), then slots j < half +# rotate against j + half: (a, b) to (round(a c - b s), round(b c + a s)). +# 2. the keys and values join the caches at pos. +# 3. query head h reads cache head g = h / (nq / nkv) at every position t +# from pos down to 0: score = (sum k q, in order) * (1 / sqrt(hd)), m their +# max, e = exp(score - m), sum = the e's added in that order, acc = the +# e * v's added in that order; o[h * hd + d] = round(round(acc * (1 / sum)) +# * round(sigmoid(gate))). +# 4. o[nq * hd + g] holds group g's sum as Array.rmsq writes it. +# Queued like Array.rmsq; use every array after the next Array.q4wait. +# Metal's exp and its tree sums may differ in the last bits, and it leaves +# y as it was. A Metal heap queues positions below 4096; past that the call +# runs on the CPU at once, so wait before it. +def Array.atq.dep(k: Nat, +n: U32, +m: U32, +d: Nat) -> Nat: + match k: + case 0n: + d + case 1n+p: + Bool.pick(Nat, U32.is_le(n, m), d, Array.atq.dep(p, n, (m + m : U32), 1n+d)) + +# x[i ..] = round(w[j ..] * (x[i ..] * inv)) over n slots. +def Array.atq.w1(r: Array & F32, a: Array, +x: F32, +inv: F32, +i: U32) -> Array & Array: + (kw, +w) = r + (Array.set(F32, a, i, Array.q4mv.round((w * (x * inv : F32) : F32))), kw) + +def Array.atq.w0(r: Array & F32, kw: Array, +inv: F32, +i: U32, +j: U32) -> Array & Array: + (a, +x) = r + Array.atq.w1(Array.gdn.ui(kw, j), a, x, inv, i) + +def Array.atq.ws(n: Nat, r: Array & Array, +inv: F32, +i: U32, +j: U32) -> Array & Array: + match n: + case 0n: + r + case 1n+p: + (a, kw) = r + Array.atq.ws(p, Array.atq.w0(Array.get(F32, a, i), kw, inv, i, j), inv, (i + 1 : U32), (j + 1 : U32)) + +def Array.atq.nf(r: Array & F32, kw: Array, +i: U32, +j: U32, +hd: U32, +eps: F32) -> Array & Array: + (a, +s) = r + Array.atq.ws(U32.to_nat(hd), (a, kw), (1.0 / F32.sqrt((s / U32.to_f32(hd) + eps : F32)) : F32), i, j) + +# The rotation of slot i + j against i + j + half by cs[j], cs[half + j]. +def Array.atq.r3(r: Array & F32, a: Array, +x: F32, +y: F32, +c: F32, +i: U32, +h: U32) -> Array & Array: + (cs, +s) = r + (Array.set(F32, Array.set(F32, a, i, Array.q4mv.round((x * c - y * s : F32))), (i + h : U32), + Array.q4mv.round((y * c + x * s : F32))), cs) + +def Array.atq.r2(r: Array & F32, a: Array, +x: F32, +y: F32, +i: U32, +j: U32, +h: U32) -> Array & Array: + (cs, +c) = r + Array.atq.r3(Array.get(F32, cs, (h + j : U32)), a, x, y, c, i, h) + +def Array.atq.r1(r: Array & F32, cs: Array, +x: F32, +i: U32, +j: U32, +h: U32) -> Array & Array: + (a, +y) = r + Array.atq.r2(Array.get(F32, cs, j), a, x, y, i, j, h) + +def Array.atq.r0(r: Array & F32, cs: Array, +i: U32, +j: U32, +h: U32) -> Array & Array: + (a, +x) = r + Array.atq.r1(Array.get(F32, a, (i + h : U32)), cs, x, i, j, h) + +def Array.atq.rot(n: Nat, r: Array & Array, +i: U32, +j: U32, +h: U32) -> Array & Array: + match n: + case 0n: + r + case 1n+p: + (a, cs) = r + Array.atq.rot(p, Array.atq.r0(Array.get(F32, a, i), cs, i, j, h), (i + 1 : U32), (j + 1 : U32), h) + +def Array.atq.hr.k(r: Array & Array, kw: Array) -> (Array & Array) & Array: + (a, cs) = r + ((a, kw), cs) + +def Array.atq.hr(r: Array & Array, cs: Array, +i: U32, +h: U32) -> (Array & Array) & Array: + (a, kw) = r + Array.atq.hr.k(Array.atq.rot(U32.to_nat(h), (a, cs), i, 0, h), kw) + +def Array.atq.hd1(q: (Array & Array) & Array, +i: U32, +j: U32, +hd: U32, +h: U32, +eps: F32) -> + (Array & Array) & Array: + ((a, kw), cs) = q + Array.atq.hr(Array.atq.nf(Array.rmsq.sqs(U32.to_nat(hd), (a, 0.0), i), kw, i, j, hd, eps), cs, i, h) + +# k heads from i, step apart: norm by kw[j ..], then rotate, in place. +def Array.atq.hds(k: Nat, q: (Array & Array) & Array, +i: U32, +step: U32, +j: U32, +hd: U32, +h: U32, + +eps: F32) -> (Array & Array) & Array: + match k: + case 0n: + q + case 1n+p: + Array.atq.hds(p, Array.atq.hd1(q, i, j, hd, h, eps), (i + step : U32), step, j, hd, h, eps) + +# c[j ..] = the bits of a[i ..] over n slots. +def Array.atq.p1(r: Array & F32, c: Array, +j: U32) -> Array & Array: + (a, +x) = r + (a, Array.set(U32, c, j, Array.q4mv.bits(x))) + +def Array.atq.put(n: Nat, r: Array & Array, +i: U32, +j: U32) -> Array & Array: + match n: + case 0n: + r + case 1n+p: + (a, c) = r + Array.atq.put(p, Array.atq.p1(Array.get(F32, a, i), c, j), (i + 1 : U32), (j + 1 : U32)) + +# acc + k * q over hd slots, in order (cache k at ki, q in a at qi). +def Array.atq.d1(r: Array & F32, kc: Array, +k: F32, +acc: F32) -> (Array & Array) & F32: + (a, +q) = r + ((kc, a), (acc + k * q : F32)) + +def Array.atq.d0(r: Array & F32, a: Array, +qi: U32, +acc: F32) -> (Array & Array) & F32: + (kc, +k) = r + Array.atq.d1(Array.get(F32, a, qi), kc, k, acc) + +def Array.atq.dot(n: Nat, r: (Array & Array) & F32, +ki: U32, +qi: U32) -> (Array & Array) & F32: + match n: + case 0n: + r + case 1n+p: + ((kc, a), +acc) = r + Array.atq.dot(p, Array.atq.d0(Array.gdn.ui(kc, ki), a, qi, acc), (ki + 1 : U32), (qi + 1 : U32)) + +def Array.atq.sc1(r: (Array & Array) & F32, ss: Array, +t: U32, +sc: F32) -> + (Array & Array) & Array: + (ka, +d) = r + (ka, Array.set(F32, ss, t, (d * sc : F32))) + +# ss[t] = score of the position t below the newest, its key at ki. +def Array.atq.scs(n: Nat, q: (Array & Array) & Array, +t: U32, +ki: U32, +qi: U32, +stride: U32, + +hd: U32, +sc: F32) -> (Array & Array) & Array: + match n: + case 0n: + q + case 1n+p: + (ka, ss) = q + Array.atq.scs(p, Array.atq.sc1(Array.atq.dot(U32.to_nat(hd), (ka, 0.0), ki, qi), ss, t, sc), (t + 1 : U32), + (ki - stride : U32), qi, stride, hd, sc) + +def Array.atq.mx1(r: Array & F32, +m: F32) -> Array & F32: + (a, +x) = r + (a, F32.max(x, m)) + +def Array.atq.mx(n: Nat, r: Array & F32, +i: U32) -> Array & F32: + match n: + case 0n: + r + case 1n+p: + (a, +m) = r + Array.atq.mx(p, Array.atq.mx1(Array.get(F32, a, i), m), (i + 1 : U32)) + +def Array.atq.ex1(r: Array & F32, +i: U32, +m: F32) -> Array: + (a, +x) = r + Array.set(F32, a, i, F32.exp((x - m : F32))) + +def Array.atq.ex(n: Nat, a: Array, +i: U32, +m: F32) -> Array: + match n: + case 0n: + a + case 1n+p: + Array.atq.ex(p, Array.atq.ex1(Array.get(F32, a, i), i, m), (i + 1 : U32), m) + +# acc[d] = acc[d] + v[d] * pr over hd slots (cache v at vi). +def Array.atq.a1(r: Array & F32, vc: Array, +v: F32, +pr: F32, +d: U32) -> Array & Array: + (acc, +x) = r + (vc, Array.set(F32, acc, d, (x + (v * pr : F32) : F32))) + +def Array.atq.a0(r: Array & F32, acc: Array, +pr: F32, +d: U32) -> Array & Array: + (vc, +v) = r + Array.atq.a1(Array.get(F32, acc, d), vc, v, pr, d) + +def Array.atq.acc(n: Nat, r: Array & Array, +vi: U32, +d: U32, +pr: F32) -> Array & Array: + match n: + case 0n: + r + case 1n+p: + (vc, acc) = r + Array.atq.acc(p, Array.atq.a0(Array.gdn.ui(vc, vi), acc, pr, d), (vi + 1 : U32), (d + 1 : U32), pr) + +def Array.atq.wt1(r: Array & F32, va: Array & Array, +vi: U32, +hd: U32) -> Array & (Array & Array): + (es, +pr) = r + (es, Array.atq.acc(U32.to_nat(hd), va, vi, 0, pr)) + +# acc = sum of e_t * v_t over the positions, newest first. +def Array.atq.wts(n: Nat, q: Array & (Array & Array), +t: U32, +vi: U32, +stride: U32, +hd: U32) -> + Array & (Array & Array): + match n: + case 0n: + q + case 1n+p: + (es, va) = q + Array.atq.wts(p, Array.atq.wt1(Array.get(F32, es, t), va, vi, hd), (t + 1 : U32), (vi - stride : U32), stride, hd) + +def Array.atq.o2(r: Array & F32, acc: Array, o: Array, +a: F32, +inv: F32, +oi: U32) -> + (Array & Array) & Array: + (y, +g) = r + ((acc, y), Array.set(F32, o, oi, Array.q4mv.round((Array.q4mv.round((a * inv : F32)) * + Array.q4mv.round(Array.gdn.sig(g)) : F32)))) + +def Array.atq.o1(r: Array & F32, y: Array, o: Array, +inv: F32, +gi: U32, +oi: U32) -> + (Array & Array) & Array: + (acc, +a) = r + Array.atq.o2(Array.get(F32, y, gi), acc, o, a, inv, oi) + +def Array.atq.outs(n: Nat, q: (Array & Array) & Array, +d: U32, +gi: U32, +oi: U32, +inv: F32) -> + (Array & Array) & Array: + match n: + case 0n: + q + case 1n+p: + ((acc, y), o) = q + Array.atq.outs(p, Array.atq.o1(Array.get(F32, acc, d), y, o, inv, gi, oi), (d + 1 : U32), (gi + 1 : U32), + (oi + 1 : U32), inv) + +# One query head's attention; its scratch arrays are dropped. +def Array.atq.h4(q: (Array & Array) & Array, ss: Array, kc: Array, vc: Array) -> + ((Array & Array) & Array) & Array: + ((acc, y), o) = q + _ = acc + _ = ss + (((y, kc), vc), o) + +def Array.atq.h3(r: Array & (Array & Array), y: Array, kc: Array, o: Array, +sum: F32, + +h: U32, +hd: U32) -> ((Array & Array) & Array) & Array: + (ss, va) = r + (vc, acc) = va + Array.atq.h4(Array.atq.outs(U32.to_nat(hd), ((acc, y), o), 0, (h * hd * 2 + hd : U32), (h * hd : U32), + (1.0 / sum : F32)), ss, kc, vc) + +def Array.atq.h2(r: Array & F32, y: Array, kc: Array, vc: Array, o: Array, +n: U32, +vi: U32, + +stride: U32, +h: U32, +hd: U32) -> ((Array & Array) & Array) & Array: + (ss, +sum) = r + Array.atq.h3(Array.atq.wts(U32.to_nat(n), (ss, (vc, Array.new(F32, Array.atq.dep(32n, hd, 1, 0n), 0.0))), 0, vi, stride, + hd), y, kc, o, sum, h, hd) + +def Array.atq.h1(r: Array & F32, y: Array, kc: Array, vc: Array, o: Array, +n: U32, +vi: U32, + +stride: U32, +h: U32, +hd: U32) -> ((Array & Array) & Array) & Array: + (ss, +m) = r + Array.atq.h2(Array.rmsq.half(U32.to_nat(n), (Array.atq.ex(U32.to_nat(n), ss, 0, m), 0.0), 0), y, kc, vc, o, n, vi, + stride, h, hd) + +def Array.atq.h0(q: (Array & Array) & Array, vc: Array, o: Array, +n: U32, +vi: U32, + +stride: U32, +h: U32, +hd: U32) -> ((Array & Array) & Array) & Array: + ((kc, y), ss) = q + Array.atq.h1(Array.atq.mx(U32.to_nat((n - 1 : U32)), Array.get(F32, ss, 0), 1), y, kc, vc, o, n, vi, stride, h, hd) + +def Array.atq.head(q: ((Array & Array) & Array) & Array, +h: U32, +pos: U32, +nq: U32, +nkv: U32, + +hd: U32) -> ((Array & Array) & Array) & Array: + (((y, kc), vc), o) = q + +n = (pos + 1 : U32) + +stride = (nkv * hd : U32) + +at = ((pos * nkv + U32.div(h, U32.div(nq, nkv))) * hd : U32) + Array.atq.h0(Array.atq.scs(U32.to_nat(n), ((kc, y), Array.new(F32, Array.atq.dep(32n, n, 1, 0n), 0.0)), 0, at, + (h * hd * 2 : U32), stride, hd, (1.0 / F32.sqrt(U32.to_f32(hd)) : F32)), vc, o, n, at, stride, h, hd) + +def Array.atq.heads(k: Nat, q: ((Array & Array) & Array) & Array, +h: U32, +pos: U32, +nq: U32, + +nkv: U32, +hd: U32) -> ((Array & Array) & Array) & Array: + match k: + case 0n: + q + case 1n+p: + Array.atq.heads(p, Array.atq.head(q, h, pos, nq, nkv, hd), (h + 1 : U32), pos, nq, nkv, hd) + +def Array.atq.fin(q: ((Array & Array) & Array) & Array, kw: Array, cs: Array, +n: U32) -> + ((((Array & Array) & Array) & Array) & Array) & Array: + (((y, kc), vc), o) = q + (((((y, kc), vc), kw), cs), Array.rmsq.gs(U32.to_nat((n >> 6n : U32)), o, 0, n)) + +def Array.atq.cv(r: Array & Array, kc: Array, kw: Array, cs: Array, o: Array, +pos: U32, + +nq: U32, +nkv: U32, +hd: U32) -> ((((Array & Array) & Array) & Array) & Array) & Array: + (y, vc) = r + Array.atq.fin(Array.atq.heads(U32.to_nat(nq), (((y, kc), vc), o), 0, pos, nq, nkv, hd), kw, cs, (nq * hd : U32)) + +def Array.atq.ck(r: Array & Array, vc: Array, kw: Array, cs: Array, o: Array, +pos: U32, + +nq: U32, +nkv: U32, +hd: U32) -> ((((Array & Array) & Array) & Array) & Array) & Array: + (y, kc) = r + +ko = (nq * hd * 2 : U32) + Array.atq.cv(Array.atq.put(U32.to_nat((nkv * hd : U32)), (y, vc), (ko + nkv * hd : U32), (pos * nkv * hd : U32)), kc, + kw, cs, o, pos, nq, nkv, hd) + +def Array.atq.rq(q: (Array & Array) & Array, kc: Array, vc: Array, o: Array, +pos: U32, + +nq: U32, +nkv: U32, +hd: U32) -> ((((Array & Array) & Array) & Array) & Array) & Array: + ((y, kw), cs) = q + Array.atq.ck(Array.atq.put(U32.to_nat((nkv * hd : U32)), (y, kc), (nq * hd * 2 : U32), (pos * nkv * hd : U32)), vc, kw, + cs, o, pos, nq, nkv, hd) + +def Array.atq(y: Array, kc: Array, vc: Array, kw: Array, cs: Array, o: Array, +pos: U32, + +nq: U32, +nkv: U32, +hd: U32, +half: U32, +eps: F32) -> + ((((Array & Array) & Array) & Array) & Array) & Array: + Array.atq.rq(Array.atq.hds(U32.to_nat(nq), Array.atq.hds(U32.to_nat(nkv), ((y, kw), cs), (nq * hd * 2 : U32), hd, hd, + hd, half, eps), 0, (hd * 2 : U32), 0, hd, half, eps), kc, vc, o, pos, nq, nkv, hd) + def Array.to_list.go(~T: Type, a: Array, acc: List) -> List: match a: case ALeaf{x}: diff --git a/bend2/comp.ts b/bend2/comp.ts index de8caa31aa..839d97fa86 100644 --- a/bend2/comp.ts +++ b/bend2/comp.ts @@ -248,6 +248,15 @@ const OPERATIONS: Record = Object.setPrototypeOf({ swap: "array_rmw($0, $1, () => $2)", size: "{$: \"Tuple\", fst: $0, snd: $0.length}", }).map(([k, JS]) => ["array_" + k, { C: null, JS }])), + array_q4mv: { C: ["$0", "$1", "q4mv(e, $0, $1, $2, $3, $4, $5, $6, 0)"] }, + array_q4go: { C: ["$0", "$1", "q4mv(e, $0, $1, $2, $3, $4, $5, $6, 1)"] }, + array_q4wait: { C: "q4wait($0)" }, + array_rmsq: { C: ["$0", "$1", "$2", "rmsq(e, $0, $1, $2, $3, $4, $5, $6)"] }, + array_swq: { C: ["$0", "swq(e, $0, $1, $2)"] }, + array_atq: { C: ["$0", "$1", "$2", "$3", "$4", + "atq(e, $0, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)"] }, + array_gdn: { C: ["$0", "$1", "$2", "$3", "$4", "$5", + "gdn(e, $0, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)"] }, array_clone: { C: ["$0", "blk_copy(e, $0)"], JS: "{$: \"Tuple\", fst: $0, snd: $0.slice()}", @@ -582,7 +591,7 @@ function tpl_deep(e: string): boolean { function tpl(t: Tpl, xs: string[]): string { return typeof t !== "string" ? t(xs) - : t.split(/\$(\d)/).map((p, i) => (i % 2 ? xs[+p] : p)).join(""); + : t.split(/\$(\d+)/).map((p, i) => (i % 2 ? xs[+p] : p)).join(""); } function view_of(e: string): string | undefined { @@ -810,7 +819,7 @@ function intr_of(k: Name, js = false): Intr | undefined { const tld = FL.book.tlds[k]; const it = tld?.$ === "Def" && tld.i === undefined && tld.b ? OPERATIONS[op_name(k)] : undefined; - return it && (js || it.C !== undefined) ? it : undefined; + return it && (js ? it.JS : it.C) !== undefined ? it : undefined; } function op_name(k: Name): string { @@ -1410,10 +1419,11 @@ function file_book(roots: Name[]): void { type_adts(s.T).forEach(src_add); } if (s.$ === "Ref") { - if (s.b) { + if (s.b || ["Array.q4go", "Array.rmsq", "Array.swq", "Array.gdn", + "Array.atq"].includes(s.k)) { FL.bangs.add(s.k); } - if (intr_of(s.k) === undefined) { + if (intr_of(s.k, FL.js) === undefined) { refs.add(s.k); FL.sites.set(s.k, (FL.sites.get(s.k) ?? 0) + 1); } @@ -3297,6 +3307,9 @@ using namespace metal; #include #include #include +#if defined(__ARM_NEON) +#include +#endif #include #include #include @@ -3567,6 +3580,21 @@ static id gpu_que; static id gpu_pso; static id gpu_buf; static id gpu_enc; +static id gpu_lib; +static id gpu_q4pso; +static id gpu_qcb; +static id gpu_qenc; +static id gpu_qlast; +static pthread_mutex_t gpu_qlock = PTHREAD_MUTEX_INITIALIZER; +static double gpu_qt0; +static u32 gpu_qk; +static id gpu_qfpso; +static id gpu_qflag; +static u32 gpu_qseq; +static id gpu_rmpso; +static id gpu_swpso; +static id gpu_gdpso; +static id gpu_atpso; #elif BEND_CUDA static CUdevice gpu_dev; static CUmodule gpu_lib; @@ -4221,6 +4249,483 @@ INLINE Term blk_new(Env e, bool arr, u64 d, u32 lgs, u32 n, THR Term* v) { return term_blk(arr, c, l); } +// Array.q4mv: rows r .. r+n of an MLX affine q4g64 matrix against x, each +// row summed in base.bend's tree order and narrowed to bfloat16 into y. + +INLINE f32 q4_half(u32 u, u32 g) { + return f32_unbox(g & 1 ? u & 0xFFFF0000u : u << 16); +} + +INLINE u32 q4_round(f32 v) { + u32 u = (u32)f32_rewrap(v); + u32 b = (u & 0x7FFFFFFFu) > 0x7F800000u ? ((u >> 16) & 0x8000u) | 0x7FFFu + : (u + 0x7FFFu + ((u >> 16) & 1u)) >> 16; + return b << 16; +} + +// A group's term: its eight lane sums meet as ((0+4)+(2+6))+((1+5)+(3+7)), +// then scale * that + bias * the group's x-sum. +INLINE f32 q4_term(DEV u32a* W, DEV u32a* X, u32 wm, u32 xm, u32 at, u32 d, + u32 g, u32 c, THR f32* L) { + f32 qx = ((L[0] + L[4]) + (L[2] + L[6])) + ((L[1] + L[5]) + (L[3] + L[7])); + return q4_half(W[at & wm], g) * qx + + q4_half(W[(at + d) & wm], g) * f32_unbox(X[(c + g) & xm]); +} + +// The 32 lane partials fold as p[l] += p[l + o] for o = 16, 8, 4, 2, 1. +INLINE u32 q4_fold(THR f32* p) { + for (u32 o = 16; o > 0; o >>= 1) { + for (u32 l = 0; l < o; l += 1) { + p[l] = p[l] + p[l + o]; + } + } + return q4_round(p[0]); +} + +// One row i: lane j of group g sums q*x over its words k = 0..7 in order +// (nibble j of word k), and group g's term goes onto partial g % 32. +INLINE u32 q4_row(DEV u32a* W, DEV u32a* X, u32 wm, u32 xm, u32 i, u32 c, + u32 cw, u32 s0, u32 sw, u32 d) { + f32 p[32]; + for (u32 l = 0; l < 32; l += 1) { + p[l] = 0.0f; + } + for (u32 g = 0; g < c >> 6; g += 1) { + f32 L[8] = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f }; + for (u32 k = 0; k < 8; k += 1) { + u32 q = W[(i * cw + g * 8 + k) & wm]; + for (u32 j = 0; j < 8; j += 1) { + L[j] = L[j] + (f32)((q >> (4 * j)) & 15) + * f32_unbox(X[(g * 64 + k * 8 + j) & xm]); + } + } + p[g & 31] = p[g & 31] + + q4_term(W, X, wm, xm, s0 + i * sw + (g >> 1), d, g, c, L); + } + return q4_fold(p); +} + +// The host hands a big product to the GPU's own kernel (a SIMD group a row) +// when a Metal heap is live; Q4_GPU_MIN weights is the break-even. +#define Q4_GPU_MIN (1ull << 22) +#ifndef Q4R +#define Q4R 4 +#endif +#if !DEVICE && BEND_METAL +static bool gpu_q4mv(u64 wl, u64 xl, u64 yl, u32 wm, u32 xm, u32 ym, u32 r, + u32 n, u32 c, u32 rows, bool q); +static void gpu_q4wait(void); +INLINE Term q4wait(Term x) { + gpu_q4wait(); + return x; +} +#else +INLINE Term q4wait(Term x) { + return x; +} +#endif + +#if !DEVICE && defined(__ARM_NEON) +// Four rows at once in NEON: vector a holds lanes 0..3, b lanes 4..7. +INLINE void q4_neon4(const u32* W, const f32* X, u32* Y, u32 i, u32 c, + u32 cw, u32 s0, u32 sw, u32 d) { + int32x4_t ra = { 0, -4, -8, -12 }; + int32x4_t rb = { -16, -20, -24, -28 }; + uint32x4_t nb = vdupq_n_u32(15); + f32 p[4][32]; + memset(p, 0, sizeof p); + for (u32 g = 0; g < c >> 6; g += 1) { + float32x4_t a[4], b[4]; + for (u32 r = 0; r < 4; r += 1) { + a[r] = b[r] = vdupq_n_f32(0.0f); + } + for (u32 k = 0; k < 8; k += 1) { + float32x4_t xa = vld1q_f32(X + g * 64 + k * 8); + float32x4_t xb = vld1q_f32(X + g * 64 + k * 8 + 4); + for (u32 r = 0; r < 4; r += 1) { + uint32x4_t q = vdupq_n_u32(W[(i + r) * cw + g * 8 + k]); + a[r] = vaddq_f32(a[r], vmulq_f32(xa, + vcvtq_f32_u32(vandq_u32(vshlq_u32(q, ra), nb)))); + b[r] = vaddq_f32(b[r], vmulq_f32(xb, + vcvtq_f32_u32(vandq_u32(vshlq_u32(q, rb), nb)))); + } + } + for (u32 r = 0; r < 4; r += 1) { + f32 L[8]; + vst1q_f32(L, a[r]); + vst1q_f32(L + 4, b[r]); + p[r][g & 31] = p[r][g & 31] + q4_term((u32a*)W, (u32a*)X, ~0u, ~0u, + s0 + (i + r) * sw + (g >> 1), d, g, c, L); + } + } + for (u32 r = 0; r < 4; r += 1) { + Y[i + r] = q4_fold(p[r]); + } +} +#endif + +// q (Array.q4go): a GPU product joins the open queue, any size, and y holds +// it only after the next Array.q4wait. +INLINE Term q4mv(Env e, Term w, Term x, Term y, u64 r, u64 n, u64 cols, + u64 rows, bool q) { + DEV u64* H = e.mem; + DEV u32a* W = blk_ptr(H, blk_loc(H, w), 0); + DEV u32a* X = blk_ptr(H, blk_loc(H, x), 0); + DEV u32a* Y = blk_ptr(H, blk_loc(H, y), 0); + u32 wm = (u32)((1ull << blk_cls(w)) - 1); + u32 xm = (u32)((1ull << blk_cls(x)) - 1); + u32 ym = (u32)((1ull << blk_cls(y)) - 1); + u32 c = (u32)cols; + u32 cw = c >> 3; + u32 sw = c >> 7; + u32 s0 = (u32)rows * cw; + u32 d = (u32)rows * sw; + u32 i = (u32)r; + u64 t = 0; + bool fit = (u64)rows * (cw + 2 * sw) <= (u64)wm + 1 + && (u64)c + (c >> 6) <= (u64)xm + 1 && r + n <= (u64)ym + 1 && r + n <= rows; +#if !DEVICE && BEND_METAL + if (fit && (q || n * cols >= Q4_GPU_MIN) && gpu_q4mv(blk_loc(H, w), + blk_loc(H, x), blk_loc(H, y), wm, xm, ym, i, (u32)n, c, (u32)rows, q)) { + return y; + } +#endif +#if !DEVICE && defined(__ARM_NEON) + for (; fit && t + 4 <= n; t += 4, i += 4) { + q4_neon4((u32*)W, (f32*)X, (u32*)Y, i, c, cw, s0, sw, d); + } +#endif + for (; t < n; t += 1, i += 1) { + Y[i & ym] = q4_row(W, X, wm, xm, i, c, cw, s0, sw, d); + } + return y; +} + +// Array.rmsq: base.bend's residual add and RMS norm into q4 activations; a +// Metal heap queues it between the products before and after it. +#if !DEVICE && BEND_METAL +static bool gpu_rmsq(u64 xl, u64 yl, u64 wl, u64 ol, u32 n, u32 eps, u32 add); +#endif +INLINE Term rmsq(Env e, Term x, Term y, Term w, Term o, u64 n, u64 eps, + u64 add) { + DEV u64* H = e.mem; + DEV u32a* X = blk_ptr(H, blk_loc(H, x), 0); + DEV u32a* Y = blk_ptr(H, blk_loc(H, y), 0); + DEV u32a* W = blk_ptr(H, blk_loc(H, w), 0); + DEV u32a* O = blk_ptr(H, blk_loc(H, o), 0); + u64 c = 1ull << blk_cls(x); + bool fit = n <= c && (n & 63) == 0 && n <= 1ull << blk_cls(w) + && (!add || n <= 1ull << blk_cls(y)) && n + (n >> 6) <= 1ull << blk_cls(o); +#if !DEVICE && BEND_METAL + if (fit && gpu_rmsq(blk_loc(H, x), blk_loc(H, y), blk_loc(H, w), + blk_loc(H, o), (u32)n, (u32)eps, (u32)add)) { + return o; + } +#endif + if (!fit) { + return o; + } + for (u32 i = 0; add && i < n; i += 1) { + X[i] = q4_round(f32_unbox(X[i]) + f32_unbox(Y[i])); + } + f32 s = 0.0f; + for (u32 i = 0; i < n; i += 1) { + s = s + f32_unbox(X[i]) * f32_unbox(X[i]); + } + f32 k = 1.0f / (f32)sqrt(s / (f32)(u32)n + f32_unbox(eps)); + for (u32 i = 0; i < n; i += 1) { + O[i] = q4_round(f32_unbox(W[i]) * (f32_unbox(X[i]) * k)); + } + for (u32 g = 0; g < n >> 6; g += 1) { + f32 a = 0.0f, b = 0.0f; + for (u32 j = 0; j < 32; j += 1) { + a = a + f32_unbox(O[g * 64 + j]); + b = b + f32_unbox(O[g * 64 + 32 + j]); + } + O[n + g] = (u32)f32_rewrap(a + b); + } + return o; +} + +// Array.swq and Array.gdn: base.bend's SwiGLU and gated DeltaNet step into +// q4 activations; a Metal heap queues them like Array.rmsq. +#if !DEVICE && BEND_METAL +static bool gpu_swq(u64 al, u64 ol, u32 n); +static bool gpu_gdn(const u64* at, const u32* p); +#endif +INLINE f32 k2_round(f32 v) { + return f32_unbox(q4_round(v)); +} + +INLINE f32 k2_sig(f32 x) { + return 1.0f / (1.0f + (f32)exp(-x)); +} + +// o[n + g]: group g's sum, each 32-value half in order, then the halves. +INLINE void k2_gsum(DEV u32a* O, u32 n) { + for (u32 g = 0; g < n >> 6; g += 1) { + f32 a = 0.0f, b = 0.0f; + for (u32 j = 0; j < 32; j += 1) { + a = a + f32_unbox(O[g * 64 + j]); + b = b + f32_unbox(O[g * 64 + 32 + j]); + } + O[n + g] = (u32)f32_rewrap(a + b); + } +} + +INLINE Term swq(Env e, Term a, Term o, u64 n) { + DEV u64* H = e.mem; + DEV u32a* A = blk_ptr(H, blk_loc(H, a), 0); + DEV u32a* O = blk_ptr(H, blk_loc(H, o), 0); + bool fit = (n & 63) == 0 && 2 * n <= 1ull << blk_cls(a) + && n + (n >> 6) <= 1ull << blk_cls(o); +#if !DEVICE && BEND_METAL + if (fit && gpu_swq(blk_loc(H, a), blk_loc(H, o), (u32)n)) { + return o; + } +#endif + for (u32 i = 0; fit && i < n; i += 1) { + f32 g = f32_unbox(A[i]); + f32 s = g / (1.0f + (f32)exp(-g)); + O[i] = q4_round(s * f32_unbox(A[n + i])); + } + if (fit) { + k2_gsum(O, (u32)n); + } + return o; +} + +#if !DEVICE +// k heads of n values from i: unit norm, then scaled by sc, in place. +static void k2_l2(f32* a, u32 i, u32 n, f32 eps, f32 sc) { + f32 s = 0.0f; + for (u32 j = 0; j < n; j += 1) { + s = s + a[i + j] * a[i + j]; + } + f32 nf = (f32)n; + f32 inv = 1.0f / (f32)sqrt(s / nf + eps / nf); + for (u32 j = 0; j < n; j += 1) { + a[i + j] = k2_round(a[i + j] * inv); + } + for (u32 j = 0; j < n; j += 1) { + a[i + j] = k2_round(a[i + j] * sc); + } +} +#endif + +INLINE Term gdn(Env e, Term y, Term w, Term s, Term cw, Term ab, Term sn, + Term o, u64 c, u64 dk, u64 dv, u64 nvh, u64 din, u64 ko, u64 eps) { + DEV u64* H = e.mem; + DEV u32a* Y = blk_ptr(H, blk_loc(H, y), 0); + DEV u32a* Wn = blk_ptr(H, blk_loc(H, w), 0); + DEV u32a* S = blk_ptr(H, blk_loc(H, s), 0); + DEV u32a* CW = blk_ptr(H, blk_loc(H, cw), 0); + DEV u32a* AB = blk_ptr(H, blk_loc(H, ab), 0); + DEV u32a* SN = blk_ptr(H, blk_loc(H, sn), 0); + DEV u32a* O = blk_ptr(H, blk_loc(H, o), 0); + bool fit = dk > 0 && dk <= 256 && (dk & 31) == 0 && dv <= 256 + && (dv & 63) == 0 && (din & 63) == 0 && nvh * dv <= din + && ko * 2 + nvh * dv <= c && nvh * dk <= ko + && c + din + 2 * nvh <= 1ull << blk_cls(y) && 3 * c <= 1ull << blk_cls(w) + && nvh * dv * dk <= 1ull << blk_cls(s) && 4 * c <= 1ull << blk_cls(cw) + && 2 * nvh <= 1ull << blk_cls(ab) && dv <= 1ull << blk_cls(sn) + && din + (din >> 6) <= 1ull << blk_cls(o); +#if !DEVICE && BEND_METAL + u64 at[7] = { blk_loc(H, y), blk_loc(H, w), blk_loc(H, s), blk_loc(H, cw), + blk_loc(H, ab), blk_loc(H, sn), blk_loc(H, o) }; + u32 p[8] = { (u32)c, (u32)dk, (u32)dv, (u32)nvh, (u32)din, (u32)eps, + (u32)ko, 0 }; + if (fit && gpu_gdn(at, p)) { + return o; + } +#endif +#if !DEVICE + if (!fit) { + return o; + } + f32 ep = f32_unbox(eps); + f32* co = (f32*)malloc(c * sizeof(f32)); + for (u32 i = 0; i < c; i += 1) { + f32 a = f32_unbox(Wn[i]), b = f32_unbox(Wn[i + c]); + f32 cc = f32_unbox(Wn[i + 2 * c]), d = f32_unbox(Y[i]); + f32 acc = a * f32_unbox(CW[i * 4]); + acc = acc + b * f32_unbox(CW[i * 4 + 1]); + acc = acc + cc * f32_unbox(CW[i * 4 + 2]); + f32 v = k2_round(acc + d * f32_unbox(CW[i * 4 + 3])); + Wn[i] = Wn[i + c]; + Wn[i + c] = Wn[i + 2 * c]; + Wn[i + 2 * c] = Y[i]; + co[i] = k2_round(v * k2_round(k2_sig(v))); + } + f32 df = (f32)(u32)dk; + for (u32 h = 0; h < ko / dk; h += 1) { + k2_l2(co, h * (u32)dk, (u32)dk, ep, k2_round(1.0f / df)); + } + for (u32 h = 0; h < ko / dk; h += 1) { + k2_l2(co, (u32)ko + h * (u32)dk, (u32)dk, ep, k2_round(1.0f / (f32)sqrt(df))); + } + for (u32 h = 0; h < nvh; h += 1) { + f32 x = k2_round(f32_unbox(Y[c + din + nvh + h]) + f32_unbox(AB[nvh + h])); + f32 l1 = 1.0f + (f32)exp(-fabsf(x)); + f32 sp = k2_round((x > 0.0f ? x : 0.0f) + (f32)log(l1)); + f32 eg = (f32)exp(sp * f32_unbox(AB[h])); + f32 beta = k2_round(k2_sig(f32_unbox(Y[c + din + h]))); + const f32* q = co + h * dk; + const f32* k = co + ko + h * dk; + const f32* v = co + 2 * ko + h * dv; + for (u32 j = 0; j < dv; j += 1) { + u64 ri = (h * dv + j) * dk; + f32 d = 0.0f; + for (u32 i = 0; i < dk; i += 1) { + d = d + (f32_unbox(S[ri + i]) * eg) * k[i]; + } + f32 dj = (v[j] - d) * beta; + f32 ov = 0.0f; + for (u32 i = 0; i < dk; i += 1) { + f32 nv = f32_unbox(S[ri + i]) * eg + dj * k[i]; + S[ri + i] = (u32)f32_rewrap(nv); + ov = ov + nv * q[i]; + } + O[h * dv + j] = (u32)f32_rewrap(ov); + } + } + for (u32 h = 0; h < nvh; h += 1) { + f32 ss = 0.0f; + for (u32 j = 0; j < dv; j += 1) { + f32 r = k2_round(f32_unbox(O[h * dv + j])); + O[h * dv + j] = (u32)f32_rewrap(r); + ss = ss + r * r; + } + f32 inv = 1.0f / (f32)sqrt(ss / df + ep); + for (u32 j = 0; j < dv; j += 1) { + f32 x = k2_round(f32_unbox(SN[j]) * (f32_unbox(O[h * dv + j]) * inv)); + f32 z = f32_unbox(Y[c + h * dv + j]); + O[h * dv + j] = q4_round(x * (z / (1.0f + (f32)exp(-z)))); + } + } + k2_gsum(O, (u32)din); + free(co); +#endif + return o; +} + +// Array.atq: base.bend's gated attention for one token over the caches; a +// Metal heap queues it (positions up to ATQ_MAX); past that it runs here at +// once, so the caller waits first. +#define ATQ_MAX 4096 +#if !DEVICE && BEND_METAL +static bool gpu_atq(const u64* at, const u32* p); +#endif +#if !DEVICE +// x[i ..] normed by w, then rotated by c, s, as base.bend's Array.atq.hds. +static void k2_rope(DEV u32a* Y, u32 i, DEV u32a* W, u32 hd, u32 hf, f32 eps, + DEV u32a* CS) { + f32 s = 0.0f; + for (u32 d = 0; d < hd; d += 1) { + s = s + f32_unbox(Y[i + d]) * f32_unbox(Y[i + d]); + } + f32 inv = 1.0f / (f32)sqrt(s / (f32)hd + eps); + for (u32 d = 0; d < hd; d += 1) { + Y[i + d] = q4_round(f32_unbox(W[d]) * (f32_unbox(Y[i + d]) * inv)); + } + for (u32 j = 0; j < hf; j += 1) { + f32 a = f32_unbox(Y[i + j]), b = f32_unbox(Y[i + j + hf]); + f32 c = f32_unbox(CS[j]), sn = f32_unbox(CS[hf + j]); + f32 ac = a * c, bs = b * sn, bc = b * c, as = a * sn; + Y[i + j] = q4_round(ac - bs); + Y[i + j + hf] = q4_round(bc + as); + } +} +#endif + +INLINE Term atq(Env e, Term y, Term kc, Term vc, Term kw, Term cs, Term o, + u64 pos, u64 nq, u64 nkv, u64 hd, u64 hf, u64 eps) { + DEV u64* H = e.mem; + DEV u32a* Y = blk_ptr(H, blk_loc(H, y), 0); + DEV u32a* KC = blk_ptr(H, blk_loc(H, kc), 0); + DEV u32a* VC = blk_ptr(H, blk_loc(H, vc), 0); + DEV u32a* KW = blk_ptr(H, blk_loc(H, kw), 0); + DEV u32a* CS = blk_ptr(H, blk_loc(H, cs), 0); + DEV u32a* O = blk_ptr(H, blk_loc(H, o), 0); + u64 st = nkv * hd; + bool fit = nkv > 0 && nq % nkv == 0 && hd > 0 && (hd & 63) == 0 + && 2 * hf <= hd && nq * hd * 2 + 2 * st <= 1ull << blk_cls(y) + && (pos + 1) * st <= 1ull << blk_cls(kc) && (pos + 1) * st <= 1ull << blk_cls(vc) + && 2 * hd <= 1ull << blk_cls(kw) && 2 * hf <= 1ull << blk_cls(cs) + && nq * hd + nq * hd / 64 <= 1ull << blk_cls(o); +#if !DEVICE && BEND_METAL + u64 at[6] = { blk_loc(H, y), blk_loc(H, kc), blk_loc(H, vc), blk_loc(H, kw), + blk_loc(H, cs), blk_loc(H, o) }; + u32 p[8] = { (u32)pos, (u32)nq, (u32)nkv, (u32)hd, (u32)hf, (u32)eps, 0, 0 }; + if (fit && hd <= 256 && pos < ATQ_MAX && gpu_atq(at, p)) { + return o; + } +#endif +#if !DEVICE + if (!fit) { + return o; + } + f32 ep = f32_unbox(eps); + u32 ko = (u32)(nq * hd * 2), n = (u32)pos + 1; + for (u32 g = 0; g < nkv; g += 1) { + k2_rope(Y, ko + g * (u32)hd, KW + hd, (u32)hd, (u32)hf, ep, CS); + } + for (u32 h = 0; h < nq; h += 1) { + k2_rope(Y, h * (u32)hd * 2, KW, (u32)hd, (u32)hf, ep, CS); + } + for (u64 d = 0; d < st; d += 1) { + KC[pos * st + d] = Y[ko + d]; + VC[pos * st + d] = Y[ko + st + d]; + } + f32* ss = (f32*)malloc(n * sizeof(f32)); + f32* acc = (f32*)malloc(hd * sizeof(f32)); + f32 sc = 1.0f / (f32)sqrt((f32)(u32)hd); + for (u32 h = 0; h < nq; h += 1) { + u64 g = h / (nq / nkv); + for (u32 t = 0; t < n; t += 1) { + u64 ki = ((pos - t) * nkv + g) * hd; + f32 dot = 0.0f; + for (u32 d = 0; d < hd; d += 1) { + f32 kq = f32_unbox(KC[ki + d]) * f32_unbox(Y[h * hd * 2 + d]); + dot = dot + kq; + } + ss[t] = dot * sc; + } + f32 m = ss[0]; + for (u32 t = 1; t < n; t += 1) { + m = ss[t] > m ? ss[t] : m; + } + f32 sum = 0.0f; + for (u32 t = 0; t < n; t += 1) { + ss[t] = (f32)exp(ss[t] - m); + } + for (u32 t = 0; t < n; t += 1) { + sum = sum + ss[t]; + } + for (u32 d = 0; d < hd; d += 1) { + acc[d] = 0.0f; + } + for (u32 t = 0; t < n; t += 1) { + u64 vi = ((pos - t) * nkv + g) * hd; + for (u32 d = 0; d < hd; d += 1) { + f32 vp = f32_unbox(VC[vi + d]) * ss[t]; + acc[d] = acc[d] + vp; + } + } + f32 inv = 1.0f / sum; + for (u32 d = 0; d < hd; d += 1) { + f32 a = k2_round(acc[d] * inv); + f32 gt = k2_round(k2_sig(f32_unbox(Y[h * hd * 2 + hd + d]))); + O[h * hd + d] = q4_round(a * gt); + } + } + free(ss); + free(acc); + k2_gsum(O, (u32)(nq * hd)); +#endif + return o; +} + // Ring // ==== @@ -4552,6 +5057,352 @@ INLINE void bank_pack(DEV u64* H, u32 lane) { } } +#ifdef __METAL_VERSION__ +// A SIMD group takes Q4R rows: lane l takes groups l, l + 32, ..., as +// q4_row's p[l], each group's 64 activations read once for all the rows. +kernel void bend_q4flag(device atomic_uint* F [[buffer(0)]], + constant u32& v [[buffer(1)]]) { + atomic_store_explicit(F, v, memory_order_relaxed); +} + +kernel void bend_q4mv(device const u32* W [[buffer(0)]], + device const u32* X [[buffer(1)]], device u32* Y [[buffer(2)]], + constant u32* P [[buffer(3)]], u32 tg [[threadgroup_position_in_grid]], + u32 sg [[simdgroup_index_in_threadgroup]], + u32 l [[thread_index_in_simdgroup]]) { + u32 k0 = (tg * 8 + sg) * Q4R; + if (k0 >= P[7]) { + return; + } + u32 nr = min((u32)Q4R, P[7] - k0); + u32 c = P[3], rows = P[4]; + u32 i0 = P[6] + k0; + u32 cw = c >> 3, sw = c >> 7; + f32 p[Q4R]; + for (u32 r = 0; r < Q4R; r += 1) { + p[r] = 0.0f; + } + for (u32 g = l; g < c >> 6; g += 32) { + f32 xv[64]; + for (u32 j = 0; j < 64; j += 1) { + xv[j] = as_type(X[g * 64 + j]); + } + f32 xg = as_type(X[c + g]); + for (u32 r = 0; r < Q4R; r += 1) { + if (r < nr) { + u32 i = i0 + r; + device const uint2* R = (device const uint2*)(W + i * cw); + float4 a = 0.0f, b = 0.0f; + for (u32 h = 0; h < 8; h += 1) { + u32 q = h & 1 ? R[g * 4 + (h >> 1)].y : R[g * 4 + (h >> 1)].x; + u32 j = h * 8; + a = a + float4(uint4(q, q >> 4, q >> 8, q >> 12) & 15u) + * float4(xv[j], xv[j + 1], xv[j + 2], xv[j + 3]); + b = b + float4(uint4(q >> 16, q >> 20, q >> 24, q >> 28) & 15u) + * float4(xv[j + 4], xv[j + 5], xv[j + 6], xv[j + 7]); + } + float4 e = a + b; + float2 f = e.xy + e.zw; + u32 at = rows * cw + i * sw + (g >> 1); + p[r] = p[r] + (q4_half(W[at], g) * (f.x + f.y) + + q4_half(W[at + rows * sw], g) * xg); + } + } + } + for (u32 r = 0; r < Q4R; r += 1) { + f32 v = p[r]; + for (u32 o = 16; o > 0; o >>= 1) { + v = v + simd_shuffle_down(v, (ushort)o); + } + if (l == 0 && r < nr) { + Y[i0 + r] = q4_round(v); + } + } +} + +// rmsq's loops as one threadgroup of 512; the sum of squares is a tree. +kernel void bend_rmsq(device u32* X [[buffer(0)]], device const u32* Y [[buffer(1)]], + device const u32* W [[buffer(2)]], device u32* O [[buffer(3)]], + constant u32* P [[buffer(4)]], u32 t [[thread_position_in_threadgroup]], + u32 sg [[simdgroup_index_in_threadgroup]], u32 l [[thread_index_in_simdgroup]]) { + threadgroup f32 ps[16]; + u32 n = P[0]; + f32 s = 0.0f; + for (u32 i = t; i < n; i += 512) { + f32 v = as_type(X[i]); + if (P[2]) { + v = as_type(q4_round(v + as_type(Y[i]))); + X[i] = as_type(v); + } + s = s + v * v; + } + s = simd_sum(s); + if (l == 0) { + ps[sg] = s; + } + threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup); + s = simd_sum(l < 16 ? ps[l] : 0.0f); + f32 k = 1.0f / sqrt(s / (f32)n + as_type(P[1])); + for (u32 g = sg; g < n >> 6; g += 16) { + u32 i = g * 64 + l; + f32 a = as_type(q4_round(as_type(W[i]) * (as_type(X[i]) * k))); + f32 b = as_type(q4_round(as_type(W[i + 32]) + * (as_type(X[i + 32]) * k))); + O[i] = as_type(a); + O[i + 32] = as_type(b); + f32 sa = simd_sum(a), sb = simd_sum(b); + if (l == 0) { + O[n + g] = as_type(sa + sb); + } + } +} + +// swq: a SIMD group per 64-value group. +kernel void bend_swq(device const u32* A [[buffer(0)]], device u32* O [[buffer(1)]], + constant u32* P [[buffer(2)]], u32 g [[threadgroup_position_in_grid]], + u32 l [[thread_index_in_simdgroup]]) { + u32 n = P[0], i = g * 64 + l; + f32 x = as_type(A[i]), z = as_type(A[i + 32]); + f32 a = as_type(q4_round(x / (1.0f + exp(-x)) * as_type(A[n + i]))); + f32 b = as_type(q4_round(z / (1.0f + exp(-z)) * as_type(A[n + i + 32]))); + O[i] = as_type(a); + O[i + 32] = as_type(b); + f32 sa = simd_sum(a), sb = simd_sum(b); + if (l == 0) { + O[n + g] = as_type(sa + sb); + } +} + +// gdn: a threadgroup of 512 per value head (head h reads key head h); every +// phase of base.bend's Array.gdn is local to the head. Rows of the delta +// rule go one to a SIMD group, dk / 32 state values to a lane. +kernel void bend_gdn(device const u32* Y [[buffer(0)]], device u32* Wn [[buffer(1)]], + device u32* S [[buffer(2)]], device const u32* CW [[buffer(3)]], + device const u32* AB [[buffer(4)]], device const u32* SN [[buffer(5)]], + device u32* O [[buffer(6)]], constant u32* P [[buffer(7)]], + u32 h [[threadgroup_position_in_grid]], u32 t [[thread_position_in_threadgroup]], + u32 sg [[simdgroup_index_in_threadgroup]], u32 l [[thread_index_in_simdgroup]]) { + threadgroup f32 q[256], k[256], v[256], o[256], sc[8]; + u32 c = P[0], dk = P[1], dv = P[2], nvh = P[3], din = P[4], ko = P[6]; + f32 eps = as_type(P[5]), df = (f32)dk; + if (t < 2 * dk + dv) { + u32 i = t < dk ? h * dk + t : t < 2 * dk ? ko + h * dk + t - dk + : 2 * ko + h * dv + t - 2 * dk; + f32 a = as_type(Wn[i]), b = as_type(Wn[i + c]); + f32 cc = as_type(Wn[i + 2 * c]), d = as_type(Y[i]); + f32 acc = a * as_type(CW[i * 4]); + acc = acc + b * as_type(CW[i * 4 + 1]); + acc = acc + cc * as_type(CW[i * 4 + 2]); + f32 x = as_type(q4_round(acc + d * as_type(CW[i * 4 + 3]))); + Wn[i] = as_type(b); + Wn[i + c] = as_type(cc); + Wn[i + 2 * c] = as_type(d); + x = as_type(q4_round(x * as_type(q4_round(1.0f / (1.0f + exp(-x)))))); + if (t < dk) { + q[t] = x; + } else if (t < 2 * dk) { + k[t - dk] = x; + } else { + v[t - 2 * dk] = x; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg < 2) { + threadgroup f32* a = sg == 0 ? q : k; + f32 s = 0.0f; + for (u32 i = l; i < dk; i += 32) { + s = s + a[i] * a[i]; + } + s = simd_sum(s); + sc[sg] = 1.0f / sqrt(s / df + eps / df); + } else if (sg == 2 && l == 0) { + f32 x = as_type(q4_round(as_type(Y[c + din + nvh + h]) + + as_type(AB[nvh + h]))); + f32 sp = as_type(q4_round(max(x, 0.0f) + log(1.0f + exp(-fabs(x))))); + sc[2] = exp(sp * as_type(AB[h])); + f32 b = as_type(Y[c + din + h]); + sc[3] = as_type(q4_round(1.0f / (1.0f + exp(-b)))); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t < dk) { + f32 s = as_type(q4_round(1.0f / df)); + q[t] = as_type(q4_round(as_type(q4_round(q[t] * sc[0])) * s)); + } else if (t < 2 * dk) { + f32 s = as_type(q4_round(1.0f / sqrt(df))); + k[t - dk] = as_type(q4_round(as_type(q4_round(k[t - dk] * sc[1])) * s)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + f32 eg = sc[2], beta = sc[3]; + for (u32 j = sg; j < dv; j += 16) { + device u32* R = S + (h * dv + j) * dk; + f32 sv[8]; + f32 d = 0.0f; + for (u32 m = 0; m < dk / 32; m += 1) { + sv[m] = as_type(R[l + 32 * m]); + d = d + (sv[m] * eg) * k[l + 32 * m]; + } + d = simd_sum(d); + f32 dj = (v[j] - d) * beta; + f32 ov = 0.0f; + for (u32 m = 0; m < dk / 32; m += 1) { + f32 nv = sv[m] * eg + dj * k[l + 32 * m]; + R[l + 32 * m] = as_type(nv); + ov = ov + nv * q[l + 32 * m]; + } + ov = simd_sum(ov); + if (l == 0) { + o[j] = as_type(q4_round(ov)); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + f32 s = 0.0f; + for (u32 i = l; i < dv; i += 32) { + s = s + o[i] * o[i]; + } + s = simd_sum(s); + sc[4] = 1.0f / sqrt(s / df + eps); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t < dv) { + f32 x = as_type(q4_round(as_type(SN[t]) * (o[t] * sc[4]))); + f32 z = as_type(Y[c + h * dv + t]); + f32 r = as_type(q4_round(x * (z / (1.0f + exp(-z))))); + O[h * dv + t] = as_type(r); + o[t] = r; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (u32 g = sg; g < dv >> 6; g += 16) { + f32 sa = simd_sum(o[g * 64 + l]), sb = simd_sum(o[g * 64 + 32 + l]); + if (l == 0) { + O[din + ((h * dv) >> 6) + g] = as_type(sa + sb); + } + } +} + +// atq: a threadgroup of 256 per query head. Each normalizes and rotates its +// q and its KV head's new k (the group's first head writes k and v to the +// caches at pos); scores, softmax and the weighted values follow, the +// current position read from the threadgroup, the older ones from the caches. +kernel void bend_atq(device const u32* Y [[buffer(0)]], device u32* KC [[buffer(1)]], + device u32* VC [[buffer(2)]], device const u32* KW [[buffer(3)]], + device const u32* CS [[buffer(4)]], device u32* O [[buffer(5)]], + constant u32* P [[buffer(6)]], u32 h [[threadgroup_position_in_grid]], + u32 t [[thread_position_in_threadgroup]], + u32 sg [[simdgroup_index_in_threadgroup]], u32 l [[thread_index_in_simdgroup]]) { + threadgroup f32 q[256], k[256], ss[ATQ_MAX], ra[8], rb[8], sc[4]; + u32 pos = P[0], nq = P[1], nkv = P[2], hd = P[3], hf = P[4]; + f32 eps = as_type(P[5]); + u32 g = h / (nq / nkv), ko = nq * hd * 2, st = nkv * hd, vo = ko + st; + u32 n = pos + 1; + f32 a = 0.0f, b = 0.0f; + for (u32 d = t; d < hd; d += 256) { + q[d] = as_type(Y[h * hd * 2 + d]); + k[d] = as_type(Y[ko + g * hd + d]); + a = a + q[d] * q[d]; + b = b + k[d] * k[d]; + } + a = simd_sum(a); + b = simd_sum(b); + if (l == 0) { + ra[sg] = a; + rb[sg] = b; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t == 0) { + f32 x = 0.0f, z = 0.0f; + for (u32 i = 0; i < 8; i += 1) { + x = x + ra[i]; + z = z + rb[i]; + } + sc[0] = 1.0f / sqrt(x / (f32)hd + eps); + sc[1] = 1.0f / sqrt(z / (f32)hd + eps); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (u32 d = t; d < hd; d += 256) { + q[d] = as_type(q4_round(as_type(KW[d]) * (q[d] * sc[0]))); + k[d] = as_type(q4_round(as_type(KW[hd + d]) * (k[d] * sc[1]))); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t < hf) { + f32 c = as_type(CS[t]), s = as_type(CS[hf + t]); + f32 x = q[t], y = q[t + hf], u = k[t], v = k[t + hf]; + q[t] = as_type(q4_round(x * c - y * s)); + q[t + hf] = as_type(q4_round(y * c + x * s)); + k[t] = as_type(q4_round(u * c - v * s)); + k[t + hf] = as_type(q4_round(v * c + u * s)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (h % (nq / nkv) == 0) { + for (u32 d = t; d < hd; d += 256) { + KC[pos * st + g * hd + d] = as_type(k[d]); + VC[pos * st + g * hd + d] = Y[vo + g * hd + d]; + } + } + f32 rs = 1.0f / sqrt((f32)hd); + for (u32 p = sg; p < n; p += 8) { + f32 dot = 0.0f; + for (u32 d = l; d < hd; d += 32) { + f32 kv = p == pos ? k[d] : as_type(KC[(p * nkv + g) * hd + d]); + dot = dot + kv * q[d]; + } + dot = simd_sum(dot); + if (l == 0) { + ss[p] = dot * rs; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + f32 m = -INFINITY; + for (u32 p = t; p < n; p += 256) { + m = max(m, ss[p]); + } + m = simd_max(m); + if (l == 0) { + ra[sg] = m; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + m = ra[0]; + for (u32 i = 1; i < 8; i += 1) { + m = max(m, ra[i]); + } + f32 s = 0.0f; + for (u32 p = t; p < n; p += 256) { + f32 e = exp(ss[p] - m); + ss[p] = e; + s = s + e; + } + s = simd_sum(s); + if (l == 0) { + rb[sg] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + s = 0.0f; + for (u32 i = 0; i < 8; i += 1) { + s = s + rb[i]; + } + f32 inv = 1.0f / s; + for (u32 d = t; d < hd; d += 256) { + f32 acc = 0.0f; + for (u32 p = 0; p < n; p += 1) { + u32 vb = p == pos ? Y[vo + g * hd + d] : VC[(p * nkv + g) * hd + d]; + acc = acc + as_type(vb) * ss[p]; + } + f32 gt = as_type(Y[h * hd * 2 + hd + d]); + f32 r = as_type(q4_round(as_type(q4_round(acc * inv)) + * as_type(q4_round(1.0f / (1.0f + exp(-gt)))))); + O[h * hd + d] = as_type(r); + q[d] = r; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (u32 gg = sg; gg < hd >> 6; gg += 8) { + f32 sa = simd_sum(q[gg * 64 + l]), sb = simd_sum(q[gg * 64 + 32 + l]); + if (l == 0) { + O[nq * hd + ((h * hd) >> 6) + gg] = as_type(sa + sb); + } + } +} +#endif + #ifdef __METAL_VERSION__ kernel void bend_dev(DEV u64* H [[buffer(0)]], constant u32& pass [[buffer(1)]], TG u32* vote [[threadgroup(0)]], @@ -4900,6 +5751,7 @@ static MTLComputePipelineDescriptor* gpu_desc(void) { if (!lib) { gpu_fail(err); } + gpu_lib = lib; MTLComputePipelineDescriptor* d = [MTLComputePipelineDescriptor new]; d.computeFunction = [lib newFunctionWithName:@"bend_dev"]; return d; @@ -4990,6 +5842,209 @@ static void gpu_pass(u32 f) { } } +static double gpu_us(void) { + struct timespec t; + clock_gettime(CLOCK_MONOTONIC, &t); + return t.tv_sec * 1e6 + t.tv_nsec / 1e3; +} + +// Commits the open queue and waits for it, or for the last one committed. +// BEND_Q4_TRACE=1 prints each wait's products and GPU and wall time. +static void gpu_q4wait(void) { + pthread_mutex_lock(&gpu_qlock); + id cb = gpu_qcb; + u32 k = gpu_qk; + u32 seq = 0; + if (cb != nil) { + seq = gpu_qseq += 1; + [gpu_qenc memoryBarrierWithScope:MTLBarrierScopeBuffers]; + [gpu_qenc setComputePipelineState:gpu_qfpso]; + [gpu_qenc setBuffer:gpu_qflag offset:0 atIndex:0]; + [gpu_qenc setBytes:&seq length:4 atIndex:1]; + [gpu_qenc dispatchThreads:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; + [gpu_qenc endEncoding]; + [cb commit]; + gpu_qlast = cb; + gpu_qcb = nil; + gpu_qenc = nil; + gpu_qk = 0; + } else { + cb = gpu_qlast; + } + pthread_mutex_unlock(&gpu_qlock); + if (cb == nil) { + return; + } + // The queue's last kernel stores seq after a barrier; spinning on it + // returns ~0.1 ms sooner than waitUntilCompleted. + volatile u32* f = (volatile u32*)[gpu_qflag contents]; + while (seq != 0 && *f != seq && [cb status] < MTLCommandBufferStatusCompleted) { + } + __atomic_thread_fence(__ATOMIC_ACQUIRE); + static int trace = -1; + if (trace < 0) { + const char* s = getenv("BEND_Q4_TRACE"); + trace = s != NULL && s[0] == '1'; + } + if (seq == 0 || trace) { + [cb waitUntilCompleted]; + } + if ([cb error]) { + gpu_fail([cb error]); + } + if (trace && k > 0) { + fprintf(stderr, "q4 wait k=%u gpu_us=%.1f wall_us=%.1f\n", k, + ([cb GPUEndTime] - [cb GPUStartTime]) * 1e6, gpu_us() - gpu_qt0); + } +} + +static id gpu_fn(NSString* name) { + NSError* err = nil; + id p = [gpu_dev newComputePipelineStateWithFunction: + [gpu_lib newFunctionWithName:name] error:&err]; + if (p == nil) { + gpu_fail(err); + } + return p; +} + +static void gpu_qinit(void) { + if (gpu_q4pso == nil) { + gpu_q4pso = gpu_fn(@"bend_q4mv"); + gpu_qfpso = gpu_fn(@"bend_q4flag"); + gpu_rmpso = gpu_fn(@"bend_rmsq"); + gpu_swpso = gpu_fn(@"bend_swq"); + gpu_gdpso = gpu_fn(@"bend_gdn"); + gpu_atpso = gpu_fn(@"bend_atq"); + gpu_qflag = [gpu_dev newBufferWithLength:64 + options:MTLResourceStorageModeShared]; + } +} + +// Opens the queue's command buffer if none is (caller holds gpu_qlock). +static void gpu_qopen(void) { + gpu_qinit(); + if (gpu_qcb == nil) { + gpu_qt0 = gpu_us(); + gpu_qcb = [gpu_que commandBuffer]; + gpu_qenc = [gpu_qcb computeCommandEncoderWithDispatchType: + MTLDispatchTypeConcurrent]; + } + gpu_qk += 1; +} + +static void gpu_qkern(id pso, const u64* at, u32 nb, + const u32* p, u32 np, u32 groups, u32 threads); + +static bool gpu_rmsq(u64 xl, u64 yl, u64 wl, u64 ol, u32 n, u32 eps, u32 add) { + if (gpu_buf == nil || gpu_lib == nil) { + return false; + } + @autoreleasepool { + pthread_mutex_lock(&gpu_qlock); + u64 at[4] = { xl, yl, wl, ol }; + u32 p[4] = { n, eps, add, 0 }; + gpu_qinit(); + gpu_qkern(gpu_rmpso, at, 4, p, 4, 1, 512); + pthread_mutex_unlock(&gpu_qlock); + } + return true; +} + +// A non-product kernel on the open queue: barriers on both sides, since the +// encoder is concurrent (caller holds gpu_qlock). +static void gpu_qkern(id pso, const u64* at, u32 nb, + const u32* p, u32 np, u32 groups, u32 threads) { + gpu_qopen(); + [gpu_qenc memoryBarrierWithScope:MTLBarrierScopeBuffers]; + [gpu_qenc setComputePipelineState:pso]; + for (u32 i = 0; i < nb; i += 1) { + [gpu_qenc setBuffer:gpu_buf offset:at[i] * 8 atIndex:i]; + } + [gpu_qenc setBytes:p length:np * 4 atIndex:nb]; + [gpu_qenc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + [gpu_qenc memoryBarrierWithScope:MTLBarrierScopeBuffers]; +} + +static bool gpu_swq(u64 al, u64 ol, u32 n) { + if (gpu_buf == nil || gpu_lib == nil) { + return false; + } + @autoreleasepool { + pthread_mutex_lock(&gpu_qlock); + u64 at[2] = { al, ol }; + u32 p[4] = { n, 0, 0, 0 }; + gpu_qinit(); + gpu_qkern(gpu_swpso, at, 2, p, 4, n >> 6, 32); + pthread_mutex_unlock(&gpu_qlock); + } + return true; +} + +static bool gpu_gdn(const u64* at, const u32* p) { + if (gpu_buf == nil || gpu_lib == nil) { + return false; + } + @autoreleasepool { + pthread_mutex_lock(&gpu_qlock); + gpu_qinit(); + gpu_qkern(gpu_gdpso, at, 7, p, 8, p[3], 512); + pthread_mutex_unlock(&gpu_qlock); + } + return true; +} + +static bool gpu_atq(const u64* at, const u32* p) { + if (gpu_buf == nil || gpu_lib == nil) { + return false; + } + @autoreleasepool { + pthread_mutex_lock(&gpu_qlock); + gpu_qinit(); + gpu_qkern(gpu_atpso, at, 6, p, 8, p[1], 256); + pthread_mutex_unlock(&gpu_qlock); + } + return true; +} + +static bool gpu_q4mv(u64 wl, u64 xl, u64 yl, u32 wm, u32 xm, u32 ym, u32 r, + u32 n, u32 c, u32 rows, bool q) { + if (gpu_buf == nil || gpu_lib == nil) { + return false; + } + @autoreleasepool { + u32 p[8] = { wm, xm, ym, c, rows, 0, r, n }; + if (q) { + pthread_mutex_lock(&gpu_qlock); + gpu_qopen(); + } else { + gpu_qinit(); + } + id cb = q ? gpu_qcb : [gpu_que commandBuffer]; + id enc = q ? gpu_qenc : [cb computeCommandEncoder]; + [enc setComputePipelineState:gpu_q4pso]; + [enc setBuffer:gpu_buf offset:wl * 8 atIndex:0]; + [enc setBuffer:gpu_buf offset:xl * 8 atIndex:1]; + [enc setBuffer:gpu_buf offset:yl * 8 atIndex:2]; + [enc setBytes:p length:sizeof p atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((n + 8 * Q4R - 1) / (8 * Q4R), 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + if (q) { + pthread_mutex_unlock(&gpu_qlock); + return true; + } + [enc endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + if ([cb error]) { + gpu_fail([cb error]); + } + } + return true; +} + #elif BEND_CUDA static u64 gpu_hash(void) { diff --git a/tests/base/array_atq.bend b/tests/base/array_atq.bend new file mode 100644 index 0000000000..7b3e771afd --- /dev/null +++ b/tests/base/array_atq.bend @@ -0,0 +1,73 @@ +# Array.atq: 4 query heads over 2 KV heads of 64, rotating 16 pairs, at +# position 5 over random caches; o (256 values, then the 4 group sums) and +# the caches' new row as F32 bits. +import Base + +def bits(+v: F32) -> U32: + match v: + case F32{+w}: + U32{w} + +def f32(+u: U32) -> F32: + match u: + case U32{+w}: + F32{w} + +def pat(+i: U32, +s: U32) -> F32: + +h = (i * 2654435761 + s : U32) + f32(((1040187392 + h % 25165824) .|. (((h >> 7n) .&. 1) << 31n) : U32)) + +def fill(n: Nat, a: Array, +i: U32, +s: U32, +k: F32) -> Array: + match n: + case 0n: + a + case 1n+p: + fill(p, Array.set(F32, a, i, (pat(i, s) * k : F32)), (i + 1 : U32), s, k) + +def ubits(a: Array) -> Array: + Array.map(~F32, ~U32, ~(v => bits(v)), a) + +def take.k(r: Array & U32, acc: List) -> Array & List: + (a, +x) = r + (a, x <> acc) + +# n slots of a ending at j, in order, onto acc. +def take(n: Nat, r: Array & List, +j: U32) -> Array & List: + match n: + case 0n: + r + case 1n+p: + (a, acc) = r + take(p, take.k(Array.get(U32, a, j), acc), (j - 1 : U32)) + +def show(r: Array & List) -> String: + (a, xs) = r + _ = a + List.show(~&1, ~U32, ~(v => U32.show(v)), xs) + +def fbits(+n: Nat, a: Array, +i: U32) -> String: + show(take(n, (ubits(a), Nil{}), (i + U32.from_nat(n) - 1 : U32))) + +def ubits.show(+n: Nat, a: Array, +i: U32) -> String: + show(take(n, (a, Nil{}), (i + U32.from_nat(n) - 1 : U32))) + +def done(o: Array, y: Array, kc: Array, vc: Array, kw: Array, cs: Array) -> IO(Unit): + _ = y + _ = kw + _ = cs + do IO: + IO.print(fbits(260n, o, 0)) + IO.print(ubits.show(128n, kc, 640)) + IO.print(ubits.show(128n, vc, 640)) + +def out(q: ((((Array & Array) & Array) & Array) & Array) & Array) -> IO(Unit): + (((((y, kc), vc), kw), cs), o) = q + done(Array.q4wait(o), y, kc, vc, kw, cs) + +def main() -> IO(Unit): + out(Array.atq(fill(768n, Array.new(F32, 10n, 0.0), 0, 5, 1.0), ubits(fill(1024n, Array.new(F32, 10n, 0.0), 0, 7, 1.0)), + ubits(fill(1024n, Array.new(F32, 10n, 0.0), 0, 9, 1.0)), ubits(fill(128n, Array.new(F32, 7n, 0.0), 0, 11, 1.0)), + fill(32n, Array.new(F32, 5n, 0.0), 0, 13, 0.5), Array.new(F32, 9n, 0.0), 5, 4, 2, 64, 16, 0.000001)) +#|[1049165824, 3179675648, 958332928, 1034158080, 3189702656, 929169408, 1045561344, 3191275520, 1009385472, 1053687808, 3195142144, 3182559232, 969015296, 1050345472, 3174432768, 995557376, 1050673152, 3171155968, 980418560, 1041498112, 3188916224, 1008402432, 1045889024, 3190816768, 3153657856, 1043922944, 1047789568, 3178758144, 3115253760, 1047920640, 3163488256, 993067008, 1047396352, 3146055680, 3135242240, 1041694720, 3184984064, 3096838144, 1042415616, 1043726336, 1012269056, 1047003136, 1043398656, 3175284736, 1035599872, 1044578304, 985464832, 3137929216, 1047396352, 3141402624, 3175022592, 1040252928, 3176595456, 3185704960, 1044512768, 1040646144, 3151691776, 1046609920, 1040711680, 1011810304, 3190947840, 1040515072, 3111649280, 3173384192, 1024589824, 1040908288, 3189506048, 3167682560, 3164471296, 3192324096, 1008205824, 3172925440, 1048576000, 1031798784, 3175415808, 1015873536, 3184197632, 1029701632, 1038352384, 3185508352, 1034027008, 1033175040, 3188850688, 993329152, 3168731136, 1045233664, 1021050880, 3172859904, 1044185088, 1039073280, 1033306112, 1023737856, 3182755840, 1024524288, 1032585216, 3184984064, 1027997696, 1038680064, 1040842752, 1015349248, 1024000000, 1041629184, 1032519680, 1031798784, 1044643840, 1038614528, 1021706240, 1026686976, 3166896128, 3155755008, 1037172736, 1037041664, 1031012352, 1041301504, 1032716288, 1026752512, 1035206656, 990052352, 1035403264, 1019740160, 1041367040, 1041301504, 3161063424, 1039400960, 3173646336, 3167485952, 1039204352, 1029701632, 3164405760, 1045495808, 1035010048, 1011351552, 3198746624, 1026228224, 1019215872, 3175809024, 1044250624, 3190751232, 3180920832, 1040515072, 995295232, 3188850688, 1042743296, 1033371648, 1013776384, 3196387328, 1026359296, 1011286016, 3198550016, 1018560512, 1017708544, 3176660992, 1038024704, 3162243072, 3186163712, 1040252928, 3150512128, 3189571584, 3191996416, 1026228224, 3191996416, 3196190720, 1023410176, 1011810304, 3196583936, 3160145920, 3156541440, 3181576192, 1032519680, 3135635456, 3186884608, 3188981760, 1005191168, 3197173760, 3188719616, 1023606784, 3196125184, 3189637120, 3160080384, 3154771968, 3195928576, 3161063424, 991559680, 3196649472, 1030946816, 1031995392, 3193765888, 3188588544, 3139108864, 3196780544, 3184984064, 1016791040, 1034813440, 1037697024, 3172466688, 1037172736, 3175153664, 3189374976, 1040646144, 1025638400, 1032060928, 3191078912, 1016528896, 1029177344, 3192651776, 3161194496, 1033633792, 3166896128, 1035075584, 3175350272, 3185442816, 1036779520, 3177709568, 3194290176, 1040252928, 1016856576, 1027211264, 3192717312, 3163684864, 1028784128, 3193176064, 3171614720, 3172925440, 3180789760, 1029767168, 3177185280, 3186425856, 1036386304, 3172270080, 3194486784, 3189047296, 3156934656, 3179216896, 3189374976, 3171287040, 3171090432, 3189047296, 3175284736, 1001521152, 3181707264, 3156803584, 1016135680, 3191078912, 3185639424, 3174825984, 3193962496, 3185246208, 3175153664, 1044971520, 3173646336, 3174105088, 3171024896, 3188981760, 3175350272, 1018560512, 3186491392, 1079626912, 1076530432, 3227461312, 3223059712] +#|[3188064256, 3184066560, 3147038720, 1005584384, 1024655360, 1037303808, 1044185088, 1003225088, 3131441152, 3180789760, 1045168128, 3137470464, 3132620800, 1039335424, 3173646336, 3140616192, 1048051712, 3171811328, 3181182976, 1035665408, 1035599872, 3170959360, 1000472576, 1037107200, 3163291648, 3215327232, 3152478208, 3171024896, 3163881472, 1050017792, 3160408064, 3172401152, 1060241408, 3181182976, 1050083328, 1040187392, 3187408896, 1054277632, 3209625600, 3191406592, 1034289152, 3206414336, 1048969216, 3187933184, 3210674176, 1054474240, 3200319488, 1043333120, 1059323904, 3205300224, 1049231360, 3202875392, 3185901568, 1053294592, 3199860736, 1043005440, 3180593152, 3205562368, 1047789568, 3178299392, 1062207488, 1053097984, 3199336448, 1042087936, 3178102784, 3146973184, 998703104, 1017905152, 3170959360, 1028063232, 1059127296, 3149594624, 3122135040, 3214082048, 3190489088, 3127640064, 1006698496, 1041367040, 3135373312, 992149504, 1038024704, 3176660992, 1032454144, 1050607616, 3181510656, 3196649472, 1016856576, 3183149056, 3152805888, 3168731136, 998572032, 3160735744, 1040384000, 1043202048, 3167223808, 1023344640, 1050542080, 1032060928, 1048838144, 1054932992, 1037631488, 1069678592, 1059979264, 1042612224, 1049231360, 1057619968, 1038680064, 1029963776, 1061748736, 1043922944, 1051459584, 1041367040, 1049690112, 1055981568, 1064304640, 1053884416, 1036189696, 1068695552, 1050345472, 1033306112, 3195731968, 1056702464, 1037631488, 1028521984, 1060765696, 1042612224, 1050542080, 1057226752] +#|[3199220357, 1046983734, 3206491623, 3201738648, 1049502025, 3209009914, 3204256939, 1043631708, 1064044557, 3206775230, 1046149999, 1041397024, 3200904913, 3196151938, 1043915315, 3203423204, 3198670229, 1046433606, 3205941495, 3201188520, 1040563289, 1060976138, 3203706811, 1043081580, 1063494429, 3197836494, 1045599871, 1040846896, 3200354785, 3195601810, 1043365187, 3202873076, 3198120101, 1062660694, 1057907719, 3200638392, 1065178985, 1060426010, 3203156683, 1042531452, 1062944301, 3197286366, 3192533391, 1040296768, 3199804657, 3195051682, 1059592275, 1054839300, 3197569973, 1062110566, 1057357591, 3200088264, 1064628857, 1059875882, 3194217947, 3189464972, 1062394173, 3196736238, 3191983263, 1064912464, 1051770881, 3194501554, 1059042147, 1054289172, 3197019845, 1061560438, 1056807463, 3191149528, 3211562377, 1059325754, 3193667819, 3188914844, 1061844045, 3196186110, 3191433135, 1055973728, 1051220753, 3193951426, 1058492019, 1053739044, 3188081109, 3208493958, 1056257335, 3190599400, 3211012249, 1058775626, 3193117691, 3188364716, 1052905309, 1048152334, 3190883007, 1055423600, 1050670625, 3193401298, 3205425539, 1053188916, 3212696805, 3207943830, 1055707207, 3190049272, 3210462121, 1049836890, 1045083915, 3187814588, 1052355181, 1047602206, 3190332879, 3202357120, 1050120497, 3209628386, 3204875411, 1052638788, 3212146677, 3207393702, 1046768471, 1042015496, 3209911993, 1049286762, 1044533787, 3212430284, 1051805053, 1047052078, 3206559967, 3201806992, 1049570369, 3209078258, 3204325283, 1052088660] diff --git a/tests/base/array_gdn.bend b/tests/base/array_gdn.bend new file mode 100644 index 0000000000..79f3dca958 --- /dev/null +++ b/tests/base/array_gdn.bend @@ -0,0 +1,76 @@ +# Array.gdn: two value heads of 64 rows over key heads of 32 (c = 256, +# din = 128); o (128 values, then the 2 group sums), 64 state values (rounded +# to bfloat16: Metal's sums move their last bits) and 64 window values as F32 +# bits. a and dt are halved, the state quartered. +import Base + +def bits(+v: F32) -> U32: + match v: + case F32{+w}: + U32{w} + +def f32(+u: U32) -> F32: + match u: + case U32{+w}: + F32{w} + +def pat(+i: U32, +s: U32) -> F32: + +h = (i * 2654435761 + s : U32) + f32(((1040187392 + h % 25165824) .|. (((h >> 7n) .&. 1) << 31n) : U32)) + +def fill(n: Nat, a: Array, +i: U32, +s: U32, +k: F32) -> Array: + match n: + case 0n: + a + case 1n+p: + fill(p, Array.set(F32, a, i, (pat(i, s) * k : F32)), (i + 1 : U32), s, k) + +def ubits(a: Array) -> Array: + Array.map(~F32, ~U32, ~(v => bits(v)), a) + +def take.k(r: Array & U32, acc: List) -> Array & List: + (a, +x) = r + (a, x <> acc) + +# n slots of a ending at j, in order, onto acc. +def take(n: Nat, r: Array & List, +j: U32) -> Array & List: + match n: + case 0n: + r + case 1n+p: + (a, acc) = r + take(p, take.k(Array.get(U32, a, j), acc), (j - 1 : U32)) + +def show(r: Array & List) -> String: + (a, xs) = r + _ = a + List.show(~&1, ~U32, ~(v => U32.show(v)), xs) + +def fbits(+n: Nat, a: Array, +i: U32) -> String: + show(take(n, (ubits(a), Nil{}), (i + U32.from_nat(n) - 1 : U32))) + +def ubits.show(+n: Nat, a: Array, +i: U32) -> String: + show(take(n, (a, Nil{}), (i + U32.from_nat(n) - 1 : U32))) + +def done(o: Array, y: Array, w: Array, s: Array, cw: Array, ab: Array, sn: Array) -> IO(Unit): + _ = y + _ = cw + _ = ab + _ = sn + do IO: + IO.print(fbits(130n, o, 0)) + IO.print(fbits(64n, Array.map(~F32, ~F32, ~(v => Array.q4mv.round(v)), s), 1000)) + IO.print(fbits(64n, w, 200)) + +def out(q: (((((Array & Array) & Array) & Array) & Array) & Array) & Array) -> IO(Unit): + ((((((y, w), s), cw), ab), sn), o) = q + done(Array.q4wait(o), y, w, s, cw, ab, sn) + +def main() -> IO(Unit): + out(Array.gdn(fill(388n, Array.new(F32, 9n, 0.0), 0, 5, 1.0), fill(768n, Array.new(F32, 10n, 0.0), 0, 7, 1.0), + fill(4096n, Array.new(F32, 12n, 0.0), 0, 9, 0.25), ubits(fill(1024n, Array.new(F32, 10n, 0.0), 0, 11, 0.5)), + ubits(fill(4n, Array.new(F32, 2n, 0.0), 0, 13, 0.5)), ubits(fill(64n, Array.new(F32, 6n, 0.0), 0, 17, 1.0)), + Array.new(F32, 8n, 0.0), 256, 32, 64, 2, 128, 64, 0.000001)) +#|[3170238464, 3168468992, 3163750400, 3172532224, 1015742464, 3167223808, 1057882112, 1035403264, 3165847552, 3204186112, 3184984064, 3165126656, 994836480, 3180265472, 1023475712, 1024065536, 3191668736, 3174498304, 1021050880, 3185508352, 1020526592, 3153657856, 1045692416, 1040384000, 3153592320, 3205758976, 1040384000, 3159359488, 1001586688, 3168731136, 1012269056, 3163553792, 3179479040, 3167027200, 3164864512, 3187736576, 1017643008, 984743936, 1037959168, 3179544576, 3148349440, 3199926272, 3185704960, 3155886080, 991166464, 1021378560, 996081664, 1010368512, 3147563008, 3163750400, 3186819072, 3185836032, 3171352576, 1024000000, 1025966080, 1018167296, 998440960, 3178627072, 3185770496, 3160080384, 1023606784, 1011548160, 1007747072, 1017380864, 3169386496, 3173777408, 3156606976, 3192913920, 3172925440, 3172401152, 1032716288, 3184656384, 3158966272, 3200712704, 3175022592, 3153330176, 3164340224, 3164798976, 1011286016, 3160670208, 3191013376, 1013841920, 1016266752, 1010565120, 3153199104, 1015414784, 1043136512, 3182100480, 3161653248, 1042677760, 1037238272, 3125346304, 3168337920, 1016856576, 3144548352, 1028849664, 3178168320, 1025769472, 1022099456, 1025835008, 3180068864, 3130327040, 3171549184, 1030815744, 3150381056, 1025638400, 1013776384, 1000341504, 3163160576, 1007157248, 3149660160, 1025114112, 3160670208, 990052352, 1015087104, 1020985344, 1023213568, 1042087936, 1022099456, 1024393216, 3158573056, 3176333312, 3152216064, 1002831872, 1015414784, 1032650752, 994312192, 1037631488, 3216881920, 3188944640] +#|[1043922944, 1030750208, 3173580800, 3185573888, 1033240576, 3176071168, 3187933184, 1035665408, 3195404288, 3190554624, 1038221312, 1025245184, 3193110528, 1040646144, 1027866624, 3195666432, 3182428160, 1030029312, 3172859904, 3185049600, 1032585216, 3175350272, 3187671040, 1035141120, 1047265280, 3189964800, 1037500416, 1024065536, 3192717312, 1030946816, 1026490368, 3195207680, 3181838336, 1029373952, 3172794368, 3184394240, 1032060928, 1044185088, 3186950144, 1034289152, 1046740992, 3189440512, 1027932160, 1024393216, 3191996416, 3179610112, 1027014656, 3194552320, 3181248512, 1028390912, 1041039360, 3183935488, 1031274496, 1043726336, 3186622464, 1033895936, 1047527424, 3188457472, 3166699520, 1049362432, 3189833728, 3161194496, 1032519680, 1039794176] +#|[1053868879, 1049115904, 3191846577, 1056387170, 1051634195, 3211142084, 1058905461, 1054152486, 3188494551, 3208907400, 1056670777, 3191012842, 3211425691, 1050800460, 3193531133, 3188778158, 1053318751, 1048565776, 3208073665, 1055837042, 1051084067, 3210591956, 3205838981, 1053602358, 3187944423, 3208357272, 1047732041, 3190462714, 3210875563, 1050250332, 1045497357, 3205005246, 1052768623, 1048015648, 3207523537, 3202770562, 1050533939, 3210041828, 3205288853, 1044663622, 3212560119, 3207807144, 1047181913, 1042428938, 3210325435, 1049700204, 1044947229, 3204455118, 1052218495, 1047465520, 3206973409, 3202220434, 1041595203, 3209491700, 3204738725, 1044113494, 1064526343, 3207257016, 1046631785, 1041878810, 3201386699, 1049150076, 1044397101, 3203904990] diff --git a/tests/base/array_q4go.bend b/tests/base/array_q4go.bend new file mode 100644 index 0000000000..89fb10fa94 --- /dev/null +++ b/tests/base/array_q4go.bend @@ -0,0 +1,55 @@ +# Array.q4go + Array.q4wait: the 13x128 case of array_q4mv.bend through the +# queue (rows 0 .. 12, then rows 3 .. 8 alone). w and y stay live until after +# the wait, because a queued product reads w and writes y until then. Slots it +# does not write keep 7.0 (1088421888). Without a Metal heap, q4go is q4mv and +# q4wait is the identity. +import Base + +def bits(+v: F32) -> U32: + match v: + case F32{+w}: + U32{w} + +def wset(n: Nat, a: Array, +i: U32) -> Array: + match n: + case 0n: + a + case 1n+p: + +s = (15360 + (i % 8) * 17 + ((15488 + (i % 5) * 19) << 16n) : U32) + +b = (48128 + (i % 6) * 23 + ((15232 + (i % 7) * 13) << 16n) : U32) + wset(p, Array.set(U32, Array.set(U32, a, (208 + i : U32), s), (221 + i : U32), b), (i + 1 : U32)) + +def cset(n: Nat, a: Array, +i: U32) -> Array: + match n: + case 0n: + wset(13n, a, 0) + case 1n+p: + cset(p, Array.set(U32, a, i, (i * 2654435761 + 12345 : U32)), (i + 1 : U32)) + +def xset(n: Nat, a: Array, +j: U32) -> Array: + match n: + case 0n: + Array.set(F32, Array.set(F32, a, 128, 1.5), 129, 2.5) + case 1n+p: + xset(p, Array.set(F32, a, j, ((U32.to_f32((j % 7 : U32)) - 3.0) * 0.25 : F32)), (j + 1 : U32)) + +def show.fin(w: Array, s: String) -> String: + _ = w + s + +# w is dropped after y is read: a queued product reads w until the wait. +def show.k(w: Array, x: Array, y: Array) -> String: + _ = x + show.fin(w, List.show(~&1, ~U32, ~(v => U32.show(v)), Array.to_list(~U32, Array.map(~F32, ~U32, ~(v => bits(v)), y)))) + +def show(q: (Array & Array) & Array) -> String: + ((w, x), y) = q + show.k(w, Array.q4wait(x), y) + +def main() -> IO(Unit): + do IO: + IO.print(show(Array.q4go(cset(208n, Array.new(U32, 8n, 0), 0), xset(128n, Array.new(F32, 8n, 0.0), 0), Array.new(F32, 4n, 7.0), 0, 13, 128, 13))) + IO.print(show(Array.q4go(cset(208n, Array.new(U32, 8n, 0), 0), xset(128n, Array.new(F32, 8n, 0.0), 0), Array.new(F32, 4n, 7.0), 3, 6, 128, 13))) + +#|[3190161408, 1045233664, 1059061760, 3189637120, 3178823680, 1039532032, 1043202048, 3171614720, 3194748928, 3205890048, 1055129600, 1045954560, 3183542272, 1088421888, 1088421888, 1088421888] +#|[1088421888, 1088421888, 1088421888, 3189637120, 3178823680, 1039532032, 1043202048, 3171614720, 3194748928, 1088421888, 1088421888, 1088421888, 1088421888, 1088421888, 1088421888, 1088421888] diff --git a/tests/base/array_q4mv.bend b/tests/base/array_q4mv.bend new file mode 100644 index 0000000000..e84ae99e3e --- /dev/null +++ b/tests/base/array_q4mv.bend @@ -0,0 +1,44 @@ +# Array.q4mv: 13 rows of 128 weights (one block of eight rows and a tail +# of five), then rows 3 .. 8 alone; outputs as F32 bits, slots it does not +# write keep 7.0 (1088421888) +import Base + +def bits(+v: F32) -> U32: + match v: + case F32{+w}: + U32{w} + +def wset(n: Nat, a: Array, +i: U32) -> Array: + match n: + case 0n: + a + case 1n+p: + +s = (15360 + (i % 8) * 17 + ((15488 + (i % 5) * 19) << 16n) : U32) + +b = (48128 + (i % 6) * 23 + ((15232 + (i % 7) * 13) << 16n) : U32) + wset(p, Array.set(U32, Array.set(U32, a, (208 + i : U32), s), (221 + i : U32), b), (i + 1 : U32)) + +def cset(n: Nat, a: Array, +i: U32) -> Array: + match n: + case 0n: + wset(13n, a, 0) + case 1n+p: + cset(p, Array.set(U32, a, i, (i * 2654435761 + 12345 : U32)), (i + 1 : U32)) + +def xset(n: Nat, a: Array, +j: U32) -> Array: + match n: + case 0n: + Array.set(F32, Array.set(F32, a, 128, 1.5), 129, 2.5) + case 1n+p: + xset(p, Array.set(F32, a, j, ((U32.to_f32((j % 7 : U32)) - 3.0) * 0.25 : F32)), (j + 1 : U32)) + +def show(q: (Array & Array) & Array) -> String: + ((w, x), y) = q + List.show(~&1, ~U32, ~(v => U32.show(v)), Array.to_list(~U32, Array.map(~F32, ~U32, ~(v => bits(v)), y))) + +def main() -> IO(Unit): + do IO: + IO.print(show(Array.q4mv(cset(208n, Array.new(U32, 8n, 0), 0), xset(128n, Array.new(F32, 8n, 0.0), 0), Array.new(F32, 4n, 7.0), 0, 13, 128, 13))) + IO.print(show(Array.q4mv(cset(208n, Array.new(U32, 8n, 0), 0), xset(128n, Array.new(F32, 8n, 0.0), 0), Array.new(F32, 4n, 7.0), 3, 6, 128, 13))) + +#|[3190161408, 1045233664, 1059061760, 3189637120, 3178823680, 1039532032, 1043202048, 3171614720, 3194748928, 3205890048, 1055129600, 1045954560, 3183542272, 1088421888, 1088421888, 1088421888] +#|[1088421888, 1088421888, 1088421888, 3189637120, 3178823680, 1039532032, 1043202048, 3171614720, 3194748928, 1088421888, 1088421888, 1088421888, 1088421888, 1088421888, 1088421888, 1088421888] diff --git a/tests/base/array_q4mv_tree.bend b/tests/base/array_q4mv_tree.bend new file mode 100644 index 0000000000..72d12ebce1 --- /dev/null +++ b/tests/base/array_q4mv_tree.bend @@ -0,0 +1,68 @@ +# Array.q4mv's summation tree: 5 rows of 2176 weights (34 groups, so groups +# 32 and 33 wrap onto partials 0 and 1), then rows 1 .. 4 alone. Lanes 0 and +# 4 of every group hold +-2^25 that cancel only when lane j sums its own +# nibble over the words before the lanes meet; x-sums of +-2^24 in groups 0 +# and 16 cancel only in the fold p[l] += p[l + 16]. Summed in index order +# (or with another lane pairing, partial count or fold) every row differs. +import Base + +def bits(+v: F32) -> U32: + match v: + case F32{+w}: + U32{w} + +def wset(n: Nat, a: Array, +i: U32) -> Array: + match n: + case 0n: + a + case 1n+p: + +s = (15360 + (i % 8) * 17 + ((15488 + (i % 5) * 19) << 16n) : U32) + wset(p, Array.set(U32, Array.set(U32, a, (1360 + i : U32), s), (1445 + i : U32), 1065369472), (i + 1 : U32)) + +def cset(n: Nat, a: Array, +i: U32) -> Array: + match n: + case 0n: + wset(85n, a, 0) + case 1n+p: + cset(p, Array.set(U32, a, i, ((1 + i % 3) * 286331153 : U32)), (i + 1 : U32)) + +def pick(b: Bool, +t: F32, +e: F32) -> F32: + match b: + case True{}: + t + case False{}: + e + +def xv(+j: U32) -> F32: + pick(U32.is_eq((j % 8 : U32), 0), 33554432.0, + pick(U32.is_eq((j % 8 : U32), 4), (0.0 - 33554432.0 : F32), (U32.to_f32((j % 5 : U32)) + 0.5 : F32))) + +def sv(+g: U32) -> F32: + pick(U32.is_eq(g, 0), 16777216.0, pick(U32.is_eq(g, 16), (0.0 - 16777216.0 : F32), pick(U32.is_eq(g, 32), 3.0, + pick(U32.is_eq(g, 1), 2.0, (U32.to_f32((g % 5 : U32)) + 0.25 : F32))))) + +def sset(n: Nat, a: Array, +g: U32) -> Array: + match n: + case 0n: + a + case 1n+p: + sset(p, Array.set(F32, a, (2176 + g : U32), sv(g)), (g + 1 : U32)) + +def xset(n: Nat, a: Array, +j: U32) -> Array: + match n: + case 0n: + sset(34n, a, 0) + case 1n+p: + xset(p, Array.set(F32, a, j, xv(j)), (j + 1 : U32)) + +def show(q: (Array & Array) & Array) -> String: + ((w, x), y) = q + List.show(~&1, ~U32, ~(v => U32.show(v)), Array.to_list(~U32, Array.map(~F32, ~U32, ~(v => bits(v)), y))) + +def main() -> IO(Unit): + do IO: + IO.print(show(Array.q4mv(cset(1360n, Array.new(U32, 11n, 0), 0), xset(2176n, Array.new(F32, 12n, 0.0), 0), Array.new(F32, 3n, 7.0), 0, 5, 2176, 5))) + IO.print(show(Array.q4mv(cset(1360n, Array.new(U32, 11n, 0), 0), xset(2176n, Array.new(F32, 12n, 0.0), 0), Array.new(F32, 3n, 7.0), 1, 4, 2176, 5))) + +#|[1128857600, 1128988672, 1128988672, 1128988672, 1129185280, 1088421888, 1088421888, 1088421888] +#|[1088421888, 1128988672, 1128988672, 1128988672, 1129185280, 1088421888, 1088421888, 1088421888] diff --git a/tests/base/array_rmsq.bend b/tests/base/array_rmsq.bend new file mode 100644 index 0000000000..52ac78f374 --- /dev/null +++ b/tests/base/array_rmsq.bend @@ -0,0 +1,57 @@ +# Array.rmsq: 128 values (two groups) with the residual add, then without; +# o (128 values, the 2 group sums, then untouched zeros) and x as F32 bits. Inputs are F32 +# bit patterns in +-[0.125, 2), the weights as F32 bits; eps is the F32 nearest 1e-6. +import Base + +def bits(+v: F32) -> U32: + match v: + case F32{+w}: + U32{w} + +def f32(+u: U32) -> F32: + match u: + case U32{+w}: + F32{w} + +def pat(+i: U32, +s: U32) -> F32: + +h = (i * 2654435761 + s : U32) + f32(((1040187392 + h % 25165824) .|. (((h >> 7n) .&. 1) << 31n) : U32)) + +def fill(n: Nat, a: Array, +i: U32, +s: U32) -> Array: + match n: + case 0n: + a + case 1n+p: + fill(p, Array.set(F32, a, i, pat(i, s)), (i + 1 : U32), s) + +def show(a: Array) -> String: + List.show(~&1, ~U32, ~(v => U32.show(v)), Array.to_list(~U32, Array.map(~F32, ~U32, ~(v => bits(v)), a))) + +# y and w stay live until the wait: the GPU reads them until then. +def fin(o: Array, x: Array, y: Array, w: Array) -> Array & Array: + _ = y + _ = w + (o, x) + +def put(r: Array & Array) -> IO(Unit): + (o, x) = r + do IO: + IO.print(show(o)) + IO.print(show(x)) + +def out(q: ((Array & Array) & Array) & Array) -> IO(Unit): + (((x, y), w), o) = q + put(fin(Array.q4wait(o), x, y, w)) + +def one(+add: U32) -> IO(Unit): + out(Array.rmsq(fill(128n, Array.new(F32, 7n, 0.0), 0, 11), fill(128n, Array.new(F32, 7n, 0.0), 0, 23), + Array.map(~F32, ~U32, ~(v => bits(v)), fill(128n, Array.new(F32, 7n, 0.0), 0, 37)), Array.new(F32, 8n, 0.0), 128, f32(897988541), add)) + +def main() -> IO(Unit): + do IO: + one(1) + one(0) +#|[1023868928, 1049231360, 3186360320, 1063518208, 1053491200, 1043857408, 1068302336, 1058013184, 1032978432, 1072955392, 889651200, 1037238272, 1027801088, 1068761088, 1041891328, 3180855296, 1056768000, 1048379392, 859242496, 1062076416, 1052508160, 1025769472, 1067384832, 3204644864, 1032060928, 1071906816, 1062797312, 1036255232, 3173908480, 1051066368, 1040646144, 856817664, 1055653888, 1046872064, 1036713984, 1060700160, 3199008768, 1024524288, 1066532864, 1056243712, 1030815744, 1070792704, 1044643840, 1035272192, 853082112, 1050148864, 1040056320, 1031602176, 1054605312, 3192782848, 1069350912, 1059323904, 874250240, 1073938432, 1065680896, 1038352384, 1029373952, 892141568, 1043267584, 1034354688, 1024000000, 1049362432, 3186360320, 1063649280, 1053556736, 871170048, 1068367872, 1058144256, 1049755648, 1073020928, 3211919360, 1037303808, 1027866624, 1068892160, 1041956864, 3180920832, 1056833536, 1048510464, 859242496, 1062207488, 1052573696, 1042677760, 1067450368, 3204644864, 1032192000, 1071906816, 1062928384, 1036255232, 1026555904, 1067909120, 1040777216, 856817664, 1055719424, 1046872064, 1036845056, 1060765696, 3199139840, 1024589824, 1066598400, 876675072, 1030946816, 1070923776, 1061421056, 1035337728, 3172663296, 1050279936, 1040056320, 1031733248, 1054605312, 3192913920, 1069350912, 1059454976, 874250240, 1074003968, 1065746432, 1055195136, 1029373952, 3217358848, 1043333120, 1034420224, 1024000000, 1049427968, 1039007744, 1030160384, 1053622272, 871235584, 1068433408, 1058144256, 1103782272, 1103740416, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] +#|[1048576000, 3208052736, 1055850496, 1067843584, 3210608640, 3205824512, 1070399488, 3213099008, 3199991808, 1072889856, 893386752, 3202482176, 1050279936, 1070661632, 3205038080, 1052770304, 1064828928, 3207528448, 3024093184, 1067319296, 3210084352, 3196911616, 1069809664, 3212574720, 3199401984, 1072365568, 1067581440, 3201957888, 1049690112, 1061748736, 3204448256, 3024093184, 1064239104, 3207004160, 3202220032, 1066795008, 3209494528, 3196387328, 1069285376, 1064566784, 3198877696, 1071775744, 1058668544, 3201368064, 3024093184, 1061158912, 3203923968, 3199139840, 1063714816, 3206414336, 3218472960, 1066205184, 884998144, 3220963328, 1068761088, 1055588352, 3198353408, 3040870400, 1058144256, 3200843776, 3196125184, 1060634624, 3203334144, 3215392768, 1063124992, 884998144, 3217883136, 1065680896, 1060896768, 3220439040, 1068171264, 1055064064, 3197763584, 3218210816, 1057554432, 3200319488, 3212312576, 1060110336, 876609536, 3214868480, 1062600704, 1057882112, 3217358848, 1065091072, 1051983872, 3219849216, 3215130624, 1054474240, 3197239296, 3217620992, 1057030144, 876609536, 3211788288, 1059520512, 1054801920, 3214278656, 1062076416, 1048903680, 3216834560, 3032481792, 1051459584, 3219324928, 3214540800, 1053949952, 3196649472, 3208708096, 1056440320, 1051721728, 3211198464, 1058996224, 1070989312, 3213754368, 3032481792, 1073545216, 3216244736, 3211526144, 1050869760, 3218800640, 3205627904, 1053425664, 1048641536, 3208183808, 1055916032, 1051131904, 3210674176, 3032481792, 1070465024, 3213164544] +#|[1023606784, 1048969216, 3185770496, 1063059456, 1053032448, 1043464192, 1067909120, 1057751040, 1032650752, 1072431104, 3211264000, 1036779520, 1027342336, 1068367872, 1041563648, 3180527616, 1056178176, 1047789568, 3184787456, 1061617664, 1052114944, 1025376256, 1067057152, 3204251648, 1031798784, 1071382528, 1062338560, 1035796480, 3173515264, 1050738688, 1040384000, 3179741184, 1055129600, 1046282240, 1036320768, 1060241408, 3198681088, 1024196608, 1066205184, 1055653888, 1030291456, 1070333952, 1044185088, 1034878976, 3172270080, 1049821184, 1039466496, 1031077888, 1054081024, 3192389632, 1068892160, 1058996224, 3197763584, 1073545216, 1065418752, 1037828096, 1028849664, 3216900096, 1042874368, 1033961472, 1023672320, 1049034752, 3185836032, 1063124992, 1053097984, 3191013376, 1067974656, 1057816576, 1049427968, 1072496640, 3211395072, 1036845056, 1027473408, 1068433408, 1041629184, 3180593152, 1056243712, 1047920640, 3184852992, 1061683200, 1052180480, 1042284544, 1067122688, 3204317184, 1031864320, 1071448064, 1062404096, 1035862016, 1026097152, 1067515904, 1040449536, 3179741184, 1055195136, 1046413312, 1036386304, 1060372480, 3198681088, 1024262144, 1066270720, 3203203072, 1030356992, 1070399488, 1061027840, 1034944512, 3172401152, 1049886720, 1039532032, 1031143424, 1054146560, 3192455168, 1068957696, 1059061760, 3197829120, 1073610752, 1065418752, 1054670848, 1028915200, 3216900096, 1042939904, 1034027008, 1023737856, 1049100288, 1038417920, 1029636096, 1053163520, 3191078912, 1068040192, 1057882112, 1101816192, 1102610304, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] +#|[1040187403, 3199695292, 1047458669, 1059482910, 3202213583, 3197460608, 1062001201, 3204731874, 3191590291, 1064519492, 3207250165, 3194108582, 1041871959, 1062284808, 3196626873, 1044390250, 1056414491, 3199145164, 1046908541, 1058932782, 3201663455, 3188521872, 1061451073, 3204181746, 3191040163, 1063969364, 1059216389, 3193558454, 1041321831, 1053346072, 3196076745, 1043840122, 1055864363, 3198595036, 3193842061, 1058382654, 3201113327, 3187971744, 1060900945, 1056147970, 3190490035, 1063419236, 1050277653, 3193008326, 1040771703, 1052795944, 3195526617, 3190773642, 1055314235, 3198044908, 3210069149, 1057832526, 3200563199, 3212587440, 1060350817, 1047209234, 3189939907, 1062869108, 1049727525, 3192458198, 3187705223, 1052245816, 3194976489, 3207000730, 1054764107, 3197494780, 3209519021, 1057282398, 1052529423, 3212037312, 1059800689, 1046659106, 3189389779, 3209802628, 1049177397, 3191908070, 3203932311, 1051695688, 3194426361, 3206450602, 1054213979, 1049461004, 3208968893, 1056732270, 1043590687, 3211487184, 3206734209, 1046108978, 3188839651, 3209252500, 1048627269, 3191357942, 3203382183, 1051145560, 1046392585, 3205900474, 1053663851, 1040522268, 3208418765, 1056182142, 1043040559, 3210937056, 3206184081, 1045558850, 3188289523, 3200313764, 1048077141, 1043324166, 3202832055, 1050595432, 1062619673, 3205350346, 1053113723, 1065137964, 3207868637, 3203115662, 1042490431, 3210386928, 3197245345, 1045008722, 1040255747, 3199763636, 1047527013, 1042774038, 3202281927, 1050045304, 1062069545, 3204800218] diff --git a/tests/base/array_swq.bend b/tests/base/array_swq.bend new file mode 100644 index 0000000000..a57a53ac94 --- /dev/null +++ b/tests/base/array_swq.bend @@ -0,0 +1,63 @@ +# Array.swq: 128 gate values then 128 up values (bits in +-[0.125, 2)); +# o (128 values, then the 2 group sums) as F32 bits. +import Base + +def bits(+v: F32) -> U32: + match v: + case F32{+w}: + U32{w} + +def f32(+u: U32) -> F32: + match u: + case U32{+w}: + F32{w} + +def pat(+i: U32, +s: U32) -> F32: + +h = (i * 2654435761 + s : U32) + f32(((1040187392 + h % 25165824) .|. (((h >> 7n) .&. 1) << 31n) : U32)) + +def fill(n: Nat, a: Array, +i: U32, +s: U32, +k: F32) -> Array: + match n: + case 0n: + a + case 1n+p: + fill(p, Array.set(F32, a, i, (pat(i, s) * k : F32)), (i + 1 : U32), s, k) + +def ubits(a: Array) -> Array: + Array.map(~F32, ~U32, ~(v => bits(v)), a) + +def take.k(r: Array & U32, acc: List) -> Array & List: + (a, +x) = r + (a, x <> acc) + +# n slots of a ending at j, in order, onto acc. +def take(n: Nat, r: Array & List, +j: U32) -> Array & List: + match n: + case 0n: + r + case 1n+p: + (a, acc) = r + take(p, take.k(Array.get(U32, a, j), acc), (j - 1 : U32)) + +def show(r: Array & List) -> String: + (a, xs) = r + _ = a + List.show(~&1, ~U32, ~(v => U32.show(v)), xs) + +def fbits(+n: Nat, a: Array, +i: U32) -> String: + show(take(n, (ubits(a), Nil{}), (i + U32.from_nat(n) - 1 : U32))) + +def ubits.show(+n: Nat, a: Array, +i: U32) -> String: + show(take(n, (a, Nil{}), (i + U32.from_nat(n) - 1 : U32))) + +def done(o: Array, a: Array) -> IO(Unit): + _ = a + IO.print(fbits(130n, o, 0)) + +def out(r: Array & Array) -> IO(Unit): + (a, o) = r + done(Array.q4wait(o), a) + +def main() -> IO(Unit): + out(Array.swq(fill(256n, Array.new(F32, 8n, 0.0), 0, 5, 1.0), Array.new(F32, 8n, 0.0), 128)) +#|[3167289344, 3188719616, 3182624768, 3183607808, 3167944704, 3188457472, 3189243904, 3172270080, 3181772800, 3195076608, 3176267776, 3178889216, 3171549184, 3189833728, 3182886912, 3176136704, 3176660992, 3188064256, 3181641728, 3182428160, 3166830592, 3176071168, 3188195328, 3171549184, 3172597760, 3193700352, 3175350272, 3177512960, 3170566144, 3178889216, 3181772800, 3175022592, 3175481344, 3187146752, 3178233856, 3181248512, 3165716480, 3166371840, 3187081216, 3170959360, 3171811328, 3192389632, 3172204544, 3176202240, 3168927744, 3194093568, 3180986368, 3172204544, 3174301696, 3185901568, 3188981760, 3180199936, 3164667904, 3183869952, 3185704960, 3165388800, 3171090432, 3191078912, 3188064256, 3174957056, 3165782016, 3192848384, 3180199936, 3184263168, 3173253120, 3184721920, 3179872256, 3179020288, 3193503744, 3183149056, 3184328704, 3181182976, 3169845248, 3189899264, 3186622464, 3173842944, 3179675648, 3191734272, 3179479040, 3174957056, 3172204544, 3187539968, 3179085824, 3177381888, 3182952448, 3182428160, 3183083520, 3180265472, 3168468992, 3179544576, 3184984064, 3173056512, 3170435072, 3190620160, 3180724224, 3174105088, 3171221504, 3176595456, 3178102784, 3176136704, 3173580800, 3181707264, 3174563840, 3179347968, 3167158272, 3172728832, 3183804416, 3174105088, 3169255424, 3189571584, 3198943232, 3173253120, 3169779712, 3196387328, 3177250816, 3169845248, 3172597760, 3180986368, 3167420416, 3177709568, 3165847552, 3188785152, 3182755840, 3173122048, 3168075776, 3188588544, 3189374976, 3172401152, 3232117760, 3231749632]