---
title: Tensor Processing Unit (TPU) Overview
url: https://www.emergentmind.com/topics/tensor-processing-unit-tpu
type: topic
---

# Tensor Processing Unit (TPU) Overview

A Tensor Processing Unit (TPU) is a domain-specific hardware accelerator architected for high-throughput linear algebra, particularly tailored for the multiply-accumulate workloads central to deep neural networks (DNNs). Originally developed for large-scale production environments such as Google data centers, as well as for edge inference, TPUs are engineered around massively parallel 2D systolic arrays of MAC (multiply-accumulate) units, on-chip memory systems designed for deterministic performance, and a suite of compiler and runtime systems that jointly optimize compute intensity, bandwidth, and energy efficiency [1704.04760][2309.08918][2407.08700].

## 1. Architectural Fundamentals

At the core of every TPU generation is a systolic array microarchitecture. Each array consists of a grid of processing elements (PEs), with each PE containing one or more MAC units and several small, single-purpose local registers. Inputs (activations), weights, and partial sums are streamed into the array along orthogonal directions, with the dataflow pattern optimized to maximize local reuse and minimize off-chip memory bandwidth.

Key architectural features include:
- **Systolic Array Composition:** Early production TPUs employed a 256×256 8-bit MAC array (65,536 MACs, 92 TOPS at 700 MHz), later evolving to multi-core chips with multiple 128×128 arrays and bfloat16 or int8 support [1704.04760][2309.08918][2112.09017].
- **Local Scratchpad Memory:** A large, software-managed on-chip SRAM (e.g., 28 MiB) provides high-bandwidth data staging, separate from host DRAM (e.g., 8 GiB DDR3/4 per chip). Exclusive reliance on scratchpads (not caches) guarantees deterministic access latencies and avoids cache coherence issues [1704.04760].
- **Instruction Delivery and Determinism:** A simple, in-order instruction pipeline issues fixed-format commands (e.g., matrix-multiply, activation) without out-of-order scheduling, SMT, or speculation. This deterministic execution is essential for tail-latency guarantees in end-to-end inference pipelines [1704.04760].
- **Off-Chip Interconnect:** In multi-chip deployments (“TPU pods”), a fast 2D toroidal or 3D twisted torus mesh links chips for distributed computation. Google’s Palomar optical circuit switches (OCSes) in TPU v4 provide dynamic topology, low-latency reconfiguration, and energy proportional scaling [2304.01433].

## 2. Dataflow and Computational Patterns

The throughput and efficiency of a TPU are shaped by its concrete dataflow: the spatio-temporal schedule of inputs, weights, and partial sums across the array. Three classical dataflow paradigms dominate:
- **Input-Stationary (IS):** Hold input feature map elements stationary within each PE, stream weights horizontally, and propagate partial sums downstream. Suited for layers or kernels with high input reuse, such as depthwise convolutions [2407.08700].
- **Weight-Stationary (WS):** Hold weights fixed per PE, stream input activations vertically. Favored in early convolutional layers with high weight reuse [2407.08700].
- **Output-Stationary (OS):** Accumulate each output pixel’s partial sum in-place, streaming both activations and weights as needed. Optimal for high-intensity deep layers [2407.08700][2503.05951].

Static (single-mode) dataflow selection underutilizes potential in workloads with heterogeneous layers, incurring up to 40% suboptimality [2407.08700]. The Flex-TPU architecture introduced a runtime-reconfigurable datapath, using an extra “stationary” register per PE and per-layer selection logic—achieving up to 2.75× speedup over baseline designs with ∼10% area/power overhead [2407.08700].

## 3. Memory Hierarchy and System Integration

TPUs decouple compute and memory bandwidth via:
- **High-Bandwidth On-Chip SRAM:** Software-managed scratchpads enable spatial tiling and deterministic data reuse, nullifying architectural jitter from hardware-managed cache hierarchies. For example, 28 MiB Unified Buffer on the original TPU enables direct compiler-controlled staging of activations, weights, and intermediate tensors [1704.04760].
- **External/Off-Chip Memory:** Bulk model weights reside in high-capacity DRAM (e.g., DDR3 or HBM). Newer generations (TPU v3/v4) integrate HBM2 (e.g., 900–1200 GB/s/chip), essential for bandwidth-bound operations [2309.08918][2304.01433].
- **Sparse Data and Embeddings:** In TPU v4, "SparseCore" tiles (∼5% die area/power) directly handle sparse embedding lookups, providing 5–7× acceleration vs. host memory placement for large DLRM-like models [2304.01433].

## 4. Programming Model and Compiler Toolchains

The TPU software stack is built around high-level tensor graph abstractions, aggressive graph optimization, and device-specific codegen:
- **Framework Integration:** TensorFlow and JAX express computation as dataflow graphs, which are then operated on by the XLA (Accelerated Linear Algebra) compiler [2309.08918].
- **Graph and Kernel Optimization:** The XLA pipeline includes HLO (high-level optimizer) transformations (constant folding, op fusion), layout assignment and tiling, lowering to LLO (low-level ops), and final device codegen. Operator fusion (e.g., Conv-ReLU-BatchNorm) is key to maximizing local reuse and reducing on-chip/off-chip transfers [2309.08918][2210.15016].
- **TPU-MLIR Compiler:** The domain-specific TPU-MLIR pass pipeline leverages MLIR dialects (TOP for operator semantics, TPU for hardware kernels), partitioning graphs into SRAM-fitting groups, assigning memory, and emitting binary code and metadata. Quantized inference (INT8) is supported via calibration and quantization passes, matching hand-optimized kernel performance within 5% [2210.15016].
- **Autotuning and Performance Modeling:** Code generation and autotuning benefit from learned performance models (GNN-based), which guide tile size and fusion heuristic choices and discover optimal configurations with minimal hardware search [2008.01040].

## 5. Scaling Performance, Efficiency, and Applications

TPUs deliver several orders of magnitude more compute per watt and per dollar than general-purpose hardware:
- **Performance Roofline:** Sustained throughput in compute–memory bandwidth balance is modeled as $P_\text{sustained} = \min(P_\text{peak},\, \mathrm{AI}\cdot B_\text{mem})$, where arithmetic intensity (AI) is kernel-flop/byte ratio [2309.08918][1704.04760]. Memory-bound models (MLPs, LSTMs) are sensitive to DRAM bandwidth, while CNNs approach compute-bound ceilings.
- **Energy Efficiency:** The original TPU delivered ∼2.3 TOPS/W, 30–80× the energy efficiency of contemporary GPUs or CPUs for real-world inference [1704.04760]. TPU v4 chips achieve 2.1× the TPU v3 throughput, 2.7× performance/W, and are 3–20× more energy/carbon efficient than typical enterprise or on-prem deployments [2304.01433].
- **Distributed Supercomputing:** Full TPU pods (e.g., 2048 v3 cores) execute classic 2D-blocked distributed GEMM, QR, and iterative matrix computations with nearly ideal weak scaling. Large problems (e.g., $N=2^{20}$) reach >20 PFLOPS fp32 throughput. The physical 2D-torus interconnect aligns with SUMMA/CAQR communication patterns for scalable matrix-algebra workloads [2112.09017].
- **Scientific Workloads:** TPUs support graph-based CFD solvers (Navier-Stokes, turbulent flows), with performance comparable to GPU/CPU clusters and good scaling up to full pod deployments [2108.11076].
- **Automation and Design Customization:** LLM-driven hardware generators (TPU-Gen) automate the RTL generation of both exact and approximate systolic TPUs, covering a spectrum of array sizes, bitwidths, and PPA-optimized designs. RTL pass rates improve from 0% (LLM alone) to >95% (LLM+RAG), with up to 92% area and 96% power reductions over manual design [2503.05951].

## 6. Advancements, Trade-offs, and Future Directions

Ongoing research explores precision-enhanced TPUs (using fractional Residue Number System arithmetics) for arbitrary wide fixed-point/floating-point computation, maintaining fixed per-MAC cost and throughput at scale [1706.03251]. Approximate arithmetic units (such as in TPU-Gen) provide further energy and area efficiency, with minimal accuracy impact tuned via per-layer profiles [2503.05951]. Hardware-software co-design continues to advance with per-layer dynamic dataflow switching (Flex-TPU), specialization for sparse workloads, and expanding compiler abstractions for verification, quantization, and device onboarding [2407.08700][2210.15016].

Key trade-offs persist:
- Static vs. dynamic dataflow selection
- Precision/energy/area trade-offs in MAC units
- Determinism vs. general-purpose flexibility
- Manual vs. automated hardware–software design flows

A plausible implication is that future TPUs will further broaden their application domains—beyond deep learning—to distributed scientific computing, custom low-power edge inference, and self-evolving automated hardware design platforms. Persistent emphasis on maximizing arithmetic intensity, minimizing data movement, and software-managed predictability will define subsequent generations of tensor accelerator architectures.

Source: https://www.emergentmind.com/topics/tensor-processing-unit-tpu