算子开发中的 Tile 策略：基于 L1/L2 缓存的访存优化

前言

在深度学习推理和训练的算子开发中，如何高效地利用（Ascend）AI 处理器上的片上存储资源，特别是 L1 和 L2 缓存，是决定算子性能的关键因素之一。架构（如 Atlas 900/800 系列）的内存层次结构设计独特，其强大的计算能力（TFLOPS）往往受到访存带宽的限制。因此，精妙的“Tile 策略”成为实现高性能算子设计的核心技术。

本文将以 CANN 社区开源项目（位于 https://atomgit.com/cann）为基础，特别是 ops-nn 仓库中的实现，深入探讨算子开发中如何运用 Tile 策略，结合 L1/L2 缓存机制，实现访存优化。

核心技术原理：内存层次与 Tile 策略

AI 处理器上的内存层次结构通常包括：

1. 全局内存 (GMEM/DDR)：容量大，带宽相对较低，延迟高。
2. L2 缓存 (Local Data Memory, LDM)：容量适中，用于存储更大规模的中间结果或数据块，其访问速度远快于 DDR。
3. L1 缓存 (Register File/Shared Memory)：容量最小，但访问速度最快，通常用于存储单次计算所需的最热数据。

Tile 策略的核心思想是将大规模的张量运算分解为一系列更小的、可以在 L1 或 L2 缓存中完全容纳的子任务（Tiles）。通过合理规划 Tile 的大小和数据加载顺序，我们可以最大化数据重用率，减少对慢速全局内存的访问次数，从而实现访存优化。

这种策略遵循著名的 “数据局部性原理”：尽可能在数据被加载到高速缓存后，立即完成对该数据的所有计算，避免数据被其他计算任务替换出缓存。

代码/架构分析：以 ops-nn 仓库为例

在 CENN 算子开发中，Tile 策略的实现通常体现在 TBE (Tensor Builder Environment) 编程模型或更底层的 AI Core 编程中。我们重点关注 ops-nn 仓库（https://atomgit.com/cann/ops-nn），该仓库包含了大量标准算子的参考实现。

1. L2 缓存（LDM）的使用

对于需要较大工作集（如矩阵乘法中的权重矩阵或特征图的局部区域）的算子，L2 缓存是首选的暂存区。

在实现中，Tile 策略通常表现为：

• 分块加载 (Blocking Load)：将大型矩阵或张量沿着维度进行分块。例如，在矩阵乘法（GEMM）中，我们会将输入矩阵 A 和 B 划分成 L2 缓存大小的块。
• 循环嵌套优化：在 TBE 代码中，我们通常会设计三层或四层循环：
  1. 外层循环：迭代 L2 Tile 的索引。
  2. 中层循环：控制 L2 Tile 内部的数据加载和计算。
  3. 内层循环：执行核心的 SIMD/向量计算。

通过这种方式，我们可以确保在 L2 缓存中，数据 A 的一个 Tile 和数据 B 的一个 Tile 可以被重复利用多次，直到计算完成该 Tile 对应的输出 Tile。

2. L1 缓存（寄存器/共享内存）的利用

L1 缓存的优化更加精细，它直接关系到 AI Core 的计算单元的效率。在 AI 处理器上，L1 缓存通常被映射为寄存器文件或片上共享内存。

• 向量化与寄存器分配：算子实现需要确保数据被加载到寄存器后，能够被充分利用。例如，在执行点积或融合操作时，Tile 的大小需要精确匹配 AI Core 的向量处理单元宽度，以实现最大程度的并行和数据重用。
• 数据移动指令：开发者需要关注 TBE 中用于数据在 L2 和 L1/寄存器之间移动的指令（如 vector_load, vector_store）。Tile 策略的成功与否，直接决定了这些指令的执行频率和效率。如果 Tile 过大，会导致 L1 缓存的频繁换入换出（Cache Thrashing）；如果 Tile 过小，则无法充分利用 L1 的局部性。

性能优化实践：Tile 大小的确定

Tile 策略的优化是一个平衡 L1/L2 容量、计算量和访存带宽的过程。

1. 确定 L2 Tile 大小

L2 Tile 的大小主要受限于 L2 缓存的总容量和算子的内存访问模式。对于 GEMM 算子，一个常见的策略是：

$$\text{TileSize}\approx \sqrt{\text{CacheCapacity}}$$

我们需要预留一部分 L2 空间给输入 A、输入 B 以及输出 C 的 Tile。一个经验法则是，L2 Tile 应该足够大，以容纳所有参与一次核心计算（如一组矩阵乘法）所需的数据，同时保证数据可以长时间停留在 L2 中。如果 L2 Tile 太小，则需要频繁地从 DDR 重新加载数据，性能将急剧下降。

2. 确定 L1 Tile 大小

L1 Tile 的大小则直接受限于 AI Core 的寄存器资源和向量宽度。

• 充分利用计算单元：L1 Tile 的设计必须保证在执行计算时，所有必要的操作数都能同时加载到寄存器中，以实现最大吞吐量。例如，如果 AI Core 支持 128 元素的向量操作，那么 L1 Tile 的数据量应能被 128 整除，并尽可能填充所有可用的计算单元。
• 避免溢出：L1 Tile 的数据量必须远小于寄存器文件的大小，以避免数据溢出到更慢的存储层级。

3. 访存模式的优化

Tile 策略不仅关乎大小，还关乎顺序。例如，在矩阵乘法中，我们通常采用 “平铺和分块” 的策略，以确保在加载 A 的一个行 Tile 时，能同时加载 B 的多个列 Tile，从而实现高重叠度的计算和访存。这要求算子开发者深入理解 AI Core 的访存单元的工作方式，确保数据以连续的、符合硬件预取机制的方式被访问。

总结

算子开发中的 Tile 策略是实现 AI 处理器高性能的关键。它本质上是对内存层次结构的一次精细化调度，旨在最大化 L1/L2 缓存的命中率和数据重用率。通过对 ops-nn 等开源仓库中算子实现的分析，我们可以看到，成功的 Tile 策略需要精确计算 L1 和 L2 的容量限制，并设计出与 AI Core 计算模式完美匹配的循环结构和数据加载顺序。掌握 L1/L2 缓存的访存优化，是每一位资深 CANN 算子架构师必须具备的核心能力。