v2_npu_tutorial

第5章：脉动阵列原理与设计

本章深入探讨脉动阵列（Systolic Array）的基本原理和设计思想，以Google TPU为例分析其架构演进。我们将从数据流动模式、计算同步机制入手，理解Weight-stationary设计的优势，并通过TPU四代产品的架构对比，掌握脉动阵列在实际应用中的优化策略。通过本章学习，读者将能够分析脉动阵列的利用率，设计高效的矩阵运算映射方案，并理解不同控制流架构的权衡。

5.1 脉动阵列基本概念

5.1.1 数据流动与计算同步

脉动阵列是一种规则的处理单元（PE）阵列结构，数据像心脏泵血一样有节奏地在PE间流动。每个PE在每个时钟周期执行相同类型的操作，通过精心设计的数据流动模式实现高效的并行计算。这种架构最早由H.T. Kung和Charles Leiserson在1978年提出，其设计理念是让数据在处理单元间有规律地”脉动”流过，每个数据项在流经系统时被多次使用，从而最大化数据重用率。

脉动阵列的核心特征：

脉动阵列的数学基础建立在矩阵运算的规律性上。考虑矩阵乘法 $C = A \times B$，其中 $C_{ij} = \sum_{k} A_{ik} \times B_{kj}$。这个计算具有高度的规律性和并行性：每个输出元素的计算相互独立，且都需要执行相同的运算序列（乘累加）。脉动阵列正是利用了这种规律性，将计算在空间（多个PE）和时间（多个周期）上展开。

深入理解脉动阵列的理论基础需要从计算的数据依赖图（Data Dependence Graph, DDG）入手。对于矩阵乘法，DDG具有高度规则的结构：每个 $C_{ij}$ 的计算依赖于第i行的所有A元素和第j列的所有B元素。这种规则性使得可以通过系统化的方法将算法映射到硬件架构上。

空间-时间映射理论： 脉动阵列的设计本质上是将算法的迭代空间（iteration space）映射到处理器的空间-时间域。对于三重循环的矩阵乘法：

for i = 0 to M-1:
    for j = 0 to N-1:
        for k = 0 to K-1:
            C[i][j] += A[i][k] * B[k][j]

这个迭代空间是一个三维空间 $(i,j,k)$，脉动阵列将其映射为：

这种映射的数学表达为： \(\begin{bmatrix} PE_x \\ PE_y \\ t \end{bmatrix} = \begin{bmatrix} 1 & 0 & 0 \\ 0 & 1 & 0 \\ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix} i \\ j \\ k \end{bmatrix}\)

这是最简单的恒等映射，实际设计中可能使用更复杂的仿射变换来优化数据重用和减少通信。

从数据流的角度看，脉动阵列实现了三种关键的数据移动模式：

  1. 单播到多播的转换：一个输入数据通过PE间的传递被多个计算单元使用。这种机制将原本的一对多通信转换为多个一对一通信，极大地简化了互连网络的设计。具体而言，如果使用总线实现一对多通信，需要驱动N个负载，这会导致延迟和功耗都随N线性增长；而脉动阵列的邻近通信只需要驱动单个负载
  2. 流水线化的累加：部分和在PE间传递并逐步累加，避免了全局归约的开销。传统的树状归约需要log(N)级的加法器和复杂的控制逻辑，而脉动阵列通过空间展开实现了线性时间的累加。更重要的是，流水线化的累加可以与新的乘法运算重叠，实现计算的连续性
  3. 时空映射：将算法的循环嵌套映射到硬件的空间维度和时间维度。这种映射将串行算法的循环迭代转换为并行硬件的空间分布和时间步进，实现了算法到硬件的自然映射

数据复用的定量分析： 脉动阵列的一个核心优势是数据复用。我们可以定量分析不同级别的数据复用：

  1. 时间复用（Temporal Reuse）：同一个PE在不同时间使用相同数据
    • 在Weight-stationary设计中，权重在PE中驻留整个计算过程
    • 复用因子：等于处理的批次数B
  2. 空间复用（Spatial Reuse）：不同PE使用相同数据
    • 输入激活在水平方向传播，被每行的所有PE使用
    • 复用因子：等于阵列的列数N
  3. 总复用因子：$R_{total} = R_{temporal} \times R_{spatial} = B \times N$

这意味着每个数据项从内存读取一次后，可以被使用 $B \times N$ 次，极大地提高了运算强度（Operational Intensity）： \(OI = \frac{计算量}{数据移动量} = \frac{2MNK}{M \times K + K \times N + M \times N} \approx \frac{2K}{1 + \frac{1}{M} + \frac{1}{N}}\)

当M和N足够大时，$OI \approx 2K$，这解释了为什么脉动阵列特别适合深度（K维度大）的矩阵乘法。

基本的一维脉动阵列可以表示为：

时刻 t=0:  [a0] → PE0 → PE1 → PE2 → PE3
时刻 t=1:  [a1] → PE0 → PE1 → PE2 → PE3
                    ↓     ↓     ↓     ↓
                   [b0]  [b1]  [b2]  [b3]

对于矩阵乘法 $C = A \times B$，二维脉动阵列的数据流动模式：

        b00  b01  b02  b03
         ↓    ↓    ↓    ↓
a00 → [PE00][PE01][PE02][PE03]
a10 → [PE10][PE11][PE12][PE13]  
a20 → [PE20][PE21][PE22][PE23]
a30 → [PE30][PE31][PE32][PE33]

每个PE执行乘累加操作：$c_{ij} = c_{ij} + a_{ik} \times b_{kj}$

脉动执行的时序分析对理解其效率至关重要。对于一个 $m \times n$ 的脉动阵列计算 $M \times K$ 和 $K \times N$ 的矩阵乘法：

启动延迟（Startup Latency）：第一个结果出现前需要等待的周期数

稳态吞吐量（Steady-State Throughput）：

总执行时间： \(T_{total} = L_{startup} + \lceil \frac{M}{m} \rceil \times \lceil \frac{K}{1} \rceil \times \lceil \frac{N}{n} \rceil\)

这个公式揭示了脉动阵列设计的一个关键权衡：更大的阵列有更高的并行度，但也有更长的启动延迟。对于小矩阵，启动延迟可能占主导地位，导致利用率下降。

效率分析的临界点： 当矩阵维度满足 $M \times N > \alpha \times (m + n)^2$ 时（其中$\alpha$为经验系数，通常为2-4），脉动阵列才能达到较高的效率。这解释了为什么脉动阵列特别适合大规模矩阵运算，而对于小矩阵可能不如简单的向量处理器高效。

在实际系统中，常用的优化技术包括：

5.1.2 Weight-stationary设计选择

脉动阵列有三种主要的数据流模式：

  1. Weight-stationary (WS)：权重固定在PE中，输入和输出流动
  2. Output-stationary (OS)：输出固定在PE中，权重和输入流动
  3. Row-stationary (RS)：每行PE保持特定数据，其他数据流动

TPU选择Weight-stationary设计的关键原因深植于深度学习工作负载的特性。在推理场景中，模型权重是固定的，而输入数据不断变化。Weight-stationary设计完美契合了这一特点。

从能量效率的角度看，Weight-stationary的优势源于数据移动能耗的基本物理定律。根据电容充放电理论，移动n位数据的能耗为 $E = \frac{1}{2} C \times V^2 \times f \times \alpha$，其中$C$为线路电容，$V$为电压，$f$为频率，$\alpha$为翻转因子。在现代工艺下，电容随着距离增加而线性增长：$C = C_0 + c_w \times L$，其中$C_0$是驱动器和接收器的固定电容，$c_w$是单位长度的线电容（典型值0.2fF/μm），$L$是连线长度。

能耗层次分析： 现代芯片中，数据移动的能耗呈现明显的层次结构：

这个能耗金字塔清楚地表明，将数据保持在尽可能接近计算单元的地方至关重要。Weight-stationary通过将权重锁定在PE的本地寄存器中（最低能耗层次），实现了能效最优化。

优势分析：

让我们通过一个具体例子深入理解Weight-stationary的优势。考虑一个全连接层，批大小B=128，输入维度I=1024，输出维度O=1024：

这个例子代表了典型的Transformer模型中的FFN层。在这种规模下，权重矩阵占据2MB存储（FP16），而每个批次的激活数据也为256KB。如果采用传统的计算方式，每次计算都需要从内存加载权重，这将成为性能瓶颈。

对于典型的神经网络层计算，设批大小为 $B$，输入特征为 $I$，输出特征为 $O$：

在Weight-stationary模式下：

总数据移动量：$D_{WS} = W + O \times A_{in} + A_{out}$

相比之下，Output-stationary的数据移动量： $D_{OS} = B \times W + A_{in} + A_{out}$

当批大小 $B$ 较大时，$D_{WS} < D_{OS}$，Weight-stationary更有优势。

实际硬件实现中，Weight-stationary还带来了其他系统级优势：

  1. 预取优化：权重的访问模式完全可预测，可以实现完美的预取，隐藏内存延迟。现代DDR4/5的延迟高达60-100ns，通过预取可以完全隐藏这些延迟
  2. 功耗门控：权重寄存器在加载后可以关闭写使能，降低动态功耗。在28nm工艺下，关闭写使能可以节省约15%的寄存器功耗
  3. 容错设计：权重存储可以使用ECC保护，因为写入频率低，校验开销可接受。SECDED ECC只增加12.5%的存储开销，但能显著提高可靠性
  4. 压缩友好：权重可以压缩存储，在加载到PE时解压，因为解压只需要做一次。典型的权重压缩率可达2-4倍，节省大量存储带宽

然而，Weight-stationary也有其局限性：

这些局限性解释了为什么一些NPU（如Nvidia的Tensor Core）选择了更灵活的Output-stationary或Row-stationary设计。不同的设计选择反映了对不同应用场景的优化重点。

5.1.3 阵列维度与利用率分析

脉动阵列的维度选择直接影响硬件利用率和性能。设阵列维度为 $M \times N$，分析不同场景下的利用率是一个多目标优化问题，需要在面积、功耗、性能之间寻找平衡。

矩阵乘法映射效率：

对于矩阵乘法 $C_{P \times Q} = A_{P \times K} \times B_{K \times Q}$，映射到 $M \times N$ 的脉动阵列涉及三个维度的分块：

  1. 完美匹配：当 $P = kM, Q = lN$（$k,l$ 为整数）时，利用率接近100%
    • 这种情况下，矩阵恰好可以分成整数个tiles
    • 每个tile完全利用脉动阵列的所有PE
    • 实际应用中很少出现完美匹配
  2. 维度不匹配：实际利用率 $U = \frac{\lceil P/M \rceil \times \lceil Q/N \rceil}{P \times Q} \times \frac{P \times Q}{M \times N}$
    • 第一项表示tiling的效率损失
    • 第二项表示最后一个tile的利用率
    • 综合利用率通常在60-85%之间
  3. 边界效应：需要padding处理
    • Padding开销：$(M - P \mod M) \times (N - Q \mod N)$
    • 有效计算比例：$\eta = \frac{P \times Q}{M \times N \times \lceil P/M \rceil \times \lceil Q/N \rceil}$
    • Padding不仅增加计算量，还增加内存访问

维度选择的设计空间探索：

选择最优的阵列维度需要考虑多个因素，这是一个复杂的多目标优化问题。

  1. 目标工作负载分析：
    • 统计常见层的维度分布：需要收集目标应用的所有层参数，建立维度分布直方图
    • 识别性能关键层：通常遵循帕累托原则，20%的层占据80%的计算量
    • 考虑未来模型趋势：Transformer模型倾向于使用更大的隐藏维度（如4096、8192）

    以BERT-Large为例，其关键维度包括：

    • 注意力矩阵乘法：[batch×seq, 1024] × [1024, 1024]
    • FFN第一层：[batch×seq, 1024] × [1024, 4096]
    • FFN第二层：[batch×seq, 4096] × [4096, 1024]

    对于batch=32, seq=512的典型配置，第一个维度达到16384，这就需要考虑如何高效地将其映射到物理阵列上。

  2. 硬件约束：
    • 芯片面积：$Area \propto M \times N$（PE数量）。在现代工艺下，每个INT8 MAC单元约需100-200平方微米，256×256阵列需要约13平方毫米
    • 布线复杂度：$Routing \propto M + N$（边界PE的连接）。随着阵列增大，全局布线变得困难，可能需要多级布线网络
    • 时钟频率：更大的阵列可能需要降低频率以满足时序。经验公式：$f_{max} \approx \frac{c}{\sqrt{M \times N}}$，其中$c$为工艺相关常数
  3. 功耗考虑：
    • 静态功耗：$P_{static} \propto M \times N$
    • 动态功耗：$P_{dynamic} = \alpha \times C \times V^2 \times f \times U$
    • 其中U是利用率，低利用率意味着功耗效率低

批处理优化：

对于批量矩阵乘法，可以通过合并批次提高利用率：

稀疏处理策略：

对于2:4结构化稀疏：

5.2 TPU v1-v4i架构演进

5.2.1 TPU v1：推理专用架构

TPU v1是Google的第一代张量处理器，专为数据中心推理工作负载设计。这个项目始于2013年，当时Google意识到如果所有用户每天使用3分钟的语音搜索，需要的计算资源将是当时数据中心容量的两倍。这促使Google开发专用硬件来加速神经网络推理。

项目背景与动机：

2013年，Google面临着一个关键挑战：深度学习模型在语音识别、图像分类等任务上取得了突破性进展，但部署成本极高。具体数据显示：

这促使Google做出了一个大胆的决定：设计专用的推理芯片。项目代号”Manhattan”，目标是在15个月内完成从架构设计到批量部署。

设计理念与目标：

TPU v1的设计哲学是”做一件事，并把它做到极致”。与GPU试图兼顾图形渲染和通用计算不同，TPU v1专注于神经网络推理这一个任务。这种专注带来了几个关键优势：

架构决策的深层考虑：

  1. 为什么选择INT8而非FP16/FP32：
    • Google的大规模实验表明，对于已训练模型，INT8量化的精度损失可控（<1%）
    • INT8乘法器面积仅为FP32的1/6，功耗为1/10
    • 内存带宽需求减少4倍，这在内存墙问题日益严重的今天尤为重要
  2. 为什么选择256×256的阵列规模：
    • 太小（如64×64）：控制开销比例过高，利用率受限
    • 太大（如512×512）：面积和功耗呈平方增长，时序收敛困难
    • 256×256是sweet spot：在28nm工艺下，可以达到700MHz主频，同时保持合理的面积（~80mm²）
  3. 为什么采用脉动阵列而非其他架构：
    • SIMD向量处理器：带宽需求过高，每个周期需要读取所有操作数
    • 多核MIMD：同步开销大，编程模型复杂
    • 脉动阵列：数据重用率高，编程模型简单，硬件效率最高

核心规格：

MXU (Matrix Multiply Unit) 设计：

      统一缓冲区 (24MB)
           ↓
    ┌──────────────┐
    │  权重FIFO    │ ← 权重加载
    └──────────────┘
           ↓
    ┌──────────────────────┐
    │                      │
    │   256×256 脉动阵列   │ ← 激活输入
    │                      │
    └──────────────────────┘
           ↓
    ┌──────────────┐
    │  累加器(4MB) │
    └──────────────┘
           ↓
       激活单元

关键设计决策：

量化策略深入分析：

TPU v1采用的INT8量化包含几个关键技术，这些技术后来成为业界标准：

  1. 非对称量化：$q = \text{round}(r/s) + z$
    • r：实数值
    • s：缩放因子（scale）
    • z：零点（zero point）

    非对称量化相比对称量化的优势在于可以更好地处理激活值的非对称分布。例如，ReLU后的激活值都是非负的，使用非对称量化可以充分利用INT8的表示范围。

  2. 逐通道量化： 每个输出通道使用独立的缩放因子，这是因为不同通道的权重分布差异很大。实验表明：
    • 逐层量化：精度损失2-5%
    • 逐通道量化：精度损失<0.5%
    • 开销：需要存储额外的缩放因子（每通道4字节）
  3. 激活函数融合： ReLU6（$y = \min(\max(x, 0), 6)$）的使用有多重考虑：
    • 自然限制输出范围到[0, 6]，简化量化
    • 防止激活值爆炸，提高数值稳定性
    • 硬件实现简单，只需要比较器和多路选择器
  4. 量化感知的编译优化： 编译器在量化时会进行以下优化：
    • 常量折叠：预计算可以离线完成的量化操作
    • 重排序：将量化/反量化操作移动到计算密集区域的边界
    • 融合：将量化与其他操作（如激活函数）融合，减少内存访问

性能分析：

实际应用案例：

5.2.2 TPU v2/v3：训练能力扩展

TPU v2标志着Google从推理专用向训练推理一体化的战略转变。这个转变背后有深刻的技术和商业考量。

从v1到v2的架构革新：

TPU v2不是v1的简单升级，而是一次彻底的架构重新设计。主要驱动因素包括：

  1. 训练需求激增：2016-2017年，Google内部的模型训练需求每3.5个月翻倍
  2. 模型规模爆炸：BERT、GPT等大模型需要数千GPU天的训练时间
  3. 软硬件协同：TensorFlow的发展需要配套的训练硬件

TPU v2关键升级：

bfloat16的设计智慧：

bfloat16（Brain Floating Point）是Google专门为深度学习设计的数据格式：

为什么选择bfloat16而非FP16？

  1. 动态范围优先：深度学习中，动态范围比精度更重要
  2. 转换简单：FP32到bfloat16只需截断尾数，无需调整指数
  3. 硬件友好：可以复用FP32的指数处理逻辑

TPU v3增强：

架构创新：

  1. 向量处理单元(VPU)：
    • 处理非矩阵运算（激活函数、归一化）
    • SIMD架构，128宽向量
    • 支持超越函数（exp、log、sqrt）
  2. 跨核心互连：
    • 2D mesh拓扑
    • 支持模型并行和数据并行
    • ICI (Inter-Core Interconnect) 带宽：50GB/s/方向
  3. 混合精度训练： 
    前向传播：bfloat16 计算
反向传播：bfloat16 计算  
权重更新：float32 累加

    误差分析： $\epsilon_{total} = \epsilon_{round} + \epsilon_{trunc} + \epsilon_{accum}$

    其中：

    • $\epsilon_{round}$：舍入误差 ≈ $2^{-8}$ (bfloat16)
    • $\epsilon_{trunc}$：截断误差 ≈ $2^{-7}$
    • $\epsilon_{accum}$：累积误差 ≈ $\sqrt{n} \times \epsilon_{round}$

5.2.3 TPU v4i：推理优化版本

TPU v4i是v4的推理优化变体，针对大规模推理部署。

优化特性：

性能规格：

推理优化技术：

  1. 动态批处理：
    • 请求合并减少延迟
    • 批大小自适应：$B_{opt} = \min(B_{max}, \lceil \frac{L_{target}}{L_{single}} \rceil)$
    • 其中 $L_{target}$ 为目标延迟，$L_{single}$ 为单样本延迟
  2. 算子融合优化： 
    Conv → BN → ReLU  融合为单个kernel
GEMM → Bias → Activation 融合执行

    内存访问减少：$M_{fused} = M_{input} + M_{output}$ 相比未融合：$M_{unfused} = M_{input} + 2 \times M_{intermediate} + M_{output}$

  3. 稀疏加速支持：
    • 2:4结构化稀疏硬件支持
    • 稀疏索引压缩：4位索引编码
    • 理论加速：2×，实际加速：1.6-1.8×

5.2.4 架构演进趋势分析

通过四代TPU的演进，可以总结出以下趋势：

  1. 专用化程度提升：
    • v1：纯推理
    • v2/v3：训练+推理
    • v4i：推理专用优化
  2. 内存系统演进：
    • 带宽增长：34GB/s → 1.2TB/s (35×)
    • 容量增长：28MB → 32GB (1000×)
    • 带宽/计算比：0.37 → 4.36 bytes/op
  3. 精度支持扩展：
    • v1：INT8
    • v2+：bfloat16、INT8、混合精度
    • v4：FP32累加器
  4. 系统级优化：
    • 单芯片 → 多核心 → Pod级系统
    • 软件栈成熟度提升
    • 编译器优化深化

5.3 控制流与指令集

5.3.1 VLIW vs MIMD vs SIMD架构对比

脉动阵列的控制方式直接影响其灵活性和效率。TPU采用了CISC风格的指令集，结合了多种控制模式的优点。理解这些控制模式的权衡对于设计高效的NPU至关重要。

控制架构的历史演进：

计算机体系结构的发展历程为我们提供了宝贵的经验。从Flynn分类法（1966年）开始，并行架构被分为SISD、SIMD、MISD和MIMD四类。NPU设计主要关注后三种，每种都有其独特的优势和适用场景。

Flynn分类法的现代解读：

在NPU设计中，纯粹的分类已经不够，现代架构往往是混合式的：

SIMD (Single Instruction Multiple Data)：

在脉动阵列中，SIMD是最自然的选择，因为矩阵乘法本质上是规则的、无分支的计算。每个PE执行相同的MAC操作，只是操作数不同。这种规律性使得SIMD控制极其高效。

SIMD在TPU中的具体实现：

  1. 统一控制器：单个控制器向所有PE广播控制信号
  2. 锁步执行：所有PE在同一时钟周期执行相同操作
  3. 条件执行处理：通过predicate mask实现部分PE的条件执行

控制开销的定量分析：

这种巨大的面积节省使得可以将更多的晶体管用于实际计算，这是脉动阵列能够达到高计算密度的关键。

脉动阵列中的SIMD体现：

MATRIX_MULTIPLY：
  所有PE执行: c += a × b
  数据不同，操作相同

VLIW (Very Long Instruction Word)：

TPU的VLIW设计：

指令格式：
[MXU_OP | VPU_OP | DMA_OP | SCALAR_OP]
   8bit     8bit     8bit      8bit

MIMD (Multiple Instruction Multiple Data)：

5.3.2 指令调度与依赖管理

TPU的指令调度采用静态调度为主、动态调整为辅的策略。这种设计选择反映了对编译时优化和运行时效率的平衡考虑。

静态调度的优势：

  1. 确定性延迟：每条指令的执行时间在编译时确定，便于性能预测
  2. 硬件简化：不需要复杂的乱序执行单元、重排序缓冲区等
  3. 功耗优化：避免了动态调度的功耗开销（约占CPU功耗的20-30%）
  4. 调试友好：执行顺序与程序顺序一致，便于调试和验证

指令级并行(ILP)分析：

设指令序列为 $I = {i_1, i_2, …, i_n}$，依赖图 $G = (V, E)$：

关键路径长度：$L_{critical} = \max_{path} \sum_{i \in path} latency(i)$

并行度量化：

在TPU中，典型的并行效率可达80-90%，这得益于：

依赖类型：

  1. RAW (Read After Write)：真依赖 
    i1: C = MATRIX_MUL(A, B)
i2: D = ACTIVATE(C)     // 必须等待i1完成
  2. WAR (Write After Read)：反依赖 
    i1: C = LOAD(addr1)
i2: STORE(addr1, D)     // 必须等待i1读取完成
  3. WAW (Write After Write)：输出依赖 
    i1: C = COMPUTE1()
i2: C = COMPUTE2()      // 必须保序

调度算法：

TPU编译器使用改进的List Scheduling算法：

def schedule_instructions(dag, resources):
    ready_list = get_ready_nodes(dag)
    schedule = []
    time = 0
    
    while ready_list:
        # 优先级函数
        priority = lambda n: (
            -critical_path_length(n),  # 关键路径优先
            -resource_usage(n),         # 资源密集优先
            node_id(n)                  # 确定性tie-breaking
        )
        
        node = max(ready_list, key=priority)
        
        if can_schedule(node, resources, time):
            schedule.append((node, time))
            update_resources(resources, node, time)
            ready_list.remove(node)
            ready_list.extend(get_newly_ready(dag, node))
        else:
            time += 1
    
    return schedule

流水线设计：

TPU采用5级流水线：

  1. IF (Instruction Fetch)：取指令
  2. ID (Instruction Decode)：解码和操作数读取
  3. EX1 (Execute 1)：地址计算/标量运算
  4. EX2 (Execute 2)：矩阵运算启动
  5. WB (Write Back)：结果写回

流水线冒险处理：

5.3.3 TPU指令集架构

TPU指令集包含以下主要类别：

1. 矩阵运算指令：

MatrixMultiply(src1_addr, src2_addr, dst_addr, M, K, N)
  - src1: M×K矩阵地址
  - src2: K×N矩阵地址  
  - dst: M×N结果地址
  - 执行周期: ceil(M/256) × ceil(N/256) × ceil(K/256)

2. 向量运算指令：

VectorOp(op_type, src_addr, dst_addr, length)
  - op_type: ADD/MUL/ReLU/Sigmoid等
  - 支持广播和element-wise操作
  - SIMD宽度: 128

3. 数据传输指令：

DMA_Load(dram_addr, sram_addr, size, stride)
DMA_Store(sram_addr, dram_addr, size, stride)  
  - 支持2D/3D传输模式
  - 最大传输: 16MB
  - 双缓冲支持

4. 同步与控制指令：

Barrier(barrier_id)     // 同步点
Branch(condition, target) // 条件跳转
Loop(count, body_addr)   // 硬件循环

指令编码示例：

31  28 27  24 23  20 19  16 15   12 11    8 7     0
┌────┬──────┬──────┬──────┬───────┬───────┬───────┐
│ Op │ Dst  │ Src1 │ Src2 │ Imm12 │ Flags │ Rsvd  │
└────┴──────┴──────┴──────┴───────┴───────┴───────┘

Op: 4位操作码
Dst/Src: 4位寄存器编号
Imm12: 12位立即数
Flags: 4位标志位(条件码等)

5.3.4 性能计数器与调试支持

TPU提供丰富的性能计数器用于性能分析和调试：

硬件性能计数器：

  1. 计算相关：
    • MXU利用率：active_cycles / total_cycles
    • VPU利用率
    • 整数运算单元利用率
  2. 内存相关：
    • HBM带宽利用率
    • Cache命中率
    • DMA传输效率
  3. 同步开销：
    • Barrier等待周期
    • Pipeline stall周期
    • 依赖等待时间

性能分析公式：

实际性能 = 理论峰值 × 利用率 × 效率

其中：

对于矩阵乘法： \(Performance = \frac{2MNK}{T_{execution}} = \frac{2MNK}{T_{compute} + T_{memory} + T_{sync}}\)

其中：

调试支持特性：

  1. 断点支持：硬件断点和watchpoint
  2. 单步执行：指令级单步
  3. 寄存器和内存dump
  4. Trace buffer：记录执行历史

本章小结

本章深入探讨了脉动阵列的基本原理和Google TPU的架构演进。我们学习了：

  1. 脉动阵列基础：
    • 数据流动的三种模式（WS/OS/RS）及其权衡
    • Weight-stationary设计在能效和带宽效率上的优势
    • 阵列维度选择对利用率的影响，边界效应的处理
  2. TPU架构演进：
    • 从推理专用（v1）到训练支持（v2/v3）再到推理优化（v4i）
    • 内存系统的大幅提升：带宽35倍增长，容量1000倍增长
    • 混合精度训练和稀疏加速的硬件支持
  3. 控制流设计：
    • SIMD/VLIW/MIMD架构的权衡
    • 静态调度为主的指令调度策略
    • 5级流水线设计和冒险处理机制

关键公式回顾：

练习题

基础题

习题5.1：矩阵乘法映射分析 给定一个128×128的脉动阵列，需要计算矩阵乘法 $C_{384 \times 512} = A_{384 \times 256} \times B_{256 \times 512}$。请计算：

  1. 需要多少个tiling块？
  2. 理论利用率是多少？
  3. 如果采用padding到最近的128倍数，padding开销是多少？

Hint：考虑如何将大矩阵分块映射到脉动阵列上。

参考答案 1. Tiling分块： - M维度：$\lceil 384/128 \rceil = 3$ 块 - N维度：$\lceil 512/128 \rceil = 4$ 块 - K维度：$\lceil 256/128 \rceil = 2$ 块 - 总块数：$3 \times 4 = 12$ 个输出块 2. 理论利用率： - 不考虑K维度的利用率：$U = \frac{384 \times 512}{3 \times 128 \times 4 \times 128} = \frac{196608}{196608} = 100\%$ - 考虑所有维度：需要 $3 \times 4 \times 2 = 24$ 次脉动阵列调用 3. Padding开销： - M维度无需padding（384 = 3×128） - N维度无需padding（512 = 4×128） - K维度无需padding（256 = 2×128） - 总padding开销：0

习题5.2：数据移动量计算 对于批大小B=32，输入特征I=1024，输出特征O=2048的全连接层，分别计算Weight-stationary和Output-stationary模式下的数据移动量。假设权重为FP16（2字节），激活为INT8（1字节）。

Hint：分别统计权重、输入激活和输出激活的读写次数。

参考答案 Weight-stationary模式： - 权重加载：$W = 1024 \times 2048 \times 2 = 4MB$（一次性加载） - 输入读取：$A_{in} = 32 \times 1024 \times 1 \times 2048 = 64MB$（每个输出通道读一次） - 输出写入：$A_{out} = 32 \times 2048 \times 1 = 64KB$ - 总数据移动：$D_{WS} = 4MB + 64MB + 64KB ≈ 68.06MB$ Output-stationary模式： - 权重读取：$W = 1024 \times 2048 \times 2 \times 32 = 128MB$（每个batch读一次） - 输入读取：$A_{in} = 32 \times 1024 \times 1 = 32KB$ - 输出写入：$A_{out} = 32 \times 2048 \times 1 = 64KB$ - 总数据移动：$D_{OS} = 128MB + 32KB + 64KB ≈ 128.19MB$ Weight-stationary节省：$(128.19 - 68.06) / 128.19 = 46.9\%$

习题5.3：稀疏加速计算 一个256×256的脉动阵列支持2:4结构化稀疏。对于一个包含75%稀疏度的权重矩阵（随机稀疏），转换为2:4结构化稀疏后，请计算：

  1. 结构化稀疏的实际稀疏度
  2. 理论加速比
  3. 考虑20%索引开销后的实际加速比

Hint：2:4稀疏要求每4个元素中恰好2个为零。

参考答案 1. 结构化稀疏的实际稀疏度： - 2:4稀疏固定为50%稀疏度 - 原始75%随机稀疏转换为2:4后，稀疏度降为50% 2. 理论加速比： - 相对于密集计算：$S_{theory} = \frac{1}{1-0.5} = 2×$ - 相对于75%随机稀疏：无法直接在硬件上加速 3. 实际加速比（考虑索引开销）： - $S_{actual} = \frac{2}{1 + 0.2} = \frac{2}{1.2} = 1.67×$ - 有效计算减少：$1 - \frac{1}{1.67} = 40\%$

挑战题

习题5.4：流水线性能分析 TPU的5级流水线中，矩阵乘法指令需要100个周期完成，向量运算需要10个周期，DMA传输需要50个周期。给定以下指令序列：

I1: DMA_Load(weights)     // 50 cycles
I2: DMA_Load(inputs)      // 50 cycles  
I3: MatrixMultiply()      // 100 cycles, 依赖I1,I2
I4: VectorAdd()           // 10 cycles, 依赖I3
I5: VectorReLU()          // 10 cycles, 依赖I4
I6: DMA_Store(outputs)    // 50 cycles, 依赖I5

请计算：

  1. 不考虑流水线的总执行时间
  2. 理想流水线执行时间（无stall）
  3. 如果I3必须等待I1和I2都完成才能开始，实际执行时间是多少？

Hint：画出流水线执行时序图。

参考答案 1. 顺序执行总时间： $T_{seq} = 50 + 50 + 100 + 10 + 10 + 50 = 270$ cycles 2. 理想流水线（假设可以完美重叠）： - I1和I2可以并行（DMA双通道）：50 cycles - I3开始：100 cycles - I4-I6顺序执行：10 + 10 + 50 = 70 cycles - 理想时间：$\max(50, 50) + 100 + 70 = 220$ cycles 3. 实际执行时间（考虑依赖）： ``` 时间 0-50: I1执行，I2并行执行 时间 50-150: I3执行（等待I1,I2完成） 时间 150-160: I4执行 时间 160-170: I5执行 时间 170-220: I6执行 ``` 总时间：220 cycles 流水线效率：$\eta = \frac{270}{220} = 1.23×$ 加速

习题5.5：TPU Pod系统性能建模 一个TPU v3 Pod包含1024个TPU芯片，采用2D torus互连拓扑。每个芯片有2个核心，每核心123 TFLOPS。对于训练一个GPT-3规模的模型（175B参数），请分析：

  1. 理论峰值算力
  2. 如果采用数据并行，每个芯片需要存储多少模型参数？
  3. 如果采用模型并行，通信开销如何估算？

Hint：考虑参数、梯度、优化器状态的存储需求。

参考答案 1. 理论峰值算力： - 单芯片：$2 \times 123 = 246$ TFLOPS - Pod总算力：$1024 \times 246 = 251,904$ TFLOPS ≈ 252 PFLOPS 2. 数据并行存储需求： - 模型参数：175B × 2 bytes (bfloat16) = 350GB - 梯度：175B × 2 bytes = 350GB - Adam优化器状态（动量+二阶矩）：175B × 4 bytes = 700GB - 总需求：350 + 350 + 700 = 1400GB per replica - 单芯片HBM：32GB - 需要：$\lceil 1400/32 \rceil = 44$ 个芯片才能存储完整模型 3. 模型并行通信开销： - 采用1024路模型并行 - 每个芯片存储：175B/1024 ≈ 171M参数 - 前向传播：每层需要All-Gather激活 - 反向传播：需要Reduce-Scatter梯度 - 通信量/step：$2 \times batch\_size \times hidden\_dim \times 2$ bytes - 对于hidden_dim=12288，batch=1：通信量≈49KB/层 - 96层transformer：总通信≈4.7MB/sample 通信时间估算（50GB/s互连）： $T_{comm} = \frac{4.7MB}{50GB/s} = 94μs/sample$

习题5.6：编译器Tiling策略优化 给定一个卷积层：输入[1,224,224,3]，卷积核[7,7,3,64]，步长2，padding 3。目标是映射到128×128的脉动阵列上。请设计：

  1. Im2col转换后的矩阵维度
  2. 最优的tiling策略（最小化数据移动）
  3. 估算总的数据移动量

Hint：考虑空间维度和通道维度的不同tiling方式。

参考答案 1. Im2col转换： - 输出尺寸：$\lfloor (224 + 2×3 - 7)/2 \rfloor + 1 = 112$ - Im2col矩阵： - 输入矩阵A：[112×112, 7×7×3] = [12544, 147] - 权重矩阵B：[7×7×3, 64] = [147, 64] - 输出矩阵C：[12544, 64] 2. Tiling策略： - 策略1：空间维度tiling - M维度（空间）：$\lceil 12544/128 \rceil = 98$ 块 - N维度（输出通道）：$\lceil 64/128 \rceil = 1$ 块 - K维度（卷积核）：$\lceil 147/128 \rceil = 2$ 块 - 策略2：优化tiling（利用64<128） - 将2个空间位置打包：[256, 147] × [147, 64] - M维度：$\lceil 12544/256 \rceil = 49$ 块 - 仍需要98次计算，但可以更好地利用阵列 3. 数据移动量估算： - 输入数据（Im2col后）：12544 × 147 × 1 byte = 1.84MB - 权重数据：147 × 64 × 1 byte = 9.4KB - 输出数据：12544 × 64 × 1 byte = 803KB 采用Weight-stationary： - 权重加载1次：9.4KB - 输入读取98次（每个M块）：1.84MB × 98 = 180MB - 输出写入1次：803KB - 总数据移动：≈181MB 优化：通过输入数据重用减少读取 - 相邻输出位置共享输入数据 - 实际数据移动可减少约30-40%

常见陷阱与错误

1. 阵列利用率陷阱

错误：假设脉动阵列总是能达到100%利用率

原因：

正确做法：

2. 内存带宽估算错误

错误：只考虑计算时间，忽略数据传输

示例：

错误估算：256×256阵列，700MHz
计算时间 = M×N×K / (256×256×700M) = ...

正确估算：
总时间 = max(计算时间, 数据传输时间)
如果数据传输时间 > 计算时间，则memory-bound

3. 精度损失累积

错误：在量化时不考虑误差累积

问题场景：

解决方案：

4. 稀疏模式不匹配

错误：假设任何稀疏模式都能加速

实际情况：

5. 编译器优化假设

错误：手动优化一定比编译器好

事实：

最佳实践检查清单

设计阶段

实现阶段

验证阶段

优化阶段

部署阶段