catlass 架构剖析 - 昇腾算子模板库设计哲学

前言

catlass 是 CANN 的算子模板库提供了丰富的算子开发模板本文深入剖析 catlass 的架构设计和核心设计理念

背景 - 为什么需要算子模板库

在 NPU 上开发算子是一个复杂的工程需要考虑内存管理数据搬运计算调度等多个方面catlass 提供了完整的算子开发模板

整体架构

catlass 的架构设计遵循分层理念

复制代码

      应用层PyTorch 等          

      AscendCL 接口层              

      catlass 模板层               
         
   GEMM  Conv  Reduce ...    
         

      opbase 基础层               

      NPU 硬件

核心设计理念

设计理念一模板化开发

Catlass 采用了模板化的开发方式

cpp 复制代码

template <typename T>
class GemmKernel {
public:
    __aicore__ inline void process() {
        load_data();
        compute();
        store_result();
    }
};

设计理念二内存层级管理

NPU 的内存分为全局内存GM和局部内存LM

cpp 复制代码

class GemmKernel {
private:
    GlobalTensor<half> a_gm;  // 全局内存
    LocalTensor<half> a_lm[128];  // 局部内存
};

设计理念三计算流水线化

通过计算和数据搬运的流水化隐藏内存访问延迟

核心模块详解

模块一GEMM 模块

GEMM 是深度学习中最重要的算子之一

cpp 复制代码

class GemmKernel {
public:
    __aicore__ inline void init() {
        M = input_shape[0];
        N = weight_shape[1];
        K = input_shape[1];
    }
    
    __aicore__ inline void compute() {
        // Tile 级别的矩阵乘法
        for (int i = 0; i < M; i += TILE_M) {
            for (int j = 0; j < N; j += TILE_N) {
                gemm_core(c_lm, a_lm, b_lm, TILE_M, TILE_N, K);
            }
        }
    }
};

模块二卷积模块

卷积是计算机视觉模型的核心算子

性能优化技巧

Tile 大小选择

cpp 复制代码

static constexpr int TILE_M = 128;
static constexpr int TILE_N = 128;
static constexpr int TILE_K = 32;

循环展开

cpp 复制代码

#pragma unroll 4
for (int k = 0; k < TILE_K; k++) {
    process(i, j, k);
}

总结

catlass 是 CANN 算子开发的核心模板库，其设计理念深刻体现了对 NPU 硬件特性和深度学习计算模式的深度适配。这些理念共同构建了一个高效、易用且可扩展的算子开发框架：

模板化开发，简化算子开发流程
- 核心理念：将算子的通用计算模式（如 GEMM、卷积、规约）抽象为可复用的 C++ 模板类。开发者只需专注于核心计算逻辑和数据排布，无需重复编写内存管理、流水线调度等底层代码。
- 带来的优势：极大降低了算子开发的门槛和代码量，保证了不同算子间实现风格和性能优化策略的一致性，提升了代码的可维护性。
内存层级管理，自动管理 GM 和 LM
- 核心理念 ：显式地区分并管理 NPU 的全局内存（GM，容量大、带宽高、延迟高）和局部内存（LM，容量小、带宽极高、延迟低）。Catlass 通过 GlobalTensor 和 LocalTensor 等封装类，提供了清晰的数据搬运接口。
- 带来的优势：引导开发者进行高效的数据复用（Data Reuse），将频繁访问的数据缓存在 LM 中，减少对 GM 的访问，从而有效克服"内存墙"瓶颈，是提升算子在 NPU 上性能的关键。
计算流水线化，隐藏内存访问延迟
- 核心理念：将"计算"与"数据搬运"这两个主要操作进行重叠（Overlap）。当一部分计算单元在处理当前数据块时，另一部分硬件资源（如 DMA）已经在为下一个数据块执行加载（Load）或存储（Store）操作。
- 带来的优势：充分利用 NPU 的硬件并行能力，将原本串行的"加载-计算-存储"流程转化为并行的流水线，从而将高延迟的内存访问时间"隐藏"在计算时间内，显著提升计算单元的利用率和整体吞吐量。

总结而言，Catlass 通过这三层设计理念，构建了一个从抽象到具体、从软件到硬件紧密协同的开发范式。它不仅是一个代码库，更是一套指导如何在 Ascend NPU 上编写高性能算子的最佳实践方法论。开发者遵循这套方法论，能够更高效地开发出充分发挥硬件潜力的高性能算子。

更多细节请参考https://atomgit.com/cann/catlass