本章深入探讨TPU编译器的核心技术，重点分析XLA（Accelerated Linear Algebra）编译器如何将高层神经网络计算图映射到脉动阵列硬件上。我们将详细剖析编译优化策略、矩阵运算映射方法以及卷积算子的高效实现。通过理解编译器的工作原理，读者将掌握如何充分发挥脉动阵列架构的计算潜力，实现接近理论峰值的性能。
XLA（Accelerated Linear Algebra）编译器是Google为TPU开发的领域特定编译器，其设计理念是通过激进的编译时优化来最大化硬件利用率。与传统的深度学习框架运行时不同，XLA采用全程序编译（whole-program compilation）策略，能够跨越算子边界进行全局优化。这种方法特别适合TPU这类具有确定性执行模型的专用硬件。
XLA的架构设计体现了几个重要原则。第一，确定性执行：所有的调度和资源分配决策都在编译时完成，运行时没有任何不确定性，这使得性能高度可预测。第二，激进优化：由于目标硬件是已知的（TPU），编译器可以采用针对性极强的优化策略，不需要考虑可移植性。第三，全局视角：能看到整个计算图，可以进行跨算子的优化，如算子融合、数据布局转换等。第四，静态内存管理：所有内存分配在编译时确定，避免了运行时的动态分配开销。
HLO（High-Level Optimizer）作为XLA的核心中间表示，具有几个关键特性。首先，它采用静态形状系统，所有张量的维度在编译时必须已知，这使得编译器可以进行精确的内存规划和优化决策。其次，HLO使用函数式编程范式，每个操作都是纯函数，没有副作用，这极大地简化了优化pass的实现。第三，HLO支持丰富的原语集合，涵盖了深度学习中的常见操作，同时保持了足够的低层控制能力。
HLO的设计借鉴了传统编译器的SSA（Static Single Assignment）形式，每个值只被赋值一次，这简化了数据流分析和优化。同时，HLO的类型系统包含了丰富的形状信息，不仅包括维度大小，还包括数据布局（如行主序、列主序、分块布局等），这对于生成高效的硬件代码至关重要。
HLO图的基本构成包括：
Input(X) Weight(W) Bias(b)
[B,H,W,C] [C,K] [K]
\ / /
\ / /
MatMul(Y=XW) /
[B,H,W,K] /
| /
BiasAdd(Z=Y+b) ------
[B,H,W,K]
|
ReLU(A=max(0,Z))
[B,H,W,K]
|
Output
HLO的优化pipeline是一个精心设计的多阶段流程，每个阶段针对特定的优化目标。优化passes的执行顺序经过精心编排，确保早期的优化为后续优化创造机会。整个优化流程可以分为三个主要阶段：
前端优化阶段：主要进行与硬件无关的高层优化，包括代数简化、常量折叠、公共子表达式消除等。这些优化减少了计算图的复杂度，为后续的硬件特定优化奠定基础。
中端优化阶段：进行算子融合、内存规划、数据布局优化等与硬件相关但不涉及具体指令生成的优化。这是XLA优化的核心阶段，大部分性能提升来自于此。
后端优化阶段：将HLO lowering到硬件特定的指令，进行指令调度、寄存器分配等低层优化。这个阶段紧密结合TPU的微架构特征。
主要的优化类别包括：
HLO的语义规范与类型系统：
HLO指令集包含约100个原语（primitives），每个都有精确的语义定义。这些原语覆盖了深度学习的主要计算模式：
HLO的类型系统建立在形状（Shape）概念之上：
Shape = (element_type, dimensions, layout)
其中：
element_type：数据类型（F16、F32、F64、S8、S16、S32、S64、U8、U16、U32、U64、PRED、C64、C128）dimensions：各维度大小的数组，如[batch, height, width, channels]layout：维度的物理存储顺序，如{3,2,1,0}表示NHWC布局形状推断与验证：
XLA在构建HLO图时执行严格的形状推断和类型检查：
前向推断（Forward Inference）： 从输入形状推导输出形状。例如，对于矩阵乘法： \(\text{Dot}([M, K], [K, N]) \rightarrow [M, N]\)
对于卷积操作，输出形状计算： \(H_{out} = \lfloor \frac{H_{in} + 2 \times \text{pad}_h - \text{dilation}_h \times (K_h - 1) - 1}{\text{stride}_h} \rfloor + 1\)
反向推断（Backward Inference）： 某些情况下从输出形状反推输入形状，用于验证和优化。
HLO的中间表示特性：
不可变性（Immutability）： HLO采用纯函数式设计，所有操作产生新值而不修改现有值。这简化了分析和优化，避免了别名分析的复杂性。
单赋值（Single Assignment）： 每个HLO指令产生一个唯一的值，该值只被定义一次。这类似于SSA形式，使得def-use链清晰明确。
显式依赖（Explicit Dependencies）： 所有数据和控制依赖都在图中显式表示，没有隐式的副作用或全局状态。
嵌套结构（Nested Structure）： HLO支持嵌套的计算（Computation），用于表示函数、循环体、条件分支等。每个Computation是一个独立的图，有自己的参数和返回值。
优化Pass的实现机制：
XLA的优化器采用Pass管理器架构，每个Pass是一个独立的转换：
class HloPass {
virtual StatusOr<bool> Run(HloModule* module) = 0;
virtual string name() const = 0;
};
Pass依赖管理： 某些Pass依赖其他Pass的结果，Pass管理器确保正确的执行顺序。例如，算子融合需要在内存分配之前完成。
Pattern Matching与重写规则：
XLA使用模式匹配来识别优化机会：
auto pattern = m::Dot(m::Op(), m::Transpose(m::Op())); // 匹配 A @ B^T
Profile-Guided Optimization（PGO）：
XLA支持基于性能剖析的优化：
算子融合（Operator Fusion）是XLA编译器最重要的优化技术之一，其核心思想是将多个细粒度算子合并为粗粒度的融合算子，从而减少内存访问开销并提高计算密度。在TPU等专用加速器上，内存带宽往往是性能瓶颈，算子融合通过减少中间结果的存储和读取，可以带来显著的性能提升。融合策略的设计需要在多个维度进行权衡：融合范围、资源约束、数值精度和硬件特性。
算子融合的理论基础源于计算强度（Computational Intensity）的概念，定义为算术运算次数与内存访问字节数的比值。通过融合，我们可以提高整体的计算强度，使其更接近硬件的峰值计算强度（由计算吞吐量与内存带宽的比值决定）。根据Roofline模型，当计算强度低于硬件峰值时，程序性能受内存带宽限制；融合通过减少内存访问，将程序推向计算受限区域，从而提高硬件利用率。
垂直融合（Producer-Consumer Fusion）： 垂直融合是最常见的融合模式，它将数据依赖链上的相邻算子合并。这种融合模式的关键在于消除中间张量的物化（materialization），即避免将中间结果写入内存。在TPU的脉动阵列架构中，垂直融合可以利用累加器（accumulator）直接传递部分结果，显著减少内存带宽需求。
融合决策需要考虑多个约束条件：
常见的垂直融合模式及其收益分析：
线性层融合链（Linear Layer Fusion Chain） \(Y = \text{Activation}(\text{Norm}(XW + b))\) 这是深度学习中最常见的模式，包含矩阵乘法、偏置加法、归一化和激活函数。未融合时需要4次内存访问（读X和W，写临时结果3次，写最终结果），融合后只需2次（读输入，写输出）。内存访问减少率：$(4-2)/4 = 50\%$
批归一化与激活融合（BatchNorm-Activation Fusion） \(Y = \text{ReLU}\left(\gamma \frac{X - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta\right)\) 批归一化涉及多个统计量和参数，融合可以避免存储归一化后的中间结果。特别是在推理阶段，$\mu$和$\sigma$是固定的，可以预计算$\frac{\gamma}{\sqrt{\sigma^2 + \epsilon}}$和$\beta - \frac{\gamma\mu}{\sqrt{\sigma^2 + \epsilon}}$，将运算简化为线性变换加激活
注意力机制融合（Attention Fusion） \(\text{Attention}(Q,K,V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V\) Flash Attention等技术通过分块融合，避免存储完整的注意力矩阵$QK^T$，将内存复杂度从$O(N^2)$降至$O(N)$
水平融合（Horizontal Fusion）： 水平融合将数据独立的并行算子打包执行，其目标是提高硬件利用率，特别是当单个算子无法充分利用硬件资源时。这种融合模式在处理小批量或小矩阵时特别有效。
水平融合的关键技术包括：
水平融合的收益模型： \(\text{Speedup} = \frac{\sum_{i} T_i}{\max_j(T_j) + T_{\text{overhead}}}\) 其中$T_i$是各算子的独立执行时间，$T_{\text{overhead}}$是融合带来的额外开销（如同步、资源竞争等）
适用场景与收益分析：
循环融合（Loop Fusion）： 循环融合是另一类重要的融合技术，它合并具有相同或兼容迭代空间的循环，减少循环开销并提高数据局部性。
循环融合的条件与约束：
循环融合的类型：
原始：for i: A[i] = B[i] + C[i]
for i: D[i] = A[i] * E[i]
融合：for i:
A[i] = B[i] + C[i]
D[i] = A[i] * E[i]
原始：for i in [0, N): A[i] = ...
for i in [M, N+M): B[i] = ...
融合：for i in [0, M): A[i] = ...
for i in [M, N): A[i] = ...; B[i] = ...
for i in [N, N+M): B[i] = ...
原始：for i: for j: A[i][j] = ...
for i: for j: B[i][j] = A[i][j] + ...
融合：for i: for j:
A[i][j] = ...
B[i][j] = A[i][j] + ...
算子融合的高级模式：
融合决策的机器学习方法：
现代编译器越来越多地使用机器学习来指导融合决策：
融合的代码生成策略：
算子融合后需要生成高效的融合kernel代码：
融合的性能建模与预测：
准确预测融合收益是编译器决策的关键：
静态分析模型： \(\text{Benefit} = \text{MemSaved} \times BW - \text{ComputeOverhead} \times \frac{1}{Throughput}\)
其中：
内存管理是NPU编译器的核心挑战之一，特别是在TPU这类片上存储容量受限的架构中。TPU采用了分层的存储体系：每个核心有私有的片上SRAM（如TPUv4i的32MB），多个核心共享HBM（144MB），以及系统级的主存。编译器必须精心编排数据在各级存储间的移动，以最大化数据重用并隐藏访存延迟。XLA采用静态内存管理策略，在编译时完成所有内存分配决策，避免了运行时的动态分配开销。
内存规划的核心问题可以形式化为一个约束优化问题：给定计算图$G=(V,E)$，其中节点$v \in V$表示张量，边$e \in E$表示数据依赖，目标是为每个张量分配内存地址，使得峰值内存使用最小化，同时满足容量约束和对齐要求。这个问题是NP困难的，实际中使用启发式算法求解。
静态内存分配的三阶段框架：
第一阶段是生命周期分析（Liveness Analysis），确定每个张量的活跃区间。张量的生命周期由三个关键时刻定义：
生命周期分析的算法实现：
for each tensor T in computation_graph:
T.birth = execution_time(T.producer)
T.uses = {execution_time(op) for op in T.consumers}
T.death = max(T.uses) + epsilon // epsilon保证在最后使用后释放
T.live_interval = [T.birth, T.death]
生命周期的精确分析需要考虑算子的执行顺序。XLA使用拓扑排序确定执行顺序，并通过依赖分析优化调度，以减少同时活跃的张量数量。生命周期重叠的张量不能共享内存，这构成了内存分配的基本约束。
生命周期优化技术：
第二阶段是冲突图构建与着色（Conflict Graph Coloring）。构建冲突图$G_c=(V_c, E_c)$，其中：
内存分配问题转化为图着色问题：为每个节点分配一个”颜色”（内存块），使得相邻节点颜色不同。使用的贪心启发式算法包括：
第三阶段是内存池管理（Memory Pool Management）。XLA使用分层的内存池设计：
内存分配的优化目标函数： \(\min \max_{t \in T} \sum_{v \in \text{Live}(t)} \text{Size}(v)\) 其中$T$是时间步集合，$\text{Live}(t)$是时刻$t$活跃的张量集合，$\text{Size}(v)$是张量$v$的大小。
双缓冲与三缓冲技术（Double/Triple Buffering）：
双缓冲是隐藏内存访问延迟的经典技术，通过重叠计算与数据传输实现流水线并行。在TPU中，双缓冲的实现涉及三个关键组件：
时序安排： Cycle 0-99: 计算(A[0]) | DMA加载(B[1]) | 空闲 Cycle 100-199: 计算(B[1]) | DMA加载(A[2]) | DMA存储(A[0]结果) Cycle 200-299: 计算(A[2]) | DMA加载(B[3]) | DMA存储(B[1]结果)
2. **同步原语设计**：
- **生产者-消费者信号量**：确保数据准备就绪后才开始计算
- **DMA完成中断**：通知CPU数据传输完成
- **栅栏指令（Barrier）**：全局同步点，确保所有操作完成
3. **性能模型与分析**：
设计算时间为$T_c$，数据传输时间为$T_m$，则：
- 无缓冲：总时间 = $N \times (T_c + T_m)$
- 双缓冲：总时间 = $T_m + N \times \max(T_c, T_m)$
- 加速比：$S = \frac{T_c + T_m}{\max(T_c, T_m)}$，理想情况下接近2
**内存布局优化（Memory Layout Optimization）**：
数据在内存中的布局方式直接影响访问效率。XLA支持多种布局转换：
1. **维度重排（Dimension Reordering）**：
- NHWC → NCHW：适应不同硬件的偏好
- 优化准则：内层循环访问的维度应连续存储
2. **内存对齐（Memory Alignment）**：
- 确保数据地址对齐到cache line（通常64字节）
- 使用padding填充，公式：$\text{AlignedSize} = \lceil \frac{\text{Size}}{\text{Alignment}} \rceil \times \text{Alignment}$
3. **Bank冲突避免**：
- TPU的SRAM通常组织为多个bank，并行访问不同bank可提高带宽
- 通过地址交织（interleaving）避免冲突：$\text{Bank}(addr) = (addr / \text{ElementSize}) \bmod \text{NumBanks}$
### 7.1.4 Tiling策略与参数选择
Tiling（分块）是将大规模张量运算分解为硬件友好的小块计算的核心技术，它在编译器优化中扮演着至关重要的角色。Tiling的本质是在计算的时间局部性和空间局部性之间寻找最优平衡点，使得工作集能够驻留在快速的片上存储中，同时最大化数据重用。在TPU等脉动阵列架构上，正确的tiling策略可以将性能提升数倍甚至数十倍。
Tiling策略的理论基础源于多面体模型（Polyhedral Model），它将嵌套循环的迭代空间表示为整数点的多面体，通过仿射变换实现循环优化。对于深度学习工作负载，tiling不仅要考虑传统的cache优化，还要适配专用硬件的特殊约束，如脉动阵列的固定维度、向量单元的SIMD宽度、以及DMA传输的突发长度要求。
**多维Tiling参数空间的形式化定义**：
对于广义的张量运算，我们定义一个n维的tiling参数向量$\vec{T} = (T_1, T_2, ..., T_n)$，其中每个$T_i$表示第i个维度的tile大小。以矩阵乘法$C_{M \times N} = A_{M \times K} \times B_{K \times N}$为例，完整的tiling参数空间包括：
1. **空间维度tiling**：
- $T_M$：输出矩阵的行维度tile大小
- $T_N$：输出矩阵的列维度tile大小
- $T_K$：归约维度（内积维度）的tile大小
2. **时间维度tiling**（多级tiling）：
- $(T_{M1}, T_{M2})$：M维度的两级tiling，外层循环步长$T_{M1}$，内层$T_{M2}$
- 多级tiling可以更好地适配多级存储层次
3. **并行维度tiling**：
- $T_P$：跨多个处理单元的并行分块大小
- 需要考虑负载均衡和通信开销
Tiling参数必须满足的约束系统：
1. **硬件资源约束**：
$$T_M \times T_N \leq SA_{rows} \times SA_{cols}$$
其中$SA_{rows}$和$SA_{cols}$是脉动阵列的行列数。这确保单个tile能够映射到硬件上。
2. **存储容量约束**：
$$\text{sizeof}(A_{tile}) + \text{sizeof}(B_{tile}) + \text{sizeof}(C_{tile}) \leq SRAM_{capacity}$$
展开为：
$$T_M \times T_K \times s_A + T_K \times T_N \times s_B + T_M \times T_N \times s_C \leq SRAM_{capacity}$$
其中$s_A$、$s_B$、$s_C$是元素大小（如FP16为2字节）
3. **数据对齐约束**：
$$T_i \equiv 0 \pmod{A_i}$$
其中$A_i$是第i维度的对齐要求（通常是向量宽度的倍数，如128）
4. **DMA传输约束**：
$$T_i \times s \geq DMA_{min\_burst}$$
确保每次传输达到DMA的最小突发长度，提高传输效率
**性能建模与代价函数**：
准确的性能模型是tiling优化的基础。我们构建一个分析模型来预测不同tiling参数下的执行时间：
$$T_{total} = T_{compute} + T_{memory} - T_{overlap}$$
模型的精确度取决于对硬件特性的准确建模：
其中：
- $T_{compute} = \frac{M \times N \times K}{T_M \times T_N \times T_K} \times T_{tile\_compute}$
- $T_{memory} = T_{load\_A} + T_{load\_B} + T_{store\_C}$
- $T_{overlap}$：计算与数据传输的重叠时间
详细的内存访问时间建模：
$$T_{load\_A} = \frac{M \times K}{R_A} \times \frac{1}{BW_{eff}}$$
其中重用因子$R_A$的计算：
$$R_A = \begin{cases}
\frac{N}{T_N} & \text{if A is reused across N dimension} \\
1 & \text{otherwise}
\end{cases}$$
有效带宽$BW_{eff}$考虑了访问模式的影响：
$$BW_{eff} = BW_{peak} \times \eta_{pattern} \times \eta_{conflict}$$
其中$\eta_{pattern}$是访问模式效率（连续访问接近1，随机访问可能低至0.1），$\eta_{conflict}$是bank冲突因子。
**访问模式效率的详细分析**：
- **连续访问（Sequential）**：$\eta_{pattern} = 0.95-1.0$，充分利用突发传输
- **跨步访问（Strided）**：$\eta_{pattern} = \frac{1}{1 + \alpha \cdot stride}$，$\alpha$是惩罚系数
- **随机访问（Random）**：$\eta_{pattern} = 0.1-0.3$，严重影响性能
- **块访问（Block）**：$\eta_{pattern} = \frac{block\_size}{block\_size + gap\_size}$
**Bank冲突分析**：
Bank冲突发生条件：多个并发访问映射到同一bank
$$\eta_{conflict} = \frac{1}{1 + \beta \cdot P_{conflict}}$$
其中$P_{conflict}$是冲突概率，$\beta$是冲突惩罚系数（通常2-4）
**自动调优框架（Auto-tuning Framework）**：
现代编译器越来越依赖机器学习技术来搜索最优的tiling参数。XLA集成了多种自动调优方法：
1. **基于搜索的方法**：
- **穷举搜索**：适用于小参数空间，保证找到最优解
- **遗传算法**：通过进化策略探索大参数空间
- **模拟退火**：允许接受次优解以跳出局部最优
2. **基于学习的方法**：
- **代价模型学习**：使用历史数据训练性能预测模型
$$\hat{T} = f_{ML}(\vec{T}, \vec{F})$$
其中$\vec{F}$是问题特征向量（矩阵大小、稀疏度等）
- **迁移学习**：将相似问题的优化经验迁移到新问题
- **强化学习**：将编译决策序列建模为马尔可夫决策过程
3. **混合策略**：
结合分析模型的指导性和机器学习的适应性：
- 使用分析模型剪枝搜索空间
* 根据硬件约束排除不可行解
* 使用理论上界指导搜索方向
* 快速筛选候选参数集
- 用机器学习微调分析模型的预测
* 收集实际执行数据作为训练集
* 使用梯度提升树（GBDT）或神经网络
* 特征工程：tile大小、数据重用率、内存访问模式等
- 在线学习持续改进模型
* 增量学习新的工作负载特征
* 动态调整模型参数
* A/B测试验证优化效果
**Tiling策略的高级优化技术**：
1. **矩形tiling vs 梯形tiling**：
- 矩形tiling简单但可能有边界浪费
* 适用于维度可整除的情况
* 控制逻辑简单，硬件实现成本低
* 可能导致10-20%的计算资源浪费
- 梯形tiling可以更好地处理非整除情况
* 根据剩余大小调整最后一个tile
* 需要额外的边界处理逻辑
* 提高硬件利用率至95%以上
- 三角tiling（特殊场景）
* 适用于三角矩阵运算
* 减少冗余计算
* 负载均衡挑战大
2. **动态tiling**：
- 根据输入大小动态调整tile参数
* 运行时分析输入特征
* 选择预编译的最佳kernel
* 避免重新编译开销
- 使用查找表存储常见大小的最优参数
* 离线构建参数表
* 使用插值处理中间大小
* 定期更新优化参数
- 自适应调整机制
* 监控实际性能指标
* 在线调整tiling参数
* 平滑过渡避免抖动
3. **协同tiling**：
- 同时优化多个相关算子的tiling
* 分析算子间的数据依赖
* 协调tile边界对齐
* 最大化数据重用
- 考虑算子间的数据传递模式
* Producer-Consumer模式
* Pipeline模式
* Wavefront模式
- 全局优化目标
* 最小化总体执行时间
* 平衡各阶段负载
* 减少中间数据存储
4. **多级tiling（Hierarchical Tiling）**：
- L1/L2/L3级别tiling
* 每级针对不同存储层次优化
* 嵌套tile结构
* 复杂度与收益权衡
- 时空联合tiling
* 空间维度：数据分块
* 时间维度：计算顺序
* 联合优化两个维度
## 7.2 矩阵乘法映射
### 7.2.1 大矩阵分块策略
将大规模矩阵乘法高效映射到有限大小的脉动阵列是编译器的核心任务。考虑矩阵乘法 $C = A \times B$，其中 $A \in \mathbb{R}^{M \times K}$，$B \in \mathbb{R}^{K \times N}$。
**分块算法设计**：
基本分块策略将计算分解为三层嵌套循环：
for m_tile in range(0, M, T_M): for n_tile in range(0, N, T_N): for k_tile in range(0, K, T_K): # 计算C[m_tile:m_tile+T_M, n_tile:n_tile+T_N]的部分和 C_tile += A_tile @ B_tile
**数据布局优化**：
1. **行主序vs列主序**：
- A矩阵采用行主序存储，便于按行streaming
- B矩阵采用列主序存储，便于按列广播
- 减少数据重排开销
2. **Z字形布局（Z-order）**：
提高空间局部性，适合2D脉动阵列
原始布局: 0 1 2 3 Z字形: 0 1 4 5 4 5 6 7 2 3 6 7 8 9 A B 8 9 C D C D E F A B E F
3. **块内连续存储**：
将tile内的数据连续存放，提高DMA效率
- 减少地址生成开销
- 提高突发传输（burst）效率
**边界处理与padding**：
当矩阵维度不能被tile大小整除时，需要特殊处理：
1. **零填充（Zero Padding）**：
$$M' = \lceil \frac{M}{T_M} \rceil \times T_M$$
$$N' = \lceil \frac{N}{T_N} \rceil \times T_N$$
优点：简化控制逻辑
缺点：浪费计算资源
2. **动态tile大小**：
最后一个tile使用较小的维度
- 需要硬件支持可变大小配置
- 控制逻辑更复杂
3. **预计算掩码（Predication）**：
使用掩码禁用超出边界的计算
- 保持规则的tile大小
- 通过使能信号控制PE
### 7.2.2 性能优化技巧
**工作负载平衡**：
确保所有PE均匀分配计算任务，避免负载不均。
1. **静态调度**：
编译时确定每个PE的计算任务
- 优点：无运行时开销
- 缺点：灵活性差
2. **工作窃取（Work Stealing）**：
空闲PE从忙碌PE窃取任务
- 需要硬件支持任务队列
- 增加控制复杂度
**数据预取优化**：
预取策略设计：
1. **软件流水线**：
加载A[i+1] | 加载B[i+1] | 计算C[i] = A[i] × B[i]
重叠下一轮的数据加载与当前计算
2. **预取距离计算**：
$$D_{prefetch} = \lceil \frac{L_{memory}}{T_{compute}} \rceil$$
其中$L_{memory}$是内存延迟，$T_{compute}$是计算时间
3. **自适应预取**：
根据运行时访问模式动态调整预取策略
### 7.2.3 批处理维度处理
现代深度学习中，批处理维度（batch dimension）的高效处理至关重要。
**批处理GEMM映射**：
对于批处理矩阵乘法 $C[b] = A[b] \times B[b]$，$b \in [0, B)$：
1. **串行处理**：
依次处理每个批次
- 简单但效率低
- 适合batch size较小的场景
2. **批次并行**：
将不同批次映射到不同的PE组
$$\text{PE\_group}[i] \leftarrow \text{Batch}[i \bmod G]$$
其中G是PE组数量
3. **批次合并**：
将批次维度展开到M或N维度
$$A_{BM \times K} = \text{reshape}(A_{B \times M \times K})$$
提高大矩阵的硬件利用率
**动态批处理优化**：
适应可变batch size的策略：
1. **Padding到2的幂次**：
$$B' = 2^{\lceil \log_2 B \rceil}$$
简化地址计算但可能浪费计算
2. **分组处理**：
将batch分为大小相近的组
- 组内使用相同的tiling参数
- 减少重配置开销
3. **在线合并（Online Batching）**：
动态合并小batch以提高利用率
- 需要考虑延迟约束
- 适合推理服务场景
## 7.3 卷积映射优化
### 7.3.1 Im2col变换
Im2col（Image to Column）是将卷积运算转换为矩阵乘法的经典方法。
**变换原理**：
对于卷积运算 $Y = X * W$，其中：
- 输入：$X \in \mathbb{R}^{B \times H \times W \times C_{in}}$
- 卷积核：$W \in \mathbb{R}^{K_h \times K_w \times C_{in} \times C_{out}}$
- 输出：$Y \in \mathbb{R}^{B \times H_{out} \times W_{out} \times C_{out}}$
Im2col变换步骤：
1. 将输入展开为矩阵：$X_{col} \in \mathbb{R}^{(B \cdot H_{out} \cdot W_{out}) \times (K_h \cdot K_w \cdot C_{in})}$
2. 将权重reshape：$W_{col} \in \mathbb{R}^{(K_h \cdot K_w \cdot C_{in}) \times C_{out}}$
3. 执行矩阵乘法：$Y_{col} = X_{col} \times W_{col}$
4. 将结果reshape回原始形状
**内存开销分析**：
Im2col的主要缺点是内存膨胀：
$$\text{膨胀率} = K_h \times K_w$$
对于3×3卷积，数据量增加9倍。优化策略：
1. **分块Im2col**：只展开当前处理的tile
2. **隐式Im2col**：通过地址生成实现虚拟展开
3. **重叠优化**：复用相邻窗口的重叠数据
### 7.3.2 Direct Convolution
直接卷积避免了Im2col的内存开销，直接在脉动阵列上实现卷积运算。
**空间映射策略**：
1. **输出静止（Output Stationary）**：
- 每个PE负责一个输出像素
- 输入和权重流经PE阵列
- 适合大卷积核
2. **权重静止（Weight Stationary）**：
- 每个PE存储部分权重
- 输入数据广播到所有PE
- 适合深度可分离卷积
3. **输入静止（Input Stationary）**：
- 每个PE缓存部分输入
- 权重在PE间传递
- 适合1×1卷积
**数据流优化**：
输入特征图
↓
[Shift Register] ← 实现滑动窗口
↓
[PE Array] × 权重
↓
输出特征图 ```
使用移位寄存器链实现滑动窗口：
Winograd算法通过数论变换减少乘法次数，特别适合小卷积核。
F(2,3)变换示例： 将3×3卷积的2×2输出tile计算从16次乘法减少到4次。
变换矩阵： \(G = \begin{bmatrix} 1 & 0 & 0 \\ 0.5 & 0.5 & 0.5 \\ 0.5 & -0.5 & 0.5 \\ 0 & 0 & 1 \end{bmatrix}\)
\[B^T = \begin{bmatrix} 1 & 0 & -1 & 0 \\ 0 & 1 & 1 & 0 \\ 0 & -1 & 1 & 0 \\ 0 & 1 & 0 & -1 \end{bmatrix}\] \[A^T = \begin{bmatrix} 1 & 1 & 1 & 0 \\ 0 & 1 & -1 & -1 \end{bmatrix}\]计算步骤：
适用性分析： Winograd优势条件：
不适用场景：
深度可分离卷积（Depthwise Separable）： 分解为depthwise和pointwise两步：
映射策略：
空洞卷积（Dilated Convolution）： 通过调整数据访问模式支持dilation： \(\text{Index}(i,j) = i \times \text{dilation}_h + j \times \text{dilation}_w\)
硬件支持：
分组卷积（Grouped Convolution）： 将输入输出通道分组，组间独立计算： \(Y_g = X_g * W_g, \quad g \in [0, G)\)
优化要点：
本章详细探讨了TPU编译器的核心技术，从XLA编译流程到具体算子的硬件映射策略。关键要点包括：
练习7.1：Tiling参数计算 给定脉动阵列大小128×128，片上SRAM容量8MB，要计算矩阵乘法C[1024×1024] = A[1024×768] × B[768×1024]，数据类型为FP16。计算最优的tiling参数(T_M, T_N, T_K)。
提示：考虑内存约束：$2(T_M \times T_K + T_K \times T_N + T_M \times T_N) \leq 8MB$
练习7.2：算子融合收益分析 考虑融合BatchNorm-ReLU序列，输入张量大小为[64, 224, 224, 128]，数据类型FP16。计算融合前后的内存访问量。
提示：BatchNorm需要读写一次中间结果，融合后可以省去这次读写。
练习7.3：Im2col内存膨胀计算 对于输入张量[32, 224, 224, 64]，使用3×3卷积，stride=1，padding=1，输出通道128。计算Im2col变换后的内存需求。
提示：Im2col后矩阵大小为$(B \times H_{out} \times W_{out}) \times (K_h \times K_w \times C_{in})$
练习7.4：Winograd数值稳定性分析 Winograd F(2,3)变换中，变换矩阵包含0.5和-0.5。如果输入数据范围是[-1, 1]，分析经过变换后的数值范围，并讨论对INT8量化的影响。
提示：分析$B^T \cdot X \cdot B$的最大值情况
练习7.5：多核负载均衡优化 有4个TPU核心，需要处理批大小为100的推理任务。设计一个负载均衡策略，使得：
提示：考虑静态分配vs动态调度的权衡
练习7.6：编译器优化决策树 设计一个决策树，根据网络层的特征（输入大小、卷积核大小、通道数等）自动选择最优的卷积实现策略（Im2col、Direct、Winograd）。
提示：考虑计算密度、内存带宽、数值精度等因素
练习7.7：流水线优化分析 分析一个3级流水线：加载权重→计算→存储结果。如果三个阶段分别需要100、200、50个周期，如何优化流水线以提高吞吐量？
提示：找出瓶颈阶段，考虑并行化或分割策略