Google TPU architecture and systolic arrays

WHY does the TPU exist?

The WHY: Consider a fully-connected layer $y = Wx$. To compute it, you must do many multiply-accumulates. On a CPU each MAC involves:

• fetch weight from register/cache,
• fetch activation,
• multiply, add to accumulator,
• write accumulator back.

The arithmetic is cheap; the data movement dominates energy and time. This is the memory wall. The fix: build hardware where a value, once loaded, is reused maximally before it leaves.

WHAT is a systolic array?

Key properties (WHAT to remember):

• Local wiring only → short wires → fast clock, low energy.
• Massive reuse → each operand read from memory once, used by a whole row/column.
• No instruction fetch per op → the dataflow is the program.

HOW does it compute a matrix multiply? (Derivation from scratch)

We want $C = A \times B$ where $A$ is $m\times k$ and $B$ is $k\times n$. Element:

$C_{ij} = \sum_{p=1}^{k} A_{ip}\,B_{pj}$

Step 1 — map the sum onto a grid. Put PE $(i,j)$ in charge of computing $C_{ij}$. It must accumulate $k$ products. So each PE needs an internal accumulator:

$\text{acc}_{ij} \leftarrow \text{acc}_{ij} + A_{ip}B_{pj}$

Why this step? Each output needs its own running sum, so accumulation must be local to the cell.

Step 2 — deliver operands without global memory. In a weight-stationary array (what the TPU uses), preload weights $B_{pj}$ into the cells (weight stays put). Then stream activations $A_{ip}$ horizontally left→right, and stream partial sums vertically top→bottom.

At each clock a cell does: $\text{partial}_{\text{out}} = \text{partial}_{\text{in}} + a_{\text{in}}\cdot w_{\text{stored}}, \qquad a_{\text{out}} = a_{\text{in}}$

Why this step? $a_{\text{out}}=a_{\text{in}}$ means the activation is passed along to the next cell — that's the one-read-many-uses trick.

Step 3 — skew the input (staggering). Because it takes cell $(i,j)$ some cycles for data to reach it, inputs are fed diagonally staggered (each row delayed by one clock). This keeps every cell busy on the right operands at the right time.

Step 4 — count the cost. A naïve implementation of an $n\times n$ multiply reads $O(n^3)$ operands from memory. The systolic array reads each of the $2n^2$ input values once and does $n^3$ MACs inside — so the arithmetic intensity is:

$\text{Arithmetic intensity} = \frac{n^3 \text{ MACs}}{2n^2 \text{ reads}} = \frac{n}{2}$

Why this matters? Higher arithmetic intensity = fewer memory accesses per compute = you break through the memory wall. For a $256\times256$ array that's ~128 MACs per byte fetched.

The full TPU block diagram (WHAT sits around the array)

• Matrix Multiply Unit (MXU) — the 256×256 systolic array. The workhorse.
• Weight FIFO — feeds weights from off-chip DRAM into the array.
• Unified Buffer (UB) — large (24 MB) on-chip SRAM holding activations; feeds the array and stores results. Keeps activations on-chip → avoids DRAM round-trips.
• Accumulators — collect the column partial sums coming out the bottom.
• Activation unit — applies ReLU / pooling etc. after the multiply.
• No cache, no branch predictor, no out-of-order — deliberately simple; all die area goes to arithmetic. This is the domain-specific bet.

Worked examples

Common mistakes (Steel-manned)

Flashcards

Connections

• GPU architecture and SIMT execution — the alternative: temporal parallelism + caches vs. TPU's spatial dataflow.
• Memory wall and arithmetic intensity — the roofline reason TPUs exist.
• Matrix multiplication algorithms — the operation being accelerated.
• ASIC vs FPGA vs general-purpose processors — TPU as a domain-specific ASIC.
• Dataflow architectures — systolic arrays are a classic dataflow design (Kung & Leiserson, 1978).
• Quantization and 8-bit inference — why TPUv1 uses INT8 MACs.
• Roofline model — visualizing when you're compute- vs memory-bound.

Concept Map

Hinglish (regional understanding)