CANN ops-math IsInf 算子实战指南在 NPU 上逐元素判定张量是否为无穷大【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathIsInf 是 CANN ops-math 数学算子库中的一个逐元素判定算子用于判断张量中哪些元素是无限大值inf、-inf。本指南以 math/is_inf/README.md 为骨架结合算子定义、AscendC 内核实现、Tiling 与测试用例完整讲解 IsInf 的产品支持情况、参数约束、aclnn 与图模式两种调用方式以及底层实现原理帮助开发者在 Atlas 训练/推理等 NPU 产品上快速完成无穷大判定并读懂其实现细节。功能说明IsInf 的算子功能是判断张量中哪些元素是无限大值即 inf、-inf。它逐元素地对输入进行判定输出与输入形状完全相同的布尔张量。计算公式如下$$ out_i(input_i \pm inf) $$即对于输入张量中的每一个元素input_i若其等于正无穷或负无穷则对应输出out_i为True1否则为False0。示例若x [9, 6, 3]计算结果为[False, False, False]。有限数值均不会被判定为无穷大。若x [[-3.14, inf], [2.7183, nan]]计算结果为[[False, True], [False, False]]。注意nan非数不等于无穷大因此判定结果为False。需要特别注意的是整型、布尔都是有界的。整型和布尔类型的数值范围天然不包含无穷大因此对该类输入执行判定时结果恒为False这一点与浮点类型FLOAT、FLOAT16、BFLOAT16 等存在本质区别。仓库中的黄金比对脚本 tests/assets/golden.py 也印证了语义aclnn 路径使用torch.isinf(x)kernel 路径使用np.isinf(x)生成期望结果。产品支持情况IsInf 算子在 CANN ops-math 仓库中的支持情况如下以 README 为准产品是否支持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_inf_def.cpp 中通过AICore().AddConfig(...)注册的ascend910b、ascend910_93、ascend950、ascend350、mc62、ascend310p、kirinx90、kirin9030配置一一对应可以相互印证。产品侧还有一份针对各 SoC 的算子二进制与简化 key 配置位于 op_host/config 目录下例如ascend950/is_inf_binary.json、kirinx90/is_inf_simplified_key.ini等。参数说明算子的输入输出参数如下针对算子级 ND 格式参数名输入/输出/属性描述数据类型数据格式x输入待进行判断的入参公式中的input_i。FLOAT、FLOAT16、BFLOAT16NDout输出待进行判断的出参公式中的out_i。BOOLND存在以下平台差异限制Atlas 训练系列产品、Atlas 推理系列产品不支持 BFLOAT16。Kirin X90/Kirin 9030 处理器系列产品不支持 BFLOAT16。从源码角度进一步印证算子定义 op_host/is_inf_def.cpp 中ascend310p与 Kirin 系列配置的输入仅注册了DT_FLOAT16、DT_FLOAT两种类型不含 BF16而ascend950等配置的输入注册了DT_FLOAT16、DT_FLOAT、DT_BF16三种类型与文档中的平台差异说明一致。约束说明约束说明无。不过在实际使用中仍需遵循文档与代码中的隐含边界例如通过 aclnn 接口调用时输入x的维度shape 的维数应不超过 8确定性计算方面aclnnIsInf默认为确定性实现。调用说明IsInf 算子支持两种调用方式调用方式调用样例说明aclnn 调用test_aclnn_is_inf.cpp通过 aclnnIsInf 接口方式调用 IsInf 算子。图模式调用test_geir_is_inf.cpp通过算子 IRis_inf_proto.h构图方式调用 IsInf 算子。下面分别对两种方式进行深入讲解。aclnn 调用两段式接口aclnnIsInf采用 CANN 算子库通用的两段式接口设计必须先调用第一段接口aclnnIsInfGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器再调用第二段接口aclnnIsInf执行计算。函数原型如下aclnnStatus aclnnIsInfGetWorkspaceSize( const aclTensor *x, const aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnIsInf( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)aclnnIsInfGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorx输入待进行判断是否是 inf 的入参公式中的input_i。-FLOAT、FLOAT16、DOUBLE、BFLOAT16、INT8、INT16、INT32、INT64、UINT8、BOOLND0-8√out输出待进行判断是否是 inf 的出参公式中的out_i。shape 与 x 相同。BOOLND0-8×workspaceSize输出返回需要在 Device 侧申请的 workspace 大小。-----executor输出返回 op 执行器包含了算子计算流程。-----需要注意aclnn 接口层对输入x的支持类型比算子级FLOAT、FLOAT16、BFLOAT16更宽额外覆盖了 DOUBLE、INT8/16/32/64、UINT8、BOOL 等类型同时支持非连续 Tensor 输入。但如前所述整型与布尔类型本身是有界的其判定结果恒为False。返回值与异常场景第一段接口会完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_INNER_TILING_ERROR561002x 维度大于 8。ACLNN_ERR_PARAM_NULLPTR161001参数 x、out 是空指针。ACLNN_ERR_PARAM_INVALID161002参数 x、out 的数据类型不在支持范围内或 x、out 的 shape 不一致。aclnnIsInf 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnIsInfGetWorkspaceSize获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。第二段接口同样返回aclnnStatus状态码。完整的 aclnn 调用示例仓库中的 examples/test_aclnn_is_inf.cpp 给出了完整的调用示例其核心流程可概括为以下 7 步固定写法的通用步骤 算子特有步骤device/stream 初始化aclInit→aclrtSetDevice→aclrtCreateStream。构造输入与输出通过aclrtMalloc申请 Device 侧内存aclrtMemcpy将 Host 数据拷入计算连续 tensor 的 strides 后调用aclCreateTensor创建xACL_FLOAT与outACL_BOOL两个aclTensor。调用第一段接口aclnnIsInfGetWorkspaceSize(x, out, workspaceSize, executor)获取 workspace 大小与执行器若workspaceSize 0用aclrtMalloc申请对应大小的 workspace 内存。调用第二段接口aclnnIsInf(workspaceAddr, workspaceSize, executor, stream)真正执行计算。同步等待aclrtSynchronizeStream(stream)等待任务执行结束。取回结果将 Device 侧输出内存通过aclrtMemcpyACL_MEMCPY_DEVICE_TO_HOST拷贝回 Host逐元素打印result[i]。释放资源aclDestroyTensor释放 tensoraclrtFree释放 Device 内存与 workspace最后aclrtDestroyStream、aclrtResetDevice、aclFinalize收尾。示例中的核心调用片段如下// 3. 调用CANN算子库API uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnIsInf第一段接口 ret aclnnIsInfGetWorkspaceSize(x, out, workspaceSize, executor); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 调用aclnnIsInf第二段接口 ret aclnnIsInf(workspaceAddr, workspaceSize, executor, stream);示例的输入为 4×4 的 FLOAT 张量数据如{0, 1.123, -2.001, 303.45, ..., -15.34023}由于均为有限值输出结果应全部为0False。若将输入中任意元素替换为inf或-inf对应位置输出即为1True。更详细的编译与执行过程可参考仓库中的编译运行样例说明文档。图模式调用基于算子 IR 构图图模式调用通过 GEGraph Engine的 IR 构图方式执行 IsInf 算子。算子原型注册于 op_graph/is_inf_proto.hREG_OP(IsInf) .INPUT(x, TensorType({DT_FLOAT16, DT_FLOAT, DT_DOUBLE, DT_BF16})) .OUTPUT(y, TensorType({DT_BOOL})) .OP_END_FACTORY_REG(IsInf)输入xfloat16、float32、double、bfloat16 类型ND 格式。输出y与 x 相同 shape 的 bool 张量表示 x 中哪些元素是 inf。该算子 IR 与 TensorFlow 的IsInf算子兼容。示例 examples/test_geir_is_inf.cpp 展示了完整的图模式调用流程通过op::IsInf(add1)创建算子节点输入 shape 为{4, 4}输入类型为DT_FLOAT输出y为DT_BOOLshape 与输入一致。构造ge::Graph将算子节点加入图中并通过graph.SetInputs(inputs).SetOutputs(outputs)绑定图输入输出。初始化 GEge::GEInitialize(global_options)其中ge.exec.deviceId指定设备号、ge.graphRunMode指定运行模式。创建ge::Session通过session-AddGraph(graph_id, graph, graph_options)添加计算图必要时可用aclgrphDumpGraph导出图结构。调用session-RunGraph(graph_id, input, output)运行整图随后将输入输出数据写入.bin文件并逐元素打印结果最后GEFinalize()收尾。图模式适合算子嵌入到整张计算图、由 GE 统一调度执行的场景与 aclnn 的单算子调用方式互为补充。源码级实现原理除了接口层面的调用方式IsInf 算子在仓库中还具备完整的 Host 侧 Tiling、Device 侧 AscendC 内核以及形状推导实现理解这些实现有助于深入掌握算子行为。形状推导输出与输入同 shape形状推导实现位于 op_host/is_inf_infershape.cpp逻辑非常简洁将输入的 shape 直接赋给输出即*outShape *inShape并通过IMPL_OP_INFERSHAPE(IsInf)完成注册。这保证了输出out与输入x形状一致与 aclnn 接口文档中“out 的 shape 与 x 相同”的说明完全对应。Tiling多核切分与 workspace 分配Host 侧 Tiling 实现位于 op_host/is_inf_tiling.cpp核心逻辑如下以 32 个数据DATA_BLOCK 32为粒度将输入数据切分到多个 AIV 核上核数needCoreNum取“数据块数”与“平台 AIV 核数”的较小值尾部不足一块的数据由tailDataCoreNum与lastCoreDataCount记录由最后一个核处理。workspace 固定申请32 * 1024 * 102432MB。通过TILING_KEY区分三种数据类型的计算分支TILING_KEY_HALF 1FP16、TILING_KEY_FLOAT 2FP32、TILING_KEY_BFLOAT16 3BF16并以SetTilingKey写入 Tiling 上下文。UB统一缓冲区的可用大小会预留 tiling 数据空间并做双缓冲与 32 数据对齐处理。此外对于支持 RegBase 的 SoC 版本会走独立的IsInfRegbaseTiling实现op_host/is_inf_tiling_arch35.cpp对应 arch35 内核。AscendC 内核位运算判定 infDevice 侧 AscendC 内核位于 op_kernel/is_inf.cpp 与 op_kernel/is_inf.h。内核入口函数is_inf根据TILING_KEY选择模板实例TILING_KEY_IS(1)IsInfhalf, SIGN_MASK, F16_INF_NUM其中F16_INF_NUM 0x7c00FP16 正无穷的位模式TILING_KEY_IS(2)IsInffloat, SIGN_MASK, FLOAT_INF_NUM其中FLOAT_INF_NUM 0x7f800000FP32 正无穷的位模式TILING_KEY_IS(3)IsInfhalf, SIGN_MASK, BF16_INF_NUM其中BF16_INF_NUM 0x7f80BF16 正无穷的位模式该分支在__NPU_ARCH__ 3003 || 3113平台上会被排除对应文档中部分平台不支持 BF16 的约束。判定算法的核心是按位运算而非浮点比较先用SIGN_MASK 0x7fff与输入做按位与屏蔽符号位得到数值的绝对值位模式再减去对应的 INF 位模式通过And、Mins、Muls、Adds等向量指令将结果归一化到 0/1最后Cast为 uint8 输出1 表示 ±inf0 表示非 inf。由于 ±inf 只差一个符号位屏蔽符号位后两者位模式一致因此一次运算即可同时判定inf与-inf。这种位级方案比浮点比较更快也天然规避了nan被误判的问题nan 的位模式不等于 inf。内核在数据搬运上采用双缓冲BUFFER_NUM 2流水CopyInX将全局内存数据拷入 VECIN 队列CompareInf完成按位比较计算CopyOut将结果写回全局内存并针对尾部数据做 32 对齐处理保证多核切分场景下的正确性。测试与黄金比对仓库为 IsInf 提供了完善的单算子测试与黄金数据生成黄金脚本 tests/assets/golden.pykernel 路径用np.isinf(x)生成期望结果aclnn 路径用torch.isinf(x)生成期望结果覆盖aclnnIsInf与is_inf两个入口。Host 侧单测tests/ut/op_host/test_is_inf_infershape.cpp、test_is_inf_tiling.cpp、test_dynamic_is_inf_impl.py 等分别验证形状推导、Tiling 结果与动态 shape 实现。Kernel 侧单测tests/ut/op_kernel/test_is_inf.cpp、test_is_inf_arch35.cpp配合 tests/ut/op_kernel/is_inf_data/gen_data.py 生成输入数据。ST 场景测试配置tests/st/aclnnIsInf/atk_aclnnIsInf.json、tests/st/arch35/ttk_aclnn_is_inf_st.csv 等覆盖 aclnn 与 kernel 两个入口的端到端用例。小结本文以 math/is_inf/README.md 为核心完整覆盖了 CANN ops-math 中 IsInf 算子的功能语义、产品支持矩阵、参数与平台差异、aclnn 两段式接口与图模式 IR 两种调用方式并结合 op_host/is_inf_tiling.cpp、op_kernel/is_inf.cpp、op_host/is_inf_infershape.cpp 等源码深入剖析了其多核切分、位运算判 inf 与形状推导的底层实现最后通过 tests/assets/golden.py 及单测/ST 配置给出了验证路径。无论是作为网络训练/推理中检测数值发散出现 inf的检查算子还是作为学习 CANN 单算子“定义—Tiling—Kernel—测试”全流程的样例IsInf 都是一个功能清晰、实现精炼的参考实现可直接参考本指南在支持的产品上快速上手使用。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考