当前位置：首页 > news >正文

CANN-Ascend-C入门-昇腾NPU上写第一个自定义算子

news 2026/5/24 2:06:33

CANN-Ascend-C入门-昇腾NPU上写第一个自定义算子

Ascend C 是昇腾NPU的算子开发语言，语法类似 CUDA C。如果你有 GPU 算子开发经验，迁移成本约 2-3 天。这篇从零写一个向量加法算子，走通编译、注册、调用的全流程。

环境准备

# CANN 开发环境exportASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest# 验证编译器可用whichascendc# 应该输出 ascendc 的路径

第一步：写 Kernel

// add_custom_kernel.cpp#include"kernel_operator.h"classAddCustomKernel{public:__aicore__inlinevoidInit(GM_ADDR x,GM_ADDR y,GM_ADDR z,uint32_ttotal_len){// 昇腾NPU的全局内存映射x_gm_.SetGlobalBuffer((__gm__ half*)x,total_len);y_gm_.SetGlobalBuffer((__gm__ half*)y,total_len);z_gm_.SetGlobalBuffer((__gm__ half*)z,total_len);// 计算分块参数：每个 Vector 单元处理 256 个元素// 这样 32 个 Vector 单元可以并行处理 8192 个元素block_len_=256;tiles_=(total_len+block_len_-1)/block_len_;}__aicore__inlinevoidProcess(){for(int32_ti=0;i<tiles_;i++){// 1. 从 HBM 读入数据到 Vector 本地缓存LocalTensor<half>x_local=x_buf_.Get<half>();LocalTensor<half>y_local=y_buf_.Get<half>();LocalTensor<half>z_local=z_buf_.Get<half>();DataCopy(x_local,x_gm_[i*block_len_],block_len_);DataCopy(y_local,y_gm_[i*block_len_],block_len_);// 2. Vector 单元做加法Add(z_local,x_local,y_local,block_len_);// 3. 写回 HBMDataCopy(z_gm_[i*block_len_],z_local,block_len_);}}private:TPipe pipe_;TBuf<QuePosition::VECCALC>x_buf_,y_buf_,z_buf_;GlobalTensor<half>x_gm_,y_gm_,z_gm_;uint32_tblock_len_;uint32_ttiles_;};// 算子入口函数（必须用这个签名）extern"C"__global__voidadd_custom_kernel(GM_ADDR x,GM_ADDR y,GM_ADDR z,GM_ADDR workspace,GM_ADDR tiling){// 从 Tiling 数据获取元素总数uint32_ttotal_len=*reinterpret_cast<uint32_t*>(tiling);AddCustomKernel op;op.Init(x,y,z,total_len);op.Process();}

几个关键概念：

__aicore__：函数跑在 AI Core 上，不是 CPU
GM_ADDR：HBM 全局内存地址
LocalTensor：Vector 单元的片上缓存
DataCopy：DMA 数据搬运
Add：Vector 单元的向量加法

第二步：写 Tiling 函数

// add_custom_tiling.cpp#include"op_host.h"ge::graphStatusAddCustomTiling(constge::Operator&op,TilingContext*ctx){autox_shape=ctx->GetInputShape(0);uint32_ttotal_len=x_shape.GetShapeSize();ctx->SetTilingData("total_len",total_len);returnge::GRAPH_SUCCESS;}

这个例子不需要复杂的分块策略——只是传递元素总数。复杂算子的 Tiling 需要根据输入 shape 计算 Cube 单元的分块大小。

第三步：注册算子

// add_custom.cpp#include"op_host.h"namespacege{IMPL_OP(AddCustom).Inputs({"x","y"}).Outputs({"z"}).Tiling(AddCustomTiling);}

第四步：编译

# 编译 kernelascendc--chip=Ascend910B2 add_custom_kernel.cpp-obuild/# 编译 host 代码g++-shared-fPICadd_custom.cpp add_custom_tiling.cpp\-I${ASCEND_HOME_PATH}/include\-L${ASCEND_HOME_PATH}/lib64-lopapi\-obuild/libadd_custom.so

第五步：在 PyTorch 中调用

importtorchimporttorch_npu# 手动调用自定义算子x=torch.randn(1024,device="npu",dtype=torch.float16)y=torch.randn(1024,device="npu",dtype=torch.float16)z=torch_npu.npu.add_custom(x,y)# 验证结果expected=x+yasserttorch.allclose(z,expected,atol=1e-3)

和 CUDA C 的差异

概念	CUDA C	Ascend C
入口函数	`__global__ void kernel(...)`	`extern "C" __global__ void kernel(...)`
线程索引	`threadIdx.x`	无直接等价（用 tile 循环）
共享内存	`__shared__`	`LocalTensor`
全局内存	指针	`GM_ADDR`+`SetGlobalBuffer`
同步	`__syncthreads()`	`pipe_barrier()`
计算单元	CUDA Core	Cube (MatMul) + Vector (逐元素)