CANN算子开发入门：从零开始写第一个Ascend C算子

前言

想给昇腾NPU写自定义算子，但不知道从哪入手。官方文档很全，但太偏参考手册------你想知道的是"第一步该干什么"、"代码怎么组织"、"怎么编译调试"。这篇文章从零开始，手把手带你写第一个Ascend C算子，并让它成功在昇腾NPU上跑起来。全程实战，不绕弯子。

准备工作：搭建开发环境

写Ascend C算子之前，得先把开发环境搭好。需要这些东西：

1. 昇腾NPU硬件：训练用Ascend 910，推理用Ascend 310P。本地开发可以用模拟器（CPU模式），不需要真实硬件。
2. CANN toolkit：昇腾的开发者工具包，包含Ascend C编译器、运行时库、调试工具。去昇腾官网下载对应版本（推荐8.0.RC1或以上）。
3. opbase组件：算子开发的基础框架，提供内存管理、tiling策略、kernel模板。这个在CANN toolkit里自带，不需要单独装。
4. CMake 3.16+：用于构建算子工程。
5. Python 3.8+：用于跑测试用例。

搭环境的具体步骤（Ubuntu 22.04为例）：

# 1. 安装 CANN toolkit
# 去昇腾官网下载 CANN 8.0.RC1 的 toolkit 包
# 假设下载到了 ~/Downloads/Ascend-cann-toolkit_8.0.RC1_linux-x86_64.run
chmod +x ~/Downloads/Ascend-cann-toolkit_8.0.RC1_linux-x86_64.run
sudo ~/Downloads/Ascend-cann-toolkit_8.0.RC1_linux-x86_64.run --install

# 2. 设置环境变量
# 把这几行加到 ~/.bashrc 里
export ASCEND_HOME=/usr/local/Ascend
export PATH=$ASCEND_HOME/bin:$PATH
export LD_LIBRARY_PATH=$ASCEND_HOME/lib64:$LD_LIBRARY_PATH
export PYTHONPATH=$ASCEND_HOME/python/site-packages:$PYTHONPATH

# 3. 验证安装
ascendc --version
# 应该输出：Ascend C Compiler 8.0.RC1

# 4. 安装 Python 依赖
pip3 install numpy torch torch_npu

环境搭好后，可以跑一个helloworld算子试试水。

第一步：创建算子工程

Ascend C算子工程有固定的目录结构。用opbase提供的工程模板可以一键生成：

# 用 opbase 的模板创建算子工程
cd ~/workspace
opbase create_operator --name=add --type=ELEMENT_WISE

# 生成的目录结构：
# add_operator/
# ├── op_kernel/          # 算子 kernel 实现
# │   └── add_kernel.cpp
# ├── op_proto/          # 算子原型定义
# │   └── add.h
# ├── op_tiling/         # Tiling 策略实现
# │   └── add_tiling.cpp
# ├── testcases/         # 测试用例
# │   └── test_add.py
# ├── CMakeLists.txt     # 构建脚本
# └── README.md          # 算子说明文档

这个结构不是随便定的，是CANN社区的标准。op_kernel/放Ascend C实现的kernel、op_proto/放算子原型（输入输出格式）、op_tiling/放Tiling策略（怎么切数据块）、testcases/放测试用例。

第二步：写算子原型（op_proto/）

算子原型定义了算子的输入输出格式和属性。这是算子跟外部交互的接口，必须写对。

op_proto/add.h的内容：

// add.h - Add 算子的原型定义
// 这个文件定义了 Add 算子的输入输出格式，以及属性参数
// 为什么要单独定义原型？因为框架（PyTorch/MindSpore）需要根据原型来验证输入格式
// 如果输入格式不对（比如数据类型不支持），框架会提前报错，而不是跑到 kernel 里才 core dump

#ifndef ADD_OP_PROTO_H
#define ADD_OP_PROTO_H

#include "op_proto.h"  // CANN 提供的原型定义基础头文件

namespace add_operator {

// Add 算子的输入描述
// 两个输入：a 和 b，类型都是 FP16 或 FP32
// 输出描述：一个输出 c，类型跟输入一致
// 属性描述：无属性（Add 算子不需要额外参数）
inline void DefineAddOpProto(OpProto& proto) {
    // 定义输入 a
    // 为什么要限制数据类型？因为 Ascend C 的 Vector 单元对不同数据类型的支持不一样
    // FP16 的吞吐是 FP32 的两倍，所以优先用 FP16
    proto.Input("a")
        .TypeConstraint({"FP16", "FP32"})  // 支持的数据类型
        .Description("第一个输入张量");
    
    // 定义输入 b
    proto.Input("b")
        .TypeConstraint({"FP16", "FP32"})
        .Description("第二个输入张量");
    
    // 定义输出 c
    // 输出类型必须跟输入类型一致，这是广播规则要求的
    proto.Output("c")
        .TypeConstraint({"FP16", "FP32"})
        .Description("输出张量（a + b）");
    
    // Add 算子没有属性参数，所以不需要定义 attr
}

}  // namespace add_operator

#endif  // ADD_OP_PROTO_H

代码里的注释解释了WHY：为什么要单独定义原型？为什么要限制数据类型？这些注释不是废话，是帮后人理解设计决策的关键。

第三步：写Tiling策略（op_tiling/）

Tiling是Ascend C编程里最关键的性能优化手段。它的作用是把大张量切成小块，让每块都能装进L1 Buffer，同时让Cube单元或Vector单元的并行度最高。

op_tiling/add_tiling.cpp的内容：

// add_tiling.cpp - Add 算子的 Tiling 策略
// Tiling 的作用是：把大张量切成小块，每块的大小要适配硬件缓存
// 为什么要 Tiling？因为 Ascend 910 的 L1 Buffer 只有 1MB
// 如果一次性处理整个张量（比如 1GB 的数据），L1 Buffer 装不下
// 必须切成小块，逐块处理

#include "tiling/add_tiling.h"
#include "op_tiling.h"

namespace add_operator {

// 计算 Tiling 参数
// 输入：numel（张量元素总数）、dtype_size（每个元素占的字节数）
// 输出：Tiling 参数（block_size、block_num、tail_size）
void CalcAddTiling(int64_t numel, int dtype_size, AddTiling& tiling) {
    // L1 Buffer 大小是 1MB = 1048576 字节
    // 我们要让每个 block 的数据刚好能装进 L1 Buffer
    // 这里留点余量，用 256KB 作为 block 大小
    const int64_t L1_CAPACITY = 256 * 1024;  // 256KB
    
    // 每个 block 能放多少个元素
    int64_t elem_per_block = L1_CAPACITY / dtype_size;
    
    // block 数量 = 向上取整(numel / elem_per_block)
    tiling.block_num = (numel + elem_per_block - 1) / elem_per_block;
    
    // 每个 block 的实际大小（最后一个 block 可能不满）
    tiling.block_size = elem_per_block;
    
    // 最后一个 block 的大小（可能小于 elem_per_block）
    tiling.tail_size = numel % elem_per_block;
    if (tiling.tail_size == 0) {
        tiling.tail_size = elem_per_block;
    }
}

// Tiling 入口函数，会被框架调用
// 为什么要提供这个入口？因为框架在跑算子之前，必须先知道怎么切数据
// 这个函数在 Host 端执行，计算出 Tiling 参数后传给 Device 端的 kernel
extern "C" __global__ void add_tiling(
    GM_ADDR tiling_buf, int64_t numel, int dtype_size)
{
    // 在 Host 端计算 Tiling 参数
    AddTiling tiling;
    CalcAddTiling(numel, dtype_size, tiling);
    
    // 把 Tiling 参数写到 tiling_buf（这块内存在 Host 和 Device 之间共享）
    // Device 端的 kernel 会读这块内存，拿到 Tiling 参数
    AddTiling* t = reinterpret_cast<AddTiling*>(tiling_buf);
    *t = tiling;
}

}  // namespace add_operator

Tiling策略的核心逻辑是：根据L1 Buffer大小算每个block该多大、一共要切几个block、最后一个block有多大。这些参数会传给Device端的kernel，指导它怎么切数据。

第四步：写算子Kernel（op_kernel/）

Kernel是算子的核心计算逻辑，跑在昇腾NPU的AI Core上。Add算子很简单，就是element-wise的加法，用Vector单元就能搞定。

op_kernel/add_kernel.cpp的内容：

// add_kernel.cpp - Add 算子的 Ascend C 实现
// 这个文件是算子的核心，跑在 NPU 的 AI Core 上
// Add 算子做 element-wise 加法：c[i] = a[i] + b[i]
// 为什么用 Vector 单元？因为 element-wise 操作不需要矩阵乘，Vector 单元的吞吐更高

#include "kernel/add_kernel.h"
#include "tuning/add_tiling.h"

namespace add_operator {

// Add 算子的 kernel 实现
// __aicore__ 是 Ascend C 的关键字，表示这个函数是跑在 AI Core 上的 kernel
// GM_ADDR 是全局内存地址（HBM 的地址）
extern "C" __global__ __aicore__ void add_kernel(
    GM_ADDR a, GM_ADDR b, GM_ADDR c,
    GM_ADDR tiling_buf, int64_t numel)
{
    // 1. 初始化 pipe（Ascend C 的内存管理对象）
    // 为什么要管理内存？因为 AI Core 上的内存（L1 Buffer）很有限
    // pipe 帮你自动管理缓冲区的分配和释放，避免内存泄漏
    TPipe pipe;
    
    // 2. 从 tiling_buf 读取 Tiling 参数
    // 为什么要先读 Tiling？因为 kernel 需要知道怎么切数据
    AddTiling* tiling = reinterpret_cast<AddTiling*>(tiling_buf);
    int64_t block_size = tiling->block_size;
    int64_t block_num = tiling->block_num;
    int64_t tail_size = tiling->tail_size;
    
    // 3. 初始化缓冲区（在 L1 Buffer 上）
    // TQue 是 Ascend C 的队列管理器，负责在 L1 Buffer 上分配缓冲区
    // 这里申请三个缓冲区：a_buf（存 a）、b_buf（存 b）、c_buf（存 c）
    // 每个缓冲区的大小是 block_size 个元素（因为每次处理一个 block）
    TQue<QuePosition::VECIN, 1> a_buf, b_buf;
    TQue<QuePosition::VECOUT, 1> c_buf;
    pipe.InitBuffer(a_buf, block_size * sizeof(half));
    pipe.InitBuffer(b_buf, block_size * sizeof(half));
    pipe.InitBuffer(c_buf, block_size * sizeof(half));
    
    // 4. 分块处理
    // 每个 AI Core 处理一个 block（由框架的调度器分配）
    // 这里用 GetBlockIdx() 拿到当前 AI Core 的编号
    // 如果编号 >= block_num，说明没有数据要处理，直接返回
    int32_t block_idx = GetBlockIdx();
    if (block_idx >= block_num) {
        return;
    }
    
    // 5. 计算当前 block 的实际大小
    int64_t cur_block_size = block_size;
    if (block_idx == block_num - 1) {
        // 最后一个 block，大小可能是 tail_size（不满 block_size）
        cur_block_size = tail_size;
    }
    
    // 6. 从 HBM 加载 a 和 b 到 L1 Buffer
    // DataCopy 是 Ascend C 提供的内存拷贝函数
    // 它会在后台发起 DMA 传输，把数据从 HBM 搬到 L1 Buffer
    LocalTensor<half> a_local = a_buf.AllocTensor<half>();
    LocalTensor<half> b_local = b_buf.AllocTensor<half>();
    DataCopy(a_local, a + block_idx * block_size, cur_block_size * sizeof(half));
    DataCopy(b_local, b + block_idx * block_size, cur_block_size * sizeof(half));
    
    // 7. Vector 单元做加法
    // vec_add 是 Ascend C 提供的向量加法接口
    // 它会在 Vector 单元上并行计算，比写 for 循环快得多
    LocalTensor<half> c_local = c_buf.AllocTensor<half>();
    vec_add(c_local, a_local, b_local, cur_block_size, 1);
    
    // 8. 把结果从 L1 Buffer 写回 HBM
    DataCopy(c + block_idx * block_size, c_local, cur_block_size * sizeof(half));
    
    // 9. 释放缓冲区
    // 为什么不手动 free？因为 pipe 会在 kernel 结束时自动回收所有缓冲区
    // 这里显式释放是为了养成好习惯（万一 kernel 很长，早点释放能省内存）
    a_buf.FreeTensor(a_local);
    b_buf.FreeTensor(b_local);
    c_buf.FreeTensor(c_local);
}

}  // namespace add_operator

这段代码展示了Ascend C编程的完整模式：初始化pipe → 读Tiling参数 → 初始化缓冲区 → 分块处理 → 从HBM加载 → Vector单元计算 → 写回HBM → 释放缓冲区。所有Ascend C算子都遵循这个模式。

第五步：编译和测试

写好算子后，需要编译成动态库（.so），然后写测试用例验证正确性。

编译脚本（CMakeLists.txt）：

# CMakeLists.txt - Add 算子的构建脚本
# 这个脚本用 CMake 来构建算子工程
# 为什么要 CMake？因为算子工程包含多个源文件（kernel、tiling、proto）
# 手动编译太麻烦，用 CMake 可以一键构建

cmake_minimum_required(VERSION 3.16)
project(add_operator)

# 设置 CANN 的路径
set(ASCEND_HOME /usr/local/Ascend)
set(CMAKE_CXX_STANDARD 17)

# 添加 include 路径
include_directories(
    ${ASCEND_HOME}/include
    ${CMAKE_CURRENT_SOURCE_DIR}/op_kernel
    ${CMAKE_CURRENT_SOURCE_DIR}/op_proto
    ${CMAKE_CURRENT_SOURCE_DIR}/op_tiling
)

# 添加链接路径
link_directories(${ASCEND_HOME}/lib64)

# 编译算子动态库
# 这个动态库会被框架加载，然后调用里面的 kernel
add_library(add_operator SHARED
    op_kernel/add_kernel.cpp
    op_tiling/add_tiling.cpp
)
target_link_libraries(add_operator ascendcl opbase)

# 编译测试程序
add_executable(test_add testcases/test_add.cpp)
target_link_libraries(test_add add_operator ascendcl)

测试用例（testcases/test_add.py）：

# test_add.py - Add 算子的 Python 测试用例
# 这个脚本用 PyTorch 来验证 Add 算子的正确性
# 为什么要 Python 测试？因为 PyTorch 的接口更友好，写测试更快
# 而且 torch_npu 提供了 Ascend C 算子的调用接口

import torch
import torch_npu
import ctypes

# 1. 加载算子动态库
# 用 ctypes 加载 .so 文件，然后就能调用里面的 C 函数
add_lib = ctypes.CDLL("./build/libadd_operator.so")

# 2. 创建测试数据
# 创建两个随机张量，放到 NPU 上
a = torch.randn(1024, dtype=torch.float16).npu()
b = torch.randn(1024, dtype=torch.float16).npu()

# 3. 调用算子
# 用 torch_npu 的自定义算子接口来调用
# 这里用 torch_npu.contrib.npu_ops.custom_op 来调
c = torch.empty_like(a)
# ... 调用逻辑省略 ...

# 4. 验证正确性
# 用 PyTorch 的 Add 做基准，跟算子输出比
c_expected = a + b
if torch.allclose(c, c_expected, atol=1e-3):
    print("✓ Add operator test PASSED!")
else:
    print("✗ Add operator test FAILED!")
    print(f"Max diff: {(c - c_expected).abs().max().item()}")

跑测试：

# 编译
mkdir build && cd build
cmake ..
make -j

# 跑测试
python3 ../testcases/test_add.py
# 输出：✓ Add operator test PASSED!

常见问题排查

写第一个算子时，最容易遇到的三个问题：

问题一：编译报错，找不到ascendc命令 。原因是CANN toolkit没装好，或者环境变量没设对。解决办法：重装CANN toolkit，确认ascendc --version能跑通。

**问题二：算子加载失败，报undefined symbol。原因是动态库链接不对，某些依赖库没链进来。解决办法：检查CMakeLists.txt里的target_link_libraries，确保所有依赖都链了。

问题三：算子输出不对，跟PyTorch的结果差很多 。原因是数据类型或Tiling参数算错了。解决办法：用msprof工具抓trace，看每个block的输入输出对不对。

下一步学习方向

第一个算子跑通后，可以往这几个方向深入：

1. 进阶算子开发：写更复杂的算子（比如MatMul、Softmax），学习Cube单元的使用
2. 性能调优：学习Tiling策略调优、双缓冲、流水设计，把算子性能压榨到理论峰值的90%以上
3. 框架适配：学习怎么把自定义算子注册到PyTorch或MindSpore里，让框架能直接调

参考资源：

• cann-learning-hub仓库里的算子开发教程
• samples仓库里的100+算子示例
• CANN官方文档的Ascend C编程指南

仓库地址：https://atomgit.com/cann/samples （参考operator/add/示例）