ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

CANN ops-math OnesLike 算子全解析:原型设计、两段式 aclnn 调用与 NPU 源码实现

CANN ops-math OnesLike 算子全解析:原型设计、两段式 aclnn 调用与 NPU 源码实现 算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载OnesLike 是 CANN ops-math 数学算子库中用于复制形状、填充全 1的基础算子它接收任意 ND 格式张量输出一个形状完全一致、所有元素均为 1 的新张量。本文以仓库中的 OnesLike 算子说明文档 与配套的 aclnnInplaceOne 接口文档 为主体结合算子定义、InferShape、Tiling 与 Kernel 源码完整讲解该算子的产品支持情况、参数原型、两段式 aclnn 调用方式并给出可编译运行的完整示例代码。读完本文你将掌握在 Atlas A2 系列产品上通过aclnnInplaceOne接口调用 OnesLike 算子的完整流程并能理解其全 1 填充在 NPU 上的底层实现原理。产品支持情况OnesLike 算子当前支持的产品如下与文档一致产品是否支持Atlas A2 训练系列产品/Atlas A2 推理系列产品√从源码看该算子在 Host 侧通过 ones_like_def.cpp 中的this-AICore().AddConfig(ascend910b)注册了 AICore 计算配置对应 Atlas A2 系列所使用的昇腾 910B 芯片架构。API 层则根据当前运行环境动态区分支持的数据类型范围详见下文数据类型支持一节。功能说明算子功能创建一个与给定张量形状相同且所有元素都为 1 的新张量。该算子与zeros_like属同一族形状克隆 常量填充算子输入张量的作用仅在于提供形状以及数据类型其元素数值不会被读取输出的每个元素都被填充为 1。从 Kernel 实现 可以看到这一填充动作在向量计算单元上通过AscendC::Duplicate指令完成——该指令会将一个标量值广播填充到整块 LocalTensor 中。算子原型设计OnesLike 算子的输入输出定义如下参数名类别描述数据类型数据格式input输入张量输入张量。FLOAT16、FLOAT32、INT32、INT8、UINT8、BF16、BOOLNDoutput输出张量与输入张量形状相同的新张量所有元素都为 1。与 input 一致ND其中 Atlas A2 训练系列产品/Atlas A2 推理系列产品支持的数据类型为 FLOAT16、FLOAT32、INT32、INT8、UINT8、BF16、BOOL。从源码验证参数定义上述原型在 ones_like_def.cpp 中通过算子注册框架声明输入xParamType(REQUIRED)必选输入支持ge::DT_FLOAT / DT_FLOAT16 / DT_INT32 / DT_UINT8 / DT_INT8 / DT_BF16 / DT_BOOL七种数据类型输出yParamType(REQUIRED)必选输出数据类型与输入一一对应输入输出均声明为FORMAT_ND含未知 shape 场景的UnknownShapeFormat并通过.AutoContiguous()声明内存自动连续化允许框架将非连续输入转为连续内存后再计算。形状推导InferShape在 ones_like_infershape.cpp 中InferShapeOnesLike将输入 shape 的维度数GetDimNum和每一维大小GetDim逐项拷贝到输出 shape因此输出与输入形状严格一致auto xShapeSize xShape-GetDimNum(); yShape-SetDimNum(xShapeSize); for (size_t i 0; i xShapeSize; i) { int64_t dim xShape-GetDim(i); yShape-SetDim(i, dim); }值得注意的是InferShape 阶段不读取输入数据的实际内容只搬运维度信息这也印证了 OnesLike形状克隆的语义。数据类型支持的两级差异虽然 README 给出的算子级支持类型为七种FLOAT16、FLOAT32、INT32、INT8、UINT8、BF16、BOOL但从源码结构看API 层还依据芯片版本做了更细的划分AICore 侧支持列表见 ones_like.cppop_api昇腾 910B / 910_93即 Atlas A2 系列DT_FLOAT、DT_FLOAT16、DT_INT8、DT_INT32、DT_UINT8、DT_BF16、DT_BOOL共 7 种其他昇腾 910 系列DT_FLOAT、DT_FLOAT16、DT_INT8、DT_INT32、DT_UINT8共 5 种不含 BF16 与 BOOL。AICPU 侧支持列表见 ones_like.hop_api还额外包含DT_INT16、DT_UINT16、DT_INT64、DT_DOUBLE、DT_COMPLEX64、DT_COMPLEX128等类型。也就是说当输入数据类型不在 AICore 支持范围内例如复数类型时API 层会自动回退到 AICPU 执行路径OnesLikeAiCPU从而在算子原型声明的七种类型之外为更多 dtype 提供兼容支持。这是该算子单原型、双执行路径的典型设计。约束说明按 README 说明该算子无约束。结合源码可以补充两点实践上的注意点维度上限API 参数校验aclnn_ones.cpp通过OP_CHECK_MAX_DIM(self, MAX_SUPPORT_DIMS_NUMS)限制最大支持 8 维超过会返回ACLNN_ERR_PARAM_INVALID非连续 TensorAPI 内部会先将输入转为连续 Tensor调用l0op::Contiguous计算完成后再通过l0op::ViewCopy将结果写回原始视图因此传入非连续 Tensor 也是被支持的接口文档中非连续 Tensor一列为 √。调用说明两段式 aclnn 接口OnesLike 算子通过aclnnInplaceOne接口对外暴露inplace 语义输入输出为同一张量self原地填充为 1。与 CANN 单算子 API 的通用约定一致该接口采用两段式调用方式先调用第一段aclnnInplaceOneGetWorkspaceSize获取 workspace 大小与执行器再调用第二段aclnnInplaceOne真正下发计算。完整函数原型摘自 aclnnInplaceOne.mdaclnnStatus aclnnInplaceOneGetWorkspaceSize( const aclTensor *self, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnInplaceOne( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)第一段接口 aclnnInplaceOneGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 Tensorself输入/输出输入/输出张量数据类型和数据格式在支持的范围之内。无FLOAT16、FLOAT32、INT32、INT8、UINT8、BF16、BOOLND0-8√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小。-----executor输出返回 op 执行器包含了算子计算流程。-----第一段接口会完成入参校验校验失败时的返回码如下返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 tensor 是空指针。ACLNN_ERR_PARAM_INVALID161002self 的数据类型和数据格式不在支持的范围之内或 self 的数据维度超过了 8 维或 self 的数据形状不一致。其他返回码的完整含义可参考 aclnn 返回码。第二段接口 aclnnInplaceOne 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnInplaceOneGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。API 内部流程从 aclnn_ones.cpp 的实现可以看到第一段接口的内部编排创建aclOpExecutor对self做空指针、数据类型、最大维度三类校验空 Tensor 直接返回workspaceSize 0调用l0op::Contiguous将输入转为连续 Tensor调用l0op::OnesLike执行实际计算内部根据芯片与 dtype 选择 AICore 或 AICPU 路径调用l0op::ViewCopy将结果写回原视图支持非连续输出返回计算所需的 workspace 大小与执行器。第二段接口aclnnInplaceOne则通过CommonOpExecutorRun将任务下发到指定 stream 上执行aclnn_ones.cpp。编译与运行完整示例以下示例代码与仓库中的 test_aclnn_ones_like.cpp 一致演示了在 shape 为{32, 4, 4, 4}、数据类型为 FLOAT32 的 Tensor 上原地填充全 1 的完整流程。编译与运行环境的具体准备步骤可参考仓库的 编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnn_ones.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } void PrintOutResult(std::vectorint64_t shape, void** deviceAddr) { auto size GetShapeSize(shape); std::vectorfloat resultData(size, 0); auto ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return); LOG_PRINT(size:%d\n, size); for (int64_t i 0; i size; i) { if (resultData[i] ! 1.f) { LOG_PRINT(error result[%ld] is: %f\n, i, resultData[i]); break; } } } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法初始化 auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); // 申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. 调用acl进行device/stream初始化 int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出 aclTensor* out nullptr; void* outDeviceAddr nullptr; std::vectorint64_t outShape {32, 4, 4, 4}; std::vectorfloat outHostData(2048, 1); ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API uint64_t workspaceSize 0; aclOpExecutor* executor; // 4. 调用aclnnInplaceOne第一段接口 ret aclnnInplaceOneGetWorkspaceSize(out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnInplaceOneGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize static_castuint64_t(0)) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 5. 调用aclnnInplaceOne第二段接口 ret aclnnInplaceOne(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnInplaceOne failed. ERROR: %d\n, ret); return ret); // 6. 同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 7. 获取输出的值将device侧内存上的结果拷贝至host侧 PrintOutResult(outShape, outDeviceAddr); // 8. 释放aclTensor aclDestroyTensor(out); // 9. 释放device资源 aclrtFree(outDeviceAddr); if (workspaceSize static_castuint64_t(0)) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); // 10. acl去初始化 aclFinalize(); return 0; }示例执行后的预期行为PrintOutResult将输出拷回 Host 并逐元素校验若发现任何元素不等于 1.0 则打印error result[i] is: x并中断全部为 1 时仅打印size:2048。整个调用链遵循标准的两段式接口范式——第一段接口完成参数校验、workspace 大小计算与执行器创建第二段接口下发计算。深入源码Tiling 与 Kernel 实现原理Tiling多核切分与 Tile 划分Tiling 逻辑位于 ones_like_tiling.cpp整体分为五步获取平台信息 → 获取 shape/属性 → 获取 workspace → 设置 tiling 数据 → 计算切分参数。关键切分策略从源码结构看核间切分默认按全部 AIV 核数coreNum切分将总元素数totalIdx平均分配到每个核。当数据量过小totalIdx * typeByte 4096 * coreNum即每个核不足 4KB 数据时会缩减使用的核数避免小算子浪费调度开销前批/尾批均衡通过formerNum / tailNum / formerLength / tailLength四个字段实现商 余数分配——余数部分的核former 核多处理 1 个元素其余核tail 核处理等量元素保证负载均衡核内 Tile 划分单核处理的数据长度可能超过 UBUnified Buffer容量因此进一步以maxTileLen ubSize / typeByte / 2预留一半 UB 空间为上限切分成多个 Tile记录formerTileNum / formerTileLength / formerLastTileLength与对应的 tail 侧参数Workspace通过GetLibApiWorkSpaceSize()获取系统 workspace 需求与用户侧 workspace 一起写入 GetWorkspaceSize 返回给框架。这些 tiling 字段在 ones_like_tiling_data.h 中统一定义为OnesLikeTilingData结构体作为 Host 与 Device 之间的数据契约。KernelDuplicate 全 1 填充Kernel 入口位于 ones_like.cppop_kernel实例化NsOnesLike::OnesLikeT后依次执行Init与Process。核心计算在 ones_like.h 的Compute中针对不同数据类型使用不同的填充方式数据类型填充方式uint8 / int8 / bool对uint16_t视图执行Duplicate(yLocal, 0x0101, (dataLength 1) / 2)两个 8bit 单元各填 1halfDuplicate(yLocal, (half)1, dataLength)int32Duplicate(yLocal, (int32_t)1, dataLength)floatDuplicate(yLocal, (float)1, dataLength)bfloat16_t对uint16_t视图执行Duplicate(yLocal, 0x3F80, dataLength)其中0x3F80即 BF16 格式的 1.0Process中的流水处理先在 UB 上填充一整块 Tile随后循环调用CopyOut将结果写回 Global Memory并在DuplicateV 单元与DataCopyMTE3 单元之间通过Sync(V, MTE3)事件同步保证访存正确性。CopyOut还针对非 32 字节对齐的尾部数据使用DataCopyPad做带参拷贝确保任意长度切分都能正确写出。测试验证仓库为该算子提供了完整的单测覆盖API 层tests/ut/op_api/test_aclnn_ones_like.cpp 验证两段式接口调用Tiling 层tests/ut/op_host/test_ones_like_tiling.cpp 验证切分参数计算Kernel 层tests/ut/op_kernel/test_oneslike.cpp 验证 Device 侧计算结果配套数据生成与比对脚本 gen_data.py 与 compare_data.py。贡献信息该算子由个人开发者 CyndiZ 于 2025/12/2 贡献至开源仓贡献内容为 OnesLike 算子适配开源仓详见 README.md 的贡献说明表。其目录结构完整覆盖了op_api对外接口、op_host原型/InferShape/Tiling、op_kernelAICore 内核、examples调用样例与tests/ut三层单测等 CANN 算子开发的标准组成可作为自定义算子接入 ops-math 开源仓的参考模板。赞分享算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载相关推荐CANN ops-math LogicalNot 算子aclnn 两段式调用接口、实现原理与源码深度解析CANN ops math LogicalNot 算子aclnn 两段式调用接口、实现原理与源码深度解析 本篇技术指南围绕 CANN 开源算子库 ops ma算子库人工智能CANNCANN ops-math 仓库 Addcmul 算子完全解析从算子原型到 aclnn 两段式调用实战CANN ops math 仓库 Addcmul 算子完全解析从算子原型到 aclnn 两段式调用实战 Addcmul 是 CANN ops math 仓库中算子库人工智能CANNRuView ESP32 CSI 传感网格从节点固件到汇聚端的分布式无摄像感知落地指南RuView ESP32 CSI 传感网格从节点固件到汇聚端的分布式无摄像感知落地指南 本文基于仓库架构决策记录 ADR 012 https://link.g人工智能算子库深度学习CANNAscend创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表