4. 内存层次与多级分块 (Memory Hierarchy & Tiling)

在解决了数据布局后，编译器必须面对现代处理器最严峻的挑战——**内存墙 (Memory Wall)**问题。

本章关注如何将庞大的数据与算力映射到硬件陡峭的存储金字塔 (Memory Pyramid)上。核心目标是掩盖不同层级存储器 (HBM/DDR L2 L1/Shared Register) 之间的延迟差异。通过Multi-level Tiling技术，编译器将计算限制在高速缓存内完成，并利用显式的异步搬运与生命周期管理技术，确保计算单元永远"有数可算"，实现带宽利用率的最大化。

内存墙问题由Wulf & McKee[11]在1995年首次系统阐述，成为计算机体系结构的基础理论。Roofline模型[12]将这一理论形式化，成为HPC性能分析的标准工具。现代AI编译器（TVM[30]、Halide[2]等）通过多级分块和显式内存管理来应对这一挑战。

4.1 多级分块 (Multi-level Tiling)

4.1.1 Register/L1/L2 Tiling

背景

为了掩盖 DRAM 的高延迟，必须将大张量切分为适应各级 Cache 大小的 Tile (图块)。Tile-local Fusion 的本质是将原本串行的"全图算子"，转换为在 Tile 粒度上紧密耦合的"子图融合"，确保生产者 (Producer) 生成的 Tile 在被驱逐出 Cache 前，立刻被消费者 (Consumer) 使用。

Wolf & Lam[1]的经典研究奠定了数据局部性优化的理论基础，Halide[2]等现代编译器将多级分块技术推向了新的高度。Auto-Scheduler[3]等研究实现了自动化的分块策略选择。

技术原理

多级分块不仅仅是循环变换，它是对硬件存储带宽金字塔的数学适配。其核心原理包含以下三个层面：

1. 带宽放大效应 (Bandwidth Amplification)：

   • 硬件的存储带宽呈倒金字塔状：。
   • Tiling 的目标是确保：越慢的内存，访问频率越低[1,11,12]。
   • 通过分块，数据从 HBM 加载一次到 L2，从 L2 加载一次到 L1，但在 L1 和寄存器之间进行成百上千次的高频读写。这使得计算单元 (ALU) 感受到的是寄存器的超高带宽，而非 HBM 的低带宽。

2. 表面积-体积比优化 (Surface-to-Volume Ratio Optimization)：

   • 以矩阵乘法为例，计算量是 (体积)，数据量是 (表面积)。
   • 如果对全图计算，数据复用率低。
   • 如果切分为 的 Tile，加载 的数据可以进行 次计算[1]。
   • 原理：只要 Tile 大小 足够大，计算密度 (Compute/Memory Ratio) 就会随 线性增长，直到填满 Cache 容量。

3. 硬件作用域映射 (Hardware Scope Mapping)： 编译器将不同层级的 Tile 映射到硬件不同的并行层级，这决定了数据的共享范围[2,30]：

   • L2 Tile (Grid Level)：映射到 GPU Grid 或 NPU Cluster。数据在不同核心间不共享 (或通过 L2 弱共享)。
   • L1/Shared Tile (Block Level)：映射到 GPU Thread Block 或 NPU AI Core。数据加载到 Shared Memory/UB，被这一组线程/Core 显式共享。
   • Register Tile (Thread Level)：映射到 GPU Thread 或 NPU Vector Unit。数据驻留在寄存器堆，仅被当前线程/指令独占，速度最快。

策略

• L2 Tiling：适应 L2 Cache 大小，优化网格 (Grid) 调度。
• L1/Shared Tiling：适应 L1/SRAM 大小，优化线程块 (Thread Block) 调度。
• Register Tiling：适应寄存器文件大小，优化指令级并行。

MLIR 实现

MLIR 通过 scf 循环嵌套表达分块逻辑，并利用 linalg.tile pass 自动生成。

// 原始算子：全量计算
%0 = linalg.matmul ins(%A, %B) outs(%C)
  
  
// 优化后：三级分块结构 (L2 -> L1 -> Register)
// 外层：L2 Cache Tile, Size=64 (对应 Grid/WorkGroup)
scf.for %i0 = 0 to 1024 step 64 {
  scf.for %j0 = 0 to 1024 step 64 {
    // 中层：L1/Shared Memory Tiling, Size=8 (对应 ThreadBlock)
    scf.for %i1 = %i0 to min(%i0 + 64, 1024) step 8 {
      scf.for %j1 = %j0 to min(%j0 + 64, 1024) step 8 {
        // 内层：Register Tiling (对应 SIMD/Vector 指令)
        // 此时数据驻留在寄存器堆，执行极致的融合计算
        %tile = linalg.matmul ... // 8x8 micro-kernel矩阵乘法
      }
    }
  }
}

---

4.2 显式内存层级管理 (Explicit Hierarchy Management)

4.2.1 Memory Promotion & Copy-Compute Overlap

背景

大多数 AI 处理器 (GPU/TPU/NPU/DSP) 都包含用户可控的片上高速缓存 (Scratchpad Memory) (如 GPU 的 Shared Memory 或 NPU 的 Local Memory)。内存提升 (Promotion) 指将频繁访问的数据从慢速的主存 (HBM/DDR) 显式搬运到片上高速缓存。

为了避免数据搬运阻塞计算单元，必须利用硬件的 DMA 机制 实现拷贝与计算的时间重叠 (Overlap)。这一技术在GPU[15]、NPU[18,19]等硬件上都有广泛应用。

技术原理

• Asynchronous Data Movement (异步数据搬运)： 利用硬件独立的 DMA 引擎 或 异步拷贝指令，在计算单元 (ALU/Tensor Core) 处理当前数据的同时，后台静默地预取下一块数据。

  • 通用性说明：在 NVIDIA GPU 上映射为 cp.async，在 Ascend NPU 上映射为 DataCopy，在 TPU 上映射为 DMA 指令。

• Double Buffering (双缓冲/乒乓机制)： 一种软件流水线 (Software Pipelining) 策略。分配两块片上缓存 (Buffer A 和 Buffer B)，当计算单元处理 Buffer A 时，DMA 引擎填充 Buffer B，交替进行。

• Bank Conflict Avoidance (存储冲突避免)： 针对采用多体存储 (Banked Memory) 架构的片上缓存，通过 Padding (填充) 或 Swizzling (地址重排) 优化数据布局，防止并行访问冲突。

MLIR 示例

// 这是一个通用的"异步拷贝 + 融合计算"模式
func.func @explicit_memory_hierarchy(%A_global, %B_global, %C_global) {
  // 1. 分配片上高速缓存 (Scratchpad/Shared Memory)
  // 这里的 memory_space 是一个通用属性，不同硬件对应不同层级
  %A_local = memref.alloc() : memref<128x32xf16, #gpu.address_space<workgroup>>
  %B_local = memref.alloc() : memref<128x32xf16, #gpu.address_space<workgroup>>

  // 2. 启动异步 DMA 搬运 (Async DMA Start)
  // 编译器后端会将其映射为特定硬件的 DMA 指令 (如 cp.async 或 dma_copy)
  %token = gpu.device_async_copy %A_global[...] to %A_local[...] 
  
  // ... (此处可以插入不依赖 A_local 的其他独立计算以掩盖延迟) ...

  // 3. 等待数据就绪 (DMA Wait)
  gpu.device_async_wait %token
  gpu.barrier

  // 4. 在高速缓存上执行融合算子 (Compute on Scratchpad)
  // 此时算子全速运行，无主存带宽瓶颈
  %acc = linalg.matmul ins(%A_local, %B_local) ... 
  
  // 5. 结果写回
  memref.store %acc, %C_global[...]
}

---

4.3 内存生命周期优化 (Lifetime Optimization)

4.3.1 Static Memory Planning & Reuse

背景

深度学习模型中，中间 Tensor (Activation) 的生命周期通常较短。内存规划旨在通过活跃度分析 (Liveness Analysis)，让互斥的 Tensor 共享同一块物理内存 (Memory Arena)，从而降低峰值内存 (Peak Memory Usage)。

这在边缘设备 (如手机 NPU、MCU) 上尤为重要，通常在编译期确定所有 Buffer 的偏移量 (Offset)。TeMCO[25]、TorchTune[26]等现代编译器实现了高效的静态内存规划优化。

技术原理

静态内存规划的核心是将时间维度上的互斥转化为空间维度上的共享。其技术实现通常包含三个步骤：

1. 活跃度分析与干涉图构建 (Liveness Analysis & Interference Graph)：

   • 编译器对计算图进行拓扑排序，确定每个张量的定义点 (Def)和最后使用点 (Last Use)，从而得到一个活跃区间 (Live Interval)。
   • 构建干涉图 (Build Interference Graph)：节点代表张量，边代表活跃区间重叠。如果两个节点之间有边 (Edge)，表示它们 "相互干涉" ，不能共享同一块物理内存。

2. 内存分配算法 (Offset Assignment / Tensor Allocation)：

   • 这是一个经典的图着色问题 (Graph Coloring Problem)。目标是用最少的"颜色"(内存偏移量) 为干涉图的所有节点着色，使得相邻节点颜色不同。
   • 由于图着色是 NP-Hard 问题，编译器通常采用 Greedy Best-Fit 算法：按张量大小排序，依次为每个张量分配一个偏移量 (Offset)。分配时，尝试寻找空闲区间 (Free Block) 中最小且满足大小的块。
   • 目标：最小化所有张量所需的总内存池大小 (Total Workspace Size)。

3. Arena 机制 (Single Arena / Workspace)：

   • 为了避免运行时频繁调用操作系统级的 malloc/free(会导致碎片和系统调用开销)，编译器计算出整个图运行所需的峰值内存 (Peak Memory)。
   • 在运行时，仅分配一块巨大的连续内存 (Arena/Workspace)。所有算子的输入输出都通过 Base_Ptr + Offset 进行访问。

MLIR 策略

• Bufferization：在将 Tensor 降级为 Memref 时，使用 BufferDeallocation pass 插入 dealloc，并合并 alloc。
• In-place Bufferization：尽可能复用输入 buffer 作为输出 buffer (如果输入不再被使用)。

// 场景：两个串行的算子 Op1 -> Op2
// Op1 的输出 %buf1 在 Op2 执行完后就不再需要
// Op3 需要一个新的 Buffer，复用 %buf1 的空间

func.func @static_memory_planning(%arena: memref<10MB>) {
  
  // 1. 静态规划结果：Buf1 分配在偏移 0
  %buf1 = memref.view %arena[0] ... : memref<1MB>
  call @op1_produce(%buf1)
  
  // 2. Op2 读取 Buf1
  call @op2_consume(%buf1)
  // --- 此时 %buf1 生命周期结束 ---

  // 3. 内存复用：Buf2 复用偏移 0 的空间
  // 编译器静态分析确认 %buf1 和 %buf2 无干涉 (Interference)
  %buf2 = memref.view %arena[0] ... : memref<1MB>
  call @op3_produce(%buf2)
  
  return
}

4.3.2 Memory-Constrained Operator Splitting

虽然本节在形式上和 4.1 多级分块 很像 (都是切分)，但两者的目标函数截然不同：

4.1 是为了快 (Cache 命中率)，而 4.3.2 是为了活下去 (避免 OOM，内存溢出)

背景

当单个算子 (即使是分块后) 所需的临时空间超过硬件限制 (如 SRAM 大小或 TPU 内存)，或者为了适配特定的内存 Bank 限制，编译器必须将算子 拆分 (Split) 为多次执行。

此技术常用于超大模型训练 (Activation Checkpointing 也是一种变体) 或受限内存嵌入式推理。Activation Checkpointing[31,32,33]可将训练时的内存占用降低到√M（M为原始内存需求），是训练大模型的关键技术。

技术原理

内存受限拆分的核心逻辑是时空置换 (Space-Time Trade-off)：通过增加少量的控制流开销 (时间)，换取峰值内存占用的显著降低 (空间)。其技术实现基于以下机制：

1. 工作集缩减 (Working Set Reduction)：

   • 对于一个大算子 (如 的 Element-wise)，其峰值内存需求是输入+输出的总和。如果这个总和超过了硬件的 片上内存 (SRAM/UB) 容量，算子将无法执行。
   • Strip Mining 将空间的维度 (Spatial Dimension) 转化为时间的维度 (Temporal Dimension)。通过将大循环切分为小循环，系统只需要分配能容纳 单个切片 (Tile) 的内存。
   • 数学效果：峰值内存从 降低到 。

2. 流式执行与流水线 (Streaming Execution & Pipelining)：

   • 拆分后的算子变成了"流式处理"模式：Load Tile -> Compute -> Store Tile。
   • 为了掩盖拆分带来的频繁 I/O 开销，编译器通常结合 Double Buffering (双缓冲) 技术。在计算第 个 Tile 时，DMA 引擎并行搬运第 个 Tile。

3. 算子裂变 (Operator Fission)：

   • 与通常追求的 Fusion 相反，这里有时需要反向操作。如果一个融合算子 (OpA + OpB) 所需的中间 Buffer 太大而无法驻留片上，编译器会选择将其裂变为两个循环，中间结果写回主存 (DRAM)，以牺牲带宽为代价换取内存可行性。

硬件视角:：Ascend NPU 的 UB 适配

在 Ascend NPU 上，这一技术是刚需。

• UB 强约束：Ascend 的 Unified Buffer (UB) 通常只有几百 KB。对于大分辨率图像 (如 4K 图) 或大语言模型 Tensor，绝对无法一次性塞入 UB。
• 自动切分：Ascend 编译器需要根据 UB 的实际大小计算最大可能的 Tile Size，自动插入循环结构。如果计算涉及多个 Tensor，必须保证它们切片后的总和 。

MLIR 示例：Strip Mining (Loop Fission)

假设我们需要处理一个巨大的张量 (如 1GB)，但硬件的片上内存 (SRAM) 仅有 1MB。直接对整个张量执行 linalg.generic 会导致内存溢出 (OOM)。 编译器通过 Strip Mining 将其重写为循环形式，每次只申请一小块缓冲区。

// 原始：大内存需求 (High Peak Memory)
// 需要分配整个 %large_result，可能导致 OOM
func.func @naive_large_op(%in: tensor<1024x1024xf32>) -> tensor<1024x1024xf32> {
  %large_result = linalg.generic ... ins(%in) ... 
  return %large_result
}

// 优化后：受限内存分割 (Low Peak Memory)，峰值内存 (Peak Memory) 降低到 1MB
func.func @split_large_op(%in: tensor<1024x1024xf32>, %out_buf: memref<1024x1024xf32>) {
  %c0 = arith.constant 0 : index
  %c1024 = arith.constant 1024 : index
  %step = arith.constant 32 : index // 拆分步长，由 SRAM 大小决定

  // 将大算子拆分为循环
  scf.for %i = %c0 to %c1024 step %step {
    // 1. Slice: 逻辑切片，不发生物理拷贝
    %sub_in = tensor.extract_slice %in[%i, 0] [32, 1024] [1, 1] 
              : tensor<1024x1024xf32> to tensor<32x1024xf32>

    // 2. Alloc: 分配微小的临时 Buffer (驻留在 SRAM)
    %small_init = tensor.empty() : tensor<32x1024xf32>
    
    // 3. Fused Compute: 在小 Buffer 上执行融合计算
    // 即使是大算子，在这一步也被局部化了，结果存在 SRAM 中
    %sub_res = linalg.generic ... ins(%sub_in) outs(%small_init) ...

    // 4. Copy Back: 将结果从 SRAM 刷回主存 (DRAM)
    // 这里的 bufferization.to_memref 示意将 Tensor 结果写入 %out_buf
    memref.subview %out_buf[%i, 0] ...
    // ... copy %sub_res content to %out_buf subview ...
  }
}

---

参考文献 (References)

多级分块与Tiling (Multi-level Tiling)

[1] M. Wolf, M.S. Lam. "A Data Locality Optimizing Algorithm." ACM PLDI, 1991. [Link]

[2] J. Ragan-Kelley, et al. "Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing." ACM PLDI, 2013. [Link]

[3] A. Adams, et al. "Learning to Optimize Halide with Tree Search and Random Search." arXiv, 2019. [Link]

[4] F. Mullapudi, et al. "Automatically Scheduling Halide Image Processing Pipelines." arXiv, 2016. [Link]

[5] (2023). "Register Tiling for Unstructured Sparsity in Neural Network." 2023. [Link]

内存墙与带宽优化 (Memory Wall & Bandwidth Optimization)

[6] W. Wulf, S. McKee. "Hitting the Memory Wall: Implications of the Obvious." ACM Computer Architecture News, 1995. [Link]

[7] S. Williams, et al. "Roofline: An Insightful Visual Performance Model for Floating-Point Programs and Multicore Architectures." CACM, 2009. [Link]

[8] A. Lockerman, et al. "Livia: Data-Centric Computing Throughout the Memory Hierarchy." ACM ASPLOS, 2020. [Link]

[9] R. Ausavarungnirun, et al. "MASK: Redesigning the GPU Memory Hierarchy to Manage Multiple Data Abstractions." ACM ASPLOS, 2018. [Link]

[10] M. Rhu, et al. "A Locality-Aware Memory Hierarchy for Energy-Efficient GPUs." IEEE MICRO, 2013. [Link]

显式内存管理 (Explicit Memory Management)

[11] (2021). "Compiler-directed Scratchpad Memory Data Transfer." Journal of Supercomputing, 2021. [Link]

[12] Y. Su, et al. "Revitalizing the Forgotten On-Chip DMA." USENIX FAST, 2023. [Link]

[13] A. Eichenberger, et al. "Optimizing Compiler for a CELL Processor." IBM, 2005. [Classic]

[14] (2021). "Compiling Halide Programs to Push-Memory Accelerators." arXiv, 2021. [Link]

[15] (NVIDIA 2020). "Controlling Data Movement on Ampere." NVIDIA Developer Blog, 2020. [Link]

异步数据搬运与双缓冲 (Async Data Movement & Double Buffering)

[16] (2024). "Ascend C Memory Transfer Art: Double Buffer & Pipeline Optimization." 2024. [Link]

[17] (2024). "Data Transfer Optimizations for Host-CPU and Accelerators." arXiv, 2024. [Link]

[18] (Huawei). "How to Use Double Buffer for Performance Optimization." Huawei Developer. [Link]

[19] (2024). "Fastmove: A Comprehensive Study of On-Chip DMA." ACM, 2024. [Link]

[20] (2025). "Pipeline Parallelism - Large Model Seminar." 2025. [Link]

静态内存规划 (Static Memory Planning)

[21] (ICPP 2024). "TeMCO: Tensor Memory Compiler Optimization." ACM ICPP, 2024. [Link]

[22] (PyTorch 2025). "TorchTune Performance Optimization." PyTorch Blog, 2025. [Link]

[23] (LLVM 2026). "LLVM Operator Fusion Optimization Improvement." LLVM Discourse, 2026. [Discussion]

[24] (ACM). "A Generalized Algorithm for Graph-Coloring Register Allocation." ACM. [Classic]

[25] (UW-Madison). "Register Allocation via Hierarchical Graph Coloring." CS Dept. [Link]

内存受限算子分割 (Memory-Constrained Operator Splitting)

[26] Z. Jiang, et al. "Optimal Sharded Data Parallel for Distributed Deep Learning." arXiv, 2022. [Link]

[27] Y. Xu, et al. "Training of Deep Learning Pipelines on Memory-Constrained GPUs." PMC, 2022. [Link]

[28] (Medium 2024). "PyTorch Activation Checkpointing Complete Guide." 2024. [Link]

[29] (ACM). "Hierarchical Memory-Constrained Operator Scheduling." ACM. [Link]

[30] (Blog 2024). "Train Large ML Models with Activation." 2024. [Link]

编译器框架与工具 (Compiler Frameworks & Tools)

[31] T. Chen, et al. "TVM: An Automated End-to-End Optimizing Compiler for Deep Learning." USENIX OSDI, 2018. [Link]

[32] (Oxford). "MLIR: An Optimizing Compiler Framework." Oxford Technical Report. [Link]

[33] (2024). "A Survey of General-purpose Polyhedral Compilers." ACM, 2024. [Link]

[34] (LLVM). "Polly - Polyhedral Optimizations for LLVM." LLVM Project. [Official]

综合参考资料

[35] (Wikipedia). "Memory Wall (Computer Architecture)." Wikipedia. [Reference]

[36] (Berkeley). "Roofline Model Tutorial." UC Berkeley. [Tutorial]