systolic-arrays-explained
← /learn · 08

Other ways to build it

GPU tensor cores against systolic arrays, Eyeriss's dataflow taxonomy with row-stationary animated, and computing in or near memory.

Loading the animation…

Concept

Weight-, output- and input-stationary are named after the one operand each keeps still. Chen, Emer and Sze's Eyeriss (ISCA 2016) asked what a dataflow should keep still when the layer is a convolution, where every number is reused several ways at once: each filter weight across every position of the image, each input pixel by every filter tap that overlaps it, each partial sum across every input channel. Their row-stationary dataflow keeps a whole row of each, and the animation runs it.

PE(i,j)(i, j) holds filter row ii and works through input row i+ji + j, one multiply per cycle: a one-dimensional convolution of two rows. When it finishes an output column it adds the partial row sum handed down by the PE above and passes the total on, so the PEs of column jj together add up the filter's rows into output row jj. Here a 3 × 3 filter runs over a 6 × 7 input on 12 PEs: 180 MACs in 17 cycles, with only 40 partial sums passed between PEs for the 20 outputs. Each filter row stays in its PE for the whole run, each input row is shared along a diagonal of PEs (rows i+ji + j), and the partial sums move one short hop at a time.

The survey by Sze, Chen, Yang and Emer (2017) sorts accelerators by what stays put: weight-stationary (the TPU's matrix unit, chapter 2), output-stationary (chapter 3), no local reuse (nothing stays in the PE: every operand comes from a shared global buffer, which can then be larger) and row-stationary. The same matrix multiply costs different amounts of buffer traffic, partial-sum movement and register traffic in each, as chapters 3 and 5 counted.

Concept

GPU tensor cores. A GPU does the same matrix multiply with a different machine. Its matrix units are small and many: a warp hands a tensor core whole tiles of AA and BB from registers through one instruction (mma.sync on the A100, wgmma on the H100), and the tile's reuse comes from the memory hierarchy, with shared memory and registers holding the tiles the threads load cooperatively from HBM. GPU Kernels Explained, chapter 7 builds that GEMM step by step.

Systolic matrix unit (TPU MXU)GPU tensor cores
Sizea few large arrays per core (128 × 128 or 256 × 256)many small units, fed by warps
Operand reuseinside the array: each value passes through a row or column of PEsin registers and shared memory, tile by tile
Who schedulesthe compiler, statically (the array is a fixed pipeline)warps, dynamically, with many in flight to hide latency
Programmed inwhole-graph compilation (XLA)kernels (CUDA, CUTLASS, Triton)
Good atlarge dense matrix multiplies at high utilisationthe same, plus irregular work the rest of the GPU does well

In their usual matrix modes both multiply low-precision inputs (bfloat16, FP16, INT8 and narrower) and accumulate in a wider format (chapter 6).

Computing in or near memory. The survey's section on near-data processing goes further: put the memory on top of the logic (3D-stacked DRAM such as HBM, with an order of magnitude more bandwidth), put logic into the memory's own die, or compute inside the memory array itself. In a resistive crossbar, each weight is a conductance; apply the inputs as voltages and Kirchhoff's current law sums the products along each wire. The survey calls it "the ultimate form of a weight stationary dataflow, as the weights are always held in place". It trades precision for that: analogue computation is sensitive to device variation, and converting between analogue and digital costs energy at every edge of the array.

Maths

For a valid 2-D convolution of an H×WH \times W input with an R×SR \times S filter, the output is E×FE \times F with E=H−R+1E = H - R + 1 and F=W−S+1F = W - S + 1:

oj,c=∑i=0R−1∑s=0S−1xi+j, c+s fi,s.o_{j,c} = \sum_{i=0}^{R-1} \sum_{s=0}^{S-1} x_{i+j,\,c+s}\, f_{i,s}.

The inner sum is a one-dimensional convolution of input row i+ji + j with filter row ii: PE(i,j)(i, j)'s job. It takes SS cycles per output column, FSFS in all, starting ii cycles late so that the partial sum from PE(i−1,j)(i - 1, j), finished one cycle earlier, is ready when PE(i,j)(i, j) finishes the same column. The run therefore takes R−1+FSR - 1 + FS cycles on R×ER \times E PEs, does R⋅E⋅F⋅SR \cdot E \cdot F \cdot S MACs, and passes (R−1)EF(R - 1) E F partial sums. The tests check the output against a direct convolution and these counts against the frames.

Code

The row-stationary PE, cut from src/lib/sa/model.ts: local cycle u=t−iu = t - i computes output column u/Su / S and tap u mod Su \bmod S:

const col = Math.floor(u / s);
const tap = u % s;
const px = x[i + j]![col + tap]!;
const pw = f[i]![tap]!;
acc[i]![j] = (tap === 0 ? 0 : acc[i]![j]!) + px * pw;