CANN ops-math 算子开发指南：aclnnIsClose 两段式接口详解与 NPU 源码级实现剖析

CANN ops-math 算子开发指南aclnnIsClose 两段式接口详解与 NPU 源码级实现剖析【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本篇技术指南聚焦 CANN 数学算子库 ops-math 中IsClose判断两个张量元素是否彼此接近算子的aclnn级调用方式完整讲解其产品支持情况、closeness 判定公式、两段式接口的函数原型与全部参数语义、错误返回码以及可编译可运行的 C 调用示例同时深入 math/is_close 模块的 op_api、op_host、op_kernel 源码与测试用例还原该算子从参数校验、Tiling 调度到 Vector 计算指令的完整实现链路。读完本文你将能够独立基于aclnnIsCloseGetWorkspaceSize/aclnnIsClose编写 NPU 上的 is_close 计算程序并理解其底层 Broadcast 调度与 NaN 语义。产品支持情况aclnnIsClose作为 ops-math 数学类基础算子在不同昇腾硬件形态上的支持情况如下表所示以 math/is_close/docs/aclnnIsClose.md 为准产品形态是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品支持Atlas 训练系列产品支持需要说明的是math/is_close/README.md 中的产品支持表格对Atlas 推理系列产品 / Atlas 训练系列产品标注为不支持两处文档口径存在差异从算子注册代码看math/is_close/op_host/is_close_def.cpp 仅为ascend950与ascend350两个芯片架构注册了 AICore 配置因此以当前仓库源码为准可以确认算子当前面向 Ascend 950 系列与 Ascend 350A3系列实现其余形态的支持范围以发布版本的官方文档说明为准。功能说明与 closeness 判定公式aclnnIsClose的接口功能是返回一个带有布尔元素的新张量逐元素判断输入self与other是否彼此接近若接近则对应位置为True否则为False。输出张量out的数据类型为BOOL。closeness 判定公式定义如下$$ \left | self_{i}-other_{i}\right | \le atol rtol\times \left | other_{i} \right | $$其中rtol相对容忍度relative tolerance公式中的rtolatol绝对容忍度absolute tolerance公式中的atol当self与other都是有限值finite时上述公式成立当self与other都是非有限值如无穷大时仅当两者完全相等才判定为 closeequalNanTrue时两个NaN被视为 closeequalNanFalse时两个NaN被视为不 close。源码级的判定实现上述语义在 Vector 内核 math/is_close/op_kernel/arch35/is_close_dag.h 的NanEqualCompare中逐指令落地判定被拆解为四路比较结果再合并结果 1完全相等Compare(x1, x2, CMPMODE::EQ)用于处理有限值相等及后续非有限值相等即 close的分支结果 2双方均为 NaNCompare(x1, x1, NE)与Compare(x2, x2, NE)分别定位 NaNNaN 与自身比较不相等再And求交得到bothNanMask当equalNan 1时该掩码参与最终结果结果 3公式判定AbsSub计算|x1 - x2|Abs取|x2|Muls计算rtol * |x2|Adds加上atol最后Compare(LE)判断|x1 - x2| atol rtol * |x2|结果 4有限性判断Duplicate(INF_CONST)构造正无穷Compare(LT)判断|x2| INF即x2为有限值。最终通过And(funcCmpMask, finiteMask)得到符合公式且有限再Or上完全相等或双方 NaN用Select(regTensorOne, regTensorZero, resultMask)产出uint8布尔结果。这一实现精确对应文档公式与 NaN 语义。两段式接口架构aclnnIsClose属于 CANN 的两段式接口详见 docs/zh/context/two_phase_api.md第一段aclnnIsCloseGetWorkspaceSize完成入参校验根据具体计算流程计算 workspace 大小并返回包含算子计算流程的aclOpExecutor执行器第二段aclnnIsClose基于第一段获取的 workspace 与 executor在指定 stream 上真正执行计算。两段必须按顺序配合调用先获取 workspace 大小并据此在 Device 侧申请内存再执行计算。函数原型两段接口的 C 函数原型如下aclnnStatus aclnnIsCloseGetWorkspaceSize( const aclTensor* self, const aclTensor* other, double rtol, double atol, bool equalNan, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnIsClose( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)对应声明位于 math/is_close/op_api/aclnn_isclose.h接口归属aclnn_math领域同时仓库还提供 Level-0 形态的l0op::IsClose见 math/is_close/op_api/isclose.h供上层框架在组装执行器时使用。aclnnIsCloseGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入输入张量公式中的 self。数据类型需与 other 一致shape 需与 other 满足 broadcast 关系。FLOAT、INT32、INT64、FLOAT16、INT16、INT8、UINT8、DOUBLE、BOOL、BFLOAT16ND0-8√otheraclTensor*输入输入张量公式中的 other。数据类型需与 self 一致shape 需与 self 满足 broadcast 关系。FLOAT、INT32、INT64、FLOAT16、INT16、INT8、UINT8、DOUBLE、BOOL、BFLOAT16ND0-8√rtoldouble输入相对容忍度公式中的 rtol。-----atoldouble输入绝对容忍度公式中的 atol。-----equalNanbool输入是否将两个 NaN 视为相等。-----outaclTensor*输出输出张量公式中的 out。数据类型为 BOOLshape 为 self 与 other broadcast 后的 shape。BOOLND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含算子计算流程。-----平台相关的数据类型差异对于 Atlas 训练系列产品与 Atlas 推理系列产品不支持 BFLOAT16。这一点在 API 层实现中也有体现math/is_close/op_api/aclnn_isclose.cpp 维护了两份 dtype 支持列表——DTYPE_SUPPORT_LIST_910无DT_BF16与DTYPE_SUPPORT_LIST_910B含DT_BF16运行时根据当前 NPU 架构DAV_2201或 RegBase 架构动态选择CheckDtypeValid会据此拒绝列表之外的数据类型。返回值与错误码aclnnStatus返回状态码完整定义参见 docs/zh/context/aclnn_return_code.md。第一段接口在入参校验阶段遇到以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self、other 或 out 是空指针。ACLNN_ERR_PARAM_INVALID161002self、other 或 out 的数据类型不在支持范围之内。ACLNN_ERR_PARAM_INVALID161002self 和 other 的数据类型不满足数据类型推导规则。ACLNN_ERR_PARAM_INVALID161002self 和 other 的 shape 无法做 broadcast。ACLNN_ERR_PARAM_INVALID161002out 的 shape 不是 self 和 other broadcast 后的 shape。上述校验与错误码的对应关系在 math/is_close/op_api/aclnn_isclose.cpp 的CheckParams中严格实现CheckNotNull空指针 →ACLNN_ERR_PARAM_NULLPTR、CheckDtypeValiddtype 越界或 self/other 不一致 →ACLNN_ERR_PARAM_INVALID、CheckShape维度超 8、无法 broadcast、out shape 不匹配 →ACLNN_ERR_PARAM_INVALID。aclnnIsClose 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnIsCloseGetWorkspaceSize 获取。executor输入op 执行器包含算子计算流程。stream输入指定执行任务的 Stream。第二段接口的返回值同样为aclnnStatus参见 docs/zh/context/aclnn_return_code.md。约束说明确定性计算aclnnIsClose默认采用确定性实现确定性计算的一般说明参见 docs/zh/context/determinism_compute.md即相同输入下计算结果可复现。其余补充约束可从源码确认rtol / atol 非负Tiling 阶段 math/is_close/op_host/arch35/is_close_tiling_arch35.cpp 会检查rtol 0 || atol 0并直接返回失败要求两者均大于等于 0属性默认值算子定义 math/is_close/op_host/is_close_def.cpp 为三个属性设置了默认值——rtol 1e-05、atol 1e-08、equal_nan false未显式传入时按默认值参与计算Tiling 层读取属性指针为空时同样回退到这些默认值self/other 数据类型必须一致Tiling 层校验x1DType ! x2DType直接返回失败且 kernel 侧仅支持 FLOAT16、FLOAT、INT32、BF16 四类输入aclnn 层通过 Cast 预处理扩展了更多 dtype 的入口能力空 tensorAPI 层对self-IsEmpty() || other-IsEmpty()的空输入直接返回workspaceSize 0空 tensor 由 kernel 支持维度上限 8CheckShape通过OP_CHECK_MAX_DIM(self, MAX_DIM_LEN)限制维度不超过 8。调用示例以下为aclnnIsClose的完整调用示例同时存在于 math/is_close/examples/test_aclnn_isclose.cpp具体编译和执行过程请参考 docs/zh/context/compile_and_run_sample.md#include cinttypes #include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_isclose.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 shape_size 1; for (auto i : shape) { shape_size * i; } return shape_size; } 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根据自己的需要处理 CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口自定义构造 std::vectorint64_t selfShape {4, 2}; std::vectorint64_t otherShape {4, 2}; std::vectorint64_t outShape {4, 2}; void* selfDeviceAddr nullptr; void* otherDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* other nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0, 1, 2, 3, 4, 5, 6, 7}; std::vectorfloat otherHostData {1, 1, 1, 2, 1, 2, 3, 3}; std::vectoruint8_t outHostData {0, 0, 0, 0, 0, 0, 0, 0}; double rtol 1.0; double atol 1.0; bool equal_nan false; // 创建other aclTensor ret CreateAclTensor(otherHostData, otherShape, otherDeviceAddr, aclDataType::ACL_FLOAT, other); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建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 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnIsClose第一段接口 ret aclnnIsCloseGetWorkspaceSize(self, other, rtol, atol, equal_nan, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnIsCloseGetWorkspaceSize 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); } // 调用aclnnIsClose第二段接口 ret aclnnIsClose(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnIsClose 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侧 auto size GetShapeSize(outShape); std::vectoruint8_t 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[% PRId64 ] is: %d\n, i, resultData[i]); } // 6. 释放aclTensor aclDestroyTensor(other); aclDestroyTensor(self); aclDestroyTensor(out); // 7. 释放device资源 aclrtFree(otherDeviceAddr); aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }运行结果预期示例中self {0,1,2,3,4,5,6,7}、other {1,1,1,2,1,2,3,3}rtol atol 1.0。以self[0]0, other[0]1为例|0-1|1 ≤ 1 1×|1|2结果为1self[1]1, other[1]1完全相等结果为1。全部 8 个元素均在容差范围内最终输出应为{1,1,1,1,1,1,1,1}bool 以uint8_t打印为 0/1。源码级实现链路从 API 到 Vector 指令API 层非连续输入与 view 语义第一段接口在 math/is_close/op_api/aclnn_isclose.cpp 中按固定模板组装执行流CheckParams完成空指针、dtype、shape 三类校验对应错误码表将double类型的rtol/atol转为float32下传给 kernelTiling 侧属性读取的也是float空 tensor 提前返回workspaceSize 0通过l0op::Contiguous将可能非连续的self/other转换为连续 tensor非连续输入说明参见 docs/zh/context/non_contiguous_tensor.md调用l0op::IsClose生成计算节点通过l0op::ViewCopy将中间结果拷贝到用户提供的out上保证out为非连续 tensor 时也能正确写出uniqueExecutor-GetWorkspaceSize()汇总整个计算流程的 workspace 需求并返回。Tiling 层按 dtype × equal_nan 组合选择计算模板math/is_close/op_host/arch35/is_close_tiling_arch35.cpp 的DoOpTiling根据输入数据类型与equal_nan取值在IsCloseDagequal_nanfalse与IsCloseEqualNanDagequal_nantrue之间选择模板并组合BroadcastSch调度模式生成tilingKey随后通过brcBaseTiling.SetScalarfloat(rtol_)/SetScalarfloat(atol_)将两个容差参数作为标量写入 Tiling 数据供 kernel 以Placeholder::Varfloat读取。TilingPrepareForIsClose负责读取 AIV 核数与 UB 大小为 Broadcast 调度做准备。算子定义与推导math/is_close/op_host/is_close_def.cpp注册IsClose算子输入x1、x2输出yBOOL三个可选属性rtol默认1e-05、atol默认1e-08、equal_nan默认false并为ascend950、ascend350注册 AICore 配置开启动态 shape、动态 rank 支持math/is_close/op_host/is_close_infershape.cpp复用Ops::Base::InferShape4Broadcast推导输出 shape即两输入 broadcast 后的 shape并将输出数据类型固定为DT_BOOL。Kernel 层AIV 单核 Vector 流水math/is_close/op_kernel/is_close_apt.cpp 以KERNEL_TYPE_AIV_ONLY模式启动将工作委托给BroadcastSch数据经CopyInBrc广播搬入、Castfloat, T, 0统一提升到 float 计算、NanEqualCompare完成逐元素判定、CopyOutuint8_t写出 BOOL 结果MemOptCfgMemLevel::LEVEL_2控制二级缓存优化策略。其中NanEqualCompare采用 256-bit 向量寄存器单次 8 个 float循环处理先Compare产出掩码再经Pack压缩后StoreAlign写出保证了大 tensor 上的吞吐。测试与验证仓库为aclnnIsClose提供了多维度的验证资产ATK 用例math/is_close/tests/st/aclnnIsClose/atk_aclnnIsClose.json 基于torch.isclose语义生成 210 条用例覆盖fp32 / fp16 / bf16 / int8 / int16 / int32 / int64 / fp64 / bool全部支持的 dtypeshape 覆盖 4 维的各种广播形态含256、257等大维度边界以及[1, 1, 257, 257]这类极端尺寸rtol/atol取0/1/2与[0,2]区间随机equal_nan真假交替并标注is_boundary边界用例ST 测试math/is_close/tests/st/arch35/ttk_kernel_is_close_st.csv 面向 arch35 的 kernel 级测试UT 测试math/is_close/tests/ut/op_host/test_is_close_infershape.cpp 与 math/is_close/tests/ut/op_api/test_aclnn_isclose.cpp 分别覆盖 infershape 推导与 API 调用路径Golden 基准math/is_close/tests/assets/golden.py 提供基准数据生成逻辑。通过这些用例可以验证dtype 一致性校验、broadcast shape 推导、equal_nan语义切换对应 kernel 中两套模板分支、非连续 tensor 支持以及各平台 BF16 支持差异。小结aclnnIsClose是 ops-math 中实现元素级容差相等比较的标准算子其两段式接口设计、|self - other| atol rtol * |other|的判定公式、非有限值与 NaN 的特殊语义以及从 aclnn API → Tiling → Broadcast Vector kernel 的完整实现链在 math/is_close 模块的源码与测试中都有完整且自洽的呈现。开发者可按本文示例直接落地调用也可参考其架构快速理解同类 broadcast 型数学算子的实现范式。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考