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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

//! Tile geometry: the hardware-dictated tile shape every lattice is built around.
//!
//! A TPU does not see a flat array. Its vector memory is addressed as a grid of
//! `(sublane, lane)` slots — classically 8 sublanes by 128 lanes — and its matrix
//! unit consumes square `mxu x mxu` blocks (classically 128x128). A
//! [`Geometry`] captures those three numbers so the rest of the crate can pad,
//! lay out, and iterate in the order the hardware expects.

use crate::error::{LatticeError, Result};

/// The hardware tile shape a lattice is padded and laid out to.
#[derive(Clone, Copy, PartialEq, Eq)]
pub struct Geometry {
    /// Rows per storage tile (the sublane count of a vector register).
    pub sublanes: usize,
    /// Columns per storage tile (the lane count of a vector register).
    pub lanes: usize,
    /// Side length of the square systolic matrix-unit block.
    pub mxu: usize,
}

impl Geometry {
    /// The canonical TPU v-class geometry: 8 sublanes, 128 lanes, 128x128 MXU.
    pub const TPU_V: Geometry = Geometry {
        sublanes: 8,
        lanes: 128,
        mxu: 128,
    };

    /// A small geometry handy for tests and documentation examples.
    pub const TINY: Geometry = Geometry {
        sublanes: 2,
        lanes: 4,
        mxu: 4,
    };

    /// Build a geometry, validating that every dimension is non-zero.
    pub fn new(sublanes: usize, lanes: usize, mxu: usize) -> Result<Self> {
        if sublanes == 0 || lanes == 0 || mxu == 0 {
            return Err(LatticeError::ZeroTileDimension);
        }
        Ok(Geometry {
            sublanes,
            lanes,
            mxu,
        })
    }

    /// Number of elements in a single storage tile (`sublanes * lanes`).
    #[inline]
    pub const fn tile_len(&self) -> usize {
        self.sublanes * self.lanes
    }

    /// Round `n` up to the next multiple of `to`.
    #[inline]
    pub const fn round_up(n: usize, to: usize) -> usize {
        n.div_ceil(to) * to
    }

    /// Pad a row count up to a whole number of sublane tiles.
    #[inline]
    pub const fn pad_rows(&self, rows: usize) -> usize {
        Geometry::round_up(rows, self.sublanes)
    }

    /// Pad a column count up to a whole number of lane tiles.
    #[inline]
    pub const fn pad_cols(&self, cols: usize) -> usize {
        Geometry::round_up(cols, self.lanes)
    }

    /// How many storage tiles tall a padded `rows` is.
    #[inline]
    pub const fn tile_rows(&self, rows: usize) -> usize {
        rows.div_ceil(self.sublanes)
    }

    /// How many storage tiles wide a padded `cols` is.
    #[inline]
    pub const fn tile_cols(&self, cols: usize) -> usize {
        cols.div_ceil(self.lanes)
    }

    /// True if both row and column counts already sit on tile boundaries.
    #[inline]
    pub const fn is_aligned(&self, rows: usize, cols: usize) -> bool {
        rows % self.sublanes == 0 && cols % self.lanes == 0
    }
}

impl Default for Geometry {
    #[inline]
    fn default() -> Self {
        Geometry::TPU_V
    }
}

impl core::fmt::Debug for Geometry {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(
            f,
            "Geometry {{ {}x{} tile, {}x{} mxu }}",
            self.sublanes, self.lanes, self.mxu, self.mxu
        )
    }
}