CANN ops-cv NMSWithMask 算子深度解析NPU 上带掩码输出的非极大值抑制实现与图模式调用指南【免费下载链接】ops-cv本项目是CANN提供的图像处理、目标检测相关的算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-cv本文以 CANN ops-cv 仓库中的 NMSWithMask 算子为对象系统讲解其在昇腾 NPU 上的功能定义、参数与约束、多核 Tiling 调度原理、掩码式 NMS 的两阶段计算流程以及通过算子 IR 构图进行图模式调用的完整方案。读完本文你将掌握该算子的全部对外接口细节并能依据源码读懂其底层并行实现直接指导目标检测后处理场景下的算子选型与调用。算子概览面向目标检测后处理的 NMSWithMaskNMSWithMask 是 CANN ops-cvimage/nms_with_mask提供的一个图像处理领域算子核心功能是对边界框执行非极大值抑制Non-Maximum SuppressionNMS一次性输出经过 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 在不同昇腾产品上的支持情况如下产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品×Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×从算子 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按得分从高到低选取候选框计算候选框与其余框之间的 IoU交并比将 IoU 超过阈值iou_threshold的低分框剔除输出最终保留的框、对应输入索引与掩码。与常规 NMS 算子不同NMSWithMask 的输出形态保留了与输入等长的结构selected_boxes与输入同 shapeselected_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定义补充细节参数名输入/输出/属性描述数据类型数据格式box_scores输入二维 Tensorshape 为 (num_boxes, 5)5 个元素依次表示 [y1, x1, y2, x2, score]num_boxes 不超过 39936FLOAT16、FLOAT、BF16NDiou_threshold属性浮点数用于判断候选框是否在 IoU 上重叠过多的阈值默认值为 0.5FLOAT-selected_boxes输出二维 Tensorshape 为 (num_boxes, 5)内容与输入 box_scores 同构过滤后的框及对应得分FLOAT16、FLOAT、BF16NDselected_idx输出一维 Tensorshape 为 (num_boxes)表示 0 到 num_boxes-1 的序列数被选中框在输入中的索引INT32NDselected_mask输出一维 Tensorshape 为 (num_boxes)表示目标框的掩码情况0/1 标记有效性UINT8ND源码视角下的接口定义原型注册文件 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_scoresREQUIRED输入支持DT_FLOAT16 / DT_FLOAT / DT_BF16格式为NDselected_boxesREQUIRED输出数据类型与输入一致selected_idxREQUIRED输出固定DT_INT32selected_maskREQUIRED输出固定DT_UINT8iou_thresholdREQUIRED属性浮点类型默认值 0.5。从 Tiling 实现的CheckDtypeimage/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp可以看到更严格的运行时校验box_scores与selected_boxes的数据类型必须一致selected_idx仅接受DT_INT32selected_mask仅接受DT_UINT8否则 Tiling 阶段直接返回失败并打印错误日志。iou_threshold 的取值范围Tiling 代码在SetTilingDataimage/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp中对属性做了显式范围校验auto iouThrPtr attrs-GetAttrPointerfloat(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 源码校验逻辑展开输入最后一维必须为 5box_scores的最后一维固定表示[y1, x1, y2, x2, score]。Tiling 的CheckInputShapeimage/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp会校验第二维必须等于 5否则报错 Input box_scores second dim must be 5。输入必须是二维box_scores的维度数必须为 2CheckInputShape中校验GetDimNum() ! DIM_NUM_TWO即失败。框数量范围num_boxes必须大于 0 且小于 39936。源码中MAX_BLOCK_NUM 39936image/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp校验条件为boxesNum_ 0 || boxesNum_ MAX_BLOCK_NUM即返回失败因此实际合法区间是(0, 39936)README 中不超过 39936的描述与之一致。输出 shape 与输入对齐selected_boxes的 shape 必须为 (num_boxes, 5)selected_idx、selected_mask必须为一维且长度等于 num_boxes。CheckOutputShapeimage/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp逐一校验了这三个输出的维度数与第一维大小任何不一致都会导致失败。README 中输出 batch 维度与输入保持一致的约束在此得到源码级印证。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 侧。此外TilingPrepare4NMSWithMaskimage/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp在编译期收集硬件信息AIV 核数GetCoreNumAiv、UB 内存大小GetCoreMemSize以及maxBoxesNum 39936供 Tiling 阶段决策使用。Device 侧 Kernel两阶段并行 NMSKernel 入口image/nms_with_mask/op_kernel/nms_with_mask_apt.cpp为nms_with_mask函数它首先校验 workspace 与用户 workspace 指针随后解析 TilingData并在 TilingKey 为 10000 时进入多核实现NMSWithMaskRegbaseMultiProcessimage/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 置 1ComputeMaskVfCalcIntersection计算结果以**位图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即为全部框的有效性标记。从该实现可以看出两个值得注意的设计点掩码位图压缩IoU 矩阵以 bit 形式存储groupSize_ / BIT_PER_BYTE字节每行配合寄存器级Reg::MaskReg位操作指令显著减少了全局内存流量面积与 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 的调用方式为图模式即通过算子 IRimage/nms_with_mask/op_graph/nms_with_mask_proto.h构图来调用具体说明如下调用方式样例代码说明图模式-通过算子IR构图方式调用 NMSWithMask 算子图模式调用的一般流程为定义算子节点使用REG_OP(NMSWithMask)生成的算子描述即nms_with_mask_proto.h中的定义在计算图中创建算子节点配置输入box_scores、属性iou_threshold以及三个输出连接上下游将目标检测网络输出的原始框形如 (num_boxes, 5) 的 Tensor接到box_scores将selected_idx/selected_mask接到后续的 gather、过滤或 ROI 算子编译执行图编译阶段由 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_hosttest_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.cpptest_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且仅计算上三角掩码避免冗余动态 shapehost 侧 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-cv NMSWithMask 算子深度解析：NPU 上带掩码输出的非极大值抑制实现与图模式调用指南
CANN ops-cv NMSWithMask 算子深度解析NPU 上带掩码输出的非极大值抑制实现与图模式调用指南【免费下载链接】ops-cv本项目是CANN提供的图像处理、目标检测相关的算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-cv本文以 CANN ops-cv 仓库中的 NMSWithMask 算子为对象系统讲解其在昇腾 NPU 上的功能定义、参数与约束、多核 Tiling 调度原理、掩码式 NMS 的两阶段计算流程以及通过算子 IR 构图进行图模式调用的完整方案。读完本文你将掌握该算子的全部对外接口细节并能依据源码读懂其底层并行实现直接指导目标检测后处理场景下的算子选型与调用。算子概览面向目标检测后处理的 NMSWithMaskNMSWithMask 是 CANN ops-cvimage/nms_with_mask提供的一个图像处理领域算子核心功能是对边界框执行非极大值抑制Non-Maximum SuppressionNMS一次性输出经过 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 在不同昇腾产品上的支持情况如下产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品×Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×从算子 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按得分从高到低选取候选框计算候选框与其余框之间的 IoU交并比将 IoU 超过阈值iou_threshold的低分框剔除输出最终保留的框、对应输入索引与掩码。与常规 NMS 算子不同NMSWithMask 的输出形态保留了与输入等长的结构selected_boxes与输入同 shapeselected_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定义补充细节参数名输入/输出/属性描述数据类型数据格式box_scores输入二维 Tensorshape 为 (num_boxes, 5)5 个元素依次表示 [y1, x1, y2, x2, score]num_boxes 不超过 39936FLOAT16、FLOAT、BF16NDiou_threshold属性浮点数用于判断候选框是否在 IoU 上重叠过多的阈值默认值为 0.5FLOAT-selected_boxes输出二维 Tensorshape 为 (num_boxes, 5)内容与输入 box_scores 同构过滤后的框及对应得分FLOAT16、FLOAT、BF16NDselected_idx输出一维 Tensorshape 为 (num_boxes)表示 0 到 num_boxes-1 的序列数被选中框在输入中的索引INT32NDselected_mask输出一维 Tensorshape 为 (num_boxes)表示目标框的掩码情况0/1 标记有效性UINT8ND源码视角下的接口定义原型注册文件 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_scoresREQUIRED输入支持DT_FLOAT16 / DT_FLOAT / DT_BF16格式为NDselected_boxesREQUIRED输出数据类型与输入一致selected_idxREQUIRED输出固定DT_INT32selected_maskREQUIRED输出固定DT_UINT8iou_thresholdREQUIRED属性浮点类型默认值 0.5。从 Tiling 实现的CheckDtypeimage/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp可以看到更严格的运行时校验box_scores与selected_boxes的数据类型必须一致selected_idx仅接受DT_INT32selected_mask仅接受DT_UINT8否则 Tiling 阶段直接返回失败并打印错误日志。iou_threshold 的取值范围Tiling 代码在SetTilingDataimage/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp中对属性做了显式范围校验auto iouThrPtr attrs-GetAttrPointerfloat(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 源码校验逻辑展开输入最后一维必须为 5box_scores的最后一维固定表示[y1, x1, y2, x2, score]。Tiling 的CheckInputShapeimage/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp会校验第二维必须等于 5否则报错 Input box_scores second dim must be 5。输入必须是二维box_scores的维度数必须为 2CheckInputShape中校验GetDimNum() ! DIM_NUM_TWO即失败。框数量范围num_boxes必须大于 0 且小于 39936。源码中MAX_BLOCK_NUM 39936image/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp校验条件为boxesNum_ 0 || boxesNum_ MAX_BLOCK_NUM即返回失败因此实际合法区间是(0, 39936)README 中不超过 39936的描述与之一致。输出 shape 与输入对齐selected_boxes的 shape 必须为 (num_boxes, 5)selected_idx、selected_mask必须为一维且长度等于 num_boxes。CheckOutputShapeimage/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp逐一校验了这三个输出的维度数与第一维大小任何不一致都会导致失败。README 中输出 batch 维度与输入保持一致的约束在此得到源码级印证。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 侧。此外TilingPrepare4NMSWithMaskimage/nms_with_mask/op_host/arch35/nms_with_mask_tiling_arch35.cpp在编译期收集硬件信息AIV 核数GetCoreNumAiv、UB 内存大小GetCoreMemSize以及maxBoxesNum 39936供 Tiling 阶段决策使用。Device 侧 Kernel两阶段并行 NMSKernel 入口image/nms_with_mask/op_kernel/nms_with_mask_apt.cpp为nms_with_mask函数它首先校验 workspace 与用户 workspace 指针随后解析 TilingData并在 TilingKey 为 10000 时进入多核实现NMSWithMaskRegbaseMultiProcessimage/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 置 1ComputeMaskVfCalcIntersection计算结果以**位图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即为全部框的有效性标记。从该实现可以看出两个值得注意的设计点掩码位图压缩IoU 矩阵以 bit 形式存储groupSize_ / BIT_PER_BYTE字节每行配合寄存器级Reg::MaskReg位操作指令显著减少了全局内存流量面积与 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 的调用方式为图模式即通过算子 IRimage/nms_with_mask/op_graph/nms_with_mask_proto.h构图来调用具体说明如下调用方式样例代码说明图模式-通过算子IR构图方式调用 NMSWithMask 算子图模式调用的一般流程为定义算子节点使用REG_OP(NMSWithMask)生成的算子描述即nms_with_mask_proto.h中的定义在计算图中创建算子节点配置输入box_scores、属性iou_threshold以及三个输出连接上下游将目标检测网络输出的原始框形如 (num_boxes, 5) 的 Tensor接到box_scores将selected_idx/selected_mask接到后续的 gather、过滤或 ROI 算子编译执行图编译阶段由 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_hosttest_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.cpptest_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且仅计算上三角掩码避免冗余动态 shapehost 侧 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考