dtype.bend
bendtorch@1.1.0.0 · 657 lines · 0 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.1.0.0/main.bend as T
15import bend-ml-tensor-array@0.1.14.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 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))
281
282# Not proved (debt, tested instead by reference/test_dtype.py):
283# - 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
284# operations for every value, which the checker does not reach by computation (`dtype_test joins` checks all 65536
285# values of each half against a second one walking every value);
286# - that fill_u and read_u give back what was written (Array.set then Array.get), and the int8 words' byte order;
287# - the rounding itself: F32 is opaque to the checker, so BF16, F16 and int8 are tested bit for bit against PyTorch.
288
289# =====================================================================================================
290# int8 with a scale per row, for inference
291# =====================================================================================================
292# Row i of an r x c weight matrix W gets the scale s_i = max_j |W[i, j]| / 127 and the integers
293# 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,
294# 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
295# activations reads q straight from the words: y[n, i] = s_i · Σ_j x[n, j] · q[i, j], summed in order of j.
296
297type Q8<-r: Nat, -c: Nat> is Type:
298 Q8{q: Array<U32>, s: Array<F32>}
299
300def pick_f(b: Bool, +x: F32, +y: F32) -> F32:
301 match b:
302 case True{}:
303 x
304 case False{}:
305 y
306
307# v rounded to the nearest integer, ties to even (torch.round), for |v| < 2^23 (here |v| <= 127 and a bit)
308def rne_f(+v: F32) -> F32:
309 +f = F32.floor(v)
310 +d = (v - f : F32)
311 +h = (f * 0.5 : F32)
312 +up = Bool.or(F32.is_gt(d, 0.5), Bool.and(F32.is_eq(d, 0.5), F32.is_ne(F32.floor(h), h)))
313 pick_f(up, (f + 1.0 : F32), f)
314
315def q8_of(+w: F32, +s: F32) -> F32:
316 pick_f(F32.is_eq(s, 0.0), 0.0, F32.clamp(rne_f((w / s : F32)), F32.neg(127.0), 127.0))
317
318# an integer -127..127 (as an F32) to its two's complement byte, and back
319def q8_byte(+q: F32) -> U32:
320 ((F32.to_u32((q + 128.0 : F32)) + 128 : U32) .&. 255 : U32)
321
322def q8_val(+b: U32) -> F32:
323 (U32.to_f32(((b + 128 : U32) .&. 255 : U32)) - 128.0 : F32)
324
325def amax_l(xs: List<&2, F32>, +acc: F32) -> F32:
326 match xs:
327 case Nil{}:
328 acc
329 case Con{+h, t}:
330 amax_l(t, F32.max(acc, F32.abs(h)))
331
332def q8_row(xs: List<&2, F32>, +s: F32) -> List<&2, U32>:
333 match xs:
334 case Nil{}:
335 Nil{}
336 case Con{+h, t}:
337 q8_byte(q8_of(h, s)) <> q8_row(t, s)
338
339type QR is Type:
340 QR{q: List<&2, U32>, s: List<&2, F32>}
341
342def qr_add(+row: List<&2, F32>, r: QR) -> QR:
343 match r:
344 case QR{qs, ss}:
345 +s = (amax_l(row, 0.0) / 127.0 : F32)
346 QR{List.append(&2, U32, q8_row(row, s), qs), s <> ss}
347
348# the rows of a flat list of k rows of c, from the last one: each row's bytes and scale
349def q8_rows(k: Nat, +c: Nat, +xs: List<&2, F32>) -> QR:
350 match k:
351 case 0n:
352 QR{Nil{}, Nil{}}
353 case 1n+p:
354 qr_add(List.take(&2, F32, xs, c), q8_rows(p, c, List.drop(&2, F32, xs, c)))
355
356# four bytes to a word, little-endian; a short last group is padded with zeros
357def quads(xs: List<&2, U32>) -> List<&2, U32>:
358 match xs:
359 case Nil{}:
360 Nil{}
361 case Con{+a, Nil{}}:
362 [a]
363 case Con{+a, Con{+b, Nil{}}}:
364 [(a .|. U32.shln(b, 8n) : U32)]
365 case Con{+a, Con{+b, Con{+c, Nil{}}}}:
366 [((a .|. U32.shln(b, 8n) : U32) .|. U32.shln(c, 16n) : U32)]
367 case Con{+a, Con{+b, Con{+c, Con{+d, t}}}}:
368 (((a .|. U32.shln(b, 8n) : U32) .|. U32.shln(c, 16n) : U32) .|. U32.shln(d, 24n) : U32) <> quads(t)
369
370def fill_f(xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
371 match xs:
372 case Nil{}:
373 a
374 case Con{h, t}:
375 fill_f(t, (i + 1 : U32), Array.set(F32, a, i, h))
376
377def q8_fin(-r: Nat, -c: Nat, +n: Nat, +rr: Nat, qr: QR) -> Q8<r, c>:
378 match qr:
379 case QR{qs, ss}:
380 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))}
381
382# W quantized by rows (the matrix is consumed)
383def Q8.quantize(+r: Nat, +c: Nat, m: TA.Mat<r, c>) -> Q8<r, c>:
384 q8_fin(r, c, Nat.mul(r, c), r, q8_rows(r, c, ml_l(r, c, TA.Mat.to_list(r, c, m))))
385
386type FL is Type:
387 FL{a: Array<F32>, l: List<&2, F32>}
388
389def read_f(k: Nat, +i: U32, acc: List<&2, F32>, ra: Array<F32> & F32) -> FL:
390 match k ra:
391 case 0n Tuple{a, v}:
392 FL{a, List.reverse(&2, F32, v <> acc)}
393 case 1n+p Tuple{a, +v}:
394 read_f(p, (i + 1 : U32), v <> acc, Array.get(F32, a, (i + 1 : U32)))
395
396def fl_list(r: FL) -> List<&2, F32>:
397 match r:
398 case FL{a, l}:
399 l
400
401# byte e of the packed words, as its integer
402def q8_at(+e: U32, +w: U32) -> F32:
403 q8_val((U32.shrn(w, U32.to_nat((8 * (e .&. 3 : U32) : U32))) .&. 255 : U32))
404
405# the words and scales, for printing and tests
406type QL is Type:
407 QL{q: List<&2, U32>, s: List<&2, F32>}
408
409def Q8.lists(+r: Nat, +c: Nat, w: Q8<r, c>) -> QL:
410 match w:
411 case Q8{q, s}:
412 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)))}
413
414# W back in F32: q[i, j] · s_i, one F32 product each
415def dq_row(k: Nat, +e: U32, +s: F32, acc: List<&2, F32>, rq: Array<U32> & U32) -> UF:
416 match k rq:
417 case 0n Tuple{q, w}:
418 UF{q, acc}
419 case 1n+p Tuple{q, +w}:
420 dq_row(p, (e + 1 : U32), s, (q8_at(e, w) * s : F32) <> acc, Array.get(U32, q, U32.shrn((e + 1 : U32), 2n)))
421
422type UF is Type:
423 UF{q: Array<U32>, l: List<&2, F32>}
424
425def dq_rows(k: Nat, +c: Nat, +cu: U32, +i: U32, ss: List<&2, F32>, r: UF) -> List<&2, F32>:
426 match k ss r:
427 case 0n _ UF{q, acc}:
428 List.reverse(&2, F32, acc)
429 case 1n+p Nil{} UF{q, acc}:
430 List.reverse(&2, F32, acc)
431 case 1n+p Con{+s, st} UF{q, acc}:
432 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))))
433
434def dq_fin(+r: Nat, +c: Nat, q: Array<U32>, ss: List<&2, F32>) -> TA.Mat<r, c>:
435 TA.Mat.from_list(r, c, dq_rows(r, c, U32.from_nat(c), 0, ss, UF{q, Nil{}}))
436
437def dq_go(+r: Nat, +c: Nat, q: Array<U32>, sl: FL) -> TA.Mat<r, c>:
438 match sl:
439 case FL{s, ss}:
440 dq_fin(r, c, q, ss)
441
442def Q8.dequantize(+r: Nat, +c: Nat, w: Q8<r, c>) -> TA.Mat<r, c>:
443 match w:
444 case Q8{q, s}:
445 dq_go(r, c, q, read_f(Nat.sub(r, 1n), 0, Nil{}, Array.get(F32, s, 0)))
446
447# ---- 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 ----
448
449type QD is Type:
450 QD{x: Array<F32>, q: Array<U32>, acc: F32}
451
452# Σ_l x[xi + l] · q[qe + l] for l < k, in order, from acc
453def q8_dot(k: Nat, +xi: U32, +qe: U32, +acc: F32, rx: Array<F32> & F32, rq: Array<U32> & U32) -> QD:
454 match k rx rq:
455 case 0n Tuple{x, v} Tuple{q, w}:
456 QD{x, q, acc}
457 case 1n+p Tuple{x, +v} Tuple{q, +w}:
458 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)))
459
460type QS is Type:
461 QS{x: Array<F32>, q: Array<U32>, s: Array<F32>, y: Array<F32>}
462
463def q8_put2(+idx: U32, +acc: F32, x: Array<F32>, q: Array<U32>, rs: Array<F32> & F32, y: Array<F32>) -> QS:
464 match rs:
465 case Tuple{s, +sc}:
466 QS{x, q, s, Array.set(F32, y, idx, (sc * acc : F32))}
467
468def q8_put(+idx: U32, +j: U32, d: QD, s: Array<F32>, y: Array<F32>) -> QS:
469 match d:
470 case QD{x, q, acc}:
471 q8_put2(idx, acc, x, q, Array.get(F32, s, j), y)
472
473# the m outputs of row i
474def q8_cols(ml: Nat, +i: U32, +j: U32, +k: Nat, +ku: U32, +mu: U32, st: QS) -> QS:
475 match ml st:
476 case 0n QS{x, q, s, y}:
477 QS{x, q, s, y}
478 case 1n+p QS{x, q, s, y}:
479 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))
480
481def q8_rowsm(nl: Nat, +i: U32, +m: Nat, +k: Nat, +ku: U32, +mu: U32, st: QS) -> QS:
482 match nl:
483 case 0n:
484 st
485 case 1n+p:
486 q8_rowsm(p, (i + 1 : U32), m, k, ku, mu, q8_cols(m, i, 0, k, ku, mu, st))
487
488type Q8Mul<-n: Nat, -k: Nat, -m: Nat> is Type:
489 Q8Mul{x: TA.Mat<n, k>, w: Q8<m, k>, y: TA.Mat<n, m>}
490
491def q8_mm_fin(-n: Nat, -k: Nat, -m: Nat, st: QS) -> Q8Mul<n, k, m>:
492 match st:
493 case QS{x, q, s, y}:
494 Q8Mul{TA.Mat{x}, Q8{q, s}, TA.Mat{y}}
495
496# y = x · Wᵀ; x and W come back
497def Q8.matmul_nt(+n: Nat, +k: Nat, +m: Nat, x: TA.Mat<n, k>, w: Q8<m, k>) -> Q8Mul<n, k, m>:
498 match x w:
499 case TA.Mat{xa} Q8{q, s}:
500 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))}))
501
502# =====================================================================================================
503# A two-F32 ("double-word") accumulator, for long sums (an experiment: Bend 2.0.37 has no F64)
504# =====================================================================================================
505# The running sum is hi + lo, with |lo| <= ulp(hi) / 2. Each addition is Knuth's TwoSum (the exact error of hi + x),
506# folded into lo, then renormalized; a product is first split exactly by Dekker's TwoProduct (Veltkamp's split with
507# 4097 = 2^12 + 1), which needs every operation rounded on its own: Bend's emitted C turns off FMA contraction.
508
509type DW is Data:
510 DW{hi: F32, lo: F32}
511
512def DW.zero() -> DW:
513 DW{0.0, 0.0}
514
515def DW.value(d: DW) -> F32:
516 match d:
517 case DW{+hi, +lo}:
518 (hi + lo : F32)
519
520# hi + lo when |hi| >= |lo| (Fast2Sum), as a renormalized pair
521def dw_norm(+s: F32, +e: F32) -> DW:
522 +h = (s + e : F32)
523 DW{h, (e - (h - s : F32) : F32)}
524
525# hi + x: TwoSum's s and its exact error, plus the old lo and an extra error term c
526def dw_fin(+s: F32, +bb: F32, +a: F32, +x: F32, +lo: F32, +c: F32) -> DW:
527 +e = ((a - (s - bb : F32) : F32) + (x - bb : F32) : F32)
528 dw_norm(s, ((e + lo : F32) + c : F32))
529
530def dw_add3(d: DW, +x: F32, +c: F32) -> DW:
531 match d:
532 case DW{+hi, +lo}:
533 +s = (hi + x : F32)
534 dw_fin(s, (s - hi : F32), hi, x, lo, c)
535
536def DW.add(d: DW, +x: F32) -> DW:
537 dw_add3(d, x, 0.0)
538
539def split_hi(+a: F32) -> F32:
540 +t = (4097.0 * a : F32)
541 (t - (t - a : F32) : F32)
542
543# a · b exactly as p + e
544def two_prod_e(+a: F32, +b: F32, +p: F32) -> F32:
545 +ah = split_hi(a)
546 +al = (a - ah : F32)
547 +bh = split_hi(b)
548 +bl = (b - bh : F32)
549 ((((((ah * bh : F32) - p : F32) + (ah * bl : F32) : F32) + (al * bh : F32) : F32)) + (al * bl : F32) : F32)
550
551# acc + a · b, with the product's rounding error kept
552def DW.fma(d: DW, +a: F32, +b: F32) -> DW:
553 +p = (a * b : F32)
554 dw_add3(d, p, two_prod_e(a, b, p))
555
556# ---- sums over arrays: a running F32 sum, and the same in double-word ----
557
558type AF is Type:
559 AF{a: Array<F32>, v: F32}
560
561def sum_f(k: Nat, +i: U32, +acc: F32, ra: Array<F32> & F32) -> AF:
562 match k ra:
563 case 0n Tuple{a, v}:
564 AF{a, acc}
565 case 1n+p Tuple{a, +v}:
566 sum_f(p, (i + 1 : U32), (acc + v : F32), Array.get(F32, a, (i + 1 : U32)))
567
568def sum_dw(k: Nat, +i: U32, acc: DW, ra: Array<F32> & F32) -> AF:
569 match k ra:
570 case 0n Tuple{a, v}:
571 AF{a, DW.value(acc)}
572 case 1n+p Tuple{a, +v}:
573 sum_dw(p, (i + 1 : U32), DW.add(acc, v), Array.get(F32, a, (i + 1 : U32)))
574
575type AF2 is Type:
576 AF2{a: Array<F32>, b: Array<F32>, v: F32}
577
578def dot_f(k: Nat, +i: U32, +acc: F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> AF2:
579 match k ra rb:
580 case 0n Tuple{a, x} Tuple{b, y}:
581 AF2{a, b, acc}
582 case 1n+p Tuple{a, +x} Tuple{b, +y}:
583 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)))
584
585def dot_dw(k: Nat, +i: U32, acc: DW, ra: Array<F32> & F32, rb: Array<F32> & F32) -> AF2:
586 match k ra rb:
587 case 0n Tuple{a, x} Tuple{b, y}:
588 AF2{a, b, DW.value(acc)}
589 case 1n+p Tuple{a, +x} Tuple{b, +y}:
590 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)))
591
592# the sums of the first n numbers of a (the array comes back)
593def Sum.f32(+n: Nat, a: Array<F32>) -> AF:
594 sum_f(n, 0, 0.0, Array.get(F32, a, 0))
595
596def Sum.dw(+n: Nat, a: Array<F32>) -> AF:
597 sum_dw(n, 0, DW.zero(), Array.get(F32, a, 0))
598
599def Dot.f32(+n: Nat, a: Array<F32>, b: Array<F32>) -> AF2:
600 dot_f(n, 0, 0.0, Array.get(F32, a, 0), Array.get(F32, b, 0))
601
602def Dot.dw(+n: Nat, a: Array<F32>, b: Array<F32>) -> AF2:
603 dot_dw(n, 0, DW.zero(), Array.get(F32, a, 0), Array.get(F32, b, 0))
604
605# =====================================================================================================
606# Mixed precision: BF16 storage around F32 arithmetic
607# =====================================================================================================
608# Two wrappers keep a layer's types (P, X, Y, S), so a model gets BF16 storage by wrapping its layers, with the same
609# parameters, checkpoints and optimizer state:
610# - CastOut(l): l's output rounded to BF16 (as if stored in BF16), and the gradient arriving at it rounded too, which
611# is what autograd does through x.to(torch.bfloat16).to(torch.float32).
612# - CastP(l), for a layer whose parameter is one matrix: the parameter stays F32 (the master weights the optimizer
613# updates); l's forward and backward passes see its BF16 rounding, and the gradient l accumulates is rounded to
614# BF16 before the optimizer sees it, as autograd does for w.to(torch.bfloat16).to(torch.float32). The rounded copy
615# is made again in the backward pass (no saved state), so S does not change.
616# Inside the layers every operation is F32. The values are BF16 values held in F32 arrays: this gives BF16's numbers,
617# not yet its memory (no kernel reads packed operands; that is the next step).
618
619def bf16_mat(+r: Nat, +c: Nat, m: TA.Mat<r, c>) -> TA.Mat<r, c>:
620 T.Ew.map(~BF16.round, r, c, m)
621
622def 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>:
623 match f:
624 case T.Fwd{p, y, s}:
625 T.Fwd{p, bf16_mat(r, c, y), s}
626
627def 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>:
628 cast_f(P, S, r, c, T.Layer.fwd(P, X, TA.Mat<r, c>, S, l, p, x))
629
630def 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>:
631 T.Layer.bwd(P, X, TA.Mat<r, c>, S, l, p, s, bf16_mat(r, c, dy), g)
632
633def 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>:
634 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)}
635
636def 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>:
637 match f:
638 case T.Fwd{wc, y, s}:
639 T.Fwd{w, y, s}
640
641def 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>:
642 match d:
643 case T.Two{w, wc}:
644 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))
645
646def 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>:
647 match b:
648 case T.Bwd{wc, dx, g}:
649 T.Bwd{w, dx, bf16_mat(r, c, g)}
650
651def 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>:
652 match d:
653 case T.Two{w, wc}:
654 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))
655
656def 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>:
657 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)}