The TPU and Systolic Arrays

Strip a modern neural network down to its engine and you find one operation doing almost all the work: matrix multiplication. A layer takes a batch of activations, multiplies them by a matrix of learned weights, and passes the result on. Training and inference are, to a first approximation, a blizzard of dense matrix multiplies — billions of multiply-then-add operations, all identical in shape. When one operation dominates a whole domain, you have the perfect target for a domain-specific architecture.

A general CPU multiplies matrices the slow, expensive way: for every scalar multiply it fetches an instruction, reads the register file, moves data through the cache hierarchy, and writes the result back. The arithmetic is a whisper; the bookkeeping is a roar. In 2015 Google shipped a chip that deleted the roar. The Tensor Processing Unit (TPU) is built around an idea from 1978 — the systolic array — and on dense matrix multiply it beat the GPUs of its day by an order of magnitude in operations per joule.

The systolic array: a heartbeat of arithmetic

A systolic array is a grid of tiny, identical cells, each one a multiply-accumulate (MAC) unit: it multiplies two numbers and adds the product to a running total. There is no instruction fetch, no decode, no register file — just arithmetic wired to its neighbours. Data does not sit in a memory waiting to be fetched; it flows through the grid in lock-step with the clock, like blood pulsing through tissue. Each value that enters is reused by every cell it passes, so a single memory read feeds a whole row of multiplies.

The figure shows the weight-stationary dataflow the TPU uses. Preload each cell with one weight and leave it there. Then stream the activations in from the left — each moves one cell to the right every clock tick — while partial sums accumulate downward. Reveal the flows one at a time:

Follow one activation: it enters on the left, gets multiplied by the stationary weight in the first cell, and the product is added to a partial sum travelling down that column. The same activation then steps right and is multiplied by the next weight, and so on across the row. One number fetched from memory, n multiplies performed. Multiply that reuse across the whole grid and the memory system barely breaks a sweat.

Why it crushes a GPU on ops per joule

The killer feature is data reuse. Multiplying two n \times n matrices takes n^3 multiply-accumulates but touches only n^2 input numbers on each side. A naive machine that re-reads operands from memory for every MAC moves O(n^3) data; the systolic array loads each operand once and lets it ripple through, moving only O(n^2). That factor-of-n cut in memory traffic is the whole game, because moving a byte costs far more energy than multiplying two of them.

a 2-D grid of simple MAC cells, each wired only to its neighbours — no fetch, no decode, no register file;
data flows rhythmically ("systolic", after the heart's pumping beat), and every value entering the array is reused by many cells;
weight-stationary dataflow: weights are preloaded and held fixed; activations stream in and partial sums accumulate as they pass through;
memory traffic for an n\times n multiply drops from O(n^3) to O(n^2) — the source of its energy win;
at peak, a k \times k array performs k^2 MACs every clock cycle.

Google's first TPU put a 256 \times 256 array of 8\text{-bit} MAC units on one chip — 65{,}536 multipliers firing together, about 92 tera-operations per second, all fed by one flow of weights and activations. A GPU of the era spread its arithmetic across thousands of general-purpose lanes, each carrying the overhead of a flexible instruction pipeline. For dense matmul, the specialised grid simply wastes less.

Counting the reuse

Let's make the reuse concrete. Multiply two n \times n matrices and count the multiply-accumulates against the number of weights you had to load. The ratio is how many useful MACs you squeeze out of each weight fetched — the systolic array's superpower.

// Dense matrix multiply, counting MACs and weight reuse. function matmulStats(n: number) { const A: number[][] = Array.from({ length: n }, (_, i) => Array.from({ length: n }, (_, j) => i + j)); const B: number[][] = Array.from({ length: n }, (_, i) => Array.from({ length: n }, (_, j) => (i === j ? 1 : 0))); // identity const C: number[][] = Array.from({ length: n }, () => new Array(n).fill(0)); let macs = 0; for (let i = 0; i < n; i++) for (let j = 0; j < n; j++) for (let k = 0; k < n; k++) { C[i][j] += A[i][k] * B[k][j]; macs++; } const weightsLoaded = n * n; // each weight loaded ONCE into its cell const reuse = macs / weightsLoaded; console.log(`n = ${n}: ${macs} MACs, ${weightsLoaded} weights loaded once`); console.log(`reuse factor = ${reuse} MACs per weight (= n)`); console.log(`A x I equals A? ${JSON.stringify(C) === JSON.stringify(A)}`); } matmulStats(4); matmulStats(8);

The reuse factor comes out to exactly n: an 8 \times 8 multiply performs 512 MACs while loading each weight only once. On a full 256\times256 array, each weight loaded is reused 256 times before it retires. That is why the memory bus, so often the bottleneck, gets to relax.

The name comes from H. T. Kung and Charles Leiserson, who described these arrays in 1978. Systole is the contraction phase of a heartbeat — the pulse that pushes blood rhythmically through your body. Kung saw the same picture in his array: a global clock beats, and on every beat data is pumped one step further through a mesh of little processors, each doing a scrap of work and passing the result along. Nothing waits in a memory to be summoned; everything moves in a coordinated pulse. Four decades later that 1978 idea is the beating heart of every TPU and the tensor cores inside modern GPUs.

The efficiency depends on keeping every cell busy every cycle with a predictable flow. Two things spoil that. First, sparsity and irregularity: if the matrix is mostly zeros, or the computation branches, the rigid grid can't skip the wasted work — it multiplies by zero at full power. Second, utilisation: a giant 256\times256 array running a tiny 3\times3 convolution leaves almost every cell idle, and the array also spends cycles filling and draining (roughly 2k-1 cycles of latency before the first result appears). The systolic array is not a general matrix machine — it is a machine for big, dense, steady matrix multiplies, and it is glorious at exactly that and nothing else.

The takeaway

The TPU is the textbook domain-specific architecture: it identified the one operation that dominates its domain, built silicon that does only that, and won on operations per joule by an order of magnitude — at the cost of doing nothing else. Next we look outward, at the networks that must feed and connect thousands of these accelerators, because a single chip, however fast, is never the whole story at scale.