算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载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)非连续Tensorself输入表示第一个输入self 与 other 的数据类型满足数据类型推导规则参见互推导关系FLOAT16、FLOAT、INT32、INT8、UINT8、BOOL、DOUBLE、INT64、INT16、UINT16、UINT32、UINT64、BFLOAT16ND-√other输入表示第二个输入other 与 self 的数据类型满足数据类型推导规则参见互推导关系FLOAT16、FLOAT、INT32、INT8、UINT8、BOOL、DOUBLE、INT64、INT16、UINT16、UINT32、UINT64、BFLOAT16ND-√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_NULLPTR161001传入的 self、other 是空指针时ACLNN_ERR_PARAM_INVALID161002self 和 other 推导后的数据类型不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self、other、out 的维度大于 8ACLNN_ERR_PARAM_INVALID161002out 的 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::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(), 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::vectorint64_t selfShape, std::vectorint64_t otherShape, std::vectorint64_t outShape, void** selfDeviceAddr, void** otherDeviceAddr, void** outDeviceAddr, aclTensor** self, aclTensor** other, aclTensor** out) { std::vectordouble selfHostData {0, 1, 2, 3, 4, 5, 6, 7}; std::vectordouble otherHostData {0, 1, 2, 3, 4, 5, 6, 7}; std::vectorchar 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::vectorint64_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::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]); } return ACL_SUCCESS; } int main() { int32_t deviceId 0; aclrtStream stream; auto ret InitAcl(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, return ret); std::vectorint64_t selfShape {4, 2}; std::vectorint64_t otherShape {4, 2}; std::vectorint64_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→ 构造输入输出 aclTensoraclrtMalloc 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 不同或两者均为空 TensorIsEmpty()则无需真正计算——直接通过l0op::Fill生成结果shape 相同且都为空时填True1否则填False0再经ViewCopy拷入 out 返回。这是对大小相同语义的快捷路径实现。类型推导与统一对输入做PromoteType互推导得到隐式公共类型若推导结果为 BF16 则提升为 FLOAT因为比较在更高精度下进行更稳妥同时规避部分平台 BF16 比较指令限制。连续化与 Cast对self、other分别执行l0op::Contiguous转为连续张量再l0op::Cast到公共类型。核心计算调用l0op::TensorEqual(selfCasted, otherCasted, ...)生成一个shape[1]、dtypeBOOL的中间结果。结果回拷经l0op::ViewCopy将中间结果写入用户提供的 outout 可能是非连续 TensorViewCopy 保证正确落位。汇总 workspace最后通过uniqueExecutor-GetWorkspaceSize()汇总整张图所需 workspace 大小并将 executor 释放给调用方。8.2 TensorEqual 算子的 AICore/AICPU 分流在 math/tensor_equal/op_api/tensor_equal.cpp 中l0op::TensorEqual会根据输入数据类型与平台架构选择执行路径定义了两张平台相关支持列表AICORE910_DTYPE_SUPPORT_LIST1980 芯片含 FLOAT/INT32/FLOAT16/INT8/UINT8/BOOL/BF16与ARCH3510_DTYPE_SUPPORT_LISTRegBase 架构覆盖 12 种常见整型/浮点类型IsAiCoreSupport()依据IsRegBase()判断当前架构后查询对应列表两个输入均在支持列表内则走TensorEqualAiCoreAICore 向量 Kernel否则走TensorEqualAiCpuAICPU 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_BOOLAICore 配置开启动态 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。三种 TilingKeyNORMAL_SHAPE_TILINGKEY 121常规同 shape 场景、DIFFER_SHAPE_TILINGKEY 111shape 不同仅需 1 核、EMPTY_SHAPE_TILINGKEY 101含空 shape仅需 1 核。切分策略为按总核数均分inputShapeSize单核处理量下限MIN_CORE_PROCESS 2048并对齐向量寄存器长度vRegFactorUB 侧基于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模板类TensorEqualKernelT先通过SetSysWorkspace获取用户 workspace依据 TilingKey 分发到正常/异形/空 shape 三条路径TensorEqualKernel::Init按blockIdx计算各核的数据偏移并仅由 0 号核初始化全局输出为1NORMAL_OUTPUT异形场景为0之后用InitGlobalMemory写初值并通过事件同步保证可见性CopyIn使用DataCopyPad双缓冲搬入 x、y 分块Compute利用AscendC::Reg寄存器级指令按向量寄存器长度分块CompareCMPMODE::NE比较 x 与 y 是否不相等得到比较掩码后逐块Or累加到bakMaskReg最后对掩码做ReduceMAX归约——若存在任一元素不等掩码非零累加结果不为 0则将全局输出置0。整型输入统一以uint8_t视角InputType conditional_tis_integral_vT, 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.cppTiling 单测见 math/tensor_equal/tests/ut/op_host/arch35/test_tensor_equal_tiling_arch35.cppKernel 单测见 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 与上述源码路径深入研读并在实际项目中复用本文的示例代码快速上手。赞分享算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载相关推荐CANN ops-math 的 aclnnNanToNum / aclnnInplaceNanToNum 算子接口详解与两段式调用实战CANN ops math 的 aclnnNanToNum / aclnnInplaceNanToNum 算子接口详解与两段式调用实战 本指南以 CANN op算子库人工智能CANNCANN ops-math 算子库 aclnnCat 接口详解基于 ConcatD 的 tensor 级联两段式调用指南CANN ops math 算子库 aclnnCat 接口详解基于 ConcatD 的 tensor 级联两段式调用指南 aclnnCat 是 CANN op算子库人工智能CANNCANN ops-math 算子详解aclnnXLogYScalarOther 与 aclnnInplaceXLogYScalarOther 两段式接口调用指南CANN ops math 算子详解aclnnXLogYScalarOther 与 aclnnInplaceXLogYScalarOther 两段式接口调用指算子库人工智能CANN创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考