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
)
}
}