dtype.bend

bendtorch@1.3.0.0 · 658 lines · 113 definitions

1# BendTorch's dtypes (plan M11): BF16 and F16 stored two per U32, int8 with a scale per row stored four per U32, a
2# two-F32 ("double-word") accumulator, and the layers that give a model BF16 storage around F32 arithmetic.
3#
4# Every number is computed in F32. A 16-bit value lives in the low 16 bits of a U32 while it is handled alone, and two
5# of them share a word in storage (element 2k in the low half of word k, element 2k + 1 in the high half: the bytes of
6# a little-endian torch.bfloat16 / torch.float16 tensor). The conversions from F32 round to nearest, ties to even, and
7# treat NaN as PyTorch 2.14.1's CPU kernels do on this machine (tested bit for bit by reference/test_dtype.py):
8#   - F32 -> BF16: every NaN becomes 0xFFFF (the vectorized kernel's answer; c10's scalar code would give 0x7FC0).
9#   - F32 -> F16: a NaN keeps its sign and the top 9 bits of its payload, with the quiet bit set (F16C's vcvtps2ph).
10#   - F16 -> F32: a NaN gets the quiet bit set; subnormals become normal F32 numbers, exactly.
11# No @unsafe and no foreign code: plain Bend over U32 and F32 bit operations.
12
13import Base
14import bendtorch@1.3.0.0/main.bend as T
15import bend-ml-tensor-array@0.1.18.0/main.bend as TA
16
17# =====================================================================================================
18# Bits
19# =====================================================================================================
20
21# U32 and F32 hold the same Word(32n)
22def u2f(u: U32) -> F32:
23  match u:
24    case U32{w}:
25      F32{w}
26
27def f2u(x: F32) -> U32:
28  match x:
29    case F32{w}:
30      U32{w}
31
32def pick_u(b: Bool, +x: U32, +y: U32) -> U32:
33  match b:
34    case True{}:
35      x
36    case False{}:
37      y
38
39# =====================================================================================================
40# BF16: the top 16 bits of an F32
41# =====================================================================================================
42
43# F32 -> BF16, round to nearest even: add 0x7FFF plus the lowest kept bit, keep the top half. The carry rounds up into
44# the exponent when the mantissa overflows, and past the largest finite value to infinity, as IEEE asks.
45def bf16_rne(+u: U32) -> U32:
46  U32.shrn(((u + 32767 : U32) + (U32.shrn(u, 16n) .&. 1 : U32) : U32), 16n)
47
48def bf16_of_bits(+u: U32) -> U32:
49  pick_u(U32.is_gt((u .&. 2147483647 : U32), 2139095040), 65535, bf16_rne(u))
50
51def BF16.of_f32(x: F32) -> U32:
52  bf16_of_bits(f2u(x))
53
54# BF16 -> F32 is exact: the 16 bits become the top half
55def BF16.to_f32(+h: U32) -> F32:
56  u2f(U32.shln((h .&. 65535 : U32), 16n))
57
58# x rounded to the nearest BF16 value, as an F32 (x.to(torch.bfloat16).to(torch.float32))
59def BF16.round(+x: F32) -> F32:
60  BF16.to_f32(BF16.of_f32(x))
61
62# =====================================================================================================
63# F16: 1 sign bit, 5 exponent bits (bias 15), 10 mantissa bits
64# =====================================================================================================
65
66# a normal F16 (|x| >= 2^-14): rebias the exponent (127 - 15 = 112), round the 23-bit mantissa to 10 bits to nearest
67# even; past 65504 the carry reaches the infinity pattern 0x7C00 and anything above it is infinity too
68def f16_normal(+a: U32) -> U32:
69  +m = (a - 939524096 : U32)
70  +r = U32.shrn(((m + 4095 : U32) + (U32.shrn(m, 13n) .&. 1 : U32) : U32), 13n)
71  pick_u(U32.is_gt(r, 31744), 31744, r)
72
73# a subnormal F16 (2^-25 < |x| < 2^-14): the value in units of 2^-24 is mant / 2^sh, sh = 126 - e in 14..24, rounded to
74# nearest even; a result of 0x400 is the smallest normal, which is the right answer
75def f16_sub2(+q: U32, +rem: U32, +half: U32) -> U32:
76  +up = Bool.or(U32.is_gt(rem, half), Bool.and(U32.is_eq(rem, half), U32.is_eq((q .&. 1 : U32), 1)))
77  (q + Bool.to_u32(up) : U32)
78
79def f16_sub(+a: U32) -> U32:
80  +mant = ((a .&. 8388607 : U32) .|. 8388608 : U32)
81  +sh = U32.to_nat((126 - U32.shrn(a, 23n) : U32))
82  f16_sub2(U32.shrn(mant, sh), (mant .&. (U32.shln(1, sh) - 1 : U32) : U32), U32.shln(1, Nat.sub(sh, 1n)))
83
84# a NaN keeps the top 9 bits of its payload and gets the quiet bit
85def f16_nan(+a: U32) -> U32:
86  (32256 .|. U32.shrn((a .&. 8388607 : U32), 13n) : U32)
87
88def f16_mag(nan: Bool, normal: Bool, tiny: Bool, +a: U32) -> U32:
89  match nan normal tiny:
90    case True{} _ _:
91      f16_nan(a)
92    case False{} True{} _:
93      f16_normal(a)
94    case False{} False{} True{}:
95      0
96    case False{} False{} False{}:
97      f16_sub(a)
98
99# F32 -> F16, round to nearest even, with the sign kept (so -0.0 and the negative tiny numbers give 0x8000)
100def f16_of_bits(+u: U32) -> U32:
101  +a = (u .&. 2147483647 : U32)
102  ((U32.shrn(u, 16n) .&. 32768 : U32) .|. f16_mag(U32.is_gt(a, 2139095040), U32.is_ge(a, 947912704), U32.is_le(a, 855638016), a) : U32)
103
104def F16.of_f32(x: F32) -> U32:
105  f16_of_bits(f2u(x))
106
107# F16 -> F32, exact. A subnormal is m · 2^-24: the product of two F32 numbers that is exact and normal.
108def f16_up(zero: Bool, top: Bool, +m: U32, +s: U32, +eu: U32) -> F32:
109  match zero top:
110    case True{} _:
111      u2f((s .|. f2u((U32.to_f32(m) * 0.000000059604644775390625 : F32)) : U32))
112    case False{} True{}:
113      u2f((((s .|. 2139095040 : U32) .|. U32.shln(m, 13n) : U32) .|. pick_u(U32.is_eq(m, 0), 0, 4194304) : U32))
114    case False{} False{}:
115      u2f(((s .|. U32.shln((eu + 112 : U32), 23n) : U32) .|. U32.shln(m, 13n) : U32))
116
117def F16.to_f32(+h: U32) -> F32:
118  +eu = (U32.shrn(h, 10n) .&. 31 : U32)
119  f16_up(U32.is_eq(eu, 0), U32.is_eq(eu, 31), (h .&. 1023 : U32), U32.shln((h .&. 32768 : U32), 16n), eu)
120
121def F16.round(+x: F32) -> F32:
122  F16.to_f32(F16.of_f32(x))
123
124# =====================================================================================================
125# Two halves in a word
126# =====================================================================================================
127
128def half_join(+lo: U32, +hi: U32) -> U32:
129  ((lo .&. 65535 : U32) .|. U32.shln(hi, 16n) : U32)
130
131def half_lo(+w: U32) -> U32:
132  (w .&. 65535 : U32)
133
134def half_hi(+w: U32) -> U32:
135  U32.shrn(w, 16n)
136
137# a pair of 16-bit values, the unit of storage
138type HW is Data:
139  HW{lo: U32, hi: U32}
140
141# the values two by two; an odd last one is paired with 0
142def pairs(xs: List<&2, U32>) -> List<&2, HW>:
143  match xs:
144    case Nil{}:
145      Nil{}
146    case Con{a, Nil{}}:
147      [HW{a, 0}]
148    case Con{a, Con{b, t}}:
149      HW{a, b} <> pairs(t)
150
151# every value of a list of pairs, in order
152def unpairs(ws: List<&2, HW>) -> List<&2, U32>:
153  match ws:
154    case Nil{}:
155      Nil{}
156    case Con{HW{a, b}, t}:
157      a <> (b <> unpairs(t))
158
159# the first n of them: the padding of an odd count is left out
160def unpairs_n(+n: Nat, ws: List<&2, HW>) -> List<&2, U32>:
161  List.take(&2, U32, unpairs(ws), n)
162
163# =====================================================================================================
164# Packed 16-bit matrices: r·c values in ceil(r·c / 2) words, element 2k in the low half of word k
165# =====================================================================================================
166
167type HMat<-r: Nat, -c: Nat> is Type:
168  HMat{d: Array<U32>}
169
170type UL is Type:
171  UL{a: Array<U32>, l: List<&2, U32>}
172
173def fill_u(xs: List<&2, U32>, +i: U32, a: Array<U32>) -> Array<U32>:
174  match xs:
175    case Nil{}:
176      a
177    case Con{h, t}:
178      fill_u(t, (i + 1 : U32), Array.set(U32, a, i, h))
179
180# the k + 1 values from index i, in order (ra holds the array and the value at i, already read)
181def read_u(k: Nat, +i: U32, acc: List<&2, U32>, ra: Array<U32> & U32) -> UL:
182  match k ra:
183    case 0n Tuple{a, v}:
184      UL{a, List.reverse(&2, U32, v <> acc)}
185    case 1n+p Tuple{a, +v}:
186      read_u(p, (i + 1 : U32), v <> acc, Array.get(U32, a, (i + 1 : U32)))
187
188def ul_list(r: UL) -> List<&2, U32>:
189  match r:
190    case UL{a, l}:
191      l
192
193def enc_l(~f: F32 -> U32, xs: List<&2, F32>) -> List<&2, U32>:
194  match xs:
195    case Nil{}:
196      Nil{}
197    case Con{h, t}:
198      f(h) <> enc_l(~f, t)
199
200def dec_l(~f: @+h: U32 -> F32, xs: List<&2, U32>) -> List<&2, F32>:
201  match xs:
202    case Nil{}:
203      Nil{}
204    case Con{+h, t}:
205      f(h) <> dec_l(~f, t)
206
207def join_l(ws: List<&2, HW>) -> List<&2, U32>:
208  match ws:
209    case Nil{}:
210      Nil{}
211    case Con{HW{+a, +b}, t}:
212      half_join(a, b) <> join_l(t)
213
214def split_l(ws: List<&2, U32>) -> List<&2, HW>:
215  match ws:
216    case Nil{}:
217      Nil{}
218    case Con{+w, t}:
219      HW{half_lo(w), half_hi(w)} <> split_l(t)
220
221# the number of words for n halves
222def hwords(+n: Nat) -> Nat:
223  Nat.div(Nat.add(n, 1n), 2n)
224
225def ml_l(-r: Nat, -c: Nat, m: TA.ML<r, c>) -> List<&2, F32>:
226  match m:
227    case TA.ML{d, xs}:
228      xs
229
230# the values of a 16-bit list packed two per word
231def Half.words(+n: Nat, hs: List<&2, U32>) -> Array<U32>:
232  fill_u(join_l(pairs(hs)), 0, Array.new(U32, T.depth(hwords(n)), 0))
233
234# the first n 16-bit values of a packed array
235def Half.values(+n: Nat, a: Array<U32>) -> List<&2, U32>:
236  unpairs_n(n, split_l(ul_list(read_u(Nat.sub(hwords(n), 1n), 0, Nil{}, Array.get(U32, a, 0)))))
237
238# a matrix encoded by f (BF16.of_f32 or F16.of_f32) and packed
239def Half.pack(~f: F32 -> U32, +r: Nat, +c: Nat, m: TA.Mat<r, c>) -> HMat<r, c>:
240  HMat{Half.words(Nat.mul(r, c), enc_l(~f, ml_l(r, c, TA.Mat.to_list(r, c, m))))}
241
242# a packed matrix decoded by g (BF16.to_f32 or F16.to_f32) into F32
243def Half.unpack(~g: @+h: U32 -> F32, +r: Nat, +c: Nat, h: HMat<r, c>) -> TA.Mat<r, c>:
244  match h:
245    case HMat{d}:
246      TA.Mat.from_list(r, c, dec_l(~g, Half.values(Nat.mul(r, c), d)))
247
248def BF16.pack(+r: Nat, +c: Nat, m: TA.Mat<r, c>) -> HMat<r, c>:
249  Half.pack(~BF16.of_f32, r, c, m)
250
251def BF16.unpack(+r: Nat, +c: Nat, h: HMat<r, c>) -> TA.Mat<r, c>:
252  Half.unpack(~BF16.to_f32, r, c, h)
253
254def F16.pack(+r: Nat, +c: Nat, m: TA.Mat<r, c>) -> HMat<r, c>:
255  Half.pack(~F16.of_f32, r, c, m)
256
257def F16.unpack(+r: Nat, +c: Nat, h: HMat<r, c>) -> TA.Mat<r, c>:
258  Half.unpack(~F16.to_f32, r, c, h)
259
260# the words of a packed matrix, for printing and tests
261def HMat.words(+r: Nat, +c: Nat, h: HMat<r, c>) -> List<&2, U32>:
262  match h:
263    case HMat{d}:
264      ul_list(read_u(Nat.sub(hwords(Nat.mul(r, c)), 1n), 0, Nil{}, Array.get(U32, d, 0)))
265
266# =====================================================================================================
267# Laws of the packing
268# =====================================================================================================
269
270# LAW (unpair_pair): packing a list two by two and reading back as many values as it had gives the list back, in
271# order, for any length (an odd last value is paired with a 0 that the read leaves out). Proof: by induction two
272# values at a time; the pair HW{a, b} unpacks to a then b, and the rest is the induction hypothesis.
273def unpair_pair(xs: List<&2, U32>) -> {List.take(&2, U32, unpairs(pairs(xs)), List.length(&2, U32, xs)) == xs : List<&2, U32>}:
274  match xs:
275    case Nil{}:
276      {==}
277    case Con{a, Nil{}}:
278      {==}
279    case Con{+a, Con{+b, +t}}:
280      # bend-lint: disable-next bend-lint-rules/performance/unused_result -- the binder of Equal.cong's motive, used in its body not-expected-because: a law, never run
281      Equal.cong(List<&2, U32>, List<&2, U32>, l => a <> (b <> l), List.take(&2, U32, unpairs(pairs(t)), List.length(&2, U32, t)), t, unpair_pair(t))
282
283# Not proved (debt, tested instead by reference/test_dtype.py):
284#   - half_lo(half_join(a, b)) == a and half_hi(half_join(a, b)) == b for 16-bit a and b: a fact about Word(32n)'s bit
285#     operations for every value, which the checker does not reach by computation (`dtype_test joins` checks all 65536
286#     values of each half against a second one walking every value);
287#   - that fill_u and read_u give back what was written (Array.set then Array.get), and the int8 words' byte order;
288#   - the rounding itself: F32 is opaque to the checker, so BF16, F16 and int8 are tested bit for bit against PyTorch.
289
290# =====================================================================================================
291# int8 with a scale per row, for inference
292# =====================================================================================================
293# Row i of an r x c weight matrix W gets the scale s_i = max_j |W[i, j]| / 127 and the integers
294# q[i, j] = round(W[i, j] / s_i), to nearest even, in -127..127 (0 for a row of zeros). q is stored four per word,
295# element 4k + b in byte b of word k, as two's complement bytes: the bytes of a torch.int8 tensor. A product with F32
296# activations reads q straight from the words: y[n, i] = s_i · Σ_j x[n, j] · q[i, j], summed in order of j.
297
298type Q8<-r: Nat, -c: Nat> is Type:
299  Q8{q: Array<U32>, s: Array<F32>}
300
301def pick_f(b: Bool, +x: F32, +y: F32) -> F32:
302  match b:
303    case True{}:
304      x
305    case False{}:
306      y
307
308# v rounded to the nearest integer, ties to even (torch.round), for |v| < 2^23 (here |v| <= 127 and a bit)
309def rne_f(+v: F32) -> F32:
310  +f = F32.floor(v)
311  +d = (v - f : F32)
312  +h = (f * 0.5 : F32)
313  +up = Bool.or(F32.is_gt(d, 0.5), Bool.and(F32.is_eq(d, 0.5), F32.is_ne(F32.floor(h), h)))
314  pick_f(up, (f + 1.0 : F32), f)
315
316def q8_of(+w: F32, +s: F32) -> F32:
317  pick_f(F32.is_eq(s, 0.0), 0.0, F32.clamp(rne_f((w / s : F32)), F32.neg(127.0), 127.0))
318
319# an integer -127..127 (as an F32) to its two's complement byte, and back
320def q8_byte(+q: F32) -> U32:
321  ((F32.to_u32((q + 128.0 : F32)) + 128 : U32) .&. 255 : U32)
322
323def q8_val(+b: U32) -> F32:
324  (U32.to_f32(((b + 128 : U32) .&. 255 : U32)) - 128.0 : F32)
325
326def amax_l(xs: List<&2, F32>, +acc: F32) -> F32:
327  match xs:
328    case Nil{}:
329      acc
330    case Con{+h, t}:
331      amax_l(t, F32.max(acc, F32.abs(h)))
332
333def q8_row(xs: List<&2, F32>, +s: F32) -> List<&2, U32>:
334  match xs:
335    case Nil{}:
336      Nil{}
337    case Con{+h, t}:
338      q8_byte(q8_of(h, s)) <> q8_row(t, s)
339
340type QR is Type:
341  QR{q: List<&2, U32>, s: List<&2, F32>}
342
343def qr_add(+row: List<&2, F32>, r: QR) -> QR:
344  match r:
345    case QR{qs, ss}:
346      +s = (amax_l(row, 0.0) / 127.0 : F32)
347      QR{List.append(&2, U32, q8_row(row, s), qs), s <> ss}
348
349# the rows of a flat list of k rows of c, from the last one: each row's bytes and scale
350def q8_rows(k: Nat, +c: Nat, +xs: List<&2, F32>) -> QR:
351  match k:
352    case 0n:
353      QR{Nil{}, Nil{}}
354    case 1n+p:
355      qr_add(List.take(&2, F32, xs, c), q8_rows(p, c, List.drop(&2, F32, xs, c)))
356
357# four bytes to a word, little-endian; a short last group is padded with zeros
358def quads(xs: List<&2, U32>) -> List<&2, U32>:
359  match xs:
360    case Nil{}:
361      Nil{}
362    case Con{+a, Nil{}}:
363      [a]
364    case Con{+a, Con{+b, Nil{}}}:
365      [(a .|. U32.shln(b, 8n) : U32)]
366    case Con{+a, Con{+b, Con{+c, Nil{}}}}:
367      [((a .|. U32.shln(b, 8n) : U32) .|. U32.shln(c, 16n) : U32)]
368    case Con{+a, Con{+b, Con{+c, Con{+d, t}}}}:
369      (((a .|. U32.shln(b, 8n) : U32) .|. U32.shln(c, 16n) : U32) .|. U32.shln(d, 24n) : U32) <> quads(t)
370
371def fill_f(xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
372  match xs:
373    case Nil{}:
374      a
375    case Con{h, t}:
376      fill_f(t, (i + 1 : U32), Array.set(F32, a, i, h))
377
378def q8_fin(-r: Nat, -c: Nat, +n: Nat, +rr: Nat, qr: QR) -> Q8<r, c>:
379  match qr:
380    case QR{qs, ss}:
381      Q8{fill_u(quads(qs), 0, Array.new(U32, T.depth(Nat.div(Nat.add(n, 3n), 4n)), 0)), fill_f(ss, 0, Array.new(F32, T.depth(rr), 0.0))}
382
383# W quantized by rows (the matrix is consumed)
384def Q8.quantize(+r: Nat, +c: Nat, m: TA.Mat<r, c>) -> Q8<r, c>:
385  q8_fin(r, c, Nat.mul(r, c), r, q8_rows(r, c, ml_l(r, c, TA.Mat.to_list(r, c, m))))
386
387type FL is Type:
388  FL{a: Array<F32>, l: List<&2, F32>}
389
390def read_f(k: Nat, +i: U32, acc: List<&2, F32>, ra: Array<F32> & F32) -> FL:
391  match k ra:
392    case 0n Tuple{a, v}:
393      FL{a, List.reverse(&2, F32, v <> acc)}
394    case 1n+p Tuple{a, +v}:
395      read_f(p, (i + 1 : U32), v <> acc, Array.get(F32, a, (i + 1 : U32)))
396
397def fl_list(r: FL) -> List<&2, F32>:
398  match r:
399    case FL{a, l}:
400      l
401
402# byte e of the packed words, as its integer
403def q8_at(+e: U32, +w: U32) -> F32:
404  q8_val((U32.shrn(w, U32.to_nat((8 * (e .&. 3 : U32) : U32))) .&. 255 : U32))
405
406# the words and scales, for printing and tests
407type QL is Type:
408  QL{q: List<&2, U32>, s: List<&2, F32>}
409
410def Q8.lists(+r: Nat, +c: Nat, w: Q8<r, c>) -> QL:
411  match w:
412    case Q8{q, s}:
413      QL{ul_list(read_u(Nat.sub(Nat.div(Nat.add(Nat.mul(r, c), 3n), 4n), 1n), 0, Nil{}, Array.get(U32, q, 0))), fl_list(read_f(Nat.sub(r, 1n), 0, Nil{}, Array.get(F32, s, 0)))}
414
415# W back in F32: q[i, j] · s_i, one F32 product each
416def dq_row(k: Nat, +e: U32, +s: F32, acc: List<&2, F32>, rq: Array<U32> & U32) -> UF:
417  match k rq:
418    case 0n Tuple{q, w}:
419      UF{q, acc}
420    case 1n+p Tuple{q, +w}:
421      dq_row(p, (e + 1 : U32), s, (q8_at(e, w) * s : F32) <> acc, Array.get(U32, q, U32.shrn((e + 1 : U32), 2n)))
422
423type UF is Type:
424  UF{q: Array<U32>, l: List<&2, F32>}
425
426def dq_rows(k: Nat, +c: Nat, +cu: U32, +i: U32, ss: List<&2, F32>, r: UF) -> List<&2, F32>:
427  match k ss r:
428    case 0n _ UF{q, acc}:
429      List.reverse(&2, F32, acc)
430    case 1n+p Nil{} UF{q, acc}:
431      List.reverse(&2, F32, acc)
432    case 1n+p Con{+s, st} UF{q, acc}:
433      dq_rows(p, c, cu, (i + 1 : U32), st, dq_row(c, (i * cu : U32), s, acc, Array.get(U32, q, U32.shrn((i * cu : U32), 2n))))
434
435def dq_fin(+r: Nat, +c: Nat, q: Array<U32>, ss: List<&2, F32>) -> TA.Mat<r, c>:
436  TA.Mat.from_list(r, c, dq_rows(r, c, U32.from_nat(c), 0, ss, UF{q, Nil{}}))
437
438def dq_go(+r: Nat, +c: Nat, q: Array<U32>, sl: FL) -> TA.Mat<r, c>:
439  match sl:
440    case FL{s, ss}:
441      dq_fin(r, c, q, ss)
442
443def Q8.dequantize(+r: Nat, +c: Nat, w: Q8<r, c>) -> TA.Mat<r, c>:
444  match w:
445    case Q8{q, s}:
446      dq_go(r, c, q, read_f(Nat.sub(r, 1n), 0, Nil{}, Array.get(F32, s, 0)))
447
448# ---- the product y = x · Wᵀ (n x m) for x (n x k) in F32 and W (m x k) in int8, reading the bytes from the words ----
449
450type QD is Type:
451  QD{x: Array<F32>, q: Array<U32>, acc: F32}
452
453# Σ_l x[xi + l] · q[qe + l] for l < k, in order, from acc
454def q8_dot(k: Nat, +xi: U32, +qe: U32, +acc: F32, rx: Array<F32> & F32, rq: Array<U32> & U32) -> QD:
455  match k rx rq:
456    case 0n Tuple{x, v} Tuple{q, w}:
457      QD{x, q, acc}
458    case 1n+p Tuple{x, +v} Tuple{q, +w}:
459      q8_dot(p, (xi + 1 : U32), (qe + 1 : U32), (acc + (v * q8_at(qe, w) : F32) : F32), Array.get(F32, x, (xi + 1 : U32)), Array.get(U32, q, U32.shrn((qe + 1 : U32), 2n)))
460
461type QS is Type:
462  QS{x: Array<F32>, q: Array<U32>, s: Array<F32>, y: Array<F32>}
463
464def q8_put2(+idx: U32, +acc: F32, x: Array<F32>, q: Array<U32>, rs: Array<F32> & F32, y: Array<F32>) -> QS:
465  match rs:
466    case Tuple{s, +sc}:
467      QS{x, q, s, Array.set(F32, y, idx, (sc * acc : F32))}
468
469def q8_put(+idx: U32, +j: U32, d: QD, s: Array<F32>, y: Array<F32>) -> QS:
470  match d:
471    case QD{x, q, acc}:
472      q8_put2(idx, acc, x, q, Array.get(F32, s, j), y)
473
474# the m outputs of row i
475def q8_cols(ml: Nat, +i: U32, +j: U32, +k: Nat, +ku: U32, +mu: U32, st: QS) -> QS:
476  match ml st:
477    case 0n QS{x, q, s, y}:
478      QS{x, q, s, y}
479    case 1n+p QS{x, q, s, y}:
480      q8_cols(p, i, (j + 1 : U32), k, ku, mu, q8_put(((i * mu : U32) + j : U32), j, q8_dot(k, (i * ku : U32), (j * ku : U32), 0.0, Array.get(F32, x, (i * ku : U32)), Array.get(U32, q, U32.shrn((j * ku : U32), 2n))), s, y))
481
482def q8_rowsm(nl: Nat, +i: U32, +m: Nat, +k: Nat, +ku: U32, +mu: U32, st: QS) -> QS:
483  match nl:
484    case 0n:
485      st
486    case 1n+p:
487      q8_rowsm(p, (i + 1 : U32), m, k, ku, mu, q8_cols(m, i, 0, k, ku, mu, st))
488
489type Q8Mul<-n: Nat, -k: Nat, -m: Nat> is Type:
490  Q8Mul{x: TA.Mat<n, k>, w: Q8<m, k>, y: TA.Mat<n, m>}
491
492def q8_mm_fin(-n: Nat, -k: Nat, -m: Nat, st: QS) -> Q8Mul<n, k, m>:
493  match st:
494    case QS{x, q, s, y}:
495      Q8Mul{TA.Mat{x}, Q8{q, s}, TA.Mat{y}}
496
497# y = x · Wᵀ; x and W come back
498def Q8.matmul_nt(+n: Nat, +k: Nat, +m: Nat, x: TA.Mat<n, k>, w: Q8<m, k>) -> Q8Mul<n, k, m>:
499  match x w:
500    case TA.Mat{xa} Q8{q, s}:
501      q8_mm_fin(n, k, m, q8_rowsm(n, 0, m, k, U32.from_nat(k), U32.from_nat(m), QS{xa, q, s, T.zeros_a(Nat.mul(n, m))}))
502
503# =====================================================================================================
504# A two-F32 ("double-word") accumulator, for long sums (an experiment: Bend 2.0.37 has no F64)
505# =====================================================================================================
506# The running sum is hi + lo, with |lo| <= ulp(hi) / 2. Each addition is Knuth's TwoSum (the exact error of hi + x),
507# folded into lo, then renormalized; a product is first split exactly by Dekker's TwoProduct (Veltkamp's split with
508# 4097 = 2^12 + 1), which needs every operation rounded on its own: Bend's emitted C turns off FMA contraction.
509
510type DW is Data:
511  DW{hi: F32, lo: F32}
512
513def DW.zero() -> DW:
514  DW{0.0, 0.0}
515
516def DW.value(d: DW) -> F32:
517  match d:
518    case DW{+hi, +lo}:
519      (hi + lo : F32)
520
521# hi + lo when |hi| >= |lo| (Fast2Sum), as a renormalized pair
522def dw_norm(+s: F32, +e: F32) -> DW:
523  +h = (s + e : F32)
524  DW{h, (e - (h - s : F32) : F32)}
525
526# hi + x: TwoSum's s and its exact error, plus the old lo and an extra error term c
527def dw_fin(+s: F32, +bb: F32, +a: F32, +x: F32, +lo: F32, +c: F32) -> DW:
528  +e = ((a - (s - bb : F32) : F32) + (x - bb : F32) : F32)
529  dw_norm(s, ((e + lo : F32) + c : F32))
530
531def dw_add3(d: DW, +x: F32, +c: F32) -> DW:
532  match d:
533    case DW{+hi, +lo}:
534      +s = (hi + x : F32)
535      dw_fin(s, (s - hi : F32), hi, x, lo, c)
536
537def DW.add(d: DW, +x: F32) -> DW:
538  dw_add3(d, x, 0.0)
539
540def split_hi(+a: F32) -> F32:
541  +t = (4097.0 * a : F32)
542  (t - (t - a : F32) : F32)
543
544# a · b exactly as p + e
545def two_prod_e(+a: F32, +b: F32, +p: F32) -> F32:
546  +ah = split_hi(a)
547  +al = (a - ah : F32)
548  +bh = split_hi(b)
549  +bl = (b - bh : F32)
550  ((((((ah * bh : F32) - p : F32) + (ah * bl : F32) : F32) + (al * bh : F32) : F32)) + (al * bl : F32) : F32)
551
552# acc + a · b, with the product's rounding error kept
553def DW.fma(d: DW, +a: F32, +b: F32) -> DW:
554  +p = (a * b : F32)
555  dw_add3(d, p, two_prod_e(a, b, p))
556
557# ---- sums over arrays: a running F32 sum, and the same in double-word ----
558
559type AF is Type:
560  AF{a: Array<F32>, v: F32}
561
562def sum_f(k: Nat, +i: U32, +acc: F32, ra: Array<F32> & F32) -> AF:
563  match k ra:
564    case 0n Tuple{a, v}:
565      AF{a, acc}
566    case 1n+p Tuple{a, +v}:
567      sum_f(p, (i + 1 : U32), (acc + v : F32), Array.get(F32, a, (i + 1 : U32)))
568
569def sum_dw(k: Nat, +i: U32, acc: DW, ra: Array<F32> & F32) -> AF:
570  match k ra:
571    case 0n Tuple{a, v}:
572      AF{a, DW.value(acc)}
573    case 1n+p Tuple{a, +v}:
574      sum_dw(p, (i + 1 : U32), DW.add(acc, v), Array.get(F32, a, (i + 1 : U32)))
575
576type AF2 is Type:
577  AF2{a: Array<F32>, b: Array<F32>, v: F32}
578
579def dot_f(k: Nat, +i: U32, +acc: F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> AF2:
580  match k ra rb:
581    case 0n Tuple{a, x} Tuple{b, y}:
582      AF2{a, b, acc}
583    case 1n+p Tuple{a, +x} Tuple{b, +y}:
584      dot_f(p, (i + 1 : U32), (acc + (x * y : F32) : F32), Array.get(F32, a, (i + 1 : U32)), Array.get(F32, b, (i + 1 : U32)))
585
586def dot_dw(k: Nat, +i: U32, acc: DW, ra: Array<F32> & F32, rb: Array<F32> & F32) -> AF2:
587  match k ra rb:
588    case 0n Tuple{a, x} Tuple{b, y}:
589      AF2{a, b, DW.value(acc)}
590    case 1n+p Tuple{a, +x} Tuple{b, +y}:
591      dot_dw(p, (i + 1 : U32), DW.fma(acc, x, y), Array.get(F32, a, (i + 1 : U32)), Array.get(F32, b, (i + 1 : U32)))
592
593# the sums of the first n numbers of a (the array comes back)
594def Sum.f32(+n: Nat, a: Array<F32>) -> AF:
595  sum_f(n, 0, 0.0, Array.get(F32, a, 0))
596
597def Sum.dw(+n: Nat, a: Array<F32>) -> AF:
598  sum_dw(n, 0, DW.zero(), Array.get(F32, a, 0))
599
600def Dot.f32(+n: Nat, a: Array<F32>, b: Array<F32>) -> AF2:
601  dot_f(n, 0, 0.0, Array.get(F32, a, 0), Array.get(F32, b, 0))
602
603def Dot.dw(+n: Nat, a: Array<F32>, b: Array<F32>) -> AF2:
604  dot_dw(n, 0, DW.zero(), Array.get(F32, a, 0), Array.get(F32, b, 0))
605
606# =====================================================================================================
607# Mixed precision: BF16 storage around F32 arithmetic
608# =====================================================================================================
609# Two wrappers keep a layer's types (P, X, Y, S), so a model gets BF16 storage by wrapping its layers, with the same
610# parameters, checkpoints and optimizer state:
611#   - CastOut(l): l's output rounded to BF16 (as if stored in BF16), and the gradient arriving at it rounded too, which
612#     is what autograd does through x.to(torch.bfloat16).to(torch.float32).
613#   - CastP(l), for a layer whose parameter is one matrix: the parameter stays F32 (the master weights the optimizer
614#     updates); l's forward and backward passes see its BF16 rounding, and the gradient l accumulates is rounded to
615#     BF16 before the optimizer sees it, as autograd does for w.to(torch.bfloat16).to(torch.float32). The rounded copy
616#     is made again in the backward pass (no saved state), so S does not change.
617# Inside the layers every operation is F32. The values are BF16 values held in F32 arrays: this gives BF16's numbers,
618# not yet its memory (no kernel reads packed operands; that is the next step).
619
620def bf16_mat(+r: Nat, +c: Nat, m: TA.Mat<r, c>) -> TA.Mat<r, c>:
621  T.Ew.map(~BF16.round, r, c, m)
622
623def cast_f(-P: Type, -S: Type, +r: Nat, +c: Nat, f: T.Fwd<P, TA.Mat<r, c>, S>) -> T.Fwd<P, TA.Mat<r, c>, S>:
624  match f:
625    case T.Fwd{p, y, s}:
626      T.Fwd{p, bf16_mat(r, c, y), s}
627
628def cast_fwd(~P: Type, ~X: Type, ~S: Type, ~r: Nat, ~c: Nat, ~l: T.Layer<P, X, TA.Mat<r, c>, S>, p: P, x: X) -> T.Fwd<P, TA.Mat<r, c>, S>:
629  cast_f(P, S, r, c, T.Layer.fwd(P, X, TA.Mat<r, c>, S, l, p, x))
630
631def cast_bwd(~P: Type, ~X: Type, ~S: Type, ~r: Nat, ~c: Nat, ~l: T.Layer<P, X, TA.Mat<r, c>, S>, p: P, s: S, dy: TA.Mat<r, c>, g: P) -> T.Bwd<P, X>:
632  T.Layer.bwd(P, X, TA.Mat<r, c>, S, l, p, s, bf16_mat(r, c, dy), g)
633
634def CastOut(~P: Type, ~X: Type, ~S: Type, ~r: Nat, ~c: Nat, ~l: T.Layer<P, X, TA.Mat<r, c>, S>) -> T.Layer<P, X, TA.Mat<r, c>, S>:
635  T.Layer{p => x => cast_fwd(~P, ~X, ~S, ~r, ~c, ~l, p, x), p => s => dy => g => cast_bwd(~P, ~X, ~S, ~r, ~c, ~l, p, s, dy, g), op => p => g => st => T.Layer.step(P, X, TA.Mat<r, c>, S, l, op, p, g, st), u => T.Layer.zeros(P, X, TA.Mat<r, c>, S, l), p => T.Layer.dump(P, X, TA.Mat<r, c>, S, l, p), p => xs => off => T.Layer.load(P, X, TA.Mat<r, c>, S, l, p, xs, off), p => xs => off => T.Layer.store(P, X, TA.Mat<r, c>, S, l, p, xs, off), u => T.Layer.count(P, X, TA.Mat<r, c>, S, l)}
636
637def castp_f2(-X: Type, -Y: Type, -S: Type, -r: Nat, -c: Nat, w: TA.Mat<r, c>, f: T.Fwd<TA.Mat<r, c>, Y, S>) -> T.Fwd<TA.Mat<r, c>, Y, S>:
638  match f:
639    case T.Fwd{wc, y, s}:
640      T.Fwd{w, y, s}
641
642def castp_f(~X: Type, ~Y: Type, ~S: Type, ~r: Nat, ~c: Nat, ~l: T.Layer<TA.Mat<r, c>, X, Y, S>, d: T.Two<TA.Mat<r, c>, TA.Mat<r, c>>, x: X) -> T.Fwd<TA.Mat<r, c>, Y, S>:
643  match d:
644    case T.Two{w, wc}:
645      castp_f2(X, Y, S, r, c, w, T.Layer.fwd(TA.Mat<r, c>, X, Y, S, l, bf16_mat(r, c, wc), x))
646
647def castp_b2(-X: Type, +r: Nat, +c: Nat, w: TA.Mat<r, c>, b: T.Bwd<TA.Mat<r, c>, X>) -> T.Bwd<TA.Mat<r, c>, X>:
648  match b:
649    case T.Bwd{wc, dx, g}:
650      T.Bwd{w, dx, bf16_mat(r, c, g)}
651
652def castp_b(~X: Type, ~Y: Type, ~S: Type, ~r: Nat, ~c: Nat, ~l: T.Layer<TA.Mat<r, c>, X, Y, S>, d: T.Two<TA.Mat<r, c>, TA.Mat<r, c>>, s: S, dy: Y, g: TA.Mat<r, c>) -> T.Bwd<TA.Mat<r, c>, X>:
653  match d:
654    case T.Two{w, wc}:
655      castp_b2(X, r, c, w, T.Layer.bwd(TA.Mat<r, c>, X, Y, S, l, bf16_mat(r, c, wc), s, dy, g))
656
657def CastP(~X: Type, ~Y: Type, ~S: Type, ~r: Nat, ~c: Nat, ~l: T.Layer<TA.Mat<r, c>, X, Y, S>) -> T.Layer<TA.Mat<r, c>, X, Y, S>:
658  T.Layer{p => x => castp_f(~X, ~Y, ~S, ~r, ~c, ~l, T.Mat.dup(r, c, p), x), p => s => dy => g => castp_b(~X, ~Y, ~S, ~r, ~c, ~l, T.Mat.dup(r, c, p), s, dy, g), op => p => g => st => T.Layer.step(TA.Mat<r, c>, X, Y, S, l, op, p, g, st), u => T.Layer.zeros(TA.Mat<r, c>, X, Y, S, l), p => T.Layer.dump(TA.Mat<r, c>, X, Y, S, l, p), p => xs => off => T.Layer.load(TA.Mat<r, c>, X, Y, S, l, p, xs, off), p => xs => off => T.Layer.store(TA.Mat<r, c>, X, Y, S, l, p, xs, off), u => T.Layer.count(TA.Mat<r, c>, X, Y, S, l)}