systile 0.1.0

A TPU-native tiled tensor data structure: the Systolic Tile Lattice. Padding-aware, sublane/lane laid out, bf16/int8 first, with a CPU reference simulator of systolic dataflow.
Documentation 
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)
    }
}