CANN ops-math 算子解析：IsPosInf 逐元素正无穷判定算子的原理、ACLNN 接口与源码实现

CANN ops-math 算子解析：IsPosInf 逐元素正无穷判定算子的原理、ACLNN 接口与源码实现

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math

导读

本文围绕 CANN ops-math 数学算子库中的IsPosInf算子展开，逐层讲解其功能语义、对外 ACLNN 两段式接口、参数与约束、编译运行示例，以及从 L2 接口、tiling 到 AiCore kernel 的完整源码实现链路。读完本文，你将能够独立在 Ascend 910B/910C 设备上调用aclnnIsPosInf完成正无穷判定，并能依据源码理解浮点路径与整数/布尔路径的差异化实现原理。

算子功能与数学语义

IsPosInf是一个逐元素（element-wise）的判定类算子，功能定义在 README.md 中：

• 浮点输入：对输入张量self的每个元素判断是否为+inf，返回布尔结果，即out_i = (self_i == +∞)；
• 非浮点有界输入：对bool/int8/uint8/int16/int32/int64等有界类型，数值域内不存在正无穷，因此直接返回与输入同 shape 的全false布尔张量，即out_i = false。

数学表达如下：

若输入为浮点型：

$$ out_i = (self_i == +\infty) $$

若输入为支持的非浮点有界类型：

$$ out_i = \mathrm{false} $$

这一设计意味着调用方无需区分输入 dtype 即可统一使用该接口：浮点输入走真实的逐元素比较计算路径，而整数、布尔等有界类型走"填充全 false"的快捷路径，两者的结果语义保持一致。需要指出的是，该语义是"仅判定正无穷"，-inf与NaN均不命中（NaN == +inf恒为 false），与同时覆盖正负无穷的IsInf类算子存在区别。

产品支持情况与运行前提

根据 README.md 与算子定义源码 is_pos_inf_def.cpp：

产品	是否支持
Atlas A2 训练系列产品 / Atlas A2 推理系列产品	√
Atlas A3 训练系列产品 / Atlas A3 推理系列产品	√

当前工程按需求仅面向Ascend 910B/910C路径生成 ACLNN 接口与 AiCore 实现，具体到算子注册层面表现为：

• host 端算子定义通过AICore().AddConfig("ascend910b", aicoreConfig)和AICore().AddConfig("ascend910_93", aicoreConfig)分别注册ascend910b与ascend910_93两类 AiCore 配置；
• L2 接口层在 aclnn_is_pos_inf.cpp 中通过CheckSocAndDtypeValid对 SoC 版本做运行期校验：仅当SocVersion::ASCEND910B <= soc <= SocVersion::ASCEND910E时才允许继续执行，否则返回ACLNN_ERR_PARAM_INVALID并记录错误日志 "aclnnIsPosInf is only supported on Ascend 910B/910C class devices"；
• BFLOAT16浮点路径同样依赖 910B/910C 类 SoC 支持。

对外接口：两段式 ACLNN 调用

IsPosInf对外暴露两段式（two-phase）ACLN 接口，函数原型定义在 aclnn_is_pos_inf.h：

aclnnStatus aclnnIsPosInfGetWorkspaceSize( const aclTensor *self, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor);

aclnnStatus aclnnIsPosInf( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream);

两段接口的分工如下：

• 第一段aclnnIsPosInfGetWorkspaceSize：完成参数校验、构建算子执行流程（executor），并返回执行所需的设备侧 workspace 大小。它不真正触发计算；
• 第二段aclnnIsPosInf：将第一段返回的 workspace、executor 以及用户 stream 传入，真正提交计算任务到设备执行。

参数说明

参数名	输入/输出	描述	数据类型	数据格式
self	输入	待判断是否为正无穷的张量	FLOAT16/FLOAT/BFLOAT16/BOOL/INT8/UINT8/INT16/INT32/INT64	ND
out	输出	逐元素判断结果	BOOL	ND

补充说明（来自 docs/aclnnIsPosInf.md）：

• self不支持 broadcast，out的 shape 必须与self完全一致，维度范围为 0~8 维，支持标量和空 Tensor；
• workspaceSize与executor均为第一段接口的输出参数，由第一段接口填充返回；
• 两个 Tensor 均支持非连续（non-contiguous）输入输出：self非连续时接口内部先转连续，out非连续时通过ViewCopy回写。

返回值与典型报错场景

第一段接口完成参数检查并构造执行流程，典型报错场景如下：

返回码	错误码	描述
ACLNN_ERR_PARAM_NULLPTR	161001	self或out是空指针
ACLNN_ERR_PARAM_INVALID	161002	self/outdtype 不支持
ACLNN_ERR_PARAM_INVALID	161002	outdtype 不是BOOL
ACLNN_ERR_PARAM_INVALID	161002	self/outshape 不一致
ACLNN_ERR_PARAM_INVALID	161002	维度数超过 8
ACLNN_ERR_PARAM_INVALID	161002	当前设备不在910B/910C支持范围

从源码看，上述检查在 aclnn_is_pos_inf.cpp 的CheckParams中依次执行：先做OP_CHECK_NULL空指针检查（返回ACLNN_ERR_PARAM_NULLPTR），再做OP_CHECK_SHAPE_NOT_EQUALshape 一致性检查、OP_CHECK_MAX_DIM最大 8 维检查（均返回ACLNN_ERR_PARAM_INVALID），最后通过CheckSocAndDtypeValid校验 SoC 版本与 dtype 合法性。其中out的 dtype 必须为BOOL（OP_CHECK_DTYPE_NOT_MATCH(out, DataType::DT_BOOL)），self支持的数据类型由ASCEND910B_DTYPE_SUPPORT_LIST限定，与 README 参数表中的 9 种类型一一对应。

调用示例：从设备初始化到结果校验

仓库在 examples/test_aclnn_is_pos_inf.cpp 中提供了完整可运行的调用样例，其核心流程如下：

1. 初始化环境：aclInit(nullptr)初始化 ACL 运行时，aclrtSetDevice(kDeviceId)指定设备（示例中kDeviceId = 0），aclrtCreateStream(&stream)创建计算流；
2. 准备数据：示例构造 shape 为{6}的FLOAT输入{+inf, -2.0f, 0.0f, 3.0f, +inf, 5.0f}，期望输出为{1, 0, 0, 0, 1, 0}，通过aclrtMalloc分配设备侧内存，aclrtMemcpy将输入拷贝到设备；
3. 构造 aclTensor：通过aclCreateTensor以ACL_FORMAT_ND格式创建输入 Tensor（ACL_FLOAT）与输出 Tensor（ACL_BOOL）；
4. 两段式调用：

ret = aclnnIsPosInfGetWorkspaceSize(inputTensor, outputTensor, &workspaceSize, &executor); // ... if (workspaceSize > 0) { ret = aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); // ... } ret = aclnnIsPosInf(workspace, workspaceSize, executor, stream);

5. 同步与校验：aclrtSynchronizeStream(stream)等待执行完成，将输出aclrtMemcpy回 host，用memcmp与期望结果比对，一致则打印is_pos_inf example passed；
6. 资源释放：依次释放 workspace、设备内存、Tensor、stream，并aclrtResetDevice、aclFinalize。

编译与运行

examples/run.sh 提供了开箱即用的编译运行脚本，核心编译命令如下（已省略版权注释，变量含义见脚本）：

g++ -std=c++17 -O2 \ test_aclnn_is_pos_inf.cpp \ -I"${CANN_ROOT}/${HOST_ARCH_DIR}/include" \ -I"${CANN_ROOT}/${HOST_ARCH_DIR}/include/aclnnop" \ -I"${CUSTOM_OPP_ROOT}/op_api/include" \ -L"${CANN_ROOT}/${HOST_ARCH_DIR}/lib64" \ -L"${CUSTOM_OPP_ROOT}/op_api/lib" \ -Wl,-rpath,"${CANN_ROOT}/${HOST_ARCH_DIR}/lib64:${CUSTOM_OPP_ROOT}/op_api/lib" \ -lcust_opapi -lnnopbase -lascendcl \ -o build/test_aclnn_is_pos_inf

运行前提说明：

• 脚本通过source "${ASCEND_HOME_PATH:-/usr/local/Ascend/cann}/set_env.sh"加载 CANN 环境；
• 头文件来自 CANN 安装目录（含aclnnop目录）与自定义算子包（ASCEND_CUSTOM_OPP_PATH，缺省回退到${CANN_ROOT}/opp/vendors/custom_math）的op_api/include；
• 链接libcust_opapi（自定义算子 API 库）、libnnopbase与libascendcl；
• 脚本按宿主机架构自动选择aarch64-linux或x86_64-linux目录。

源码级实现剖析：从 L2 接口到 AiCore Kernel

分层调用链总览

IsPosInf的实现遵循 CANN 算子典型的"L2 接口 → L0 算子（l0op）→ AiCore kernel"三层结构：

aclnnIsPosInfGetWorkspaceSize / aclnnIsPosInf (op_api/aclnn_is_pos_inf.cpp) │ ▼ l0op::IsPosInf / l0op::Contiguous / l0op::Fill / l0op::ViewCopy │ (op_api/is_pos_inf.cpp) ▼ IsPosInf AiCore Kernel (op_kernel/is_pos_inf.cpp / .h) ▲ │ host 侧 tiling 决策 IsPosInfTilingFunc (op_host/is_pos_inf_tiling.cpp)

L2 接口：浮点与非浮点两条路径

第一段接口aclnnIsPosInfGetWorkspaceSize在完成参数校验后，针对空 Tensor 与 dtype 类型分别处理（见 aclnn_is_pos_inf.cpp）：

• 空 Tensor 快捷返回：若self->IsEmpty() || out->IsEmpty()，则*workspaceSize = 0并直接返回ACLNN_SUCCESS，不触发任何计算；
• 非浮点有界类型（非IsFloatingType且非IsComplexType）：调用FillBoolTensor(out, false, executor)构造全false结果。该函数内部通过executor->ConvertToTensor构造 dims/形状参数，并调用l0op::Fill用false填充输出 Tensor；
• 浮点类型：先执行l0op::Contiguous(self)保证输入连续，再调用l0op::IsPosInf(selfContiguous, executor)完成逐元素比较；
• 统一回写：两条路径的结果最终都通过l0op::ViewCopy(result, out, executor)写回用户输出，从而支持out为非连续 Tensor 的场景。

第二段接口aclnnIsPosInf则直接委托CommonOpExecutorRun(workspace, workspaceSize, executor, stream)执行由第一段构造好的算子流。

L0 层：dtype 支持范围收窄

L0 层实现位于 op_api/is_pos_inf.cpp。其内部维护AICORE_DTYPE_SUPPORT_LIST = {FLOAT16, FLOAT, BF16}，即只有三种浮点类型会真正进入 AiCore 计算路径：

• IsAiCoreSupport用CheckType判断输入 dtype 是否在列表中；
• 若不在列表中（即整数、布尔等有界类型），l0op::IsPosInf直接返回nullptr，此时由 L2 层兜底走Fill全 false 路径，这正是 README 所述"非浮点返回全 false"语义的实现机制；
• 浮点输入则executor->AllocTensor分配DT_BOOL、FORMAT_ND输出，并通过ADD_TO_LAUNCHER_LIST_AICORE(IsPosInf, ...)将算子加入 AiCore 启动列表。

算子注册与 shape 推导

Host 端算子定义见 is_pos_inf_def.cpp：

• 通过OP_ADD(IsPosInf)注册算子IsPosInf；
• 输入命名为x（REQUIRED），支持 9 种 dtype，FORMAT_ND，并设置AutoContiguous()（自动连续化）；
• 输出命名为y（REQUIRED），dtype 恒为DT_BOOL；
• AiCore 配置开启DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)、PrecisionReduceFlag(true)，并指定 kernel 文件为is_pos_inf。

Shape 推导见 is_pos_inf_infershape.cpp：InferShape4IsPosInf直接将输入 shape 拷贝给输出（*outputShape = *inputShape），与 README 中"shape 必须完全一致、不支持 broadcast"的约束一致。

Tiling 策略：按核切分与 Cache Line 对齐

Tiling 逻辑位于 is_pos_inf_tiling.cpp，核心思路是"先按核均分、再按 Cache Line 对齐、最后按 UB 容量切 tile"：

1. 获取平台信息：GetPlatformInfo通过PlatformAscendC获取 AIV 核数coreNum、UB 内存大小ubSize与库函数 workspace 大小wsSysSize；wsSysSize会写入算子 workspace（currentWorkspace[0] = wsSysSize）；
2. 合法性校验：ValidateInput检查维度数不超过 8、dtype 必须为三种浮点之一（DT_FLOAT/DT_FLOAT16/DT_BF16），并换算元素总数与单元素字节数；
3. 按核分配：totalLengthCore = ceil(totalLength / coreNum)，再按CACHE_LINE_BYTE_LENGTH(512) / dtypeSize对齐得到totalLengthCoreAlign；实际使用核数usedCoreNum由对齐后的长度反推；
4. tile 切分：以ubSize / FLOAT_BUFFER_COEFFICIENT(10)为单 tile 元素上限，并按256 / dtypeSize对齐；FP32 与 BF16 分别额外限制在DEBUG_FP32_TILE_ELEMENTS = 192、DEBUG_BF16_TILE_ELEMENTS = 192以内；
5. 落盘：SaveTilingData将formerNum/formerLength/tailLength/tileLength/dtypeId写入 tiling 数据（结构定义见 is_pos_inf_tiling_data.h），并通过context->SetBlockDim(usedCoreNum)设置启动核数。

AiCore Kernel：三条计算路径

Kernel 入口位于 op_kernel/is_pos_inf.cpp，通过GET_TILING_DATA读取 tiling 数据，按dtypeId分派到三条实现（详细代码见 op_kernel/is_pos_inf.h）：

• FP16 通用模板路径（KernelIsPosInf<half>）：以half实例化模板，采用Duplicate+Compare(CMPMODE::EQ)+Select+Cast的向量指令流水实现逐元素比较；
• FP32 统一路径（KernelIsPosInfFp32Unified）：Compare比较后，通过BuildMaskOutput把比较掩码Select映射为1/0的half值，再Cast成uint8_t输出；
• BF16 统一路径（KernelIsPosInfBf16Unified）：由于 BF16 比较指令受限，先把bfloat16_tCast成float，再走与 FP32 相同的比较与掩码输出流程；
• 数据搬运：输入输出均采用双缓冲（BUFFER_NUM = 2）TQue，CopyInTile用DataCopyPad搬运 tile（含尾块 padding），CopyOutTile回写uint8_t结果；ProcessTiles按ceil(blockLength / tileLength)循环处理每个 tile；
• 核间分工：InitKernelContext依据GetBlockIdx()判断当前核属于前formerNum个核（各处理formerLength元素）还是尾核（处理tailLength元素），并从对应全局偏移建立 GM 视图。

约束说明汇总

综合 README.md、docs/aclnnIsPosInf.md 与源码，IsPosInf的完整约束如下：

• self和out不能为空指针；
• self与outshape 必须完全一致，不支持 broadcast（inferShape 直接拷贝 shape 佐证）；
• 维度数不超过 8，支持标量（tiling 侧通过SCALAR_SHAPE = {1}把 0 维 shape 归一为 1 元素）和空 Tensor（L2 接口对空 Tensor 直接返回成功，workspaceSize = 0）；
• outdtype 必须为BOOL；
• self非连续时先转连续（l0op::Contiguous）；out非连续时通过ViewCopy回写；
• BFLOAT16浮点路径依赖910B/910C类 SoC 支持；
• 算子为确定性实现（deterministic），无随机性。

测试与验证

算子单元测试位于 tests/ut/op_api/test_aclnn_is_pos_inf.cpp，基于 gtest 框架覆盖了以下关键场景：

测试用例	验证点
case_support_float32	FP32 输入{+inf, -1, 0, 1, +inf, 2}输出{true, false, false, false, true, false}
case_support_float16	FP16 输入含两个+inf的判定
case_support_bf16	BF16 输入含两个+inf的判定
case_support_integer_all_false	INT32 输入全false输出
case_support_bool_all_false	BOOL 输入全false输出
case_non_contiguous_input_output	非连续 Tensor（自定义 stride）场景
case_null_self/case_null_out	空指针返回ACLNN_ERR_PARAM_NULLPTR

此外仓库还提供了 ST（系统测试）用例目录 tests/st/aclnnIsPosInf（含用例配置 JSON 与执行器脚本），以及对应测试工程的 tests/CMakeLists.txt。这些测试与示例共同验证了本文前述的功能语义、非浮点全 false 行为、非连续支持与异常路径，是复现与二次开发时的最佳参照。

小结

IsPosInf是 ops-math 中实现简洁但语义边界清晰的逐元素判定算子：对外以标准两段式 ACLNN 接口呈现，兼容 9 种输入 dtype；对内依据 dtype 分流——浮点走"Contiguous → AiCore 比较 kernel → ViewCopy"，有界类型走"Fill 全 false"，并在 tiling 层按核、按 Cache Line、按 UB 容量三级切分。理解其实现链路，可以举一反三地掌握 ops-math 中同类 element-wise 判定类算子的通用工程范式。

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math