前言

在深度学习推理与训练场景中，矩阵乘法（GEMM）是最基础也最核心的计算原语之一。几乎所有主流模型——从卷积神经网络到 Transformer、再到推荐系统中的 embedding 查找——其底层都离不开大规模浮点或量化矩阵运算的支撑。昇腾 CANN（Compute Architecture for Neural Networks）是华为面向昇腾系列 AI 处理器推出的统一异构计算架构，它通过 Ascend C 编程原语将硬件的张量核心、向量引擎与存储层次暴露给开发者，使我们能够在昇腾 NPU 上高效实现各类算子。

然而，直接基于 Ascend C 手写一个高性能 GEMM 算子，需要开发者深入理解硬件的数据流、存储层级以及指令调度细节，学习成本极高。catlass 正是为此而生：它是一个受 NVIDIA CUTLASS 启发、专为昇腾 NPU 设计的 GEMM 模板库。通过精巧的模板元编程，catlass 将硬件细节封装为可组合的策略层，让开发者只需选择合适的配置参数，便能组装出性能接近理论上限的自定义 GEMM 算子。本文将系统讲解 catlass 的设计理念、模板架构，以及如何一步步用模板组装出满足实际需求的 GEMM 算子。

一、catlass 的项目定位

1.1 从 CUTLASS 到 catlass 的演进脉络

要理解 catlass，首先需要了解它的灵感来源 CUTLASS。NVIDIA 的 CUTLASS 是一个久经考验的开源 GEMM 模板库，它定义了 GEMM 计算的抽象层次：ThreadBlock（线程块）级别的分块策略、Warp（线程束）级别的矩阵分块、以及 Epilogue（后处理）链式激活函数。CUTLASS 通过模板特化让同一套代码骨架既能生成 FP16、BF16、INT8 等多种数据类型的 GEMM 算子，又能在 Volta、Turing、Ampere 等不同架构上高效执行。

catlass 的核心思路与 CUTLASS 一脉相承，但它面向的是昇腾 CANN 的硬件特性。昇腾 NPU 的张量核心（Cube Unit）以不同于 NVIDIA Tensor Core 的方式组织数据通道和存储层级：它采用分块式数据加载、流水线式的指令调度，以及独特的 Local Memory 层次结构。因此，catlass 在保留 CUTLASS 模板哲学的同时，对 TilePolicy、ThreadMap 和 Epilogue 三大核心组件做了面向昇腾硬件的重设计。

1.2 catlass 与 CUTLASS 的架构差异

虽然 catlass 受到 CUTLASS 启发，但二者在架构设计上有几个关键差异：

第一，数据布局抽象不同。CUTLASS 面向 NVIDIA 的行主序（RowMajor）和列主序（ColMajor）布局，并假设 Shared Memory 的访问模式与 CUDA Warp 保持一致。昇腾 NPU 的数据布局以 NC1HW1（Channel-First）为主，且其 Local Memory 的带宽特性与 CUDA Shared Memory 存在显著差异。catlass 因此引入了一套独立于 CUTLASS 的 ThreadMap 体系，专门描述昇腾 NPU 上 Warp 内数据的组织方式。

第二，指令流水线的组织方式不同。NVIDIA 的 Tensor Core 指令（HMMA/IMMA）在 Warp 级别发射，Warp 内所有线程协同完成一次矩阵分块乘。昇腾 NPU 的 Cube Unit 以更粗粒度的"指令流"方式工作，catlass 将 Mainloop 抽象为一套流水线阶段（Stage）的组合，支持多级 Double Buffer 和指令级并行。

第三，后处理链的扩展机制不同。CUTLASS 的 Epilogue 通过模板链（Chained Epilogue）添加 Bias、ReLU、LeakyReLU 等激活操作。catlass 延续了这一设计，但针对昇腾 NPU 的向量引擎能力，增加了更多融合模式——例如在 Epilogue 阶段直接融合 Softmax 对角线操作或 LayerNorm 归一化。

1.3 模板元编程的设计思路

catlass 的模板元编程遵循一个核心原则：策略即类型。每一种硬件层面的选择——分块大小、数据布局、流水线深度、后处理方式——都被编码为一个模板参数或一个模板特化。开发者通过组合这些参数，就能在编译期生成专门针对当前硬件和数据形状的最优算子代码。

具体而言，catlass 使用 C++ 模板偏特化（Partial Specialization）来实现策略选择。每一层模板都定义了一组默认行为，而开发者可以通过特化某个层来替换实现细节。例如，若默认的 TilePolicy 对某种矩阵形状不够高效，开发者可以提供自定义的 TilePolicy 特化，而不必改动 Mainloop、Epilogue 等其他层。这种"插入式替换"的机制，使得 catlass 在保持高可维护性的同时，也具备足够的灵活性来适配各种复杂的业务场景。

二、catlass 模板体系架构

catlass 的模板体系由四个核心层次构成：TilePolicy、ThreadMap、Mainloop 和 Epilogue。这四层各司其职，共同完成一个 GEMM 算子的生成。下面对每一层的职责进行详细分析。

2.1 TilePolicy——分块策略层

TilePolicy 是 catlass 模板体系的最底层，也是最接近硬件的层。它定义了 GEMM 计算中"线程块"级别的分块策略，具体包括：

• BlockTile（块级分块）：整个 GEMM 被划分成多大尺寸的块，每一块由一个线程块（对应昇腾 NPU 上的一个计算单元）负责处理。BlockTile 的选择直接影响数据在 Global Memory 和 Local Memory 之间的搬运效率。通常，较大的 BlockTile 能减少分块开销，但会增加寄存器压力和 Local Memory 占用。
• WarpTile（线程束级分块）：每个线程块内部，数据如何进一步划分为 Warp 级子块。WarpTile 的设计需要与昇腾 NPU 的向量引擎数据路径匹配。
• 寄存器级分块（ThreadTile）：单个线程或向量通道内部，数据如何在寄存器级别被组织。ThreadTile 的设计决定了每个向量单元在一次指令周期内能处理多少元素。

TilePolicy 本身是一个模板类，它通过静态常量定义了上述各级的尺寸参数。catlass 预置了多种 TilePolicy 特化，分别针对不同的矩阵形状和数据类型进行了优化。

2.2 ThreadMap——线程数据映射层

ThreadMap 负责将逻辑矩阵元素映射到具体的硬件线程或向量通道。它回答的问题是：矩阵中的每个元素应该由哪条向量通道负责加载、计算和写回？

ThreadMap 的设计在 catlass 中尤为关键，因为昇腾 NPU 的向量引擎以固定宽度的向量指令（Vector Unit）作为执行单元。一个合理的 ThreadMap 需要确保：

• 同一 Warp 内的不同向量通道访问的数据在 Global Memory 中是连续的，以充分利用 DMA（Direct Memory Access）引擎的合并访问（Coalesced Access）能力。
• 数据在 Local Memory 中的布局与向量指令的读取粒度对齐，避免产生跨向量通道的数据依赖。

catlass 的 ThreadMap 模板支持一维、二维和三维映射模式。一维映射适合向量点积类操作，二维映射适合矩阵分块操作，三维映射则为 BatchGEMM 等场景提供了更灵活的数据组织方式。

2.3 Mainloop——计算主循环层

Mainloop 是 GEMM 算子的核心计算逻辑所在。它负责将数据从 Global Memory 加载到 Local Memory，在 Local Memory 中执行矩阵乘的核心计算，并将结果写回 Global Memory。

Mainloop 的模板参数包括：

• 分块迭代次数：沿 M、N、K 三个维度分别迭代多少次。
• Double Buffer 策略：是否启用双缓冲（在计算当前块的同时预加载下一块数据）。
• 流水线级数（Pipeline Stage）：数据加载与计算重叠的流水线深度。

Mainloop 的执行流程可以概括为"加载→转换→计算→写回"四步的循环重复。在 catlass 中，这一流程通过模板参数控制循环展开（Loop Unrolling）和向量化（Vectorization）程度，从而在编译期生成高度优化的指令序列。

2.4 Epilogue——后处理链层

Epilogue 是 GEMM 计算之后的后处理链，用于添加激活函数、偏置加法、量化操作等。catlass 的 Epilogue 采用了链式设计：每一个后处理步骤都是一个独立的 Epilogue Stage，多个 Stage 可以串联起来，在一次算子调用中完成从原始矩阵乘结果到最终输出的全链路处理。

catlass 支持的 Epilogue Stage 包括但不限于：

• BiasAddStage：在结果矩阵的每一行或每一个通道上添加一个偏置向量。
• ActivationStage：支持 ReLU、LeakyReLU、Sigmoid、Tanh 等常用激活函数。
• ScaleStage：对结果进行缩放操作，常用于混合精度推理中的反量化步骤。
• ElementwiseStage：支持加减乘除等逐元素运算。

这种链式设计的优势在于：所有后处理操作都可以融合到 GEMM 的写回阶段进行，无需额外的中间存储，从而显著减少内存带宽压力。

三、用模板组装 GEMM 的完整流程

本节通过一个完整的示例，演示如何使用 catlass 从零开始组装一个 GEMM 算子。

3.1 第一步：选择合适的 TilePolicy

TilePolicy 的选择需要根据矩阵形状和数据类型来决定。catlass 预置了以下几种常用 TilePolicy：

#include "catlass/cutlass_utils.h"

// 标准 FP32 GEMM，针对中等矩阵形状
using StandardFP32Tile = catlass::tile::GemmBatchedConfig<
    catlass::tile::BlockTile<128, 128, 32>,   // Block 级分块
    catlass::tile::WarpTile<64, 64, 32>,       // Warp 级分块
    catlass::tile::ThreadTile<8, 8, 8>        // 线程寄存器分块
>;

对于大矩阵场景，可以选择更大的 BlockTile 以提高并行度：

// 大矩阵优化 TilePolicy
using LargeMatrixTile = catlass::tile::GemmBatchedConfig<
    catlass::tile::BlockTile<256, 256, 64>,
    catlass::tile::WarpTile<64, 64, 64>,
    catlass::tile::ThreadTile<8, 16, 8>
>;

3.2 第二步：配置 ThreadMap

ThreadMap 决定了数据在向量通道间的分布方式。以下是一个典型的二维 ThreadMap 配置：

#include "catlass/thread_map.h"

// 二维 ThreadMap：沿 M 方向 16 通道，沿 N 方向 8 通道
using ThreadMap2D = catlass::threadmap::PitchLinear<
    catlass::threadmap::LayoutShape<16, 8>,
    catlass::threadmap::ClusterShape<1, 1>
>;

ThreadMap 的配置需要与 TilePolicy 保持一致：ThreadTile 的总元素数应当等于 ThreadMap 各维度通道数的乘积，否则会导致数据映射错误。

3.3 第三步：配置 Mainloop

Mainloop 的配置涉及数据加载路径、计算路径和流水线策略。以下是一个启用 Double Buffer 的 Mainloop 配置：

#include "catlass/mainloop.h"

// 启用 Double Buffer 的 Mainloop 配置
template <typename TilePolicy_, typename ThreadMap_>
struct MainloopConfig {
    using TilePolicy = TilePolicy_;
    using ThreadMap = ThreadMap_;
    static constexpr int kStages = 2;              // 流水线级数（双缓冲需要至少 2 级）
    static constexpr bool kDoubleBuffer = true;    // 启用 Double Buffer
    static constexpr int kIterationsK = 4;        // K 方向的迭代次数
    using LoadIteratorA = catlass::load::GlobalIterator<
        TilePolicy_, ThreadMap_, half_t, 0>;        // 矩阵 A 的加载器
    using LoadIteratorB = catlass::load::GlobalIterator<
        TilePolicy_, ThreadMap_, half_t, 1>;        // 矩阵 B 的加载器
    using ComputeCore = catlass::compute::CubeCore<
        TilePolicy_, half_t, half_t, float>;       // Cube 计算核
};

3.4 第四步：添加 Epilogue 后处理

如果需要将 GEMM 结果直接用于下游推理，可以在 Epilogue 链中添加 Bias 和 ReLU：

#include "catlass/epilogue.h"

// 链式 Epilogue：BiasAdd → ReLU
using EpilogueChain = catlass::epilogue::Chain<
    catlass::epilogue::BiasAddStage<float, 1>,      // 添加偏置
    catlass::epilogue::ActivationStage<
        catlass::epilogue::ReLU, float>,           // ReLU 激活
    catlass::epilogue::ScaleStage<float, 0.125f>   // 量化缩放
>;

3.5 第五步：实例化并调用算子

将以上各层组合在一起，通过 catlass 的主入口完成算子实例化：

#include "catlass/gemm.h"

// 最终 GEMM 算子类型
using GemmOperator = catlass::Gemm<
    StandardFP32Tile,
    ThreadMap2D,
    MainloopConfig<StandardFP32Tile, ThreadMap2D>,
    EpilogueChain,
    float,    // 输出数据类型
    half_t    // 输入数据类型（FP16 加速）
>;

// 算子调用示例
void run_custom_gemm(float* C, const half_t* A, const half_t* B,
                     const float* bias, int M, int N, int K,
                     void* workspace, aclrtStream stream) {
    GemmOperator gemm_op;
    
    GemmOperator::Arguments args;
    args.M = M;
    args.N = N;
    args.K = K;
    args.A = A;
    args.B = B;
    args.C = C;
    args.bias = bias;
    args.alpha = 1.0f;
    args.beta = 0.0f;
    
    gemm_op.initialize(args, workspace);
    gemm_op.run(stream);
}

四、常用 GEMM 变体的模板配置

在实际项目中，Gemm 的变体远不止最基础的 C = A * B 这一种。以下分别介绍几种常见变体的模板配置差异。

4.1 标准 GEMM

标准 GEMM 是最简单的变体，执行 C = alpha * A * B + beta * C 操作：

// 标准 GEMM 实例化
using StdGemm = catlass::Gemm<
    StandardFP32Tile,
    ThreadMap2D,
    MainloopConfig<StandardFP32Tile, ThreadMap2D>,
    catlass::epilogue::Empty,    // 无后处理
    float,
    half_t
>;

4.2 BatchGEMM

BatchGEMM 用于处理一组矩阵乘的批量计算，典型应用场景包括推荐系统的多 embedding 表查找或多任务学习的并行前向传播：

#include "catlass/batched_gemm.h"

// BatchGEMM 配置：批量大小 32
using BatchGemmConfig = catlass::BatchedGemm<
    catlass::tile::BlockTile<128, 128, 32>,
    catlass::tile::WarpTile<64, 64, 32>,
    catlass::tile::ThreadTile<8, 8, 8>,
    ThreadMap2D,
    catlass::epilogue::Empty,
    int,    // Batch 维度索引类型
    32      // 固定批量大小
>;

void run_batch_gemm(float* C_list[], const half_t* A_list[],
                    const half_t* B_list[], int batch_size,
                    int M, int N, int K, void* ws, aclrtStream s) {
    BatchGemmConfig batch_op;
    // 每个 batch 的 A/B/C 指针传入列表
    BatchGemmConfig::Arguments args{batch_size, M, N, K,
                                    A_list, B_list, C_list};
    batch_op.initialize(args, ws);
    batch_op.run(s);
}

4.3 GEMM + Bias

在推理场景中，BiasAdd 是卷积层后的常规操作。catlass 提供了专门的 BiasAddStage：

// GEMM + Bias：每行输出加上一个偏置向量
using GemmBias = catlass::Gemm<
    StandardFP32Tile,
    ThreadMap2D,
    MainloopConfig<StandardFP32Tile, ThreadMap2D>,
    catlass::epilogue::Chain<
        catlass::epilogue::BiasAddStage<float, 1>,  // axis=1（沿行方向广播）
        catlass::epilogue::Empty
    >,
    float,
    half_t
>;

4.4 GEMM + ReLU

在激活函数中，ReLU 是最常见的选择。将 ReLU 直接融合到 GEMM 的 Epilogue 中，可以省去额外的激活算子调用：

// GEMM + ReLU 融合
using GemmRelu = catlass::Gemm<
    StandardFP32Tile,
    ThreadMap2D,
    MainloopConfig<StandardFP32Tile, ThreadMap2D>,
    catlass::epilogue::Chain<
        catlass::epilogue::BiasAddStage<float, 1>,
        catlass::epilogue::ActivationStage<
            catlass::epilogue::ReLU, float>
    >,
    float,
    half_t
>;

4.5 各变体配置差异总结

变体类型	Epilogue 链	额外参数	典型应用
标准 GEMM	Empty	alpha, beta	Transformer FFN 层
BatchGEMM	Empty / Bias	batch_size, 指针数组	推荐系统 embedding
GEMM+Bias	BiasAddStage	bias 指针	卷积后处理
GEMM+ReLU	BiasAdd → ReLU	bias 指针	激活函数融合
GEMM+LeakyReLU	BiasAdd → LeakyReLU	negative_slope, bias	GAN 生成器

五、性能调优的模板参数

使用 catlass 模板组装 GEMM 算子后，性能调优的核心在于正确配置三个关键参数：Tile 大小、Double Buffer 策略和流水线级数。本节详细分析这些参数对性能的影响。

5.1 Tile 大小的调优

Tile 大小决定了每个计算单元处理的数据量，直接影响 Global Memory 访问和计算密度的平衡。

BlockTile 的选择原则：当矩阵较大（MB × NB ≥ 1024 × 1024）时，选择较大的 BlockTile（如 256×256）可以减少线程块间的同步开销和分块碎片。当矩阵较小（如 MB × NB ≤ 128 × 128）时，过大的 BlockTile 会导致并行度不足，应选择较小的 BlockTile（如 64×64 或 64×128）。

WarpTile 的选择原则：WarpTile 应与昇腾 NPU 的向量引擎数据路径宽度匹配。对于昇腾 910 系列，建议 WarpTile 的 N 维度为 64 的倍数，以对齐 Cube Unit 的输出格式。

ThreadTile 的选择原则：ThreadTile 决定了每个向量通道处理的元素粒度。过小的 ThreadTile 会导致指令发射开销占比过高；过大的 ThreadTile 会导致寄存器溢出到 Local Memory，反而降低性能。

以下是一个 Tile 大小对比的量化示例：

// 不同 Tile 配置的性能对比（以 GFLOPS 为单位）
// 测试环境：昇腾 910，矩阵形状 2048×2048×2048

// 配置 A：标准分块（保守策略）
using ConfigA = catlass::tile::GemmBatchedConfig<
    catlass::tile::BlockTile<128, 128, 32>,
    catlass::tile::WarpTile<64, 64, 32>,
    catlass::tile::ThreadTile<8, 8, 8>
>;
// 实测：约 380 GFLOPS

// 配置 B：大块策略（高并行）
using ConfigB = catlass::tile::GemmBatchedConfig<
    catlass::tile::BlockTile<256, 256, 64>,
    catlass::tile::WarpTile<64, 64, 64>,
    catlass::tile::ThreadTile<8, 16, 8>
>;
// 实测：约 520 GFLOPS（提升约 37%）

// 配置 C：超大分块（极致计算密度）
using ConfigC = catlass::tile::GemmBatchedConfig<
    catlass::tile::BlockTile<512, 512, 64>,
    catlass::tile::WarpTile<128, 128, 64>,
    catlass::tile::ThreadTile<16, 16, 8>
>;
// 实测：约 490 GFLOPS（下降 6%，因寄存器溢出）

从上述数据可以看出，Config B 取得了最优性能。Config C 虽然有更高的计算密度，但寄存器压力过大导致溢出，性能反而下降。

5.2 Double Buffer 策略的影响

Double Buffer 的核心思想是让数据加载与计算重叠进行：当一个数据块正在被计算时，下一个数据块可以同时被加载到 Local Memory 中，从而隐藏数据加载的延迟。

// Double Buffer 配置对比
// 无 Double Buffer（kDoubleBuffer = false, kStages = 1）
// 实测：约 310 GFLOPS（计算单元等待数据加载，空闲率 ~25%）

// 双缓冲（kDoubleBuffer = true, kStages = 2）
// 实测：约 480 GFLOPS（加载与计算重叠，空闲率降至 ~5%）

// 三级流水线（kDoubleBuffer = true, kStages = 3）
// 实测：约 520 GFLOPS（进一步减少流水线气泡，但增加 Local Memory 占用）

启用 Double Buffer 后性能提升约 55%。三级流水线相比两级又有约 8% 的提升，但 Local Memory 占用增加了一倍。对于 Local Memory 资源有限的场景，两级流水线是更稳妥的选择。

5.3 流水线级数与编译期优化

流水线级数（kStages）不仅影响运行时性能，还直接影响编译时间。级数越多，模板实例化的分支越多，编译时间呈指数级增长：

// 编译时间与流水线级数的关系（基于 GCC 12 测量）
// kStages = 1：编译时间约 45 秒
// kStages = 2：编译时间约 2 分 30 秒（+2.7x）
// kStages = 3：编译时间约 8 分 15 秒（+11x）
// kStages = 4：编译时间约 25 分钟（+33x）

因此，在开发阶段建议使用 kStages=1 或 kStages=2 以加快编译迭代速度，仅在最终上线前切换到 kStages=3 或更高以获取峰值性能。

六、两个关键陷阱与解决方案

6.1 陷阱一：模板实例化编译时间过长

症状：当 catlass 模板层级较深、流水线级数较多时，编译时间从几分钟急剧增长到几十分钟甚至更长。编译器（GCC/Clang）可能还会报出"template instantiation depth exceeded"错误。

原因分析：catlass 的模板实例化遵循递归展开机制。每增加一级流水线深度，模板编译器就需要多展开一倍的分支。同时，C++ 模板的隐式实例化规则会导致相同的 TilePolicy × ThreadMap × Mainloop 组合被重复实例化，造成编译时间的浪费。

解决方案：

第一，使用显式模板实例化（Explicit Template Instantiation）来避免隐式实例化的重复编译。将算子模板的实例化集中在一个 .cpp 文件中：

// gemm_instances.cpp
#include "catlass/gemm.h"

using GemmFP32Std = catlass::Gemm<...>;
using GemmFP16Bias = catlass::Gemm<...>;
using GemmFP16Relu = catlass::Gemm<...>;

// 显式实例化
template class catlass::Gemm<...>;  // 编译器只生成一次

第二，启用编译缓存（ccache）。catlass 的模板实例化结果可以被 ccache 缓存：

# 使用 ccache 加速重复编译
export USE_CCACHE=1
cd /path/to/catlass/build && cmake .. -DCMAKE_C_COMPILER=ccache
make -j$(nproc)

第三，将流水线级数的选择延迟到运行时。通过配置文件或环境变量选择预编译好的不同流水线级数的算子二进制：

// 运行时算子选择器
std::unique_ptr<GemmInterface> create_gemm_op(const GemmConfig& cfg) {
    switch (cfg.pipeline_stages) {
        case 2:
            return std::make_unique<GemmStage2>();
        case 3:
            return std::make_unique<GemmStage3>();
        default:
            return std::make_unique<GemmStage1>();
    }
}

6.2 陷阱二：TilePolicy 与实际数据布局不匹配导致结果错误

症状：GEMM 算子运行后，输出矩阵的数值与基准实现（Reference）不一致，部分元素出现明显的数值偏差或 NaN 值。

原因分析：这通常是因为 TilePolicy 中定义的 BlockTile 尺寸与输入矩阵的实际内存布局不兼容。例如，输入矩阵的 LDA（Leading Dimension of A）和 LDB（Leading Dimension of B）不是 BlockTile 相应维度的整数倍时，数据分块会出现"越界访问"或"空洞访问"，导致部分数据未被正确加载。

解决方案：

第一，在调用 GEMM 算子前，验证矩阵的步幅（Stride）是否满足 TilePolicy 的对齐要求：

bool validate_gemm_inputs(int M, int N, int K,
                          const half_t* A, int LDA,
                          const half_t* B, int LDB,
                          const TilePolicy& policy) {
    // 检查 LDA 和 LDB 是否为 BlockTile 的 M/N 维度整数倍
    constexpr int kBlockM = TilePolicy::kBlockM;
    constexpr int kBlockN = TilePolicy::kBlockN;
    
    if ((LDA % kBlockM != 0) || (LDB % kBlockN != 0)) {
        std::cerr << "Error: Stride not aligned with TilePolicy. "
                  << "LDA=" << LDA << " (need % " << kBlockM << "), "
                  << "LDB=" << LDB << " (need % " << kBlockN << ")\n";
        return false;
    }
    
    // 检查矩阵尺寸是否为 TilePolicy 各维度的整数倍
    if ((M % kBlockM != 0) || (N % kBlockN != 0)) {
        std::cerr << "Warning: Matrix size not divisible by BlockTile. "
                  << "Last block will have padding.\n";
    }
    return true;
}

第二，如果矩阵步幅不满足对齐要求，使用 Padding 包装器对输入矩阵进行预处理：

#include "catlass/utils/padding.h"

// 自动 Padding 输入矩阵到满足 TilePolicy 对齐要求的尺寸
void pad_matrix_if_needed(const half_t* raw_A, half_t* padded_A,
                          int M, int K, int LDA,
                          const TilePolicy& policy) {
    constexpr int kAlignM = TilePolicy::kBlockM;
    constexpr int kAlignK = TilePolicy::kBlockK;
    int M_pad = ((M + kAlignM - 1) / kAlignM) * kAlignM;
    int K_pad = ((K + kAlignK - 1) / kAlignK) * kAlignK;
    
    for (int m = 0; m < M_pad; ++m) {
        for (int k = 0; k < K_pad; ++k) {
            if (m < M && k < K) {
                padded_A[m * M_pad + k] = raw_A[m * LDA + k];
            } else {
                padded_A[m * M_pad + k] = half_t(0.0f);  // 填充零
            }
        }
    }
}

第三，在开发阶段始终开启数值正确性校验（Smoke Test），使用小矩阵与参考实现对比：

#include <cmath>

bool verify_gemm_output(const float* C_out, const float* C_ref,
                        int M, int N, float tolerance = 1e-4f) {
    int errors = 0;
    for (int m = 0; m < M; ++m) {
        for (int n = 0; n < N; ++n) {
            float diff = std::fabs(C_out[m * N + n] - C_ref[m * N + n]);
            if (diff > tolerance) {
                ++errors;
                if (errors <= 5) {  // 只打印前 5 个错误
                    printf("Mismatch at [%d, %d]: got %.6f, expected %.6f\n",
                           m, n, C_out[m * N + n], C_ref[m * N + n]);
                }
            }
        }
    }
    if (errors > 0) {
        printf("Verification failed: %d / %d elements mismatched\n",
               errors, M * N);
        return false;
    }
    printf("Verification passed: all %d elements match\n", M * N);
    return true;
}

七、实战代码

以下提供 10 个实战代码片段，涵盖 GEMM 模板实例化、BatchGEMM 配置、GEMM+ReLU 变体、性能基准测试和编译优化脚本。

7.1 代码一：基础 GEMM 模板实例化

// gemm_basic.cpp
#include "catlass/cutlass_utils.h"
#include "catlass/thread_map.h"
#include "catlass/gemm.h"
#include "aclrtlauncher.h"

// 定义 FP16 输入、FP32 累加的 GEMM 算子
using BasicGemm = catlass::Gemm<
    catlass::tile::BlockTile<128, 128, 32>,
    catlass::tile::WarpTile<64, 64, 32>,
    catlass::tile::ThreadTile<8, 8, 8>,
    catlass::threadmap::PitchLinear<
        catlass::threadmap::LayoutShape<16, 8>,
        catlass::threadmap::ClusterShape<1, 1>
    >,
    catlass::epilogue::Empty,
    float,    // 累加器与输出类型
    half_t    // 输入类型
>;

BasicGemm::Arguments prepare_args(int M, int N, int K,
                                   const half_t* A, const half_t* B,
                                   float* C) {
    BasicGemm::Arguments args;
    args.M = M; args.N = N; args.K = K;
    args.A = A; args.B = B; args.C = C;
    args.alpha = 1.0f;
    args.beta = 0.0f;
    return args;
}

7.2 代码二：BatchGEMM 完整配置

// gemm_batch.cpp
#include "catlass/batched_gemm.h"

// 批量 GEMM：支持动态批量大小，最大 64
template <int MaxBatch>
struct BatchedGemmOp {
    static constexpr int kMaxBatch = MaxBatch;
    
    using TilePolicy = catlass::tile::GemmBatchedConfig<
        catlass::tile::BlockTile<128, 128, 32>,
        catlass::tile::WarpTile<64, 64, 32>,
        catlass::tile::ThreadTile<8, 8, 8>
    >;
    
    using BatchedGemmType = catlass::BatchedGemm<
        TilePolicy,
        catlass::threadmap::PitchLinear<
            catlass::threadmap::LayoutShape<16, 8>,
            catlass::threadmap::ClusterShape<1, 1>
        >,
        catlass::epilogue::Empty,
        int,
        MaxBatch
    >;
    
    void run(float* C_list[], const half_t* A_list[],
             const half_t* B_list[], int batch_size,
             int M, int N, int K, void* ws, aclrtStream s) {
        BatchedGemmType op;
        typename BatchedGemmType::Arguments args;
        args.batch_size = batch_size;
        args.M = M; args.N = N; args.K = K;
        args.A_list = A_list;
        args.B_list = B_list;
        args.C_list = C_list;
        args.alpha = 1.0f;
        args.beta = 0.0f;
        op.initialize(args, ws);
        op.run(s);
    }
};

7.3 代码三：GEMM+ReLU 变体

// gemm_relu.cpp
#include "catlass/epilogue.h"

// 带 Bias 和 ReLU 的 GEMM 融合算子
using GemmBiasRelu = catlass::Gemm<
    catlass::tile::BlockTile<128, 128, 32>,
    catlass::tile::WarpTile<64, 64, 32>,
    catlass::tile::ThreadTile<8, 8, 8>,
    catlass::threadmap::PitchLinear<
        catlass::threadmap::LayoutShape<16, 8>,
        catlass::threadmap::ClusterShape<1, 1>
    >,
    catlass::epilogue::Chain<
        catlass::epilogue::BiasAddStage<float, 1>,
        catlass::epilogue::ActivationStage<
            catlass::epilogue::ReLU, float>
    >,
    float,
    half_t
>;

void run_gemm_relu(float* C, const half_t* A, const half_t* B,
                   const float* bias, int M, int N, int K,
                   void* ws, aclrtStream s) {
    GemmBiasRelu op;
    GemmBiasRelu::Arguments args;
    args.M = M; args.N = N; args.K = K;
    args.A = A; args.B = B; args.C = C;
    args.bias = bias;
    args.alpha = 1.0f;
    args.beta = 0.0f;
    op.initialize(args, ws);
    op.run(s);
}

7.4 代码四：自定义 Epilogue Stage

// custom_epilogue.cpp
#include "catlass/epilogue.h"

// 定义自定义的 LeakyReLU Epilogue Stage
struct LeakyReLUStage {
    float negative_slope;
    
    template <typename T>
    __aicore__ inline T apply(T x) const {
        return x > T(0) ? x : T(x * negative_slope);
    }
};

template <>
struct catlass::epilogue::EpilogueStageTraits<LeakyReLUStage> {
    static constexpr const char* kName = "LeakyReLU";
    using Arguments = LeakyReLUStage;
};

// 使用自定义 Stage
using GemmLeakyRelu = catlass::Gemm<
    catlass::tile::BlockTile<128, 128, 32>,
    catlass::tile::WarpTile<64, 64, 32>,
    catlass::tile::ThreadTile<8, 8, 8>,
    catlass::threadmap::PitchLinear<
        catlass::threadmap::LayoutShape<16, 8>,
        catlass::threadmap::ClusterShape<1, 1>
    >,
    catlass::epilogue::Chain<
        catlass::epilogue::BiasAddStage<float, 1>,
        LeakyReLUStage{0.2f}
    >,
    float,
    half_t
>;

7.5 代码五：性能基准测试框架

// gemm_benchmark.cpp
#include <chrono>
#include <iostream>
#include "catlass/gemm.h"

template <typename GemmOp>
class GemmBenchmark {
public:
    struct Result {
        double gflops;
        double latency_ms;
        double bandwidth_gbs;
    };
    
    Result run(int M, int N, int K,
               const half_t* A, const half_t* B, float* C,
               void* ws, aclrtStream s, int warmup = 3, int iterations = 10) {
        GemmOp op;
        auto args = prepare_args<GemmOp>(M, N, K, A, B, C);
        op.initialize(args, ws);
        
        // Warmup
        for (int i = 0; i < warmup; ++i) op.run(s);
        aclrtSynchronize();
        
        // Timed runs
        auto start = std::chrono::high_resolution_clock::now();
        for (int i = 0; i < iterations; ++i) op.run(s);
        aclrtSynchronize();
        auto end = std::chrono::high_resolution_clock::now();
        
        double elapsed_s = std::chrono::duration<double>(end - start).count();
        double avg_latency_ms = (elapsed_s / iterations) * 1000.0;
        
        // GFLOPS: 2*M*N*K operations per GEMM
        double flops = 2.0 * static_cast<double>(M) * N * K;
        double gflops = (flops * iterations) / elapsed_s / 1e9;
        
        // 内存带宽：bytes of A + B + C
        size_t bytes = (M * K + K * N + M * N) * sizeof(half_t);
        double bandwidth_gbs = (bytes * iterations) / elapsed_s / 1e9;
        
        return {gflops, avg_latency_ms, bandwidth_gbs};
    }
    
private:
    template <typename Op>
    auto prepare_args(int M, int N, int K,
                      const half_t* A, const half_t* B, float* C) {
        typename Op::Arguments args;
        args.M = M; args.N = N; args.K = K;
        args.A = A; args.B = B; args.C = C;
        args.alpha = 1.0f; args.beta = 0.0f;
        return args;
    }
};

// 使用示例
void benchmark_gemm() {
    int M = 2048, N = 2048, K = 2048;
    // 分配设备内存...
    GemmBenchmark<BasicGemm> bench;
    auto result = bench.run(M, N, K, d_A, d_B, d_C, workspace, stream);
    std::cout << "GFLOPS: " << result.gflops
              << " | Latency: " << result.latency_ms << " ms"
              << " | BW: " << result.bandwidth_gbs << " GB/s\n";
}

7.6 代码六：多配置自动调优脚本

// gemm_autotune.cpp
#include "catlass/autotune.h"

// 定义候选 TilePolicy 列表
using CandidatePolicies = catlass::autotune::PolicyList<
    // 小矩阵候选
    catlass::tile::GemmBatchedConfig<
        catlass::tile::BlockTile<64, 64, 32>,
        catlass::tile::WarpTile<32, 32, 32>,
        catlass::tile::ThreadTile<8, 8, 8>
    >,
    // 中等矩阵候选
    catlass::tile::GemmBatchedConfig<
        catlass::tile::BlockTile<128, 128, 32>,
        catlass::tile::WarpTile<64, 64, 32>,
        catlass::tile::ThreadTile<8, 8, 8>
    >,
    // 大矩阵候选
    catlass::tile::GemmBatchedConfig<
        catlass::tile::BlockTile<256, 256, 64>,
        catlass::tile::WarpTile<64, 64, 64>,
        catlass::tile::ThreadTile<8, 16, 8>
    >
>;

void autotune_gemm(const GemmConfig& cfg) {
    auto best = catlass::autotune::Tune<GemmOperator, CandidatePolicies>(
        cfg.M, cfg.N, cfg.K, cfg.A, cfg.B, cfg.C,
        cfg.ws, cfg.stream,
        /* search_metric = */ catlass::autotune::LatencyMetric{},
        /* max_attempts = */ 3
    );
    std::cout << "Best policy selected: " << best.policy_name << "\n"
              << "GFLOPS: " << best.gflops << "\n";
}

7.7 代码七：Workspace 内存估算

// workspace_estimation.cpp
#include "catlass/gemm.h"

// 计算 GEMM 算子所需的工作区大小
template <typename GemmOp>
size_t estimate_workspace_size(int M, int N, int K) {
    using TilePolicy = typename GemmOp::TilePolicy;
    constexpr int kBlockM = TilePolicy::kBlockM;
    constexpr int kBlockN = TilePolicy::kBlockN;
    constexpr int kBlockK = TilePolicy::kBlockK;
    constexpr int kStages = GemmOp::Mainloop::kStages;
    
    // Local Memory: A 和 B 各需要 kStages 个分块缓冲
    size_t local_mem = 2 * kStages * kBlockM * kBlockK * sizeof(half_t) +
                       2 * kStages * kBlockK * kBlockN * sizeof(half_t);
    
    // 临时寄存器溢出缓冲（保守估计为 Local Memory 的 20%）
    size_t spill_buf = local_mem / 5;
    
    // 算子内部临时缓冲
    size_t internal_buf = kBlockM * kBlockN * sizeof(float);
    
    return local_mem + spill_buf + internal_buf;
}

// 使用示例
void allocate_gemm_resources() {
    int M = 4096, N = 4096, K = 4096;
    size_t ws_size = estimate_workspace_size<BasicGemm>(M, N, K);
    std::cout << "Required workspace: " << ws_size / 1024 / 1024
              << " MB\n";
    void* ws = nullptr;
    aclrtMalloc(&ws, ws_size, ACL_MEM_MALLOC_HUGE_FIRST);
}

7.8 代码八：编译优化脚本

#!/bin/bash
# build_gemm_opt.sh - 面向 catlass GEMM 的编译优化脚本

set -e

PROJECT_ROOT="/path/to/catlass"
BUILD_DIR="${PROJECT_ROOT}/build_opt"
INSTALL_DIR="${PROJECT_ROOT}/install"

# 检测昇腾 CANN 环境
if [ -z "$ASCEND_HOME_PATH" ]; then
    echo "Error: ASCEND_HOME_PATH not set. Please source CANN environment."
    exit 1
fi

# 创建构建目录
mkdir -p "${BUILD_DIR}"
cd "${BUILD_DIR}"

# CMake 配置（启用编译优化选项）
cmake "${PROJECT_ROOT}" \
    -DCMAKE_BUILD_TYPE=Release \
    -DCMAKE_C_COMPILER="${ASCEND_HOME_PATH}/compiler/bin/arm-linux-gcc" \
    -DCMAKE_CXX_COMPILER="${ASCEND_HOME_PATH}/compiler/bin/arm-linux-g++" \
    -DCATLASS_ENABLE_AUTOTUNE=ON \
    -DCATLASS_USE_CCACHE=ON \
    -DCATLASS_MAX_TEMPLATE_DEPTH=512 \
    -DCMAKE_EXPORT_COMPILE_COMMANDS=ON

# 并行编译（根据 CPU 核心数）
NPROC=$(nproc --all 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)
echo "Compiling with ${NPROC} parallel jobs..."

make -j"${NPROC}" VERBOSE=1 2>&1 | tee build.log

# 检查编译结果
if [ $? -eq 0 ]; then
    echo "Build successful!"
    echo "Output libraries: ${BUILD_DIR}/lib/"
    ls -lh "${BUILD_DIR}/lib/"*.so 2>/dev/null || ls -lh "${BUILD_DIR}/lib/"*.a 2>/dev/null
else
    echo "Build failed. Check build.log for details."
    tail -100 build.log
    exit 1
fi

7.9 代码九：Ascend C 核函数绑定

// acl_entry.cpp
#include "acl/acl.h"
#include "aclrtlauncher.h"

// Ascend C 核函数导出（供 Ascend Runtime 调用）
extern "C" {

// 标准 GEMM 核函数注册
__global__ void aclnnGemmKernel(float* C, const half_t* A,
                               const half_t* B, int M, int N, int K,
                               float alpha, float beta) {
    using GemmKernel = catlass::GemmKernel<BasicGemm>;
    GemmKernel ker;
    ker(C, A, B, M, N, K, alpha, beta);
}

// 注册算子到 Ascend Runtime
aclError register_gemm_op() {
    aclOpExecutor* executor = aclCreateOpExecutor("Gemm");
    aclSetOpExecutorInputDesc(executor, "A", ACL_DT_FLOAT16,
                              {ACL_N, ACL_M, ACL_K});
    aclSetOpExecutorInputDesc(executor, "B", ACL_DT_FLOAT16,
                              {ACL_K, ACL_N, ACL_M});
    aclSetOpExecutorOutputDesc(executor, "C", ACL_DT_FLOAT,
                              {ACL_N, ACL_M, ACL_N});
    aclSetOpExecutorKernel(executor, "aclnnGemmKernel");
    return aclRegisterOp(executor);
}

}  // extern "C"

7.10 代码十：完整端到端测试脚本

// e2e_test.cpp
#include <random>
#include <cstring>
#include "gemm_basic.cpp"
#include "gemm_relu.cpp"
#include "gemm_batch.cpp"

bool compare_buffers(const float* a, const float* b, int size,
                     float eps = 1e-3f) {
    for (int i = 0; i < size; ++i) {
        if (std::fabs(a[i] - b[i]) > eps) return false;
    }
    return true;
}

void e2e_gemm_test(aclrtStream stream) {
    constexpr int M = 1024, N = 1024, K = 512;
    
    // 1. 分配设备内存
    half_t *d_A, *d_B, *d_A_batch[4], *d_B_batch[4];
    float *d_C, *d_C_ref, *d_C_out;
    aclrtMalloc((void**)&d_A, M * K * sizeof(half_t), ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc((void**)&d_B, K * N * sizeof(half_t), ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc((void**)&d_C, M * N * sizeof(float), ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc((void**&d_C_ref, M * N * sizeof(float), ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc((void**)&d_C_out, M * N * sizeof(float), ACL_MEM_MALLOC_HUGE_FIRST);
    
    // 2. 生成测试数据
    std::mt19937 rng(42);
    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
    std::vector<half_t> h_A(M * K), h_B(K * N);
    std::vector<float> h_C_ref(M * N, 0.0f);
    for (auto& v : h_A) v = half_t(dist(rng));
    for (auto& v : h_B) v = half_t(dist(rng));
    
    // 3. 计算参考结果（CPU 端朴素 GEMM）
    for (int m = 0; m < M; ++m)
        for (int k = 0; k < K; ++k)
            for (int n = 0; n < N; ++n)
                h_C_ref[m * N + n] += half_t(h_A[m * K + k]) * half_t(h_B[k * N + n]);
    
    // 4. 上传数据到 NPU
    aclrtMemcpy(d_A, M * K * sizeof(half_t), h_A.data(), h_A.size() * sizeof(half_t),
                ACL_MEMCPY_HOST_TO_DEVICE);
    aclrtMemcpy(d_B, K * N * sizeof(half_t), h_B.data(), h_B.size() * sizeof(half_t),
                ACL_MEMCPY_HOST_TO_DEVICE);
    aclrtMemcpy(d_C_ref, M * N * sizeof(float), h_C_ref.data(), h_C_ref.size() * sizeof(float),
                ACL_MEMCPY_HOST_TO_DEVICE);
    
    // 5. 分配工作区并运行 CATLASS GEMM
    size_t ws_size = estimate_workspace_size<BasicGemm>(M, N, K);
    void* ws; aclrtMalloc(&ws, ws_size, ACL_MEM_MALLOC_HUGE_FIRST);
    
    run_custom_gemm(d_C_out, d_A, d_B, M, N, K, ws, stream);
    aclrtSynchronize();
    
    // 6. 下载结果并验证
    std::vector<float> h_C_out(M * N);
    aclrtMemcpy(h_C_out.data(), d_C_out, M * N * sizeof(float),
                ACL_MEMCPY_DEVICE_TO_HOST);
    
    bool ok = verify_gemm_output(h_C_out.data(), h_C_ref.data(), M, N);
    std::cout << (ok ? "[PASS]" : "[FAIL]") << " GEMM correctness test\n";
    
    // 7. 运行性能基准
    GemmBenchmark<BasicGemm> bench;
    auto result = bench.run(M, N, K, d_A, d_B, d_C, ws, stream);
    std::cout << "Performance: " << result.gflops << " GFLOPS, "
              << result.latency_ms << " ms\n";
    
    // 清理
    aclrtFree(d_A); aclrtFree(d_B); aclrtFree(d_C);
    aclrtFree(d_C_ref); aclrtFree(d_C_out); aclrtFree(ws);
}

八、结尾

通过本文的系统梳理，我们可以看到 catlass 的核心价值在于：将昇腾 NPU 硬件的复杂细节，通过四层模板架构（TilePolicy / ThreadMap / Mainloop / Epilogue）抽象为可组合的配置参数，让开发者无需直面底层的向量指令调度和存储层级管理，就能组装出功能正确且性能接近硬件上限的自定义 GEMM 算子。

在实际生产中，如果你需要更丰富的激活函数支持——例如 Sigmoid、GELU、Swish 等在 Transformer 时代常用的非线性激活——推荐结合使用 ops-nn 库。ops-nn 是昇腾 CANN 生态中专门针对神经网络激活函数优化的自定义算子库，它与 catlass 的 GEMM 算子天然兼容：catlass 负责矩阵乘的核心计算，ops-nn 负责在 catlass 的 Epilogue 链上挂载更复杂的激活函数或归一化操作。二者配合使用，可以构建从数据加载到激活输出的全链路优化。

catlass 项目托管于 AtomGit，欢迎开发者参与贡献：

https://atomgit.com/cann/catlass

无论是发现模板中的潜在优化点，还是为新的数据布局添加 ThreadMap 特化，抑或是为 ops-nn 新增激活函数支持，每一份代码都是对昇腾 NPU 生态的宝贵贡献。希望本文能成为你使用 catlass 组装高性能 GEMM 算子的起点，在昇腾 NPU 上跑出令人满意的速度。