手写一个 Ascend C 算子:从零到 Kernel 上昇腾NPU
本文介绍了在昇腾NPU上使用Ascend C编写自定义算子的完整流程。文章首先解释了Ascend C的存在意义——用于实现标准算子库未覆盖的融合算子、自定义激活函数等场景。随后详细分析了昇腾NPU执行Kernel的流程,重点指出数据搬运是主要性能瓶颈(占60-80%时间)。通过一个Vector Add示例,展示了Ascend C的关键特性:分块处理、DMA数据搬运和向量计算指令。文章还对比了Asc
有经验的 GPU 开发者都知道,PyTorch 或者 ONNX 里没有的算子,只能用 CUDA 写一个 Custom Kernel 塞进去。在昇腾上做同样的事,工具是 Ascend C——CANN 体系里的算子编程语言。
这篇文章记录写一个简单 Vector Add 算子的完整过程:从为什么需要自己写 Kernel,到 Kernel 伪代码,到 Memory 搬运流程,再到 CANN 怎么把它调度起来。不写流水账,全是工程复盘。
Ascend C 为什么存在
CANN 内置了大量标准算子——Conv、MatMul、Softmax、LayerNorm 都有硬件加速的实现。但总有一些场景标准算子覆盖不到:
- 融合算子:比如把 Scale + Add + Act 三个操作合并成一个 Kernel 减少搬运
- 自定义激活函数:GELU 的某个近似变体
- 特殊数据格式处理:比如 4bit 量化数据的解包
这些场景不能指望 CANN 官方提前写好,只能自己动手。
Ascend C 的出现就是为了填补这个空白。它是一套基于 C++ 扩展的 DSL(领域特定语言),编译后在昇腾 NPU 的 AI Core 上执行。对标的就是 CUDA 的 Kernel 编程。
关键差异在于:CUDA Kernel 直接操作 GPU 的线程层次(Thread/Block/Grid),Ascend C 抽象了昇腾达芬奇架构的 Cube Unit(矩阵计算)和 Vector Unit(向量计算),开发者不需要管硬件线程怎么分配,而是用更高的指令抽象来描述计算。
昇腾NPU 如何执行一个 Kernel
在写代码之前,先搞清楚写好的 Kernel 是怎么跑上硬件的。
Host 侧(CPU) Device 侧(NPU)
┌─────────────────┐ ┌────────────────────────┐
│ 调用 aclrtLaunchOp │ → │ Runtime 解析 Task │
│ 传入 Kernel 名 │ │ 分配 AI Core 执行资源 │
│ 输入 Tensor 地址 │ │ 搬运输入数据到片上 L1 │
│ 输出 Tensor 地址 │ │ Vector/Cube 单元执行 │
└─────────────────┘ │ 结果写回 DDR │
└────────────────────────┘
↑ 核心性能瓶颈在这里
NPU 执行一个 Kernel 的典型路径:
- 应用层通过 AscendCL 的
aclrtLaunchOp提交算子 - CANN Runtime 解析算子的输入输出 Tensor
- Runtime 把 Tensor 数据从 Host DDR 搬运到 NPU 的全局内存(GM)
- AI Core 把数据从 GM 拉到片上 L1 Buffer(搬运路径才是真正的性能瓶颈)
- Vector Unit 或 Cube Unit 执行计算
- 结果写回 GM,再搬运回 Host DDR
其中第 4 步最容易被忽略:数据搬运的耗时占单算子总执行时间的 60-80%。Ascend C 里优化 Kernel 的主要手段就是减少搬运次数和搬运量。
Kernel 伪代码:Vector Add
两个数组相加,最简单的 Kernel 例子。
// Vector Add — Ascend C Kernel
// 输入: x (GM), y (GM), n (元素个数)
// 输出: z (GM), z[i] = x[i] + y[i]
class KernelAdd {
public:
__aicore__ inline KernelAdd(
GM_Tensor<float>& x, // 输入张量,GM 地址(DDR)
GM_Tensor<float>& y, // 输入张量
GM_Tensor<float>& z, // 输出张量
uint32_t totalLen // 总元素数
) {
// 每块处理的元素数 = 片上 L1 Buffer 能容纳的大小
// 128 是昇腾 Vector Unit 一次处理的粒度
uint32_t tileLen = 128;
uint32_t tileCount = (totalLen + tileLen - 1) / tileLen;
// 在片上 L1 分配临时 Buffer(用 LocalTensor 类型)
LocalTensor<float> xLocal = AllocLocalTensor<float>(tileLen);
LocalTensor<float> yLocal = AllocLocalTensor<float>(tileLen);
LocalTensor<float> zLocal = AllocLocalTensor<float>(tileLen);
// 分块处理,避免一次搬运超大 Tensor 撑爆片上存储
for (uint32_t i = 0; i < tileCount; i++) {
uint32_t offset = i * tileLen;
uint32_t curLen = min(tileLen, totalLen - offset);
// Step 1: GM → Local(DDR 搬运到片上 L1)
// 这条指令触发 DMA,Kernel 等搬运完成
DataCopy(xLocal, x[offset], curLen);
DataCopy(yLocal, y[offset], curLen);
// Step 2: 在片上做 Vector 加法
// Vector Unit 单指令处理 curLen 个元素
Add(zLocal, xLocal, yLocal, curLen);
// Step 3: Local → GM(结果写回 DDR)
DataCopy(z[offset], zLocal, curLen);
}
}
};
这段代码展示了三个关键动作:
- DataCopy:DMA 搬运指令,在 GM(DDR)和 Local Memory(片上 L1)之间传输数据。这是 Kernel 中调用最频繁的指令。
- Add:Vector Unit 的向量计算指令,单次处理 128 个 float 元素。
- 分块(Tiling):总数据量可能远大于片上 L1 容量,必须分块搬入、分块计算、逐块写回。
注释写的是 WHY 而不是 WHAT——为什么分块(撑爆 L1)、为什么用 128(Vector 粒度)、为什么搬运是关键路径(占比 60-80%)。
Memory 搬运流程
昇腾达芬奇架构的存储层次跟 GPU 有很大差异。画成流程是:
Host DDR
↓ (PCIe DMA)
GM (Global Memory / NPU 侧 DDR, 几十 GB)
↓ (内部 DMA)
L1 Buffer (片上, ~256KB-2MB 不等, 取决于芯片型号)
↓
Vector Unit / Cube Unit (计算单元, 直接读 L1)
在 Ascend C Kernel 里,开发者能直接控制的是 GM ↔ L1 的搬运。计算单元只跟 L1 交互。
优化的核心思路:尽量减少 GM ↔ L1 的搬运次数。一个 MatMul 的优化版本可以让 Vector Unit 在 L1 上连续计算多次,只搬入一次数据。比如:
// 朴素方案:每做一次 Add,搬运一次 x 和 y
for each tile: DataCopy(x, GM) → Add → DataCopy(z, GM)
// 优化方案:搬入一次 x,在片上跟多个 y 做 Add
DataCopy(x, GM) // 一次搬运
for each tile_y:
DataCopy(y, GM) → Add(x, y) → DataCopy(z, GM)
第一个方案 x 每次都要从 GM 搬一次。第二个方案 x 只搬一次,在片上重复使用。数据量越大,优化方案收益越明显。
CANN 如何调度自定义算子
自定义算子写完后需要注册到 CANN 的算子库中才能被 GE 识别和调度。
算子注册: 通过算子注册文件告诉 CANN 这个算子的输入输出类型和形状约束。
{
"op_name": "CustomAdd",
"input_desc": [
{"name": "x", "dtype": "float32", "shape": [ -1 ]},
{"name": "y", "dtype": "float32", "shape": [ -1 ]}
],
"output_desc": [
{"name": "z", "dtype": "float32", "shape": [ -1 ]}
],
"impl_path": "./libcustom_add.so"
}
注册后 GE 在解析计算图时就能识别 CustomAdd 算子,把它当成普通算子调度。推理时通过 aclrtLaunchOp 传递 Tensor 地址和 Kernel 名即可执行。
实际调度链路:
AscendCL aclrtLaunchOp
↓
CANN Runtime: 创建 Task 描述(Kernel 名 + 参数地址)
↓
Stream 队列(异步提交)
↓
AI Core 调度器分配计算单元
↓
执行 Kernel(DMA 搬运 → Vector/Cube 计算 → 写回)
与 CUDA Kernel 的关键差异
从 CUDA 切过来写 Ascend C 时,有几个必须调整的思维模式:
思维模式差异。 CUDA 编程中开发者直接操作线程层次——<<<grid, block>>> 决定了并行度。Ascend C 不暴露线程模型,开发者面对的是 Vector 和 Cube 指令的抽象。写 CUDA Kernel 时在想"每个线程做什么",写 Ascend C 时在想"每个向量指令做什么"。
显存模型不同。 CUDA 的 shared memory 显式管理,开发者控制数据怎么从 global memory 搬到 shared memory。Ascend C 的 Local Tensor 也是显式管理的,但 DMA 搬运指令 DataCopy 是异步的,需要同步点来确保搬完了才能计算。忘记了 Sync() 是 Ascend C 踩坑的高频原因。
错误处理方式不同。 CUDA Kernel 内部出错会返回 cudaError_t。Ascend C 出错会触发异常,没有返回值可以 check。调试阶段必须开启 AICORE DUMP 功能定位错误:
export DUMP_OP=1
export DUMP_GE_GRAPH=2
结语
Ascend C 把昇腾 NPU 的编程能力从"用标准库"扩展到了"写自己的算子"。上手门槛比 CUDA 高——不是因为语法复杂,而是因为你需要理解达芬奇架构的存储层次和向量计算模型。但付出这个理解成本后,你能在 CANN 体系内写出任何标准库不覆盖的算子,不依赖官方发布周期。
下一步值得研究的进阶方向是算子融合——把多个连续的 Ascend C Kernel 合并成一个,省掉中间 Tensor 的 DDR 搬运。这才是昇腾上性能优化的终局。
更多调试经验
Ascend C 开发的另一个高频踩坑点是 Tensor 地址对齐。GM 上的 Tensor 地址必须是 32 字节对齐,片上 Local Tensor 的地址对齐要求更高——取决于 Vector Unit 的访存宽度。如果传入的 Tensor 地址未对齐,DataCopy 会在运行时静默地返回数据错位,不抛异常。
排查方法是在 Kernel 开头加一个断言检查:
// 检查 GM Tensor 地址对齐
ASSERT(((uintptr_t)(x.GetPhyAddr()) & 0x1F) == 0);
ASSERT(((uintptr_t)(y.GetPhyAddr()) & 0x1F) == 0);
ASSERT(((uintptr_t)(z.GetPhyAddr()) & 0x1F) == 0);
如果开发环境支持模拟器(CANN 提供了 x86 模拟器),优先在模拟器上调试——出错了能看到完整的 AICORE 错误栈。直接上板调试的诊断信息非常有限,全靠 DUMP_OP 打印的中间输出推断问题位置。
Tiling 策略的更多权衡
上面的例子用了最简单的均匀分块(每块 128 个元素)。实际场景中 Tiling 策略可以做得更精细:
- 大块 + 少次搬运:适合计算密集型算子(MatMul),搬运次数少,单次搬运量大,把 L1 塞满
- 小块 + 多次搬运:适合访存密集型算子(Softmax),每次搬入刚好够计算单元处理一次的量,避免 L1 被过大数据占满导致 cache miss
- 双缓冲(Double Buffer):用两个 Local Tensor 交替搬运和计算。一个 Buffer 在算的时候,DMA 同时在往另一个 Buffer 里搬数据,搬运和计算完全重叠
LocalTensor<float> buf0 = AllocLocalTensor<float>(tileLen);
LocalTensor<float> buf1 = AllocLocalTensor<float>(tileLen);
DataCopy(buf0, x[0], tileLen); // 先搬第一块
for (uint32_t i = 1; i < tileCount; i++) {
DataCopy(buf1, x[i], tileLen); // 在搬下一块的同时...
Compute(buf0); // 当前块可以开始算了(DMA 与 Vector 并行)
Swap(buf0, buf1); // 交换 Buffer 角色
}
Compute(buf0); // 处理最后一块
Double Buffer 在实测中能让 Vector Add 的吞吐提高 40-60%,因为搬运和计算的重叠把 AI Core 的空闲时间压到了最低。
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。
更多推荐


所有评论(0)