1. 项目概述
最近在华为云开发者空间上体验了基于CANN 8.0.0.alpha003的Ascend C算子开发,实现了一个固定shape为8*2048的Add算子。这个项目让我深入理解了Ascend C编程模型和自定义算子开发的完整流程。下面我将详细分享这个Add算子的开发过程,包括环境搭建、算子设计、代码实现和部署验证等关键环节。
2. 开发环境准备
2.1 华为云开发者空间
我选择了华为云开发者空间作为开发环境,这个环境有几个显著优势:
- 预装了Ascend开发工具链
- 提供免费的Ascend NPU资源(每天2小时)
- 无需本地配置环境,开箱即用
访问地址:https://developer.huaweicloud.com/space/home
2.2 环境配置
环境关键信息:
- NPU型号:Ascend910B3
- CANN版本:8.0.0.alpha003
安装CANN工具包的命令如下:
bash复制wget https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/Milan-ASL/Milan-ASL%20V100R001C20SPC703/Ascend-cann-toolkit_8.0.0.alpha003_linux-aarch64.run
bash Ascend-cann-toolkit_8.0.0.alpha003_linux-aarch64.run --full
注意:安装前请确保系统满足最低要求,特别是内存和存储空间。安装过程可能需要10-15分钟。
3. 算子设计与实现
3.1 算子功能定义
Add算子的数学表达式很简单:
code复制z = x + y
但要在Ascend NPU上高效实现,需要考虑内存管理、并行计算等底层细节。
3.2 计算逻辑设计
Ascend C的矢量计算接口操作的是LocalTensor,因此需要设计三级流水线:
- CopyIn:将Global Memory数据搬运到Local Memory
- Compute:在Local Memory执行加法运算
- CopyOut:将结果写回Global Memory
这种设计能充分利用NPU的存储层次结构,提高计算效率。
3.3 算子原型定义
首先创建算子原型JSON文件add_custom.json:
json复制[
{
"op": "AddCustom",
"input_desc": [
{
"name": "x",
"param_type": "required",
"format": ["ND"],
"type": ["fp16"]
},
{
"name": "y",
"param_type": "required",
"format": ["ND"],
"type": ["fp16"]
}
],
"output_desc": [
{
"name": "z",
"param_type": "required",
"format": ["ND"],
"type": ["fp16"]
}
]
}
]
使用msopgen工具生成工程骨架:
bash复制/usr/local/Ascend/ascend-toolkit/latest/python/site-packages/bin/msopgen gen -i add_custom.json -c ai_core-Ascend910B3 -lan cpp -out ./addCustom
关键参数说明:
-i:算子原型定义文件路径-c:指定AI Core型号(通过npu-smi info查询)
4. 核函数实现
4.1 核函数框架
核函数是算子实现的核心,主要包含三个部分:
- 内存初始化
- 三级流水线处理
- 资源释放
基本框架如下:
cpp复制extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) {
GET_TILING_DATA(tiling_data, tiling);
KernelAdd op;
op.Init(x, y, z, tiling_data.totalLength, tiling_data.tileNum);
op.Process();
}
4.2 KernelAdd类实现
完整的KernelAdd类实现包含三个关键方法:
4.2.1 Init方法
cpp复制__aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength, uint32_t tileNum) {
this->blockLength = totalLength / AscendC::GetBlockNum();
this->tileNum = tileNum;
this->tileLength = this->blockLength / tileNum / BUFFER_NUM;
xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
zGm.SetGlobalBuffer((__gm__ DTYPE_Z*)z + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileLength * sizeof(DTYPE_Y));
pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(DTYPE_Z));
}
4.2.2 Process方法
cpp复制__aicore__ inline void Process() {
int32_t loopCount = this->tileNum * BUFFER_NUM;
for (int32_t i = 0; i < loopCount; i++) {
CopyIn(i);
Compute(i);
CopyOut(i);
}
}
4.2.3 三级流水线实现
cpp复制// CopyIn阶段
__aicore__ inline void CopyIn(int32_t progress) {
AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
AscendC::LocalTensor<DTYPE_Y> yLocal = inQueueY.AllocTensor<DTYPE_Y>();
AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
AscendC::DataCopy(yLocal, yGm[progress * this->tileLength], this->tileLength);
inQueueX.EnQue(xLocal);
inQueueY.EnQue(yLocal);
}
// Compute阶段
__aicore__ inline void Compute(int32_t progress) {
AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
AscendC::LocalTensor<DTYPE_Y> yLocal = inQueueY.DeQue<DTYPE_Y>();
AscendC::LocalTensor<DTYPE_Z> zLocal = outQueueZ.AllocTensor<DTYPE_Z>();
AscendC::Add(zLocal, xLocal, yLocal, this->tileLength);
outQueueZ.EnQue<DTYPE_Z>(zLocal);
inQueueX.FreeTensor(xLocal);
inQueueY.FreeTensor(yLocal);
}
// CopyOut阶段
__aicore__ inline void CopyOut(int32_t progress) {
AscendC::LocalTensor<DTYPE_Z> zLocal = outQueueZ.DeQue<DTYPE_Z>();
AscendC::DataCopy(zGm[progress * this->tileLength], zLocal, this->tileLength);
outQueueZ.FreeTensor(zLocal);
}
5. Host侧实现
5.1 Tiling参数定义
在add_custom_tiling.h中定义Tiling参数:
cpp复制BEGIN_TILING_DATA_DEF(TilingData)
TILING_DATA_FIELD_DEF(uint32_t, totalLength); // 总计算数据量
TILING_DATA_FIELD_DEF(uint32_t, tileNum); // 每个核上的分块数
END_TILING_DATA_DEF;
5.2 Tiling函数实现
在op_host/add_custom.cpp中实现Tiling逻辑:
cpp复制static ge::graphStatus TilingFunc(gert::TilingContext* context) {
TilingData tiling;
uint32_t totalLength = context->GetInputShape(0)->GetOriginShape().GetShapeSize();
context->SetBlockDim(BLOCK_DIM);
tiling.set_totalLength(totalLength);
tiling.set_tileNum(TILE_NUM);
tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
context->GetRawTilingData()->GetCapacity());
context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
size_t *currentWorkspace = context->GetWorkspaceSizes(1);
currentWorkspace[0] = 0;
return ge::GRAPH_SUCCESS;
}
5.3 算子注册
完整的算子注册实现:
cpp复制class AddCustom : public OpDef {
public:
explicit AddCustom(const char* name) : OpDef(name) {
this->Input("x")
.ParamType(REQUIRED)
.DataType({ ge::DT_FLOAT16 })
.Format({ ge::FORMAT_ND });
this->Input("y")
.ParamType(REQUIRED)
.DataType({ ge::DT_FLOAT16 })
.Format({ ge::FORMAT_ND });
this->Output("z")
.ParamType(REQUIRED)
.DataType({ ge::DT_FLOAT16 })
.Format({ ge::FORMAT_ND });
this->SetInferShape(ge::InferShape);
this->AICore()
.SetTiling(optiling::TilingFunc)
.AddConfig("ascend910b");
}
};
6. 编译与部署
6.1 编译算子工程
执行编译脚本:
bash复制./build.sh
编译过程会生成:
- 算子二进制文件
- 安装包(.run文件)
6.2 部署算子
运行生成的安装包:
bash复制./custom_opp_openEuler_aarch64.run
部署成功后,算子就可以在模型中使用。
7. 经验总结
在实际开发过程中,我总结了几个关键点:
-
内存管理:Ascend C采用显式内存管理,必须确保每个AllocTensor都有对应的FreeTensor,否则会导致内存泄漏。
-
流水线设计:三级流水线(CopyIn-Compute-CopyOut)是Ascend C的典型模式,合理设置BUFFER_NUM可以提高并行度。
-
调试技巧:可以使用ICPU_RUN_KF宏在CPU侧调试核函数逻辑,验证正确性后再移植到NPU。
-
性能优化:通过调整tileNum和BUFFER_NUM参数可以找到最佳性能点,需要针对具体硬件进行调优。
这个项目让我对Ascend C编程有了更深入的理解,特别是对NPU的存储层次和计算模式有了直观认识。虽然开始有些挑战,但通过官方文档和示例代码,逐步掌握了开发方法。