CANN/elec-ops-inspection UniqueV3算子

UniqueV3

【免费下载链接】elec-ops-inspectionelec-ops-inspection 是 CANN 社区 Electrical Engineering SIG（电力行业兴趣小组）旗下的电力装备巡检算子库， 覆盖 CV 视觉检测与具身智能两大技术路线，面向输电线路、变电设备、配电设施等电力装备的智能化巡检场景， 基于华为昇腾（Ascend）硬件平台进行深度优化。项目地址: https://gitcode.com/cann/elec-ops-inspection

项目简介

本算子为基于华为 CANN Ascend C 框架开发的高性能去重算子，功能上对齐torch.unique(sorted=True, return_inverse=True, return_counts=True)，在昇腾 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 进行排序与去重，输出时还原为原始类型。
• 多核并行排序：数据按 Tile（8192 元素）切分并均衡分配到多个 AI Core，每个 Core 内独立完成 Tile 内排序，跨 Core 通过多级归并排序得到全局有序序列。
• 位掩码去重：利用 Compare + GatherMask 生成首次出现掩码，一次性提取所有唯一值，避免逐元素比较的串行开销。
• 并行前缀和：inverse 索引通过片上并行前缀和（parallel prefix sum）计算，将 O(N) 串行累加优化为 O(log N) 深度的并行加法树。
• 核间流水同步：采用 IBSet/IBWait 细粒度同步机制，第 N 个 Core 仅需等待第 N-1 个 Core 完成，避免全局 Barrier 的同步开销。

应用场景

电力行业赋能

场景：电力设备巡检（elec-ops-inspection）

业务背景： 在特高压换流站、输电通道等电力区域的“地空协同”巡检中，无人机或四足机器狗需要对极其复杂的物理设施（如密集的管线、绝缘子串）进行避障导航与全景勘探。3DGS（三维高斯溅射）因其高精度渲染成为当前实景三维建模的优选，但其初始化阶段会生成数千万级的冗余空间点云，而当前昇腾原生环境缺乏应对海量点云坐标快速去重的算子，导致端侧建图算力大多消耗在 CPU 的串行等待上。

• 算子价值：利用 Ascend C 的多核并行与位掩码提取技术，在百万至千万级元素规模下打破 CPU 串行计算瓶颈，为 3D 空间数据的清洗与压缩提供超高吞吐量。
• 对应模型：3DGS（三维高斯溅射）异构建图模型、PointNet++ 点云分割模型、GNN 图神经网络大模型。
• 应用落地：
  1. 大场景实景点云提纯（应用：换流站阀厅高精度三维重建）：无人机在阀厅环绕飞行采集的 SfM 初始点云存在大量视差重叠。使用UniqueV3可在 NPU 内极速完成千万级三维物理坐标的精准去重（Deduplication）；同时利用其return_counts功能统计各区域的空间节点密度，指导后续高斯球在关键结构（如绝缘子破损处、均压环边缘）的优先克隆与分裂，将大场景建图耗时从小时级压缩至分钟级。
  2. 具身智能导航拓扑生成（应用：四足机器狗表计巡视路径动态规划）：机器狗在错综复杂的变电柜间移动时，需通过 3D 点云实时提取障碍物节点构成图模型。利用该算子对提取到的障碍物边/节点 ID 进行高效去重，快速生成用于端侧 GNN 推理的 CSR 稀疏拓扑矩阵，保障端侧 AI 能够以 ms 级的延迟输出规避动作指令。

价值说明

为什么需要 Unique 专用融合算子？

1. 内存搬运优化

传统torch.unique的 CPU/GPU 实现通常分为多个独立步骤，每步均需读写 HBM：

1. 读取输入 → 排序 → 写回排序结果（2 次搬运）
2. 读取排序结果 → 相邻比较生成掩码 → 写回掩码（2 次搬运）
3. 读取排序结果 + 掩码 → 压缩提取唯一值 → 写回唯一值（3 次搬运）
4. 读取排序结果 + 掩码 → 计算 inverse 索引 → 写回 inverse（3 次搬运）
5. 读取排序结果 + 掩码 → 计算 counts → 写回 counts（3 次搬运）

总计：13+ 次 HBM 搬运，中间结果（排序数组、掩码数组）全部占据 HBM。

本算子融合优化：

• 排序完成后，去重、inverse、counts 三个计算阶段均在 UB 片上缓存内完成
• 中间掩码、移位数组、前缀和均为 UB 临时变量，不写回 HBM
• 一次读取排序数据到 UB，在 UB 内完成 Compare → GatherMask → 前缀和全流程，一次写回最终结果

2. 计算优化

3. 精度保证

参数说明

约束说明

• 输入必须为一维张量（或可展平为一维）。
• 输入数据类型必须为 BF16、FP16、INT16、FP32、INT32 或 INT64 之一。
• 输入元素总数 N 应大于 0。
• 当输入类型为 INT64 时，值域需在 FP32 可精确表示的范围内（即 $|x| \leq 2^{24}$），否则 Cast 转换可能引入精度损失。
• output 和 inverse 的 shape 与 input 相同，实际有效长度由 uniqueCnt 指定。
• 当flag_inverse=False时，inverse 输出内容未定义；当flag_counts=False时，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 → MrgSort4（TileSort） │ 2. 块内排序：BlockSortV2（多 Tile 4 路归并） │ 3. 全局排序：GlobalSortV2（跨核 4 路归并 + IBSet/IBWait 流水同步） │ 4. [可选] Counts：CalculateCounts（移位比较 + 差分计数 + 跨核修正） │ 5. [可选] Inverse：CalculateInverse（首现掩码 + 并行前缀和 + 跨核累加修正） │ 6. 去重：TileUnique（ConsecutiveUnique → 位掩码提取唯一值） │ 7. 聚合：CopyOut（跨核偏移累加 + 写回 output/uniqueCnt/counts/inverse） └──────────────────────────────────────────

调用说明

import torch import torch_npu import custom_ops # 输入：一维张量（乱序、含重复） x = torch.tensor([9, 9, 9, 1, 1, 2, 3, 9, 15, 1000, 998, 123, 123], dtype=torch.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(), sorted=True, return_inverse=True, return_counts=True ) 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)}")

算子特性

精度测试

与torch.unique(sorted=True, return_inverse=True, return_counts=True)CPU 双精度参考实现对比：

精度验证标准：排序去重为精确算法，output / inverse / counts 三项均要求与torch.unique结果逐元素完全一致。

性能数据

测试环境：CPU 为 Intel Xeon Platinum 8378A @ 3.0GHz（单线程torch.unique），GPU 为 NVIDIA A100-SXM4-80GB（PyTorch 2.1 + CUDA 12.1），NPU 为 Atlas 910B（CANN 8.0）。输入数据类型为 FP32，数据分布为随机整数乱序排列，return_inverse=True, return_counts=True。

说明：

• 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