Systolic Arrays — Building the Heart of the TPU From Scratch

Piling on multipliers doesn't make a chip faster, because moving data costs more than the arithmetic does. This piece takes the grid at the center of the TPU — the systolic array — and shows exactly how it maps the triple loop of matrix multiplication onto rows, columns and clock ticks, from a 2×2 hand trace to a cycle-accurate simulator to the batch-size and shape rules you'll actually tune.

In-Datacenter Performance Analysis of a Tensor Processing Unit

Run to the well, or pass the bucket

Tell a hundred workers to fetch water from a well and pour it on a field. Each of them could grab a bucket and run back and forth. But the fast arrangement is a line stretching from well to field, buckets moving hand to hand. Nobody runs. Each person only passes to a neighbor. Water still arrives without a gap.

Chips work the same way. Multipliers are tiny in silicon, and you can pack in as many as you like. But for one of them to do one multiply-accumulate, it has to read two inputs and write one result. When those reads and writes go to a register file or SRAM, moving the data costs more time and more energy than the arithmetic itself. Grow to a thousand multipliers and, unless you can feed all thousand every cycle, most of them sit idle. That "the delivery side is the bottleneck" pattern is the subject of The Memory Wall, and it is the starting point of AI chip design.

The systolic array is the bucket-brigade answer. Pack the arithmetic units into a grid and let each unit exchange data only with its immediate neighbors. Data that enters the chip once never goes back to memory: it flows across the grid from edge to edge, putting every unit along its path to work before it leaves. Systolic refers to the contraction of the heart — data advances one square per beat, like a pulse — and H. T. Kung and Charles Leiserson named it that in the late 1970s. The TPU is that forty-year-old idea revived and aimed at exactly one operation: matrix multiplication.

What one square of the grid does

A single square — a PE (Processing Element) — is startlingly simple inside. One multiplier, one adder, and a handful of registers to hold a value for a single beat. That's it.

It has three inputs: an activation $a$ arriving from the left neighbor, a partial sum $p$ descending from the neighbor above, and a weight $w$ that just sits there inside it. Every cycle, the PE does two things at once.

Said in words: take the number from above, add the product of the number from the left and your own weight, drop it downward — and pass the number from the left straight through to the right. Here $p_{\text{in}}$ is the running total handed down by the PE above, $a$ is the input value that arrived from the left, and $w$ is the weight this PE is holding.

Three things follow. A PE is wired only to its neighbors, so there are no wires crossing the chip and the whole thing can run at a high clock. The weight $w$ never moves once placed, so weight reads disappear. And partial sums are handed downward from PE to PE, so intermediate results never get written back to memory. Same logic as nobody running to the well.

"Never moves once placed" is literal. Before the multiplication begins you stream the weights into the grid once and burn them into the PEs' registers; after that you can push thousands of inputs through without touching the weights again. Neural network inference applies the same weights to enormous numbers of inputs, so that assumption holds comfortably. The flip side is that any workload which swaps weights often throws this advantage away.

The "multiply and add" a PE ticks out every beat is exactly one term of a dot product.

Mapping a matrix product onto grid and time

Now the real question: how does this grid compute $XW$? Let $X$ be a $T \times N$ matrix holding $T$ input vectors stacked vertically, and $W$ an $N \times M$ weight matrix.

Unpacked in words, a single number in row $t$, column $j$ of the output is what you get by walking along one input vector and one column of weights together, multiplying each pair, and adding all of it up. Here $t$ says which input vector, $j$ says which element of the output, and $i$ indexes the terms being summed. Written as a program, that's a triple loop.

The whole trick of the systolic array is that those three loops are assigned to rows, columns and time. Put $i$ (the summation index) on the vertical axis of the grid, $j$ (the output element) on the horizontal axis, and $t$ (which input vector) on the clock.