1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
//! Affine int8 quantisation.
//!
//! The other dtype a TPU loves is int8: an int8 matmul runs at roughly four times
//! the throughput of bf16. To use it you quantise an f32 tensor to int8 with an
//! affine map `q = round(x / scale) + zero_point`, run the integer matmul, then
//! dequantise. This module provides the per-tensor affine map and the lattice
//! conversions, preserving the hardware tiling throughout.
use crate::error::Result;
use crate::lattice::PaddedTileLattice;
/// Parameters of an affine int8 quantisation: `real = scale * (q - zero_point)`.
#[derive(Clone, Copy, PartialEq, Debug)]
pub struct QuantParams {
/// The size of one quantisation step in real units.
pub scale: f32,
/// The int8 value that maps to real zero.
pub zero_point: i8,
}
impl QuantParams {
/// Derive symmetric parameters (zero_point = 0) from a value's magnitude bound.
///
/// `abs_max` is the largest absolute value the tensor takes. The int8 range is
/// treated as `-127..=127` so that negation stays representable.
pub fn symmetric(abs_max: f32) -> Self {
let scale = if abs_max == 0.0 { 1.0 } else { abs_max / 127.0 };
QuantParams {
scale,
zero_point: 0,
}
}
/// Derive asymmetric parameters from an observed `[min, max]` range, mapping
/// it onto the full `-128..=127` int8 interval.
pub fn asymmetric(min: f32, max: f32) -> Self {
let span = (max - min).max(f32::MIN_POSITIVE);
let scale = span / 255.0;
// Solve for the zero point that sends `min` to -128.
let zp = (-128.0 - min / scale).round().clamp(-128.0, 127.0);
QuantParams {
scale,
zero_point: zp as i8,
}
}
/// Quantise a single real value to int8 with saturation.
#[inline]
pub fn quantize(&self, value: f32) -> i8 {
let q = (value / self.scale).round() + self.zero_point as f32;
q.clamp(-128.0, 127.0) as i8
}
/// Dequantise a single int8 value back to real units.
#[inline]
pub fn dequantize(&self, q: i8) -> f32 {
self.scale * (q as f32 - self.zero_point as f32)
}
}
impl PaddedTileLattice<f32> {
/// Find the largest absolute logical value, useful for symmetric calibration.
pub fn abs_max(&self) -> f32 {
let mut m = 0.0f32;
for (_, _, v) in self.iter_logical() {
m = m.max(v.abs());
}
m
}
/// Quantise this f32 lattice to an int8 lattice with the given parameters,
/// preserving geometry and padding layout.
pub fn quantize(&self, params: QuantParams) -> Result<PaddedTileLattice<i8>> {
let mut out = PaddedTileLattice::<i8>::zeroed(self.rows(), self.cols(), *self.geometry())?;
for (row, col, v) in self.iter_logical() {
out.set(row, col, params.quantize(v))?;
}
Ok(out)
}
}
impl PaddedTileLattice<i8> {
/// Dequantise this int8 lattice back to f32, preserving geometry and layout.
pub fn dequantize(&self, params: QuantParams) -> Result<PaddedTileLattice<f32>> {
let mut out = PaddedTileLattice::<f32>::zeroed(self.rows(), self.cols(), *self.geometry())?;
for (row, col, q) in self.iter_logical() {
out.set(row, col, params.dequantize(q))?;
}
Ok(out)
}
}