nx_quant.nx source
↩ module page · 661 lines · 25392 B
1// quant.nx -- quantization primitives for AI inference.
2//
3// Packs / unpacks tensor data between fp16/fp32 and the
4// low-precision formats nxgguf supports (int8, int4, int2,
5// ternary, fp8 E4M3). These are the actual VRAM-saving
6// transforms that take a 32GB Llama 70B fp16 model down to
7// 8GB int4 or 4GB int2.
8//
9// Reference quantization schemes:
10// int8 SmoothQuant (Xiao et al. 2022)
11// int4 GPTQ (Frantar et al. 2022)
12// int4 AWQ (Lin et al. 2023)
13// GGUF k-quants (llama.cpp project)
14// BitNet b1.58 (Wang et al. 2024) -- ternary {-1, 0, +1}
15//
16// v0.0.1 ships symmetric quantization with per-tensor scale.
17// k-quants (per-block scale + zero-point) follow in v0.1.0.
18
19// nx_safety_envelope:
20// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
21// sil_target: SIL1
22// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
23// verdict: NOT_YET_EVALUATED
24
25import "nx_syscalls.nx"
26const K_MAGIC_2046: i64 = 2046
27
28// === fp32 literal encoding ===========================================
29//
30// Convert a lexer-split decimal literal (whole + fractional-digits +
31// number-of-fractional-digits) to IEEE 754 binary32 bit pattern.
32// Called from parse.nx when it sees TK_FLOAT and needs an IR
33// constant with a concrete bit pattern.
34//
35// Algorithm:
36// value = (whole * 10^frac_digits + frac_num) / 10^frac_digits
37// scale the numerator by 2^40 before the divide to preserve
38// precision, locate the highest set bit to derive the exponent,
39// then pack mantissa+exponent into the 32-bit layout.
40//
41// Precision: accurate to ~7 decimal digits (f32 is 24-bit
42// significand). Exact for common cases like 0.0, 0.5, 1.0, 2.0,
43// 1.5, 3.14, 0.1 (all tested). Overflow saturates to +inf;
44// underflow flushes to zero.
45//
46// Negative literals handled by parse.nx prepending a unary minus
47// separately -- the lexer emits only the unsigned magnitude.
48
49func fp32_from_parts(whole: i64, frac_num: i64, frac_digits: i64) -> i64 {
50 // Zero literal -> bit pattern 0.
51 if whole == 0 {
52 if frac_num == 0 { return 0 }
53 }
54
55 // 10^frac_digits.
56 var denom: i64 = 1
57 var i: i64 = 0
58 while i < frac_digits {
59 denom = denom * 10
60 i = i + 1
61 }
62
63 // Scaled numerator. 2^40 gives room for a 24-bit mantissa above
64 // denom across the common literal range -- 40-bit of margin
65 // drowns rounding error for 7-digit literals.
66 let num: i64 = whole * denom + frac_num
67 let scale: i64 = 40
68 let scaled: i64 = num << scale
69 let quot: i64 = scaled / denom
70
71 // Locate highest set bit of quot.
72 var high_bit: i64 = 0
73 var q: i64 = quot
74 while q > 1 {
75 q = q >> 1
76 high_bit = high_bit + 1
77 }
78
79 // True exponent = high_bit - scale.
80 let true_exp: i64 = high_bit - scale
81
82 // Align quot so the implicit 1 bit lands at bit 23, then mask.
83 var mant: i64 = 0
84 if high_bit >= 23 {
85 mant = quot >> (high_bit - 23)
86 }
87 if high_bit < 23 {
88 mant = quot << (23 - high_bit)
89 }
90 mant = mant & 0x7FFFFF
91
92 // Bias exponent by 127. Saturate on overflow, flush on underflow.
93 let biased: i64 = true_exp + 127
94 if biased < 0 { return 0 }
95 if biased > 254 { return 0x7F800000 }
96 return (biased << 23) | mant
97}
98
99// === fp64 literal encoding ===========================================
100//
101// Convert a lexer-split decimal literal to IEEE 754 binary64 bit
102// pattern. Mirrors fp32_from_parts; differs in:
103// - mantissa width 52 (vs 23)
104// - exponent bias 1023 (vs 127)
105// - exponent saturation at 2046 / +inf bit pattern 0x7FF0000000000000
106// - scale 52 (vs 40) -- gives 52-bit mantissa headroom in the
107// scaled-numerator divide. Caps decodable literal magnitude at
108// ~3 digits whole + ~3 digits frac before the i64 multiply
109// overflows; matches fp32_from_parts's same pragmatic limit.
110// Larger-magnitude literals need i128 multiword arithmetic
111// (queued: u128/i128 substrate primitives, see roadmap doc 16).
112//
113// Precision: accurate to ~7 decimal digits today (limited by the
114// scale=52 i64 budget, not by the algorithm itself). Exact for
115// 0.0, 0.5, 1.0, 2.0, 1.5, 4.0, 8.0; ULP-close for 0.1, 3.14.
116//
117// Negative literals: parse.nx prepends unary minus separately --
118// the lexer emits unsigned magnitude only.
119//
120// IEEE 754 binary64 layout: 1 sign + 11 exp + 52 mant.
121
122// LN36 + LN40 (2026-09-03): the decimal literal machine now hands the parser (m, e10, sticky) -- value = m x
123// 10^e10 with m the first 18 significant digits -- and fp64_from_dec below packs it. This legacy entry keeps
124// its (whole, frac, digits) contract for the fp32 path and any other caller; both share fp64_from_ratio.
125func fp64_from_parts(whole: i64, frac_num: i64, frac_digits: i64) -> i64 {
126 // Zero literal -> bit pattern 0.
127 if whole == 0 {
128 if frac_num == 0 { return 0 }
129 }
130 var denom0: i64 = 1
131 var idx0: i64 = 0
132 while idx0 < frac_digits {
133 denom0 = denom0 * 10
134 idx0 = idx0 + 1
135 }
136 return fp64_from_ratio(whole * denom0 + frac_num, denom0, 0)
137}
138
139// THE EXACT CORE: num/denom (num > 0, 1 <= denom < 2^62, num/denom < 2^53) -> binary64 bits, round-to-nearest-
140// even by bit-by-bit long division, sticky_in OR-ed into the sticky bit so digits the lexer dropped still
141// break a tie correctly. The q_int > 0 arm needs q_int < 2^53 (its mantissa shift is 52 - high_bit); callers
142// with a larger numerator take fp64_from_sig, which normalises any 63-bit integer exactly.
143func fp64_from_ratio(num: i64, denom: i64, sticky_in: i64) -> i64 {
144 if num == 0 { return 0 }
145
146 // Build the rational num/denom representing the literal exactly.
147 // denom = 10^frac_digits
148 // num = whole*denom + frac_num
149 // (No scaling, no magic constants -- num and denom are the exact
150 // integer-ratio form of the decimal literal.)
151 // Step 1: integer and fractional parts of num/denom.
152 var q_int: i64 = num / denom
153 var r: i64 = num - q_int * denom
154
155 // Step 2: locate the implicit-1 bit and seed true_exp + mantissa.
156 //
157 // Case A: q_int > 0 -- the implicit 1 is the MSB of q_int. Bits
158 // below it become the top of the mantissa; remaining
159 // mantissa bits come from r/denom by bit-extraction.
160 // Case B: q_int == 0 -- the value is sub-1. Walk r through
161 // "double-and-subtract" until the implicit 1 appears,
162 // counting leading binary zeros as negative true_exp.
163 var true_exp: i64 = 0
164 var mant: i64 = 0
165 var bits_to_fill: i64 = 52
166
167 if q_int > 0 {
168 var high_bit: i64 = 0
169 var q: i64 = q_int
170 while q > 1 {
171 q = q >> 1
172 high_bit = high_bit + 1
173 }
174 true_exp = high_bit
175 // Bits below the implicit 1 become the top of the mantissa.
176 let mant_top: i64 = q_int & ((1 << high_bit) - 1)
177 mant = mant_top << (52 - high_bit)
178 bits_to_fill = 52 - high_bit
179 } else {
180 // q_int == 0, r > 0. Find the leading-1 bit position by
181 // doubling r until r >= denom. Each doubling consumes one
182 // negative exponent step. Bounded by ~log2(denom) iterations
183 // (e.g. 10 for frac_digits=3) so loop terminates quickly.
184 true_exp = -1
185 r = r + r
186 while r < denom {
187 true_exp = true_exp - 1
188 r = r + r
189 }
190 // r >= denom now: the implicit 1 is present. Consume it.
191 r = r - denom
192 bits_to_fill = 52
193 }
194
195 // Step 3: extract `bits_to_fill` mantissa bits from r/denom via
196 // bit-by-bit long division (double, compare, subtract). Each
197 // iteration peels one binary bit off the rational r/denom.
198 var pos: i64 = bits_to_fill - 1
199 while pos >= 0 {
200 r = r + r
201 if r >= denom {
202 mant = mant | (1 << pos)
203 r = r - denom
204 }
205 pos = pos - 1
206 }
207
208 // Step 4: round-to-nearest-even using the next bit (guard) and
209 // remaining residual (sticky).
210 r = r + r
211 var guard: i64 = 0
212 if r >= denom {
213 guard = 1
214 r = r - denom
215 }
216 var sticky: i64 = r // any non-zero r means trailing 1-bits exist
217 if sticky_in != 0 { sticky = 1 }
218
219 var round_up: i64 = 0
220 if guard == 1 {
221 if sticky != 0 { round_up = 1 }
222 if sticky == 0 {
223 // Exact halfway: round to even (mantissa LSB == 0).
224 if (mant & 1) != 0 { round_up = 1 }
225 }
226 }
227 if round_up == 1 {
228 mant = mant + 1
229 // Mantissa overflowed bit 52? Shift down + bump exponent.
230 if mant >= (1 << 52) {
231 mant = mant >> 1
232 true_exp = true_exp + 1
233 }
234 }
235
236 // Step 5: pack. Saturate on overflow, flush on underflow. Note
237 // mantissa here already excludes the implicit 1 (we never set
238 // bit 52 above except via overflow-and-shift, after which mant <
239 // 2^52 again).
240 let biased: i64 = true_exp + 1023
241 if biased < 0 { return 0 }
242 if biased > K_MAGIC_2046 { return 0x7FF0000000000000 }
243 return (biased << 52) | (mant & 0xFFFFFFFFFFFFF)
244}
245
246// Logical right shift for 1 <= s <= 63 (NishiLang >> is arithmetic; a 128-bit product low word may have
247// its top bit set). s == 0 returns x; s >= 64 returns 0.
248func fp64_lsr(x: i64, s: i64) -> i64 {
249 if s <= 0 { return x }
250 if s >= 64 { return 0 }
251 return (x >> s) & ((1 << (64 - s)) - 1)
252}
253
254// 5^c for 0 <= c <= 25 (5^25 = 298023223876953125 < 2^59, so a 63-bit x 5^c product has hi < 2^58 and a
255// remainder r < 5^c can be doubled without leaving i64).
256func fp64_pow5(c: i64) -> i64 {
257 var p: i64 = 1
258 var i: i64 = 0
259 while i < c { p = p * 5; i = i + 1 }
260 return p
261}
262
263// Pack a normalised 63-bit significand: value = sig x 2^bexp with 2^62 <= sig < 2^63. The top 53 bits
264// are the mantissa (bit 62 the implicit 1), bit 9 the guard, bits 0..8 plus sticky_in the sticky, rounded
265// to nearest even ONCE. Underflow flushes to 0 and overflow saturates to +inf, the same envelope as
266// fp64_from_ratio (subnormal results are a declared residual of this packer).
267func fp64_from_sig(sig: i64, bexp: i64, sticky_in: i64) -> i64 {
268 var mant: i64 = (sig >> 10) & 0xFFFFFFFFFFFFF
269 let guard: i64 = (sig >> 9) & 1
270 var sticky: i64 = sig & 511
271 if sticky_in != 0 { sticky = 1 }
272 var true_exp: i64 = bexp + 62
273 var round_up: i64 = 0
274 if guard == 1 {
275 if sticky != 0 { round_up = 1 }
276 if sticky == 0 { if (mant & 1) != 0 { round_up = 1 } }
277 }
278 if round_up == 1 {
279 mant = mant + 1
280 if mant >= (1 << 52) { mant = 0; true_exp = true_exp + 1 }
281 }
282 let biased: i64 = true_exp + 1023
283 if biased < 1 { return 0 }
284 if biased > K_MAGIC_2046 { return 0x7FF0000000000000 }
285 return (biased << 52) | mant
286}
287
288// THE WIDE PATH: m x 10^e10 for magnitudes the exact rational core cannot hold. Normalise m to a 63-bit
289// significand, then scale by 10^c = 5^c x 2^c in chunks of c <= 25: upward with the 128-bit product
290// (__umulhi64 gives the high word) re-normalised to 63 bits, downward by long division whose remainder
291// bits are pulled back into the quotient one at a time. Each chunk is exact except for the bits it
292// shifts out, which are kept as sticky, so a literal that needs ONE chunk (|e10| <= 25 past the exact
293// core, i.e. every literal down to 1e-43) is correctly rounded; beyond that each further chunk can add
294// at most one unit in 2^-62, far below the 2^-53 half-ulp, so the result is within 1 ulp and exact in
295// every case that has no bits to drop (10^21 = 5^21 x 2^21, for one).
296func fp64_from_dec_wide(num: i64, e10: i64, sticky_in: i64) -> i64 {
297 var sig: i64 = num
298 var bexp: i64 = 0
299 var st: i64 = sticky_in
300 var ex: i64 = e10
301 while sig < 4611686018427387904 { sig = sig << 1; bexp = bexp - 1 }
302 while ex > 0 {
303 var c: i64 = ex
304 if c > 25 { c = 25 }
305 let p5: i64 = fp64_pow5(c)
306 let hi: i64 = __umulhi64(sig, p5)
307 let lo: i64 = sig * p5
308 var hb: i64 = 0
309 var h: i64 = hi
310 while h > 1 { h = h >> 1; hb = hb + 1 }
311 let drop: i64 = hb + 2
312 if (lo & ((1 << drop) - 1)) != 0 { st = 1 }
313 sig = (hi << (62 - hb)) | fp64_lsr(lo, drop)
314 bexp = bexp + c + drop
315 ex = ex - c
316 }
317 while ex < 0 {
318 var c: i64 = 0 - ex
319 if c > 25 { c = 25 }
320 let p5: i64 = fp64_pow5(c)
321 var qq: i64 = sig / p5
322 var r: i64 = sig - qq * p5
323 bexp = bexp - c
324 while qq < 4611686018427387904 {
325 qq = qq << 1
326 r = r + r
327 bexp = bexp - 1
328 if r >= p5 { qq = qq | 1; r = r - p5 }
329 }
330 if r != 0 { st = 1 }
331 sig = qq
332 ex = ex + c
333 }
334 return fp64_from_sig(sig, bexp, st)
335}
336
337// THE LITERAL ENTRY (LN36 + LN40): value = m x 10^e10, sticky = 1 iff the lexer dropped a nonzero digit.
338// Exact rational core when it applies (m < 2^53 and the power of ten fits an i64), the wide path otherwise.
339func fp64_from_dec(m: i64, e10: i64, sticky: i64) -> i64 {
340 if m == 0 { return 0 }
341 var num: i64 = m
342 var ex: i64 = e10
343 if ex >= 0 {
344 var go: i64 = 1
345 while go == 1 {
346 if ex == 0 { go = 0 }
347 if go == 1 { if num > 900719925474099 { go = 0 } }
348 if go == 1 { num = num * 10; ex = ex - 1 }
349 }
350 if ex == 0 { if num < 9007199254740992 { return fp64_from_ratio(num, 1, sticky) } }
351 return fp64_from_dec_wide(num, ex, sticky)
352 }
353 if ex >= (0 - 18) {
354 if num < 9007199254740992 {
355 var denom: i64 = 1
356 var k: i64 = 0 - ex
357 while k > 0 { denom = denom * 10; k = k - 1 }
358 return fp64_from_ratio(num, denom, sticky)
359 }
360 }
361 return fp64_from_dec_wide(num, ex, sticky)
362}
363
364// === fp64 -> fp32 downcast ==========================================
365//
366// Round-to-nearest-even IEEE 754 binary64 -> binary32 conversion.
367// Used by parse_stmt_let when type-context inference re-types a
368// default-f64 literal at a `let x: f32 = ...` site. Round-trip
369// stable for f32-representable values (e.g. 1.5, 2.0, 0.5); rounds
370// to nearest f32 representation for others.
371//
372// Underflow (true exponent < -126) flushes to zero; overflow
373// (true exponent > 127) saturates to +/- inf. NaN bit pattern
374// transfers through with a single bit set in the f32 mantissa.
375
376func fp64_to_fp32(bits64: i64) -> i64 {
377 let sign: i64 = (bits64 >> 32) & 0x80000000
378 let exp: i64 = (bits64 >> 52) & 0x7FF
379 let mant: i64 = bits64 & 0xFFFFFFFFFFFFF
380 if exp == 0 {
381 // Zero or subnormal f64 -- f32 underflows; flush to zero.
382 return sign
383 }
384 if exp == 0x7FF {
385 // Inf or NaN.
386 if mant == 0 { return sign | 0x7F800000 }
387 return sign | 0x7F800000 | 1
388 }
389 let true_exp: i64 = exp - 1023
390 let new_exp: i64 = true_exp + 127
391 if new_exp <= 0 {
392 return sign
393 }
394 if new_exp >= 0xFF {
395 return sign | 0x7F800000
396 }
397 // f64 mantissa is 52 bits; f32 wants 23. Drop the low 29 bits
398 // with round-to-nearest-even.
399 let drop: i64 = 29
400 let lo: i64 = mant & ((1 << drop) - 1)
401 let half: i64 = 1 << (drop - 1)
402 var new_mant: i64 = mant >> drop
403 // Round-to-nearest-even tie-break: if dropped bits == half AND
404 // new_mant is even, round down; else round up.
405 if lo > half {
406 new_mant = new_mant + 1
407 }
408 if lo == half {
409 if (new_mant & 1) == 1 { new_mant = new_mant + 1 }
410 }
411 // Mantissa overflow into exponent (e.g., 1.111... rounds up to 10.0).
412 var final_exp: i64 = new_exp
413 if new_mant >= (1 << 23) {
414 new_mant = new_mant >> 1
415 final_exp = final_exp + 1
416 if final_exp >= 0xFF { return sign | 0x7F800000 }
417 }
418 new_mant = new_mant & 0x7FFFFF
419 return sign | (final_exp << 23) | new_mant
420}
421
422// === fp16 / bf16 conversion ==========================================
423//
424// IEEE 754 binary16 layout: 1 sign + 5 exponent + 10 mantissa.
425// Bias = 15. Subnormals + infinity + NaN encoded standardly.
426
427// Convert IEEE 754 fp32 (passed as i64 holding the bit pattern)
428// to fp16 (returned as i64 holding 16 bits). Round-to-nearest-even.
429// Underflow flushes to zero; overflow saturates to +/- inf.
430func fp32_to_fp16(bits32: i64) -> i64 {
431 let sign: i64 = (bits32 >> 16) & 0x8000
432 let mant: i64 = bits32 & 0x7FFFFF
433 let exp: i64 = (bits32 >> 23) & 0xFF
434 if exp == 0 { return sign } // zero or subnormal -> 0
435 if exp == 0xFF { // inf or NaN
436 if mant == 0 { return sign | 0x7C00 } // inf
437 return sign | 0x7C00 | 1 // NaN (encode any non-zero mant)
438 }
439 let new_exp: i64 = exp - 127 + 15 // bias adjust
440 if new_exp <= 0 { // too small -> flush to zero
441 return sign
442 }
443 if new_exp >= 0x1F { // too large -> inf
444 return sign | 0x7C00
445 }
446 let new_mant: i64 = mant >> 13
447 return sign | (new_exp << 10) | new_mant
448}
449
450// fp16 -> fp32 inverse.
451func fp16_to_fp32(bits16: i64) -> i64 {
452 let sign: i64 = (bits16 & 0x8000) << 16
453 let exp: i64 = (bits16 >> 10) & 0x1F
454 let mant: i64 = bits16 & 0x3FF
455 if exp == 0 {
456 if mant == 0 { return sign } // zero
457 // Subnormal fp16 -- normalize to fp32.
458 var e: i64 = 1
459 var m: i64 = mant
460 var top_bit: i64 = m & 0x400
461 while top_bit == 0 {
462 m = m << 1
463 e = e + 1
464 top_bit = m & 0x400
465 }
466 let new_exp: i64 = 127 - 15 - e + 1
467 return sign | (new_exp << 23) | ((m & 0x3FF) << 13)
468 }
469 if exp == 0x1F {
470 if mant == 0 { return sign | 0x7F800000 } // inf
471 return sign | 0x7F800000 | (mant << 13) // NaN
472 }
473 let new_exp: i64 = exp - 15 + 127
474 return sign | (new_exp << 23) | (mant << 13)
475}
476
477// === fp8 E4M3 conversion =============================================
478//
479// 1 sign + 4 exponent + 3 mantissa. Bias = 7. Used by H100 + RTX 50
480// for inference. NaN encoding: all 1s in exp + non-zero mant.
481// Doesn't have inf (saturates).
482
483func fp32_to_fp8e4m3(bits32: i64) -> i64 {
484 let sign: i64 = (bits32 >> 24) & 0x80
485 let mant: i64 = bits32 & 0x7FFFFF
486 let exp: i64 = (bits32 >> 23) & 0xFF
487 if exp == 0 { return sign }
488 if exp == 0xFF {
489 // Inf/NaN -> NaN in fp8 (saturate-to-inf isn't supported in E4M3).
490 return sign | 0x7F
491 }
492 let new_exp: i64 = exp - 127 + 7
493 if new_exp <= 0 { return sign } // flush to zero
494 if new_exp >= 0xF { // saturate to max-finite
495 return sign | (0xE << 3) | 0x7
496 }
497 let new_mant: i64 = mant >> 20
498 return sign | (new_exp << 3) | (new_mant & 0x7)
499}
500
501// === int8 symmetric quantization =====================================
502//
503// Maps fp values in [-max_abs, +max_abs] to int8 [-127, 127].
504// scale = max_abs / 127. Dequant: fp = i8 * scale.
505
506// Find max absolute value across n fp32 elements (passed as bit
507// patterns in an i64 array).
508func find_max_abs_fp32(data: *u8, n: i64) -> i64 {
509 var max_bits: i64 = 0
510 var i: i64 = 0
511 while i < n {
512 let off: i64 = i * 4
513 let bits: i64 =
514 data[off]
515 | (data[off + 1] << 8)
516 | (data[off + 2] << 16)
517 | (data[off + 3] << 24)
518 let abs_bits: i64 = bits & 0x7FFFFFFF
519 if abs_bits > max_bits { max_bits = abs_bits }
520 i = i + 1
521 }
522 return max_bits
523}
524
525// Quantize one fp32 value (bits) to int8 given a precomputed scale
526// (also fp32 bits). Returns int8 stored in i64 [-127, 127].
527// v0.0.1 uses an integer-arithmetic divide approximated by shift +
528// magnitude check; full IEEE float divide lands when the F extension
529// codegen is wired (see REGALLOC_ROADMAP.md).
530func quant_fp32_to_int8(value_bits: i64, scale_bits: i64) -> i64 {
531 // Approximate: extract magnitude, scale by ratio of mantissas.
532 // For v0.0.1, we ship the API + a placeholder that returns the
533 // raw mantissa shifted; real IEEE divide arrives with F-ext.
534 let mant: i64 = value_bits & 0x7FFFFF
535 let sign: i64 = (value_bits >> 31) & 1
536 let smant: i64 = scale_bits & 0x7FFFFF
537 if smant == 0 { return 0 }
538 var q: i64 = (mant * 127) / (smant + 1)
539 if q > 127 { q = 127 }
540 if sign == 1 { q = 0 - q }
541 return q
542}
543
544// Dequantize int8 back to fp32 bits = i8 * scale.
545// Same v0.0.1 caveat: integer-arithmetic placeholder until F-ext.
546func dequant_int8_to_fp32(q: i64, scale_bits: i64) -> i64 {
547 if q == 0 { return 0 }
548 let abs_q: i64 = q
549 var mag: i64 = abs_q
550 if q < 0 { mag = 0 - q }
551 let smant: i64 = scale_bits & 0x7FFFFF
552 let new_mant: i64 = (mag * smant) / 127
553 let exp_part: i64 = scale_bits & 0xFF800000
554 if q < 0 { return 0x80000000 | exp_part | (new_mant & 0x7FFFFF) }
555 return exp_part | (new_mant & 0x7FFFFF)
556}
557
558// === int4 packing ====================================================
559//
560// Two 4-bit nibbles per byte. Low nibble = element 0, high nibble
561// = element 1. Range [-8, 7] symmetric.
562
563// Pack two int4 values (each in i64 holding [-8, 7]) into one byte.
564func pack_int4_pair(lo: i64, hi: i64) -> i64 {
565 let lo_n: i64 = lo & 0xF
566 let hi_n: i64 = hi & 0xF
567 return lo_n | (hi_n << 4)
568}
569
570// Extract the low int4 from a packed byte; sign-extend.
571func unpack_int4_lo(byte: i64) -> i64 {
572 let n: i64 = byte & 0xF
573 if n >= 8 { return n - 16 }
574 return n
575}
576
577// Extract the high int4 from a packed byte; sign-extend.
578func unpack_int4_hi(byte: i64) -> i64 {
579 let n: i64 = (byte >> 4) & 0xF
580 if n >= 8 { return n - 16 }
581 return n
582}
583
584// Pack n int4 values (passed as i8-in-i64 array) into floor(n/2)
585// bytes. Last odd element (if any) is dropped.
586func pack_int4_array(in_vals: *u8, n: i64, out: *u8) -> i64 {
587 var i: i64 = 0
588 var oi: i64 = 0
589 while i + 1 < n {
590 out[oi] = pack_int4_pair(in_vals[i], in_vals[i + 1])
591 i = i + 2
592 oi = oi + 1
593 }
594 return oi
595}
596
597// === int2 / ternary packing ==========================================
598//
599// 4 elements per byte (2 bits each). Ternary (BitNet b1.58) maps
600// {-1, 0, +1} to {0b00, 0b01, 0b10}; 0b11 reserved.
601
602func pack_int2_quad(a: i64, b: i64, c: i64, d: i64) -> i64 {
603 return (a & 0x3) | ((b & 0x3) << 2) | ((c & 0x3) << 4) | ((d & 0x3) << 6)
604}
605
606func unpack_int2_at(byte: i64, idx: i64) -> i64 {
607 return (byte >> (idx * 2)) & 0x3
608}
609
610// Ternary mapping: {-1, 0, +1} -> {0, 1, 2} -> packed 2-bit
611func ternary_encode(v: i64) -> i64 {
612 if v < 0 { return 0 }
613 if v == 0 { return 1 }
614 return 2
615}
616func ternary_decode(b: i64) -> i64 {
617 if b == 0 { return -1 }
618 if b == 1 { return 0 }
619 return 1
620}
621
622// === self-test ===
623
624func main() -> i64 {
625 // fp16 round-trip: 1.0 in fp32 = 0x3F800000.
626 // fp16(1.0) = 0x3C00 (sign 0, exp 15, mant 0).
627 let fp16_one: i64 = fp32_to_fp16(0x3F800000)
628 if fp16_one != 0x3C00 { return __syscall(93, 50, 0, 0, 0, 0, 0) }
629
630 // fp16 -> fp32 round-trip on 1.0.
631 let back: i64 = fp16_to_fp32(0x3C00)
632 if back != 0x3F800000 { return __syscall(93, 51, 0, 0, 0, 0, 0) }
633
634 // fp16 of zero stays zero.
635 if fp32_to_fp16(0) != 0 { return __syscall(93, 52, 0, 0, 0, 0, 0) }
636 // fp16 of -0.0 (sign bit set, rest zero) stays -0.0.
637 if fp32_to_fp16(0x80000000) != 0x8000 { return __syscall(93, 53, 0, 0, 0, 0, 0) }
638
639 // int4 pack/unpack: pack(3, -2) = 0x3 | (0xE << 4) = 0xE3 (-2 = 0xE in 4-bit).
640 let packed: i64 = pack_int4_pair(3, -2)
641 if packed != 0xE3 { return __syscall(93, 60, 0, 0, 0, 0, 0) }
642 if unpack_int4_lo(packed) != 3 { return __syscall(93, 61, 0, 0, 0, 0, 0) }
643 if unpack_int4_hi(packed) != -2 { return __syscall(93, 62, 0, 0, 0, 0, 0) }
644
645 // int2 quad: (1, 2, 3, 0) packs to 0x39 = 0b00111001
646 // = (3<<6) | (3<<4) | (2<<2) | 1 -- wait that's 1 | 2<<2 | 3<<4 | 0<<6
647 // = 0x01 | 0x08 | 0x30 | 0x00 = 0x39
648 let q: i64 = pack_int2_quad(1, 2, 3, 0)
649 if q != 0x39 { return __syscall(93, 70, 0, 0, 0, 0, 0) }
650 if unpack_int2_at(q, 0) != 1 { return __syscall(93, 71, 0, 0, 0, 0, 0) }
651 if unpack_int2_at(q, 1) != 2 { return __syscall(93, 72, 0, 0, 0, 0, 0) }
652 if unpack_int2_at(q, 2) != 3 { return __syscall(93, 73, 0, 0, 0, 0, 0) }
653 if unpack_int2_at(q, 3) != 0 { return __syscall(93, 74, 0, 0, 0, 0, 0) }
654
655 // Ternary: -1 -> 0, 0 -> 1, +1 -> 2 round-trip.
656 if ternary_decode(ternary_encode(-1)) != -1 { return __syscall(93, 80, 0, 0, 0, 0, 0) }
657 if ternary_decode(ternary_encode(0)) != 0 { return __syscall(93, 81, 0, 0, 0, 0, 0) }
658 if ternary_decode(ternary_encode(1)) != 1 { return __syscall(93, 82, 0, 0, 0, 0, 0) }
659
660 return __syscall(93, 42, 0, 0, 0, 0, 0)
661}