Dynamic ref NMS (#1288)

Has NMS op output a dynamic shape (ONNX spec behavior)
Allows for dynamic input shape to NMS op

src/include/migraphx/onnx.hpp

@@ -35,17 +35,13 @@ struct onnx_options
 {
     /// Old way to set default fixed dimension size
     std::size_t default_dim_value = 0;
-    /*!
-     * Default dynamic dimension size (if both default_dim_value and default_dyn_dim_value
-     * set parser throws)
-     */
+    /// Default dynamic dimension size (if both default_dim_value and default_dyn_dim_value set
+    /// parser throws)
     shape::dynamic_dimension default_dyn_dim_value = {1, 1, 0};
     /// Explicitly specify the dims of an input
     std::unordered_map<std::string, std::vector<std::size_t>> map_input_dims = {};
-    /*!
-     * Explicitly specify dynamic dims of an input (if both map_input_dims and
-     * map_dyn_input_dims set parser throws)
-     */
+    /// Explicitly specify dynamic dims of an input (if both map_input_dims and map_dyn_input_dims
+    /// set parser throws)
     std::unordered_map<std::string, std::vector<shape::dynamic_dimension>> map_dyn_input_dims = {};
     /// Continue parsing onnx file if an unknown operator is found
     bool skip_unknown_operators = false;
@@ -53,6 +49,8 @@ struct onnx_options
     bool print_program_on_error = false;
     /// Max iter num for the loop operator
     int64_t max_loop_iterations = 10;
+    /// Use dynamic output for operators when available
+    bool use_dyn_output = false;
 };
 
 /// Create a program from an onnx file

src/include/migraphx/op/convert.hpp

@@ -45,7 +45,15 @@ struct convert : unary<convert>
     shape compute_shape(std::vector<shape> inputs) const
     {
         check_shapes{inputs, *this}.has(1);
-        return {target_type, inputs.at(0).lens(), inputs.at(0).strides()};
+        auto input = inputs.at(0);
+        if(input.dynamic())
+        {
+            return {target_type, input.dyn_dims()};
+        }
+        else
+        {
+            return {target_type, input.lens(), input.strides()};
+        }
     }
 
     std::string point_op() const

src/include/migraphx/op/nonmaxsuppression.hpp

@@ -45,11 +45,13 @@ namespace op {
 struct nonmaxsuppression
 {
     bool center_point_box = false;
+    bool use_dyn_output   = false;
 
     template <class Self, class F>
     static auto reflect(Self& self, F f)
     {
-        return pack(f(self.center_point_box, "center_point_box"));
+        return pack(f(self.center_point_box, "center_point_box"),
+                    f(self.use_dyn_output, "use_dyn_output"));
     }
 
     std::string name() const { return "nonmaxsuppression"; }
@@ -57,27 +59,81 @@ struct nonmaxsuppression
     shape compute_shape(std::vector<shape> inputs) const
     {
         // requires at least 2 inputs
-        check_shapes{{inputs.at(0), inputs.at(1)}, *this}.only_dims(3);
-        auto lens = inputs.front().lens();
+        check_shapes{{inputs.at(0), inputs.at(1)}, *this, true}.only_dims(3).same_ndims();
+        auto boxes_max_lens = inputs.at(0).max_lens();
+        // num batches * num boxes
+        const auto max_num_boxes = boxes_max_lens.at(0) * boxes_max_lens.at(1);
 
-        // check input shape
-        if(lens[1] != inputs.at(1).lens()[2])
+        auto fixed_shape_error_check = [&]() {
+            auto lens = inputs.front().lens();
+            if(lens[1] != inputs.at(1).lens()[2])
+            {
+                MIGRAPHX_THROW(
+                    "NonMaxSuppression: spatial dimension mismatch between boxes and scores input");
+            }
+            if(lens[0] != inputs.at(1).lens()[0])
+            {
+                MIGRAPHX_THROW(
+                    "NonMaxSuppression: number of batches mismatch between boxes and scores input");
+            }
+        };
+
+        if(use_dyn_output)
         {
-            MIGRAPHX_THROW(
-                "NonMaxSuppression: spatial dimension mismatch between boxes and scores input");
+            if(inputs.at(0).dynamic())
+            {
+                // both boxes and scores should be dynamic
+                // check dynamic dimensions are consistent
+                const auto boxes_dims  = inputs.at(0).dyn_dims();
+                const auto scores_dims = inputs.at(1).dyn_dims();
+                if(boxes_dims.at(1) != scores_dims.at(2))
+                {
+                    MIGRAPHX_THROW("NonMaxSuppression: dynamic spatial dimension mismatch between "
+                                   "boxes and scores input");
+                }
+                if(boxes_dims.at(0) != scores_dims.at(0))
+                {
+                    MIGRAPHX_THROW("NonMaxSuppression: dynamic number of batches mismatch between "
+                                   "boxes and scores input");
+                }
+            }
+            else if(inputs.at(1).dynamic())
+            {
+                // scores has dynamic shape, boxes fixed shape
+                // check that it is only a dynamic number of classes
+                const auto scores_dims = inputs.at(1).dyn_dims();
+                const auto boxes_lens  = inputs.at(0).lens();
+                if(not scores_dims.at(0).is_fixed() or scores_dims.at(0).max != boxes_lens.at(0))
+                {
+                    MIGRAPHX_THROW("NonMaxSuppression: scores dynamic num_classes; num_batches not "
+                                   "fixed or mismatched");
+                }
+                if(not scores_dims.at(2).is_fixed() or scores_dims.at(2).max != boxes_lens.at(1))
+                {
+                    MIGRAPHX_THROW("NonMaxSuppression: scores dynamic num_classes; "
+                                   "spatial_dimension not fixed or mismatches");
+                }
+            }
+            else
+            {
+                fixed_shape_error_check();
+            }
+            std::vector<shape::dynamic_dimension> out_lens = {};
+            out_lens.push_back({0, max_num_boxes, 0});
+            out_lens.push_back({3, 3, 0});
+            return {shape::int64_type, out_lens};
         }
-
-        // check batch sizes
-        if(lens[0] != inputs.at(1).lens()[0])
+        else
         {
-            MIGRAPHX_THROW(
-                "NonMaxSuppression: number of batches mismatch between boxes and scores input");
+            if(inputs.at(0).dynamic() or inputs.at(1).dynamic())
+            {
+                MIGRAPHX_THROW(
+                    "NonMaxSuppression: dynamic input shape with use_dyn_output set to false");
+            }
+            fixed_shape_error_check();
+            std::vector<std::size_t> out_lens = {max_num_boxes, 3};
+            return {shape::int64_type, out_lens};
         }
-
-        std::vector<int64_t> out_lens(2);
-        out_lens.at(0) = lens.at(1);
-        out_lens.at(1) = 3;
-        return {shape::int64_type, out_lens};
     }
 
     struct box
@@ -181,13 +237,13 @@ struct nonmaxsuppression
     }
 
     template <class Output, class Boxes, class Scores>
-    void compute_nms(Output output,
-                     Boxes boxes,
-                     Scores scores,
-                     const shape& output_shape,
-                     std::size_t max_output_boxes_per_class,
-                     double iou_threshold,
-                     double score_threshold) const
+    std::size_t compute_nms(Output output,
+                            Boxes boxes,
+                            Scores scores,
+                            const shape& max_output_shape,
+                            std::size_t max_output_boxes_per_class,
+                            double iou_threshold,
+                            double score_threshold) const
     {
         std::fill(output.begin(), output.end(), 0);
         const auto& lens       = scores.get_shape().lens();
@@ -197,7 +253,7 @@ struct nonmaxsuppression
         // boxes of a class with NMS applied [score, index]
         std::vector<std::pair<double, int64_t>> selected_boxes_inside_class;
         std::vector<int64_t> selected_indices;
-        selected_boxes_inside_class.reserve(output_shape.elements());
+        selected_boxes_inside_class.reserve(max_output_shape.elements());
         // iterate over batches and classes
         shape comp_s{shape::double_type, {num_batches, num_classes}};
         shape_for_each(comp_s, [&](auto idx) {
@@ -237,11 +293,14 @@ struct nonmaxsuppression
             }
         });
         std::copy(selected_indices.begin(), selected_indices.end(), output.begin());
+        return selected_indices.size() / 3;
     }
 
     argument compute(const shape& output_shape, std::vector<argument> args) const
     {
-        argument result{output_shape};
+        // make buffer of maximum size
+        shape max_output_shape = {output_shape.type(), output_shape.max_lens()};
+        argument result{max_output_shape};
 
         std::size_t max_output_boxes_per_class =
             (args.size() > 2) ? (args.at(2).at<std::size_t>()) : 0;
@@ -249,22 +308,29 @@ struct nonmaxsuppression
         {
             return result;
         }
-        double iou_threshold   = (args.size() > 3) ? (args.at(3).at<double>()) : 0.0f;
-        double score_threshold = (args.size() > 4) ? (args.at(4).at<double>()) : 0.0f;
+        double iou_threshold     = (args.size() > 3) ? (args.at(3).at<double>()) : 0.0f;
+        double score_threshold   = (args.size() > 4) ? (args.at(4).at<double>()) : 0.0f;
+        std::size_t num_selected = 0;
 
         result.visit([&](auto output) {
             visit_all(args[0], args[1])([&](auto boxes, auto scores) {
-                compute_nms(output,
-                            boxes,
-                            scores,
-                            output_shape,
-                            max_output_boxes_per_class,
-                            iou_threshold,
-                            score_threshold);
+                num_selected = compute_nms(output,
+                                           boxes,
+                                           scores,
+                                           max_output_shape,
+                                           max_output_boxes_per_class,
+                                           iou_threshold,
+                                           score_threshold);
             });
         });
-
-        return result;
+        if(use_dyn_output)
+        {
+            return result.reshape({output_shape.type(), {num_selected, 3}});
+        }
+        else
+        {
+            return result;
+        }
     }
 };
 

src/onnx/include/migraphx/onnx/onnx_parser.hpp

@@ -97,6 +97,7 @@ struct onnx_parser
     shape::dynamic_dimension default_dyn_dim_value = {1, 1, 0};
     std::unordered_map<std::string, std::vector<std::size_t>> map_input_dims;
     std::unordered_map<std::string, std::vector<shape::dynamic_dimension>> map_dyn_input_dims;
+    bool use_dyn_output         = false;
     bool skip_unknown_operators = false;
     int64_t max_loop_iterations = 10;
     int64_t opset_version       = 13;

src/onnx/onnx.cpp

@@ -60,8 +60,14 @@ program parse_onnx_from(const onnx_options& options, Ts&&... xs)
     {
         parser.default_dyn_dim_value = options.default_dyn_dim_value;
     }
+    if(not options.map_input_dims.empty() and not options.map_dyn_input_dims.empty())
+    {
+        MIGRAPHX_THROW("PARSE_ONNX_FROM: both map_input_dims and map_dyn_input_dims non-empty, only"
+                       "one should be used");
+    }
     parser.skip_unknown_operators = options.skip_unknown_operators;
     parser.max_loop_iterations    = options.max_loop_iterations;
+    parser.use_dyn_output         = options.use_dyn_output;
 
     if(options.print_program_on_error)
     {
@@ -80,6 +86,7 @@ program parse_onnx_from(const onnx_options& options, Ts&&... xs)
     {
         parser.parse_from(std::forward<Ts>(xs)...);
     }
+
     return std::move(parser.prog);
 }
 

src/onnx/onnx_parser.cpp

@@ -256,11 +256,6 @@ int64_t onnx_parser::get_opset_version(const onnx::ModelProto& model)
 
 void onnx_parser::parse_graph(module* mod, const onnx::GraphProto& graph)
 {
-    if(not map_input_dims.empty() and not map_dyn_input_dims.empty())
-    {
-        MIGRAPHX_THROW("PARSE_GRAPH: both map_input_dims and map_dyn_input_dims non-empty, only"
-                       "one should be used");
-    }
     std::unordered_map<std::string, instruction_ref> mod_insts;
     for(auto&& f : graph.initializer())
     {

src/onnx/parse_generic_op.cpp

@@ -58,7 +58,6 @@ struct parse_generic_op : op_parser<parse_generic_op>
                 {"Log", "log"},
                 {"LRN", "lrn"},
                 {"Neg", "neg"},
-                {"NonMaxSuppression", "nonmaxsuppression"},
                 {"Reciprocal", "recip"},
                 {"Relu", "relu"},
                 {"Round", "round"},
@@ -75,7 +74,7 @@ struct parse_generic_op : op_parser<parse_generic_op>
 
     bool needs_contiguous(const std::string& op_name) const
     {
-        return contains({"flatten", "gather", "nonmaxsuppression", "scatter"}, op_name);
+        return contains({"flatten", "gather", "scatter"}, op_name);
     }
 
     instruction_ref parse(const op_desc& opd,

src/onnx/parse_nonmaxsuppression.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <migraphx/onnx/op_parser.hpp>
+#include <migraphx/ranges.hpp>
+#include <migraphx/make_op.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace onnx {
+
+struct parse_nonmaxsuppression : op_parser<parse_nonmaxsuppression>
+{
+    std::vector<op_desc> operators() const { return {{"NonMaxSuppression", "nonmaxsuppression"}}; }
+
+    instruction_ref parse(const op_desc& opd,
+                          const onnx_parser& parser,
+                          const onnx_parser::node_info& info,
+                          const std::vector<instruction_ref>& args) const
+    {
+        auto op = parser.load(opd.op_name, info);
+        op.from_value({{"use_dyn_output", parser.use_dyn_output}});
+        return info.add_instruction(op, args);
+    }
+};
+
+} // namespace onnx
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

test/onnx/gen_onnx.py

@@ -3589,7 +3589,7 @@ def nms_test():
     st = helper.make_tensor_value_info('score_threshold', TensorProto.FLOAT,
                                        [1])
     out = helper.make_tensor_value_info('selected_indices', TensorProto.INT64,
-                                        [6, 3])
+                                        [None, 3])
 
     node = onnx.helper.make_node('NonMaxSuppression',
                                  inputs=[
@@ -3603,6 +3603,108 @@ def nms_test():
     return ([node], [b, s, mo, iou, st], [out])
 
 
+@onnx_test
+def nms_use_dyn_output_false_test():
+    b = helper.make_tensor_value_info('boxes', TensorProto.FLOAT, [1, 6, 4])
+    s = helper.make_tensor_value_info('scores', TensorProto.FLOAT, [1, 1, 6])
+    mo = helper.make_tensor_value_info('max_output_boxes_per_class',
+                                       TensorProto.INT64, [1])
+    iou = helper.make_tensor_value_info('iou_threshold', TensorProto.FLOAT,
+                                        [1])
+    st = helper.make_tensor_value_info('score_threshold', TensorProto.FLOAT,
+                                       [1])
+    out = helper.make_tensor_value_info('selected_indices', TensorProto.INT64,
+                                        [None, 3])
+
+    node = onnx.helper.make_node('NonMaxSuppression',
+                                 inputs=[
+                                     'boxes', 'scores',
+                                     'max_output_boxes_per_class',
+                                     'iou_threshold', 'score_threshold'
+                                 ],
+                                 outputs=['selected_indices'],
+                                 use_dyn_output=0)
+
+    return ([node], [b, s, mo, iou, st], [out])
+
+
+@onnx_test
+def nms_dynamic_batch_test():
+    b = helper.make_tensor_value_info('boxes', TensorProto.FLOAT, [None, 6, 4])
+    s = helper.make_tensor_value_info('scores', TensorProto.FLOAT,
+                                      [None, 1, 6])
+    mo = helper.make_tensor_value_info('max_output_boxes_per_class',
+                                       TensorProto.INT64, [1])
+    iou = helper.make_tensor_value_info('iou_threshold', TensorProto.FLOAT,
+                                        [1])
+    st = helper.make_tensor_value_info('score_threshold', TensorProto.FLOAT,
+                                       [1])
+    out = helper.make_tensor_value_info('selected_indices', TensorProto.INT64,
+                                        [None, 3])
+
+    node = onnx.helper.make_node('NonMaxSuppression',
+                                 inputs=[
+                                     'boxes', 'scores',
+                                     'max_output_boxes_per_class',
+                                     'iou_threshold', 'score_threshold'
+                                 ],
+                                 outputs=['selected_indices'],
+                                 center_point_box=1,
+                                 use_dyn_output=1)
+
+    return ([node], [b, s, mo, iou, st], [out])
+
+
+@onnx_test
+def nms_dynamic_boxes_test():
+    b = helper.make_tensor_value_info('boxes', TensorProto.FLOAT, [1, None, 4])
+    s = helper.make_tensor_value_info('scores', TensorProto.FLOAT,
+                                      [1, 1, None])
+    mo = helper.make_tensor_value_info('max_output_boxes_per_class',
+                                       TensorProto.INT64, [1])
+    iou = helper.make_tensor_value_info('iou_threshold', TensorProto.FLOAT,
+                                        [1])
+    st = helper.make_tensor_value_info('score_threshold', TensorProto.FLOAT,
+                                       [1])
+    out = helper.make_tensor_value_info('selected_indices', TensorProto.INT64,
+                                        [None, 3])
+
+    node = onnx.helper.make_node('NonMaxSuppression',
+                                 inputs=[
+                                     'boxes', 'scores',
+                                     'max_output_boxes_per_class',
+                                     'iou_threshold', 'score_threshold'
+                                 ],
+                                 outputs=['selected_indices'])
+
+    return ([node], [b, s, mo, iou, st], [out])
+
+
+@onnx_test
+def nms_dynamic_classes_test():
+    b = helper.make_tensor_value_info('boxes', TensorProto.FLOAT, [1, 6, 4])
+    s = helper.make_tensor_value_info('scores', TensorProto.FLOAT,
+                                      [1, None, 6])
+    mo = helper.make_tensor_value_info('max_output_boxes_per_class',
+                                       TensorProto.INT64, [1])
+    iou = helper.make_tensor_value_info('iou_threshold', TensorProto.FLOAT,
+                                        [1])
+    st = helper.make_tensor_value_info('score_threshold', TensorProto.FLOAT,
+                                       [1])
+    out = helper.make_tensor_value_info('selected_indices', TensorProto.INT64,
+                                        [None, 3])
+
+    node = onnx.helper.make_node('NonMaxSuppression',
+                                 inputs=[
+                                     'boxes', 'scores',
+                                     'max_output_boxes_per_class',
+                                     'iou_threshold', 'score_threshold'
+                                 ],
+                                 outputs=['selected_indices'])
+
+    return ([node], [b, s, mo, iou, st], [out])
+
+
 @onnx_test
 def not_test():
     x = helper.make_tensor_value_info('0', TensorProto.INT32, [4])

test/onnx/onnx_test.cpp

@@ -3378,13 +3378,127 @@ TEST_CASE(nms_test)
     auto st = mm->add_parameter("score_threshold", sst);
 
     auto ret = mm->add_instruction(
-        migraphx::make_op("nonmaxsuppression", {{"center_point_box", 1}}), b, s, mo, iou, st);
+        migraphx::make_op("nonmaxsuppression", {{"center_point_box", true}}), b, s, mo, iou, st);
     mm->add_return({ret});
 
     auto prog = migraphx::parse_onnx("nms_test.onnx");
     EXPECT(p == prog);
 }
 
+TEST_CASE(nms_dynamic_batch_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape sb{migraphx::shape::float_type, {{1, 10, 0}, {6, 6, 0}, {4, 4, 0}}};
+    auto b = mm->add_parameter("boxes", sb);
+    migraphx::shape ss{migraphx::shape::float_type, {{1, 10, 0}, {1, 1, 0}, {6, 6, 0}}};
+    auto s = mm->add_parameter("scores", ss);
+    migraphx::shape smo{migraphx::shape::int64_type, {1}};
+    auto mo = mm->add_parameter("max_output_boxes_per_class", smo);
+    migraphx::shape siou{migraphx::shape::float_type, {1}};
+    auto iou = mm->add_parameter("iou_threshold", siou);
+    migraphx::shape sst{migraphx::shape::float_type, {1}};
+    auto st  = mm->add_parameter("score_threshold", sst);
+    auto ret = mm->add_instruction(
+        migraphx::make_op("nonmaxsuppression",
+                          {{"center_point_box", true}, {"use_dyn_output", true}}),
+        b,
+        s,
+        mo,
+        iou,
+        st);
+    mm->add_return({ret});
+
+    migraphx::onnx_options options;
+    options.default_dyn_dim_value = {1, 10, 0};
+    options.use_dyn_output        = true;
+
+    auto prog = migraphx::parse_onnx("nms_dynamic_batch_test.onnx", options);
+    EXPECT(p == prog);
+}
+
+TEST_CASE(nms_dynamic_boxes_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape sb{migraphx::shape::float_type, {{1, 1, 0}, {6, 20, 0}, {4, 4, 0}}};
+    auto b = mm->add_parameter("boxes", sb);
+    migraphx::shape ss{migraphx::shape::float_type, {{1, 1, 0}, {1, 1, 0}, {6, 20, 0}}};
+    auto s = mm->add_parameter("scores", ss);
+    migraphx::shape smo{migraphx::shape::int64_type, {1}};
+    auto mo = mm->add_parameter("max_output_boxes_per_class", smo);
+    migraphx::shape siou{migraphx::shape::float_type, {1}};
+    auto iou = mm->add_parameter("iou_threshold", siou);
+    migraphx::shape sst{migraphx::shape::float_type, {1}};
+    auto st  = mm->add_parameter("score_threshold", sst);
+    auto ret = mm->add_instruction(
+        migraphx::make_op("nonmaxsuppression", {{"use_dyn_output", true}}), b, s, mo, iou, st);
+    mm->add_return({ret});
+
+    migraphx::onnx_options options;
+    options.default_dyn_dim_value = {6, 20, 0};
+    options.use_dyn_output        = true;
+
+    auto prog = migraphx::parse_onnx("nms_dynamic_boxes_test.onnx", options);
+    EXPECT(p == prog);
+}
+
+TEST_CASE(nms_dynamic_classes_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape sb{migraphx::shape::float_type, {1, 6, 4}};
+    auto b = mm->add_parameter("boxes", sb);
+    migraphx::shape ss{migraphx::shape::float_type, {{1, 1, 0}, {1, 10, 0}, {6, 6, 0}}};
+    auto s = mm->add_parameter("scores", ss);
+    migraphx::shape smo{migraphx::shape::int64_type, {1}};
+    auto mo = mm->add_parameter("max_output_boxes_per_class", smo);
+    migraphx::shape siou{migraphx::shape::float_type, {1}};
+    auto iou = mm->add_parameter("iou_threshold", siou);
+    migraphx::shape sst{migraphx::shape::float_type, {1}};
+    auto st  = mm->add_parameter("score_threshold", sst);
+    auto ret = mm->add_instruction(
+        migraphx::make_op("nonmaxsuppression", {{"use_dyn_output", true}}), b, s, mo, iou, st);
+    mm->add_return({ret});
+
+    migraphx::onnx_options options;
+    options.default_dyn_dim_value = {1, 10, 0};
+    options.use_dyn_output        = true;
+
+    auto prog = migraphx::parse_onnx("nms_dynamic_classes_test.onnx", options);
+    EXPECT(p == prog);
+}
+
+TEST_CASE(nms_overwrite_use_dyn_output_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape sb{migraphx::shape::float_type, {1, 6, 4}};
+    auto b = mm->add_parameter("boxes", sb);
+
+    migraphx::shape ss{migraphx::shape::float_type, {1, 1, 6}};
+    auto s = mm->add_parameter("scores", ss);
+
+    migraphx::shape smo{migraphx::shape::int64_type, {1}};
+    auto mo = mm->add_parameter("max_output_boxes_per_class", smo);
+
+    migraphx::shape siou{migraphx::shape::float_type, {1}};
+    auto iou = mm->add_parameter("iou_threshold", siou);
+
+    migraphx::shape sst{migraphx::shape::float_type, {1}};
+    auto st = mm->add_parameter("score_threshold", sst);
+
+    auto ret = mm->add_instruction(
+        migraphx::make_op("nonmaxsuppression", {{"use_dyn_output", true}}), b, s, mo, iou, st);
+    mm->add_return({ret});
+
+    migraphx::onnx_options options;
+    options.use_dyn_output = true;
+
+    auto prog = migraphx::parse_onnx("nms_use_dyn_output_false_test.onnx", options);
+    EXPECT(p == prog);
+}
+
 TEST_CASE(nonzero_dynamic_test)
 {
     migraphx::program p;

test/op_shape_test.cpp

@@ -1135,6 +1135,149 @@ TEST_CASE(multinomial)
     throws_shape(migraphx::make_op("multinomial", {{"dtype", dtype}}), s, s);
 }
 
+TEST_CASE(nms_shape)
+{
+    // use_dyn_output == false
+    migraphx::shape boxes_s{migraphx::shape::float_type, {1, 6, 4}};
+    migraphx::shape scores_s{migraphx::shape::float_type, {1, 1, 6}};
+    migraphx::shape max_out_s{migraphx::shape::int64_type, {1}};
+    migraphx::shape iou_thres_s{migraphx::shape::float_type, {1}};
+    migraphx::shape score_thres_s{migraphx::shape::float_type, {1}};
+    migraphx::shape output_s{migraphx::shape::int64_type, {6, 3}};
+    expect_shape(output_s,
+                 migraphx::make_op("nonmaxsuppression",
+                                   {{"center_point_box", true}, {"use_dyn_output", false}}),
+                 boxes_s,
+                 scores_s,
+                 max_out_s,
+                 iou_thres_s,
+                 score_thres_s);
+
+    // use_dyn_output == true
+    output_s = {migraphx::shape::int64_type, {{0, 6, 0}, {3, 3, 0}}};
+    expect_shape(output_s,
+                 migraphx::make_op("nonmaxsuppression",
+                                   {{"center_point_box", true}, {"use_dyn_output", true}}),
+                 boxes_s,
+                 scores_s,
+                 max_out_s,
+                 iou_thres_s,
+                 score_thres_s);
+
+    // dynamic batches
+    boxes_s  = {migraphx::shape::float_type, {{1, 3, 0}, {6, 6, 0}, {4, 4, 0}}};
+    scores_s = {migraphx::shape::float_type, {{1, 3, 0}, {1, 1, 0}, {6, 6, 0}}};
+    output_s = {migraphx::shape::int64_type, {{0, 18, 0}, {3, 3, 0}}};
+    expect_shape(output_s,
+                 migraphx::make_op("nonmaxsuppression",
+                                   {{"center_point_box", true}, {"use_dyn_output", true}}),
+                 boxes_s,
+                 scores_s,
+                 max_out_s,
+                 iou_thres_s,
+                 score_thres_s);
+
+    // dynamic num boxes
+    boxes_s  = {migraphx::shape::float_type, {{1, 1, 0}, {6, 20, 0}, {4, 4, 0}}};
+    scores_s = {migraphx::shape::float_type, {{1, 1, 0}, {1, 1, 0}, {6, 20, 0}}};
+    output_s = {migraphx::shape::int64_type, {{0, 20, 0}, {3, 3, 0}}};
+    expect_shape(output_s,
+                 migraphx::make_op("nonmaxsuppression",
+                                   {{"center_point_box", true}, {"use_dyn_output", true}}),
+                 boxes_s,
+                 scores_s,
+                 max_out_s,
+                 iou_thres_s,
+                 score_thres_s);
+
+    // use_dyn_output false with dynamic input shape
+    throws_shape(migraphx::make_op("nonmaxsuppression",
+                                   {{"center_point_box", true}, {"use_dyn_output", false}}),
+                 boxes_s,
+                 scores_s,
+                 max_out_s,
+                 iou_thres_s,
+                 score_thres_s);
+
+    // dynamic classes
+    boxes_s  = {migraphx::shape::float_type, {{1, 1, 0}, {6, 6, 0}, {4, 4, 0}}};
+    scores_s = {migraphx::shape::float_type, {{1, 1, 0}, {1, 3, 0}, {6, 6, 0}}};
+    output_s = {migraphx::shape::int64_type, {{0, 6, 0}, {3, 3, 0}}};
+    expect_shape(output_s,
+                 migraphx::make_op("nonmaxsuppression",
+                                   {{"center_point_box", true}, {"use_dyn_output", true}}),
+                 boxes_s,
+                 scores_s,
+                 max_out_s,
+                 iou_thres_s,
+                 score_thres_s);
+
+    // fixed mismatch batches
+    boxes_s  = {migraphx::shape::float_type, {2, 6, 4}};
+    scores_s = {migraphx::shape::float_type, {1, 1, 6}};
+    throws_shape(migraphx::make_op("nonmaxsuppression",
+                                   {{"center_point_box", true}, {"use_dyn_output", true}}),
+                 boxes_s,
+                 scores_s,
+                 max_out_s,
+                 iou_thres_s,
+                 score_thres_s);
+
+    // fixed mismatch num boxes
+    boxes_s  = {migraphx::shape::float_type, {1, 6, 4}};
+    scores_s = {migraphx::shape::float_type, {1, 1, 4}};
+    throws_shape(migraphx::make_op("nonmaxsuppression",
+                                   {{"center_point_box", true}, {"use_dyn_output", true}}),
+                 boxes_s,
+                 scores_s,
+                 max_out_s,
+                 iou_thres_s,
+                 score_thres_s);
+
+    // dynamic mismatch batches
+    boxes_s  = {migraphx::shape::float_type, {{1, 4, 0}, {6, 6, 0}, {4, 4, 0}}};
+    scores_s = {migraphx::shape::float_type, {{2, 8, 0}, {1, 1, 0}, {6, 6, 0}}};
+    throws_shape(migraphx::make_op("nonmaxsuppression",
+                                   {{"center_point_box", true}, {"use_dyn_output", true}}),
+                 boxes_s,
+                 scores_s,
+                 max_out_s,
+                 iou_thres_s,
+                 score_thres_s);
+
+    // dynamic mismatch num boxes
+    boxes_s  = {migraphx::shape::float_type, {{1, 1, 0}, {6, 8, 0}, {4, 4, 0}}};
+    scores_s = {migraphx::shape::float_type, {{1, 1, 0}, {1, 1, 0}, {3, 9, 0}}};
+    throws_shape(migraphx::make_op("nonmaxsuppression",
+                                   {{"center_point_box", true}, {"use_dyn_output", true}}),
+                 boxes_s,
+                 scores_s,
+                 max_out_s,
+                 iou_thres_s,
+                 score_thres_s);
+
+    // dynamic number of classes, fixed boxes_s, mismatch batches
+    boxes_s  = {migraphx::shape::float_type, {1, 6, 4}};
+    scores_s = {migraphx::shape::float_type, {{1, 3, 0}, {1, 3, 0}, {6, 6, 0}}};
+    throws_shape(migraphx::make_op("nonmaxsuppression",
+                                   {{"center_point_box", true}, {"use_dyn_output", true}}),
+                 boxes_s,
+                 scores_s,
+                 max_out_s,
+                 iou_thres_s,
+                 score_thres_s);
+    // dynamic number of classes, fixed boxes_s, mismatch num boxes
+    boxes_s  = {migraphx::shape::float_type, {1, 6, 4}};
+    scores_s = {migraphx::shape::float_type, {{1, 1, 0}, {1, 3, 0}, {4, 8, 0}}};
+    throws_shape(migraphx::make_op("nonmaxsuppression",
+                                   {{"center_point_box", true}, {"use_dyn_output", true}}),
+                 boxes_s,
+                 scores_s,
+                 max_out_s,
+                 iou_thres_s,
+                 score_thres_s);
+}
+
 TEST_CASE(pooling_shape)
 {
     migraphx::shape output{migraphx::shape::float_type, {4, 3, 1, 1}};

test/ref_ops_test.cpp

@@ -3624,6 +3624,174 @@ TEST_CASE(neg_test)
     EXPECT(migraphx::verify_range(result_vector, gold));
 }
 
+TEST_CASE(nms_dynamic_out_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape boxes_s{migraphx::shape::float_type, {1, 6, 4}};
+    std::vector<float> boxes_vec = {0.5, 0.5,  1.0, 1.0, 0.5, 0.6,  1.0, 1.0, 0.5, 0.4,   1.0, 1.0,
+                                    0.5, 10.5, 1.0, 1.0, 0.5, 10.6, 1.0, 1.0, 0.5, 100.5, 1.0, 1.0};
+
+    migraphx::shape scores_s{migraphx::shape::float_type, {1, 1, 6}};
+    std::vector<float> scores_vec = {0.9, 0.75, 0.6, 0.95, 0.5, 0.3};
+
+    auto boxes_l         = mm->add_literal(migraphx::literal(boxes_s, boxes_vec));
+    auto scores_l        = mm->add_literal(migraphx::literal(scores_s, scores_vec));
+    auto max_out_l       = mm->add_literal(int64_t{4});
+    auto iou_threshold   = mm->add_literal(0.5f);
+    auto score_threshold = mm->add_literal(0.0f);
+
+    auto r = mm->add_instruction(
+        migraphx::make_op("nonmaxsuppression",
+                          {{"center_point_box", true}, {"use_dyn_output", true}}),
+        boxes_l,
+        scores_l,
+        max_out_l,
+        iou_threshold,
+        score_threshold);
+    mm->add_return({r});
+
+    p.compile(migraphx::ref::target{});
+    auto output = p.eval({}).back();
+    std::vector<int64_t> result;
+    output.visit([&](auto out) { result.assign(out.begin(), out.end()); });
+    std::vector<int64_t> gold = {0, 0, 3, 0, 0, 0, 0, 0, 5};
+    EXPECT(migraphx::verify_range(result, gold));
+}
+
+TEST_CASE(nms_dynamic_batch_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape boxes_s{migraphx::shape::float_type, {{1, 3, 0}, {6, 6, 0}, {4, 4, 0}}};
+
+    migraphx::shape scores_s{migraphx::shape::float_type, {{1, 3, 0}, {1, 1, 0}, {6, 6, 0}}};
+
+    auto boxes_p         = mm->add_parameter("boxes", boxes_s);
+    auto scores_p        = mm->add_parameter("scores", scores_s);
+    auto max_out_l       = mm->add_literal(int64_t{4});
+    auto iou_threshold   = mm->add_literal(0.5f);
+    auto score_threshold = mm->add_literal(0.0f);
+
+    auto r = mm->add_instruction(
+        migraphx::make_op("nonmaxsuppression",
+                          {{"center_point_box", true}, {"use_dyn_output", true}}),
+        boxes_p,
+        scores_p,
+        max_out_l,
+        iou_threshold,
+        score_threshold);
+    mm->add_return({r});
+
+    p.compile(migraphx::ref::target{});
+
+    std::vector<float> boxes_vec  = {0.5, 0.5,  1.0, 1.0, 0.5, 0.6,  1.0, 1.0, 0.5, 0.4,   1.0, 1.0,
+                                    0.5, 10.5, 1.0, 1.0, 0.5, 10.6, 1.0, 1.0, 0.5, 100.5, 1.0, 1.0,
+                                    0.5, 0.5,  1.0, 1.0, 0.5, 0.6,  1.0, 1.0, 0.5, 0.4,   1.0, 1.0,
+                                    0.5, 10.5, 1.0, 1.0, 0.5, 10.6, 1.0, 1.0, 0.5, 100.5, 1.0, 1.0};
+    std::vector<float> scores_vec = {
+        0.9, 0.75, 0.6, 0.95, 0.5, 0.3, 0.9, 0.75, 0.6, 0.95, 0.5, 0.3};
+
+    migraphx::shape input_fixed_shape0{migraphx::shape::float_type, {2, 6, 4}};
+    migraphx::shape input_fixed_shape1{migraphx::shape::float_type, {2, 1, 6}};
+    migraphx::parameter_map params0;
+    params0["boxes"]  = migraphx::argument(input_fixed_shape0, boxes_vec.data());
+    params0["scores"] = migraphx::argument(input_fixed_shape1, scores_vec.data());
+    auto output       = p.eval(params0).back();
+
+    std::vector<int64_t> result;
+    output.visit([&](auto out) { result.assign(out.begin(), out.end()); });
+    std::vector<int64_t> gold = {0, 0, 3, 0, 0, 0, 0, 0, 5, 1, 0, 3, 1, 0, 0, 1, 0, 5};
+    EXPECT(migraphx::verify_range(result, gold));
+}
+
+TEST_CASE(nms_dynamic_boxes_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape boxes_s{migraphx::shape::float_type, {{1, 1, 0}, {4, 20, 0}, {4, 4, 0}}};
+
+    migraphx::shape scores_s{migraphx::shape::float_type, {{1, 1, 0}, {1, 1, 0}, {4, 20, 0}}};
+
+    auto boxes_p         = mm->add_parameter("boxes", boxes_s);
+    auto scores_p        = mm->add_parameter("scores", scores_s);
+    auto max_out_l       = mm->add_literal(int64_t{4});
+    auto iou_threshold   = mm->add_literal(0.5f);
+    auto score_threshold = mm->add_literal(0.0f);
+
+    auto r = mm->add_instruction(
+        migraphx::make_op("nonmaxsuppression",
+                          {{"center_point_box", true}, {"use_dyn_output", true}}),
+        boxes_p,
+        scores_p,
+        max_out_l,
+        iou_threshold,
+        score_threshold);
+    mm->add_return({r});
+
+    p.compile(migraphx::ref::target{});
+
+    std::vector<float> boxes_vec  = {0.5, 0.5,  1.0, 1.0, 0.5, 0.6,  1.0, 1.0, 0.5, 0.4,   1.0, 1.0,
+                                    0.5, 10.5, 1.0, 1.0, 0.5, 10.6, 1.0, 1.0, 0.5, 100.5, 1.0, 1.0};
+    std::vector<float> scores_vec = {0.9, 0.75, 0.6, 0.95, 0.5, 0.3};
+
+    migraphx::shape input_fixed_shape0{migraphx::shape::float_type, {1, 6, 4}};
+    migraphx::shape input_fixed_shape1{migraphx::shape::float_type, {1, 1, 6}};
+    migraphx::parameter_map params0;
+    params0["boxes"]  = migraphx::argument(input_fixed_shape0, boxes_vec.data());
+    params0["scores"] = migraphx::argument(input_fixed_shape1, scores_vec.data());
+    auto output       = p.eval(params0).back();
+
+    std::vector<int64_t> result;
+    output.visit([&](auto out) { result.assign(out.begin(), out.end()); });
+    std::vector<int64_t> gold = {0, 0, 3, 0, 0, 0, 0, 0, 5};
+    EXPECT(migraphx::verify_range(result, gold));
+}
+
+TEST_CASE(nms_dynamic_classes_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape boxes_s{migraphx::shape::float_type, {{1, 1, 0}, {6, 6, 0}, {4, 4, 0}}};
+
+    migraphx::shape scores_s{migraphx::shape::float_type, {{1, 1, 0}, {1, 3, 0}, {6, 6, 0}}};
+
+    auto boxes_p         = mm->add_parameter("boxes", boxes_s);
+    auto scores_p        = mm->add_parameter("scores", scores_s);
+    auto max_out_l       = mm->add_literal(int64_t{2});
+    auto iou_threshold   = mm->add_literal(0.5f);
+    auto score_threshold = mm->add_literal(0.0f);
+
+    auto r = mm->add_instruction(
+        migraphx::make_op("nonmaxsuppression",
+                          {{"center_point_box", true}, {"use_dyn_output", true}}),
+        boxes_p,
+        scores_p,
+        max_out_l,
+        iou_threshold,
+        score_threshold);
+    mm->add_return({r});
+
+    p.compile(migraphx::ref::target{});
+
+    std::vector<float> boxes_vec  = {0.0, 0.0,  1.0, 1.0,  0.0, 0.1,   1.0, 1.1,
+                                    0.0, -0.1, 1.0, 0.9,  0.0, 10.0,  1.0, 11.0,
+                                    0.0, 10.1, 1.0, 11.1, 0.0, 100.0, 1.0, 101.0};
+    std::vector<float> scores_vec = {
+        0.9, 0.75, 0.6, 0.95, 0.5, 0.3, 0.9, 0.75, 0.6, 0.95, 0.5, 0.3};
+    migraphx::shape input_fixed_shape0{migraphx::shape::float_type, {1, 6, 4}};
+    migraphx::shape input_fixed_shape1{migraphx::shape::float_type, {1, 2, 6}};
+    migraphx::parameter_map params0;
+    params0["boxes"]  = migraphx::argument(input_fixed_shape0, boxes_vec.data());
+    params0["scores"] = migraphx::argument(input_fixed_shape1, scores_vec.data());
+    auto output       = p.eval(params0).back();
+
+    std::vector<int64_t> result;
+    output.visit([&](auto out) { result.assign(out.begin(), out.end()); });
+    std::vector<int64_t> gold = {0, 0, 3, 0, 0, 0, 0, 1, 3, 0, 1, 0};
+    EXPECT(migraphx::verify_range(result, gold));
+}
+
 TEST_CASE(nms_not_center_test)
 {
     migraphx::program p;
@@ -3642,12 +3810,14 @@ TEST_CASE(nms_not_center_test)
     auto iou_threshold   = mm->add_literal(0.5f);
     auto score_threshold = mm->add_literal(0.0f);
 
-    auto r = mm->add_instruction(migraphx::make_op("nonmaxsuppression"),
-                                 boxes_l,
-                                 scores_l,
-                                 max_out_l,
-                                 iou_threshold,
-                                 score_threshold);
+    // set use_dyn_output back to false in operator map
+    auto r =
+        mm->add_instruction(migraphx::make_op("nonmaxsuppression", {{"use_dyn_output", false}}),
+                            boxes_l,
+                            scores_l,
+                            max_out_l,
+                            iou_threshold,
+                            score_threshold);
     mm->add_return({r});
 
     p.compile(migraphx::ref::target{});
@@ -3675,12 +3845,13 @@ TEST_CASE(nms_test)
     auto iou_threshold   = mm->add_literal(0.5f);
     auto score_threshold = mm->add_literal(0.0f);
 
-    auto r = mm->add_instruction(migraphx::make_op("nonmaxsuppression", {{"center_point_box", 1}}),
-                                 boxes_l,
-                                 scores_l,
-                                 max_out_l,
-                                 iou_threshold,
-                                 score_threshold);
+    auto r =
+        mm->add_instruction(migraphx::make_op("nonmaxsuppression", {{"center_point_box", true}}),
+                            boxes_l,
+                            scores_l,
+                            max_out_l,
+                            iou_threshold,
+                            score_threshold);
     mm->add_return({r});
 
     p.compile(migraphx::ref::target{});
@@ -3712,12 +3883,13 @@ TEST_CASE(nms_transpose1_test)
 
     auto transpose_boxes = mm->add_instruction(
         migraphx::make_op("transpose", {{"permutation", {0, 2, 1}}}), t_boxes_l);
-    auto r = mm->add_instruction(migraphx::make_op("nonmaxsuppression", {{"center_point_box", 1}}),
-                                 transpose_boxes,
-                                 scores_l,
-                                 max_out_l,
-                                 iou_threshold,
-                                 score_threshold);
+    auto r =
+        mm->add_instruction(migraphx::make_op("nonmaxsuppression", {{"center_point_box", true}}),
+                            transpose_boxes,
+                            scores_l,
+                            max_out_l,
+                            iou_threshold,
+                            score_threshold);
     mm->add_return({r});
 
     p.compile(migraphx::ref::target{});
@@ -3749,12 +3921,13 @@ TEST_CASE(nms_transpose2_test)
 
     auto transpose_boxes = mm->add_instruction(
         migraphx::make_op("transpose", {{"permutation", {1, 2, 0}}}), t_boxes_l);
-    auto r = mm->add_instruction(migraphx::make_op("nonmaxsuppression", {{"center_point_box", 1}}),
-                                 transpose_boxes,
-                                 scores_l,
-                                 max_out_l,
-                                 iou_threshold,
-                                 score_threshold);
+    auto r =
+        mm->add_instruction(migraphx::make_op("nonmaxsuppression", {{"center_point_box", true}}),
+                            transpose_boxes,
+                            scores_l,
+                            max_out_l,
+                            iou_threshold,
+                            score_threshold);
     mm->add_return({r});
 
     p.compile(migraphx::ref::target{});