Systolic Array: 2D Weight-Stationary Dataflow

Array Configuration

Array size N × N 5 × 5

Clock speed 1.5 cyc/s

Live Readouts

Phase

–

Global cycle

0 / 0

MACs computed

0 / 0

Latency: weight-stat. vs output-stat.

–

How it works (2D model)

The 3D version of this simulator is an output-stationary array: both operand matrices stream in diagonally and each PE's accumulator sits fixed in place for the whole run. This 2D counterpart is a genuinely different, independently-computed design: a weight-stationary array — the dataflow Google's original TPU MXU used. Here the weight matrix B is loaded once into the grid and never moves; only activations and partial sums travel.

Phase 1 — load (real extra latency this design pays):
  weight B[k][j] shifts in diagonally,
  arriving at PE(k,j) at load-cycle L = k+j
  → full grid loaded after 2N−2 cycles

Phase 2 — stream + drain:
  activation A[i][k] enters row k at
  compute-cycle t = i+k  → moves right
  partial sum enters column j at row 0,
  moves down, accumulating:
    psum(i,k,j) = Σ A[i][k']·B[k'][j], k'=0..k
  reaches PE(k,j) at t = i+k+j
  C[i][j] finalizes when k reaches N−1

Total latency = (2N−2) + (3N−2) = 5N−4
  vs. 3N−2 for the 3D output-stationary array

Amber cells — the stationary weight grid; a cell lights up the instant its weight has finished shifting into place during the load phase.
Green dot — an activation value A[i][k] travelling rightward through row k this cycle.
Blue dot — the running partial sum for one output cell, travelling downward through column j; brightness scales with its current magnitude.
Output readout (bottom strip) — cell (i,j) fills in the instant its dot product finishes draining out of row N−1 — the weight grid above never stores a finished output at all, only in-flight partial sums.
The compute-phase schedule t=i+j+k is the same combinatorial wavefront as the 3D model — that's a mathematical fact about any valid systolic schedule for a 3-index contraction, not a coincidence. What genuinely differs is which value is stationary and which two move, and the real cost that weight-stationary designs must pay: a full 2N−2-cycle reload stall every time the weight matrix changes, which output-stationary arrays avoid entirely by streaming both operands from cycle zero.

Stationary weight B[k][j] Activation moving right Partial sum moving down Finished output C[i][j]