[Issue]: `nonmaxsuppression` node produces incorrect outputs when `nonmaxsupression::use_dyn_output == False`

[mferencevic]

Problem description

When nonmaxsupression::use_dyn_output == False the node will return max_num_boxes boxes (equal to num_batches * total_boxes_in_batch_element).
If the nonmaxsupression algorithm selects less than max_num_boxes boxes, the rest of the output will be padded with zeros.
Output size is determined here:

AMDMIGraphX/src/include/migraphx/op/nonmaxsuppression.hpp

Line 187 in 0681c63

std::vector<std::size_t> out_lens = {max_num_boxes, 3};

Padding with 0 happens here:

AMDMIGraphX/src/include/migraphx/op/nonmaxsuppression.hpp

Line 324 in 0681c63

std::fill(output.begin(), output.end(), 0);

This means that the first box of the first image in the batch will be included in the output up to max_num_boxes times, while it should be included at most once.

This behavior breaks compatibility with the ONNX NonMaxSupression node (https://onnx.ai/onnx/operators/onnx__NonMaxSuppression.html) whose output only contains the selected boxes and isn't padded. This causes common post-NMS steps, like Top-K filtering based on the scores of the kept boxes, to be completely wrong.

When nonmaxsuppression::use_dyn_output == true you correctly shrink the output to only include the selected boxes:

AMDMIGraphX/src/include/migraphx/op/nonmaxsuppression.hpp

Line 403 in 0681c63

return result.reshape({output_shape.type(), {num_selected, 3}});

This doesn't happen when nonmaxsuppression::use_dyn_output == false.

There are a couple more problems with nonmaxsuppression::use_dyn_output:

1. use_dyn_output is never set when using the Python bindings for MIGraphX. nonmaxsuppression::use_dyn_output is set from onnx_parser::use_dyn_output:

   AMDMIGraphX/src/onnx/parse_nonmaxsuppression.cpp

   Line 42 in 0681c63

   op.from_value({{"use_dyn_output", parser.use_dyn_output}});

   Which is in turn set from onnx_options::use_dyn_output:

   AMDMIGraphX/src/onnx/onnx.cpp

   Line 73 in 0681c63

   parser.use_dyn_output = options.use_dyn_output;

   This option is never set when using the Python bindings:

   AMDMIGraphX/src/py/migraphx_py.cpp

   Line 675 in 0681c63

   migraphx::onnx_options options;

2. Even if use_dyn_output were set to true, nonmaxsuppression::compute_shape would fail unless the node's input shapes are also dynamic, which isn't the case when the entire model is compiled with static shapes. If we just change the onnx_options::use_dyn_output (while keeping static input shapes) we get the following error: parse: PARSE_EXPAND: dynamic input tensor with fixed dims input not supported. We often have to compile the whole graph with static shapes, as otherwise we'd get dynamic-shape-related compilation errors.

I'd suggest to completely remove the nonmaxsuppression::use_dyn_output == false code paths, as this option would always produce wrong outputs.
This would also require changing the node so that nonmaxsuppression::use_dyn_output == true works with static input shapes.

Steps to reproduce

This behavior is already encoded in the reference tests, in the nms_test test case:

AMDMIGraphX/test/ref/nonmaxsuppression.cpp

Line 272 in 0681c63

TEST_CASE(nms_test)

As you can see the expected outputs have been padded with 0, 0, 0 triplets which doesn't conform to the ONNX operator specification.

Environment

OS: Debian GNU/Linux 12 (bookworm)
GPU: AMD Radeon AI PRO R9700
ROCm version: 7.2.0
MIGraphX version: 2.16.0.dev+20250912-17-299-g0681c63ab

[adityas-amd]

@mferencevic thank you for the report, I have created a small test script repro.py:

import migraphx
import numpy as np
import onnx
from onnx import helper, TensorProto

print("Building NMS ONNX model...")

# Define inputs
boxes_in      = helper.make_tensor_value_info('boxes', TensorProto.FLOAT, [1, 6, 4])
scores_in     = helper.make_tensor_value_info('scores', TensorProto.FLOAT, [1, 1, 6])
max_boxes_in  = helper.make_tensor_value_info('max_output_boxes_per_class', TensorProto.INT64, [1])
iou_in        = helper.make_tensor_value_info('iou_threshold', TensorProto.FLOAT, [1])
score_in      = helper.make_tensor_value_info('score_threshold', TensorProto.FLOAT, [1])
output        = helper.make_tensor_value_info('selected_indices', TensorProto.INT64, None)

# NMS node
nms_node = helper.make_node(
    'NonMaxSuppression',
    inputs=[
        'boxes', 'scores',
        'max_output_boxes_per_class',
        'iou_threshold',
        'score_threshold'
    ],
    outputs=['selected_indices']
)

# Build and save model
graph = helper.make_graph(
    [nms_node], 'nms_graph',
    [boxes_in, scores_in, max_boxes_in, iou_in, score_in],
    [output]
)
model = helper.make_model(
    graph,
    opset_imports=[helper.make_opsetid('', 11)]
)
onnx.save(model, 'nms_test.onnx')
print("ONNX model saved!")

# Load with MIGraphX
print("\nLoading model with MIGraphX...")
prog = migraphx.parse_onnx('nms_test.onnx')
prog.compile(migraphx.get_target("gpu"))
print("Model compiled!")

# Input: 6 boxes, only 3 should survive NMS
boxes_data = np.array([[
    [0.0,  0.0,   1.0,  1.0],   # box 0 - score 0.90
    [0.0,  0.1,   1.0,  1.1],   # box 1 - score 0.75 (overlaps 0)
    [0.0, -0.1,   1.0,  0.9],   # box 2 - score 0.60 (overlaps 0)
    [0.0,  10.0,  1.0, 11.0],   # box 3 - score 0.95 (separate)
    [0.0,  10.1,  1.0, 11.1],   # box 4 - score 0.50 (overlaps 3)
    [0.0, 100.0,  1.0, 101.0]   # box 5 - score 0.30 (separate)
]], dtype=np.float32)

scores_data  = np.array([[[0.9, 0.75, 0.6, 0.95, 0.5, 0.3]]], dtype=np.float32)
max_boxes    = np.array([3], dtype=np.int64)
iou_thresh   = np.array([0.5], dtype=np.float32)
score_thresh = np.array([0.0], dtype=np.float32)

params = {
    'boxes':                      migraphx.argument(boxes_data),
    'scores':                     migraphx.argument(scores_data),
    'max_output_boxes_per_class': migraphx.argument(max_boxes),
    'iou_threshold':              migraphx.argument(iou_thresh),
    'score_threshold':            migraphx.argument(score_thresh)
}

print("\nRunning inference...")
result = prog.run(params)
output_data = np.array(result[0])

print("\n" + "="*50)
print("MIGraphX NMS Output:")
print(output_data)
print(f"Output shape: {output_data.shape}")
print("="*50)

# Check result
print("\nRESULT ANALYSIS:")
print(f"   Number of output boxes: {output_data.shape[0]}")
print(f"   Expected (ONNX spec):   3 boxes")

if output_data.shape[0] == 6:
    print("\n BUG CONFIRMED!")
    print("   Output is padded with fake [0,0,0] zeros")
    print("   This means box 0 appears multiple times!")
    print("   Downstream Top-K will be completely wrong!")
elif output_data.shape[0] == 3:
    print("\nOutput is CORRECT!")
    print("   Only selected boxes returned — ONNX spec compliant!")
else:
    print(f"\n  Unexpected output size: {output_data.shape[0]}")

print("\nExpected output per ONNX spec:")
print("[[0, 0, 3],   <- batch 0, class 0, box 3 (score 0.95)")
print(" [0, 0, 0],   <- batch 0, class 0, box 0 (score 0.90)")
print(" [0, 0, 5]]   <- batch 0, class 0, box 5 (score 0.30)")

Also I have drafted changes and testing it locally, here are preliminary results, will push a PR once its done.

current version:

Building NMS ONNX model...
ONNX model saved!

Loading model with MIGraphX...
 Model compiled!

Running inference...

==================================================
MIGraphX NMS Output:
[[0 0 3]
 [0 0 0]
 [0 0 5]
 [0 0 0]
 [0 0 0]
 [0 0 0]]
Output shape: (6, 3)
==================================================

RESULT ANALYSIS:
   Number of output boxes: 6
   Expected (ONNX spec):   3 boxes

BUG CONFIRMED!
   Output is padded with fake [0,0,0] zeros
   0 appears multiple times!

Expected output per ONNX spec:
[[0, 0, 3],   <- batch 0, class 0, box 3 (score 0.95)
 [0, 0, 0],   <- batch 0, class 0, box 0 (score 0.90)
 [0, 0, 5]]   <- batch 0, class 0, box 5 (score 0.30)

After changes

Building NMS ONNX model... 
ONNX model saved!

Loading model with MIGraphX... 
Model compiled!

Running inference...
MIGraphX NMS Output:
[[0 0 3] 
[0 0 0] 
[0 0 5]] 
Output shape: (3, 3)

RESULT ANALYSIS: 
 Number of output boxes: 3
 Expected (ONNX spec): 3 boxes

Output is CORRECT! Only selected boxes returned — ONNX spec compliant!

Expected output per ONNX spec: 
[[0, 0, 3], <- batch 0, class 0, box 3 (score 0.95)
 [0, 0, 0], <- batch 0, class 0, box 0 (score 0.90)
 [0, 0, 5]] <- batch 0, class 0, box 5 (score 0.30)

[CharlieL7]

I'll be working on making a GPU implementation of NMS and figuring out if we can continue to use the static output shape version (the zero-padded version) for the models we're interested in. If use_dyn_output is truly required for everything, then the only way for this to be performant is to have full dynamic shape support.