diff --git a/AGENTS.md b/AGENTS.md index 30ca974..fc5c460 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,3 +1,5 @@ + + # quantize A simple Rust library for learning and experimenting with quantization techniques. diff --git a/Cargo.toml b/Cargo.toml index 7d0232e..2c3ee0d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,8 @@ [package] name = "quantize" -version = "0.1.0" +version = "0.2.0" edition = "2021" +rust-version = "1.87" description = "A simple quantization library." license = "MIT" repository = "https://github.com/aksheyd/quantize" @@ -21,6 +22,10 @@ exclude = [ all-features = true rustdoc-args = ["--cfg", "docsrs"] +[features] +default = ["std"] +std = [] + [dependencies] half = "2" @@ -28,6 +33,9 @@ half = "2" candle-core = "0.9" rand = "0.8" +[lints.clippy] +all = { level = "deny", priority = 10 } + [[example]] name = "ch01_simple" path = "chapters/ch01_simple.rs" @@ -48,6 +56,14 @@ path = "chapters/ch04_block.rs" name = "ch05_asymmetric" path = "chapters/ch05_asymmetric.rs" +[[example]] +name = "ch06_adaptive" +path = "chapters/ch06_adaptive.rs" + +[[example]] +name = "ch07_learned" +path = "chapters/ch07_learned.rs" + [[example]] name = "compare" path = "benchmarks/compare.rs" @@ -56,3 +72,6 @@ path = "benchmarks/compare.rs" name = "update_readme" path = "benchmarks/update_readme.rs" +[[example]] +name = "throughput" +path = "benchmarks/throughput.rs" diff --git a/Justfile b/Justfile index 1e75d70..aaaab39 100644 --- a/Justfile +++ b/Justfile @@ -2,7 +2,13 @@ c: cargo run --release --example compare l: - cargo fmt && cargo clippy -- -D warnings + cargo fmt && cargo clippy --all-targets -- -D warnings + +t: + cargo test + +b: + cargo run --release --example throughput ur: - cargo run --release --example update_readme \ No newline at end of file + cargo run --release --example update_readme diff --git a/README.md b/README.md index 9ef63cd..9fd0a77 100644 --- a/README.md +++ b/README.md @@ -5,32 +5,56 @@ A simple quantization library ## use as a library ```rust -use quantize::{quantize, dequantize}; +use quantize::{adaptive, asymmetric, quantize, Scheme}; let weights = [0.42_f32, -0.10, 0.70, -0.50]; -let (scales, codes) = quantize::(&weights); -let back = dequantize::<_, 32>(&scales, &codes); + +let q = quantize::(&weights).unwrap(); +let _ = asymmetric::quantize::(&weights).unwrap(); +let _ = adaptive::quantize::(&weights, 0.001).unwrap(); +let _ = Scheme::Q4_32.quantize::(&weights).unwrap(); + +let back = q.dequantize(); +let _ = q.dot(&weights); ``` --- ### comparison +1024x1024 matrix, 50 iterations. + +#### quality + ``` cargo run --release --example compare ``` - +Reconstruct, then matmul. -| method | bits/elt | mse (mean) | cosine (mean) | -| --- | ---: | ---: | ---: | -| quantize 4b×32 | 4.5 | 0.008339 | 0.995434 | -| candle Q4_0 | 4.5 | 0.007543 | 0.995825 | -| quantize 5b×32 | 5.5 | 0.001808 | 0.998995 | -| candle Q5_0 | 5.5 | 0.001730 | 0.999037 | -| quantize 8b×32 | 8.5 | 0.000025 | 0.999986 | -| candle Q8_0 | 8.5 | 0.000025 | 0.999986 | + -_matrix size: 128x128, runs: 10_ +| bits/value | quantize mse | candle mse | +| ---: | ---: | ---: | +| 4.5 | 0.066669 | 0.060276 | +| 5.5 | 0.014429 | 0.013859 | +| 8.5 | 0.000201 | 0.000201 | + +#### speed + +``` +cargo run --release --example throughput +``` + +Pack and unpack. Apple M4. + +| kernel | quantize ns/value | candle ns/value | +| --- | ---: | ---: | +| 4-bit quant | 0.29 | 0.46 | +| 8-bit quant | 0.24 | 0.30 | +| 4-bit dequant | 0.09 | 0.29 | +| 8-bit dequant | 0.07 | 0.26 | + +_ns = nanosecond, a billionth of a second. smaller is faster._ diff --git a/benchmarks/compare.rs b/benchmarks/compare.rs index 5de9013..9469cdf 100644 --- a/benchmarks/compare.rs +++ b/benchmarks/compare.rs @@ -1,12 +1,12 @@ //! Quantization comparison playground. //! //! Runs `RUNS` fresh matmuls for every method registered in -//! `harness/methods.rs` and prints MSE / cosine as mean ± std. +//! `harness/methods.rs` and prints mean MSE. mod harness; use harness::{Comparison, Harness}; -const MATRIX_SIZE: usize = 128; +const MATRIX_SIZE: usize = 1024; const RUNS: usize = 50; fn main() -> candle_core::Result<()> { @@ -16,26 +16,11 @@ fn main() -> candle_core::Result<()> { } fn print_report(r: &Comparison) { - println!("matrix_size = {}, runs = {}\n", r.matrix_size, r.runs); - println!( - "{:<16}{:>10}{:>22}{:>22}", - "method", "bits/elt", "mse (mean ± std)", "cosine (mean ± std)", - ); - println!("{:-<16}{:->10}{:->22}{:->22}", "", "", "", ""); - - let mut prev_bits = 0.0_f32; + println!("matrix_size = {MATRIX_SIZE}, runs = {RUNS}\n"); + println!("{:<12}{:>14}", "bits/value", "mse"); + println!("{:-<12}{:->14}", "", ""); for m in &r.methods { - // Blank line between precision tiers (>1 bit jump) for easier scanning. - if prev_bits > 0.0 && m.bits_per_element - prev_bits > 1.0 { - println!(); - } - prev_bits = m.bits_per_element; - let mse = format!("{:.6} ± {:.6}", m.stats.mse_mean, m.stats.mse_std); - let cos = format!("{:.6} ± {:.6}", m.stats.cosine_mean, m.stats.cosine_std); - println!( - "{:<16}{:>10.2}{:>22}{:>22}", - m.name, m.bits_per_element, mse, cos - ); + println!("{:<12.1}{:>14.6}", m.bits_per_element, m.mse); } - println!("\nbits/elt = total storage per element (data + amortized scale)."); + println!("\nbits/value = storage for one number, including its scale."); } diff --git a/benchmarks/harness/methods.rs b/benchmarks/harness/methods.rs index 4bebab7..72fc001 100644 --- a/benchmarks/harness/methods.rs +++ b/benchmarks/harness/methods.rs @@ -29,7 +29,6 @@ impl Bits { } pub(super) struct Method { - pub name: &'static str, pub bits_per_element: Bits, pub eval: EvalFn, } @@ -39,11 +38,11 @@ pub(super) fn methods() -> Vec { use Bits::*; use GgmlDType::{Q4_0, Q5_0, Q8_0}; vec![ - Method { name: "quantize 4b×32", bits_per_element: Quantize(4), eval: eval_quantize::<4> }, - Method { name: "candle Q4_0", bits_per_element: Candle(Q4_0), eval: eval_q4_0 }, - Method { name: "quantize 5b×32", bits_per_element: Quantize(5), eval: eval_quantize::<5> }, - Method { name: "candle Q5_0", bits_per_element: Candle(Q5_0), eval: eval_q5_0 }, - Method { name: "quantize 8b×32", bits_per_element: Quantize(8), eval: eval_quantize::<8> }, - Method { name: "candle Q8_0", bits_per_element: Candle(Q8_0), eval: eval_q8_0 }, + Method { bits_per_element: Quantize(4), eval: eval_quantize::<4> }, + Method { bits_per_element: Candle(Q4_0), eval: eval_q4_0 }, + Method { bits_per_element: Quantize(5), eval: eval_quantize::<5> }, + Method { bits_per_element: Candle(Q5_0), eval: eval_q5_0 }, + Method { bits_per_element: Quantize(8), eval: eval_quantize::<8> }, + Method { bits_per_element: Candle(Q8_0), eval: eval_q8_0 }, ] } diff --git a/benchmarks/harness/metrics.rs b/benchmarks/harness/metrics.rs index 7ed7c65..309fdd9 100644 --- a/benchmarks/harness/metrics.rs +++ b/benchmarks/harness/metrics.rs @@ -1,8 +1,4 @@ //! Error metrics comparing a predicted vector against the expected one. -//! The smaller the error, the better the prediction. -//! -//! predicted refers to the output of the model, -//! while expected is the ground truth. pub(super) fn mse(predicted: &[f32], expected: &[f32]) -> f32 { predicted @@ -12,14 +8,3 @@ pub(super) fn mse(predicted: &[f32], expected: &[f32]) -> f32 { .sum::() / predicted.len() as f32 } - -pub(super) fn cosine(predicted: &[f32], expected: &[f32]) -> f32 { - let dot: f32 = predicted.iter().zip(expected).map(|(p, e)| p * e).sum(); - let norm_predicted: f32 = predicted.iter().map(|p| p * p).sum::().sqrt(); - let norm_expected: f32 = expected.iter().map(|e| e * e).sum::().sqrt(); - if norm_predicted == 0.0 || norm_expected == 0.0 { - 0.0 - } else { - dot / (norm_predicted * norm_expected) - } -} diff --git a/benchmarks/harness/mod.rs b/benchmarks/harness/mod.rs index b371dc6..f87ab93 100644 --- a/benchmarks/harness/mod.rs +++ b/benchmarks/harness/mod.rs @@ -1,5 +1,5 @@ //! Test harness. `Harness::new(matrix_size, runs).run()` generates `runs` -//! fresh random matrices, evaluates every method, returns mean ± std stats. +//! fresh random matrices, evaluates every method, returns mean MSE. mod methods; mod metrics; @@ -7,27 +7,15 @@ mod new; mod quant; mod run; mod sample; -mod stats; use candle_core::{Device, Tensor}; -#[allow(dead_code)] -pub struct Stats { - pub mse_mean: f32, - pub mse_std: f32, - pub cosine_mean: f32, - pub cosine_std: f32, -} - pub struct MethodReport { - pub name: &'static str, pub bits_per_element: f32, - pub stats: Stats, + pub mse: f32, } pub struct Comparison { - pub matrix_size: usize, - pub runs: usize, pub methods: Vec, } @@ -37,8 +25,6 @@ pub struct Harness { device: Device, } -// One run's worth of fresh data. Submodules can see the private fields -// because Rust lets descendant modules access the parent's private items. struct Sample { matrix_size: usize, matrix_a: Vec, diff --git a/benchmarks/harness/quant.rs b/benchmarks/harness/quant.rs index 94c53fa..591b78b 100644 --- a/benchmarks/harness/quant.rs +++ b/benchmarks/harness/quant.rs @@ -8,11 +8,12 @@ use candle_core::{ Device, Result, Tensor, }; use half::f16; -use quantize::{dequantize, quantize}; +use quantize::quantize; fn roundtrip_block(values: &[f32]) -> Vec { - let (scales, codes) = quantize::(values); - dequantize::<_, 32>(&scales, &codes) + quantize::(values) + .expect("valid bits/block") + .dequantize() } fn matmul(a: Vec, b: Vec, n: usize, d: &Device) -> Result> { diff --git a/benchmarks/harness/run.rs b/benchmarks/harness/run.rs index 74d2a36..1f17bb2 100644 --- a/benchmarks/harness/run.rs +++ b/benchmarks/harness/run.rs @@ -1,12 +1,7 @@ //! Outer loop: for each of `runs` fresh samples, evaluate every method, -//! accumulate (mse, cosine), then collapse each method's samples to mean ± std. +//! then collapse each method's MSE samples to a mean. -use super::{ - methods::methods, - metrics::{cosine, mse}, - stats::mean_std, - Comparison, Harness, MethodReport, Stats, -}; +use super::{methods::methods, metrics::mse, Comparison, Harness, MethodReport}; use candle_core::Result; impl Harness { @@ -14,32 +9,27 @@ impl Harness { let methods = methods(); let elements = self.matrix_size * self.matrix_size; let mut mses: Vec> = vec![Vec::with_capacity(self.runs); methods.len()]; - let mut coss: Vec> = vec![Vec::with_capacity(self.runs); methods.len()]; for _ in 0..self.runs { let s = self.sample()?; for (i, m) in methods.iter().enumerate() { let predicted = (m.eval)(&s, &self.device)?; mses[i].push(mse(&predicted, &s.ground_truth)); - coss[i].push(cosine(&predicted, &s.ground_truth)); } } - #[rustfmt::skip] - let methods = methods.iter().enumerate().map(|(i, m)| { - let (mse_mean, mse_std) = mean_std(&mses[i]); - let (cosine_mean, cosine_std) = mean_std(&coss[i]); - MethodReport { - name: m.name, - bits_per_element: m.bits_per_element.evaluate(elements), - stats: Stats { mse_mean, mse_std, cosine_mean, cosine_std }, - } - }).collect(); + let methods = methods + .iter() + .enumerate() + .map(|(i, m)| { + let n = mses[i].len() as f32; + MethodReport { + bits_per_element: m.bits_per_element.evaluate(elements), + mse: mses[i].iter().sum::() / n, + } + }) + .collect(); - Ok(Comparison { - matrix_size: self.matrix_size, - runs: self.runs, - methods, - }) + Ok(Comparison { methods }) } } diff --git a/benchmarks/harness/stats.rs b/benchmarks/harness/stats.rs deleted file mode 100644 index e1a548d..0000000 --- a/benchmarks/harness/stats.rs +++ /dev/null @@ -1,8 +0,0 @@ -//! Sample mean and (population) standard deviation of a slice of floats. - -pub(super) fn mean_std(values: &[f32]) -> (f32, f32) { - let n = values.len() as f32; - let mean = values.iter().sum::() / n; - let var = values.iter().map(|v| (v - mean).powi(2)).sum::() / n; - (mean, var.sqrt()) -} diff --git a/benchmarks/throughput.rs b/benchmarks/throughput.rs new file mode 100644 index 0000000..5e7b2bd --- /dev/null +++ b/benchmarks/throughput.rs @@ -0,0 +1,83 @@ +//! Quantize / dequantize / fused-dot throughput vs candle. +//! +//! Run: `cargo run --release --example throughput` + +use candle_core::{ + quantized::{GgmlDType, QTensor}, + Device, Tensor, +}; +use half::f16; +use quantize::quantize; +use std::hint::black_box; +use std::time::Instant; + +const SIDE: usize = 1024; +const N: usize = SIDE * SIDE; +const ITERS: usize = 50; + +fn main() -> candle_core::Result<()> { + let mut seed = 0x1234_5678u32; + let values: Vec = (0..N) + .map(|_| { + seed = seed.wrapping_mul(1664525).wrapping_add(1013904223); + (seed as f32 / u32::MAX as f32) * 2.0 - 1.0 + }) + .collect(); + + println!("n = {N} ({SIDE}x{SIDE}), iters = {ITERS}\n"); + println!("{:<22}{:>12}{:>14}", "kernel", "ns/elt", "GB/s (in)"); + println!("{:-<22}{:->12}{:->14}", "", "", ""); + + bench("quantize 4b×32", || { + black_box(quantize::(&values).unwrap()); + }); + bench("quantize 8b×32", || { + black_box(quantize::(&values).unwrap()); + }); + + let q4 = quantize::(&values).unwrap(); + let q8 = quantize::(&values).unwrap(); + let mut out = vec![0.0f32; N]; + bench("dequant 4b×32", || { + q4.dequantize_into(&mut out).unwrap(); + black_box(&out); + }); + bench("dequant 8b×32", || { + q8.dequantize_into(&mut out).unwrap(); + black_box(&out); + }); + bench("dot 8b×32", || { + black_box(q8.dot(&values).unwrap()); + }); + + let device = Device::Cpu; + let t = Tensor::from_vec(values.clone(), N, &device)?; + bench("candle Q4_0 quant", || { + black_box(QTensor::quantize(&t, GgmlDType::Q4_0).unwrap()); + }); + bench("candle Q8_0 quant", || { + black_box(QTensor::quantize(&t, GgmlDType::Q8_0).unwrap()); + }); + let cq4 = QTensor::quantize(&t, GgmlDType::Q4_0)?; + let cq8 = QTensor::quantize(&t, GgmlDType::Q8_0)?; + bench("candle Q4_0 dequant", || { + black_box(cq4.dequantize(&device).unwrap()); + }); + bench("candle Q8_0 dequant", || { + black_box(cq8.dequantize(&device).unwrap()); + }); + Ok(()) +} + +fn bench(name: &str, mut f: impl FnMut()) { + for _ in 0..4 { + f(); + } + let t0 = Instant::now(); + for _ in 0..ITERS { + f(); + } + let ns = t0.elapsed().as_secs_f64() * 1e9 / (ITERS as f64 * N as f64); + let gbs = (N as f64 * 4.0) / (ns * N as f64); // bytes in / time, in GB/s + println!("{name:<22}{ns:>12.3}{gbs:>14.2}"); +} diff --git a/benchmarks/update_readme.rs b/benchmarks/update_readme.rs index 75771d6..3da6d74 100644 --- a/benchmarks/update_readme.rs +++ b/benchmarks/update_readme.rs @@ -1,4 +1,4 @@ -//! Run the harness and rewrite the comparison table in README.md. +//! Run the harness and rewrite the quality table in README.md. //! //! Usage: `cargo run --release --example update_readme` @@ -9,25 +9,22 @@ use std::fs; const START: &str = ""; const END: &str = ""; -const MATRIX_SIZE: usize = 128; -const RUNS: usize = 10; +const MATRIX_SIZE: usize = 1024; +const RUNS: usize = 50; fn main() -> candle_core::Result<()> { let report = Harness::new(MATRIX_SIZE, RUNS)?.run()?; - let mut rows = String::from("| method | bits/elt | mse (mean) | cosine (mean) |\n"); - rows.push_str("| --- | ---: | ---: | ---: |\n"); - for m in &report.methods { + let mut rows = String::from("| bits/value | quantize mse | candle mse |\n"); + rows.push_str("| ---: | ---: | ---: |\n"); + for pair in report.methods.chunks(2) { + let Some([q, c]) = pair.get(..2) else { break }; rows.push_str(&format!( - "| {} | {:.1} | {:.6} | {:.6} |\n", - m.name, m.bits_per_element, m.stats.mse_mean, m.stats.cosine_mean, + "| {:.1} | {:.6} | {:.6} |\n", + q.bits_per_element, q.mse, c.mse, )); } - let table = format!( - "{START}\n\n{rows}\n_matrix size: {n}x{n}, runs: {runs}_\n\n{END}", - n = report.matrix_size, - runs = report.runs, - ); + let table = format!("{START}\n\n{rows}\n{END}"); let readme = fs::read_to_string("README.md").expect("README.md not found"); let start = readme.find(START).expect("missing start marker"); diff --git a/chapters/ch06_adaptive.rs b/chapters/ch06_adaptive.rs new file mode 100644 index 0000000..bf483ad --- /dev/null +++ b/chapters/ch06_adaptive.rs @@ -0,0 +1,49 @@ +//! # Chapter 6 — adaptive mixed precision +//! +//! **Previously** (`ch05_asymmetric`): a zero-point stretches the grid to +//! `[min, max]`, but every block still uses the same bit width. +//! +//! **Problem**: a nearly-constant block does not need 8 bits. Paying 8 bits +//! for a 0.002 range wastes memory that a wild block actually needs. +//! +//! **Fix**: pick bits from a *calibrated tolerance*. The half-step of the +//! integer grid must stay `<= tol`. Quiet blocks drop to 2–3 bits; busy +//! blocks keep 8. +//! +//! **Still wrong**: scale and zero-point are computed from min/max, not from +//! the reconstruction error we actually care about. They can be *learned*. +//! +//! Run it: `cargo run --release --example ch06_adaptive` + +fn choose_bits(range: f32, tol: f32) -> u32 { + if range <= 0.0 { + return 2; + } + for b in 2..=8 { + if range / ((1u32 << b) - 1) as f32 / 2.0 <= tol { + return b; + } + } + 8 +} + +fn main() { + let tol = 0.001_f32; + let tensor = [ + 0.500, 0.501, 0.499, 0.5005, // quiet + 0.10, 0.30, 0.70, 1.10, // busy + ]; + + println!("tol = {tol}\n"); + println!("{:>8} {:>6} {:>5}", "range", "bits", "block"); + println!("{:>8} {:>6} {:>5}", "-----", "----", "-----"); + for (i, block) in tensor.chunks(4).enumerate() { + let rmin = block.iter().copied().fold(f32::INFINITY, f32::min); + let rmax = block.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let bits = choose_bits(rmax - rmin, tol); + println!("{:>8.4} {:>6} {:>5}", rmax - rmin, bits, i); + } + + println!("\nSame tensor, two precisions. Chapter 7 (`ch07_learned`) treats"); + println!("scale and zero-point as a 1-neuron layer and fits them."); +} diff --git a/chapters/ch07_learned.rs b/chapters/ch07_learned.rs new file mode 100644 index 0000000..1da7cf8 --- /dev/null +++ b/chapters/ch07_learned.rs @@ -0,0 +1,75 @@ +//! # Chapter 7 — learned scale and zero-point +//! +//! **Previously** (`ch06_adaptive`): bit width follows a tolerance, but +//! scale/zero-point still come from min/max of the block. +//! +//! **Problem**: min/max fit the *range*, not the *error*. Outliers set the +//! scale; the rest of the block pays for it. +//! +//! **Fix**: freeze the integer codes and treat dequant as a line: +//! `value ≈ scale * code + offset`, with `offset = -scale * zero_point`. +//! One best-fit line per block. +//! +//! **Still wrong**: codes are frozen. Jointly learning codes is the next +//! research step, not this chapter. +//! +//! Run it: `cargo run --release --example ch07_learned` + +fn fit_scale_and_zero_point(values: &[f32], codes: &[i32]) -> (f32, f32) { + let count = values.len() as f32; + let mut sum_codes = 0.0; + let mut sum_values = 0.0; + let mut sum_code_squared = 0.0; + let mut sum_code_times_value = 0.0; + for (&value, &code) in values.iter().zip(codes) { + let code = code as f32; + sum_codes += code; + sum_values += value; + sum_code_squared += code * code; + sum_code_times_value += code * value; + } + let mean_code = sum_codes / count; + let mean_value = sum_values / count; + let code_spread = sum_code_squared - sum_codes * mean_code; + let scale = (sum_code_times_value - sum_codes * mean_value) / code_spread; + let offset = mean_value - scale * mean_code; + (scale, -offset / scale) +} + +fn mse(predicted: &[f32], expected: &[f32]) -> f32 { + predicted + .iter() + .zip(expected) + .map(|(a, b)| (a - b) * (a - b)) + .sum::() + / predicted.len() as f32 +} + +fn main() { + // One outlier dominates a min/max scale; a best-fit line ignores it better. + let values = [0.10_f32, 0.12, 0.11, 0.13, 4.0]; + let codes = [10, 12, 11, 13, 127]; + + let max = values.iter().copied().fold(0.0_f32, f32::max); + let minmax_scale = max / 127.0; + let minmax_back: Vec = codes + .iter() + .map(|&code| code as f32 * minmax_scale) + .collect(); + + let (scale, zero_point) = fit_scale_and_zero_point(&values, &codes); + let fitted: Vec = codes + .iter() + .map(|&code| scale * (code as f32 - zero_point)) + .collect(); + + println!( + "minmax scale={minmax_scale:.5} mse={:.5}", + mse(&minmax_back, &values) + ); + println!( + "fitted scale={scale:.5} zero_point={zero_point:.3} mse={:.5}", + mse(&fitted, &values) + ); + println!("\nDequant is a line. Fit the line; keep the codes."); +} diff --git a/src/asymmetric.rs b/src/asymmetric.rs deleted file mode 100644 index cfc6d52..0000000 --- a/src/asymmetric.rs +++ /dev/null @@ -1,110 +0,0 @@ -//! Asymmetric quantization. -//! -//! A scale and a zero-point per group. The integer codes are shifted -//! so the full range maps exactly to the group's actual min and max. -//! This avoids wasting codes on the unused side of zero. -//! -//! A "group" can be a fixed-size block or the entire tensor. -//! -//! When using blocks, you can supply a different bit width per block. -//! This is the "precision-aware" part: blocks with small ranges can use -//! fewer bits while still meeting a target error tolerance. - -use crate::Scale; - -fn max_int(bits: u32) -> i32 { - (1_i32 << (bits - 1)) - 1 -} -fn min_int(bits: u32) -> i32 { - -(1_i32 << (bits - 1)) -} - -/// Quantize into fixed-size blocks, using the given bit width for each block. -/// `bits` must have one entry per block (or fewer; missing blocks default to 8). -/// Returns `(scales, zero_points, codes)`. -pub fn quantize( - values: &[f32], - bits: &[u32], -) -> (Vec, Vec, Vec) { - let mut scales = Vec::with_capacity(values.len() / BLOCK + 1); - let mut zps = Vec::with_capacity(values.len() / BLOCK + 1); - let mut codes = Vec::with_capacity(values.len()); - for (i, chunk) in values.chunks(BLOCK).enumerate() { - let b = *bits.get(i).unwrap_or(&8); - let rmin = chunk.iter().copied().fold(f32::INFINITY, f32::min); - let rmax = chunk.iter().copied().fold(f32::NEG_INFINITY, f32::max); - let (sf, zpf) = if rmin >= rmax { - (1.0, 0.0) - } else { - let qmin = min_int(b) as f32; - let qmax = max_int(b) as f32; - let scale = (rmax - rmin) / (qmax - qmin); - let zp = qmin - rmin / scale; - (scale, zp) - }; - scales.push(S::from_f32(sf)); - zps.push(S::from_f32(zpf)); - let qmin = min_int(b); - let qmax = max_int(b); - codes.extend( - chunk - .iter() - .map(|&x| ((x / sf + zpf).round() as i32).clamp(qmin, qmax)), - ); - } - (scales, zps, codes) -} - -/// Reconstruct from per-block scales and zero-points. -pub fn dequantize( - scales: &[S], - zero_points: &[S], - codes: &[i32], -) -> Vec { - codes - .chunks(BLOCK) - .zip(scales.iter().zip(zero_points)) - .flat_map(|(blk, (&s, &zp))| { - let sf = s.to_f32(); - let zpf = zp.to_f32(); - blk.iter().map(move |&q| (q as f32 - zpf) * sf) - }) - .collect() -} - -/// Quantize the entire tensor with one scale and one zero-point. -/// `bits` sets the integer width for the whole tensor. -/// Returns `(scale, zero_point, codes)`. -pub fn quantize_tensor(values: &[f32], bits: u32) -> (S, S, Vec) { - if values.is_empty() { - let z = S::from_f32(0.0); - return (S::from_f32(1.0), z, vec![]); - } - let rmin = values.iter().copied().fold(f32::INFINITY, f32::min); - let rmax = values.iter().copied().fold(f32::NEG_INFINITY, f32::max); - let (sf, zpf) = if rmin >= rmax { - (1.0, 0.0) - } else { - let qmin = min_int(bits) as f32; - let qmax = max_int(bits) as f32; - let scale = (rmax - rmin) / (qmax - qmin); - let zp = qmin - rmin / scale; - (scale, zp) - }; - let s = S::from_f32(sf); - let zp = S::from_f32(zpf); - let qmin = min_int(bits); - let qmax = max_int(bits); - let codes = values - .iter() - .map(|&x| ((x / sf + zpf).round() as i32).clamp(qmin, qmax)) - .collect(); - (s, zp, codes) -} - -/// Reconstruct the entire tensor from a single scale and zero-point. -pub fn dequantize_tensor(scale: S, zero_point: S, codes: &[i32]) -> Vec { - let sf = scale.to_f32(); - let zpf = zero_point.to_f32(); - codes.iter().map(|&q| (q as f32 - zpf) * sf).collect() -} diff --git a/src/kernels/block.rs b/src/kernels/block.rs new file mode 100644 index 0000000..3758b05 --- /dev/null +++ b/src/kernels/block.rs @@ -0,0 +1,97 @@ +//! Arbitrary bit-width and asymmetric loops. 4/8-bit symmetric bypasses this. + +use crate::packed::Packed; +use crate::params::{asymmetric_params, largest_code, smallest_code, symmetric_scale}; + +use super::i4::pack_sym_i4; +use super::i8::pack_sym_i8; +use super::reduce::{abs_max, min_max}; + +pub(crate) fn quantize_sym_packed(values: &[f32], bits: u32, block: usize) -> (Vec, Packed) { + match bits { + 8 => pack_sym_i8(values, block), + 4 => pack_sym_i4(values, block), + _ => pack_sym_general(values, bits, block), + } +} + +fn pack_sym_general(values: &[f32], bits: u32, block: usize) -> (Vec, Packed) { + let mut scales = Vec::with_capacity(values.len().div_ceil(block)); + let mut codes = Vec::with_capacity(values.len()); + for chunk in values.chunks(block) { + let scale = symmetric_scale(abs_max(chunk), bits); + let one_over_scale = 1.0 / scale; + let code_min = smallest_code(bits) as f32; + let code_max = largest_code(bits) as f32; + for &value in chunk { + codes.push((value * one_over_scale).round().clamp(code_min, code_max) as i32); + } + scales.push(scale); + } + (scales, Packed::from_i32s(&codes, bits)) +} + +pub(crate) fn quantize_asym_block(block: &[f32], bits: u32, codes: &mut Vec) -> (f32, f32) { + let (lowest, highest) = min_max(block); + let (scale, zero_point) = asymmetric_params(lowest, highest, bits); + let one_over_scale = 1.0 / scale; + let code_min = smallest_code(bits); + let code_max = largest_code(bits); + for &value in block { + let code = (value * one_over_scale + zero_point).round() as i32; + codes.push(code.clamp(code_min, code_max)); + } + (scale, zero_point) +} + +pub(crate) fn dequant_sym_into(scales: &[f32], packed: &Packed, block: usize, out: &mut [f32]) { + let mut codes = vec![0i32; packed.len()]; + packed.unpack_into(&mut codes); + for (index, &code) in codes.iter().enumerate() { + out[index] = code as f32 * scales[index / block]; + } +} + +pub(crate) fn dequant_asym_into( + scales: &[f32], + zero_points: &[f32], + packed: &Packed, + block: usize, + out: &mut [f32], +) { + let mut codes = vec![0i32; packed.len()]; + packed.unpack_into(&mut codes); + for (index, &code) in codes.iter().enumerate() { + let scale = scales[index / block]; + let zero_point = zero_points[index / block]; + out[index] = (code as f32 - zero_point) * scale; + } +} + +pub(crate) fn dot_sym(scales: &[f32], packed: &Packed, block: usize, rhs: &[f32]) -> f32 { + let mut codes = vec![0i32; packed.len()]; + packed.unpack_into(&mut codes); + let mut total = 0.0_f32; + for (index, &code) in codes.iter().enumerate() { + total += code as f32 * scales[index / block] * rhs[index]; + } + total +} + +pub(crate) fn dot_asym( + scales: &[f32], + zero_points: &[f32], + packed: &Packed, + block: usize, + rhs: &[f32], +) -> f32 { + let mut codes = vec![0i32; packed.len()]; + packed.unpack_into(&mut codes); + let mut total = 0.0_f32; + for (index, &code) in codes.iter().enumerate() { + let scale = scales[index / block]; + let zero_point = zero_points[index / block]; + total += (code as f32 - zero_point) * scale * rhs[index]; + } + total +} diff --git a/src/kernels/i4.rs b/src/kernels/i4.rs new file mode 100644 index 0000000..d3d6a4b --- /dev/null +++ b/src/kernels/i4.rs @@ -0,0 +1,147 @@ +//! Symmetric 4-bit: two codes per byte, low nibble first. + +use crate::packed::{nbytes, Packed}; +use crate::params::symmetric_scale; + +use super::reduce::abs_max; + +pub(crate) fn pack_sym_i4(values: &[f32], block: usize) -> (Vec, Packed) { + let mut scales = Vec::with_capacity(values.len().div_ceil(block)); + let mut bytes = vec![0u8; nbytes(values.len(), 4)]; + let mut i = 0usize; + for chunk in values.chunks(block) { + let scale = symmetric_scale(abs_max(chunk), 4); + scales.push(scale); + quant_chunk(chunk, scale, &mut bytes, &mut i); + } + (scales, Packed::from_raw(bytes, 4, values.len())) +} + +fn quant_chunk(values: &[f32], scale: f32, bytes: &mut [u8], value_index: &mut usize) { + let one_over_scale = 1.0 / scale; + let mut i = 0; + #[cfg(target_arch = "aarch64")] + if value_index.is_multiple_of(2) { + // SAFETY: 16 floats → 8 packed bytes. + unsafe { + while i + 16 <= values.len() { + quant_16( + values.as_ptr().add(i), + one_over_scale, + bytes.as_mut_ptr().add(*value_index / 2), + ); + i += 16; + *value_index += 16; + } + } + } + while i < values.len() { + let code = (values[i] * one_over_scale).round().clamp(-8.0, 7.0) as i32; + let byte = *value_index / 2; + if value_index.is_multiple_of(2) { + bytes[byte] = (code as u8) & 0x0F; + } else { + bytes[byte] |= ((code as u8) & 0x0F) << 4; + } + *value_index += 1; + i += 1; + } +} + +#[cfg(target_arch = "aarch64")] +unsafe fn quant_16(src: *const f32, inv: f32, dst: *mut u8) { + use core::arch::aarch64::*; + let vinv = vdupq_n_f32(inv); + let vmin = vdupq_n_f32(-8.0); + let vmax = vdupq_n_f32(7.0); + let q = |v| { + vcvtq_s32_f32(vmaxq_f32( + vminq_f32(vrndaq_f32(vmulq_f32(v, vinv)), vmax), + vmin, + )) + }; + let p0 = vcombine_s16( + vmovn_s32(q(vld1q_f32(src))), + vmovn_s32(q(vld1q_f32(src.add(4)))), + ); + let p1 = vcombine_s16( + vmovn_s32(q(vld1q_f32(src.add(8)))), + vmovn_s32(q(vld1q_f32(src.add(12)))), + ); + let codes = vcombine_s8(vmovn_s16(p0), vmovn_s16(p1)); + let masked = vandq_u8(vreinterpretq_u8_s8(codes), vdupq_n_u8(0x0F)); + vst1_u8( + dst, + vget_low_u8(vorrq_u8( + vuzp1q_u8(masked, masked), + vshlq_n_u8(vuzp2q_u8(masked, masked), 4), + )), + ); +} + +pub(crate) fn dequant_i4_blocks(scales: &[f32], bytes: &[u8], block: usize, out: &mut [f32]) { + let mut i = 0usize; + for (bi, chunk) in out.chunks_mut(block).enumerate() { + let s = scales[bi]; + let mut j = 0; + #[cfg(target_arch = "aarch64")] + if i.is_multiple_of(2) { + // SAFETY: 32 codes = 16 packed bytes. + unsafe { + while j + 32 <= chunk.len() { + dequant_32(bytes.as_ptr().add(i / 2), s, chunk.as_mut_ptr().add(j)); + i += 32; + j += 32; + } + } + } + while j < chunk.len() { + let byte = bytes[i / 2]; + let nib = if i.is_multiple_of(2) { + byte & 0x0F + } else { + byte >> 4 + }; + chunk[j] = ((((nib as i8) << 4) >> 4) as i32 as f32) * s; + i += 1; + j += 1; + } + } +} + +#[cfg(target_arch = "aarch64")] +unsafe fn dequant_32(src: *const u8, scale: f32, dst: *mut f32) { + use core::arch::aarch64::*; + let raw = vld1q_u8(src); + let lo = vshrq_n_s8( + vshlq_n_s8(vreinterpretq_s8_u8(vandq_u8(raw, vdupq_n_u8(0x0F))), 4), + 4, + ); + let hi = vshrq_n_s8(vshlq_n_s8(vreinterpretq_s8_u8(vshrq_n_u8(raw, 4)), 4), 4); + store_i8x16(vzip1q_s8(lo, hi), scale, dst); + store_i8x16(vzip2q_s8(lo, hi), scale, dst.add(16)); +} + +#[cfg(target_arch = "aarch64")] +unsafe fn store_i8x16(q: core::arch::aarch64::int8x16_t, scale: f32, dst: *mut f32) { + use core::arch::aarch64::*; + let lo = vmovl_s8(vget_low_s8(q)); + let hi = vmovl_s8(vget_high_s8(q)); + let vs = vdupq_n_f32(scale); + vst1q_f32( + dst, + vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo))), vs), + ); + vst1q_f32( + dst.add(4), + vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo))), vs), + ); + vst1q_f32( + dst.add(8), + vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi))), vs), + ); + vst1q_f32( + dst.add(12), + vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi))), vs), + ); +} diff --git a/src/kernels/i8.rs b/src/kernels/i8.rs new file mode 100644 index 0000000..de300e4 --- /dev/null +++ b/src/kernels/i8.rs @@ -0,0 +1,130 @@ +//! Symmetric 8-bit: one i8 per value. NEON convert on aarch64. + +use crate::packed::Packed; +use crate::params::symmetric_scale; + +use super::reduce::abs_max; + +pub(crate) fn pack_sym_i8(values: &[f32], block: usize) -> (Vec, Packed) { + let mut scales = Vec::with_capacity(values.len().div_ceil(block)); + let mut bytes = vec![0u8; values.len()]; + let mut off = 0; + for chunk in values.chunks(block) { + let scale = symmetric_scale(abs_max(chunk), 8); + scales.push(scale); + quant_chunk(chunk, scale, &mut bytes[off..off + chunk.len()]); + off += chunk.len(); + } + (scales, Packed::from_raw(bytes, 8, values.len())) +} + +fn quant_chunk(values: &[f32], scale: f32, out: &mut [u8]) { + let one_over_scale = 1.0 / scale; + let mut i = 0; + #[cfg(target_arch = "aarch64")] + { + // SAFETY: `i + 16 <= len`. + unsafe { + while i + 16 <= values.len() { + quant_16( + values.as_ptr().add(i), + one_over_scale, + out.as_mut_ptr().add(i), + ); + i += 16; + } + } + } + while i < values.len() { + out[i] = (values[i] * one_over_scale).round().clamp(-128.0, 127.0) as i32 as i8 as u8; + i += 1; + } +} + +#[cfg(target_arch = "aarch64")] +unsafe fn quant_16(src: *const f32, inv: f32, dst: *mut u8) { + use core::arch::aarch64::*; + let vinv = vdupq_n_f32(inv); + let vmin = vdupq_n_f32(-128.0); + let vmax = vdupq_n_f32(127.0); + let q = |v| { + vcvtq_s32_f32(vmaxq_f32( + vminq_f32(vrndaq_f32(vmulq_f32(v, vinv)), vmax), + vmin, + )) + }; + let p0 = vcombine_s16( + vmovn_s32(q(vld1q_f32(src))), + vmovn_s32(q(vld1q_f32(src.add(4)))), + ); + let p1 = vcombine_s16( + vmovn_s32(q(vld1q_f32(src.add(8)))), + vmovn_s32(q(vld1q_f32(src.add(12)))), + ); + vst1q_s8(dst.cast(), vcombine_s8(vmovn_s16(p0), vmovn_s16(p1))); +} + +pub(crate) fn dequant_i8_blocks(scales: &[f32], bytes: &[u8], block: usize, out: &mut [f32]) { + let mut off = 0; + for (bi, chunk) in out.chunks_mut(block).enumerate() { + dequant_chunk(&bytes[off..off + chunk.len()], scales[bi], chunk); + off += chunk.len(); + } +} + +fn dequant_chunk(bytes: &[u8], scale: f32, out: &mut [f32]) { + let mut i = 0; + #[cfg(target_arch = "aarch64")] + { + // SAFETY: `i + 16 <= len`. + unsafe { + while i + 16 <= out.len() { + dequant_16(bytes.as_ptr().add(i), scale, out.as_mut_ptr().add(i)); + i += 16; + } + } + } + while i < out.len() { + out[i] = (bytes[i] as i8 as f32) * scale; + i += 1; + } +} + +#[cfg(target_arch = "aarch64")] +unsafe fn dequant_16(src: *const u8, scale: f32, dst: *mut f32) { + use core::arch::aarch64::*; + let q = vld1q_s8(src.cast()); + let lo = vmovl_s8(vget_low_s8(q)); + let hi = vmovl_s8(vget_high_s8(q)); + let vs = vdupq_n_f32(scale); + vst1q_f32( + dst, + vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo))), vs), + ); + vst1q_f32( + dst.add(4), + vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo))), vs), + ); + vst1q_f32( + dst.add(8), + vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi))), vs), + ); + vst1q_f32( + dst.add(12), + vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi))), vs), + ); +} + +pub(crate) fn dot_i8_blocks(scales: &[f32], bytes: &[u8], block: usize, rhs: &[f32]) -> f32 { + let mut acc = 0.0_f32; + let mut off = 0; + for (bi, chunk) in rhs.chunks(block).enumerate() { + let mut inner = 0.0_f32; + for (j, &x) in chunk.iter().enumerate() { + inner += (bytes[off + j] as i8 as f32) * x; + } + acc += scales[bi] * inner; + off += chunk.len(); + } + acc +} diff --git a/src/kernels/mod.rs b/src/kernels/mod.rs new file mode 100644 index 0000000..a66c472 --- /dev/null +++ b/src/kernels/mod.rs @@ -0,0 +1,15 @@ +//! Hot loops, split the same way the schemes are: reduce, then 4-bit, 8-bit, +//! then the slow general path. + +mod block; +mod i4; +mod i8; +mod reduce; + +pub(crate) use block::{ + dequant_asym_into, dequant_sym_into, dot_asym, dot_sym, quantize_asym_block, + quantize_sym_packed, +}; +pub(crate) use i4::dequant_i4_blocks; +pub(crate) use i8::{dequant_i8_blocks, dot_i8_blocks}; +pub(crate) use reduce::min_max; diff --git a/src/kernels/reduce.rs b/src/kernels/reduce.rs new file mode 100644 index 0000000..9c279d3 --- /dev/null +++ b/src/kernels/reduce.rs @@ -0,0 +1,106 @@ +//! Per-block range: max-abs (symmetric) and min/max (asymmetric). + +#[inline] +pub(crate) fn abs_max(xs: &[f32]) -> f32 { + #[cfg(target_arch = "aarch64")] + { + abs_max_neon(xs) + } + #[cfg(not(target_arch = "aarch64"))] + { + let mut m = 0.0_f32; + for &x in xs { + m = m.max(x.abs()); + } + m + } +} + +#[inline] +pub(crate) fn min_max(xs: &[f32]) -> (f32, f32) { + #[cfg(target_arch = "aarch64")] + { + min_max_neon(xs) + } + #[cfg(not(target_arch = "aarch64"))] + { + let mut lo = f32::INFINITY; + let mut hi = f32::NEG_INFINITY; + for &x in xs { + lo = lo.min(x); + hi = hi.max(x); + } + (lo, hi) + } +} + +#[cfg(target_arch = "aarch64")] +fn abs_max_neon(xs: &[f32]) -> f32 { + use core::arch::aarch64::*; + let n = xs.len(); + let mut i = 0; + // SAFETY: loads stay inside `xs`. + unsafe { + let mut acc = vdupq_n_f32(0.0); + while i + 16 <= n { + let p = xs.as_ptr().add(i); + acc = vmaxq_f32(acc, vabsq_f32(vld1q_f32(p))); + acc = vmaxq_f32(acc, vabsq_f32(vld1q_f32(p.add(4)))); + acc = vmaxq_f32(acc, vabsq_f32(vld1q_f32(p.add(8)))); + acc = vmaxq_f32(acc, vabsq_f32(vld1q_f32(p.add(12)))); + i += 16; + } + let mut m = vmaxvq_f32(acc); + while i < n { + m = m.max(xs[i].abs()); + i += 1; + } + m + } +} + +#[cfg(target_arch = "aarch64")] +fn min_max_neon(xs: &[f32]) -> (f32, f32) { + use core::arch::aarch64::*; + if xs.is_empty() { + return (f32::INFINITY, f32::NEG_INFINITY); + } + let n = xs.len(); + let mut i = 0; + // SAFETY: loads stay inside `xs`. + unsafe { + let mut vlo = vdupq_n_f32(f32::INFINITY); + let mut vhi = vdupq_n_f32(f32::NEG_INFINITY); + while i + 8 <= n { + let p = xs.as_ptr().add(i); + let a = vld1q_f32(p); + let b = vld1q_f32(p.add(4)); + vlo = vminq_f32(vlo, vminq_f32(a, b)); + vhi = vmaxq_f32(vhi, vmaxq_f32(a, b)); + i += 8; + } + let mut lo = vminvq_f32(vlo); + let mut hi = vmaxvq_f32(vhi); + while i < n { + lo = lo.min(xs[i]); + hi = hi.max(xs[i]); + i += 1; + } + (lo, hi) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn abs_max_matches_iterator() { + assert_eq!(abs_max(&[0.1, -3.5, 2.0, -0.25]), 3.5); + } + + #[test] + fn min_max_matches_iterator() { + assert_eq!(min_max(&[0.1, -3.5, 2.0]), (-3.5, 2.0)); + } +} diff --git a/src/lib.rs b/src/lib.rs index f761808..6a7f858 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,34 +1,40 @@ //! # quantize //! -//! A tiny, readable quantization library — block-wise symmetric or asymmetric, any bit width. +//! A tiny, readable quantization library — block-wise symmetric or asymmetric, +//! any bit width. //! //! ## Example //! //! ``` -//! use quantize::{quantize, dequantize}; +//! use quantize::quantize; //! //! let weights = [0.42_f32, -0.10, 0.70, -0.50]; //! //! // 8-bit, block-size-32, f32 scales -//! let (scales, codes) = quantize::(&weights); -//! let back = dequantize::<_, 32>(&scales, &codes); +//! let q = quantize::(&weights).unwrap(); +//! let back = q.dequantize(); //! //! assert!((back[0] - weights[0]).abs() < 0.01); //! ``` //! -//! `BITS` and `BLOCK` are const generics, so `quantize::<4, 32>(...)`, -//! `quantize::<8, 64>(...)`, etc. all compile to specialized code with -//! zero runtime overhead. -//! -//! See the `symmetric` and `asymmetric` modules for tensor-granularity variants -//! and precision-aware per-block bit widths. +//! `BITS` and `BLOCK` are const generics, so `quantize::(...)`, +//! `quantize::(...)`, etc. all compile to specialized code. //! +//! See `symmetric`, `asymmetric`, and `adaptive` for the other schemes. //! To learn how the library got here, please see `chapters/`. -pub mod asymmetric; -pub mod scale; -pub mod symmetric; +mod kernels; +mod methods; +mod shared; + +pub use methods::{adaptive, asymmetric, learned, symmetric}; + +pub use shared::error::{Error, Result}; +pub use shared::packed::Packed; +pub use shared::scale::Scale; +pub use shared::scheme::Scheme; +pub use shared::tensor::Quantized; +pub use shared::{error, packed, params, scale, scheme, tensor}; +pub use symmetric::{quantize, quantize_tensor}; -pub use asymmetric::{dequantize as dequantize_asym, quantize as quantize_asym}; -pub use scale::Scale; -pub use symmetric::{dequantize, quantize}; +pub(crate) use shared::decode; diff --git a/src/methods/adaptive/mod.rs b/src/methods/adaptive/mod.rs new file mode 100644 index 0000000..4e8f63c --- /dev/null +++ b/src/methods/adaptive/mod.rs @@ -0,0 +1,108 @@ +//! Mixed-precision: pick bits per block from a tolerance. + +use crate::error::{check_block, Error, Result}; +use crate::kernels::{min_max, quantize_asym_block}; +use crate::packed::Packed; +use crate::params::choose_bits; +use crate::scale::Scale; +use crate::tensor::Quantized; + +/// Quantize with a per-block bit width chosen so the half-step `<= tolerance`. +/// +/// Each block is packed at its own width and concatenated. +/// +/// # Errors +/// +/// [`Error::InvalidTolerance`] if `tolerance` is not finite and `> 0`. +/// [`Error::InvalidBlock`] if `BLOCK == 0`. +pub fn quantize( + values: &[f32], + tolerance: f32, +) -> Result> { + quantize_with::(values, BLOCK, tolerance) +} + +/// Runtime-block variant of [`quantize`]. +pub fn quantize_with( + values: &[f32], + block: usize, + tolerance: f32, +) -> Result> { + check_block(block)?; + if !(tolerance.is_finite() && tolerance > 0.0) { + return Err(Error::InvalidTolerance); + } + if values.is_empty() { + return Ok(Quantized::Adaptive { + scales: Vec::new(), + zero_points: Vec::new(), + bytes: Vec::new(), + bits: Vec::new(), + block, + len: 0, + }); + } + + let n_blocks = values.len().div_ceil(block); + let mut scales = Vec::with_capacity(n_blocks); + let mut zero_points = Vec::with_capacity(n_blocks); + let mut bits = Vec::with_capacity(n_blocks); + let mut bytes = Vec::new(); + let mut codes = Vec::new(); + + for chunk in values.chunks(block) { + let (lowest, highest) = min_max(chunk); + let bit_width = choose_bits(highest - lowest, tolerance); + codes.clear(); + let (scale, zero_point) = quantize_asym_block(chunk, bit_width, &mut codes); + scales.push(S::from_f32(scale)); + zero_points.push(S::from_f32(zero_point)); + bits.push(bit_width); + bytes.extend_from_slice(Packed::from_i32s(&codes, bit_width).as_bytes()); + } + + Ok(Quantized::Adaptive { + scales, + zero_points, + bytes, + bits, + block, + len: values.len(), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tiny_range_uses_fewer_bits_than_wide_range() { + let tiny = [0.500_f32, 0.501, 0.499, 0.5005]; + let wide = [0.10_f32, 0.30, 0.70, 1.10]; + let qt = quantize::(&tiny, 0.001).unwrap(); + let qw = quantize::(&wide, 0.001).unwrap(); + let bt = qt.block_bits().unwrap(); + let bw = qw.block_bits().unwrap(); + assert!(bt[0] < bw[0], "tiny={} wide={}", bt[0], bw[0]); + } + + #[test] + fn quiet_block_packs_tighter_than_eight_bit() { + let tiny = [0.500_f32; 32]; + let q = quantize::(&tiny, 0.001).unwrap(); + assert!( + q.codes().len() < 32, + "expected packed codes < 32 bytes, got {}", + q.codes().len() + ); + assert!(matches!(q, Quantized::Adaptive { .. })); + } + + #[test] + fn rejects_non_positive_tolerance() { + assert_eq!( + quantize::(&[1.0], 0.0).unwrap_err(), + Error::InvalidTolerance + ); + } +} diff --git a/src/methods/asymmetric/mod.rs b/src/methods/asymmetric/mod.rs new file mode 100644 index 0000000..83aa45d --- /dev/null +++ b/src/methods/asymmetric/mod.rs @@ -0,0 +1,66 @@ +//! Asymmetric quantization: scale and zero-point per group. + +use crate::error::{check_bits, check_block, Result}; +use crate::kernels::quantize_asym_block; +use crate::packed::Packed; +use crate::scale::Scale; +use crate::tensor::Quantized; + +/// Quantize into blocks of `BLOCK` using one bit width for every block. +pub fn quantize( + values: &[f32], +) -> Result> { + quantize_with::(values, BITS, BLOCK) +} + +/// Runtime-width variant of [`quantize`]. +pub fn quantize_with(values: &[f32], bits: u32, block: usize) -> Result> { + check_bits(bits)?; + check_block(block)?; + if values.is_empty() { + return Ok(Quantized::Asymmetric { + scales: Vec::new(), + zero_points: Vec::new(), + codes: Packed::from_i32s(&[], bits), + block, + len: 0, + }); + } + let mut scales = Vec::with_capacity(values.len().div_ceil(block)); + let mut zero_points = Vec::with_capacity(values.len().div_ceil(block)); + let mut codes = Vec::with_capacity(values.len()); + for chunk in values.chunks(block) { + let (s, z) = quantize_asym_block(chunk, bits, &mut codes); + scales.push(S::from_f32(s)); + zero_points.push(S::from_f32(z)); + } + Ok(Quantized::Asymmetric { + scales, + zero_points, + codes: Packed::from_i32s(&codes, bits), + block, + len: values.len(), + }) +} + +/// Quantize the entire tensor with one scale and one zero-point. +pub fn quantize_tensor(values: &[f32]) -> Result> { + quantize_with::(values, BITS, values.len().max(1)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn positive_block_uses_full_grid() { + let w = [0.10_f32, 0.30, 0.70, 1.10]; + let q = quantize::(&w).unwrap(); + assert!(matches!(q, Quantized::Asymmetric { .. })); + assert!(!q.zero_points().is_empty()); + let back = q.dequantize(); + for (a, b) in w.iter().zip(&back) { + assert!((a - b).abs() < 0.02, "{a} vs {b}"); + } + } +} diff --git a/src/methods/learned/mod.rs b/src/methods/learned/mod.rs new file mode 100644 index 0000000..bb665fe --- /dev/null +++ b/src/methods/learned/mod.rs @@ -0,0 +1,107 @@ +//! After codes are chosen, pick a better scale and zero-point. +//! +//! We want: `original ≈ scale * (code - zero_point)`. +//! Same as: `original ≈ scale * code + offset`, with +//! `offset = -scale * zero_point`. Codes stay fixed. + +use crate::decode::unpack_codes; +use crate::scale::Scale; +use crate::tensor::Quantized; + +/// Best-fit `scale` and `zero_point` for `values ≈ scale * (codes - zero_point)`. +/// +/// Returns `(1.0, 0.0)` when every code is the same (nothing to fit). +pub fn fit_scale_and_zero_point(values: &[f32], codes: &[i32]) -> (f32, f32) { + debug_assert_eq!(values.len(), codes.len()); + let count = values.len() as f32; + if count == 0.0 { + return (1.0, 0.0); + } + + let mut sum_codes = 0.0_f32; + let mut sum_values = 0.0_f32; + let mut sum_code_squared = 0.0_f32; + let mut sum_code_times_value = 0.0_f32; + for (&value, &code) in values.iter().zip(codes) { + let code = code as f32; + sum_codes += code; + sum_values += value; + sum_code_squared += code * code; + sum_code_times_value += code * value; + } + + let mean_code = sum_codes / count; + let mean_value = sum_values / count; + let code_spread = sum_code_squared - sum_codes * mean_code; + if code_spread.abs() < 1e-12 { + return (1.0, 0.0); + } + + // Line of best fit: value ≈ scale * code + offset. + let scale = (sum_code_times_value - sum_codes * mean_value) / code_spread; + let offset = mean_value - scale * mean_code; + if scale.abs() < 1e-12 { + return (1.0, 0.0); + } + let zero_point = -offset / scale; + (scale, zero_point) +} + +/// Recompute each block's scale and zero-point. Codes do not change. +/// +/// Empty input is left as-is. [`Quantized::Symmetric`] becomes +/// [`Quantized::Asymmetric`]. +pub fn refine(quantized: &mut Quantized, values: &[f32]) { + if quantized.is_empty() || values.len() != quantized.len() { + return; + } + let block = quantized.block(); + let mut codes = vec![0i32; quantized.len()]; + unpack_codes(quantized, &mut codes); + + let mut scales = Vec::new(); + let mut zero_points = Vec::new(); + for (block_index, block_values) in values.chunks(block).enumerate() { + let start = block_index * block; + let block_codes = &codes[start..start + block_values.len()]; + let (scale, zero_point) = fit_scale_and_zero_point(block_values, block_codes); + scales.push(S::from_f32(scale)); + zero_points.push(S::from_f32(zero_point)); + } + + *quantized = match quantized { + Quantized::Adaptive { + bytes, bits, len, .. + } => Quantized::Adaptive { + scales, + zero_points, + bytes: bytes.clone(), + bits: bits.clone(), + block, + len: *len, + }, + Quantized::Symmetric { codes, len, .. } | Quantized::Asymmetric { codes, len, .. } => { + Quantized::Asymmetric { + scales, + zero_points, + codes: codes.clone(), + block, + len: *len, + } + } + }; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn fit_recovers_known_line() { + let codes = [0, 1, 2, 3, 4]; + let values: Vec = codes.iter().map(|&code| 0.5 * code as f32 + 1.0).collect(); + let (scale, zero_point) = fit_scale_and_zero_point(&values, &codes); + assert!((scale - 0.5).abs() < 1e-5); + assert!((zero_point + 2.0).abs() < 1e-5); + } +} diff --git a/src/methods/mod.rs b/src/methods/mod.rs new file mode 100644 index 0000000..a8ac746 --- /dev/null +++ b/src/methods/mod.rs @@ -0,0 +1,6 @@ +//! The four ways to pick a ruler: symmetric, asymmetric, adaptive, learned. + +pub mod adaptive; +pub mod asymmetric; +pub mod learned; +pub mod symmetric; diff --git a/src/methods/symmetric/mod.rs b/src/methods/symmetric/mod.rs new file mode 100644 index 0000000..e0aefc2 --- /dev/null +++ b/src/methods/symmetric/mod.rs @@ -0,0 +1,111 @@ +//! Symmetric quantization: one scale per group, codes centered on zero. + +use crate::error::{check_bits, check_block, Result}; +use crate::kernels::quantize_sym_packed; +use crate::packed::Packed; +use crate::scale::Scale; +use crate::tensor::Quantized; + +/// Quantize `values` into fixed-size blocks of `BLOCK` with `BITS`-wide codes. +/// +/// # Errors +/// +/// [`crate::Error::InvalidBits`] or [`crate::Error::InvalidBlock`]. +pub fn quantize( + values: &[f32], +) -> Result> { + quantize_with::(values, BITS, BLOCK) +} + +/// Runtime-width variant of [`quantize`]. +pub fn quantize_with(values: &[f32], bits: u32, block: usize) -> Result> { + check_bits(bits)?; + check_block(block)?; + if values.is_empty() { + return Ok(Quantized::Symmetric { + scales: Vec::new(), + codes: Packed::from_raw(Vec::new(), bits, 0), + block, + len: 0, + }); + } + let (scales_f, codes) = quantize_sym_packed(values, bits, block); + Ok(Quantized::Symmetric { + scales: scales_f.into_iter().map(S::from_f32).collect(), + codes, + block, + len: values.len(), + }) +} + +/// Quantize the entire tensor with a single scale. +pub fn quantize_tensor(values: &[f32]) -> Result> { + quantize_with::(values, BITS, values.len().max(1)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn eight_bit_roundtrip_stays_within_half_step() { + let w = [0.42_f32, -0.10, 0.70, -0.50]; + let q = quantize::(&w).unwrap(); + let back = q.dequantize(); + for (a, b) in w.iter().zip(&back) { + assert!((a - b).abs() < 0.01, "{a} vs {b}"); + } + } + + #[test] + fn packed_four_bit_uses_half_byte_per_code() { + let w = [0.1_f32; 32]; + let q = quantize::(&w).unwrap(); + assert_eq!(q.codes().len(), 16); + } + + #[test] + fn remainder_block_roundtrips() { + let w: Vec = (0..40).map(|i| (i as f32) * 0.01 - 0.2).collect(); + let q = quantize::(&w).unwrap(); + assert_eq!(q.len(), 40); + let back = q.dequantize(); + for (a, b) in w.iter().zip(&back) { + assert!((a - b).abs() < 0.01, "{a} vs {b}"); + } + } + + #[test] + fn dequantize_into_rejects_wrong_length() { + let w = [0.1_f32; 8]; + let q = quantize::(&w).unwrap(); + let mut out = [0.0f32; 3]; + assert!(matches!( + q.dequantize_into(&mut out), + Err(crate::Error::LengthMismatch { + expected: 8, + got: 3 + }) + )); + } + + #[test] + fn four_bit_remainder_roundtrips() { + let w: Vec = (0..40).map(|i| (i as f32) * 0.02 - 0.4).collect(); + let q = quantize::(&w).unwrap(); + let back = q.dequantize(); + for (a, b) in w.iter().zip(&back) { + assert!((a - b).abs() < 0.08, "{a} vs {b}"); + } + } + + #[test] + fn fused_dot_matches_dequant_then_dot() { + let w: Vec = (0..64).map(|i| (i as f32) * 0.01 - 0.3).collect(); + let q = quantize::(&w).unwrap(); + let recon = q.dequantize(); + let naive: f32 = recon.iter().zip(&w).map(|(a, b)| a * b).sum(); + let fused = q.dot(&w).unwrap(); + assert!((naive - fused).abs() < 1e-4, "{naive} vs {fused}"); + } +} diff --git a/src/shared/decode.rs b/src/shared/decode.rs new file mode 100644 index 0000000..f3702bf --- /dev/null +++ b/src/shared/decode.rs @@ -0,0 +1,123 @@ +//! Reconstruct f32 from packed codes. + +use crate::kernels::{ + dequant_asym_into, dequant_i4_blocks, dequant_i8_blocks, dequant_sym_into, dot_asym, + dot_i8_blocks, dot_sym, +}; +use crate::packed::{nbytes, Packed}; +use crate::scale::Scale; +use crate::tensor::Quantized; + +pub(crate) fn as_f32(values: &[S]) -> Vec { + values.iter().copied().map(S::to_f32).collect() +} + +pub(crate) fn dequant_sym(scales: &[S], codes: &Packed, block: usize, out: &mut [f32]) { + let scales = as_f32(scales); + match codes.bits() { + 8 => dequant_i8_blocks(&scales, codes.as_bytes(), block, out), + 4 => dequant_i4_blocks(&scales, codes.as_bytes(), block, out), + _ => dequant_sym_into(&scales, codes, block, out), + } +} + +pub(crate) fn dequant_asym( + scales: &[S], + zero_points: &[S], + codes: &Packed, + block: usize, + out: &mut [f32], +) { + dequant_asym_into(&as_f32(scales), &as_f32(zero_points), codes, block, out); +} + +pub(crate) fn dequant_adaptive( + scales: &[S], + zero_points: &[S], + bytes: &[u8], + bits: &[u32], + block: usize, + len: usize, + out: &mut [f32], +) { + let mut byte_offset = 0; + let mut value_index = 0; + for (block_index, &bit_width) in bits.iter().enumerate() { + let count = (len - value_index).min(block); + let byte_count = nbytes(count, bit_width); + let mut codes = vec![0i32; count]; + Packed::unpack_slice( + &bytes[byte_offset..byte_offset + byte_count], + bit_width, + &mut codes, + count, + ); + let scale = scales[block_index].to_f32(); + let zero_point = zero_points[block_index].to_f32(); + for (slot, &code) in out[value_index..value_index + count].iter_mut().zip(&codes) { + *slot = (code as f32 - zero_point) * scale; + } + byte_offset += byte_count; + value_index += count; + } +} + +pub(crate) fn dot_of(quantized: &Quantized, rhs: &[f32]) -> f32 { + match quantized { + Quantized::Symmetric { + scales, + codes, + block, + .. + } if codes.bits() == 8 => dot_i8_blocks(&as_f32(scales), codes.as_bytes(), *block, rhs), + Quantized::Symmetric { + scales, + codes, + block, + .. + } => dot_sym(&as_f32(scales), codes, *block, rhs), + Quantized::Asymmetric { + scales, + zero_points, + codes, + block, + .. + } => dot_asym(&as_f32(scales), &as_f32(zero_points), codes, *block, rhs), + Quantized::Adaptive { .. } => quantized + .dequantize() + .iter() + .zip(rhs) + .map(|(left, right)| left * right) + .sum(), + } +} + +pub(crate) fn unpack_codes(quantized: &Quantized, out: &mut [i32]) { + match quantized { + Quantized::Symmetric { codes, .. } | Quantized::Asymmetric { codes, .. } => { + codes.unpack_into(out); + } + Quantized::Adaptive { + bytes, + bits, + block, + len, + .. + } => { + let mut byte_offset = 0; + let mut value_index = 0; + for &bit_width in bits { + let count = (*len - value_index).min(*block); + let byte_count = nbytes(count, bit_width); + Packed::unpack_slice( + &bytes[byte_offset..byte_offset + byte_count], + bit_width, + &mut out[value_index..value_index + count], + count, + ); + byte_offset += byte_count; + value_index += count; + } + } + } +} diff --git a/src/shared/error.rs b/src/shared/error.rs new file mode 100644 index 0000000..e46fb20 --- /dev/null +++ b/src/shared/error.rs @@ -0,0 +1,99 @@ +//! Recoverable failures from quantization and dequantization. + +use core::fmt; + +/// An error produced by a fallible quantization API. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Error { + /// `bits` is outside the supported `2..=16` range. + InvalidBits { + /// Requested bit width. + bits: u32, + }, + /// Block length must be at least 1. + InvalidBlock { + /// Requested block length. + block: usize, + }, + /// Reconstruction tolerance must be finite and strictly positive. + InvalidTolerance, + /// Output or partner buffer length does not match the quantized tensor. + LengthMismatch { + /// Length required by the quantized tensor. + expected: usize, + /// Length the caller actually passed. + got: usize, + }, +} + +impl fmt::Display for Error { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::InvalidBits { bits } => { + write!(f, "bit width {bits} is outside the supported range 2..=16") + } + Self::InvalidBlock { block } => { + write!(f, "block size {block} must be at least 1") + } + Self::InvalidTolerance => { + write!(f, "tolerance must be a finite number greater than 0") + } + Self::LengthMismatch { expected, got } => { + write!(f, "length mismatch: expected {expected}, got {got}") + } + } + } +} + +#[cfg(feature = "std")] +impl std::error::Error for Error {} + +/// Result alias for this crate. +pub type Result = core::result::Result; + +pub(crate) fn check_bits(bits: u32) -> Result<()> { + if (2..=16).contains(&bits) { + Ok(()) + } else { + Err(Error::InvalidBits { bits }) + } +} + +pub(crate) fn check_block(block: usize) -> Result<()> { + if block == 0 { + Err(Error::InvalidBlock { block }) + } else { + Ok(()) + } +} + +pub(crate) fn check_len(expected: usize, got: usize) -> Result<()> { + if expected == got { + Ok(()) + } else { + Err(Error::LengthMismatch { expected, got }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn check_bits_rejects_one_and_seventeen() { + assert!(matches!(check_bits(1), Err(Error::InvalidBits { bits: 1 }))); + assert!(matches!( + check_bits(17), + Err(Error::InvalidBits { bits: 17 }) + )); + } + + #[test] + fn display_mentions_expected_length() { + let err = Error::LengthMismatch { + expected: 4, + got: 1, + }; + assert_eq!(err.to_string(), "length mismatch: expected 4, got 1"); + } +} diff --git a/src/shared/mod.rs b/src/shared/mod.rs new file mode 100644 index 0000000..ba3d189 --- /dev/null +++ b/src/shared/mod.rs @@ -0,0 +1,10 @@ +//! Pieces every scheme uses: errors, scales, packing, the output type. + +pub mod error; +pub mod packed; +pub mod params; +pub mod scale; +pub mod scheme; +pub mod tensor; + +pub(crate) mod decode; diff --git a/src/shared/packed.rs b/src/shared/packed.rs new file mode 100644 index 0000000..88b05d8 --- /dev/null +++ b/src/shared/packed.rs @@ -0,0 +1,188 @@ +//! Bit-packed signed integer codes. +//! +//! Codes are stored as `bits`-wide two's-complement fields, packed LSB-first +//! into a `Vec`. 4-bit and 8-bit paths are specialized; other widths use +//! a general bit-buffer. + +/// Packed signed codes plus the bit width they were written with. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Packed { + bytes: Vec, + bits: u32, + len: usize, +} + +impl Packed { + /// Pack `codes` using `bits` bits each. + pub fn from_i32s(codes: &[i32], bits: u32) -> Self { + let mut p = Self { + bytes: vec![0u8; nbytes(codes.len(), bits)], + bits, + len: codes.len(), + }; + match bits { + 8 => pack_i8(&mut p.bytes, codes), + 4 => pack_i4(&mut p.bytes, codes), + _ => pack_general(&mut p.bytes, codes, bits), + } + p + } + + /// Number of codes. + #[inline] + pub fn len(&self) -> usize { + self.len + } + + /// Whether there are no codes. + #[inline] + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Bit width of each code. + #[inline] + pub fn bits(&self) -> u32 { + self.bits + } + + /// Raw packed bytes. + #[inline] + pub fn as_bytes(&self) -> &[u8] { + &self.bytes + } + + /// Wrap already-packed bytes. `bytes` must hold `len` codes of `bits`. + pub(crate) fn from_raw(bytes: Vec, bits: u32, len: usize) -> Self { + Self { bytes, bits, len } + } + + /// Unpack into `out`, which must be at least [`len`](Self::len) long. + pub fn unpack_into(&self, out: &mut [i32]) { + debug_assert!(out.len() >= self.len); + match self.bits { + 8 => unpack_i8(&self.bytes, out, self.len), + 4 => unpack_i4(&self.bytes, out, self.len), + _ => unpack_general(&self.bytes, out, self.len, self.bits), + } + } + + /// Unpack `n` codes of width `bits` from a raw byte slice. + pub(crate) fn unpack_slice(bytes: &[u8], bits: u32, out: &mut [i32], n: usize) { + match bits { + 8 => unpack_i8(bytes, out, n), + 4 => unpack_i4(bytes, out, n), + _ => unpack_general(bytes, out, n, bits), + } + } +} + +#[inline] +pub(crate) fn nbytes(len: usize, bits: u32) -> usize { + (len * bits as usize).div_ceil(8) +} + +fn pack_i8(bytes: &mut [u8], codes: &[i32]) { + for (b, &q) in bytes.iter_mut().zip(codes) { + *b = q as i8 as u8; + } +} + +fn unpack_i8(bytes: &[u8], out: &mut [i32], n: usize) { + for i in 0..n { + out[i] = bytes[i] as i8 as i32; + } +} + +fn pack_i4(bytes: &mut [u8], codes: &[i32]) { + for (i, chunk) in codes.chunks(2).enumerate() { + let lo = (chunk[0] as u8) & 0x0F; + let hi = chunk.get(1).copied().unwrap_or(0) as u8 & 0x0F; + bytes[i] = lo | (hi << 4); + } +} + +fn unpack_i4(bytes: &[u8], out: &mut [i32], n: usize) { + for i in 0..n { + let byte = bytes[i / 2]; + let nib = if i.is_multiple_of(2) { + byte & 0x0F + } else { + byte >> 4 + }; + out[i] = (((nib as i8) << 4) >> 4) as i32; + } +} + +fn pack_general(bytes: &mut [u8], codes: &[i32], bits: u32) { + for (i, &q) in codes.iter().enumerate() { + write_code(bytes, i, bits, q); + } +} + +fn unpack_general(bytes: &[u8], out: &mut [i32], n: usize, bits: u32) { + for (i, slot) in out.iter_mut().enumerate().take(n) { + *slot = read_code(bytes, i, bits); + } +} + +fn write_code(bytes: &mut [u8], index: usize, bits: u32, q: i32) { + let mask = (1u32 << bits) - 1; + let val = (q as u32) & mask; + let bit = index * bits as usize; + let byte = bit / 8; + let off = bit % 8; + let wide = (val as u64) << off; + bytes[byte] |= wide as u8; + if off + bits as usize > 8 { + bytes[byte + 1] |= (wide >> 8) as u8; + } + if off + bits as usize > 16 { + bytes[byte + 2] |= (wide >> 16) as u8; + } +} + +fn read_code(bytes: &[u8], index: usize, bits: u32) -> i32 { + let bit = index * bits as usize; + let byte = bit / 8; + let off = bit % 8; + let mut wide = bytes[byte] as u32 >> off; + if off + bits as usize > 8 { + wide |= (bytes[byte + 1] as u32) << (8 - off); + } + if off + bits as usize > 16 { + wide |= (bytes[byte + 2] as u32) << (16 - off); + } + let mask = (1u32 << bits) - 1; + let u = wide & mask; + let sign = 1u32 << (bits - 1); + if u & sign != 0 { + (u | !mask) as i32 + } else { + u as i32 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn four_bit_roundtrip_preserves_signed_codes() { + let codes = [-8, -1, 0, 7, 3, -4]; + let p = Packed::from_i32s(&codes, 4); + let mut out = [0i32; 6]; + p.unpack_into(&mut out); + assert_eq!(out, codes); + assert_eq!(p.as_bytes().len(), 3); + } + + #[test] + fn five_bit_roundtrip_preserves_signed_codes() { + let codes = [-16, -1, 0, 15, 7]; + let p = Packed::from_i32s(&codes, 5); + let mut out = [0i32; 5]; + p.unpack_into(&mut out); + assert_eq!(out, codes); + } +} diff --git a/src/shared/params.rs b/src/shared/params.rs new file mode 100644 index 0000000..581b1b1 --- /dev/null +++ b/src/shared/params.rs @@ -0,0 +1,74 @@ +//! Integer grids, scale selection, and mixed-precision bit choice. + +/// Largest integer code for this bit width. 8-bit → 127, 4-bit → 7. +#[inline] +pub const fn largest_code(bits: u32) -> i32 { + (1_i32 << (bits - 1)) - 1 +} + +/// Smallest integer code for this bit width. 8-bit → -128, 4-bit → -8. +#[inline] +pub const fn smallest_code(bits: u32) -> i32 { + -(1_i32 << (bits - 1)) +} + +/// Tick size so the biggest absolute value lands on [`largest_code`]. +#[inline] +pub fn symmetric_scale(max_abs: f32, bits: u32) -> f32 { + if max_abs > 0.0 { + max_abs / largest_code(bits) as f32 + } else { + 1.0 + } +} + +/// Scale and zero-point that stretch `[lowest, highest]` onto the integer grid. +#[inline] +pub fn asymmetric_params(lowest: f32, highest: f32, bits: u32) -> (f32, f32) { + if lowest >= highest { + return (1.0, 0.0); + } + let code_min = smallest_code(bits) as f32; + let code_max = largest_code(bits) as f32; + let scale = (highest - lowest) / (code_max - code_min); + let zero_point = code_min - lowest / scale; + (scale, zero_point) +} + +/// Smallest bit width in `2..=8` whose half-step is `<= tolerance`. +/// +/// A flat block (range `0`) always returns 2. +pub fn choose_bits(range: f32, tolerance: f32) -> u32 { + if range <= 0.0 { + return 2; + } + for bits in 2..=8 { + let tick_count = ((1u32 << bits) - 1) as f32; + let half_step = range / tick_count / 2.0; + if half_step <= tolerance { + return bits; + } + } + 8 +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn four_bit_codes_run_from_minus_eight_to_seven() { + assert_eq!(largest_code(4), 7); + assert_eq!(smallest_code(4), -8); + } + + #[test] + fn choose_bits_picks_two_for_tiny_range() { + assert_eq!(choose_bits(0.001, 0.001), 2); + } + + #[test] + fn choose_bits_saturates_at_eight() { + assert_eq!(choose_bits(10.0, 0.0001), 8); + } +} diff --git a/src/scale.rs b/src/shared/scale.rs similarity index 54% rename from src/scale.rs rename to src/shared/scale.rs index 70d68f3..f88c71d 100644 --- a/src/scale.rs +++ b/src/shared/scale.rs @@ -1,38 +1,46 @@ -//! Scale trait — defines how per-block scale factors are stored. +//! Scale trait — how per-block scale factors are stored. //! -//! Implement this for any type that can round-trip through f32. -//! The library ships impls for `f32` (lossless) and `f16` (GGML-compatible). +//! Implement this for any type that can round-trip through `f32`. +//! The crate ships impls for `f32` (lossless), `f16`, and `bf16`. use half::{bf16, f16}; -/// A type that can serve as a per-block scale factor. +/// A type that can serve as a per-block scale (or zero-point) factor. pub trait Scale: Copy { + /// Convert from the working-precision `f32` value. fn from_f32(v: f32) -> Self; + /// Convert back to `f32` for arithmetic. fn to_f32(self) -> f32; } impl Scale for f32 { + #[inline] fn from_f32(v: f32) -> Self { v } + #[inline] fn to_f32(self) -> f32 { self } } impl Scale for f16 { + #[inline] fn from_f32(v: f32) -> Self { f16::from_f32(v) } + #[inline] fn to_f32(self) -> f32 { f16::to_f32(self) } } impl Scale for bf16 { + #[inline] fn from_f32(v: f32) -> Self { bf16::from_f32(v) } + #[inline] fn to_f32(self) -> f32 { bf16::to_f32(self) } diff --git a/src/shared/scheme.rs b/src/shared/scheme.rs new file mode 100644 index 0000000..c61e75f --- /dev/null +++ b/src/shared/scheme.rs @@ -0,0 +1,67 @@ +//! Runtime scheme selection — one entry point, scheme-specific I/O inside. + +use crate::error::Result; +use crate::scale::Scale; +use crate::tensor::Quantized; +use crate::{adaptive, asymmetric, symmetric}; + +/// Which algorithm to run. Pick this from config or a CLI flag; the returned +/// [`Quantized`] variant is the scheme that ran. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Scheme { + /// Symmetric block quantization. + Symmetric { + /// Integer width, `2..=16`. + bits: u32, + /// Elements per scale. + block: usize, + }, + /// Asymmetric block quantization. + Asymmetric { + /// Integer width, `2..=16`. + bits: u32, + /// Elements per scale. + block: usize, + }, + /// Per-block bit width from a reconstruction tolerance. + Adaptive { + /// Elements per scale / bit-width decision. + block: usize, + /// Maximum half-step of the integer grid. + tolerance: f32, + }, +} + +impl Scheme { + /// Symmetric 8-bit blocks of 32 — the common default. + pub const Q8_32: Self = Self::Symmetric { bits: 8, block: 32 }; + + /// Symmetric 4-bit blocks of 32 — GGML Q4_0-shaped storage. + pub const Q4_32: Self = Self::Symmetric { bits: 4, block: 32 }; + + /// Run this scheme on `values`. + pub fn quantize(self, values: &[f32]) -> Result> { + match self { + Self::Symmetric { bits, block } => symmetric::quantize_with::(values, bits, block), + Self::Asymmetric { bits, block } => asymmetric::quantize_with::(values, bits, block), + Self::Adaptive { block, tolerance } => { + adaptive::quantize_with::(values, block, tolerance) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn scheme_enum_matches_direct_call() { + let w = [0.42_f32, -0.10, 0.70, -0.50]; + let via = Scheme::Symmetric { bits: 8, block: 4 } + .quantize::(&w) + .unwrap(); + let direct = symmetric::quantize::(&w).unwrap(); + assert_eq!(via.dequantize(), direct.dequantize()); + } +} diff --git a/src/shared/tensor.rs b/src/shared/tensor.rs new file mode 100644 index 0000000..af2cca4 --- /dev/null +++ b/src/shared/tensor.rs @@ -0,0 +1,153 @@ +//! One enum, one variant per scheme. + +use crate::decode::{dequant_adaptive, dequant_asym, dequant_sym, dot_of}; +use crate::error::{check_len, Result}; +use crate::packed::Packed; +use crate::scale::Scale; + +/// Packed codes and the scheme that produced them. +#[derive(Clone, Debug)] +pub enum Quantized { + /// One scale per block. + Symmetric { + scales: Vec, + codes: Packed, + block: usize, + len: usize, + }, + /// Scale and zero-point per block. + Asymmetric { + scales: Vec, + zero_points: Vec, + codes: Packed, + block: usize, + len: usize, + }, + /// Per-block bit width; codes packed at that width. + Adaptive { + scales: Vec, + zero_points: Vec, + bytes: Vec, + bits: Vec, + block: usize, + len: usize, + }, +} + +impl Quantized { + pub fn len(&self) -> usize { + match self { + Self::Symmetric { len, .. } + | Self::Asymmetric { len, .. } + | Self::Adaptive { len, .. } => *len, + } + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub fn block(&self) -> usize { + match self { + Self::Symmetric { block, .. } + | Self::Asymmetric { block, .. } + | Self::Adaptive { block, .. } => *block, + } + } + + pub fn scales(&self) -> &[S] { + match self { + Self::Symmetric { scales, .. } + | Self::Asymmetric { scales, .. } + | Self::Adaptive { scales, .. } => scales, + } + } + + pub fn zero_points(&self) -> &[S] { + match self { + Self::Symmetric { .. } => &[], + Self::Asymmetric { zero_points, .. } | Self::Adaptive { zero_points, .. } => { + zero_points + } + } + } + + pub fn codes(&self) -> &[u8] { + match self { + Self::Symmetric { codes, .. } | Self::Asymmetric { codes, .. } => codes.as_bytes(), + Self::Adaptive { bytes, .. } => bytes, + } + } + + pub fn block_bits(&self) -> Option<&[u32]> { + match self { + Self::Adaptive { bits, .. } => Some(bits), + _ => None, + } + } + + pub fn nbytes(&self) -> usize { + let extra = match self { + Self::Adaptive { bits, .. } => core::mem::size_of_val(bits.as_slice()), + _ => 0, + }; + self.codes().len() + + core::mem::size_of_val(self.scales()) + + core::mem::size_of_val(self.zero_points()) + + extra + } + + pub fn bits_per_element(&self) -> f32 { + if self.is_empty() { + 0.0 + } else { + self.nbytes() as f32 * 8.0 / self.len() as f32 + } + } + + pub fn dequantize(&self) -> Vec { + let mut out = vec![0.0; self.len()]; + let _ = self.dequantize_into(&mut out); + out + } + + pub fn dequantize_into(&self, out: &mut [f32]) -> Result<()> { + check_len(self.len(), out.len())?; + if self.is_empty() { + return Ok(()); + } + match self { + Self::Symmetric { + scales, + codes, + block, + .. + } => dequant_sym(scales, codes, *block, out), + Self::Asymmetric { + scales, + zero_points, + codes, + block, + .. + } => dequant_asym(scales, zero_points, codes, *block, out), + Self::Adaptive { + scales, + zero_points, + bytes, + bits, + block, + len, + } => dequant_adaptive(scales, zero_points, bytes, bits, *block, *len, out), + } + Ok(()) + } + + pub fn dot(&self, rhs: &[f32]) -> Result { + check_len(self.len(), rhs.len())?; + Ok(if self.is_empty() { + 0.0 + } else { + dot_of(self, rhs) + }) + } +} diff --git a/src/symmetric.rs b/src/symmetric.rs deleted file mode 100644 index 5afc20e..0000000 --- a/src/symmetric.rs +++ /dev/null @@ -1,81 +0,0 @@ -//! Symmetric quantization. -//! -//! One scale per group, chosen from the max absolute value. -//! Integer codes are always symmetric around zero. -//! -//! A "group" can be either a fixed-size block or the entire tensor. -//! Use the block form when you want to limit the impact of outliers within the tensor. -//! Use the tensor form for the simplest possible case (one scale total). -//! -//! This is the baseline before introducing zero-points (asymmetric). - -use crate::Scale; - -const fn max_int() -> i32 { - (1_i32 << (BITS - 1)) - 1 -} -const fn min_int() -> i32 { - -(1_i32 << (BITS - 1)) -} - -fn choose_scale(group: &[f32]) -> f32 { - let max_abs = group.iter().map(|v| v.abs()).fold(0.0_f32, f32::max); - if max_abs > 0.0 { - max_abs / max_int::() as f32 - } else { - 1.0 - } -} - -/// Quantize into fixed-size blocks. One scale per block. -/// Returns `(scales, codes)`. -pub fn quantize( - values: &[f32], -) -> (Vec, Vec) { - let mut scales = Vec::with_capacity(values.len() / BLOCK + 1); - let mut codes = Vec::with_capacity(values.len()); - for chunk in values.chunks(BLOCK) { - let s = S::from_f32(choose_scale::(chunk)); - scales.push(s); - let sf = s.to_f32(); - let lo = min_int::() as f32; - let hi = max_int::() as f32; - codes.extend(chunk.iter().map(|&x| (x / sf).round().clamp(lo, hi) as i32)); - } - (scales, codes) -} - -/// Reconstruct from per-block scales and codes. -pub fn dequantize(scales: &[S], codes: &[i32]) -> Vec { - codes - .chunks(BLOCK) - .zip(scales) - .flat_map(|(blk, s)| { - let sf = s.to_f32(); - blk.iter().map(move |&q| q as f32 * sf) - }) - .collect() -} - -/// Quantize the entire tensor with a single scale. -/// Returns `(scale, codes)`. -pub fn quantize_tensor(values: &[f32]) -> (S, Vec) { - if values.is_empty() { - return (S::from_f32(1.0), vec![]); - } - let s = S::from_f32(choose_scale::(values)); - let sf = s.to_f32(); - let lo = min_int::() as f32; - let hi = max_int::() as f32; - let codes = values - .iter() - .map(|&x| (x / sf).round().clamp(lo, hi) as i32) - .collect(); - (s, codes) -} - -/// Reconstruct the entire tensor from a single scale and codes. -pub fn dequantize_tensor(scale: S, codes: &[i32]) -> Vec { - let sf = scale.to_f32(); - codes.iter().map(|&q| q as f32 * sf).collect() -}