Systolic Arrays vs. GPUs: Architecture and Programming

Recommended resources

The freely available resources this lesson will direct you to.

Introduction

Welcome to our new module, "Emerging Architectures and Future Trends." In the previous module, we explored how to optimize and manage inference workloads at scale, culminating in a discussion of generalist inference servers like Triton. We saw how such systems abstract away the underlying hardware to manage a diverse fleet of models. Now, we will peel back that abstraction layer and examine hardware architectures designed specifically for AI, moving beyond the general-purpose nature of GPUs.

This lesson addresses the learning outcome: Describe how systolic array-based accelerators (e.g., Google TPUs) differ from GPUs in their architecture and programming model. We will investigate the fundamental principles of systolic arrays, see how they are realized in Google's Tensor Processing Unit (TPU), and conduct a detailed comparison with the GPU architecture you're now familiar with. This will illuminate the critical trade-off between general-purpose flexibility and domain-specific acceleration.

The Von Neumann Bottleneck and the Motivation for Specialization

As you know from your computer science background, the traditional von Neumann architecture separates processing units (like a CPU's ALU) from memory. Data must be constantly shuttled back and forth between them. This creates a performance and energy bottleneck, aptly named the "von Neumann bottleneck," as the system spends a significant amount of time and power simply moving data.

For the dense matrix multiplications at the heart of deep learning, this bottleneck is particularly severe. A GPU, with its thousands of cores, exacerbates this problem by creating an immense appetite for data. While GPUs have sophisticated memory hierarchies and high-bandwidth memory to mitigate this, the fundamental limitation remains.

This challenge motivated a return to an older idea: what if we could design an architecture where data flows through the processors, minimizing data movement to and from main memory? This is the core idea behind systolic arrays.

To begin, let's explore the fundamental concepts and motivation behind systolic arrays.

Computer Architecture - Lecture 27: Systolic Arrays (ETH Zürich, Fall 2020)

This first video segment, from a computer architecture lecture by Professor Onur Mutlu, introduces the concept of systolic arrays. It explains the core philosophy: bringing data from memory once and performing as much computation on it as possible before writing it back.

Please watch from 04:36 to 12:12. Focus on: The motivation for designing specialized accelerators. The basic principle of replacing a single processing element with a regular array of PEs. The analogy of the systolic array to the heart pumping blood through the body, which gives the architecture its name.

As the video explains, the key principle is to orchestrate a rhythmic flow of data through a grid of simple processing elements (PEs), maximizing the computation performed for each byte read from memory. This design is simple, regular, and highly concurrent—ideal for hardware implementation.

Systolic Arrays vs. Pipelining

You might notice a similarity to instruction pipelining in a CPU. However, there are crucial differences:

PE Function: In a pipeline, stages perform a sub-operation of a single instruction (e.g., fetch, decode, execute). In a systolic array, each PE performs a complete operation (like a multiply-accumulate) on a data element.
Dimensionality & Connectivity: CPU pipelines are typically linear. Systolic arrays can be 1D, 2D, or even 3D, with more complex, multi-directional data flows.

Professor Mutlu briefly clarifies this distinction in the next clip.

Computer Architecture - Lecture 27: Systolic Arrays (ETH Zürich, Fall 2020)

This short clip explicitly contrasts systolic arrays with traditional pipelining.

Watch from 12:12 to 13:20. This will solidify your understanding of how this paradigm differs from the pipelining you're familiar with in general-purpose processors.

A Concrete Example: 1D Convolution

Let's make this concrete. Convolution is a core operation in many neural networks. The following video segment demonstrates how a simple 1D systolic array, as originally proposed by H.T. Kung, can implement this operation efficiently.

Computer Architecture - Lecture 27: Systolic Arrays (ETH Zürich, Fall 2020)

This segment walks through the design of a 1D systolic array for convolution. It details the data flow and the simple multiply-accumulate (MAC) function of each processing element.

Watch from 13:20 to 28:40. Don't worry about the fine details of the neural network architectures mentioned; focus on the following: How the input data (x), weights (w), and partial sums (y) flow in opposite directions through the array. The simple, repeated function performed by each processing element (PE). The importance of careful data orchestration (timing) to make the system work.

This example makes the benefits clear: data is reused extensively, and the design is highly regular and parallel. The primary drawback, as noted in the video, is its specialization. It's incredibly efficient for convolution or matrix multiplication, but useless for general-purpose tasks.

Google's TPU: A Modern Systolic Array

The principles of systolic arrays, though developed decades ago, found their killer application in deep learning. Google's Tensor Processing Unit (TPU) is the most prominent modern example of an Application-Specific Integrated Circuit (ASIC) built around a large-scale systolic array.

The heart of the TPU is its Matrix Multiplier Unit (MXU), which is a large 2D systolic array of multiply-accumulate units.

GPU vs. TPU Matrix Multiplication Architectures — This image contrasts the data flow for matrix multiplication on a GPU versus a TPU. The GPU (left) shows a general parallel computation model where ALUs fetch data from registers or memory. The TPU (right) illustrates the systolic array model, where data streams through a grid of PEs, with the output of one PE becoming the input for the next, drastically reducing memory access.

To understand the TPU's architecture, we'll turn to a detailed blog post from Google that accompanied the seminal paper on their first-generation TPU.

An in-depth look at Google's first Tensor Processing Unit (TPU)

This article provides an in-depth look at the first-generation TPU. It explains the motivation, the choice of data types, and the core architectural components, including the systolic array.

Please read the following sections: 'Prediction with neural networks': This sets the stage by emphasizing the dominance of matrix multiplication. 'RISC, CISC and the TPU instruction set': This explains the TPU's high-level, CISC-style instructions that map directly to neural network operations. 'The heart of the TPU: A systolic array': This is the most critical section. It explains how the MXU works and directly contrasts its matrix-processing approach with the scalar/vector processing of CPUs/GPUs.

As the reading highlights, the TPU is not just a raw systolic array; it's a complete co-processor with its own instruction set and on-chip memory, designed to be controlled by a host CPU.

GPU vs. TPU: A Comparative Analysis

Now we can directly compare GPUs and TPUs across their architecture and programming models.

Attribute	GPU (e.g., NVIDIA H100)	TPU (e.g., Google TPU v4/v5)
Core Compute Unit	Thousands of flexible CUDA Cores (ALUs) organized into Streaming Multiprocessors (SMs). Designed for general-purpose parallel tasks (SIMT).	A single, massive Systolic Array (MXU) of simple, fixed-function MAC units. Designed for one task: dense matrix multiplication.
Design Philosophy	General-Purpose Parallelism. Includes complex features like caches, schedulers, and branch prediction to handle diverse workloads efficiently.	Domain-Specific Minimalism. Strips away nearly all general-purpose features. No caches, no branch prediction. Every transistor is dedicated to matrix math.
Programming Model	Highly Flexible & Programmable. Programmed via CUDA, OpenCL, HIP. Supports a wide ecosystem of frameworks like PyTorch, TensorFlow, JAX, etc. The software defines the computation.	Optimized & Restricted. Programmed primarily via TensorFlow or JAX. Code is compiled by the XLA (Accelerated Linear Algebra) compiler into low-level TPU instructions. The hardware defines the computation.
Data Flow	Load/Store Architecture. Cores fetch data from registers/shared memory, perform an operation, and write the result back. Relies heavily on memory bandwidth.	Dataflow Architecture. Data is streamed through the systolic array. Results from one PE are passed directly to the next, bypassing memory. Relies on orchestrating the data flow.
Flexibility	High. Can be reprogrammed for new algorithms, graphics, scientific computing, etc. Adapts to changes in the field via software.	Low. Extremely fast and power-efficient for its target workload (matrix multiplication). If the fundamental math of AI were to change, the hardware would become obsolete.

The following short video provides a very clear summary of this fundamental architectural difference and the resulting trade-off.

TPUs Are BETTER. But Why No One Uses Them?

This video concisely explains how the TPU's systolic array solves the von Neumann bottleneck that GPUs still face, but at the cost of flexibility.

Please watch from 00:31 to 03:04. This will reinforce the comparison between the GPU's von Neumann design and the TPU's systolic array, and it clearly frames the ASIC vs. general-purpose processor trade-off.

To round out your understanding, the following reading offers a comprehensive technical comparison, covering aspects like memory, interconnects, and the programming toolchains.

TPU vs GPU: A Comprehensive Technical Comparison

This article from Wevolver provides a great summary of the architectural and practical differences between GPUs and TPUs.

Please read the sections 'Architectural Differences' and 'Practical Implementation Considerations'. This will solidify the comparison points regarding compute design, memory, and programming models.

Conclusion

In this lesson, we explored an alternative to the general-purpose GPU architecture for accelerating AI workloads. We saw how systolic arrays, an architectural pattern that prioritizes data flow over random access, provide a powerful solution to the memory bottleneck inherent in deep learning's massive matrix multiplications.

Key Takeaways:

Systolic Arrays are grids of simple processing elements that reduce memory traffic by having data flow through the processors, maximizing computation on each piece of data.
Google's TPU is a prime example of an ASIC that implements a large-scale systolic array (the MXU) specifically for accelerating neural network operations.
GPU Architecture is based on a general-purpose parallel (SIMT) model. It is highly programmable and flexible, but its performance is fundamentally tied to memory bandwidth (the von Neumann bottleneck).
TPU Architecture is based on a specialized dataflow model. It is extremely power-efficient and high-performance for its target task (matrix multiplication) but lacks the flexibility of a GPU.
The Programming Model reflects this difference: GPUs offer broad, flexible programming via frameworks like CUDA, while TPUs rely on a more constrained model where frameworks like TensorFlow/JAX use the XLA compiler to target the specific hardware capabilities.

This comparison reveals a core theme in computer architecture: the trade-off between generality and specialization. The GPU is a triumph of flexible, programmable parallel computing, while the TPU is a testament to the raw efficiency gained by designing hardware for a specific domain.

Preview of the Next Lesson:

The TPU represents one form of specialization. But what if you want to create specialized systems without the massive cost and risk of designing a single monolithic ASIC? In our next lesson, we will explore the chiplet-based approach to processor design. This paradigm allows for building complex, heterogeneous systems by combining smaller, specialized dies, offering a middle ground between general-purpose chips and fully custom ASICs.