CANN ops-math 算子库 aclnnEqual 接口详解：Tensor 全量相等性判定与两段式调用实践-平芜编程栈

算子库
人工智能
CANN

【免费下载链接】ops-math

本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。

aclnnEqual 是 CANN ops-math 数学算子库中 TensorEqual 算子面向昇腾 AI 处理器（NPU）的 AscendCL 算子层（aclnn）编程接口，用于判断两个 Tensor 是否具有相同的大小（shape）且全部元素相等，并返回一个 Bool 标量。本文基于 math/tensor_equal/docs/aclnnEqual.md 文档，结合 math/tensor_equal 目录下的 op_api、op_host、op_kernel 源码与测试用例，完整讲解接口语义、两段式调用流程、参数与返回码、约束限制，并给出可编译运行的完整示例与底层实现剖析，帮助读者在 NPU 上正确、高效地完成张量相等性判定。

一、功能说明与计算语义

aclnnEqual 的接口功能为：计算两个 Tensor 是否具有相同的大小和元素，返回一个 Bool 类型的结果。它与逐元素比较算子（如逐位==）不同，输出只有一个元素：仅当两个输入在形状与全部元素值上都一致时才为True，否则为False。计算表达式为：

$$ out = (self == other) \ ? \ True : False $$

从源码实现看，该算子即 PyTorchtorch.equal的 NPU 等价实现。仓库中的 golden 参考脚本 math/tensor_equal/tests/assets/golden.py 直接使用torch.equal生成期望结果（aclnn_equal_golden与tensor_equal_golden两个函数），并在输入端将 bfloat16 提升为 float32 后再比较，这与接口层"先做数据类型互推导、再统一 Cast 后比较"的实现思路一致。

二、产品支持情况

原文档明确列出的产品支持矩阵如下（当前仓库源码配置也仅注册了对应芯片架构的 AICore 实现，见下文"底层实现"一节）：

产品	是否支持
Ascend 950PR / Ascend 950DT	支持
Atlas A3 训练系列产品 / Atlas A3 推理系列产品	支持
Atlas A2 训练系列产品 / Atlas A2 推理系列产品	支持
Atlas 200I/500 A2 推理产品	不支持
Atlas 推理系列产品	支持
Atlas 训练系列产品	支持

在仓库源码侧，算子 AICore 配置仅注册了ascend950与ascend350两个芯片代次（见 math/tensor_equal/op_host/tensor_equal_def.cpp 中的this->AICore().AddConfig("ascend950", aicoreConfig)与AddConfig("ascend350", aicoreConfig)），分别对应 Ascend 950 与 Atlas A3 系列产品；而 Atlas A2 训练/推理系列产品（昇腾 910B）通过接口层的 AICore/AICPU 分流逻辑提供支持（详见"底层实现"一节），Atlas 200I/500 A2 推理产品不受支持。

三、两段式接口与函数原型

每个 aclnn 算子都遵循两段式接口设计：必须先调用第一段接口aclnnEqualGetWorkspaceSize完成入参校验并获取计算所需 workspace 大小及包含算子计算流程的执行器，再调用第二段接口aclnnEqual实际执行计算。

第一段接口原型：

aclnnStatus aclnnEqualGetWorkspaceSize( const aclTensor* self, const aclTensor* other, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)

第二段接口原型：

aclnnStatus aclnnEqual( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

两段式设计的价值在于：第一段是"纯 Host 侧"操作，可提前做参数校验、构建算子图（l0op 组合）、计算 workspace 需求，让用户在第一段返回后统一申请 Device 内存，避免执行阶段出现资源不足；第二段才是真正的任务下发。具体实现见 math/tensor_equal/op_api/aclnn_equal.cpp 中的aclnnEqualGetWorkspaceSize与aclnnEqual两个函数。

四、aclnnEqualGetWorkspaceSize 参数说明

第一段接口的参数如下表所示：

参数名	输入/输出	描述	使用说明	数据类型	数据格式	维度(shape)	非连续Tensor
self	输入	表示第一个输入	self 与 other 的数据类型满足数据类型推导规则（参见互推导关系）	FLOAT16、FLOAT、INT32、INT8、UINT8、BOOL、DOUBLE、INT64、INT16、UINT16、UINT32、UINT64、BFLOAT16	ND	-	√
other	输入	表示第二个输入	other 与 self 的数据类型满足数据类型推导规则（参见互推导关系）	FLOAT16、FLOAT、INT32、INT8、UINT8、BOOL、DOUBLE、INT64、INT16、UINT16、UINT32、UINT64、BFLOAT16	ND	-	√
out	输出	表示输出。输出一个数据类型为 BOOL、一维包含一个元素的 Tensor	-	-	-	-	-
workspaceSize	输出	返回需要在 Device 侧申请的 workspace 大小	-	-	-	-	-
executor	输出	返回 op 执行器，包含了算子计算流程	-	-	-	-	-

需要特别说明两点：

数据类型支持范围：Atlas 训练系列产品、Atlas 推理系列产品（即昇腾 910 系列）不支持 BFLOAT16数据类型。这一差异在源码中有明确体现：math/tensor_equal/op_api/aclnn_equal.cpp 中定义了DTYPE_SUPPORT_910B_LIST（含DT_BF16）与DTYPE_SUPPORT_910_LIST（不含DT_BF16）两张列表，并通过CheckSocVersionGe910B()（判断当前 NPU 架构是否为DAV_2201或 RegBase 架构）在运行时选择使用哪张校验列表。
非连续 Tensor 支持：两个输入均支持非连续 Tensor（表格中"√"）。从源码看，接口内部会对输入统一执行l0op::Contiguous转连续，因此用户无需自行调用aclrtMemset或拷贝预整理。

返回码与入参校验

第一段接口返回aclnnStatus状态码，具体参见aclnn返回码。第一段接口完成入参校验，出现以下场景时报错：

返回码	错误码	描述
ACLNN_ERR_PARAM_NULLPTR	161001	传入的 self、other 是空指针时
ACLNN_ERR_PARAM_INVALID	161002	self 和 other 推导后的数据类型不在支持的范围之内
ACLNN_ERR_PARAM_INVALID	161002	self、other、out 的维度大于 8
ACLNN_ERR_PARAM_INVALID	161002	out 的 shape 不是 [1]

这些校验项与源码中CheckParams的检查顺序一一对应（见 math/tensor_equal/op_api/aclnn_equal.cpp）：先CheckNotNull（空指针检查，返回ACLNN_ERR_PARAM_NULLPTR），再CheckDtypeValid（数据类型校验，含PromoteType互推导）、CheckMaxShape（最大维度检查，DIM_SUPPORT_MAX = 8）与CheckOutShape（输出 shape 必须为[1]），后三者均返回ACLNN_ERR_PARAM_INVALID。

五、aclnnEqual 参数说明

第二段接口参数如下：

参数名	输入/输出	描述
workspace	输入	在 Device 侧申请的 workspace 内存地址
workspaceSize	输入	在 Device 侧申请的 workspace 大小，由第一段接口 aclnnEqualGetWorkspaceSize 获取
executor	输入	op 执行器，包含了算子计算流程
stream	输入	指定执行任务的 Stream

第二段接口同样返回aclnnStatus状态码。其内部实现非常简洁：调用框架的CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成计算（见 math/tensor_equal/op_api/aclnn_equal.cpp 第 217-222 行），并埋入L2_DFX_PHASE_2(aclnnEqual)打点用于算子级 DFX 观测。

六、约束说明

原文档给出的约束如下：

确定性计算：aclnnEqual 默认确定性实现，即相同输入多次运行输出结果确定，无随机性。
超时风险场景：如果计算量过大可能会导致算子执行超时（aicore error 类型报错，errorStr 为timeout or trap error），典型场景为最后 2 轴合轴小于 16、前面的轴合轴超大。该约束提示用户：当两个输入 shape 接近但元素规模极大时，应评估任务切分与资源配比，必要时拆分输入或调整计算规模。

七、调用示例

下面给出完整的示例代码，其中数据为self = {0,1,2,3,4,5,6,7}（shape[4,2]）与other = {0,1,2,3,4,5,6,7}（shape[4,2]），两输入完全相等，预期输出为result[0] is: 1。编译与执行过程请参考编译与运行样例。仓库中还提供了采用 RAII 智能指针管理资源的增强版示例 math/tensor_equal/examples/test_aclnn_equal.cpp，以及完整的 aclnn 调用方式说明（见 math/tensor_equal/README.md）。

#include <iostream> #include <vector> #include "acl/acl.h" #include "aclnnop/aclnn_equal.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::vector<int64_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::vector<T>& hostData, const std::vector<int64_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::vector<int64_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(), shape.size(), aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } aclError InitAcl(int32_t deviceId, aclrtStream* stream) { auto ret = Init(deviceId, stream); CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret); return ACL_SUCCESS; } aclError CreateInputs( std::vector<int64_t>& selfShape, std::vector<int64_t>& otherShape, std::vector<int64_t>& outShape, void** selfDeviceAddr, void** otherDeviceAddr, void** outDeviceAddr, aclTensor** self, aclTensor** other, aclTensor** out) { std::vector<double> selfHostData = {0, 1, 2, 3, 4, 5, 6, 7}; std::vector<double> otherHostData = {0, 1, 2, 3, 4, 5, 6, 7}; std::vector<char> outHostData = {0}; auto ret = CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_DOUBLE, self); CHECK_RET(ret == ACL_SUCCESS, return ret); ret = CreateAclTensor(otherHostData, otherShape, otherDeviceAddr, aclDataType::ACL_DOUBLE, other); CHECK_RET(ret == ACL_SUCCESS, return ret); ret = CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_BOOL, out); CHECK_RET(ret == ACL_SUCCESS, return ret); return ACL_SUCCESS; } aclError ExecOpApi( aclTensor* self, aclTensor* other, aclTensor* out, void** workspaceAddrOut, uint64_t& workspaceSize, void* outDeviceAddr, std::vector<int64_t>& outShape, aclrtStream stream) { aclOpExecutor* executor; auto ret = aclnnEqualGetWorkspaceSize(self, other, out, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnEqualGetWorkspaceSize 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); } *workspaceAddrOut = workspaceAddr; ret = aclnnEqual(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnEqual 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::vector<char> 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]); } return ACL_SUCCESS; } int main() { int32_t deviceId = 0; aclrtStream stream; auto ret = InitAcl(deviceId, &stream); CHECK_RET(ret == ACL_SUCCESS, return ret); std::vector<int64_t> selfShape = {4, 2}; std::vector<int64_t> otherShape = {4, 2}; std::vector<int64_t> outShape = {1}; void* selfDeviceAddr = nullptr; void* otherDeviceAddr = nullptr; void* outDeviceAddr = nullptr; aclTensor* self = nullptr; aclTensor* other = nullptr; aclTensor* out = nullptr; ret = CreateInputs( selfShape, otherShape, outShape, &selfDeviceAddr, &otherDeviceAddr, &outDeviceAddr, &self, &other, &out); CHECK_RET(ret == ACL_SUCCESS, return ret); uint64_t workspaceSize = 0; void* workspaceAddr = nullptr; ret = ExecOpApi(self, other, out, &workspaceAddr, workspaceSize, outDeviceAddr, outShape, stream); CHECK_RET(ret == ACL_SUCCESS, return ret); // 释放 aclDestroyTensor(self); aclDestroyTensor(other); aclDestroyTensor(out); aclrtFree(selfDeviceAddr); aclrtFree(otherDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize > 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }

示例代码的关键执行流程可归纳为：初始化（aclInit → aclrtSetDevice → aclrtCreateStream）→ 构造输入输出 aclTensor（aclrtMalloc + aclrtMemcpy + aclCreateTensor）→ 第一段接口获取 workspace 与 executor → 申请 workspace → 第二段接口执行 → 同步 Stream → 结果拷回 Host → 释放资源。注意 workspaceSize 为 0 时无需申请 workspace，直接传nullptr即可。

八、底层实现原理：从接口到 Kernel 的完整调用链

8.1 接口层的算子组合流程

aclnnEqualGetWorkspaceSize的核心逻辑（见 math/tensor_equal/op_api/aclnn_equal.cpp 注释中的流程示意图）是构建一条由多个 l0op 算子拼接而成的计算图：

self other | | \ / Contiguous(workspace_0) Contiguous(workspace_1) \ / TensorEqual(workspace_2) | ViewCopy | result

具体步骤如下：

形状特判：若self与other的视图 shape 不同，或两者均为空 Tensor（IsEmpty()），则无需真正计算——直接通过l0op::Fill生成结果：shape 相同且都为空时填True（1），否则填False（0），再经ViewCopy拷入 out 返回。这是对"大小相同"语义的快捷路径实现。
类型推导与统一：对输入做PromoteType互推导得到隐式公共类型，若推导结果为 BF16 则提升为 FLOAT（因为比较在更高精度下进行更稳妥，同时规避部分平台 BF16 比较指令限制）。
连续化与 Cast：对self、other分别执行l0op::Contiguous转为连续张量，再l0op::Cast到公共类型。
核心计算：调用l0op::TensorEqual(selfCasted, otherCasted, ...)生成一个shape=[1]、dtype=BOOL的中间结果。
结果回拷：经l0op::ViewCopy将中间结果写入用户提供的 out（out 可能是非连续 Tensor，ViewCopy 保证正确落位）。
汇总 workspace：最后通过uniqueExecutor->GetWorkspaceSize()汇总整张图所需 workspace 大小，并将 executor 释放给调用方。

8.2 TensorEqual 算子的 AICore/AICPU 分流

在 math/tensor_equal/op_api/tensor_equal.cpp 中，l0op::TensorEqual会根据输入数据类型与平台架构选择执行路径：

定义了两张平台相关支持列表：AICORE910_DTYPE_SUPPORT_LIST（1980 芯片，含 FLOAT/INT32/FLOAT16/INT8/UINT8/BOOL/BF16）与ARCH3510_DTYPE_SUPPORT_LIST（RegBase 架构，覆盖 12 种常见整型/浮点类型）；
IsAiCoreSupport()依据IsRegBase()判断当前架构后查询对应列表；两个输入均在支持列表内则走TensorEqualAiCore（AICore 向量 Kernel），否则走TensorEqualAiCpu（AICPU Kernel，通过ADD_TO_LAUNCHER_LIST_AICPU注册）；
输出 Tensor 由executor->AllocTensor分配，shape 固定为[1]、类型固定为DT_BOOL。

这解释了文档中"Atlas 训练系列产品、Atlas 推理系列产品不支持 BFLOAT16"以及 aclnn 层数据类型支持列表在不同产品上的差异——底层 Kernel 的数据类型能力决定了接口层的校验范围。

8.3 Host 侧 Tiling 切分

算子注册与 Tiling 逻辑位于 math/tensor_equal/op_host/tensor_equal_def.cpp 与 math/tensor_equal/op_host/arch35/tensor_equal_tiling_arch35.cpp：

算子定义：TensorEqual注册两个输入input_x、input_y（均要求 REQUIRED，支持 12 种 ND 格式数据类型）与一个输出output_z（仅DT_BOOL）；AICore 配置开启动态 shape、动态 rank 支持，ExtendCfgInfo指向 kernel 实现文件tensor_equal_apt，并在ascend950、ascend350两个代次注册。
Tiling 流程：继承TilingBaseClass后按 8 个阶段完成：GetPlatformInfo（获取 CoreNum 与 UB 大小，来自TilingPrepare4TensorEqual填充的TensorEqualCompileInfo）→GetShapeAttrsInfo（读取输入 shape 与 dtype，计算 shape 大小）→DoOpTiling（切分）→DoLibApiTiling→GetTilingKey→GetWorkspaceSize（返回固定 32 字节WORKSPACE_SAVE_SIZE的 save 空间）→PostTiling（设置SetBlockDim、调度模式与 TilingKey）→DumpTilingInfo。
三种 TilingKey：NORMAL_SHAPE_TILINGKEY = 121（常规同 shape 场景）、DIFFER_SHAPE_TILINGKEY = 111（shape 不同，仅需 1 核）、EMPTY_SHAPE_TILINGKEY = 101（含空 shape，仅需 1 核）。切分策略为：按总核数均分inputShapeSize，单核处理量下限MIN_CORE_PROCESS = 2048并对齐向量寄存器长度vRegFactor；UB 侧基于ubSize扣除保留空间后按双缓冲、双输入占用计算ubFactor，进而推出每核循环次数与尾部因子（perCoreLoopTimes/tailCoreLoopTimes/perCoreTailFactor/tailCoreTailFactor）。

8.4 Kernel 侧向量化比较

AICore Kernel 位于 math/tensor_equal/op_kernel/tensor_equal_apt.cpp（入口tensor_equal）与 math/tensor_equal/op_kernel/arch35/tensor_equal.h（模板类TensorEqualKernel<T>）：

先通过SetSysWorkspace获取用户 workspace，依据 TilingKey 分发到正常/异形/空 shape 三条路径；
TensorEqualKernel::Init按blockIdx计算各核的数据偏移，并仅由 0 号核初始化全局输出为1（NORMAL_OUTPUT，异形场景为0），之后用InitGlobalMemory写初值并通过事件同步保证可见性；
CopyIn使用DataCopyPad双缓冲搬入 x、y 分块；
Compute利用AscendC::Reg寄存器级指令按向量寄存器长度分块：Compare<CMPMODE::NE>比较 x 与 y 是否不相等，得到比较掩码后逐块Or累加到bakMaskReg，最后对掩码做Reduce<MAX>归约——若存在任一元素不等，掩码非零，累加结果不为 0，则将全局输出置0。整型输入统一以uint8_t视角（InputType = conditional_t<is_integral_v<T>, uint8_t, T>）参与比较与归约，保证位级语义一致；
每个核的归约结果在Process中累加进saveBuf的saveLocal，最终若有核发现不等（saveLocal(0) != 0且全局输出仍为 1），则把输出写成 0 并做DataCacheCleanAndInvalid缓存一致性处理。

可见 aclnnEqual 采用的是"逐元素比较 + 跨核归约取反"的并行策略：不等即置 0，全部相等才保持初始的 1，与文档所述语义完全吻合。

九、测试与验证

仓库为 aclnnEqual 提供了多层验证手段，可用于自行验证接口行为：

UT 单测：op_api 层测试见 math/tensor_equal/tests/ut/op_api/test_aclnn_tensor_equal.cpp；Tiling 单测见 math/tensor_equal/tests/ut/op_host/arch35/test_tensor_equal_tiling_arch35.cpp；Kernel 单测见 math/tensor_equal/tests/ut/op_kernel/test_tensor_equal.cpp。
ST 系统测试：ATK 用例定义与执行脚本见 math/tensor_equal/tests/st/aclnnEqual/atk_aclnnEqual.json 与 math/tensor_equal/tests/st/aclnnEqual/executor_aclnnEqual.py；架构 35 平台的 ST 用例矩阵见 math/tensor_equal/tests/st/arch35/ttk_aclnn_equal_st.csv 与 math/tensor_equal/tests/st/arch35/ttk_kernel_tensor_equal_st.csv。
golden 参照：math/tensor_equal/tests/assets/golden.py 使用torch.equal生成 aclnn 与 kernel 两级期望结果，同时覆盖了 bfloat16 输入先提升到 float32 再比较的行为。

十、总结

aclnnEqual 是 CANN ops-math 中实现 Tensor 全量相等判定的标准接口，语义等价于torch.equal：形状与元素全部一致才返回True。使用上必须遵循两段式调用（先aclnnEqualGetWorkspaceSize后aclnnEqual），注意 910 系列产品不支持 BFLOAT16、输出 shape 必须为[1]、维度上限 8 等约束；实现上由接口层完成连续化、类型推导与统一 Cast，由 Tiling 层完成多核与 UB 切分，由 Kernel 层以"逐元素不等比较 + 跨核归约"方式得到布尔结果。读者可结合 math/tensor_equal/docs/aclnnEqual.md、math/tensor_equal/README.md 与上述源码路径深入研读，并在实际项目中复用本文的示例代码快速上手。