📚 课程目录第3章 TPU架构详解

3. TPU架构详解

Google TPU（Tensor Processing Unit）自2015年首次部署以来，其架构设计始终围绕一个核心目标：以最高的能效比完成大规模矩阵乘法运算。与GPU采用大规模并行线程调度不同，TPU选择了一条更为极致的路径——将芯片面积和功耗预算几乎全部投入到脉动阵列（Systolic Array）与片上存储的协同设计中。理解TPU架构，不能仅停留在“它是一个矩阵乘法加速器”的层面，而需要深入其数据流组织方式、控制逻辑的简化哲学，以及如何通过“空间计算”替代“时间计算”来消除指令开销。本章将从脉动阵列的物理原理出发，逐步拆解矩阵乘法单元（MXU）、统一缓冲器（UB）与控制流的设计逻辑，并揭示其设计哲学背后的工程权衡。

上图展示了TPU v1的核心数据流架构。与CPU或GPU中“取指-译码-执行”的复杂控制流不同，TPU将计算组织为一条清晰的空间流水线：权重从片外DRAM加载后预置于脉动阵列的每个处理单元（PE）中，激活数据从阵列左侧逐周期流入，部分和在阵列底部逐周期流出。这种设计使得数据在物理空间中的移动替代了传统的时钟周期内多次访存操作，从而大幅降低了功耗和延迟。

3.1 脉动阵列的核心原理：数据复用与空间计算

脉动阵列（Systolic Array）并非Google首创，其概念可追溯至20世纪80年代卡内基梅隆大学孔祥重（H.T. Kung）教授的研究。但TPU将其工程化并推向极致。其核心思想是：让数据在多个处理单元（PE）之间规律地“脉动”流动，每个PE只负责简单的乘加运算，从而最大化数据复用，减少寄存器堆和访存带宽需求。

3.1.1 数据复用模式

在TPU的256×256脉动阵列中，每个PE内部包含一个乘法器和一个加法器。矩阵乘法C = A × B被分解为：权重矩阵B预先加载到PE中（每个PE持有B矩阵的一个元素），激活矩阵A从左侧逐列流入，部分和C从上方逐行流入并在PE内累加后从下方流出。这种设计实现了两种数据复用：

权重静态复用：每个权重元素被固定在PE中，被多个激活数据重复使用256次（对应A矩阵的一行）。
激活流式复用：每个激活元素在流入过程中被同一行的256个PE依次使用，实现广播式复用。

这种复用模式将访存次数从O(N³)降低到O(N²)，对于256×256的矩阵乘法，访存带宽需求降低了约256倍。这是TPU能效比远超GPU的关键所在。

3.1.2 判断方法：何时适合脉动阵列

并非所有计算都适合脉动阵列。判断一个算子是否适合映射到脉动阵列，需要考察以下三个条件：

计算密集度：算术运算与访存操作的比例是否足够高（通常需要大于10:1）。
数据复用性：输入数据是否可以被多个计算单元重复使用（如卷积的滑动窗口、矩阵乘法的行/列复用）。
规则性：计算模式是否规整，是否存在大量分支或数据依赖。

例如，Transformer中的Self-Attention计算（QK^T和AV）完全满足上述条件，因此TPU在NLP任务上表现出色。而动态形状的稀疏计算（如推荐系统中的Embedding查找）则不适合脉动阵列，因为其数据访问模式不规则，会导致大量PE空闲。

3.1.3 案例拆解：卷积运算的脉动映射

以3×3卷积为例，假设输入特征图为14×14×256，输出通道为256。在TPU上，卷积被转化为矩阵乘法：将输入特征图按滑动窗口展开为矩阵（Im2Col），权重作为另一个矩阵。具体映射过程如下：

步骤	操作	数据流方向	PE利用率
1	将256个3×3卷积核权重加载到脉动阵列的256行中	权重静态加载	100%
2	将输入特征图按行展开，逐列送入阵列左侧	激活从左向右流动	100%
3	部分和在阵列底部逐行输出，经累加器处理后写入UB	部分和从上向下流动	100%

在实际部署中，TPU的脉动阵列并非直接处理原始卷积，而是通过Im2Col变换将卷积转换为标准矩阵乘法。这一过程由软件栈（XLA编译器）自动完成，硬件本身不感知卷积的存在。这种软硬件协同设计使得TPU能够以统一的硬件结构处理多种算子。

经验提醒：在评估TPU性能时，不能仅看峰值算力（TOPS），更要关注“有效算力”。脉动阵列的利用率高度依赖矩阵的维度匹配。例如，当矩阵维度小于256时，阵列无法填满，实际吞吐量会显著下降。因此，在实际部署中，通常需要将小矩阵Padding到256的整数倍，或通过Batch合并来增大矩阵维度。

3.2 矩阵乘法单元（MXU）的微架构设计

MXU是TPU的心脏，其微架构设计直接决定了芯片的面积效率和时钟频率。TPU v1的MXU采用8位整数乘法（INT8），累加器为32位浮点（FP32）。这一选择基于以下工程考量：

INT8乘法器面积仅为FP16的1/4，FP32的1/16，使得在相同面积下可以集成更多PE。
32位累加器防止精度溢出：256个INT8乘积累加的最大值为256×255²≈16.6M，需要至少24位整数表示，FP32提供了足够的动态范围。

3.2.1 PE内部结构

每个PE包含三个关键部件：一个8位乘法器、一个32位加法器、以及一组寄存器用于存储权重和部分和。PE之间的连接采用“近邻互联”方式——每个PE只与上下左右四个邻居通信，这种布局使得布线长度最短，时钟频率可以更高。在物理实现上，256×256的PE阵列被划分为4个128×128的子阵列，每个子阵列有独立的时钟门控，以降低空闲功耗。

3.2.2 操作流程：一次矩阵乘法的完整生命周期

// 伪代码描述TPU执行矩阵乘法C = A × B的过程
1. 权重加载阶段：
   for i in 0..255:  // 行
     for j in 0..255:  // 列
       PE[i][j].weight = B[i][j]  // 从UB读取权重，写入PE寄存器
2. 计算阶段：
   for k in 0..255:  // 激活矩阵的列索引
     for i in 0..255:  // 行
       // 激活数据从左侧流入，每个周期流入一列
       activation = A[i][k]  // 从UB读取
       // 激活沿行方向广播到所有PE
       for j in 0..255:
         PE[i][j].partial_sum += activation * PE[i][j].weight
3. 结果回写阶段：
   for i in 0..255:
     for j in 0..255:
       C[i][j] = PE[i][j].partial_sum  // 从阵列底部流出，写入累加器

值得注意的是，上述三个阶段并非完全串行。TPU采用流水线设计，当第k列的激活正在计算时，第k+1列的激活已经可以开始流入，权重加载也可以与计算重叠进行。这种深度流水线使得MXU的利用率接近100%。

3.2.3 常见误区：关于“脉动”的误解

风险说明：许多初学者误以为脉动阵列是“数据像波浪一样在PE间传播”，实际上TPU的脉动阵列是“权重静止、激活流动”的模式。激活数据在每个时钟周期只移动一个PE的位置，而不是广播到所有PE。这种设计避免了高扇出带来的时钟延迟，但代价是计算延迟增加（需要256个周期才能填满流水线）。因此，TPU更适合大批量矩阵乘法，对于小矩阵（如1×1卷积），流水线填充时间会成为瓶颈。

3.3 片上存储体系：统一缓冲器与累加器的协同

TPU的片上存储体系由三部分组成：统一缓冲器（UB，24MB）、累加器（4MB）和指令缓冲器（16KB）。与GPU的复杂多级缓存层次不同，TPU采用极简的存储设计，将SRAM直接暴露给软件管理。

3.3.1 统一缓冲器的角色

UB承担了三个职责：存储激活输入、存储计算中间结果、存储输出数据。这种“统一”设计避免了多级缓存间的数据拷贝，但要求软件（编译器）精确管理数据的生命周期。UB的带宽为每周期256字节（即每周期可读取256个8位激活），恰好匹配脉动阵列的输入需求。

3.3.2 累加器的特殊设计

累加器并非简单的寄存器堆，而是具有“读-改-写”功能的专用存储。当矩阵乘法需要累加多个子块时（例如大矩阵分块计算），累加器可以保留中间结果，避免写回UB再读出的开销。累加器的位宽为32位，但支持INT32和FP32两种模式，由指令控制。

存储单元	容量	带宽	访问模式	功耗占比
统一缓冲器 (UB)	24MB	256 B/cycle	单端口，读写分时	约40%
累加器	4MB	512 B/cycle	双端口，读+写	约15%
指令缓冲器	16KB	—	只读	<1%

从功耗占比可以看出，UB是芯片的功耗大户。因此，TPU的设计哲学是“减少数据移动，而非减少计算量”。编译器通过循环分块（Tiling）和重排，尽可能让数据在UB中驻留更长时间，减少与片外DRAM的交互。

3.3.3 案例：Transformer推理中的存储优化

以BERT-Base推理为例，其核心计算为12层Transformer，每层包含Self-Attention和FFN。在TPU上执行时，编译器将整个计算图划分为多个“编译块”（Compilation Block），每个块内的数据完全驻留在UB中。对于Self-Attention的QK^T计算，编译器将Q矩阵（序列长度×64）和K矩阵（64×序列长度）分块，使得每个块的大小不超过UB容量。通过精细的循环展开，UB的利用率可达85%以上，而GPU的L2缓存命中率通常仅为60%-70%。

  关键洞察：TPU的存储设计体现了“软件定义硬件”的理念。硬件只提供简单的SRAM，但编译器通过静态调度确保数据在正确的时间出现在正确的位置。这种设计牺牲了灵活性（无法运行动态形状的模型），但换取了极高的能效比——TPU v1的能效比（TOPS/W）是同期GPU的5-10倍。

3.4 控制流设计：极简指令集与静态调度

TPU的控制流设计与其计算单元同样激进。它没有采用复杂的乱序执行、分支预测或超标量流水线，而是使用一个极简的VLIW（超长指令字）风格指令集，由片上控制器顺序执行。

3.4.1 指令集架构

TPU的指令集仅包含约10条指令，分为三类：

数据传输指令：Read_Host（从主机读取数据）、Write_Host（写回主机）、Load_Weights（从UB加载权重到MXU）。
计算指令：Matrix_Multiply（执行矩阵乘法）、Activate（执行激活函数，如ReLU）。
控制指令：Jump（条件跳转）、NoOp（空操作）。

每条指令包含多个字段，用于指定源地址、目标地址、数据长度等参数。控制器每周期解码一条指令，并生成相应的控制信号。由于指令执行时间可预测（每个操作都有固定的时钟周期数），编译器可以精确计算整个程序的执行时间，无需硬件动态调度。

3.4.2 判断方法：静态调度 vs 动态调度

选择静态调度还是动态调度，取决于工作负载的可预测性。TPU选择静态调度的原因在于：

深度学习计算图在编译时完全已知，不存在运行时分支。
计算模式规整，每个操作的执行时间固定。
能效优先：动态调度需要额外的硬件（如Scoreboard、Reorder Buffer），消耗功耗和面积。

这种设计的代价是：一旦遇到动态形状（如变长序列），编译器必须生成多个版本的内核，或使用Padding将形状固定。这解释了为什么TPU在静态形状的推理任务中表现优异，但在训练任务中（梯度计算涉及动态形状）需要额外的软件优化。

3.4.3 实战场景：TPU与GPU的控制流对比

在GPU上运行一个简单的矩阵乘法，硬件需要管理数千个线程的调度、同步和内存访问合并。而在TPU上，编译器将整个计算转换为一条Matrix_Multiply指令，控制器只需按顺序执行即可。这种差异在功耗上体现明显：GPU的线程调度器功耗约占芯片总功耗的10%-15%，而TPU的控制器功耗占比不足1%。

3.4.4 常见误区与复盘

风险提示：初学者常误以为TPU的极简控制流意味着“不需要编译器优化”。实际上，TPU的性能高度依赖XLA编译器的调度质量。如果编译器未能有效利用UB空间或未正确安排数据流，TPU的实际性能可能下降50%以上。因此，在使用TPU时，应关注编译日志中的“利用率”指标，而非仅看峰值算力。

总结而言，TPU的架构设计哲学可以概括为“以空间换时间，以编译换硬件”。脉动阵列通过数据复用减少访存，MXU通过简化运算单元提升面积效率，UB通过软件管理减少缓存层次，控制流通过静态调度降低硬件复杂度。这四者相互配合，形成了一个针对深度学习推理高度优化的专用计算系统。理解这一设计哲学，不仅有助于深入掌握TPU的编程优化方法，也为设计其他领域专用加速器（如NPU、DPU）提供了重要的参考范式。