CANN/ops-tensor算子开发指南

算子开发指南【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor目录结构开发一个算子需要以下文件src/ ├── add/ # 示例Add 算子 │ ├── add.cpp # Host Kernel 实现 │ ├── add_struct.h # Tiling 数据结构可选 │ ├── CMakeLists.txt # 编译配置 │ └── tests/ # 测试目录强烈推荐 │ ├── add_test.h │ └── add_test.cpp ├── ... # 其他算子 └── CMakeLists.txt说明Host 和 Kernel 可以合并为一个.cpp文件TilingData 可以定义在.cpp文件中也可以独立为_struct.harch35/目录仅在需要区分不同 SOC 架构时使用当前版本仅支持 Ascend950即 arch35测试文件强烈推荐但不是必需的文件说明1. 解决方案实现op_solution.cpp作用实现解决方案执行函数框架会自动调用计算 Tiling 参数管理设备内存调用核函数注册解决方案到全局注册表基本结构#include acl/acl.h #include kernel_operator.h #define GM_ADDR uint8_t* // Tiling 数据结构 namespace OpNameOp { struct OpNameTilingData { int64_t totalLength; int64_t usedCoreNum; // ... 其他 Tiling 参数 }; } // 解决方案执行函数框架自动调用 /* 1. 参数获取和检查 2. 计算 Tiling 数据 3. 分配设备TilingData内存并拷贝 Tiling 数据 4. 调用核函数 5. 同步异步执行安全 6. 释放TilingData设备内存 */ static acltensorStatus_t ExecuteOpNameSolution(const ElementwiseArgs args) { // 1. 参数获取和检查 const void* A args.bufferA; const void* C args.bufferC; void* D args.bufferD; if (A nullptr || ...) { ... // 异常返回 } // 2. 计算 Tiling 数据可用函数封装 OpNameOp::OpNameTilingData tilingData; tilingData.totalLength size; tilingData.usedCoreNum CalculateCoreNum(size); // 自定义函数 // ... 其他 Tiling 参数 // 3. 分配设备TilingData内存并拷贝 Tiling 数据 uint8_t *tilingDevice; aclrtMalloc((void**)tilingDevice, sizeof(tilingData), ACL_MEM_MALLOC_HUGE_FIRST); aclrtMemcpy(tilingDevice, sizeof(tilingData), tilingData, sizeof(tilingData), ACL_MEMCPY_HOST_TO_DEVICE); // 4. 调用核函数 Block, workspace, stream op_name_kernel_do(inputDevice, outputDevice, tilingDevice, nullptr, tilingData.usedCoreNum, stream); // 5. 同步 aclrtSynchronizeStream(stream); // 6. 释放设备内存 aclrtFree(tilingDevice); return ACL_SUCCESS; } // Kernel 部分核函数实现 using namespace AscendC; extern C __global__ __aicore__ void op_name(GM_ADDR input, GM_ADDR output, GM_ADDR tiling) { // Kernel 类型声明 KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); // 初始化 TPipe pipe; // ... 初始化 LocalTensor、GlobalTensor、TQue 等 // 解析 Tiling 数据 auto tilingData (OpNameOp::OpNameTilingData*)tiling; // 核心计算逻辑 for (int i 0; i tilingData-blockLoopCnt; i) { // 1. DataCopy: GM - LocalTensor // 2. 计算 // 3. DataCopy: LocalTensor - GM } } // 核函数封装 void op_name_kernel_do(GM_ADDR input, GM_ADDR output, GM_ADDR tiling, GM_ADDR workspace, uint32_t numBlocks, void *stream) { op_namenumBlocks, workspace, stream(input, output, tiling); } // 注册解决方案到全局注册表 namespace { std::shared_ptrElementwiseSolution CreateOpNameSolution() { SolutionUid uid{ACLTENSOR_OP_OP_NAME, ACLTENSOR_R_32F, 0}; // 0 表示通用维度 return std::make_sharedElementwiseSolution(uid, ExecuteOpNameSolution); } struct OpNameSolutionRegistrar { OpNameSolutionRegistrar() { auto solution CreateOpNameSolution(); ElementwiseSolutionRegistry::instance().registerSolution(solution); } }; static OpNameSolutionRegistrar g_op_nameSolutionRegistrar; }解决方案关键点定义 TilingData 结构体或 include 独立的_struct.h计算 Tiling 参数实现解决方案执行函数签名static acltensorStatus_t ExecuteOpNameSolution(const ElementwiseArgs args)注册解决方案到全局注册表使用静态注册器自动注册Kernel 部分关键点(Kernel部分可提取成独立的_kernel.cpp)使用__global__ __aicore__标记核函数实现 GM ↔ LocalTensor 的数据搬运实现核心计算逻辑注册机制说明SolutionUid- 解决方案唯一标识符由{操作符, 数据类型, 维度数}三元组组成ACLTENSOR_OP_OP_NAME- 操作符类型ADD、SUB、MUL、DIV 等ACLTENSOR_R_32F- 数据类型当前仅支持 float320- 维度数0 表示通用解决方案适配任意维度Create Solution- 创建解决方案对象返回std::shared_ptrElementwiseSolution智能指针传入ExecuteOpNameSolution函数指针作为执行接口静态注册器- 利用全局变量初始化自动注册程序启动时自动执行将解决方案注册到全局单例ElementwiseSolutionRegistry框架在执行时根据操作符、数据类型、维度数从注册表查询解决方案2. Tiling 数据结构可选文件op_name_struct.h作用定义 Host 传递给 Kernel 的 Tiling 参数基本结构#ifndef OP_NAME_STRUCT_H #define OP_NAME_STRUCT_H #include cstdint namespace OpNameOp { struct OpNameTilingData { int64_t totalLength; int64_t usedCoreNum; int64_t blockFormer; int64_t blockLoopCnt; int64_t blockTail; // ... 其他 Tiling 参数 }; } // namespace OpNameOp #endif说明这是一个简单的 C 结构体只包含基本数据类型int64_t 等Host 计算参数Kernel 读取参数也可以直接定义在.cpp文件中3. 编译配置文件CMakeLists.txtregister_operator( NAME op_name ARCH_DIR arch35 # 当前版本仅支持 arch35 (Ascend950) )4. 测试文件强烈推荐参见 测试编写指南。说明虽然测试文件不是必需的但强烈建议为每个算子编写单元测试以确保算子实现的正确性。开发流程步骤 1创建目录和文件mkdir -p src/op_name/tests touch src/op_name/op_name.cpp touch src/op_name/CMakeLists.txt可选独立的 struct 文件touch src/op_name/op_name_struct.h可选测试文件touch src/op_name/tests/op_name_test.h touch src/op_name/tests/op_name_test.cpp步骤 2编写解决方案实现在op_name_solution.cpp中定义 TilingData 结构体或 include 独立的_struct.h实现解决方案执行函数ExecuteOpNameSolution计算并注册解决方案到全局注册表在op_name_kernel.cpp中实现核函数使用__global__ __aicore__步骤 3配置编译在CMakeLists.txt中注册算子。步骤 4编写测试推荐参考 测试编写指南。步骤 5编译验证./build.sh --opsop_name --run关键概念解决方案 vs Kernel层面运行位置职责解决方案CPUTiling 计算、内存管理、调用 Kernel、注册到全局表KernelNPU AI Core实际计算逻辑Tiling目的将大任务切分成适合 NPU 执行的小块关键参数usedCoreNum- 使用多少个 AI CoreblockFormer- 每次迭代处理多少数据blockLoopCnt- 每个核迭代多少次计算原则充分利用 AI Core 并行能力数据不超过 Unified Buffer 容量对齐到 32 字节边界 核函数调用kernel_funcnumBlocks, workspace, stream(args...);参数说明numBlocks- 使用多少个 AI CoreBlockworkspace- 共享内存指针通常设置为nullptrstream- ACL 执行流完整示例参见src/add/目录add_solution.cpp- 解决方案实现Tiling 计算、解决方案执行函数、注册add_kernel.cpp- Kernel 端实现核函数逻辑arch35/add_struct.h- Tiling 数据结构tests/add_test.cpp- 单元测试CMakeLists.txt- 编译配置相关文档测试编写指南build 参数说明【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考