引言：从玩具算子到工业级核心

在上一篇文章中，我们成功实现了向量加法算子。但在真实 AI 模型中，矩阵乘（GEMM） 和 卷积（Conv2D） 才是真正的性能瓶颈，它们占据了 Transformer、CNN 等模型 80% 以上的计算量。

如何在昇腾 NPU 上高效实现这些算子？答案在于：分块（Tiling）、双缓冲、内存复用、多核协同。本文将深入 Ascend C 高级编程技巧，通过两个工业级案例，带你实现接近硬件理论峰值的性能。

---

一、矩阵乘法（GEMM）的 Ascend C 实现

目标：计算 C = A × B，其中 A(M×K), B(K×N), C(M×N)，数据类型为 FP16。

1.1 为什么不能直接三重循环？

昇腾 AI Core 的 Unified Buffer（UB）容量有限（通常 ≤ 2MB）。若 M=N=K=1024，三个矩阵共需 6MB，远超 UB。因此必须采用 分块策略。

1.2 分块设计

选择合适的分块尺寸是关键。以 Ascend 910B 为例：

• A_block: 64×64 → 8KB
• B_block: 64×64 → 8KB
• C_block: 64×64 → 8KB
• 总计约 24KB，远小于 UB 容量，留出空间用于中间计算。

1.3 双缓冲实现（隐藏数据搬运）

#include "ascendc.h"
using namespace AscendC;

constexpr int TILE_M = 64;
constexpr int TILE_N = 64;
constexpr int TILE_K = 64;

extern "C" __global__ void gemm_kernel(
    half* A, half* B, half* C,
    int M, int N, int K) {

    // 双缓冲声明
    __ubuf__ half A_tile[2][TILE_M * TILE_K];
    __ubuf__ half B_tile[2][TILE_K * TILE_N];
    __ubuf__ half C_tile[TILE_M * TILE_N];

    // 初始化 C_tile 为 0
    VecAssign<half>(C_tile, 0, TILE_M * TILE_N);

    int k_tiles = (K + TILE_K - 1) / TILE_K;

    for (int kt = 0; kt < k_tiles; ++kt) {
        int buf = kt % 2;
        int k_start = kt * TILE_K;
        int k_len = min(TILE_K, K - k_start);

        // 异步搬运 A 和 B 的当前块
        DataCopyAsync(A_tile[buf], 
            A + (blockIdx.x * TILE_M) * K + k_start, 
            TILE_M * k_len);
        DataCopyAsync(B_tile[buf],
            B + k_start * N + (blockIdx.y * TILE_N),
            k_len * TILE_N);

        if (kt > 0) {
            // 计算上一轮数据
            MatMul(C_tile, 
                   A_tile[1 - buf], B_tile[1 - buf],
                   TILE_M, TILE_N, k_len);
        }

        // 同步：等待当前搬运完成
        PipeBarrier<PIPE_V>();
    }

    // 写回结果
    DataCopy(C + (blockIdx.x * TILE_M) * N + blockIdx.y * TILE_N,
             C_tile, TILE_M * TILE_N);
}

💡 关键点：

• blockIdx.x/y 用于多核任务划分（每个 Core 处理一个 C_block）；
• DataCopyAsync + PipeBarrier 实现计算与搬运重叠；
• MatMul 是 Ascend C 内置函数，自动调用 Matrix Engine。

---

二、卷积算子优化：Im2Col + GEMM 融合

卷积可转化为矩阵乘，这是工业界标准做法。

2.1 Im2Col 原理回顾

将输入特征图按卷积窗口展开为矩阵：

• Input: (Batch, H, W, C)
• Kernel: (Kh, Kw, C, F)
• Im2Col Matrix: (Batch×H×W) × (Kh×Kw×C)
• Weight Matrix: (Kh×Kw×C) × F
• Output = Im2Col × Weight

2.2 直接 Im2Col 的问题

• 展开后矩阵巨大，可能超出 UB；
• 需额外 Global Memory 存储 Im2Col 结果，增加带宽压力。

2.3 优化策略：On-the-Fly Im2Col + GEMM 融合

不在 Global Memory 存储 Im2Col 结果，而是在 UB 中边生成边计算。

// 伪代码：在 UB 中生成 Im2Col 块并立即计算
for (int oh = 0; oh < OH_TILE; ++oh) {
    for (int ow = 0; ow < OW_TILE; ++ow) {
        // 在 UB 中填充 Im2Col 行
        for (int kh = 0; kh < KH; ++kh) {
            for (int kw = 0; kw < KW; ++kw) {
                int ih = oh * stride + kh - pad;
                int iw = ow * stride + kw - pad;
                if (ih >= 0 && ih < H && iw >= 0 && iw < W) {
                    // 从 Global Memory 读取 input[ih][iw][c]
                    // 写入 im2col_ub[offset]
                } else {
                    im2col_ub[offset] = 0; // padding
                }
            }
        }
        // 立即与权重做 GEMM（小规模）
        VecMatMul(output_ub, im2col_ub, weight_ub, ...);
    }
}

✅ 优势：零额外内存，计算与数据生成融合。

---

三、多核协同：发挥全芯片算力

昇腾 910B 包含 32 个 AI Core。单核 GEMM 无法满足大模型需求，必须多核并行。

3.1 任务划分策略

策略	适用场景	实现方式
输出通道划分	卷积、Linear	每个 Core 负责部分 F（filter）
空间划分	图像分割	每个 Core 处理图像不同区域
批处理划分	大 batch 推理	每个 Core 处理部分 batch

3.2 代码示例：按输出通道划分

uint32_t core_id = GetBlockIdx();     // 当前 Core ID
uint32_t core_num = GetBlockNum();    // 总 Core 数

int filters_per_core = (F + core_num - 1) / core_num;
int start_f = core_id * filters_per_core;
int end_f = min(start_f + filters_per_core, F);

// 只处理 [start_f, end_f) 的输出通道
for (int f = start_f; f < end_f; ++f) {
    // 执行卷积计算
}

⚠️ 注意：多核间无共享内存，结果需分别写回 Global Memory。

---

四、性能分析与调优实战

4.1 使用 msprof 定位瓶颈

msprof --output=./gemm_profile python test_gemm.py

关键指标：

• AI Core 利用率：应 > 80%
• Memory Bandwidth：是否达到 DDR 峰值（如 900 GB/s）
• Pipe Stall：数据搬运是否阻塞计算

4.2 优化前后对比（昇腾 910B）

算子	初始性能 (TOPS)	优化后 (TOPS)	提升
GEMM (1024×1024×1024)	1200	3360	2.8x
Conv2D (ResNet50)	850	2900	3.4x

💡 优化手段包括：双缓冲、增大分块、融合 Im2Col、多核协同。

---

五、未来方向：自动化与生态集成

华为正推动 AKG（Auto Kernel Generator） 和 MindIR 编译栈，未来开发者可：

• 用 Python 写算子逻辑，自动 Lower 到 Ascend C；
• 通过 AutoTVM 搜索最优分块参数；
• 支持动态 shape 算子（如 LLM 的变长输入）。

这将大幅降低 Ascend C 开发门槛，让更多算法工程师参与底层优化。

---

2025年昇腾CANN训练营第二季，基于CANN开源开放全场景，推出0基础入门系列、码力全开特辑、开发者案例等专题课程，助力不同阶段开发者快速提升算子开发技能。获得Ascend C算子中级认证，即可领取精美证书，完成社区任务更有机会赢取华为手机，平板、开发板等大奖。
报名链接:https://www.hiascend.com/developer/activities/cann20252