CATLASS Tiling 模板化设计在算子开发中的应用

我是Fanstuck，致力于将复杂的技术知识以易懂的方式传递给读者，热衷于分享最新的行业动向和技术趋势。如果你对大模型的创新应用、AI技术发展以及实际落地实践感兴趣的话，敬请关注。

背景简介

在高性能算子开发中，Tiling 模板化是一种通过将计算任务分块并参数化的方式，提高算子在不同数据形状和硬件上的性能的方法。以Ascend CANN 平台为例，算子的性能高度依赖于计算块大小（tile size）与内存布局的匹配程度。

当 tile 大小恰好适配硬件的缓存（如 L1、L0 缓存）和运算阵列（如 Cube 单元）的特性时，才能发挥出最佳性能。因此，如果能将这些 tiling 策略抽象为模板，让开发者方便地针对不同问题实例化合适的 tile 参数，将大大提升开发效率和性能可移植性。

传统算子库&CATLASS 算子库

传统的算子库往往为某一类算子预先做通用优化，希望在大多数场景下取得不错的开箱即用性能。然而，这种泛化优化可能无法针对特殊输入形状做到性能最优。相比之下，Ascend 平台提供的 CATLASS 算子模板库采取了不同思路：它针对 GEMM 类算子提供了一系列模板化样例，允许开发者根据不同输入场景快速定制出高性能实现。

换言之，CATLASS 模板库并不追求一个模板适配所有场景的“普适”最佳性能，而是提供多种模板，让开发者按需选择，实现“一次编写，多处复用”。例如，对于矩阵乘法（matmul），模板库提供了 basic-matmul、optimized-matmul、splitk-matmul、padding-splitk-matmul 等不同版本，分别适用于不同尺寸或格式的矩阵。

这种模板化设计为算子开发者、编译器和模板库开发人员提供了灵活性：面对新的形状或硬件，只需调整模板参数即可生成高效代码，而无需完全从头手工优化。

模板设计理念

Tiling 模板化设计的核心是在算法实现中引入参数化的“分块”（tile）策略。开发者预先设计好算子的计算逻辑，将计算分成大小可调的块，之后通过模板参数来控制块的尺寸、数据排布和调度策略。

这样，不同的问题规模、张量形状甚至不同硬件，都可以通过调整这些模板参数来适配。举例来说，在矩阵乘法、卷积等典型算子中，我们可以将矩阵的 M、N、K 三个维度上的计算切分为小块（tiles），并以这些 tile 的尺寸作为模板参数；同时，将矩阵的数据布局（如行优先Row Major、列优先Column Major或其他分块格式如 zN）也设计为可配置的模板参数。

下面是一个小规模（5个tile，每个tile在K维上2个block）的切分示意图：

在 Ascend 平台上，可以借助诸如 PyPTO 或 TileLang 这样的工具或 DSL（领域专用语言）来定义泛化的计算图模板。这些工具允许开发者使用高级的 Python 或声明式语法描述计算的 tile 划分和调度。

例如，可以用几行代码描述“将矩阵 A 切分为 L1 大小的块，每块 MxK，矩阵 B 切分为 KxN 的块，然后依次加载到片上缓存，进行矩阵乘加运算”的过程。模板设计时，会将块尺寸、数据布局等作为参数，使得同一份计算逻辑可以针对不同硬件资源约束进行调整。一方面，要确保 tile 大小不超过硬件资源限制（例如 Ascend910 的 L1 缓存 512KB、各级 L0 缓存大小等），并满足硬件对矩阵运算的对齐要求（如长度为16的倍数）。

• 场景一：fp16输入输出，L1TileShape<128,256,256>，L0TileShape<128,256,64>

对于输入输出是fp16场景，为了计算精度，cube核进行mmad输出到L0C上的结果是fp32数据类型，fixpipe随路传回gm时cast到fp16数据类型。

 L1大小512K
 L1实际占用 = L1::M * L1::K * 2(Byte) * 2(doubleBuffer) + L1::K * L1::N * 2(Byte) * 2(doubleBuffer)
         = 128 * 256 * 2 * 2 + 256 * 256 * 2 * 2
         = 393216 B = 384 KB = 3/4 L1_SIZE
 ​
 L0A大小64K
 L0A实际占用 = L0::M * L0::K * 2(Byte) * 2(doubleBuffer)
         = 128 * 64 * 2 * 2
         = 32768 B = 32 KB = 1/2 L0A_SIZE
 ​
 L0B大小64K
 L0B实际占用 = L0::K * L0::N * 2(Byte) * 2(doubleBuffer)
         = 64 * 256 * 2 * 2
         = 65536 B = 64 KB = 1 L0B_SIZE
 ​
 L0C大小128K
 L0C实际占用 = L0::M * L0::N * 4(Byte)
         = 128 * 256 * 4
         = 131072 B = 128 KB = 1 L0C_SIZE

另一方面，也需要考虑不同硬件单元的特点：比如 Ascend 的 Cube 数矩阵乘单元要求16×16的矩阵片作为基本运算单元，那么 tile 尺寸通常取16的倍数以充分利用计算阵列。

模板的设计目标是在保持灵活性的同时，尽量覆盖常见的算子模式。以矩阵乘法的模板为例，设计者通常提供一组参数默认值，适合“大多数”典型场景，同时允许高级用户根据需要调参。CATLASS 模板库的设计理念体现了这一点：它提供的 basic_matmul、optimized_matmul 等范例各自针对不同场景做了权衡。

例如 basic_matmul 使用Ping-Pong 双缓冲策略，在 L1 和 L0 缓存中各准备两块 tile，计算与数据搬运流水并行；optimized_matmul 则进一步启用了预加载（Preload）和K轴乱序（ShuffleK）等优化，减少流水中断、降低Bank冲突；

 template <bool ENABLE_UNIT_FLAG_ = false, bool ENABLE_SHUFFLE_K_ = false>
 struct MmadAtlasA2Preload : public MmadAtlasA2 {
     static constexpr uint32_t STAGES = 2;
     static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
     static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_;
 };

padding_matmul 针对某些内存对齐要求做了额外填充；splitk_matmul 模板允许将计算在 K 轴方向拆分到多个核心上并最终汇总结果。这些模板背后的理念是：“一次编写，按需组合”：将不同优化策略作为模块化组件嵌入模板，根据算子需求进行组合，生成贴合特定场景的实现。

需要强调的是，模板设计并不局限于固定硬件。一个好的模板应适配不同硬件约束并容易移植。例如，当换到另一款Ascend芯片，L0/L1缓存大小变化时，只需调整 tile 大小参数，就能让同一模板继续高效运行。这种设计大大提高了算子性能的可移植性：开发者无需为每种硬件各写一套内核，只要调优模板参数即可。

关键路径解读

为了说明模板化算子从定义到执行的流程，我们以矩阵乘法算子（GEMM）的一个具体模板——BasicMatmul——为例，解析其关键路径。BasicMatmul 是 CATLASS 模板库中最基础的矩阵乘内核，它体现了模板生成代码的典型结构。开发者通过组合模板库提供的组件来定义这个算子：

• 矩阵块定义：BasicMatmul 将计算拆分为矩阵 C 的小块。假设我们选定的 L1 级 tile 尺寸为 <128, 256, 256>（对应 M=128, N=256, K=256），L0 级 tile 尺寸为 <128, 256, 64>。这意味着每个计算核心（ASCEND 的 AI Core）处理 C 矩阵的一个子块，其大小最多为 128×256。对输入矩阵 A、B，会根据这个 tile 大小进行分块加载。例如，上述参数下，矩阵 A 每次加载 128×256 的子矩阵到 L1，矩阵 B 加载 256×256 的子矩阵到 L1。因为开启了双缓冲（Ping-Pong），每个 AI Core 在 L1/L0A/L0B 都保留两块 tile 空间，这样可以在计算当前 tile 的同时预取下一 tile 数据。

任务分配（Swizzle 调度）：CATLASS 使用 Swizzle 策略来决定将哪些 tile 分配给哪个核心计算。。以 C 矩阵的 M×N 方向划分的基本块为单位，Swizzle 策略按照一定规则编号这些块并映射到硬件的多个核心上并行执行。比如 BasicMatmul 默认采用 <3,0> 的 swizzle（GemmIdentityBlockSwizzle<3,0>），这代表一种特定的遍历顺序。其效果是当矩阵维度 M 大于 N 时，每个核心尽量处理不同行块，减少多个核心竞争同一行的数据。

 template <bool ENABLE_UNIT_FLAG_ = false>
 struct MmadAtlasA2Pingpong : public MmadAtlasA2  {
     static constexpr uint32_t STAGES = 2;
     static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
 };

相应地，如果 M < N，则可以切换为 <3,1> 模式，优化列方向的并行。SwizzleOffset 则影响任务在核心间分配的平衡性。例如，在一个案例中，使用 swizzle <3,1> 导致某些核心承担了多于其他核心的 tile 数，负载不均；而改用 <4,1> 后，各核心分工更加均衡，性能从 40.6µs 提升到 35.3µs。这说明合理的任务调度对流水线饱满度和整体性能有直接影响。

Kernel 主体执行：BasicMatmul 的内核由模板库提供的模块拼装而成。主要包括 BlockMmad（矩阵乘累加模块）、BlockScheduler（如前述 Swizzle）、以及一些可选的 BlockEpilogue（后处理，如无则省略）。

当执行开始时，每个核心按照调度分配开始处理自己对应的 C块：通过 DMA 将 A、B 相应 tile 从全局内存（GM）搬运到片上 L1，然后从 L1 分发到各自的 L0A（存放A片段）和 L0B（存放B片段）缓存中。

在 Ascend 架构中，每个 AI Core 的 L0A/L0B/L0C 分别用于存储矩阵片 A、B 和计算结果C。BasicMatmul 利用双缓冲机制：当第一个 tile的数据在 L0 上计算时，下一个 tile的数据在另一缓冲区并行加载。这样，计算（Cube上进行矩阵乘加）与数据搬运可以重叠，尽量避免等待。

流水线与计算图展开：模板展开的计算图本质上和手工优化的代码逻辑类似：它会显式地展开成一系列循环与内存拷贝操作。

以 K 维度为例，因为 L0A/L0B 容纳的 K 阶长度有限（如上例 L0TileShape::K=64），一次计算只能处理 K=64 的片段，所以需要循环迭代 K-dimension。

模板生成的代码会在内部产生一个for (k_tile = 0; k_tile < K; k_tile += 64)这样的循环，每次在 Cube 上执行一个 128×256×64 的矩阵乘累加（利用 Cube 16×16 阵列一次计算16个结果并累加）。这些循环和调度逻辑都是由模板库根据参数自动生成的，而开发者无需手写这些底层细节。

值得一提的是，类似的模板化思想也可应用在其他复杂算子上。例如在 Lightning Indexer（一种稀疏注意力机制中的索引计算）中，涉及对大规模索引分数的并行计算筛选。

综上所述，Tiling 模板化设计在 Ascend 算子开发中展现出了强大的生命力。它让高性能算子开发从“手工艺”走向了“工业化”，既保证了性能，又提高了可复用性和开发效率。当然，这并不意味着人工优化将完全消失——相反，模板的进步凝聚了无数人工优化的智慧。

CATLASS 以及类似的模板化框架正在朝这个方向迈进，也许在不久的将来，“模板化算子开发”将成为高性能计算领域的主流范式，为工程师解除重复优化的重担，释放出更多精力投入创新。