从零开始学昇腾Ascend C算子开发- 第十篇：完整算子开发流程

本文介绍了Ascend C自定义算子AddCustom的实现结构，主要包括工程目录设计和Kernel侧核心代码实现。工程采用标准CANN算子结构，分为op_kernel（核函数）、op_host（Host端）和framework（框架集成）三个目录。Kernel侧实现采用流水线设计，包含数据加载(CopyIn)、计算(Compute)和结果写回(CopyOut)三个阶段，通过双缓冲机制优化性能。核

红目香薰

1054人浏览 · 2025-12-21 20:54:05

红目香薰 · 2025-12-21 20:54:05 发布

10.1 算子工程结构

结构大致情况。
在这里插入图片描述
定义变量

具体操作：

10.1.1 目录结构

AddCustom算子工程采用标准的CANN算子工程结构，包含三个主要目录：

AddCustom/
├── op_kernel/              // Kernel侧实现
│   └── add_custom.cpp      // 核函数实现
├── op_host/                // Host侧实现
│   ├── add_custom.cpp      // Tiling函数、形状推导、算子注册
│   └── add_custom_tiling.h // Tiling数据结构定义
└── framework/              // 框架集成
    └── tf_plugin/          // TensorFlow插件
        └── tensorflow_add_custom_plugin.cc

op_kernel目录：存放核函数实现，这些代码会在AI Core上执行。核函数是算子的核心计算逻辑，使用Ascend C编写。

op_host目录：存放Host端代码，包括Tiling策略、形状推导、数据类型推导、算子注册等。这些代码在CPU上执行，负责为核函数准备参数。

framework目录：存放框架集成代码，用于将自定义算子集成到第三方框架（如TensorFlow、PyTorch）中。

10.1.2 算子描述文件

AddCustom.json是算子的描述文件，定义了算子的输入输出规格：

[
    {
        "op": "AddCustom",
        "input_desc": [
            {
                "name": "x",
                "param_type": "required",
                "format": ["ND"],
                "type": ["float16"]
            },
            {
                "name": "y",
                "param_type": "required",
                "format": ["ND"],
                "type": ["float16"]
            }
        ],
        "output_desc": [
            {
                "name": "z",
                "param_type": "required",
                "format": ["ND"],
                "type": ["float16"]
            }
        ]
    }
]

这个文件描述了AddCustom算子有两个输入（x和y），一个输出（z），都是float16类型，ND格式。msOpGen工具会根据这个文件生成算子工程的框架代码。

10.2 Kernel实现详解

10.2.1 KernelAdd类设计

op_kernel/add_custom.cpp实现了KernelAdd类，这是算子的核心计算逻辑：

class KernelAdd {
public:
    __aicore__ inline KernelAdd() {}
    
    // 初始化函数：设置GlobalTensor、初始化队列
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, 
                                uint32_t totalLength, uint32_t tileNum)
    {
        // 计算每个Core处理的数据长度
        this->blockLength = totalLength / AscendC::GetBlockNum();
        this->tileNum = tileNum;
        // 计算每个tile的长度（考虑双缓冲）
        this->tileLength = this->blockLength / tileNum / BUFFER_NUM;
        
        // 设置GlobalTensor，每个Core处理不同的数据块
        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);
        
        // 初始化队列，每个队列有BUFFER_NUM个缓冲区（双缓冲）
        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));
    }
    
    // 处理函数：执行流水线处理
    __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);  // 结果写回
        }
    }

Init函数：负责初始化，包括：

计算数据分片：根据总长度和Core数量，计算每个Core处理的数据长度
设置GlobalTensor：每个Core处理不同的数据块，通过GetBlockIdx()获取当前Core的索引
初始化队列：为每个队列分配缓冲区，BUFFER_NUM=2实现双缓冲

Process函数：执行流水线处理，循环调用CopyIn、Compute、CopyOut三个阶段。

10.2.2 数据加载（CopyIn）

__aicore__ inline void CopyIn(int32_t progress)
{
    // 从队列分配LocalTensor
    AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
    AscendC::LocalTensor<DTYPE_Y> yLocal = inQueueY.AllocTensor<DTYPE_Y>();
    
    // 从Global Memory拷贝数据到Local Memory
    AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], 
                      this->tileLength);
    AscendC::DataCopy(yLocal, yGm[progress * this->tileLength], 
                      this->tileLength);
    
    // 将LocalTensor入队
    inQueueX.EnQue(xLocal);
    inQueueY.EnQue(yLocal);
}

CopyIn阶段负责从Global Memory加载数据到Local Memory：

从队列分配LocalTensor：AllocTensor从队列的缓冲区池中分配一个LocalTensor
数据拷贝：使用DataCopy API从Global Memory拷贝到Local Memory
入队：将LocalTensor放入队列，供后续计算使用

10.2.3 计算（Compute）

__aicore__ inline void Compute(int32_t progress)
{
    // 从队列取出LocalTensor
    AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
    AscendC::LocalTensor<DTYPE_Y> yLocal = inQueueY.DeQue<DTYPE_Y>();
    
    // 分配输出LocalTensor
    AscendC::LocalTensor<DTYPE_Z> zLocal = outQueueZ.AllocTensor<DTYPE_Z>();
    
    // 执行Add计算
    AscendC::Add(zLocal, xLocal, yLocal, this->tileLength);
    
    // 将结果入队
    outQueueZ.EnQue<DTYPE_Z>(zLocal);
    
    // 释放输入LocalTensor
    inQueueX.FreeTensor(xLocal);
    inQueueY.FreeTensor(yLocal);
}

Compute阶段执行实际的计算：

出队：从输入队列取出LocalTensor
分配输出：为输出分配LocalTensor
执行计算：使用Add API进行向量加法
入队：将结果放入输出队列
释放：释放输入LocalTensor，归还到缓冲区池

10.2.4 结果写回（CopyOut）

__aicore__ inline void CopyOut(int32_t progress)
{
    // 从队列取出结果
    AscendC::LocalTensor<DTYPE_Z> zLocal = outQueueZ.DeQue<DTYPE_Z>();
    
    // 从Local Memory拷贝结果到Global Memory
    AscendC::DataCopy(zGm[progress * this->tileLength], zLocal, 
                      this->tileLength);
    
    // 释放LocalTensor
    outQueueZ.FreeTensor(zLocal);
}

CopyOut阶段将结果写回Global Memory：

出队：从输出队列取出结果LocalTensor
数据拷贝：使用DataCopy API从Local Memory拷贝到Global Memory
释放：释放LocalTensor，归还到缓冲区池

10.2.5 流水线执行

三个阶段的流水线执行实现了计算和数据传输的重叠：

时间轴：
Core 0: [CopyIn0] [Compute0] [CopyOut0] [CopyIn1] [Compute1] [CopyOut1] ...
Core 1: [CopyIn0] [Compute0] [CopyOut0] [CopyIn1] [Compute1] [CopyOut1] ...
...

双缓冲优化：
tile0: [CopyIn] -> [Compute] -> [CopyOut]
tile1:         [CopyIn] -> [Compute] -> [CopyOut]

当Compute在执行tile0的计算时，CopyIn可以同时加载tile1的数据，实现流水线并行。

10.2.6 Kernel函数

extern "C" __global__ __aicore__ void add_custom(
    GM_ADDR x, GM_ADDR y, GM_ADDR z, 
    GM_ADDR workspace, GM_ADDR tiling)
{
    // 获取Tiling数据
    GET_TILING_DATA(tiling_data, tiling);
    
    // 创建KernelAdd对象
    KernelAdd op;
    
    // 初始化
    op.Init(x, y, z, tiling_data.totalLength, tiling_data.tileNum);
    
    // 执行处理
    op.Process();
}

Kernel函数是核函数的入口点：

__global__：表示这是核函数
__aicore__：表示在AI Core上执行
参数：包括输入输出地址、workspace、tiling数据
GET_TILING_DATA：宏，用于从tiling参数中提取TilingData

10.3 Host端实现详解

10.3.1 Tiling数据结构

op_host/add_custom_tiling.h定义了Tiling数据结构：

namespace optiling {
BEGIN_TILING_DATA_DEF(TilingData)
TILING_DATA_FIELD_DEF(uint32_t, totalLength);  // 总数据长度
TILING_DATA_FIELD_DEF(uint32_t, tileNum);      // Tile数量
END_TILING_DATA_DEF;

REGISTER_TILING_DATA_CLASS(AddCustom, TilingData)
} // namespace optiling

Tiling数据结构用于在Host和Kernel之间传递参数：

totalLength：输入数据的总元素数
tileNum：将数据分成多少个tile处理
BEGIN_TILING_DATA_DEF/END_TILING_DATA_DEF：宏，用于定义Tiling数据结构
REGISTER_TILING_DATA_CLASS：注册Tiling数据类

10.3.2 Tiling函数

op_host/add_custom.cpp中的TilingFunc函数负责计算Tiling参数：

namespace optiling {
const uint32_t BLOCK_DIM = 8;  // 使用8个Core
const uint32_t TILE_NUM = 8;   // 每个Core处理8个tile

static ge::graphStatus TilingFunc(gert::TilingContext *context)
{
    // 创建TilingData对象
    TilingData tiling;
    
    // 获取输入形状的总元素数
    uint32_t totalLength = context->GetInputShape(0)
                          ->GetOriginShape()
                          .GetShapeSize();
    
    // 设置Block维度（Core数量）
    context->SetBlockDim(BLOCK_DIM);
    
    // 设置Tiling参数
    tiling.set_totalLength(totalLength);
    tiling.set_tileNum(TILE_NUM);
    
    // 保存Tiling数据到缓冲区
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), 
                       context->GetRawTilingData()->GetCapacity());
    
    // 设置实际使用的数据大小
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    
    // 设置workspace大小（Add算子不需要workspace，设为0）
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    
    return ge::GRAPH_SUCCESS;
}
} // namespace optiling

Tiling函数的作用：

获取输入形状：从context中获取输入tensor的形状信息
计算Tiling参数：根据输入大小和策略，计算totalLength和tileNum
设置Block维度：告诉框架使用多少个Core并行执行
保存Tiling数据：将Tiling参数保存到缓冲区，传递给Kernel
设置workspace：如果算子需要临时内存，在这里设置大小

10.3.3 形状推导

namespace ge {
static graphStatus InferShape(gert::InferShapeContext *context)
{
    // 获取输入形状
    const gert::Shape *x1_shape = context->GetInputShape(0);
    
    // 获取输出形状
    gert::Shape *y_shape = context->GetOutputShape(0);
    
    // 输出形状等于输入形状（Add算子的输出形状和输入相同）
    *y_shape = *x1_shape;
    
    return GRAPH_SUCCESS;
}
} // namespace ge

InferShape函数用于推导输出tensor的形状。对于Add算子，输出形状和输入形状相同。框架在编译图时会调用这个函数，用于验证图的正确性和分配内存。

10.3.4 数据类型推导

namespace ge {
static graphStatus InferDataType(gert::InferDataTypeContext *context)
{
    // 获取输入数据类型
    const auto inputDataType = context->GetInputDataType(0);
    
    // 输出数据类型等于输入数据类型
    context->SetOutputDataType(0, inputDataType);
    
    return GRAPH_SUCCESS;
}
} // namespace ge

InferDataType函数用于推导输出tensor的数据类型。对于Add算子，输出数据类型和输入相同。

10.3.5 算子注册

namespace ops {
class AddCustom : public OpDef {
public:
    explicit AddCustom(const char *name) : OpDef(name)
    {
        // 定义输入x
        this->Input("x")
            .ParamType(REQUIRED)           // 必需参数
            .DataType({ge::DT_FLOAT16})    // 支持float16
            .Format({ge::FORMAT_ND});       // 支持ND格式
        
        // 定义输入y
        this->Input("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND});
        
        // 定义输出z
        this->Output("z")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND});
        
        // 设置形状和数据类型推导函数
        this->SetInferShape(ge::InferShape)
            .SetInferDataType(ge::InferDataType);
        
        // 配置AI Core执行
        this->AICore()
            .SetTiling(optiling::TilingFunc)  // 设置Tiling函数
            .AddConfig("ascend910")           // 支持ascend910
            .AddConfig("ascend310p")          // 支持ascend310p
            .AddConfig("ascend310b")           // 支持ascend310b
            .AddConfig("ascend910b");         // 支持ascend910b
    }
};

// 注册算子
OP_ADD(AddCustom);
} // namespace ops

算子注册是算子开发的关键步骤：

OpDef类：所有自定义算子都继承自OpDef
Input/Output：定义算子的输入输出规格
ParamType：参数类型，REQUIRED表示必需，OPTIONAL表示可选
DataType：支持的数据类型列表
Format：支持的数据格式列表
SetInferShape/SetInferDataType：设置推导函数
AICore配置：设置Tiling函数和支持的芯片型号
OP_ADD宏：注册算子到框架

10.4 框架集成

10.4.1 TensorFlow插件

framework/tf_plugin/tensorflow_add_custom_plugin.cc实现了TensorFlow框架集成：

#include "register/register.h"

namespace domi {
// 注册算子信息到GE（Graph Engine）
REGISTER_CUSTOM_OP("AddCustom")
    .FrameworkType(TENSORFLOW)        // 框架类型：TensorFlow
    .OriginOpType("AddCustom")        // TensorFlow中的算子名称
    .ParseParamsByOperatorFn(AutoMappingByOpFn);  // 参数映射函数
}  // namespace domi

这个文件将AddCustom算子注册到TensorFlow框架：

REGISTER_CUSTOM_OP：注册自定义算子
FrameworkType：指定框架类型
OriginOpType：TensorFlow中使用的算子名称
ParseParamsByOperatorFn：参数映射函数，AutoMappingByOpFn会自动映射同名参数

10.4.2 框架集成流程

当TensorFlow调用AddCustom算子时：

TensorFlow识别到AddCustom算子
通过插件找到对应的CANN算子
GE（Graph Engine）调用Tiling函数计算参数
调用InferShape和InferDataType推导形状和类型
编译生成执行图
运行时调用Kernel函数执行计算

10.5 编译和部署流程

10.5.1 使用msOpGen生成工程

msOpGen是CANN提供的算子工程生成工具，可以根据JSON描述文件生成算子工程框架：

# 使用AddCustom.json生成算子工程
msOpGen -i AddCustom.json -o CustomOp

生成的工程包含：

目录结构：op_kernel、op_host、framework等
框架代码：基础的Kernel函数、Tiling函数框架
编译脚本：CMakeLists.txt等

10.5.2 实现算子逻辑

生成工程后，需要实现：

Kernel逻辑：在op_kernel/add_custom.cpp中实现KernelAdd类
Tiling策略：在op_host/add_custom.cpp中实现TilingFunc
形状推导：实现InferShape和InferDataType
算子注册：定义AddCustom类并注册

10.5.3 编译算子

# 进入算子工程目录
cd CustomOp

# 创建构建目录
mkdir build && cd build

# 配置CMake
cmake .. -DCMAKE_INSTALL_PREFIX=/path/to/install

# 编译
make -j8

# 安装
make install

编译会生成：

算子库：libcust_opapi.so
头文件：aclnn_add_custom.h等
安装包：custom_opp_*.run

10.5.4 部署算子

# 安装算子包
./custom_opp_ubuntu_x86_64.run

# 验证安装
ls $ASCEND_OPP_PATH/vendors/customize/op_api/

安装后，算子会被部署到ASCEND_OPP_PATH/vendors/customize/op_api/目录，可以在应用中使用。

10.6 关键设计要点

10.6.1 数据分片策略

AddCustom使用简单的均匀分片策略：

总数据长度除以Core数量，得到每个Core处理的数据长度
每个Core再分成多个tile，实现流水线处理
双缓冲（BUFFER_NUM=2）实现计算和数据传输重叠

对于更复杂的算子，可能需要：

考虑数据对齐
处理边界情况
优化内存访问模式

10.6.2 流水线设计

三个阶段的流水线设计：

CopyIn：数据加载，可以和其他阶段并行
Compute：计算，是性能瓶颈
CopyOut：结果写回，可以和其他阶段并行

双缓冲确保每个阶段都有数据可处理，最大化硬件利用率。

10.6.3 多核并行

通过GetBlockIdx()和GetBlockNum()实现多核并行：

每个Core处理不同的数据块
Tiling函数设置BLOCK_DIM=8，使用8个Core
数据自动分配到不同Core

10.6.4 内存管理

使用TPipe和TQue管理Local Memory：

TPipe：管理整个流水线的内存
TQue：管理队列的缓冲区池
AllocTensor/FreeTensor：从缓冲区池分配和释放

这种设计避免了手动内存管理，减少了内存碎片。

10.7 与简化版的对比

10.7.1 AddCustom vs AddCustomTiny

AddCustom是完整的算子工程，包含：

完整的Host端实现（Tiling、形状推导、算子注册）
框架集成（TensorFlow插件）
完整的编译和部署流程

AddCustomTiny是极简版本，只包含：

基本的Kernel实现
简单的Tiling函数
算子注册

适合快速验证和学习。

10.7.2 适用场景

使用AddCustom：

生产环境
需要框架集成
需要完整的错误处理
需要支持多种芯片

使用AddCustomTiny：

学习理解
快速原型
简单算子

学习检查点

学完这一篇，你应该能做到这些：

理解算子工程的完整结构，知道每个目录和文件的作用。掌握Kernel实现的三个阶段（CopyIn、Compute、CopyOut），理解流水线设计。理解Tiling函数的作用，能够根据输入计算Tiling参数。掌握形状和数据类型推导，能够实现InferShape和InferDataType。理解算子注册机制，能够使用OpDef注册自定义算子。了解框架集成方式，知道如何将算子集成到TensorFlow等框架。掌握编译和部署流程，能够独立完成算子的编译和安装。