Visual walkthrough — Google TPU architecture and systolic arrays

We only need one running example, so let us fix it now and never change it:

$$A=\left(\begin{array}{cc}1& 2\\ 3& 4\end{array}\right),B=\left(\begin{array}{cc}5& 6\\ 7& 8\end{array}\right)$$

$A$ and $B$ are just grids of numbers. $A$ has 2 rows and 2 columns; so does $B$. We want their product $C=A\times B$, which is also a $2\times 2$ grid. Our whole mission: build a little machine of cells that spits out $C$ without ever re-reading a number from memory.

---

Step 1 — What "multiply two matrices" even asks for

WHAT. The rule for matrix multiplication says: the number in row $i$, column $j$ of $C$ — written $C_{ij}$ (read "C-eye-jay", the entry at row $i$ and column $j$) — is built by walking across row $i$ of $A$ and down column $j$ of $B$, multiplying the pairs you meet, and adding them up.

$$C_{ij}=\underset{\text{1st pair}}{\underbrace{A_{i1}{B}_{1j}}}+\underset{\text{2nd pair}}{\underbrace{A_{i2}{B}_{2j}}}$$

Here $A_{i1}$ means "the entry of $A$ in row $i$, column $1$." The little numbers are just addresses on the grid.

WHY. Before we build hardware, we must know exactly what sum each output is. Everything the machine does is a physical realisation of this one formula. If we get the target wrong, no clever wiring saves us.

PICTURE. The figure shows $C_{11}$ being formed: the top row of $A$ (orange) slides across the left column of $B$ (teal); the crossings are the products we add.

Concretely, $C_{11}=1\cdot 5+2\cdot 7=5+14=19$. Hold that number; the machine must reproduce it.

---

Step 2 — The one cell that does a single "multiply-and-pile-up"

WHAT. We build the whole array out of one repeated part: a processing element (PE) — a tiny box that can do exactly one multiply and one add per tick of the clock. "One tick of the clock" (one cycle) is the heartbeat: every cell acts once, all together, then again. Inside the cell lives a stored weight $w$ and a running total called the accumulator, written $\text{acc}$.

Each cycle the cell does:

$${\text{acc}}_{\text{new}}=\underset{\text{pile so far}}{\underbrace{{\text{acc}}_{\text{old}}}}+\underset{\text{activation arriving}}{\underbrace{a_{\text{in}}}}\cdot \underset{\text{stored weight}}{\underbrace{w}}$$

• $a_{\text{in}}$ = the number that just walked in from the left (an entry of $A$).
• $w$ = the number parked inside this cell (an entry of $B$).
• $\text{acc}$ = the growing sum that will become one entry of $C$.

WHY. Look back at Step 1: an output is a sum of products. A "multiply-then-add-to-a-total" box is the smallest machine that grows such a sum. We pile up in place so the partial answer never has to leave the cell — that is how we dodge the memory wall.

PICTURE. One cell, labelled: weight $w$ frozen inside, $a_{\text{in}}$ entering left, and the accumulator ticking up.

---

Step 3 — Passing the activation along (read once, use many times)

WHAT. A cell doesn't just consume $a_{\text{in}}$; it also hands it to the neighbour on its right, unchanged:

$$a_{\text{out}}=a_{\text{in}}$$

That equation looks trivial, but it is the magic. It says: the same activation, after being used here, keeps travelling and gets reused by the next cell one cycle later.

WHY. In the matmul rule, the number $A_{i1}$ shows up in every entry of row $i$ of $C$ ($C_{i1},C_{i2},\dots$). So $A_{i1}$ should be read from memory once and then physically walked across all the cells that need it. $a_{\text{out}}=a_{\text{in}}$ is that walk. Reading once and reusing $N$ times is exactly what raises the arithmetic intensity.

PICTURE. A single activation (orange dot) marching left→right through a row of cells, being multiplied by a different parked weight in each, while partial sums fall downward.

---

Step 4 — Park the weights: laying $B$ into the grid

WHAT. Build a $2\times 2$ grid of cells. Before any activation flows, we load $B$ into the cells, one weight per cell, matched to position: cell in row $p$, column $j$ stores $B_{pj}$.

$$\begin{array}{ccc}& \text{col }1& \text{col }2\\ \text{row }1& B_{11}=5& B_{12}=6\\ \text{row }2& B_{21}=7& B_{22}=8\end{array}$$

Read the layout carefully: column $j$ of the array holds column $j$ of $B$. That is deliberate — column $j$ of the array will produce column $j$ of $C$.

WHY. Each output $C_{ij}$ needs the pair $(B_{1j},B_{2j})$. By stacking column $j$ of $B$ vertically in array-column $j$, an activation streaming down/through that column meets exactly the weights it must multiply. The geometry does the bookkeeping for us.

PICTURE. The $2\times 2$ grid with $5,6,7,8$ frozen in the four cells, colour-coded by which column of $C$ they serve.

---

Step 5 — Stream the activations and watch $C_{11}$ grow

WHAT. Now push a row of $A$ into the grid. Take row 1 of $A$, which is $(A_{11},A_{12})=(1,2)$, and feed it into the cells of array-column 1 (the ones holding $5$ and $7$). The accumulator in that column collects:

$$\text{acc}=A_{11}{B}_{11}+A_{12}{B}_{21}=1\cdot 5+2\cdot 7=5+14=19=C_{11}✓$$

Feed the same row 1 into array-column 2 (holding $6$ and $8$):

$$\text{acc}=A_{11}{B}_{12}+A_{12}{B}_{22}=1\cdot 6+2\cdot 8=6+16=22=C_{12}✓$$

Each partial sum (${\text{partial}}_{\text{in}}$) enters a cell from the top, gets one product added, and leaves at the bottom as ${\text{partial}}_{\text{out}}$:

$${\text{partial}}_{\text{out}}={\text{partial}}_{\text{in}}+a_{\text{in}}\cdot w$$

WHY. Compare with Step 1's target $C_{11}=1\cdot 5+2\cdot 7$. The array reproduced it without a single instruction fetch — the wiring is the program. The sum that used to unfold in time (loop iterations on a CPU) now unfolds in space (down a column of cells).

PICTURE. Row $(1,2)$ entering column 1; the partial sum accumulating $5$, then $5+14=19$ as it drops through the two cells.

---

Step 6 — The skew: why inputs enter diagonally (the timing edge case)

WHAT. There is a subtlety we glossed over: cells act all at once each cycle, but data needs time to travel. If we shoved every activation in on the same cycle, the partial sums arriving from above would meet the wrong activation. The fix is to stagger (skew) the inputs: row 2 of the data enters one cycle later than row 1, row 3 one later still, forming a slanted diagonal wavefront.

$$\text{delay of row }r=(r-1)\text{cycles}$$

WHY. For a product to be correct, the activation and the partial sum must rendezvous in the same cell on the same cycle. Because a partial sum takes one hop per cycle to descend, and an activation takes one hop per cycle to cross, we must offset their launch times so their paths cross at the right instant. Skewing by one cycle per row synchronises the whole grid.

PICTURE. A slanted "wavefront" of inputs: the top row of data already deep inside the array while the next row is only just entering — the tell-tale diagonal staircase.

---

Step 7 — Counting the cost: fill, drain, and why arrays crave big batches

WHAT. For an $N\times N$ array (here $N=2$; TPUv1 has $N=256$), the very first result cannot appear until data has propagated diagonally across the whole grid. That warm-up is the fill; emptying the pipeline afterward is the drain. Together they cost about $2N-1$ cycles. After the pipeline is full, one fresh row of results emerges every cycle. So pushing $M$ rows of input through takes:

$$\text{cycles}\approx \underset{\text{fill \& drain, fixed}}{\underbrace{(2N-1)}}+\underset{\text{one row per cycle}}{\underbrace{M}}$$

Efficiency — the fraction of cycles doing useful streaming — is:

$$\text{efficiency}=\frac{M}{M+(2N-1)}$$

WHY. Look at the shape of that fraction. If $M$ is tiny, the fixed $2N-1$ dominates and most cycles are wasted warming up. If $M$ is huge, the fixed cost vanishes into the average. That is precisely why the TPU is a throughput machine for large batches, not a low-latency machine — an idea the Roofline model and arithmetic-intensity view make quantitative.

Numbers. For $N=256$: fill/drain $\approx 2\cdot 256-1=511$. To hit $90\%$ efficiency:

$$\frac{M}{M+511}\ge 0.9\Rightarrow M\ge 0.9M+459.9\Rightarrow 0.1M\ge 459.9\Rightarrow M\ge 4599$$

so about 4600 rows (the parent's "$\ge 4608$" using the rounder $2N=512$).

PICTURE. The efficiency curve rising from near-zero toward $1$ as batch size $M$ grows, with the $90\%$ line marked.

---

The one-picture summary

Everything above collapses into a single frame: weights parked in the grid, a skewed diagonal of activations flowing in from the left, partial sums cascading down, and finished entries of $C$ dropping out the bottom.

Recall Feynman retelling — say it back in plain words

We wanted to multiply two grids of numbers. Every answer-box is just a row-times-column "dot product" — multiply the matching pairs and pile them up. So we built one tiny box that does one multiply-and-add per heartbeat, and we glued a whole grid of them together. We parked the second matrix's numbers inside the boxes (each stays put — weight-stationary). Then we fed the first matrix in from the left, and here's the trick: each box, after using an incoming number, hands it to its neighbour unchanged — so every number is read from memory once but reused all across a row. Partial sums slide downward through each column, growing until they equal an answer-box, then fall out the bottom. To make the incoming number and the falling sum meet in the right box at the right beat, we release the input rows in a slanted staircase (one beat later per row). Warming the pipeline up and draining it costs a fixed number of beats, so tiny jobs waste time warming up — that's why the TPU loves big batches, where one answer-row pops out every single beat.

Recall Quick self-test

Why is the activation passed unchanged to the right neighbour? ::: So it is read from memory once yet reused by every cell in the row — the reuse trick that beats the memory wall ($a_{\text{out}}=a_{\text{in}}$). Why must inputs enter diagonally staggered? ::: So each activation and its descending partial sum arrive in the same cell on the same cycle; skewing by one cycle per row synchronises the rendezvous. For a $256\times 256$ array, roughly what batch $M$ gives $90\%$ efficiency? ::: About $4600$ rows (fill/drain $\approx 511$, solve $M/(M+511)\ge 0.9$). Why does throughput scale as $N^2$, not $N$? ::: There are $N^2$ cells, each doing a MAC every cycle — the parallelism is spatial, so peak $=2N^{2}f$ FLOP/s.

---

Related: ASIC vs FPGA vs general-purpose processors · Quantization and 8-bit inference · GPU architecture and SIMT execution