CANN ops-nn 算子库 aclnnUnique 接口详解：NPU 全局去重与逆索引计算实战指南

CANN ops-nn 算子库 aclnnUnique 接口详解NPU 全局去重与逆索引计算实战指南【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn本文基于 CANN ops-nn 仓库 aclnnUnique 接口文档 编写深入讲解aclnnUnique全局去重算子的两段式调用方式、参数语义、错误码与平台约束并结合 op_api 源码 与 测试用例 揭示其底层基于 Sort、UniqueConsecutive、ScatterElements 等算子的组合实现原理。读完本文你将掌握在 NPUAscend 平台上正确调用aclnnUnique完成取唯一值 可选逆索引计算的完整方法并能理解其与torch.unique的行为对应关系。功能说明返回输入张量中的唯一元素aclnnUnique是 CANN ops-nn 算子库中负责全局去重的接口核心功能是返回输入张量self中的唯一元素并支持两个可选行为sorted是否对去重结果valueOut按升序排序returnInverse是否同时返回原输入self中各个元素在valueOut中的位置下标即逆索引。该接口语义与 PyTorch 的torch.unique(input, sorted, return_inverse)对齐。在仓库的 ST 测试中atk_aclnnUnique.json 将aclnnUnique与torch.unique一一对照而 executor_aclnnUnique.py 中 CPU 基准实现即为valueOut, inverseOut torch.unique(input, sorted, returnInverse, False)aclnnUnique是index/unique目录提供的两个全局去重接口之一另一个为aclnnUnique2可额外返回每个唯一元素的出现次数两者共用 op_api/unique_common.cpp 中的内部实现详见 index/unique/README.md。产品支持情况产品支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持两段式接口与函数原型aclnnUnique遵循 CANN 算子库的两段式接口规范必须先调用aclnnUniqueGetWorkspaceSize获取计算所需的 workspace 大小以及封装了算子计算流程的执行器再调用aclnnUnique真正执行计算。aclnnStatus aclnnUniqueGetWorkspaceSize( const aclTensor* self, bool sorted, bool returnInverse, aclTensor* valueOut, aclTensor* inverseOut, uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnUnique( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)两个接口在 op_api/aclnn_unique.h 中以extern C方式导出头文件即用户侧所需引入的aclnnop/aclnn_unique.h。头文件中的注释还给出了算子计算的基本路径aclnn_unique.hSelf ──▶ Contiguous ──▶ UniqueWithCountsAndSorting ──▶ valueOut / inverseOut ▲ sorted ───────┘ returnInverse ─┘aclnnUniqueGetWorkspaceSize 参数说明第一段接口完成入参校验与 workspace 规划其参数定义如下参数名输入/输出描述使用说明数据类型数据格式维度shape非连续 Tensorself输入待去重的输入张量-BOOL、FLOAT、FLOAT16、DOUBLE、UINT8、INT8、UINT16、INT16、INT32、UINT32、UINT64、INT64、BFLOAT16ND≤8-sorted输入表示是否对 valueOut 按升序进行排序。true 时排序false 时乱序-----returnInverse输入表示是否返回输入数据中各个元素在 valueOut 中的下标-----valueOut输出第一个输出张量输入张量中的唯一元素-BOOL、FLOAT、FLOAT16、DOUBLE、UINT8、INT8、UINT16、INT16、INT32、UINT32、UINT64、INT64、BFLOAT16---inverseOut输出第二个输出张量当 returnInverse 为 True 时有意义返回 self 中各元素在 valueOut 中出现的位置下标-----workspaceSize输出返回需要在 Device 侧申请的 workspace 大小-----executor输出返回 op 执行器包含了算子计算流程-----从源码看参数校验逻辑对照 op_api/aclnn_unique.cpp 中的CheckParams实现第一段接口的入参检查分为三步空指针检查CheckNotNull3Tensorself、valueOut、inverseOut任一为空即返回ACLNN_ERR_PARAM_NULLPTR数据类型检查CheckDtypeValidself必须在支持列表内且inverseOut必须为DT_INT64见OP_CHECK_DTYPE_NOT_MATCH(inverseOut, DT_INT64)Shape 检查CheckShapeValidself维度不能超过 8MAX_SUPPORT_DIMS_NUMS当returnInverse为 true 时inverseOut与self的 shape 必须一致。此外源码中还做了空张量短路处理若self-IsEmpty()直接返回workspaceSize 0无需构建计算图见 aclnn_unique.cpp。值得注意的是虽然文档表格中非连续 Tensor一列标注为空但接口内部会先通过l0op::Contiguous将输入转换为连续张量aclnn_unique.cpp因此调用方无需自行做内存连续性处理。返回值与错误码aclnnStatus返回状态码完整定义参见 aclnn返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self 或 valueOut 或 inverseOut 是空指针时ACLNN_ERR_PARAM_INVALID161002self 或 valueOut 的数据类型不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self 为非连续张量ACLNN_ERR_PARAM_INVALID161002returnInverse 为 True且 inverseOut 与 self shape 不一致ACLNN_ERR_PARAM_INVALID161002returnInverse 为 TrueinverseOut 的数据类型不为 INT64说明错误码 161001/161002 的数值对应关系可在 docs/zh/context/aclnn_return_code.md 中查询参数校验规则与上文源码分析一致。aclnnUnique 第二段接口参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnUniqueGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream第二段接口的实现非常简洁源码中仅包含L2_DFX_PHASE_2打点与一行CommonOpExecutorRun(workspace, workspaceSize, executor, stream)见 aclnn_unique.cpp即把第一段构建好的执行器与 workspace 交给框架统一调度执行这也是所有 L2 级算子接口的固定写法。约束说明确定性计算aclnnUnique默认为确定性实现相关背景可参考 确定性计算说明。Ascend 950PR / Ascend 950DT 平台由于去重算法实现差异当满足下列所有条件时算子将无视sorted入参的值固定对输出结果进行升序排序self输入为 1Dself的数据类型为FLOAT、FLOAT16、UINT8、INT8、UINT16、INT16、INT32、UINT32、UINT64、INT64、BFLOAT16。Atlas A2 / Atlas A3 系列产品在输入self包含 0 的情况下算子的输出中会包含正 0 和负 0而非只输出一个 0。底层实现原理从源码看去重与逆索引的计算链路index/unique模块的特点是不包含独立 Kernel而是通过组合底层算子完成功能见 index/unique/README.md。其计算路径在 aclnn_unique.cpp 中根据平台动态分支路径一AI Core 组合路径Ascend 950 等支持 AI Core 去重的平台当UniqueCommon::SupportAicore4Unique返回 true即 regbase 架构且数据类型满足支持列表时走ComputeUniqueViaAicore组合链路核心步骤为FlattenAndSortunique_common.cpp用l0op::Reshape将任意维输入展平为一维用l0op::Sort对展平数据升序排序同时得到排序后的值sortedValues与排序索引sortedIndices。UniqueConsecutive对排序后的相邻重复值做连续去重得到valueOut排序后重复值必然相邻因此连续去重等价于全局去重。ComputeInverseIndicesunique_common.cpp——仅在returnInverse为 true 时执行l0op::AdjacentDifference计算相邻元素差值用于标记每个唯一值首次出现的位置l0op::Cumsum对差值做前缀和得到每个元素对应的唯一值序号l0op::ScatterElements按sortedIndices把序号 scatter 回原输入位置得到逆索引l0op::Reshapel0op::ViewCopy将一维逆索引还原为输入原始多维 shape 并写入inverseOut。路径二AI CPU 路径其他场景当不满足 AI Core 条件时直接调用l0op::UniqueWithCountsAndSorting完成去重与逆索引aclnn_unique.cpp该算子在index目录下有对应的 AI CPU 实现。从上述实现可以推断无论走哪条路径排序都是去重的前置步骤这也解释了为什么 950 平台 1D 输入会固定输出升序结果——AI Core 路径本身以升序排序为中间产物。同时测试用例 atk_aclnnUnique.json 中覆盖了 fp64、bf16、fp16、int8/int16/int32/int64 等多种 dtype 以及大量 4 维 shape含 131073、80900 等超大维度与sorted/returnInverse的四种组合与文档声明的数据类型支持范围一致。调用示例以下示例完整演示了两段式接口的调用流程与仓库 examples/test_aclnn_unique.cpp 内容一致可直接作为工程参考。编译与运行环境准备请参见编译与运行样例。#include cstdio #include vector #include acl/acl.h #include aclnnop/aclnn_unique.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 {8}; std::vectorint64_t valueShape {8}; std::vectorint64_t inverseShape {8}; void* selfDeviceAddr nullptr; void* valueDeviceAddr nullptr; void* inverseDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* valueOut nullptr; aclTensor* inverseOut nullptr; std::vectorfloat selfHostData {1, 2, 3, 4, 5, 6, 7, 8}; std::vectorfloat valueHostData {0, 0, 0, 0, 0, 0, 0, 0}; std::vectorint64_t inverseHostData {0, 0, 0, 0, 0, 0, 0, 0}; bool sorted true; bool returnInverse true; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建valueOut aclTensor ret CreateAclTensor(valueHostData, valueShape, valueDeviceAddr, aclDataType::ACL_FLOAT, valueOut); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建inverseOut aclTensor ret CreateAclTensor(inverseHostData, inverseShape, inverseDeviceAddr, aclDataType::ACL_INT64, inverseOut); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnUnique第一段接口 ret aclnnUniqueGetWorkspaceSize(self, sorted, returnInverse, valueOut, inverseOut, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnUniqueGetWorkspaceSize 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); } // 调用aclnnUnique第二段接口 ret aclnnUnique(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnUnique 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侧需要根据具体API的接口定义修改 auto size GetShapeSize(valueShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), valueDeviceAddr, size * sizeof(float), 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: %f\n, i, resultData[i]); } std::vectorint64_t resultData1(size, 0); ret aclrtMemcpy(resultData1.data(), resultData1.size() * sizeof(resultData1[0]), inverseDeviceAddr, size * sizeof(int64_t), 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(result1[%ld] is: %ld\n, i, resultData1[i]); } // 6. 释放aclTensor和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(valueOut); aclDestroyTensor(inverseOut); // 7.释放device资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(valueDeviceAddr); aclrtFree(inverseDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例要点拆解初始化aclInit→aclrtSetDevice→aclrtCreateStream为所有 ACL 程序的固定前置流程。Tensor 构造通过aclrtMalloc申请 Device 内存、aclrtMemcpy拷入数据、aclCreateTensor以 ND 格式创建aclTensor。注意inverseOut使用ACL_INT64类型——这与源码中inverseOut必须是DT_INT64的校验严格对应。两段式调用第一段aclnnUniqueGetWorkspaceSize返回workspaceSize与executor当workspaceSize 0时用aclrtMalloc申请 workspace再调用aclnnUnique执行。同步与取数aclrtSynchronizeStream等待任务完成随后将valueDeviceAddr唯一值与inverseDeviceAddr逆索引拷贝回 Host 侧打印。资源释放依次释放aclTensor、Device 内存、workspace、Stream并aclrtResetDeviceaclFinalize。测试与验证仓库为aclnnUnique提供了完整的 ST 与 UT 测试ST 测试atk_aclnnUnique.json 定义了 200 组用例输入 dtype 覆盖 fp64/bf16/fp16/int8/int16/int32/int64shape 覆盖 4 维各类形态含 131073、80900 等大元素数场景sorted与returnInverse的取值组合齐全executor_aclnnUnique.py 实现了 CPU 基准torch.unique与 NPU 执行后的结果对齐逻辑valueOut排序后对比inverseOut与全零张量对比。UT 测试tests/ut/op_host/test_aclnn_unique.cpp 覆盖 host 侧参数校验与接口调用。需要说明的是ST 用例中inverseOut的基准比对采用zero_()处理这是因为逆索引数值依赖于具体去重结果布局实际精度校验以官方测试框架为准应用侧请按语义自行验证逆索引正确性。延伸阅读两段式接口通用机制两段式接口说明返回码定义aclnn返回码样例编译与运行编译与运行样例相关概念基本概念、确定性计算兄弟接口aclnnUnique2支持返回出现次数aclnnUnique2 接口文档连续去重接口aclnnUniqueConsecutiveunique_consecutive 模块说明【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考