Why systolic?

The memory wall, and Kung's answer: fetch each number once and pass it from processor to processor, so a grid of n × n multipliers needs only 2n words a cycle.

Loading the animation…

Concept

A neural network spends almost all of its arithmetic in matrix multiplies: $C = AB$ , where every output $c_{mn}$ is a dot product of a row of $A$ and a column of $B$ . The multiplications are cheap to build. Feeding them is not: fetching an 8-bit operand from a large on-chip memory costs far more energy than multiplying it, and an off-chip memory cannot deliver operands nearly as fast as a chip full of multipliers can consume them. This is the memory wall.

H. T. Kung's answer, set out with Charles Leiserson in 1978 and argued in his 1982 paper Why systolic architectures?, is to fetch each number once and then pass it from one processing element (PE) to its neighbour, every clock cycle, like blood pumped through a body (hence systolic). Each PE does one multiply-accumulate (MAC) per cycle on whatever arrives and passes the operands on. Only the PEs on the edges talk to memory; every wire inside the array is short and connects two neighbours.

The animation above is the model's weight-stationary array, the arrangement in the first TPU. Weights (vermillion) sit in the PEs; activations (blue) enter from the left and step right one PE per cycle; partial sums (sky) step down one PE per cycle, picking up one product in each PE; finished results (purple) leave from the bottom. Its 8 × 4 by 4 × 4 multiply takes 18 cycles, including loading the weights. Chapter 2 runs the same machine with every register's value on screen.

Loading the animation…

Concept

The second animation makes the argument with numbers. A memory delivers $\beta$ words per cycle to a grid of $n \times n$ multipliers. If every MAC fetches its own two operands, at most $\beta/2$ MACs can run per cycle however many multipliers you build: with $\beta =$ 32 and a 16 × 16 grid, 16 of its 256 multipliers are busy (6.3%). In a systolic array a word entering at the left edge is used by every PE along its row, and a word entering at the top by every PE down its column, so the whole grid needs only $2n =$ 32 words per cycle and all 256 MACs run.

The ratio of MACs to words fetched, the arithmetic intensity at the array's edge, grows as $n/2$ : 8 MACs per word at $n = 16$ . That is why matrix units are big. The first TPU's matrix unit is a 256 × 256 grid of 65,536 8-bit MACs (Jouppi et al., 2017), so each operand that enters it is used 256 times.

Maths

For an output-stationary array (chapter 3), PE $(i, j)$ computes $c_{ij} = \sum_k a_{ik} b_{kj}$ . Element $a_{ik}$ enters row $i$ and visits the $n$ PEs of that row; element $b_{kj}$ enters column $j$ and visits its $n$ PEs. Per cycle the array therefore reads $n$ words of $A$ and $n$ of $B$ and performs $n^2$ MACs: an intensity of $n^2 / 2n = n/2$ MACs per word read. Without reuse each MAC reads 2 words, an intensity of $\tfrac12$ .

Given a memory bandwidth of $\beta$ words per cycle, the sustainable rate is the smaller of the compute peak and the bandwidth times the intensity:

\text{rate} = \min\Bigl(n^2,\ \beta \cdot \frac{n}{2}\Bigr),

which is the roofline model with the array's edge as the memory interface. A 4 × 4 array reads 8 words per cycle where 16 independent MACs would read 32.

Code

The animation's states, cut from src/lib/sa/model.ts (throughputs are stored doubled so they stay integers):

export function reuseSteps(beta: number, nMax: number): ReuseStep[] {
  const out: ReuseStep[] = [];
  for (let n = 1; n <= nMax; n++) {
    const peak = n * n;
    out.push({
      n,
      peak,
      naive2: Math.min(2 * peak, beta),
      systolic2: Math.min(2 * peak, beta * n),
      wordsPerCycle: 2 * n,
      macsPerWord2: n,
    });
  }
  return out;
}