企业官网建设流程全解析

CANN ops-cv NMSWithMask 算子深度解析：NPU 上带掩码输出的非极大值抑制实现与图模式调用指南

【免费下载链接】ops-cv本项目是CANN提供的图像处理、目标检测相关的算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-cv

本文以 CANN ops-cv 仓库中的 NMSWithMask 算子为对象，系统讲解其在昇腾 NPU 上的功能定义、参数与约束、多核 Tiling 调度原理、掩码式 NMS 的两阶段计算流程，以及通过算子 IR 构图进行图模式调用的完整方案。读完本文，你将掌握该算子的全部对外接口细节，并能依据源码读懂其底层并行实现，直接指导目标检测后处理场景下的算子选型与调用。

算子概览：面向目标检测后处理的 NMSWithMask

NMSWithMask 是 CANN ops-cv（image/nms_with_mask）提供的一个图像处理领域算子，核心功能是对边界框执行非极大值抑制（Non-Maximum Suppression，NMS），一次性输出经过 NMS 过滤后的选中框（selected_boxes）、选中索引（selected_idx）以及掩码（selected_mask）。它广泛应用于目标检测的后处理阶段，用于剔除同一目标上冗余重叠的检测框，保留置信度最高的框。

从仓库中的算子原型定义（image/nms_with_mask/op_graph/nms_with_mask_proto.h）可以看到其语义描述：

Iteratively removes lower scoring boxes which have an IoU greater than iou_threshold with higher scoring box according to their intersection-over-union (IoU).

即：按照 IoU（交并比）判定标准，迭代地移除与高分框重叠过大的低分框。该算子不仅返回过滤后的框，还额外输出每个输入框的索引与有效性掩码，方便下游任务直接通过掩码进行 gather 等操作，避免二次计算。

产品支持情况

根据 image/nms_with_mask/README.md 中的产品支持矩阵，NMSWithMask 在不同昇腾产品上的支持情况如下：

从算子 host 侧注册代码（image/nms_with_mask/op_host/nms_with_mask_def.cpp）可以印证这一点：该算子的 AICore 配置仅通过this->AICore().AddConfig("ascend950", aicoreConfig)与this->AICore().AddConfig("ascend350", aicoreConfig)注册了ascend950（对应 Ascend 950 系列）与ascend350（对应 Atlas A2 系列）两类芯片平台，与 README 的支持矩阵一一对应。

功能说明与算法语义

NMSWithMask 对输入的一组带置信度得分的边界框执行标准 NMS：

1. 按得分从高到低选取候选框；
2. 计算候选框与其余框之间的 IoU（交并比）；
3. 将 IoU 超过阈值iou_threshold的低分框剔除；
4. 输出最终保留的框、对应输入索引与掩码。

与常规 NMS 算子不同，NMSWithMask 的输出形态保留了与输入等长的结构：selected_boxes与输入同 shape，selected_idx记录每个输出框在输入中的原始序号，selected_mask用 0/1 标记哪些位置的框是有效（被选中）的。这种"框 + 索引 + 掩码"三位一体的输出设计，使得后续的 ROI Align、Gather、掩码过滤等下游算子可以直接消费其结果，无需再次排序或比对。

参数说明

NMSWithMask 的完整参数定义如下表（继承自 image/nms_with_mask/README.md，并结合算子原型 image/nms_with_mask/op_graph/nms_with_mask_proto.h 的REG_OP定义补充细节）：

源码视角下的接口定义

原型注册文件 image/nms_with_mask/op_graph/nms_with_mask_proto.h 中的REG_OP(NMSWithMask)给出了权威定义：

REG_OP(NMSWithMask) .INPUT(box_scores, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16})) .OUTPUT(selected_boxes, TensorType({DT_FLOAT, DT_FLOAT16, DT_BF16})) .OUTPUT(selected_idx, TensorType({DT_INT32})) .OUTPUT(selected_mask, TensorType({DT_UINT8})) .ATTR(iou_threshold, Float, 0.5) .OP_END_FACTORY_REG(NMSWithMask)

对应的 host 侧 OpDef 注册（image/nms_with_mask/op_host/nms_with_mask_def.cpp）进一步明确：

• box_scores：REQUIRED输入，支持DT_FLOAT16 / DT_FLOAT / DT_BF16，格式为ND；
• selected_boxes：REQUIRED输出，数据类型与输入一致；
• selected_idx：REQUIRED输出，固定DT_INT32；
• selected_mask：REQUIRED输出，固定DT_UINT8；
• iou_threshold：REQUIRED属性，浮点类型，默认值 0.5。

从 Tiling 实现的CheckDtype（image/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp）可以看到更严格的运行时校验：box_scores与selected_boxes的数据类型必须一致，selected_idx仅接受DT_INT32，selected_mask仅接受DT_UINT8，否则 Tiling 阶段直接返回失败并打印错误日志。

iou_threshold 的取值范围

Tiling 代码在SetTilingData（image/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp）中对属性做了显式范围校验：

auto iouThrPtr = attrs->GetAttrPointer<float>(ATTR_IOU_THR_INDEX); iouThreshold_ = *iouThrPtr; if (iouThreshold_ < 0.0f || iouThreshold_ > 1.0f) { OP_LOGE(tilingContext_, "iou_threshold_ must be in [0, 1], got %f", iouThreshold_); return ge::GRAPH_FAILED; }

即iou_threshold的合法范围为[0, 1]，超出该区间会在编译/构图阶段被拒绝。阈值越大，允许保留的框越多；阈值越小，抑制越激进，保留的框越少。

约束说明

使用 NMSWithMask 时需要遵守以下约束（继承自 image/nms_with_mask/README.md，并结合 Tiling 源码校验逻辑展开）：

1. 输入最后一维必须为 5：box_scores的最后一维固定表示[y1, x1, y2, x2, score]。Tiling 的CheckInputShape（image/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp）会校验第二维必须等于 5，否则报错 "Input box_scores' second dim must be 5"。

2. 输入必须是二维：box_scores的维度数必须为 2（CheckInputShape中校验GetDimNum() != DIM_NUM_TWO即失败）。

3. 框数量范围：num_boxes必须大于 0 且小于 39936。源码中MAX_BLOCK_NUM = 39936（image/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp），校验条件为boxesNum_ <= 0 || boxesNum_ >= MAX_BLOCK_NUM即返回失败，因此实际合法区间是(0, 39936)，README 中"不超过 39936"的描述与之一致。

4. 输出 shape 与输入对齐：selected_boxes的 shape 必须为 (num_boxes, 5)，selected_idx、selected_mask必须为一维且长度等于 num_boxes。CheckOutputShape（image/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp）逐一校验了这三个输出的维度数与第一维大小，任何不一致都会导致失败。README 中"输出 batch 维度与输入保持一致"的约束在此得到源码级印证。

5. FLOAT16/BF16 场景的精度说明：在 FLOAT16 或 BF16 输入下，算子进行排序和 IoU 计算对比标杆（如 CPU 上的参考实现）时可能引入计算误差，属于低精度浮点类型固有的数值特性，工程上应结合精度对比结果决定是否需要对输入做类型提升。

实现原理：多核 Tiling 与两阶段 NMS 计算

NMSWithMask 在昇腾 NPU 上的实现采用"host 侧 Tiling 计算 + device 侧多核并行 Kernel"的标准 AscendC 算子架构，理解其实现有助于判断算子的性能特征与适用规模。

Host 侧 Tiling：分块、分核与 workspace 规划

Tiling 主流程（image/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp）依次执行CheckShape→CheckDtype→SetTilingData：

• 分块策略：groupSize_ = BASE_BLOCK_SIZE_FOR_MULTICORE = 256，即每 256 个框为一组；groupNum_ = ceil(boxesNum_ / 256)；两两分块组合的配对总数为blockNum_ = groupNum_ * (groupNum_ + 1) / 2（上三角配对，含对角线）。
• 分核策略：usedCoreNum_ = min(blockNum_, vectorCoreNum)，即实际使用核数取配对块数与 AIV 核数的较小值；随后计算headCoreNum_与blockPerHead_，将配对块尽量均匀地分配到每个核上；最终通过tilingContext_->SetBlockDim(usedCoreNum_)设定核数，并设置SCHEDULE_MODE = 1。
• workspace 规划：workspace 由两部分组成——系统库工作空间GetLibApiWorkSpaceSize()加上 IoU 掩码临时缓冲blockNum_ * groupSize_ * groupSize_ / 8字节（每个 block 的掩码按位存储，256×256 bit = 8192 字节）。这一缓冲区在 PreProcess 阶段用于暂存两两分块间的 IoU 判定结果。
• TilingData 结构（image/nms_with_mask/op_kernel/arch35/nms_with_mask_tiling_data.h）：包含boxesNum、usedCoreNum、groupSize、groupNum、blockNum、headCoreNum、blockPerHead、iouThreshold八个字段，全部由 host 侧计算完成后通过memcpy_s写入 raw tiling data，并以TILING_KEY_FOR_MULTICORE = 10000作为 TilingKey 传递给 device 侧。

此外，TilingPrepare4NMSWithMask（image/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp）在编译期收集硬件信息：AIV 核数（GetCoreNumAiv）、UB 内存大小（GetCoreMemSize）以及maxBoxesNum = 39936，供 Tiling 阶段决策使用。

Device 侧 Kernel：两阶段并行 NMS

Kernel 入口（image/nms_with_mask/op_kernel/nms_with_mask_apt.cpp）为nms_with_mask函数，它首先校验 workspace 与用户 workspace 指针，随后解析 TilingData，并在 TilingKey 为 10000 时进入多核实现NMSWithMaskRegbaseMultiProcess（image/nms_with_mask/op_kernel/arch35/nms_with_mask_regbase_multiprocess.h）。多核实现的核心思想是把 NMS 拆成两个阶段，中间通过全局 workspace 同步：

阶段一：PreProcess（并行计算 IoU 掩码矩阵）

• 将框集合按 256 一组切分为groupNum_个 group，两两组合形成上三角配对块（含对角线块），每个块由特定核负责；
• 每个核在其负责的块内：通过DataCopy将 reference（列方向）与 destination（行方向）两个 group 的框数据搬入 UB，计算各框面积（ComputeRefArea），然后逐行逐元素计算 IoU，判定intersection > iou_threshold * union时将该 bit 置 1（ComputeMaskVf+CalcIntersection）；
• 计算结果以**位图（bitmask）**形式写入 workspace 的临时缓冲区tempMaskGm_，每个 bit 代表一对框是否重叠超限。代码中仅计算上三角部分（only upper triangle part of the iou matrix will be calculated），避免一半的重复计算；
• 对角线块还会同步把 selected_boxes 原样拷贝、用CreateVecIndex生成连续索引，作为初始的选中结果输出。

阶段二：PostProcess（沿对角线传播抑制结果）

• 真正的 NMS 递推过程按行推进：每行先由 0 号核处理对角线块（ComputeNMSForDiagonal），得到该行参考框的最终选中状态；
• SyncAll()全局同步后，其余核按行重新分核处理该行的非对角线块（ComputeNMSForNormal）：读取 PreProcess 阶段生成的 IoU 位掩码，将"参考框已被抑制"的信息与 IoU 掩码做And运算，再通过Select指令把需要抑制的目标框掩码清零；
• 每处理完一行，将更新后的selected_mask写回全局内存（CopyOut），下一行的计算依赖上一行的结果，从而保证 NMS 的递推语义正确；
• 循环直至所有groupNum_行处理完毕，最终selected_mask即为全部框的有效性标记。

从该实现可以看出两个值得注意的设计点：

1. 掩码位图压缩：IoU 矩阵以 bit 形式存储（groupSize_ / BIT_PER_BYTE字节每行），配合寄存器级Reg::MaskReg位操作指令，显著减少了全局内存流量；
2. 面积与 IoU 的向量化计算：ComputeMaskVf每次迭代处理 2 个向量长度（VL）的目标框数据，并与参考框做向量广播比较，通过MaskDeInterleave将 32 位掩码交织为 16 位对齐格式写回，充分利用了向量单元的计算能力。

InferShape 与动态 shape 支持

InferShape 实现（image/nms_with_mask/op_host/nms_with_mask_infershape.cpp）逻辑非常直接：selected_boxes输出为二维 (N, 5)，selected_idx与selected_mask输出为一维 N，其中 N 直接取自输入box_scores的第 0 维。OpDef 中配置了DynamicShapeSupportFlag(true)与DynamicRankSupportFlag(true)（image/nms_with_mask/op_host/nms_with_mask_def.cpp），说明该算子支持动态 shape 场景，运行时再通过 Tiling 完成 shape 解析与分块决策。

调用说明：基于算子 IR 的图模式调用

NMSWithMask 的调用方式为图模式，即通过算子 IR（image/nms_with_mask/op_graph/nms_with_mask_proto.h）构图来调用，具体说明如下：

图模式调用的一般流程为：

1. 定义算子节点：使用REG_OP(NMSWithMask)生成的算子描述（即nms_with_mask_proto.h中的定义）在计算图中创建算子节点，配置输入box_scores、属性iou_threshold以及三个输出；
2. 连接上下游：将目标检测网络输出的原始框（形如 (num_boxes, 5) 的 Tensor）接到box_scores，将selected_idx/selected_mask接到后续的 gather、过滤或 ROI 算子；
3. 编译执行：图编译阶段由 host 侧 InferShape 完成输出 shape 推导，Tiling 完成分块分核决策；执行阶段由 device 侧 Kernel 完成多核并行 NMS 计算。

如果需要在其他算子定义中复用 NMSWithMask 的接口约束，可以参考nms_with_mask_def.cpp中 OpDef 的注册方式（输入输出数据类型、ND 格式、iou_threshold默认值 0.5 的声明），保持对外接口的一致性。

测试与验证

仓库为 NMSWithMask 提供了 host 侧单元测试，位于 image/nms_with_mask/tests/ut/op_host：

• test_nms_with_mask_infershape.cpp：使用InfershapeContextPara构造算子上下文，验证未知 rank（-2动态维度）场景下三个输出的 shape 推导结果（{{-2, 5}, {-2}, {-2}}），并检查执行返回GRAPH_SUCCESS。测试覆盖了 FLOAT 输入 +iou_threshold = 0.5的默认属性组合（test_nms_with_mask_infershape.cpp）；
• test_nms_with_mask_tiling.cpp：针对 Tiling 逻辑的单元测试，验证 shape/dtype 校验与 TilingData 计算；
• 测试工程文件 image/nms_with_mask/tests/ut/op_host/CMakeLists.txt 组织用例编译。

这些测试用例可以直接作为自定义算子接入时的行为参考：例如验证非法 shape（非 2 维、第二维非 5、框数超界）与非法 dtype 会被 Tiling 拒绝，验证iou_threshold越界（<0 或 >1）时的报错行为。

总结

NMSWithMask 是 CANN ops-cv 面向目标检测后处理场景提供的专用 NMS 算子，其核心价值在于：

• 接口完整：一次调用同时产出过滤框、输入索引与有效性掩码，天然适配后续 gather/过滤类算子；
• 多核并行：通过"分块配对 + 上三角 IoU 位图 + 对角线递推传播"的两阶段设计，在 Ascend 950 系列与 Atlas A2 系列产品上实现多核并行 NMS，且仅计算上三角掩码避免冗余；
• 动态 shape：host 侧 InferShape 与 Tiling 均支持动态 shape，框数量上限为 39936，满足主流检测模型单张图片的候选框规模需求。

使用前请重点核对三点：box_scores最后一维必须为 5 且整体为二维、num_boxes必须落在 (0, 39936) 区间、iou_threshold取值必须在 [0, 1]。在 FLOAT16/BF16 精度下，建议结合精度对比结果评估误差影响。如需深入了解实现细节，可继续阅读 nms_with_mask_regbase_multiprocess.h 中的向量化 IoU 计算与寄存器掩码操作，以及 nms_with_mask_tiling_arch35.cpp 中的分块分核决策。

【免费下载链接】ops-cv本项目是CANN提供的图像处理、目标检测相关的算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-cv