How Do TPU Systolic Arrays Reuse Matrix Weights?

Systolic arrays in Tensor Processing Units reuse matrix multiplication weights by preloading them into local registers inside a two-dimensional grid of processing elements and keeping them stationary during computation. Instead of repeatedly fetching weights from a central register file for every operation, activations flow rhythmically from neighbor to neighbor across the array while partial sums accumulate along the orthogonal axis. This weight-stationary dataflow drastically reduces memory bandwidth demands and energy consumption compared to classical von Neumann processors.

The Register File Bottleneck in Classical Processors

In classical CPUs and GPUs, executing a matrix multiplication \(C = A \times B\) requires each arithmetic logic unit (ALU) to fetch operands from a shared register file or cache, perform a multiply-accumulate (MAC) operation, and write the result back. For a matrix multiply of size \(N \times N\), there are \(N^3\) total MAC operations. Even with cache hierarchies, repeatedly reading weights and writing partial results through a centralized register file creates a severe energy and bandwidth bottleneck, often consuming more power for data movement than for the arithmetic itself.

Weight-Stationary Dataflow in a Systolic Array

A TPU systolic array replaces the centralized fetch-execute-writeback cycle with a tightly coupled 2D grid of simple processing elements (PEs), each containing a multiplier, an adder, and a small local register. In a weight-stationary architecture, the weight matrix \(W\) is loaded into the local registers of the PEs before the computation begins, with each PE holding a single weight \(W_{i,j}\).

Once the weights are in place, they remain stationary throughout the entire matrix multiplication. Input activations \(X\) are streamed into the left edge of the array and passed horizontally from one PE to the next on every clock cycle.

Rhythmic Data Movement and Partial Sum Accumulation

As an activation \(X_{k,i}\) arrives at the PE holding \(W_{i,j}\), the PE computes the product \(X_{k,i} \cdot W_{i,j}\) and adds it to an incoming partial sum \(P_{\text{in}}\) arriving from the PE directly above it:

\[P_{\text{out}} = P_{\text{in}} + X_{k,i} \cdot W_{i,j}\]

On the next clock cycle, two things happen simultaneously:

• The activation \(X_{k,i}\) is forwarded to the right neighbor PE.
• The updated partial sum \(P_{\text{out}}\) is passed downward to the PE below.

To ensure that the correct activation meets the correct partial sum at every PE, the input rows of \(X\) are skewed in time (staggered by one clock cycle per row). By the time a partial sum exits the bottom edge of the systolic array, it has accumulated the complete dot product for one element of the output matrix.

Why This Eliminates Repeated Register File Accesses

Because each weight \(W_{i,j}\) stays inside its PE's local register while an entire batch of activations flows past it, the weight is read from main memory or the weight FIFO only once per batch. All subsequent reuse happens locally within the PE's own flip-flops, and all intermediate activations and partial sums travel exclusively over short, point-to-point wires between adjacent PEs. This eliminates the need for a large, power-hungry central register file to service \(O(N^3)\) operand fetches, enabling TPUs to achieve high arithmetic density and energy efficiency.