《Ascend C 高级优化实战：实现高性能卷积算子》

int batch;int kH, kW;通过 Im2Col + GEMM 架构，我们成功在 Ascend C 中实现了高性能卷积算子。虽然开发复杂度高，但掌握了这一模式后，可快速迁移到其他线性算子（如 Linear、BatchMatmul）。关键经验数据布局决定性能：尽量减少非连续访问Cube 是性能核心：确保 GEMM 分块对齐多核协同：合理划分任务避免负载不均Ascend C 是一把“双刃剑

2501_94602287

963人浏览 · 2025-12-10 22:19:14

2501_94602287 · 2025-12-10 22:19:14 发布

引言

卷积（Convolution）是计算机视觉任务的核心操作，其计算密集性使其成为 AI 芯片性能的“试金石”。在昇腾 NPU 上，虽然已有高度优化的内置卷积算子，但在某些特殊场景（如非标准卷积核、稀疏卷积、自定义归一化）下，仍需通过 Ascend C 实现定制化高性能卷积。

本文将深入剖析卷积在昇腾上的实现原理，并手把手教你用 Ascend C 编写一个 Im2Col + GEMM 架构的 2D 卷积算子。我们将重点讲解 内存重排、Cube 单元调度、多核并行 等高级优化技术，最终实现接近理论峰值的性能。

前置知识：熟悉卷积原理、了解 GEMM、掌握 Ascend C 基础语法。

一、卷积在昇腾上的挑战

1.1 计算模式

标准卷积：
Yn,c,h,w=∑kc,kh,kwXn,ci,h′,w′⋅Wco,ci,kh,kw

问题：

输入数据非连续（滑动窗口）
权重复用率低
输出写冲突

1.2 昇腾硬件限制

Cube 单元要求输入为 16x16 矩阵块（FP16）
UB 容量有限（通常 2MB/core）
GM 带宽瓶颈

二、优化策略：Im2Col + GEMM

2.1 Im2Col 原理

将卷积转换为矩阵乘：

将输入图像按卷积窗口展开为大矩阵 A（Height × Width × Kernel_H × Kernel_W → M × K）
权重 reshape 为矩阵 B（Output_Ch × Input_Ch × K_H × K_W → N × K）
输出 = A × B^T

优点：可复用高度优化的 GEMM

2.2 Ascend C 实现难点

Im2Col 需大量数据重排（GM → UB → GM）
GEMM 需满足 Cube 对齐要求
需多级分块（Block, Tile, Fragment）

三、Ascend C 卷积算子实现

3.1 数据结构定义

struct ConvParam {
    int batch;
    int inC, outC;
    int inH, inW;
    int kH, kW;
    int padH, padW;
    int strideH, strideW;
    int outH, outW;
};

3.2 核心类设计

class Conv2dForward {
public:
    __aicore__ inline void Init(
        GM_ADDR input, GM_ADDR weight, GM_ADDR output,
        const ConvParam& param) {
        // 绑定 GM
        inputGm.SetGlobalBuffer((__gm__ half*)input, ...);
        weightGm.SetGlobalBuffer((__gm__ half*)weight, ...);
        outputGm.SetGlobalBuffer((__gm__ half*)output, ...);
        this->param = param;

        // 计算 Im2Col 后的矩阵维度
        M = param.batch * param.outH * param.outW;
        K = param.inC * param.kH * param.kW;
        N = param.outC;

        // 分配 UB 缓冲区
        AllocBuffer();
    }

    __aicore__ inline void Process() {
        // 分块处理：按输出通道分块
        for (int oc = 0; oc < N; oc += TILE_N) {
            int nSize = min(TILE_N, N - oc);
            ProcessOutputChannel(oc, nSize);
        }
    }

private:
    void ProcessOutputChannel(int ocStart, int nSize) {
        // Step 1: 从 GM 拷贝权重块到 UB（已 reshape 为 N×K）
        CopyWeightTile(ocStart, nSize);

        // Step 2: 按 M 维度分块，执行 Im2Col + GEMM
        for (int m = 0; m < M; m += TILE_M) {
            int mSize = min(TILE_M, M - m);

            // Im2Col: 从 inputGm 构造 A 矩阵（M×K）
            Im2Col(m, mSize);

            // GEMM: C = A * B^T
            GemmCompute(m, mSize, nSize);
        }
    }

    void Im2Col(int mStart, int mSize) {
        // 实现滑动窗口展开
        // 注意：需处理 padding 和 stride
        for (int i = 0; i < mSize; ++i) {
            int idx = mStart + i;
            int n = idx / (param.outH * param.outW);
            int hw = idx % (param.outH * param.outW);
            int h = hw / param.outW;
            int w = hw % param.outW;

            for (int ci = 0; ci < param.inC; ++ci) {
                for (int kh = 0; kh < param.kH; ++kh) {
                    for (int kw = 0; kw < param.kW; ++kw) {
                        int ih = h * param.strideH + kh - param.padH;
                        int iw = w * param.strideW + kw - param.padW;

                        half val = 0;
                        if (ih >= 0 && ih < param.inH && iw >=0 && iw < param.inW) {
                            val = inputGm[n * param.inC * param.inH * param.inW +
                                         ci * param.inH * param.inW +
                                         ih * param.inW + iw];
                        }
                        // 写入 im2colUb[i][ci * kH * kW + kh * kW + kw]
                        WriteIm2Col(i, ci, kh, kw, val);
                    }
                }
            }
        }
    }

    void GemmCompute(int mSize, int nSize) {
        // 调用 Cube GEMM
        // A: im2colUb (M×K), B: weightUb (N×K), C: outputUb (M×N)
        Matmul(outputUb, im2colUb, weightUb, mSize, nSize, K);
    }

private:
    GlobalTensor<half> inputGm, weightGm, outputGm;
    TBuf<half> im2colUb, weightUb, outputUb;
    ConvParam param;
    int M, N, K;
};

3.3 Cube GEMM 调用

Ascend C 提供 Matmul 指令，但需满足：

M, N, K 必须是 16 的倍数（FP16）
UB 缓冲区按 16-byte 对齐

void Matmul(TBuf<half>& C, TBuf<half>& A, TBuf<half>& B,
            int M, int N, int K) {
    // 分块为 16x16
    for (int m = 0; m < M; m += 16) {
        for (int n = 0; n < N; n += 16) {
            Cube gemm;
            gemm.Init(M, N, K, 16, 16, 16);
            gemm.SetInput(A, m, 0);
            gemm.SetWeight(B, n, 0);
            gemm.SetOutput(C, m, n);
            gemm.Compute();
        }
    }
}

四、多核并行优化

昇腾 910 有 32 个 AI Core。我们可按 输出通道 划分任务：

extern "C" __global__ void conv2d_forward(...) {
    SetSysCtrl();
    int coreId = GetBlockIdx();
    int coreNum = blockDim.x;

    // 每个 core 处理一部分输出通道
    int channelsPerCore = (param.outC + coreNum - 1) / coreNum;
    int startOc = coreId * channelsPerCore;
    int endOc = min(startOc + channelsPerCore, param.outC);

    if (startOc >= param.outC) return;

    Conv2dForward op;
    op.Init(...);
    op.ProcessRange(startOc, endOc);
}