CANN ops-math 的 IsNegInf 算子全解析从 aclnn 接口到昇腾 NPU 内核实现【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathIsNegInf 是 CANN ops-math 数学算子库中的元素级判断算子用于逐元素检测输入张量是否为负无穷-inf。本文以 math/is_neg_inf/README.md 及配套的 aclnnIsNegInf 接口文档 为主体结合仓库中的 op_api、op_host、op_kernel 源码与单测用例完整讲解该算子的功能语义、产品支持矩阵、参数约束、aclnn 两段式调用方式与底层实现原理读完即可独立编写并运行一个 IsNegInf 的 aclnn 调用样例。算子功能与语义功能定义IsNegInf对应 aclnn 接口aclnnIsNegInf的功能是判断输入张量中哪些元素是负无穷大-inf对输入input逐元素计算$$ out_i (input_i -\infty)即当元素取值为负无穷时输出 true否则输出 false输出张量 y 的元素类型为 BOOL。该算子常见于数值稳定性处理场景例如在训练/推理前筛除或保护含 -inf 的中间结果可与 is_inf、is_pos_inf、is_nan 等判断类算子配合使用同族算子可在 [math](https://link.gitcode.com/i/2801dcad7e1dba44d7588848769d24a2) 目录下查看。 ### 计算示例 - 若 x [9, 6, -inf]则 isNegInf(x) 的结果为 [False, False, True]。 ### 产品支持情况 | 产品 | 是否支持 | | :--- | :---: | | termAscend 950PR/Ascend 950DT/term | √ | | termAtlas A3 训练系列产品/Atlas A3 推理系列产品/term | √ | | termAtlas A2 训练系列产品/Atlas A2 推理系列产品/term | √ | | termAtlas 200I/500 A2 推理产品/term | × | | termAtlas 推理系列产品/term | √ | | termAtlas 训练系列产品/term | × | 需要说明的是上述矩阵来自算子 README 的公开声明而从 [op_host/is_neg_inf_def.cpp](https://link.gitcode.com/i/ad7e0be6de16731e9ad11fb23605bc17) 的算子定义看当前源码中仅为 ascend950 与 ascend350 两个平台注册了 AICore 配置AICore().AddConfig(ascend950, ...) 与 AICore().AddConfig(ascend350, ...)可以推断仓库中该算子的内核实现当前主要面向 arch35Ascend 950 系列平台其他平台的支持情况以实际发布版本为准。 ## 参数说明 ### 算子级参数IR 定义 | 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 | | :--- | :---: | :--- | :--- | :---: | | x | 输入 | 待进行判断的入参公式中的 input_i | FLOAT、FLOAT16、BFLOAT16 | ND | | y | 输出 | 待进行判断的出参公式中的 out_i | BOOL | ND | 算子注册于 [op_host/is_neg_inf_def.cpp](https://link.gitcode.com/i/ad7e0be6de16731e9ad11fb23605bc17#L18-L43)输入 x 支持 ge::DT_FLOAT16 / ge::DT_FLOAT / ge::DT_BF16 三种浮点类型输出 y 固定为 ge::DT_BOOL格式均为 ND并支持动态 shape、动态 rank0~8 维。 ### 为什么输入只支持浮点、输出必为 BOOL - 内核侧is_neg_inf 的判断语义“等于 -inf”只对浮点数有意义整型不存在负无穷概念。 - aclnn 侧有一个值得注意的优化在 [op_api/aclnn_isneginf.cpp](https://link.gitcode.com/i/4141bc1c625e44e2dcc1fa53839f1563) 中若输入 self 为整型!IsFloatingType !IsComplexType则直接跳过实际计算通过 Fill 算子将输出整体填充为 false因为整型必然有界、不可能等于 -inf。这也解释了 README 中虽然列出的输入类型为 FLOAT/FLOAT16/BFLOAT16但接口层仍对整型输入做了兜底处理见下文“接口层实现细节”。 ## aclnn 两段式调用接口原型与参数详解 与 CANN 其他 aclnn 算子一致aclnnIsNegInf 采用**两段式接口**设计先调用 GetWorkspaceSize 接口完成入参校验、构图并获取 workspace 大小与执行器再调用执行接口真正下发计算。两段式接口的通用约定可参考 [docs/zh/context/two_phase_api.md](https://link.gitcode.com/i/2bd1e695ef89c62f59411e88b1c93c78)。 ### 第一段aclnnIsNegInfGetWorkspaceSize Cpp aclnnStatus aclnnIsNegInfGetWorkspaceSize( const aclTensor* self, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入输入张量公式中的 self-FLOAT、FLOAT16、DOUBLE、BFLOAT16ND0-8√outaclTensor*输出输出张量公式中的 outshape 与 self 相同BOOLND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----平台差异Atlas 训练系列产品、Atlas 推理系列产品不支持 BFLOAT16 输入该限制在接口文档中通过条件注释声明实际是否生效取决于所运行的昇腾软件栈版本。第二段aclnnIsNegInfaclnnStatus aclnnIsNegInf( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口获取executor输入op 执行器包含算子计算流程stream输入指定执行任务的 Stream两段接口的返回值均为aclnnStatus状态码通用返回码说明见 docs/zh/context/aclnn_return_code.md。第一段接口的错误码第一段接口完成入参校验以下场景会报错返回码错误码描述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_isneginf.cpp 的CheckParams中有对应的实现先做空指针检查再比对self与out的 shapeOP_CHECK_SHAPE_NOT_EQUAL、检查维度上限MAX_DIM_LEN 8最后按平台校验数据类型。约束说明确定性计算aclnnIsNegInf默认确定性实现即相同输入在相同环境下多次执行结果可复现。确定性计算的通用约定可参考 docs/zh/context/determinism_compute.md。shape 一致out的 shape 必须与self完全一致支持标量即 0 维。非连续 Tensorself与out均支持非连续 Tensor见参数表中“非连续Tensor”列接口内部会做连续性转换。接口层实现细节一条从 aclnn 到 l0 算子的调用链阅读 op_api/aclnn_isneginf.cpp 可以还原aclnnIsNegInfGetWorkspaceSize内部的完整构图流程这对理解算子行为与参数约束很有帮助创建执行器与参数校验CREATE_EXECUTOR()创建aclOpExecutor随后调用CheckParams完成空指针、shape、维度、dtype 校验见上文错误码表。空 Tensor 短路若self或out为空 Tensor直接返回workspaceSize 0不触发实际计算。连续性转换对self调用l0op::Contiguous确保内存连续为后续逐元素计算做准备非连续 Tensor 说明见 docs/zh/context/non_contiguous_tensor.md。整型输入快速路径若self为非浮点、非复数类型调用FillTensor(out, false, ...)将输出整体填充为false——整型没有 -inf无需真正计算。浮点输入主路径对浮点输入调用 l0 算子l0op::IsNegInf实现在 op_api/isneginf.cpp该算子按输入 shape 分配DT_BOOL输出 Tensor并通过ADD_TO_LAUNCHER_LIST_AICORE挂载 AICore 内核。从该文件注释“目前无 AiCPU 实现默认走 AiCore”可以推断IsNegInf 当前只在 AICore矢量核上执行。输出视图还原若out是非连续 Tensor调用l0op::ViewCopy将连续计算结果写回out的视图。workspace 汇总通过executor-GetWorkspaceSize()汇总整条构图所需的 workspace 大小并返回。第二段接口aclnnIsNegInf则比较简单直接调用CommonOpExecutorRun由框架完成执行器运行op_api/aclnn_isneginf.cpp。调用示例完整可运行的 aclnn 样例仓库在 math/is_neg_inf/examples/test_aclnn_isneginf.cpp 提供了一份可直接编译运行的调用样例aclnnIsNegInf 文档 中也给出了对应示例。下面基于仓库样例讲解关键步骤完整编译与执行流程请参考 docs/zh/context/compile_and_run_sample.md。头文件与基础设施#include acl/acl.h #include aclnnop/aclnn_isneginf.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)环境初始化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; }构造 aclTensorCreateAclTensor完成三步工作aclrtMalloc申请 Device 侧内存ACL_MEM_MALLOC_HUGE_FIRST→aclrtMemcpy将 Host 数据拷入 → 按连续 Tensor 计算 strides 后用aclCreateTensor创建aclTensorND 格式。样例中输入 shape 为{4, 2}数据为{0, 1, 2, 3, 4, 5, 6, 7}float输出为 BOOL 类型std::vectorint64_t selfShape {4, 2}; std::vectorint64_t outShape {4, 2}; std::vectorfloat selfHostData {0, 1, 2, 3, 4, 5, 6, 7}; std::vectorchar outHostData {1, 1, 1, 1, 0, 0, 0, 0}; ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_BOOL, out);两段式调用与结果回读uint64_t workspaceSize 0; aclOpExecutor* executor; ret aclnnIsNegInfGetWorkspaceSize(self, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnIsNegInfGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); 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); } ret aclnnIsNegInf(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnIsNegInf failed. ERROR: %d\n, ret); return ret); ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); auto size GetShapeSize(outShape); std::vectorchar resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(char), 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]); }最后按顺序释放资源aclDestroyTensor销毁 Tensor、aclrtFree释放 Device 内存与 workspace、aclrtDestroyStream销毁流、aclrtResetDevice复位设备、aclFinalize去初始化。仓库示例 math/is_neg_inf/examples/test_aclnn_isneginf.cpp 中使用了 RAII 智能指针封装StreamPtr、DeviceMemPtr、TensorPtr来管理资源可作参考。图融合与 shape 推导算子如何接入计算图图侧数据类型推导op_graph/is_neg_inf_graph_infer.cpp 通过IMPL_OP(IsNegInf).InferDataType(...)在构图阶段固定输出数据类型为DT_BOOL。Host 侧 shape 推导op_host/is_neg_inf_infershape.cpp 复用Ops::Base::InferShape4Elewise完成逐元素elewiseshape 推导——输入输出 shape 一致。对应的单测 tests/ut/op_host/test_is_neg_inf_infershape.cpp 覆盖了静态 shape如{3, 4}与动态 shape含-1维度两种场景验证推导结果均为与输入一致的 shape 且输出 dtype 为 BOOL。内核实现arch35 上的 Tiling 与 KernelTiling 逻辑op_host/arch35/is_neg_inf_tiling_arch35.cpp 实现了 arch35 平台的 Tiling 逻辑核心流程RunTiling依次完成入参校验CalcInputDtype仅接受 FLOAT16/BF16/FLOATCalcOutputDtype要求输出为 BOOLCheckShape限制维度不超过 8 且输入输出 shape 必须一致。逐元素 Tiling按输入 dtype 分别以IsNegInfDAGhalf、IsNegInfDAGbfloat16_t、IsNegInfDAGfloat为模板调用elewiseBaseTiling.DoTiling生成分块调度信息。Workspace 申请固定申请ASCEND_WORKSPACE 16 * 1024 * 102416 MB的 Device workspace。TilingKey 与 BlockDimGET_TPL_TILING_KEY(tiling-scheMode, dType)组合出内核分派键SetBlockDim(tiling-blockNum)设置参与计算的核数。Tiling 准备阶段TilingPrepareForIsNegInf从平台信息中读取 AIV 核数GetCoreNumAiv与 UB 内存大小供调度使用。Kernel 逻辑op_kernel/is_neg_inf_apt.cpp 是 AscendC 内核入口声明为__global__ __aicore__仅使用 AIV 核KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY)。内核按模板参数dTypeFP16/BF16/FP32实例化通过REGISTER_TILING_DEFAULT(EleBaseTilingDataV2)读取 Tiling 数据再以ElementwiseSchschMode, IsNegInfOp::IsNegInfDAGT::OpDag调度器完成分块搬运与逐元素比较最终将 BOOL 结果写入输出。具体到每个元素“是否为 -inf”的比较在 op_kernel/arch35/is_neg_inf_dag.h 的算子 DAG 描述中定义数据布局结构见 op_kernel/arch35/is_neg_inf_struct.h。对应平台的 tiling 单测位于 tests/ut/op_host/arch35/test_is_neg_inf_tiling_arch35.cpp算子 API 层单测见 tests/ut/op_api/test_aclnn_isneginf.cppST 测试用例配置见 tests/st/aclnnIsNegInf/atk_aclnnIsNegInf.json 与 tests/st/arch35/ttk_kernel_is_neg_inf_st.csv可用于回归验证。小结IsNegInf 是一个语义简单、实现路径清晰的逐元素判断算子输入浮点张量输出同 shape 的 BOOL 张量逐元素标记 -inf 位置。通过本文可以掌握其功能语义out_i (input_i -inf)与产品支持矩阵aclnnIsNegInfGetWorkspaceSize/aclnnIsNegInf两段式接口的完整参数与错误码语义从 aclnn 构图Contiguous → l0op::IsNegInf → ViewCopy到 Host 侧 Tilingdtype/shape 校验 elewise 分块再到 AICore 内核ElementwiseSchIsNegInfDAG的完整调用链一份可直接参考的、可编译运行的调用样例与对应的单测/ST 用例位置。如需进一步了解两段式接口通用约定、aclnn 返回码或样例编译方法可继续阅读 docs/zh/context/two_phase_api.md、docs/zh/context/aclnn_return_code.md 与 docs/zh/context/compile_and_run_sample.md。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考