Google TPU architecture and systolic arrays: Exercises

Formula The three tools you will reuse

Let $N$ = the side length of a square array (so it has $N \times N$ cells), $f$ = clock frequency in hertz (cycles per second), $M$ = number of rows in the batch we stream through.

1. Peak throughput — how many arithmetic operations per second when every cell is busy: $Peak = 2 N^{2} f FLOP/s$ The $N^{2}$ is the count of cells; the $2$ is because one multiply-accumulate (MAC) is two operations — one multiply and one add.

2. Latency (cycle count) — how many clock ticks to push $M$ rows through: $cycles \approx (2 N - 1) + M$ The $(2 N - 1)$ is the fill/drain cost: data must ripple diagonally across the whole array before the first answer falls out the bottom.

3. Arithmetic intensity — MACs done per byte (or per operand) fetched from memory: $AI = \frac{MACs performed}{operands read}$ A bigger number means you lean less on memory and more on your silicon — you climb over the memory wall.

Level 1 — Recognition

Recall Solution

WHAT: "Systolic" comes from systole, the heartbeat — the contraction that pumps blood. WHY the metaphor: just as a heartbeat pushes blood rhythmically from chamber to chamber, the array pushes data rhythmically, one hop per clock, from cell to neighbouring cell. No cell reaches out to global memory mid-computation; everything is passed hand-to-hand. The key property captured: local, rhythmic dataflow with maximal reuse.

Recall Solution

Weights stay put (that's what "weight-stationary" means) — each cell holds one weight $B_{p j}$ for the whole computation.
Activations move horizontally (left → right), passed along via $a_{out} = a_{in}$ .
Partial sums move vertically (top → bottom), each cell adding its product: $partial_{out} = partial_{in} + a_{in} \cdot w_{stored}$ .

Level 2 — Application

Recall Solution

WHAT: plug into $Peak = 2 N^{2} f$ . $Peak = 2 \cdot 12 8^{2} \cdot (1 \times 1 0^{9}) = 2 \cdot 16384 \cdot 1 0^{9} = 3.2768 \times 1 0^{13} ops/s .$ WHY the factor 2: each cell finishes one MAC per clock = 1 multiply + 1 add = 2 ops. Converting to TOPS ( $1 TOPS = 1 0^{12}$ ops/s): $\approx 32.8 TOPS$ .

Recall Solution

WHAT: $cycles \approx (2 N - 1) + M = (2 \cdot 128 - 1) + 1000 = 255 + 1000 = 1255$ cycles. Wasted fraction: the fill/drain part is $255$ cycles, so $\frac{255}{1255} \approx 0.203 = 20.3% .$ WHY it matters: one fifth of the run went into just filling the pipeline — a warning that small batches on this array are inefficient (we tackle that head-on in L3).

Recall Solution

WHAT — count the MACs: an $N \times N \times N \times N$ product does $N^{3}$ MACs, here $51 2^{3}$ . Naïve reads: each MAC reads 1 element of $A$ and 1 of $B$ → $2 \cdot 51 2^{3}$ reads. Systolic reads: each of the $2 N^{2} = 2 \cdot 51 2^{2}$ inputs is read exactly once. $reuse = \frac{2 \cdot 51 2 ^{3}}{2 \cdot 51 2 ^{2}} = 512 = N .$ WHY: the reuse factor equals the array dimension $N$ — every operand marched across $N$ cells before leaving.

Level 3 — Analysis

Recall Solution

WHAT: solve $\frac{M}{M + 511} \geq 0.95$ . Multiply out: $M \geq 0.95 M + 0.95 \cdot 511$ , so $0.05 M \geq 485.45$ , giving $M \geq \frac{485.45}{0.05} = 9709.$ Round up: $M \geq 9709$ rows. WHY: the fixed 511-cycle fill is only "diluted" once the useful stream $M$ vastly outnumbers it. Higher utilisation targets demand exponentially larger batches — this is why TPUs are throughput machines, not latency machines.

Recall Solution

WHAT — cycles, then seconds (time = cycles ÷ $f$ ):

A: cycles $= (2 \cdot 128 - 1) + 2000 = 255 + 2000 = 2255$ . Time $= 2255/ (1 \times 1 0^{9}) = 2.255 μ s$ .
B: cycles $= (2 \cdot 256 - 1) + 2000 = 511 + 2000 = 2511$ . Time $= 2511/ (0.25 \times 1 0^{9}) = 10.044 μ s$ . Winner: A, and it is ~4.45× faster in wall-clock time. WHY: equal peak throughput hides the clock. B reaches its peak only in perfect steady state; its slow clock stretches every cycle, and its larger array pays a bigger fill. Peak FLOP/s is a ceiling, not a delivered time.

Figure — Google TPU architecture and systolic arrays

Level 4 — Synthesis

Recall Solution

WHAT — pick N: $N = 65536 = 256$ . (Square is forced here since the whole budget is one square array.) Peak: $2 \cdot 25 6^{2} \cdot (0.7 \times 1 0^{9}) = 2 \cdot 65536 \cdot 0.7 \times 1 0^{9} = 9.175 \times 1 0^{13} \approx 91.75 TOPS$ — exactly the TPUv1 ballpark. Utilisation for $M = 1024$ : fill $= 2 N - 1 = 511$ ; $\frac{1024}{1024 + 511} = \frac{1024}{1535} \approx 0.667 = 66.7% .$ The tension: with a fixed cell budget, a bigger $N$ raises peak throughput (∝ $N^{2}$ ) but also raises the fill cost ( $2 N - 1$ ), which hurts utilisation on modest batches. At $M = 1024$ this 256-array runs only two-thirds efficient — you would either enlarge the batch or, if latency-bound, prefer a smaller array. There is no free lunch; you are trading peak ceiling against fill overhead. See Dataflow architectures for the general form of this trade.

Recall Solution

WHAT: array AI $= N /2 = 256/2 = 128$ MACs/byte. Compare with the balance point $50$ : since $128 > 50$ , you are compute-bound (a good place to be — you are limited by your silicon, not by DRAM). Margin: $128/50 = 2.56 \times$ past the balance point. WHY: the whole reason for the systolic layout is to push AI up high enough to clear the memory wall; here you clear it by 2.56×, so adding more memory bandwidth would not speed you up — only more/faster cells would.

Level 5 — Mastery

Recall Solution

WHAT — build it: useful MACs in a batch $= M \cdot N^{2}$ (each of $M$ rows drives all $N^{2}$ cells once). Total cycles $= (2 N - 1) + M$ , and time $= cycles / f$ . So $delivered = \frac{2 \cdot M N ^{2}}{( 2 N - 1 + M ) / f} = 2 N^{2} f \cdot utilisation η \frac{M}{M + 2 N - 1} .$ WHY the shape: it is exactly peak × utilisation. As $M \to \infty$ , $η \to 1$ and delivered $\to 2 N^{2} f$ (the ceiling). At $M = 1$ , $η = \frac{1}{2 N} \to$ tiny — the single row barely amortises the fill. Numbers: peak $= 91.75$ TOPS (from 4.1); $η = \frac{1024}{1024 + 511} = 0.6671$ ; $delivered = 91.75 \times 0.6671 \approx 61.2 TOPS .$ Reading it: you paid for 91.75 TOPS of silicon but delivered 61.2 — the missing third leaked into pipeline fill. That gap is the L3/L4 lesson made quantitative.

Recall Solution

(a) Why 8-bit suffices for inference: an 8-bit multiplier is far smaller and lower-energy than a 32-bit float multiplier, so you fit many more MAC cells in the same silicon → higher $N^{2}$ → higher peak. Inference only needs a forward pass, and a well-quantised network tolerates 8-bit weights/activations with negligible accuracy loss (see Quantization and 8-bit inference). More cells + adequate precision = the TPUv1 bet. (b) Why 8-bit fails for training: training uses gradient descent; gradients can be tiny and are accumulated over many steps. Rounding them to 8-bit integers throws away the small values that carry learning signal — the model stops converging. Training needs floating point with wide dynamic range (bfloat16 keeps the 8 exponent bits of float32), which is why TPU v2+ switched. This is the same precision-vs-range tension the quantization note develops. The synthesis: precision is a dial, not a constant — you spend the minimum precision the task tolerates and buy throughput with the savings. Inference tolerates less; training tolerates far less rounding of gradients.