昇腾CANN ops-nn 里的 MatMul 为什么能跑这么快

雨季666

339人浏览 · 2026-05-20 22:53:09

雨季666 · 2026-05-20 22:53:09 发布

在昇腾NPU 上跑一个矩阵乘，CANN ops-nn 仓库里的 MatMul 算子比标准实现快出 2-3 倍，靠的不是更快的计算单元——Cube 单元是一样的——靠的是对数据搬运的极致压缩。

理解 ops-nn 的设计，等于理解昇腾达芬奇架构上高性能算子的全部秘密。

从 HBM 到 Cube：搬一次就够

矩阵乘有三个输入张量：A、B 和 bias（可选），一个输出 C。朴素实现的做法是：把 A 和 B 从 HBM 搬到 L1 缓存，Cube 算完 C，再把 C 写回 HBM。如果 A 的尺寸超过了 L1 的容量，就要把 A 切成块，分多次搬运。

这个过程中，HBM 的带宽是瓶颈。Ascend 910 的 HBM 带宽是 1.2 TB/s，Cube 峰值算力是 256 TFLOPS（FP16）。每 byte 数据实际能做的计算量是 256T / 1.2T ≈ 213 次运算。如果搬运次数太多，Cube 就得等数据，利用率往下掉。

ops-nn 的 MatMul 核心优化就一句话：每加载一次 A 和 B 的块，Cube 做满所有能做的乘法，不浪费一次搬运。

实际实现比这句话复杂。核心分三层：

第一层：分块大小选择

分块不是随便切的。昇腾NPU 的 L1 缓存有固定大小，A 块和 B 块加起来必须塞得进去。ops-nn 的分块策略基于输入矩阵的 shape 和数据类型自动选最优的 Br（A 的行块）和 Bc（B 的列块）：

// ops-nn 分块大小的自动选择逻辑
struct MatMulTileConfig {
    uint32_t Br;  // A 块行数
    uint32_t Bc;  // B 块列数
    uint32_t Kr;  // 累加循环块大小
    
    static MatMulTileConfig auto_select(uint32_t M, uint32_t N, uint32_t K, DataType dtype) {
        // L1 可用容量 / 元素大小 = 能放的个数
        // 确保 Br * Kr（A 块）+ Kr * Bc（B 块）+ Br * Bc（C 块）<= L1 容量
        size_t elem_size = (dtype == DTYPE_FP16) ? 2 : 4;
        size_t l1_capacity = 196 * 1024 / elem_size;  // L1 约 196K 元素（FP16）
        
        // 理论上最优：Br ≈ Bc ≈ sqrt(l1_capacity / 3)
        float optimal = std::sqrt(l1_capacity / 3.0f);
        
        // 对齐到 16（Cube 单元的限制：一次算 16×16 的矩阵乘）
        Br = align_up(std::min(optimal, (float)M), 16);
        Bc = align_up(std::min(optimal, (float)N), 16);
        Kr = align_up(std::min(optimal, (float)K), 16);
        
        return {Br, Bc, Kr};
    }
};

第二层：双缓冲流水

搬运和计算不是顺序的。一块在 Cube 上算的时候，下一块已经在往另一个 L1 buffer 里搬了。这个叫双缓冲，ops-nn 用 C++ 模板把它抽象成通用机制：

template <typename T, int NUM_BUFFERS = 2>
struct DoubleBufferLoader {
    T buffer[NUM_BUFFERS];  // ping-pong 两个 buffer
    int active = 0;         // 当前 Cube 在算哪个
    
    void load_next(const T* hbm_src) {
        int preload = (active + 1) % NUM_BUFFERS;
        // 异步搬运：不等 Cube 完成，直接开始搬下一块
        aclrtMemcpyAsync(
            &buffer[preload], hbm_src, 
            sizeof(T), ACL_MEMCPY_DEVICE_TO_DEVICE, copy_stream
        );
    }
    
    T& wait_and_get() {
        // 等 Cube 完成当前这块
        // 同时让预取的那块就绪
        aclrtSynchronizeStream(compute_stream);
        active = (active + 1) % NUM_BUFFERS;
        return buffer[active];
    }
};

第三层：Cube 指令排布

Cube 单元一次算 16×16 的矩阵乘——这是昇腾达芬奇架构的硬件限制。ops-nn 在分块内部，把 Br×Bc 的块再切成 16×16 的小块，用 Mmad 指令反复调用。但这种反复调用不能是朴素的 for 循环——每次 Mmad 之间有指令流水线延迟（latency），必须把连续几条 Mmad 的参数预先排好，打满流水线：

__aicore____inline__ void mmad_pipelined(
    LocalTensor<FP16>& C, 
    LocalTensor<FP16>& A, 
    LocalTensor<FP16>& B,
    int step  // 当前是第几步（0 到 num_steps-1）
) {
    // 三条 Mmad 排流水线
    // step 0：Mmad(step0) 开始执行
    // step 1：Mmad(step0) 还在执行，Mmad(step1) 已经可以开始（用不同的寄存器组）
    // step 2：Mmad(step0) 完成，Mmad(step1) 在执行，Mmad(step2) 开始
    Mmad(C_local[step % 3], A_local[step % 3], B_local[step % 3], {16, 16, 16});
}

这三个层次叠在一起，ops-nn 的 MatMul 能做到 Cube 利用率 85% 以上——也就是说大部分时间 Cube 都在干活，没有空转等数据。

融合算子：不只是加法

ops-nn 的价值不止于单个 MatMul。它真正的威力在于算子融合——把 MatMul + BiasAdd + Activation 三个串行算子和成一个 kernel。

// ops-nn 的 MatMul + Bias + GELU 融合算子
// 三个步骤在同一个 kernel 里一趟跑完

__aicore__ void fused_matmul_gelu(
    GlobalTensor<FP16>& C,      // 输出
    GlobalTensor<FP16>& A,      // 权重
    GlobalTensor<FP16>& B,      // 输入激活值
    GlobalTensor<FP16>& bias,    // bias 向量
    uint32_t M, uint32_t N, uint32_t K
) {
    // 1. MatMul：Cube 单元算
    Mmad(C_local, A_tile, B_tile, {Br, Bc, Kr});
    
    // 2. BiasAdd：Vector 单元做加法，结果还在片上
    // 不写回 HBM——这意味着省了一次 HBM 的读和一次 HBM 的写
    for (int i = 0; i < Br; ++i) {
        for (int j = 0; j < Bc; ++j) {
            C_local(i, j) = C_local(i, j) + bias(j);
        }
    }
    
    // 3. GELU：Vector 单元做激活
    // 也不写回 HBM——C_local 从 MatMul 的结果变成 bias 结果再变成 GELU 结果
    // 全程在片上，三个步骤只有最终的 GELU 结果才写回 HBM
    for (int i = 0; i < Br; ++i) {
        for (int j = 0; j < Bc; ++j) {
            FP16 x = C_local(i, j);
            // GELU ≈ 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
            C_local(i, j) = 0.5f * x * (1.0f + tanh_fp16(0.797885f * x * (1.0f + 0.044715f * x * x)));
        }
    }
    
    // 最后才写回 HBM：三个步骤只写一次
    DataCopy(C, C_local, {Br, Bc});
}

三个独立算子各跑各的，共需要 3 次 HBM 读和 3 次 HBM 写。融合后只需要 2 次读（A 和 B）和 1 次写（C），HBM 流量直接砍到三分之一。

对于 Transformer 这种每层十几个矩阵乘的模型，融合后的端到端收益在 20-30%——不是单个算子变快了，是搬运的活少了。

内存格式：NCHW 还是 NHWC

op-nn 做卷积时优先用 NHWC 格式而非 NCHW。原因和矩阵乘分块一脉相承——NHWC 下，通道维 C 是连续的，Cube 单元算张量乘时不需要跨步访问；NCHW 下每个通道的数据被空间维度隔开了，读起来不连续，HBM 带宽浪费在跨步寻址上。

换上 NHWC 后 Conv2D 的 HBM 读取效率从 ~65% 提到 ~90%，同样的时间内 Cube 拿到更多数据去算。

用起来

ops-nn 的算子在 CANN 中已经作为内置算子库部署。PyTorch 侧通过 torch_npu 自动路由，不用手动调：

import torch
import torch_npu

# Linear 层自动调用 ops-nn 的 MatMul 算子
linear = torch.nn.Linear(768, 3072).npu().half()
x = torch.randn(2, 512, 768).npu().half()
y = linear(x)  # 底层调用 ops-nn 的融合 MatMul + BiasAdd

# 显式指定融合模式可以手动触发 MatMul + GELU 融合
with torch_npu.npu.fused_mode(enable=True):
    x = torch.nn.functional.linear(x, weight, bias)
    x = torch.nn.functional.gelu(x)  # 这两个算子会被自动合并