ARTICLE DETAIL

资讯详情

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

CANN ops-math 算子解析:Round 四舍五入算子的功能、aclnn 调用与源码实现

CANN ops-math 算子解析:Round 四舍五入算子的功能、aclnn 调用与源码实现 CANN ops-math 算子解析Round 四舍五入算子的功能、aclnn 调用与源码实现【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读Round 是 CANN ops-math 仓库中位于experimental/math/round目录下的一类数学基础算子用于对输入张量的每一个元素按指定小数位执行四舍五入运算。本文以该算子目录下的 README.md 为骨架结合仓库内的接口定义、Host 侧 Tiling 逻辑、Kernel 侧计算实现与单元测试完整梳理 Round 算子的功能语义、参数约束、aclnn 两段式调用方式以及底层实现原理帮助开发者快速理解并上手使用该算子。功能说明Round 算子的功能是对输入张量x的每一个元素按照指定的小数位数decimals进行四舍五入round half away from zero 语义输出与输入形状一致的张量y。其计算公式为$$ out_i round(input_i, decimals) $$其中decimals表示四舍五入后保留的小数位数。README 中给出的示例Round(3.56, 0) 4.0不保留小数3.56 四舍五入为 4Round(3.56, 1) 3.5保留 1 位小数3.56 四舍五入为 3.6 还是 3.5 取决于四舍五入的取整规则README 以3.5作为示例输出实际行为以 Kernel 实现为准。从底层实现看decimals的语义在 Tiling 阶段被转换为缩放因子10^decimalsKernel 先将输入放大10^decimals倍执行AscendC::Round取整再缩小同样倍数回写从而实现对任意小数位的四舍五入见 round_tiling.cpp 中tiling-decimals pow(10, decimals)的处理。产品支持情况产品是否支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品√在算子定义中这一支持范围与硬件配置项对应round_def.cpp中通过this-AICore().AddConfig(ascend910b)注册了ascend910b的 AICore 配置见 round_def.cpp。可以推断Atlas A2 系列对应ascend910b平台其余平台暂未开放该算子。参数说明参数名输入/输出/属性描述数据类型数据格式x输入入参公式中的 input对该数据进行四舍五入bfloat16float16floatint32NDdecimals属性入参公式中的 decimals指定四舍五入的位数intNDy输出待进行 round 计算的出参公式中的 ybfloat16float16floatint32ND参数定义与默认值的源码佐证在 round_def.cpp 中x与y均声明为REQUIRED必选参数支持的数据类型为ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_INT32数据格式均为FORMAT_ND属性decimals为OPTIONAL可选且默认值为0。在 aclnn_round.cpp 中aclnnRound接口路径支持的数据类型列表更宽ASCEND910支持 float / float16 / double / int32 / int64ASCEND910B额外支持 bf16其中double、int64是 aclnn 接口层经 Cast 等预处理后间接支持的类型算子原生计算类型以round_def.cpp的声明为准。输入与输出必须形状一致、数据类型一致接口层通过CheckSameShape1In1Out与OP_CHECK_DTYPE_NOT_MATCH进行校验见 aclnn_round.cppHost 侧 round_infershape.cpp 也直接将输出 shape 的维度数与每一维大小拷贝自输入即输出y与输入x形状完全一致0–8 维见 aclnnRound 文档。约束说明原文档明确说明 Round 算子无约束。需要说明的边界如下当输入张量为空时aclnnRoundGetWorkspaceSize会直接返回workspaceSize 0并成功退出不会触发实际计算见 aclnn_round.cpp数据格式层面仅支持非私有格式接口层CheckFormat会拒绝私有格式支持 ND、NCHW、NHWC、HWCN、NDHWC、NCDHW 等公共格式见 aclnn_round.cppREADME 中数据格式统一为 ND若输入为 int32 类型四舍五入操作退化为原样拷贝因为整型无需取整见 round.h。调用说明Round 算子通过 CANN 标准的两段式 aclnn 接口aclnnRound调用完整可运行样例位于 test_aclnn_round.cpp接口说明详见 aclnnRound.md。调用方式调用样例说明aclnn 调用test_aclnn_round.cpp通过aclnnRound接口方式调用 Round 算子两段式接口函数原型每个 aclnn 算子都遵循两段式接口two-phase API模式先调用aclnnRoundGetWorkspaceSize获取计算所需 workspace 大小及包含算子计算流程的执行器再调用aclnnRound执行计算。aclnnStatus aclnnRoundGetWorkspaceSize( const aclTensor* self, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor ) aclnnStatus aclnnRound( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)第一段接口 aclnnRoundGetWorkspaceSize 参数参数名输入/输出描述使用说明数据类型数据格式维度shapeself输入待进行 round 计算的入参公式中的 self无bfloat16float16floatint32ND0-8out输出待进行 round 计算的出参公式中的 outshape 与 self 相同bfloat16float16floatint32ND0-8workspaceSize输出返回需要在 Device 侧申请的 workspace 大小----executor输出返回 op 执行器包含了算子计算流程----第二段接口 aclnnRound 参数参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnRoundGetWorkspaceSize获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream两段接口均返回aclnnStatus状态码用于判断调用是否成功。调用流程示例以下代码摘自 test_aclnn_round.cpp 的核心调用逻辑仅保留关键步骤#include acl/acl.h #include aclnn_round.h int main() { // 1. 调用 acl 进行 device/stream 初始化 int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); // 2. 构造输入与输出形状为 {2, 8} 的 float 张量随机初始化 aclTensor* selfX nullptr; void* selfXDeviceAddr nullptr; std::vectorint64_t selfXShape {2, 8}; std::vectorfloat selfXHostData(16, 1); for (int i 0; i 16; i) { selfXHostData[i] (static_castfloat(rand()) / RAND_MAX) * 10.0f; } ret CreateAclTensor(selfXHostData, selfXShape, selfXDeviceAddr, aclDataType::ACL_FLOAT, selfX); aclTensor* out nullptr; void* outDeviceAddr nullptr; std::vectorint64_t outShape {2, 8}; std::vectorfloat outHostData(16, 1); ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); // 3. 调用第一段接口获取 workspace 大小与执行器 uint64_t workspaceSize 0; aclOpExecutor* executor; ret aclnnRoundGetWorkspaceSize(selfX, out, workspaceSize, executor); // 4. 按需申请 workspace 内存 void* workspaceAddr nullptr; if (workspaceSize static_castuint64_t(0)) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 5. 调用第二段接口执行计算 ret aclnnRound(workspaceAddr, workspaceSize, executor, stream); // 6. 同步等待任务执行结束 ret aclrtSynchronizeStream(stream); // 7. 将结果从 device 拷贝到 host 并打印 PrintOutResult(outShape, outDeviceAddr); // 8. 释放 aclTensor 与 device 资源 aclDestroyTensor(selfX); aclDestroyTensor(out); aclrtFree(selfXDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize static_castuint64_t(0)) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }整体调用链路可归纳为acl 初始化 → 构造输入/输出 aclTensor → 两段式接口GetWorkspaceSize 执行→ stream 同步 → 结果回拷 → 资源释放。示例中CreateAclTensor使用aclCreateTensor以 ND 格式、连续 strides 创建张量并完成 host 到 device 的数据拷贝是使用 aclnn 接口的标准前置步骤。源码实现纵深接口层参数校验与 l0 算子编排aclnn_round.cpp 实现了完整的 aclnn 接口逻辑CheckParamsRound依次校验参数非空CheckNotNull2Tensor、dtype 在支持列表内CheckDtypeValid、self 与 out 形状一致、格式非私有格式第一段接口内部通过l0op::Contiguous将非连续输入转为连续张量再调用l0op::RoundDecimals(selfContiguous, 0, executor)aclnnRound固定decimals 0最后用l0op::ViewCopy将连续结果写回可能非连续的out同一文件还提供了带小数位的变体接口aclnnRoundDecimalsdecimals作为显式入参以及两个原地inplace版本aclnnInplaceRound/aclnnInplaceRoundDecimals原地版本直接将selfRef同时作为输出接口层原型见 round.h。Host 侧shape 推导与 Tilingshape 推导InferShapeRound将输入 shape 的维数与各维大小逐一拷贝到输出保证y与x形状一致见 round_infershape.cpp。Tiling 计算TilingFunc根据平台信息获取 UB 大小与可用 AI Core 数量结合输入总元素数、每个元素字节数typeLength按照数据块block32 字节→ 每个核 → 每个 tile的层次拆分任务并依据不同数据类型与是否带decimals决定 UB 内变量占用数量int32 为 4float16/bf16 无 decimals 为 8、有 decimals 为 10float 无 decimals 为 5、有 decimals 为 6同时采用BUFFER_NUM 2的 double buffer 优化见 round_tiling.cppTiling 结果通过 round_tiling_data.h 中的RoundTilingData结构体传给 Kernel其中包含大核/小核数据量、tile 数据量、尾块数据量、搬运次数与decimals缩放因子。Kernel 侧按数据类型分派的计算逻辑round.h 中的KernelRound::Compute展示了核心计算分支int32直接DataCopy原样拷贝无需取整float16 / bf16先Cast到 float 中间缓冲区若decimals 1.0f则Duplicate生成缩放向量并执行Mul → Round → Div否则直接Round最后以CAST_RINT模式Cast回原类型float与 float16/bf16 类似但直接复用输入缓冲区完成Mul → Round → Div减少一次中间缓冲区开销。Kernel 的CopyIn/Compute/CopyOut采用流水式循环处理最后一轮以tailDataNum处理尾块见 round.h入口函数 round.cpp 根据 TilingKey 实例化KernelRound并执行Init与Process。测试验证仓库为 Round 算子提供了三层单元测试op_api 层test_aclnn_round.cpp 验证 aclnn 接口的调用正确性op_host 层test_round_tiling.cpp 配合 round_tiling.h 验证 Tiling 数据切分逻辑op_kernel 层test_round.cpp 验证 Kernel 计算结果测试数据由 gen_data.py 生成、compare_data.py 完成对比构建入口见 CMakeLists.txt。贡献说明贡献者贡献方贡献算子贡献时间贡献内容番茄土豆南京亚信软件有限公司Round2025/12/11Round 算子适配开源仓小结Round 算子是 ops-math 中实现简单但链路完整的典型数学算子功能上支持按任意小数位对 bfloat16 / float16 / float / int32 张量逐元素四舍五入使用上遵循标准 aclnn 两段式接口可参考 test_aclnn_round.cpp 快速接入实现上则贯穿 OpDef 注册、shape 推导、Tiling 切分到 AscendC Kernel 的全栈链路decimals属性通过10^decimals缩放因子在 Kernel 内以放大—取整—缩小完成值得作为了解 CANN 数学类基础算子开发与调用模式的入门范本。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表