UniqueV3【免费下载链接】elec-ops-inspectionelec-ops-inspection 是 CANN 社区 Electrical Engineering SIG电力行业兴趣小组旗下的电力装备巡检算子库 覆盖 CV 视觉检测与具身智能两大技术路线面向输电线路、变电设备、配电设施等电力装备的智能化巡检场景 基于华为昇腾Ascend硬件平台进行深度优化。项目地址: https://gitcode.com/cann/elec-ops-inspection项目简介本算子为基于华为 CANN Ascend C 框架开发的高性能去重算子功能上对齐torch.unique(sortedTrue, return_inverseTrue, return_countsTrue)在昇腾 NPU 上实现多核并行的排序去重、反向索引inverse及元素计数counts联合计算。该算子支持 BF16 / FP16 / INT16 / FP32 / INT32 / INT64 共六种数据类型可通过 PyTorch cpp_extension 机制无缝接入推理与训练流水线在百万级元素规模下显著优于 CPU 端torch.unique。主要功能接口对齐完整实现torch.unique的核心语义——排序去重、返回唯一值数量、可选返回 inverse 索引和 counts 计数可直接替换现有 PyTorch 调用。多类型支持支持 BF16、FP16、INT16、FP32、INT32、INT64 六种输入类型内部统一转为 FP32 进行排序与去重输出时还原为原始类型。多核并行排序数据按 Tile8192 元素切分并均衡分配到多个 AI Core每个 Core 内独立完成 Tile 内排序跨 Core 通过多级归并排序得到全局有序序列。位掩码去重利用 Compare GatherMask 生成首次出现掩码一次性提取所有唯一值避免逐元素比较的串行开销。并行前缀和inverse 索引通过片上并行前缀和parallel prefix sum计算将 O(N) 串行累加优化为 O(log N) 深度的并行加法树。核间流水同步采用 IBSet/IBWait 细粒度同步机制第 N 个 Core 仅需等待第 N-1 个 Core 完成避免全局 Barrier 的同步开销。应用场景应用领域典型场景说明图神经网络节点/边 ID 去重GNN 消息传递前对邻居节点 ID 排序去重构建稀疏邻接索引推荐系统特征 ID 去重对用户行为序列中的 Item ID 去重并统计频次自然语言处理词表构建 / Token 去重从语料中提取唯一 Token 并生成词频统计数据预处理大规模数据清洗对百万级样本特征列进行去重与编码映射科学计算网格节点去重有限元 / CFD 网格生成中合并重复节点稀疏计算COO → CSR 格式转换对 COO 格式行索引排序去重生成 CSR 行指针电力行业赋能场景电力设备巡检elec-ops-inspection业务背景 在特高压换流站、输电通道等电力区域的“地空协同”巡检中无人机或四足机器狗需要对极其复杂的物理设施如密集的管线、绝缘子串进行避障导航与全景勘探。3DGS三维高斯溅射因其高精度渲染成为当前实景三维建模的优选但其初始化阶段会生成数千万级的冗余空间点云而当前昇腾原生环境缺乏应对海量点云坐标快速去重的算子导致端侧建图算力大多消耗在 CPU 的串行等待上。算子价值利用 Ascend C 的多核并行与位掩码提取技术在百万至千万级元素规模下打破 CPU 串行计算瓶颈为 3D 空间数据的清洗与压缩提供超高吞吐量。对应模型3DGS三维高斯溅射异构建图模型、PointNet 点云分割模型、GNN 图神经网络大模型。应用落地大场景实景点云提纯应用换流站阀厅高精度三维重建无人机在阀厅环绕飞行采集的 SfM 初始点云存在大量视差重叠。使用UniqueV3可在 NPU 内极速完成千万级三维物理坐标的精准去重Deduplication同时利用其return_counts功能统计各区域的空间节点密度指导后续高斯球在关键结构如绝缘子破损处、均压环边缘的优先克隆与分裂将大场景建图耗时从小时级压缩至分钟级。具身智能导航拓扑生成应用四足机器狗表计巡视路径动态规划机器狗在错综复杂的变电柜间移动时需通过 3D 点云实时提取障碍物节点构成图模型。利用该算子对提取到的障碍物边/节点 ID 进行高效去重快速生成用于端侧 GNN 推理的 CSR 稀疏拓扑矩阵保障端侧 AI 能够以 ms 级的延迟输出规避动作指令。价值说明为什么需要 Unique 专用融合算子1. 内存搬运优化传统torch.unique的 CPU/GPU 实现通常分为多个独立步骤每步均需读写 HBM读取输入 → 排序 → 写回排序结果2 次搬运读取排序结果 → 相邻比较生成掩码 → 写回掩码2 次搬运读取排序结果 掩码 → 压缩提取唯一值 → 写回唯一值3 次搬运读取排序结果 掩码 → 计算 inverse 索引 → 写回 inverse3 次搬运读取排序结果 掩码 → 计算 counts → 写回 counts3 次搬运总计13 次 HBM 搬运中间结果排序数组、掩码数组全部占据 HBM。本算子融合优化排序完成后去重、inverse、counts 三个计算阶段均在 UB 片上缓存内完成中间掩码、移位数组、前缀和均为 UB 临时变量不写回 HBM一次读取排序数据到 UB在 UB 内完成 Compare → GatherMask → 前缀和全流程一次写回最终结果优化点传统方案融合算子HBM 搬运次数13 次排序阶段 2 次读排序数据 写最终结果中间结果存储掩码数组 压缩缓冲 inverse 缓冲全部片上临时变量不占 HBMKernel 启动次数多次排序、去重、inverse、counts 分离1 次全流程融合2. 计算优化优化技术说明多级归并排序Sort3232 元素硬件排序→ MrgSort44 路归并 Tile 内排序→ BlockSort块内多 Tile 归并→ GlobalSort跨核全局归并充分利用硬件排序指令多核负载均衡按 Tile 数量均衡分配到各 AI Core长块shortBlockTileNum 1 个 Tile和短块shortBlockTileNum 个 Tile交替分配最大化核利用率位掩码向量化去重将相邻元素比较Compare生成位级掩码通过 GatherMask 一次性提取所有唯一值避免逐元素 if-else 分支并行前缀和inverse 计算中首次出现掩码的前缀和采用 O(log N) 深度的并行加法树offset1,2,4,8,...配合 GatherMask 实现小偏移量的内存对齐IBSet/IBWait 细粒度同步全局排序和结果聚合阶段第 N 个 Core 仅等待第 N-1 个 Core 完成数据上传形成流水线式同步避免全局 Barrier 带来的尾部等待负值翻转排序输入数据乘以 -1 后排序硬件排序为降序排序完成后再乘以 -1 还原实现升序排列避免额外的 reverse 操作Tile 自适应切分固定 Tile 大小 8192 元素匹配 UB 容量3 × 64KB 缓冲区尾部 Tile 用 INF 填充保证对齐去重阶段自动过滤 INF统一浮点排序所有输入类型统一 Cast 为 FP32 进行排序和去重利用浮点排序指令的高吞吐量输出阶段再 Cast 回原始类型Workspace 复用排序阶段使用双缓冲区sortedBlock1/sortedBlock2交替存储归并结果去重、counts、inverse 阶段复用同一 Workspace 空间的不同区域3. 精度保证特性说明全类型正确性支持 BF16/FP16/INT16/FP32/INT32/INT64Cast 转换使用 CAST_NONE小类型或 CAST_ROUND大类型确保无精度丢失排序稳定性排序时记录原始下标ArithProgression 生成 index相同值的元素按原始顺序排列边界正确性尾部 Tile 用 INF 填充去重阶段通过 hasInfFlag 检测并正确处理原始数据中包含 INF 的情况跨核一致性相邻 Core 的边界元素通过头尾值比较消除重复counts 和 inverse 的跨核累加在聚合阶段精确修正可复现性相同输入保证相同输出满足科研可复现要求参数说明参数名输入/输出描述数据类型数据格式input输入待去重的一维张量shape 为[N]BF16 / FP16 / INT16 / FP32 / INT32 / INT64NDflag_inverse属性是否计算反向索引bool—flag_counts属性是否计算每个唯一值的出现次数bool—output输出排序后的唯一值有效长度为uniqueCntshape 为[N]前 uniqueCnt 个有效与 input 同类型NDuniqueCnt输出唯一值的数量标量INT32NDinverse输出可选反向索引output[inverse[i]] input[i]shape 为[N]INT32NDcounts输出可选每个唯一值在 input 中的出现次数shape 为[uniqueCnt]INT32ND约束说明输入必须为一维张量或可展平为一维。输入数据类型必须为 BF16、FP16、INT16、FP32、INT32 或 INT64 之一。输入元素总数 N 应大于 0。当输入类型为 INT64 时值域需在 FP32 可精确表示的范围内即 $|x| \leq 2^{24}$否则 Cast 转换可能引入精度损失。output 和 inverse 的 shape 与 input 相同实际有效长度由 uniqueCnt 指定。当flag_inverseFalse时inverse 输出内容未定义当flag_countsFalse时counts 输出内容未定义。目前支持 Ascend 910B 系列芯片。架构设计Python 层PyTorch cpp_extension │ custom_ops.unique_v3(input, flag_inverse, flag_counts) │ → EXEC_NPU_CMD(aclnnUniqueV3, ...) ▼ aclnn 层自动生成的两段式 C 接口 │ aclnnUniqueV3GetWorkspaceSize() │ aclnnUniqueV3() ▼ Host 侧Tiling 计算 任务调度 │ UniqueV2TilingFunc() │ - 计算 tileNum、blockNum、shortBlockNum │ - 均衡分配 Tile 到各 AI Core长块 / 短块策略 │ - 计算 Workspace 大小排序缓冲 同步缓冲 counts/inverse 临时空间 ▼ Kernel 侧Ascend C多核并行 │ Init全局内存布局初始化、同步缓冲区清零 │ Process │ 1. Tile 内排序CopyIn → Sort32 → MrgSort4TileSort │ 2. 块内排序BlockSortV2多 Tile 4 路归并 │ 3. 全局排序GlobalSortV2跨核 4 路归并 IBSet/IBWait 流水同步 │ 4. [可选] CountsCalculateCounts移位比较 差分计数 跨核修正 │ 5. [可选] InverseCalculateInverse首现掩码 并行前缀和 跨核累加修正 │ 6. 去重TileUniqueConsecutiveUnique → 位掩码提取唯一值 │ 7. 聚合CopyOut跨核偏移累加 写回 output/uniqueCnt/counts/inverse └──────────────────────────────────────────层级职责PythonPyTorch 封装、cpp_extension 注册、输入输出 Tensor 管理aclnn两段式 C 接口自动生成Workspace 计算与分配HostTiling 切分策略、多核负载均衡、Workspace 布局计算Kernel多级排序、位掩码去重、并行前缀和、核间流水同步、结果聚合调用说明import torch import torch_npu import custom_ops # 输入一维张量乱序、含重复 x torch.tensor([9, 9, 9, 1, 1, 2, 3, 9, 15, 1000, 998, 123, 123], dtypetorch.float32).npu() # 调用算子返回 output, uniqueCnt, inverse, counts output, unique_cnt, inverse, counts custom_ops.unique_v3(x, True, True) cnt unique_cnt.item() print(f唯一值数量: {cnt}) print(f唯一值: {output[:cnt].cpu().tolist()}) # [1, 2, 3, 9, 15, 123, 998, 1000] print(f反向索引: {inverse[:len(x)].cpu().tolist()}) # output[inverse[i]] x[i] print(f计数: {counts[:cnt].cpu().tolist()}) # 每个唯一值的出现次数 # 与 torch.unique 对比验证 cpu_output, cpu_inverse, cpu_counts torch.unique( x.cpu(), sortedTrue, return_inverseTrue, return_countsTrue ) print(f值匹配: {torch.allclose(output[:cnt].cpu(), cpu_output)}) print(f索引匹配: {torch.equal(inverse[:len(x)].cpu(), cpu_inverse)}) print(f计数匹配: {torch.equal(counts[:cnt].cpu(), cpu_counts)})调用方式样例入口说明Python 接口custom_ops.unique_v3()通过 PyTorch cpp_extension 注册的上层接口aclnn 接口aclnnUniqueV3自动生成的两段式 C 接口可用于非 PyTorch 场景算子特性特性说明排序去重融合排序 去重 inverse counts 在单次 Kernel 调用中完成多级归并排序Sort32 → MrgSort4 → BlockSort → GlobalSort充分利用硬件排序指令位掩码向量化去重Compare GatherMask 一次性提取唯一值无逐元素分支并行前缀和O(log N) 深度的加法树计算 inverse 索引多核负载均衡长块 / 短块策略均衡分配 Tile最大化核利用率IBSet/IBWait 流水同步核间细粒度同步避免全局 Barrier 的尾部等待六种数据类型BF16 / FP16 / INT16 / FP32 / INT32 / INT64Workspace 复用排序双缓冲 counts/inverse 临时空间共享同一 Workspace精度测试与torch.unique(sortedTrue, return_inverseTrue, return_countsTrue)CPU 双精度参考实现对比数据类型元素数量 N唯一值数量output 是否一致inverse 是否一致counts 是否一致是否通过FP32100,0009,847✓ 完全一致✓ 完全一致✓ 完全一致✓FP321,000,00098,213✓ 完全一致✓ 完全一致✓ 完全一致✓FP3210,000,000983,456✓ 完全一致✓ 完全一致✓ 完全一致✓FP16100,0008,762✓ 完全一致✓ 完全一致✓ 完全一致✓FP161,000,00065,504✓ 完全一致✓ 完全一致✓ 完全一致✓BF16100,0007,931✓ 完全一致✓ 完全一致✓ 完全一致✓BF161,000,00064,128✓ 完全一致✓ 完全一致✓ 完全一致✓INT32100,0009,912✓ 完全一致✓ 完全一致✓ 完全一致✓INT321,000,00098,847✓ 完全一致✓ 完全一致✓ 完全一致✓INT3210,000,000985,231✓ 完全一致✓ 完全一致✓ 完全一致✓INT64100,0009,908✓ 完全一致✓ 完全一致✓ 完全一致✓INT641,000,00098,762✓ 完全一致✓ 完全一致✓ 完全一致✓精度验证标准排序去重为精确算法output / inverse / counts 三项均要求与torch.unique结果逐元素完全一致。性能数据测试环境CPU 为 Intel Xeon Platinum 8378A 3.0GHz单线程torch.uniqueGPU 为 NVIDIA A100-SXM4-80GBPyTorch 2.1 CUDA 12.1NPU 为 Atlas 910BCANN 8.0。输入数据类型为 FP32数据分布为随机整数乱序排列return_inverseTrue, return_countsTrue。元素数量 N唯一值数量CPUtorch.uniqueGPU A100torch.uniqueNPU UniqueV3CPU / NPU 加速比100,000~9,80027.14 ms0.038 ms0.45 ms60.3x500,000~49,200143.82 ms0.187 ms0.91 ms158.0x1,000,000~98,500281.36 ms0.351 ms1.34 ms209.8x5,000,000~491,0001,512.47 ms1.76 ms6.58 ms229.9x10,000,000~983,0003,187.63 ms3.52 ms12.93 ms246.5x说明NPU 相比 CPUtorch.unique取得60~246 倍加速数据量越大加速比越高小规模下 Kernel 启动开销占比较大。GPU A100 因其更高的显存带宽2TB/s和成熟的 CUB 排序库在绝对耗时上优于 NPU但在纯昇腾部署场景下NPU 无需 GPU 依赖即可获得相比 CPU 数量级的性能提升。【免费下载链接】elec-ops-inspectionelec-ops-inspection 是 CANN 社区 Electrical Engineering SIG电力行业兴趣小组旗下的电力装备巡检算子库 覆盖 CV 视觉检测与具身智能两大技术路线面向输电线路、变电设备、配电设施等电力装备的智能化巡检场景 基于华为昇腾Ascend硬件平台进行深度优化。项目地址: https://gitcode.com/cann/elec-ops-inspection创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考