How Google TPUs Use Systolic Arrays to Accelerate AI

This article explores how Google's Tensor Processing Units (TPUs) leverage a specialized hardware architecture known as a systolic array to dramatically accelerate artificial intelligence and machine learning workloads. By addressing the primary bottleneck of modern computing—memory access—systolic arrays allow TPUs to perform massive matrix multiplications with exceptional speed and energy efficiency. The following sections break down the mechanics of systolic arrays, why they are uniquely suited for deep learning, and how Google implements them to power advanced neural networks.

The AI Bottleneck: Matrix Multiplication and Memory Access

At the core of modern artificial intelligence, particularly deep learning, lies a fundamental mathematical operation: matrix multiplication. Neural networks require millions or even billions of these calculations to process data, recognize patterns, and generate predictions.

Traditional processors, such as CPUs and GPUs, handle these operations using a von Neumann architecture. For every single calculation, the processor must fetch data from memory (registers), route it to an Arithmetic Logic Unit (ALU) to perform the math, and write the result back to memory. When multiplying massive matrices, moving data back and forth between memory and the compute units consumes far more time and electricity than the actual mathematical calculations. This limitation is known as the von Neumann bottleneck.

What is a Systolic Array?

A systolic array is a hardware architecture designed specifically to eliminate the memory access bottleneck. Named by analogy to the human circulatory system—where the heart (systole) pumps blood rhythmically through a network of vessels—a systolic array pumps data rhythmically through a 2D grid of interconnected processing elements.

In a systolic array, individual compute units, called Multiply-Accumulate (MAC) cells, are wired directly to their immediate neighbors. Instead of reading from and writing to a central register after every single operation, a MAC cell receives data from an upstream neighbor, performs a calculation, and passes the result directly to a downstream neighbor on the next clock cycle. Memory is only accessed at the very beginning to load the initial data into the edges of the array, and at the very end to collect the final output.

How Google TPUs Implement Systolic Arrays

Google engineered the Tensor Processing Unit around the Matrix Multiply Unit (MXU), which is fundamentally a massive systolic array. Depending on the TPU generation, a single MXU contains tens of thousands of MAC cells arranged in a dense grid (such as 128x128 or 256x256).

The TPU executes neural network calculations through a specific dataflow mechanism within this grid:

  1. Weight Stationary Setup: Before data processing begins, the parameters of the neural network (known as weights) are pre-loaded into the MAC cells of the systolic array. They remain stationary inside the cells throughout the operation.
  2. Pumping the Data: Input data (such as image pixels or text tokens, known as activations) is fed into the array from the left edge.
  3. Rhythmic Execution: With each tick of the processor's clock, the input data moves one step to the right. Simultaneously, each MAC cell multiplies the incoming input by its stationary weight and adds that result to the partial sum flowing down from the cell above it.
  4. Collecting the Output: The updated partial sums flow downward to the next row of cells. By the time the data reaches the bottom edge of the array, the massive matrix multiplication is complete, and the final results are written back to memory.

To manage the continuous, staggered wave of data flowing through the array, the input data is fed diagonally. This ensures that the right inputs meet the right partial sums at the exact right microsecond inside the grid.

Key Advantages for AI Workloads

By replacing general-purpose processing pathways with systolic arrays, Google TPUs achieve three critical advantages for AI acceleration:

  • Drastic Reduction in Power Consumption: Accessing memory requires significantly more energy than executing arithmetic in silicon. By passing data directly between adjacent MAC cells, TPUs eliminate millions of register reads and writes, drastically lowering the energy cost per calculation.
  • Increased Compute Density: Because systolic arrays do not require complex routing logic, instruction decoders, or large register files for every single ALU, Google can pack vastly more MAC cells onto a single piece of silicon. More cells equate to higher parallel processing power.
  • High Throughput: Once the systolic array pipeline is full, it produces a massive wave of completed calculations on every single clock cycle. This sustained throughput allows TPUs to train large language models and run inference workloads in a fraction of the time required by traditional architectures.