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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
//! A CPU reference simulator of weight-stationary systolic matmul.
//!
//! A TPU multiplies matrices by streaming one operand through a square grid of
//! multiply-accumulate cells while the other operand sits stationary in the grid.
//! This module reproduces the *blocking and accumulation order* of that dataflow
//! — weights loaded one `mxu x mxu` block at a time, products accumulated into an
//! f32 accumulator — so a result computed here matches what the hardware returns
//! bit-for-bit in the f32 case and closely in the bf16 case. It also reports the
//! work performed so you can reason about utilisation before deploying.
use crate::bf16::Bf16;
use crate::error::{LatticeError, Result};
use crate::geometry::Geometry;
use crate::lattice::PaddedTileLattice;
/// A numeric element the systolic simulator can stream and accumulate.
pub trait Numeric: Copy + Default {
/// Widen to the f32 accumulator domain.
fn to_acc(self) -> f32;
/// Narrow an f32 accumulator back to the element type.
fn from_acc(value: f32) -> Self;
}
impl Numeric for f32 {
#[inline]
fn to_acc(self) -> f32 {
self
}
#[inline]
fn from_acc(value: f32) -> Self {
value
}
}
impl Numeric for Bf16 {
#[inline]
fn to_acc(self) -> f32 {
self.to_f32()
}
#[inline]
fn from_acc(value: f32) -> Self {
Bf16::from_f32(value)
}
}
/// Work counters describing a single simulated matmul.
#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
pub struct SystolicStats {
/// Total scalar multiply-accumulate operations performed.
pub macs: u64,
/// Number of stationary `mxu x mxu` weight blocks loaded into the array.
pub weight_block_loads: u64,
/// Number of `mxu x mxu` output blocks produced.
pub output_blocks: u64,
/// MAC slots that were spent on padding rather than logical data.
pub padding_macs: u64,
}
impl SystolicStats {
/// Total MAC slots the array spun through, useful plus padding.
pub fn total_macs(&self) -> u64 {
self.macs + self.padding_macs
}
/// Fraction of MAC work that did useful (non-padding) computation.
pub fn utilisation(&self) -> f64 {
let total = self.total_macs();
if total == 0 {
0.0
} else {
self.macs as f64 / total as f64
}
}
}
impl<T: Numeric> PaddedTileLattice<T> {
/// Multiply two lattices, returning the product and dataflow statistics.
///
/// `self` is `m x k`, `rhs` is `k x n`, the result is `m x n`. The product is
/// accumulated entirely in f32 (like the hardware accumulator) and narrowed to
/// `T` exactly once per output element.
pub fn matmul_with_stats(
&self,
rhs: &PaddedTileLattice<T>,
) -> Result<(PaddedTileLattice<T>, SystolicStats)> {
if self.cols() != rhs.rows() {
return Err(LatticeError::ContractionMismatch {
lhs_cols: self.cols(),
rhs_rows: rhs.rows(),
});
}
if self.geometry() != rhs.geometry() {
return Err(LatticeError::GeometryMismatch);
}
let geom: Geometry = *self.geometry();
let m = self.rows();
let k = self.cols();
let n = rhs.cols();
let mxu = geom.mxu;
// Dense f32 accumulator, exactly like a TPU's matrix accumulator banks.
let mut acc = vec![0.0f32; m * n];
let mut stats = SystolicStats::default();
for i0 in (0..m).step_by(mxu) {
for j0 in (0..n).step_by(mxu) {
stats.output_blocks += 1;
for k0 in (0..k).step_by(mxu) {
stats.weight_block_loads += 1;
let i_end = (i0 + mxu).min(m);
let j_end = (j0 + mxu).min(n);
let k_end = (k0 + mxu).min(k);
for i in i0..i_end {
for j in j0..j_end {
let mut sum = acc[i * n + j];
for kk in k0..k_end {
let a = self.get(i, kk).unwrap().to_acc();
let b = rhs.get(kk, j).unwrap().to_acc();
sum += a * b;
stats.macs += 1;
}
acc[i * n + j] = sum;
}
}
}
}
}
// Account for the padding MACs the hardware would also have to spin through:
// the array always runs whole mxu-blocks, padding included.
let padded_macs = (Geometry::round_up(m, mxu) as u64)
* (Geometry::round_up(n, mxu) as u64)
* (Geometry::round_up(k, mxu) as u64);
stats.padding_macs = padded_macs.saturating_sub(stats.macs);
let dense: Vec<T> = acc.into_iter().map(T::from_acc).collect();
let out = PaddedTileLattice::from_dense(m, n, &dense, geom)?;
Ok((out, stats))
}
/// Multiply two lattices, discarding the dataflow statistics.
pub fn matmul(&self, rhs: &PaddedTileLattice<T>) -> Result<PaddedTileLattice<T>> {
self.matmul_with_stats(rhs).map(|(out, _)| out)
}
}