CANN自定义GEMM算子（Ascend C手写高性能矩阵乘法）

💡 为什么 GEMM 是 AI 计算的基石？

在深入代码之前，我们需要明确为什么要死磕 GEMM（通用矩阵乘法）：

Convolution (卷积) = im2col + GEMM + col2im
Attention (注意力机制) = GEMM (Q×K) + Softmax + GEMM (score×V)
Linear (全连接层) = GEMM
Transformer = 无数个 GEMM + Elementwise

在大模型训练中，GEMM 占据了 60%~70% 的计算时间。因此，GEMM 的性能直接决定了整个模型的训练和推理速度。

🧠 昇腾910 AI Core 内部结构回顾

写 Ascend C 算子就像是在给一个精密的工厂排班，必须了解每个车间的能力：

Cube Unit（矩阵计算单元） ：擅长矩阵乘法、卷积。单次能处理 16×16 或 32×32 的矩阵块运算，吞吐极高。
Vector Unit（向量计算单元）：擅长逐元素操作（Add, ReLU, Softmax等）。
Unified Buffer (UB，片上存储) ：容量约 2MB，带宽远高于 HBM（全局内存），用于暂存高频访问的中间数据。
DMA Engine（数据搬运引擎）：负责 GM（全局内存/HBM）与 UB 之间的数据搬运，且能与计算单元并行工作（这是双缓冲优化的基础）。

🚀 版本一：最简单的 GEMM（暴力直搬版）

先写一个逻辑最简单、能跑通但性能较差的版本，目的是理解基本流程。

cpp 复制代码

// gemm_v1.cpp - 最简教学版本（性能差，仅用于理解流程）
#include "kernel_operator.h"
using namespace AscendC;

class GemmKernelV1 {
public:
    __aicore__ inline GemmKernelV1() {}
    
    // 初始化：绑定地址，分配UB空间
    __aicore__ inline void Init(GM_ADDR a, GM_ADDR b, GM_ADDR c,
                                int32_t M, int32_t N, int32_t K) {
        this->M = M; this->N = N; this->K = K;
        
        // 将全局内存（GM）地址绑定到 GlobalTensor
        a_gm.SetGlobalBuffer((__gm__ half*)a, M * K);
        b_gm.SetGlobalBuffer((__gm__ half*)b, K * N);
        c_gm.SetGlobalBuffer((__gm__ half*)c, M * N);
        
        // 在 UB 里分配3块空间：A, B, C
        // 为什么用 half？因为 Cube 单元在 fp16 模式下吞吐翻倍
        pipe.InitBuffer(a_ub, M * K * sizeof(half)); 
        pipe.InitBuffer(b_ub, K * N * sizeof(half)); 
        pipe.InitBuffer(c_ub, M * N * sizeof(half)); 
    }
    
    // 核心计算逻辑
    __aicore__ inline void Process() {
        // 获取 LocalTensor（指向 UB 中的实际空间）
        LocalTensor<half> a_local = a_ub.Get<half>();
        LocalTensor<half> b_local = b_ub.Get<half>();
        LocalTensor<half> c_local = c_ub.Get<half>();
        
        // 步骤1：把 A 和 B 从 GM 搬到 UB（异步 DMA 搬运）
        DataCopy(a_local, a_gm, M * K);
        DataCopy(b_local, b_gm, K * N);
        
        // 关键：pipe_barrier 确保数据搬运完成后再开始计算
        pipe_barrier();
        
        // 步骤2：调用 Cube 单元做矩阵乘累加 (Mmad = Matrix Multiply and Accumulate)
        Mmad(c_local, a_local, b_local, M, N, K);
        
        // 步骤3：把结果从 UB 搬回 GM
        DataCopy(c_gm, c_local, M * N);
        pipe_barrier();
    }
    
private:
    GlobalTensor<half> a_gm, b_gm, c_gm; // GM 地址
    TBuf<UB> a_ub, b_ub, c_ub;           // UB buffer
    int32_t M, N, K;
    TPipe pipe;                          // 管理流水线与内存
};

❌ 版本一的致命缺陷：

假设 M=1024, N=1024, K=1024，单精度半浮点（half）下：

A 需要 2MB，B 需要 2MB，C 需要 2MB，总共 6MB。
而昇腾的 UB 只有 2MB。
结果：直接报错"内存不足"。这个版本只能跑极小的矩阵。

⚡ 版本二：带 Tiling（分块）的高性能 GEMM

为了解决 UB 装不下的问题，我们需要引入 Tiling（分块） 策略：把大矩阵切成小块，每次只把一小块搬进 UB 计算，算完再搬下一块。

cpp 复制代码

// gemm_v2.cpp - Tiling 优化版本
#include "kernel_operator.h"
using namespace AscendC;

class GemmKernelV2 {
public:
    __aicore__ inline GemmKernelV2() {}
    
    __aicore__ inline void Init(GM_ADDR a, GM_ADDR b, GM_ADDR c,
                                int32_t M, int32_t N, int32_t K) {
        this->M = M; this->N = N; this->K = K;
        
        a_gm.SetGlobalBuffer((__gm__ half*)a, M * K);
        b_gm.SetGlobalBuffer((__gm__ half*)b, K * N);
        c_gm.SetGlobalBuffer((__gm__ half*)c, M * N);
        
        // ★ Tiling 参数设定
        // UB = 2MB，我们要放 A_tile, B_tile, C_tile 三个块
        // 保守起见，取 256×256 = 65536 个元素 ≈ 128KB，三个块完全能放下
        TILE_M = 256;
        TILE_N = 256;
        TILE_K = 256;
        
        // 分配 UB 空间（只分配一个 tile 的大小）
        pipe.InitBuffer(a_tile_ub, TILE_M * TILE_K * sizeof(half));
        pipe.InitBuffer(b_tile_ub, TILE_K * TILE_N * sizeof(half));
        pipe.InitBuffer(c_tile_ub, TILE_M * TILE_N * sizeof(half));
    }
    
    __aicore__ inline void Process() {
        LocalTensor<half> a_tile = a_tile_ub.Get<half>();
        LocalTensor<half> b_tile = b_tile_ub.Get<half>();
        LocalTensor<half> c_tile = c_tile_ub.Get<half>();
        
        // ★ 三重循环：按 tile 遍历矩阵
        // 外层循环：遍历输出矩阵的行 (M维度)
        for (int32_t m_start = 0; m_start < M; m_start += TILE_M) {
            int32_t cur_m = min(TILE_M, M - m_start); // 处理边界
            
            // 中层循环：遍历输出矩阵的列 (N维度)
            for (int32_t n_start = 0; n_start < N; n_start += TILE_N) {
                int32_t cur_n = min(TILE_N, N - n_start);
                
                // ★ 初始化 C tile 为零（因为是累加操作 C += A*B）
                Duplicate(c_tile, (half)0.0, cur_m * cur_n);
                
                // 内层循环：遍历累加维度 (K维度)
                for (int32_t k_start = 0; k_start < K; k_start += TILE_K) {
                    int32_t cur_k = min(TILE_K, K - k_start);
                    
                    // ★ 步骤1：搬入 A 的分块 [cur_m, cur_k]
                    // 计算 GM 偏移量：m_start * K + k_start
                    DataCopy(a_tile, a_gm[m_start * K + k_start], cur_m * cur_k);
                    
                    // ★ 步骤2：搬入 B 的分块 [cur_k, cur_n]
                    // 计算 GM 偏移量：k_start * N + n_start
                    DataCopy(b_tile, b_gm[k_start * N + n_start], cur_k * cur_n);
                    
                    // ★ 关键：等待 DMA 搬运完成
                    pipe_barrier();
                    
                    // ★ 步骤3：Cube 计算
                    Mmad(c_tile, a_tile, b_tile, cur_m, cur_n, cur_k);
                }
                
                // ★ 步骤4：把计算好的 C tile 写回 GM
                DataCopy(c_gm[m_start * N + n_start], c_tile, cur_m * cur_n);
                pipe_barrier(); // 确保写回完成，防止后续覆盖
            }
        }
    }
    
private:
    GlobalTensor<half> a_gm, b_gm, c_gm;
    TBuf<UB> a_tile_ub, b_tile_ub, c_tile_ub;
    int32_t M, N, K;
    int32_t TILE_M, TILE_N, TILE_K;
    TPipe pipe;
};

🔑 版本二的核心优化点解析

打破内存墙：通过 Tiling，我们不再试图一次性搬运几 MB 的数据，而是每次只搬运 128KB 左右的小块。这让程序能够处理任意维度的超大矩阵。
适配硬件架构：Tiling 的大小（如 256）通常是 Cube 单元单次处理大小（16 或 32）的整数倍，这能让 Cube 单元满负荷运转，避免算力浪费。
同步与异步的博弈 ：DataCopy 是异步的（由 DMA 执行），Mmad 是同步的（由 Cube 执行）。pipe_barrier() 就像一道栅栏，强制 CPU 等待前面的搬运任务全部完成后，才允许启动后续的 Mmad 计算，保证了数据的正确性。

🚀 下一步进阶方向

版本二虽然解决了内存问题，但依然存在"串行"缺陷：DMA 搬运时，Cube 单元在闲置；Cube 计算时，DMA 在闲置。

要达到极致的性能（逼近理论峰值），下一步就是引入 双缓冲（Double Buffering） 技术：在 UB 中开辟两套缓冲区，当 Cube 正在计算 Buffer 0 中的数据时，DMA 同时去搬运下一批数据到 Buffer 1。通过这种"计算与搬运并行"的策略，可以将原本串行的耗时完美掩盖掉！