引言：从“使用”到“创造”——解锁算子开发的白盒能力

在深度学习框架中，算子（Operator）是构建模型的基本单元。传统方式下，开发者依赖框架提供的内置算子，或通过 CUDA/OpenCL 从零编写底层 Kernel。前者缺乏灵活性，后者开发成本高、调试困难。

catlass 作为 CANN 社区提供的高性能算子模板库，采用 分层抽象设计，将矩阵类算子拆解为可复用、可替换的模块。开发者无需从零开始，只需 组合/修改现有模板，即可快速实现定制化高性能算子。

本文将手把手教你如何基于 catlass 扩展开发自定义算子，涵盖 环境搭建、模板结构解析、GEMM 定制、融合算子开发、性能调优 全流程。

---

一、catlass 的分层架构与核心思想

1.1 四层抽象模型

catlass 将算子开发分为四层，每层职责清晰：

✅ 核心思想：上层关注“做什么”，下层关注“怎么做”。

1.2 模板化优势

传统开发	catlass 模板开发
从零编写 Kernel	复用底层硬件原语
修改即全盘重写	局部替换算法模块
调优需深入汇编	调整 Tiling 参数即可
难以支持多硬件	抽象硬件差异

---

二、开发环境搭建与项目结构

2.1 环境依赖

• 操作系统：Linux (Ubuntu 20.04+)
• 编译器：GCC >= 7.5, < 13.0
• 构建工具：CMake >= 3.16
• Python：>= 3.8, < 3.12（用于测试）

2.2 项目目录结构

catlass/
├── cmake/              # CMake 构建脚本
├── docs/               # 开发文档
├── examples/           # 算子样例
│   ├── 00_basic_matmul/    # 基础 GEMM
│   ├── 01_matmul_bias/     # 融合 Bias
│   └── ... 
├── include/
│   ├── catlass/        # 核心模板头文件
│   │   ├── gemm/       # GEMM 相关模块
│   │   ├── layout/     # 数据布局
│   │   └── tile/       # 分块策略
│   └── tla/            # 基础数据结构
├── scripts/
│   └── build.sh        # 编译脚本
└── tests/              # 测试用例

📌 关键目录：examples/ 是学习起点，include/catlass/gemm/ 是扩展核心。

---

三、第一步：运行一个基础算子样例

以 00_basic_matmul 为例：

3.1 Host 侧调用代码 (examples/00_basic_matmul/basic_matmul.cpp)

#include "catlass/gemm/device_gemm.h"

int main() {
    // 初始化输入输出
    std::vector<half> A(M * K), B(K * N), C(M * N);
    // ... 填充数据

    // 调用 catlass GEMM
    using Gemm = cutlass::gemm::device::Gemm<
        half, cutlass::layout::RowMajor,  // A 类型与布局
        half, cutlass::layout::ColumnMajor, // B 类型与布局
        float, cutlass::layout::RowMajor   // C 类型与布局
    >;

    Gemm gemm_op;
    typename Gemm::Arguments args{
        {M, N, K},          // 问题规模
        {A.data(), K},      // A 及其 leading dimension
        {B.data(), K},      // B 及其 leading dimension
        {C.data(), N},      // C 及其 leading dimension
        {C.data(), N}       // 输出 C
    };

    gemm_op(args);
    return 0;
}

3.2 编译与运行

cd examples/00_basic_matmul
mkdir build && cd build
cmake .. && make
./basic_matmul

✅ 成功标志：输出 Verification PASSED。

---

四、第二步：理解模板参数 —— 定制你的 GEMM

catlass 的核心是 模板参数化。以 GEMM 为例：

template <
  typename ElementA,               // A 元素类型 (e.g., half)
  typename LayoutA,                // A 布局 (e.g., RowMajor)
  typename ElementB,
  typename LayoutB,
  typename ElementC,
  typename LayoutC,
  typename ElementAccumulator = ElementC, // 累加器类型
  typename OperatorClass = cutlass::arch::OpClassSimt, // 计算类型 (SIMT/TensorCore)
  int ThreadblockShapeM = 128,     // Threadblock 分块 M
  int ThreadblockShapeN = 128,     // Threadblock 分块 N
  int ThreadblockShapeK = 8,       // Threadblock 分块 K
  // ... 更多参数
>
class Gemm;

4.1 关键参数说明

参数	作用	典型值
ElementA/B/C	数据类型	half, float
LayoutA/B/C	内存布局	RowMajor, ColumnMajor
ThreadblockShape	分块大小	(128,128,8)
WarpShape	Warp 分块	(64,64,8)
InstructionShape	指令形状	(16,8,8) for TensorCore

4.2 自定义一个 FP16 GEMM

// 自定义分块策略：更大的 K 以适应长序列
using MyGemm = cutlass::gemm::device::Gemm<
    half, cutlass::layout::RowMajor,
    half, cutlass::layout::ColumnMajor,
    float, cutlass::layout::RowMajor,
    float,                          // Accumulator
    cutlass::arch::OpClassTensorOp, // 使用 TensorCore
    64, 64, 64,                     // Threadblock: 64x64x64
    32, 32, 64,                     // Warp: 32x32x64
    16, 8, 16                       // Instruction: 16x8x16
>;

💡 提示：分块大小需满足硬件约束（如 Warp 大小是 Threadblock 的因数）。

---

五、第三步：开发融合算子 —— 以 MatMul + Bias 为例

融合算子（Fused Operator）将多个操作合并为一个 Kernel，减少内存读写。

5.1 融合算子的模板结构

catlass 通过 Epilogue 模块实现融合。Epilogue 负责 GEMM 后的处理。

5.2 自定义 Epilogue

// 定义融合操作：LinearCombinationBias
template <
  typename ElementOutput,
  typename ElementAccumulator,
  typename ElementBias,
  int ElementsPerAccess
>
struct LinearCombinationBias {
  using FragmentOutput = Array<ElementOutput, ElementsPerAccess>;
  using FragmentAccumulator = Array<ElementAccumulator, ElementsPerAccess>;
  using FragmentBias = Array<ElementBias, ElementsPerAccess>;

  __device__ void operator()(
      FragmentOutput &output,
      FragmentAccumulator const &accumulator,
      FragmentBias const &bias) {
    
    // D = alpha * accumulator + beta * bias
    CUTLASS_PRAGMA_UNROLL
    for (int i = 0; i < ElementsPerAccess; ++i) {
      output[i] = accumulator[i] + bias[i]; // 简化版
    }
  }
};

5.3 组装融合 GEMM

// 使用自定义 Epilogue
using EpilogueOp = LinearCombinationBias<float, float, float, 4>;

using FusedGemm = cutlass::gemm::device::Gemm<
    half, cutlass::layout::RowMajor,
    half, cutlass::layout::ColumnMajor,
    float, cutlass::layout::RowMajor,
    float,
    cutlass::arch::OpClassTensorOp,
    128, 128, 8,
    EpilogueOp  // 注入自定义 Epilogue
>;

✅ 优势：Bias 仅加载一次，避免额外全局内存访问。

---

六、第四步：高级定制 —— 自定义 Tiling 策略

Tiling 策略直接影响性能。catlass 允许完全自定义分块逻辑。

6.1 Tiling 策略的组成

• Threadblock Tiling：GPU Block 级分块；
• Warp Tiling：Warp 级分块；
• Instruction Tiling：硬件指令级分块。

6.2 为长序列优化的 Tiling

场景：$M=1,N=32768,K=4096$（典型 LLM 推理）

问题：标准 Tiling $(128,128,8)$ 导致大量 Warp 空闲。

解决方案：增大 N 方向分块。

// 自定义 Tiling 策略
using ThreadblockShape = cutlass::gemm::GemmShape<64, 256, 64>;
using WarpShape = cutlass::gemm::GemmShape<32, 128, 64>;
using InstructionShape = cutlass::gemm::GemmShape<16, 8, 16>;

using LongSeqGemm = cutlass::gemm::device::Gemm<
    half, cutlass::layout::RowMajor,
    half, cutlass::layout::ColumnMajor,
    float, cutlass::layout::RowMajor,
    float,
    cutlass::arch::OpClassTensorOp,
    ThreadblockShape,
    WarpShape,
    InstructionShape
>;

📊 收益：Warp 利用率从 40% 提升至 90%。

---

七、第五步：性能调优与验证

7.1 性能分析工具

catlass 提供 profiler 工具：

# 编译 profiler
cd tools/profiler && mkdir build && cd build && cmake .. && make

# 运行性能测试
./gemm_profiler --m=4096 --n=4096 --k=4096 --type=half

输出：

Problem Size: 4096x4096x4096
Throughput: 125.3 TFLOPS (82% of peak)

7.2 正确性验证

使用 testbed 框架：

// tests/gemm/test_custom_gemm.cu
TEST(CustomGemm, LongSeq) {
    Testbed<GemmTraits> testbed;
    testbed.run(
        {1, 32768, 4096},  // M, N, K
        {1.0f, 0.0f}        // alpha, beta
    );
    EXPECT_TRUE(testbed.verify());
}

---

八、实战案例：开发一个自定义 Swish-GEMM 融合算子

Swish 激活函数：$\text{Swish}(x)=x\cdot \sigma (x)$

8.1 定义 Swish Epilogue

template <typename T, int N>
struct SwishEpilogue {
  using FragmentOutput = Array<T, N>;
  using FragmentAccumulator = Array<T, N>;

  __device__ T sigmoid(T x) {
    return 1.0f / (1.0f + expf(-x));
  }

  __device__ void operator()(
      FragmentOutput &output,
      FragmentAccumulator const &accumulator) {
    CUTLASS_PRAGMA_UNROLL
    for (int i = 0; i < N; ++i) {
      T x = accumulator[i];
      output[i] = x * sigmoid(x);
    }
  }
};

8.2 组装算子

using SwishGemm = cutlass::gemm::device::Gemm<
    half, cutlass::layout::RowMajor,
    half, cutlass::layout::ColumnMajor,
    half, cutlass::layout::RowMajor, // 输出为 half
    float,                           // 累加器为 float
    cutlass::arch::OpClassTensorOp,
    128, 128, 8,
    SwishEpilogue<half, 8>           // 8 个元素 per access
>;

8.3 性能收益

算子	耗时 (ms)	内存带宽 (GB/s)
GEMM + Swish (分离)	5.2 + 1.8 = 7.0	1200
Swish-GEMM (融合)	5.5	950

✅ 结论：融合后总耗时减少 21%，带宽需求降低。

---

九、常见问题与调试技巧

9.1 编译错误：模板参数不匹配

• 现象：static_assert failed: "ThreadblockShape must be divisible by WarpShape"
• 解决：检查分块大小是否满足整除关系。

9.2 性能不佳：Shared Memory Bank Conflict

• 现象：性能远低于理论峰值。
• 解决：调整 ThreadblockShape 或使用 Swizzled Shared Memory 布局。

9.3 数值错误：累加器精度不足

• 现象：FP16 GEMM 结果误差大。
• 解决：将 ElementAccumulator 设为 float。

---

十、贡献与社区参与

catlass 欢迎社区贡献：

1. Fork 仓库：https://atomgit.com/cann/catlass
2. 在 examples/ 添加新算子样例；
3. 提交 PR，附带测试和文档。

---

结语

算子开发不再是少数专家的专利。catlass 通过 分层模板化设计，让开发者能像搭积木一样构建高性能算子。

无论你是想优化现有模型，还是探索新型算子，catlass 都为你提供了强大的白盒工具。现在就开始你的扩展开发之旅吧！

---

🔗 相关链接：

• CANN 组织主页：https://atomgit.com/cann
• catlass 仓库地址：https://atomgit.com/cann/catlass