CANN ops-math 算子实战：IsFinite 有限性判断算子的接口、源码实现与调用示例

CANN ops-math 算子实战IsFinite 有限性判断算子的接口、源码实现与调用示例【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本文基于 CANN ops-math 开源仓库的 IsFinite 算子文档 展开系统讲解数学类基础算子 IsFinite 的功能定义、产品支持矩阵、输入输出约束并结合仓库中的算子注册、图 IR 定义、aclnn 两段式 API 实现与 kernel 源码深入剖析其底层计算流程。读者学完后既能通过 aclnnIsFinite 接口在单算子aclnn与图模式GEIR两种方式下完成调用与编译运行也能理解该算子在 NPU 上的完整实现链路与工程化细节。算子功能与计算公式IsFinite 是 CANN ops-math 仓库math/is_finite 目录中的数学类基础算子其功能是逐元素判断输入张量中哪些元素是有限数值——即该元素既不是inf、-inf也不是nan非数。该功能在深度学习训练与推理中常用于数值稳定性检查例如在梯度计算、Loss 归一化之前过滤非法数值。计算公式如下$$ y_i(x_i \neq \pm inf) : and : (x_i \neq nan) $$即对输入张量x的每一个元素x_i同时满足不等于正负无穷且不等于 NaN时输出对应位置的y_i为True1否则为False0。输出张量y与输入x形状一致数据类型固定为 BOOL。从 op_graph/is_finite_proto.h 中的算子 IR 定义可以看到该算子在 GE 图中的注册同时声明了与 TensorFlowIsFinite算子的兼容性输入支持 1~8 维 ND 张量REG_OP(IsFinite) .INPUT(x, TensorType({DT_BF16, DT_FLOAT16, DT_FLOAT, DT_DOUBLE, DT_BOOL, DT_UINT8, DT_INT8, DT_UINT16, DT_INT16, DT_INT32, DT_UINT32, DT_UINT64, DT_INT64})) .OUTPUT(y, TensorType({DT_BOOL})) .OP_END_FACTORY_REG(IsFinite)产品支持情况原文档给出了 IsFinite 算子的产品支持矩阵产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品√Kirin X90 处理器系列产品√Kirin 9030 处理器系列产品√从 op_host/is_finite_def.cpp 的算子注册源码可以印证上述支持情况注册的 AICore 配置包含ascend910b对应 Atlas A2 系列、ascend910_93对应 Atlas A3 系列、ascend950对应 Ascend 950 系列、ascend350、ascend310pAtlas 推理系列以及kirinx90、kirin9030Kirin 处理器系列并分别针对不同架构配置了动态编译、动态 shape、动态 rank 等能力标志。同时源码中针对不同产品对数据类型的支持做了差异化配置详见下文数据类型约束。参数说明IsFinite 算子有两个张量参数均为 ND 格式参数名输入/输出/属性描述数据类型数据格式x输入公式中的输入张量 xFLOAT、FLOAT16、BFLOAT16NDy输出公式中的输出张量 yBOOLND数据类型约束Atlas 训练系列产品不支持 BFLOAT16。Kirin X90/Kirin 9030 处理器系列产品不支持 BFLOAT16。这两条约束在 is_finite_def.cpp 中有对应的实现证据通用配置ascend910b、ascend910_93、ascend950、ascend350允许输入FLOAT16/FLOAT/BF16三种类型而ascend310p与 Kirin 系列GetKirinCoreConfig()的输入数据类型列表仅包含FLOAT16、FLOATconfig310p.Input(x) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT16, ge::DT_FLOAT}) .Format({ge::FORMAT_ND, ge::FORMAT_ND}) .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});约束说明原文档明确本算子无额外约束shape 一致性、维度上限等由接口层校验详见下文返回码章节。aclnn 两段式接口调用在 aclnnAscend CANN 轻量级单算子调用方式下IsFinite 对应接口为aclnnIsFinite接口声明位于 op_api/aclnn_isfinite.h。根据仓库 docs/zh/context/two_phase_api.md 中描述的两段式接口规范必须先调用aclnnIsFiniteGetWorkspaceSize获取计算所需 workspace 大小及执行器再调用aclnnIsFinite执行计算。函数原型aclnnStatus aclnnIsFiniteGetWorkspaceSize( const aclTensor* self, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnIsFinite( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)aclnnIsFiniteGetWorkspaceSize 参数说明第一段接口完成参数校验与计算流程编排各参数说明如下完整说明见 aclnnIsFinite.md参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入输入张量公式中的 self-FLOAT、FLOAT16、DOUBLE、BFLOAT16、INT32、INT64、INT16、INT8、UINT8、BOOLND0-8√outaclTensor*输出输出张量公式中的 out数据类型为 BOOLshape 与 self 相同BOOLND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----返回值错误码第一段接口完成入参校验出现以下场景时报错完整返回码定义参见 aclnn 返回码返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self 或 out 是空指针ACLNN_ERR_PARAM_INVALID161002self 或 out 的数据类型不在支持范围之内ACLNN_ERR_PARAM_INVALID161002self 和 out 的维度超过 8 维ACLNN_ERR_PARAM_INVALID161002self 和 out 的 shape 不一致这些校验逻辑在 op_api/aclnn_is_finite.cpp 的CheckParams中有完整的代码实现依次执行空指针检查CheckNotNull、数据类型检查CheckDtypeValid、shape 约束检查CheckShape维度上限MAX_DIM_LEN 8且OP_CHECK_SHAPE_NOT_EQUAL校验 self 与 out 的 shape 一致性。aclnnIsFinite 参数说明第二段接口用于实际执行计算参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnIsFiniteGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream约束说明aclnn 接口确定性计算aclnnIsFinite 默认确定性实现。调用示例方式一aclnn 单算子调用仓库在 examples/test_aclnn_is_finite.cpp 提供了完整的单算子调用示例同目录下还有使用 RAII 智能指针管理资源的变体示例其完整流程如下#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_isfinite.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; } 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); // 调用aclrtMalloc申请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.固定写法device/stream初始化参考acl API手册 // 根据自己的实际device填写deviceId 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. 构造输入与输出需要根据API的接口自定义构造 std::vectorint64_t selfShape {4, 4}; std::vectorint64_t outShape {4, 4}; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0, 1.123, -2.001, 303.45, 40009, -50.1234, 60.666, -7.6543, 8000, -9.009, 1024, -11.23345, 12, 1356, -14.99, -15.34023}; std::vectorchar outHostData {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建out aclTensor ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_BOOL, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API需要修改为具体的API名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnIsFinite第一段接口 ret aclnnIsFiniteGetWorkspaceSize(self, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnIsFiniteGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize 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); } // 调用aclnnIsFinite第二段接口 ret aclnnIsFinite(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnIsFinite failed. ERROR: %d\n, ret); return ret); // 4.固定写法同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果拷贝至host侧需要根据具体API的接口定义修改 auto size GetShapeSize(outShape); std::vectorchar resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, 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 ret); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %d\n, i, resultData[i]); } // 6. 释放aclTensor和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(out); // 7. 释放device资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例代码的执行可划分为 7 个固定步骤① ACL 运行时初始化aclInit/aclrtSetDevice/aclrtCreateStream② 构造输入输出aclTensor通过aclrtMalloc申请 Device 内存、aclrtMemcpy拷贝数据、aclCreateTensor创建张量句柄③ 调用两段式接口完成计算先aclnnIsFiniteGetWorkspaceSize再按需申请 workspace 后调用aclnnIsFinite④aclrtSynchronizeStream同步等待⑤ 将结果从 Device 拷贝回 Host 并打印⑥ 释放aclTensor⑦ 释放 Device 内存与 Stream 并aclFinalize。示例的完整编译与运行过程参考仓库 编译与运行样例。方式二图模式调用图模式调用通过 GEGraph Engine的算子 IR 构图方式完成示例见 examples/test_geir_is_finite.cpp。核心流程为初始化 GE通过ge::GEInitialize(global_options)设置全局选项例如{ge.exec.deviceId, 0}, {ge.graphRunMode, 1}构建计算图创建Graph graph(tc_ge_irrun_test)使用 op_graph/is_finite_proto.h 注册的算子 IR通过op::IsFinite(add1)构造单算子节点利用宏为算子添加输入 placeholder 与输出描述输入 x 为{4, 4}的 FLOAT 张量输出 y 为同 shape 的DT_BOOL构建会话并执行创建ge::Sessionsession-AddGraph(graph_id, graph, graph_options)添加计算图session-RunGraph(graph_id, input, output)运行图并获取输出结果落盘与清理将输入输出以二进制文件形式写出tc_ge_irrun_test_0008_npu_input_0.bin/npu_output_0.bin最后调用ge::GEFinalize()收尾。底层实现原理aclnn 接口层的双场景分流从 op_api/aclnn_is_finite.cpp 的aclnnIsFiniteGetWorkspaceSize实现可以看到接口层根据输入数据类型对计算流程做了分支优化场景一非浮点数对非浮点类型输入元素必然有界无需真正参与计算直接通过l0op::Fill生成全True值为 1 的 BOOL张量。GetTrueTensor内部利用 executor 分配标量1并转换为 BOOL 张量再按self的 view shape 生成dimTensor调用l0op::Fill得到结果场景二浮点数对浮点类型输入先通过l0op::Contiguous将非连续输入规整为连续张量再调用l0op::IsFinite完成逐元素有限性判断。两种场景的计算结果最终都通过l0op::ViewCopy拷贝到输出out上——这一设计保证了out即使是非连续 Tensor 也能正确写入。workspaceSize由uniqueExecutor-GetWorkspaceSize()汇总返回。若self为空 Tensorself-IsEmpty()第一段接口直接返回workspaceSize 0并提前结束。该分流逻辑与 aclnn_isfinite.h 头文件中的计算图注释完全一致非浮点场景为self → l0op::Fill → l0op::ViewCopy → out浮点场景为self → l0op::Contiguous → l0op::IsFinite → l0op::ViewCopy → out。另外值得注意的是接口层支持的数据类型FLOAT、FLOAT16、BFLOAT16、DOUBLE、INT32、INT64、INT16、INT8、UINT8、BOOL比算子 IR 层更宽泛且CheckDtypeValid会根据当前 NPU 架构910B 及 RegBase 架构选择不同的支持列表DTYPE_SUPPORT_LIST_910B含 BF16DTYPE_SUPPORT_LIST_910不含。Shape 推导与 Tiling 实现Shape 推导IsFinite 是逐元素elewise算子op_host/is_finite_infershape.cpp 直接复用基础工具Ops::Base::InferShape4Elewise完成输出 shape 推导即输出 shape 与输入一致。Tiling切分计算仓库针对不同 NPU 架构提供了独立的 tiling 实现例如 op_host/arch35/is_finite_tiling_arch35.cpp对应 Ascend 950/Atlas A3 的 RegBase 架构与 op_host/arch22/is_finite_tiling_arch22.cpp。tiling 逻辑会校验输入数据类型仅允许FLOAT/FLOAT16/BF16见CalcInputDtype与输出类型必须为BOOL见CalcOutputDtype、校验 x/y shape 一致CheckShape并通过SetTilingData计算 workspace 大小、构造 tiling key 与切分参数。Kernel 计算实现Kernel 入口位于 op_kernel/is_finite.cpp是一个典型的__global__ __aicore__函数templateint D_T_X, int D_T_Y __global__ __aicore__ void is_finite(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) { if (workspace nullptr) { return; } REGISTER_TILING_DEFAULT(IsFiniteTilingData); GET_TILING_DATA_WITH_STRUCT(IsFiniteTilingData, tilingData, tiling); IsFiniteKernelImplD_T_X, D_T_Y(x, y, tilingData); }kernel 首先通过 tiling 数据获取切分信息再调用模板化的IsFiniteKernelImplD_T_X, D_T_Y完成对全局内存中 x 的逐元素遍历与有限性判断将 BOOL 结果写入 y。针对不同架构如 arch35还提供了 DAG 形式的 kernel 实现op_kernel/arch35/is_finite_dag.h 与 is_finite_struct_arch35.h由 tiling key 中的TPL_EXTRA/TPL_SCH_MODE_1等参数驱动选择。测试与验证仓库为 IsFinite 提供了覆盖接口层、host 层与 kernel 层的完整测试kernel 层 UTtests/ut/op_kernel/test_is_finite.cpp 基于 gtest 与 tikicpulibCPU 仿真运行测试数据由 tests/ut/op_kernel/is_finite_data/gen_data.py 与 compare_data.py 生成与比对用例覆盖不同 dtype、不同 shape 的切分场景host 层 UTtests/ut/op_host/test_is_finite_infershape.cpp、test_is_finite_tiling_arch35.cpp 分别验证 shape 推导与 tiling 结果op_api 层 UTtests/ut/op_api/test_aclnn_is_finite.cpp 验证 aclnn 接口的完整调用链路并配套 ST 用例 tests/st/aclnnIsFinite/atk_aclnnIsFinite.json 与 golden 参考实现 tests/assets/golden.py按上述公式计算期望输出。小结IsFinite 是 CANN ops-math 中实现简单但工程链路完整的数学基础算子文档层面给出了明确的功能公式与产品支持矩阵源码层面从 GE 算子 IR 注册is_finite_proto.h、Host 侧定义与 shape 推导is_finite_def.cpp、is_finite_infershape.cpp、aclnn 两段式接口aclnn_is_finite.cpp到按架构拆分的 tiling 与 kernelis_finite_tiling_arch35.cpp、is_finite.cpp一应俱全并且针对非浮点输入直接填 True做了计算优化。开发者既可以参考本文的调用示例快速接入单算子或图模式场景也可以沿上述源码路径深入理解 CANN 算子的通用开发范式。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考