Merge branch 'develop' into mlir-attention

src/include/migraphx/op/round.hpp → src/include/migraphx/op/nearbyint.hpp

 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2023 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
@@ -21,24 +21,28 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  * THE SOFTWARE.
  */
-#ifndef MIGRAPHX_GUARD_OPERATORS_ROUND_HPP
-#define MIGRAPHX_GUARD_OPERATORS_ROUND_HPP
+#ifndef MIGRAPHX_GUARD_OPERATORS_NEARBYINT_HPP
+#define MIGRAPHX_GUARD_OPERATORS_NEARBYINT_HPP
 
 #include <migraphx/op/unary.hpp>
 #include <migraphx/config.hpp>
+#include <fenv.h>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 namespace op {
-
-struct round : unary<round>
+struct nearbyint : unary<nearbyint>
 {
     auto apply() const
     {
-        return [](auto x) { return std::round(x); };
+        return [](auto x) {
+            auto rounding_mode = fegetround();
+            fesetround(FE_TONEAREST);
+            return std::nearbyint(x);
+            fesetround(rounding_mode);
+        };
     }
 };
-
 } // namespace op
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/include/migraphx/op/normalize_attribute.hpp

@@ -40,6 +40,8 @@ namespace op {
  * 2. use_rank (default) vs use_len:
  *  `use_rank` sets the max value/index of the attribute as the rank of lens.
  *  `use_lens` sets the max value/index as the corresponding value in lens at the axes index.
+ *      Uses the dynamic_dimension.max value for dynamic shapes. Returns the original vector
+ *      (no normalization) if any of dynamic_dimension[axes] are not fixed.
  * 3. `clip_min` vs. `not_clip_min` (default):
  *  Clip values less than the minimum to the minimum or not.
  * 4. `include_min` vs. `exclude_min` (default):

src/include/migraphx/op/prefix_scan_op.hpp

@@ -22,6 +22,12 @@
  * THE SOFTWARE.
  */
 
+/**
+ * Parent struct for prefix scan ops.  A prefix scan is a mathematical entity useful
+ * in parallelizing various computations.  Given a list of numbers, a prefix scan
+ * op returns an equal size list of running totals of the values.  Other operations
+ * besides addition can be supported by child ops.
+ */
 #ifndef MIGRAPHX_GUARD_OPERATORS_SCAN_OP_HPP
 #define MIGRAPHX_GUARD_OPERATORS_SCAN_OP_HPP
 

src/include/migraphx/op/quantizelinear.hpp

@@ -30,11 +30,11 @@
 #include <migraphx/par_for.hpp>
 #include <migraphx/value.hpp>
 #include <cmath>
+#include <fenv.h>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 namespace op {
-
 struct quantizelinear
 {
     std::string name() const { return "quantizelinear"; }
@@ -71,26 +71,26 @@ struct quantizelinear
         {
             y_zero_point = args.at(2);
         }
-
         argument result{output_shape};
+        auto rounding_mode = fegetround();
+        fesetround(FE_TONEAREST);
         visit_all(result, y_zero_point)([&](auto output, auto zero_pts) {
             visit_all(x, y_scale)([&](auto input, auto scales) {
                 using quant_type = typename decltype(output)::value_type;
                 auto min_value   = std::numeric_limits<quant_type>::min();
                 auto max_value   = std::numeric_limits<quant_type>::max();
                 par_for(output_shape.elements(), [&](auto i) {
-                    int64_t quantized = static_cast<int64_t>(std::round(input[i] / scales[i])) +
+                    int64_t quantized = static_cast<int64_t>(std::nearbyint(input[i] / scales[i])) +
                                         static_cast<int64_t>(zero_pts[i]);
                     output[i] = std::max(static_cast<int64_t>(min_value),
                                          std::min(static_cast<int64_t>(max_value), quantized));
                 });
             });
         });
-
+        fesetround(rounding_mode);
         return result;
     }
 };
-
 } // namespace op
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/include/migraphx/op/random_uniform.hpp

@@ -65,11 +65,10 @@ struct random_uniform
         return inputs.at(1);
     }
 
-    argument compute(const shape&, std::vector<argument> args) const
+    argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
     {
         // Output goes into the passed buffer, not the shape output.
-        auto result = args[1];
-
+        argument result{dyn_out.computed_shape};
         uint64_t local_seed = args[0].at<uint64_t>(0);
         std::mt19937 gen(local_seed);
 

src/include/migraphx/op/slice.hpp

@@ -38,6 +38,18 @@ namespace op {
 
 /**
  * Slice operator that accepts variable axes, starts and ends.
+ * All of `starts`, `ends`, and `axes` must be supplied by either
+ * their attribute or an input (but not both).
+ *
+ * Valid calls:
+ * slice(input); axes, starts, ends set
+ * slice(input, starts); axes, ends set
+ * slice(input, ends); starts, axes set
+ * slice(input, axes); starts, ends set
+ * slice(input, starts, ends); axes set
+ * slice(input, starts, axes); ends set
+ * slice(input, ends, axes); starts set
+ * slice(input, start, ends, axes); none set
  *
  * Attributes:
  * axes: constant axes to slice over (optional)
@@ -46,8 +58,8 @@ namespace op {
  *
  * Parameters:
  * data: the input tensor to slice (dynamic or static shape)
- * input_starts: starting indicies of slice (optional, static shape)
- * input_ends: ending indicies of slice (optional, static shape)
+ * input_starts: starting indices of slice (optional, static shape)
+ * input_ends: ending indices of slice (optional, static shape)
  * input_axes: axes to slice over (optional, static shape)
  */
 struct slice
@@ -56,6 +68,18 @@ struct slice
     std::vector<int64_t> starts{};
     std::vector<int64_t> ends{};
 
+    /**
+     * Named arrays for the set attribute possibilities.
+     */
+    static constexpr std::array<bool, 3> all_set     = {true, true, true};
+    static constexpr std::array<bool, 3> ends_axes   = {false, true, true};
+    static constexpr std::array<bool, 3> starts_axes = {true, false, true};
+    static constexpr std::array<bool, 3> starts_ends = {true, true, false};
+    static constexpr std::array<bool, 3> axes_only   = {false, false, true};
+    static constexpr std::array<bool, 3> ends_only   = {false, true, false};
+    static constexpr std::array<bool, 3> starts_only = {true, false, false};
+    static constexpr std::array<bool, 3> none_set    = {false, false, false};
+
     template <class Self, class F>
     static auto reflect(Self& self, F f)
     {
@@ -63,24 +87,26 @@ struct slice
     }
 
     /**
-     * Ensure that attribute vectors axes, starts, and ends are all the same size and values are
-     * within limits.
+     * Ensure that attribute axes is within limits.
+     * Will attempt to normalize starts and ends; but will use the dynamic_dimension.max
+     * values for dynamic shapes. This makes it so you have to renormalize for
+     * non-fixed dynamic_dimensions.
      */
     value attributes() const
     {
-        value normalize     = value::object{};
-        normalize["axes"]   = value::array{normalize_attribute::include_min};
-        normalize["starts"] = value::array{normalize_attribute::clip_max,
-                                           normalize_attribute::clip_min,
-                                           normalize_attribute::include_max,
-                                           normalize_attribute::use_len,
-                                           normalize_attribute::include_min};
-        normalize["ends"]   = value::array{normalize_attribute::clip_max,
-                                         normalize_attribute::clip_min,
-                                         normalize_attribute::include_max,
-                                         normalize_attribute::use_len,
-                                         normalize_attribute::include_min};
-        return {{"normalize_axes", normalize}};
+        value normalize_axes     = value::object{};
+        normalize_axes["axes"]   = value::array{normalize_attribute::include_min};
+        normalize_axes["starts"] = value::array{normalize_attribute::clip_max,
+                                                normalize_attribute::clip_min,
+                                                normalize_attribute::include_max,
+                                                normalize_attribute::use_len,
+                                                normalize_attribute::include_min};
+        normalize_axes["ends"]   = value::array{normalize_attribute::clip_max,
+                                              normalize_attribute::clip_min,
+                                              normalize_attribute::include_max,
+                                              normalize_attribute::use_len,
+                                              normalize_attribute::include_min};
+        return {{"normalize_axes", normalize_axes}};
     }
 
     std::string name() const { return "slice"; }
@@ -88,7 +114,7 @@ struct slice
     /**
      * Computes the slice output shape dimensions for given starts, ends,and axes.
      * Templated to also handle tensor views.
-     * Possibily different type between [in_starts, in_ends] and [in_axes] if in_axes is this
+     * Possibly different type between [in_starts, in_ends] and [in_axes] if in_axes is this
      * object's axes attribute. Assumes in_starts and in_ends are normalized; in_axes are valid.
      */
     template <class A, class B>
@@ -104,62 +130,160 @@ struct slice
         return new_lens;
     }
 
-    shape normalize_compute_shape(std::vector<shape> inputs) const
+    /// Get the attributes that are non-empty
+    std::array<bool, 3> get_set_attributes() const
     {
-        check_shapes{inputs, *this, true}.has(1, 3, 4);
-        auto input_shape = inputs[0];
-        if(inputs.size() == 1)
+        std::array<std::vector<int64_t>, 3> attrs = {this->starts, this->ends, this->axes};
+        std::array<bool, 3> bool_vec;
+        std::transform(
+            attrs.cbegin(), attrs.cend(), bool_vec.begin(), [](auto a) { return not a.empty(); });
+        return bool_vec;
+    }
+
+    /// Helper function for normalize_compute_shape()
+    shape compute_two_or_more(std::vector<shape> inputs) const
+    {
+        auto input_shape    = inputs[0];
+        auto set_attributes = get_set_attributes();
+        // check that inputs [1, end) are all 1D, have the same
+        // dimension, and are static
+        check_shapes{inputs.begin() + 1,
+                     inputs.end(),
+                     std::string("SLICE: inputs (starts, ends, and input_axes)"),
+                     false}
+            .only_dims(1)
+            .same_dims();
+        auto dds = input_shape.to_dynamic().dyn_dims();
+        if(inputs.size() == 2)
         {
-            auto t = input_shape.type();
-            if(input_shape.dynamic() and std::any_of(axes.begin(), axes.end(), [&](auto axis) {
-                   return not input_shape.dyn_dims()[axis].is_fixed();
-               }))
+            if(set_attributes == ends_axes)
             {
-                MIGRAPHX_THROW("SLICE: slicing is not allowed on non-fixed dynamic input axis ");
+                // attr ends and axes set; inputs are (data, input_starts)
+                if(inputs[1].lens().at(0) != axes.size())
+                {
+                    MIGRAPHX_THROW("SLICE: 2 input and attributes mismatch");
+                }
+                std::for_each(axes.cbegin(), axes.cend(), [&](const auto& axis) {
+                    dds.at(axis) = {0, dds.at(axis).max};
+                });
             }
-            if(input_shape.dynamic())
+            else if(set_attributes == starts_axes)
             {
-                return shape{t,
-                             lens_calc(input_shape.min_lens(), starts, ends, axes),
-                             lens_calc(input_shape.max_lens(), starts, ends, axes),
-                             {}};
+                // attr starts and axes set; inputs are (data, input_ends)
+                if(inputs[1].lens().at(0) != axes.size())
+                {
+                    MIGRAPHX_THROW("SLICE: 2 input and attributes mismatch");
+                }
+                std::for_each(axes.cbegin(), axes.cend(), [&](const auto& axis) {
+                    dds.at(axis) = {0, dds.at(axis).max};
+                });
+            }
+            else if(set_attributes == starts_ends)
+            {
+                // attr starts and ends set; inputs are (data, input_axes)
+                if(inputs[1].lens().at(0) != starts.size())
+                {
+                    MIGRAPHX_THROW("SLICE: 2 input and attributes mismatch");
+                }
+                std::transform(dds.begin(), dds.end(), dds.begin(), [](auto dd) {
+                    return shape::dynamic_dimension{0, dd.max};
+                });
             }
             else
             {
-                return shape{
-                    t, lens_calc(input_shape.lens(), starts, ends, axes), input_shape.strides()};
+                MIGRAPHX_THROW("SLICE: Invalid 2 input and attributes configuration");
             }
         }
-        else
+        else if(inputs.size() == 3)
         {
-            // check that starts, ends, and optionally input_axes are all 1D, have the same
-            // dimension, and are static
-            check_shapes{inputs.begin() + 1,
-                         inputs.end(),
-                         std::string("SLICE: inputs (starts, ends, and input_axes)"),
-                         false}
-                .only_dims(1)
-                .same_dims();
-            auto dds = input_shape.to_dynamic().dyn_dims();
-            if(inputs.size() == 3)
+            if(set_attributes == axes_only)
             {
+                // attr axes set; inputs are (data, input_starts, input_ends)
                 if(inputs[1].lens().at(0) != axes.size())
                 {
-                    MIGRAPHX_THROW("SLICE: inputs starts and ends do not have the same dimension "
-                                   "as the axes attribute");
+                    MIGRAPHX_THROW("SLICE: 3 input and attributes mismatch");
                 }
                 std::for_each(axes.cbegin(), axes.cend(), [&](const auto& axis) {
                     dds.at(axis) = {0, dds.at(axis).max};
                 });
             }
-            else
+            else if(set_attributes == ends_only)
+            {
+                // attr ends set; inputs are (data, input_starts, input_axes)
+                if(inputs[1].lens().at(0) != ends.size())
+                {
+                    MIGRAPHX_THROW("SLICE: 3 input and attributes mismatch");
+                }
+                std::transform(dds.begin(), dds.end(), dds.begin(), [](auto dd) {
+                    return shape::dynamic_dimension{0, dd.max};
+                });
+            }
+            else if(set_attributes == starts_only)
+
             {
-                // if axes is an input, then all the output dimensions could be 0 to the max value
+                // attr starts set; inputs are (data, input_ends, input_axes)
+                if(inputs[1].lens().at(0) != starts.size())
+                {
+                    MIGRAPHX_THROW("SLICE: 3 input and attributes mismatch");
+                }
                 std::transform(dds.begin(), dds.end(), dds.begin(), [](auto dd) {
                     return shape::dynamic_dimension{0, dd.max};
                 });
             }
-            return shape{input_shape.type(), dds};
+            else
+            {
+                MIGRAPHX_THROW("Invalid 3 input and attributes configuration");
+            }
+        }
+        else
+        {
+            // all 4 inputs (data, inputs_starts, input_ends, input_axes)
+            std::transform(dds.begin(), dds.end(), dds.begin(), [](auto dd) {
+                return shape::dynamic_dimension{0, dd.max};
+            });
+        }
+        return shape{input_shape.type(), dds};
+    }
+
+    // uses the normalize_axes flag to normalize axes, starts, and ends
+    shape normalize_compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs, *this, true}.has(1, 2, 3, 4);
+        if(inputs.size() == 1)
+        {
+            auto input_shape    = inputs[0];
+            auto set_attributes = get_set_attributes();
+            if(set_attributes != all_set)
+            {
+                MIGRAPHX_THROW("SLICE 1_arg: Invalid 1 input and attributes configuration");
+            }
+            // NOTE: make sure to update how normalization works here if this type of slicing is
+            // changed to be allowed
+            if(input_shape.dynamic() and std::any_of(axes.begin(), axes.end(), [&](auto axis) {
+                   return not input_shape.dyn_dims()[axis].is_fixed();
+               }))
+            {
+                MIGRAPHX_THROW(
+                    "SLICE 1_arg: slicing is not allowed on non-fixed dynamic input axis ");
+            }
+            if(input_shape.dynamic())
+            {
+                return shape{
+                    input_shape.type(),
+                    lens_calc(input_shape.min_lens(), this->starts, this->ends, this->axes),
+                    lens_calc(input_shape.max_lens(), this->starts, this->ends, this->axes),
+                    {}};
+            }
+            else
+            {
+                return shape{input_shape.type(),
+                             lens_calc(input_shape.lens(), this->starts, this->ends, this->axes),
+                             input_shape.strides()};
+            }
+        }
+        else
+        {
+            return compute_two_or_more(inputs);
         }
     }
 
@@ -194,14 +318,14 @@ struct slice
 
     /**
      * Calculates the starting offset for the sliced tensor (for aliasing).
-     * Used when the starts and/or the axes are inputs.
+     * Used for 2-4 inputs to `slice.
      *
      * \param s static input shape
      * \param input_starts starting indices of slice
      * \param ax_vec axes to slice on
      */
-    template <class IndView, class Axes>
-    auto compute_offset(const shape& s, const IndView& input_starts, const Axes& ax_vec) const
+    template <class T>
+    auto compute_offset(const shape& s, const T& input_starts, const T& ax_vec) const
     {
         auto ret = 0;
         for(std::size_t i = 0; i < ax_vec.size(); ++i)
@@ -212,106 +336,168 @@ struct slice
         return ret * s.type_size();
     }
 
-    std::unordered_map<std::string, std::vector<int64_t>>
-    normalize_inputs(const shape& input_shape,
-                     const std::vector<int64_t>& input_starts,
-                     const std::vector<int64_t>& input_ends) const
-    {
-        auto attrs = this->attributes().at("normalize_axes");
-        return {{"input_starts",
-                 normalize_indices(input_starts,
-                                   this->axes,
-                                   input_shape,
-                                   attrs.at("starts"),
-                                   "Slice variable input_starts")},
-                {"input_ends",
-                 normalize_indices(input_ends,
-                                   this->axes,
-                                   input_shape,
-                                   attrs.at("ends"),
-                                   "Slice variable input_ends")}};
-    }
-
     /**
-     * Three input version of the normalize_inputs.
-     * This one also checks that the input_axes are valid.
+     * If given, normalize the inputs. Otherwise get from operator attributes.
+     * Return the values in a map.
+     *
+     * Parameters
+     * input_shape: static shape of the input
+     * input_starts: optional
+     * input_ends: optional
+     * input_ends: optional
      */
     std::unordered_map<std::string, std::vector<int64_t>>
-    normalize_inputs(shape input_shape,
-                     const std::vector<int64_t>& input_starts,
-                     const std::vector<int64_t>& input_ends,
-                     const std::vector<int64_t>& input_axes) const
+    normalize_starts_ends_axes(shape input_shape,
+                               const optional<std::vector<int64_t>>& input_starts,
+                               const optional<std::vector<int64_t>>& input_ends,
+                               const optional<std::vector<int64_t>>& input_axes) const
     {
-        auto attrs = this->attributes().at("normalize_axes");
-        auto norm_axes =
-            normalize_axes(input_axes, input_shape, attrs.at("axes"), "Slice variable input_axes");
-        return {{"input_starts",
-                 normalize_indices(input_starts,
-                                   norm_axes,
-                                   input_shape,
-                                   attrs.at("starts"),
-                                   "Slice variable input_starts")},
-                {"input_ends",
-                 normalize_indices(input_ends,
-                                   norm_axes,
-                                   input_shape,
-                                   attrs.at("ends"),
-                                   "Slice variable input ends")},
-                {"input_axes", norm_axes}};
+        auto axes_attrs = this->attributes().at("normalize_axes");
+        std::vector<int64_t> norm_starts;
+        std::vector<int64_t> norm_ends;
+        std::vector<int64_t> norm_axes;
+        if(input_axes)
+        {
+            norm_axes = normalize_axes(input_axes.value(),
+                                       input_shape,
+                                       axes_attrs.at("axes"),
+                                       "Slice variable input_axes");
+        }
+        else
+        {
+            norm_axes = this->axes;
+        }
+        if(input_starts)
+        {
+            norm_starts = normalize_indices(input_starts.value(),
+                                            norm_axes,
+                                            input_shape,
+                                            axes_attrs.at("starts"),
+                                            "Slice variable input_starts");
+        }
+        else
+        {
+            norm_starts = this->starts;
+        }
+        if(input_ends)
+        {
+            norm_ends = normalize_indices(input_ends.value(),
+                                          norm_axes,
+                                          input_shape,
+                                          axes_attrs.at("ends"),
+                                          "Slice variable input ends");
+        }
+        else
+        {
+            norm_ends = this->ends;
+        }
+        return {{"norm_starts", norm_starts}, {"norm_ends", norm_ends}, {"norm_axes", norm_axes}};
     }
 
     argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
     {
         auto input       = args[0];
         auto input_shape = input.get_shape();
-        switch(args.size())
+        if(args.size() == 1)
         {
-        case 1: {
             std::size_t offset = compute_offset(input_shape);
             return {dyn_out.computed_shape, [=] { return input.data() + offset; }};
         }
-        case 3: {
-            shape calc_shape;
-            std::size_t offset = 0;
-            visit_all(args[1], args[2])([&](auto input_starts, auto input_ends) {
-                auto norm_inputs = normalize_inputs(input_shape,
-                                                    input_starts.template to_vector<int64_t>(),
-                                                    input_ends.template to_vector<int64_t>());
-                offset = compute_offset(input_shape, norm_inputs.at("input_starts"), this->axes);
-                calc_shape = {input_shape.type(),
-                              lens_calc(input_shape.lens(),
-                                        norm_inputs.at("input_starts"),
-                                        norm_inputs.at("input_ends"),
-                                        this->axes),
-                              input_shape.strides()};
-            });
-            return {calc_shape, [=] { return input.data() + offset; }};
-        }
-        case 4: {
-            shape calc_shape;
-            std::size_t offset = 0;
-            visit_all(args[1], args[2], args[3])(
-                [&](auto input_starts, auto input_ends, auto input_axes) {
-                    auto norm_inputs = normalize_inputs(input_shape,
-                                                        input_starts.template to_vector<int64_t>(),
-                                                        input_ends.template to_vector<int64_t>(),
-                                                        input_axes.template to_vector<int64_t>());
-                    offset           = compute_offset(
-                        input_shape, norm_inputs.at("input_starts"), norm_inputs.at("input_axes"));
-                    calc_shape = shape{input_shape.type(),
-                                       lens_calc(input_shape.lens(),
-                                                 norm_inputs.at("input_starts"),
-                                                 norm_inputs.at("input_ends"),
-                                                 norm_inputs.at("input_axes")),
-                                       input_shape.strides()};
+        else
+        {
+            // Note that we re-normalize both the attributes and inputs because of the non-fixed
+            // dynamic input shape case. It's possible to only re-normalize if slicing over
+            // non-fixed dynamic_dimensions.
+            auto set_attributes = get_set_attributes();
+            std::unordered_map<std::string, std::vector<int64_t>> norm_inputs;
+            if(set_attributes == ends_axes)
+            {
+                // attr ends and axes set; inputs are (data, input_starts)
+                args[1].visit([&](auto input_starts) {
+                    norm_inputs =
+                        normalize_starts_ends_axes(input_shape,
+                                                   input_starts.template to_vector<int64_t>(),
+                                                   this->ends,
+                                                   this->axes);
+                });
+            }
+            else if(set_attributes == starts_axes)
+            {
+                // attr starts and axes set; inputs are (data, input_ends)
+                args[1].visit([&](auto input_ends) {
+                    norm_inputs =
+                        normalize_starts_ends_axes(input_shape,
+                                                   this->starts,
+                                                   input_ends.template to_vector<int64_t>(),
+                                                   this->axes);
+                });
+            }
+            else if(set_attributes == starts_ends)
+            {
+                // attr starts and ends set; inputs are (data, input_axes)
+                args[1].visit([&](auto input_axes) {
+                    norm_inputs =
+                        normalize_starts_ends_axes(input_shape,
+                                                   this->starts,
+                                                   this->ends,
+                                                   input_axes.template to_vector<int64_t>());
                 });
+            }
+            else if(set_attributes == axes_only)
+            {
+                // attr axes set; inputs are (data, input_starts, input_ends)
+                visit_all(args[1], args[2])([&](auto input_starts, auto input_ends) {
+                    norm_inputs =
+                        normalize_starts_ends_axes(input_shape,
+                                                   input_starts.template to_vector<int64_t>(),
+                                                   input_ends.template to_vector<int64_t>(),
+                                                   this->axes);
+                });
+            }
+            else if(set_attributes == ends_only)
+            {
+                // attr ends set; inputs are (data, input_starts, input_axes)
+                visit_all(args[1], args[2])([&](auto input_starts, auto input_axes) {
+                    norm_inputs =
+                        normalize_starts_ends_axes(input_shape,
+                                                   input_starts.template to_vector<int64_t>(),
+                                                   this->ends,
+                                                   input_axes.template to_vector<int64_t>());
+                });
+            }
+            else if(set_attributes == starts_only)
+            {
+                // attr starts set; inputs are (data, input_ends, input_axes)
+                visit_all(args[1], args[2])([&](auto input_ends, auto input_axes) {
+                    norm_inputs =
+                        normalize_starts_ends_axes(input_shape,
+                                                   this->starts,
+                                                   input_ends.template to_vector<int64_t>(),
+                                                   input_axes.template to_vector<int64_t>());
+                });
+            }
+            else
+            {
+                // no attr set, all inputs
+                visit_all(args[1], args[2], args[3])(
+                    [&](auto input_starts, auto input_ends, auto input_axes) {
+                        norm_inputs =
+                            normalize_starts_ends_axes(input_shape,
+                                                       input_starts.template to_vector<int64_t>(),
+                                                       input_ends.template to_vector<int64_t>(),
+                                                       input_axes.template to_vector<int64_t>());
+                    });
+            }
+            auto offset = compute_offset(
+                input_shape, norm_inputs.at("norm_starts"), norm_inputs.at("norm_axes"));
+            shape calc_shape = shape{input_shape.type(),
+                                     lens_calc(input_shape.lens(),
+                                               norm_inputs.at("norm_starts"),
+                                               norm_inputs.at("norm_ends"),
+                                               norm_inputs.at("norm_axes")),
+                                     input_shape.strides()};
             return {calc_shape, [=] { return input.data() + offset; }};
         }
-        default: {
-            // Should never get here; covering in case some code change occurs
-            MIGRAPHX_THROW("SLICE: invalid number of inputs");
-        }
-        }
     }
 
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }

src/include/migraphx/operators.hpp

@@ -84,6 +84,7 @@
 #include <migraphx/op/mod.hpp>
 #include <migraphx/op/mul.hpp>
 #include <migraphx/op/multibroadcast.hpp>
+#include <migraphx/op/nearbyint.hpp>
 #include <migraphx/op/neg.hpp>
 #include <migraphx/op/nonmaxsuppression.hpp>
 #include <migraphx/op/nonzero.hpp>
@@ -110,7 +111,6 @@
 #include <migraphx/op/rnn_variable_seq_lens.hpp>
 #include <migraphx/op/rnn_var_sl_last_output.hpp>
 #include <migraphx/op/roialign.hpp>
-#include <migraphx/op/round.hpp>
 #include <migraphx/op/rsqrt.hpp>
 #include <migraphx/op/scalar.hpp>
 #include <migraphx/op/scatter_add.hpp>

src/normalize_attributes.cpp

@@ -66,15 +66,15 @@ auto tune_attribute(const std::vector<int64_t>& vec,
     {
         if(input_shape.dynamic())
         {
+            // return the unchanged `vec` if the dynamic_dimensions at `axes` are not fixed
+            if(std::any_of(axes.begin(), axes.end(), [&](auto ax) {
+                   return not input_shape.dyn_dims().at(ax).is_fixed();
+               }))
+            {
+                return vec;
+            }
             std::transform(axes.begin(), axes.end(), max_vals.begin(), [&](auto i) {
-                const auto& dd = input_shape.dyn_dims().at(i);
-                if(not dd.is_fixed())
-                {
-                    MIGRAPHX_THROW(
-                        "NORMALIZE_ATTR: 'use_lens' on a non-fixed dynamic dimension, axis=" +
-                        std::to_string(i));
-                }
-                return dd.max;
+                return input_shape.dyn_dims().at(i).max;
             });
         }
         else

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

@@ -97,10 +97,11 @@ struct onnx_parser
     shape::dynamic_dimension default_dyn_dim_value = {1, 1};
     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;
+    bool use_dyn_output          = false;
+    bool skip_unknown_operators  = false;
+    int64_t max_loop_iterations  = 10;
+    int64_t limit_max_iterations = std::numeric_limits<uint16_t>::max();
+    int64_t opset_version        = 13;
 
     std::unordered_map<std::string, op_func> ops;
 

src/onnx/onnx.cpp

@@ -67,6 +67,7 @@ program parse_onnx_from(const onnx_options& options, Ts&&... xs)
     }
     parser.skip_unknown_operators = options.skip_unknown_operators;
     parser.max_loop_iterations    = options.max_loop_iterations;
+    parser.limit_max_iterations   = options.limit_max_iterations;
     parser.use_dyn_output         = options.use_dyn_output;
 
     if(options.print_program_on_error)

src/onnx/parse_clip.cpp

 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2023 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

src/onnx/parse_generic_op.cpp

@@ -60,7 +60,7 @@ struct parse_generic_op : op_parser<parse_generic_op>
                 {"Neg", "neg"},
                 {"Reciprocal", "recip"},
                 {"Relu", "relu"},
-                {"Round", "round"},
+                {"Round", "nearbyint"},
                 {"Sigmoid", "sigmoid"},
                 {"Sign", "sign"},
                 {"Sin", "sin"},

src/onnx/parse_isinf.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2023 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>
+#include <migraphx/instruction.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace onnx {
+
+struct parse_isinf : op_parser<parse_isinf>
+{
+    std::vector<op_desc> operators() const { return {{"IsInf", "isinf"}}; }
+
+    instruction_ref parse(const op_desc& /*opd*/,
+                          const onnx_parser& parser,
+                          onnx_parser::node_info info,
+                          const std::vector<instruction_ref>& args) const
+    {
+        bool detect_negative = true;
+        bool detect_positive = true;
+        if(contains(info.attributes, "detect_negative"))
+        {
+            detect_negative = static_cast<bool>(
+                parser.parse_value(info.attributes.at("detect_negative")).at<int>());
+        }
+
+        if(contains(info.attributes, "detect_positive"))
+        {
+            detect_positive = static_cast<bool>(
+                parser.parse_value(info.attributes.at("detect_positive")).at<int>());
+        }
+
+        auto x_shape = args[0]->get_shape();
+        if(not detect_negative and not detect_positive)
+        {
+            return info.add_instruction(
+                make_op("multibroadcast", {{"out_lens", x_shape.lens()}}),
+                info.add_literal(migraphx::literal{migraphx::shape{shape::bool_type}, {false}}));
+        }
+
+        auto is_inf = info.add_instruction(make_op("isinf"), args[0]);
+        if(detect_negative and detect_positive)
+        {
+            return is_inf;
+        }
+
+        auto zero_l = info.add_literal(migraphx::literal{migraphx::shape{x_shape.type()}, {0}});
+        auto mb_zero =
+            info.add_instruction(make_op("multibroadcast", {{"out_lens", x_shape.lens()}}), zero_l);
+
+        auto cond = info.add_broadcastable_binary_op(
+            detect_negative ? "less" : "greater", args[0], mb_zero);
+        if(cond->get_shape().type() != shape::bool_type)
+        {
+            cond =
+                info.add_instruction(make_op("convert", {{"target_type", shape::bool_type}}), cond);
+        }
+        return info.add_instruction(make_op("logical_and"), is_inf, cond);
+    }
+};
+
+} // namespace onnx
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/onnx/parse_loop.cpp

@@ -58,6 +58,16 @@ struct parse_loop : op_parser<parse_loop>
             }
         }
 
+        // cap max_iter because loop uses static shapes with max_iter size and huge numbers
+        // here can cause overflow
+        if(max_iterations > parser.limit_max_iterations)
+        {
+            std::cerr << "WARNING: PARSE_LOOP max_iterations exceeds the maximum loop "
+                         "iterations limit, it will be changed from "
+                      << max_iterations << " to " << parser.limit_max_iterations << ".\n";
+            max_iterations = parser.limit_max_iterations;
+        }
+
         // condition input is empty
         if(args.at(1)->name() == "undefined")
         {

src/onnx/parse_multinomial.cpp

 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2023 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
@@ -41,6 +41,9 @@ struct parse_multinomial : op_parser<parse_multinomial>
                           const onnx_parser::node_info& info,
                           std::vector<instruction_ref> args) const
     {
+        if(args.empty())
+            MIGRAPHX_THROW("PARSE_MULTINOMIAL: no arguments given");
+
         int dtype = 6;
         if(contains(info.attributes, "dtype"))
             dtype = info.attributes.at("dtype").i();
@@ -49,35 +52,90 @@ struct parse_multinomial : op_parser<parse_multinomial>
         size_t sample_size = 1;
         if(contains(info.attributes, "sample_size"))
             sample_size = info.attributes.at("sample_size").i();
+        else
+            MIGRAPHX_THROW("PARSE_MULTINOMIAL: sample_size not given");
+
+        // Use logarithmic math to scale probabilities while avoiding division by very
+        // small numbers.  Scaling by the maximum makes very tiny ranges more
+        // tractable; any constant factor gives equivalent distr. since the Multinomial op.
+        // normalizes at runtime.
 
         // Subtract the per-batch maximum log-probability, making the per-batch max 0
         auto maxes =
             info.add_instruction(migraphx::make_op("reduce_max", {{"axes", {1}}}), args[0]);
-        auto mb_maxes = info.add_instruction(
-            migraphx::make_op("multibroadcast", {{"out_lens", args[0]->get_shape().lens()}}),
-            maxes);
-        auto cdf = info.add_instruction(migraphx::make_op("sub"), args[0], mb_maxes);
+        auto cdf = info.add_common_op("sub", args[0], maxes);
         // Take the element-wise exponent to get probabilities in the range (0, 1]
         cdf = info.add_instruction(migraphx::make_op("exp"), cdf);
-        // Compute the cumulative density function
+        // Compute the cumulative distribution function
         cdf = info.add_instruction(
             migraphx::make_op("prefix_scan_sum", {{"axis", 1}, {"exclusive", false}}), cdf);
 
-        // Pre-compute random distribution
-        std::mt19937 gen(std::chrono::high_resolution_clock::now().time_since_epoch().count());
+        instruction_ref seed_input;
         if(contains(info.attributes, "seed"))
-            gen.seed(info.attributes.at("seed").f());
+        {
+            float seed = info.attributes.at("seed").f();
+            migraphx::shape s{migraphx::shape::float_type, {1}};
+            std::vector<float> data = {seed};
+            seed_input              = info.add_literal(migraphx::literal(s, data));
+        }
+        else
+        {
+            seed_input = info.add_instruction(migraphx::make_op("random_seed"));
+        }
+        instruction_ref randoms;
+
+        shape s0 = args[0]->get_shape();
+
+        if(s0.dynamic())
+        {
+            //  Dynamic batch_size will be taken from args[0].  The input argument to this should
+            // have a second dimension of sample_size.
+            std::vector<shape::dynamic_dimension> dyn_dim_set;
+            dyn_dim_set.emplace_back(s0.dyn_dims().front());
+            dyn_dim_set.emplace_back(shape::dynamic_dimension{sample_size, sample_size});
+
+            // read the input dimensions
+            auto dim_of =
+                info.add_instruction(migraphx::make_op("dimensions_of", {{"end", 2}}), args[0]);
+
+            // The next two operations insert the value sample_size into the second array position
+
+            // make an argument of (1, 0)
+            shape s(shape::int64_type, {2});
+            std::vector<int64_t> data1{1, 0};
+            auto l1        = info.add_literal(s, data1);
+            auto batch_arg = info.add_instruction(migraphx::make_op("mul"), dim_of, l1);
+            std::vector<int64_t> data2(2, 0);
+            // make an argument of (0, sample_size)
+            data2[1]         = sample_size;
+            auto l2          = info.add_literal(s, data2);
+            auto alloc_shape = info.add_instruction(migraphx::make_op("add"), batch_arg, l2);
+            // alloc_shape should contain the input-based shape dimensions as its values at runtime,
+            // and its own shape is {2}
+
+            // compile_shape is the shape used when compiling the Allocate op, and may be dynamic
+            migraphx::shape compile_shape =
+                migraphx::shape(s0.type(), {s0.dyn_dims().front(), {sample_size, sample_size}});
 
-        std::uniform_real_distribution<> dis(0.0, 1.0);
-        size_t batch_size = args[0]->get_shape().lens().front();
-        migraphx::shape dist_shape{migraphx::shape::float_type, {batch_size, sample_size}};
+            // Allocate on-device storage for the random values
+            auto alloc = info.add_instruction(
+                migraphx::make_op("allocate", {{"shape", to_value(compile_shape)}}), alloc_shape);
+            randoms = info.add_instruction(migraphx::make_op("random_uniform"), seed_input, alloc);
+        }
+        else
+        {
+            // use literal.  The array populated by random_uniform may have any shape, as long its
+            // number of elements is batch_size * sample_size .
+            size_t batch_size = s0.lens().front();
+            auto rand_dummy   = info.add_literal(
+                migraphx::literal{migraphx::shape::float_type, {batch_size * sample_size}});
 
-        std::vector<float> random_dist(batch_size * sample_size);
-        std::generate(random_dist.begin(), random_dist.end(), [&]() { return dis(gen); });
-        auto dist_lit = info.add_literal(migraphx::literal{dist_shape, random_dist});
+            randoms =
+                info.add_instruction(migraphx::make_op("random_uniform"), seed_input, rand_dummy);
+        }
 
         return info.add_instruction(
-            migraphx::make_op("multinomial", {{"dtype", output_type}}), cdf, dist_lit);
+            migraphx::make_op("multinomial", {{"dtype", output_type}}), cdf, randoms);
     }
 };
 

src/onnx/parse_resize.cpp

@@ -181,6 +181,76 @@ static std::string get_nearest_mode(const onnx_parser::attribute_map& attr)
     return nearest_mode;
 }
 
+static std::vector<double> get_scales(const onnx_parser::attribute_map& attr)
+{
+    std::vector<double> scales;
+    if(contains(attr, "scales"))
+    {
+        copy(attr.at("scales").floats(), std::back_inserter(scales));
+    }
+
+    return scales;
+}
+
+static void parse_args(const std::vector<instruction_ref>& args,
+                       const std::vector<size_t>& in_lens,
+                       const std::string& op_name,
+                       std::vector<double>& vec_scale,
+                       std::vector<std::size_t>& out_lens)
+{
+    for(const auto& arg : args)
+    {
+        if(arg->name() == "undefined" or arg == args.front())
+        {
+            continue;
+        }
+
+        // skipped empty input
+        auto lens = arg->get_shape().lens();
+        if(lens.empty())
+        {
+            continue;
+        }
+
+        auto type = arg->get_shape().type();
+        // output size
+        if(type == shape::int64_type)
+        {
+            auto arg_out_s = arg->eval();
+            check_arg_empty(arg_out_s,
+                            "PARSE_" + op_name + ": dynamic output size is not supported!");
+            arg_out_s.visit([&](const auto& ol) { out_lens.assign(ol.begin(), ol.end()); });
+
+            if(out_lens.size() != in_lens.size())
+            {
+                MIGRAPHX_THROW("PARSE_" + op_name +
+                               ": specified output size does not match input size");
+            }
+
+            // compute the scale
+            vec_scale.resize(in_lens.size());
+            std::transform(in_lens.begin(),
+                           in_lens.end(),
+                           out_lens.begin(),
+                           vec_scale.begin(),
+                           [](auto iss, auto oss) { return 1.0 * oss / iss; });
+        }
+        else
+        {
+
+            // scale input
+            if(lens[0] == in_lens.size())
+            {
+                auto arg_scale = arg->eval();
+                check_arg_empty(arg_scale,
+                                "PARSE_" + op_name + ": dynamic input scale is not supported!");
+
+                arg_scale.visit([&](const auto& v) { vec_scale.assign(v.begin(), v.end()); });
+            }
+        }
+    }
+}
+
 struct parse_resize : op_parser<parse_resize>
 {
     std::vector<op_desc> operators() const { return {{"Resize"}, {"Upsample"}}; }
@@ -214,72 +284,30 @@ struct parse_resize : op_parser<parse_resize>
         std::vector<std::size_t> out_lens(in_lens.size());
 
         // scale
-        std::vector<double> vec_scale;
+        std::vector<double> vec_scale = get_scales(info.attributes);
 
-        for(const auto& arg : args)
+        // If `scales` was not an attribute, it must be an input
+        if(vec_scale.empty())
         {
-            if(arg->name() == "undefined" or arg == args.front())
-            {
-                continue;
-            }
-
-            // skipped empty input
-            auto lens = arg->get_shape().lens();
-            if(lens.empty())
-            {
-                continue;
-            }
-
-            auto type = arg->get_shape().type();
-            // output size
-            if(type == shape::int64_type)
-            {
-                auto arg_out_s = arg->eval();
-                check_arg_empty(arg_out_s,
-                                "PARSE_" + opd.op_name + ": dynamic output size is not supported!");
-                arg_out_s.visit([&](const auto& ol) { out_lens.assign(ol.begin(), ol.end()); });
-
-                if(out_lens.size() != in_lens.size())
-                {
-                    MIGRAPHX_THROW("PARSE_" + opd.op_name +
-                                   ": specified output size does not match input size");
-                }
+            // Depending on the args, it *must* populate the `vec_scale`, and might populate
+            // `out_lens`
+            parse_args(args, in_lens, opd.op_name, vec_scale, out_lens);
+        }
 
-                // compute the scale
-                vec_scale.resize(in_lens.size());
-                std::transform(in_lens.begin(),
-                               in_lens.end(),
-                               out_lens.begin(),
-                               vec_scale.begin(),
-                               [](auto iss, auto oss) { return 1.0 * oss / iss; });
-            }
-            else
-            {
+        if(in_lens.size() != vec_scale.size())
+        {
+            MIGRAPHX_THROW("PARSE_" + opd.op_name + ": ranks of input and scale are different!");
+        }
 
-                // scale input
-                if(lens[0] == in_lens.size())
-                {
-                    auto arg_scale = arg->eval();
-                    check_arg_empty(arg_scale,
-                                    "PARSE_" + opd.op_name +
-                                        ": dynamic input scale is not supported!");
-
-                    arg_scale.visit([&](const auto& v) { vec_scale.assign(v.begin(), v.end()); });
-                    if(in_lens.size() != vec_scale.size())
-                    {
-                        MIGRAPHX_THROW("PARSE_" + opd.op_name +
-                                       ": ranks of input and scale are different!");
-                    }
-
-                    std::transform(in_lens.begin(),
-                                   in_lens.end(),
-                                   vec_scale.begin(),
-                                   out_lens.begin(),
-                                   [&](auto idx, auto scale) {
-                                       return static_cast<std::size_t>(idx * scale);
-                                   });
-                }
-            }
+        // if the output was not calculated yet, we update it based on the scales
+        if(all_of(out_lens.cbegin(), out_lens.cend(), [](auto o) { return o == 0; }))
+        {
+            std::transform(
+                in_lens.begin(),
+                in_lens.end(),
+                vec_scale.begin(),
+                out_lens.begin(),
+                [&](auto idx, auto scale) { return static_cast<std::size_t>(idx * scale); });
         }
 
         shape out_s{in_s.type(), out_lens};
@@ -288,7 +316,6 @@ struct parse_resize : op_parser<parse_resize>
 
         // reshape input to one-dimension
         std::vector<int64_t> rsp_lens = {static_cast<int64_t>(in_s.elements())};
-        args[0]                       = info.make_contiguous(args[0]);
         auto rsp = info.add_instruction(make_op("reshape", {{"dims", rsp_lens}}), args[0]);
 
         if(mode == "nearest")

src/onnx/parse_slice.cpp

@@ -46,6 +46,9 @@ struct parse_slice : op_parser<parse_slice>
 
         void always_insert(instruction_ref arg) { op_args.insert(op_args.begin(), arg); }
 
+        /**
+         * Either insert argument into `this->op_args` or return the constant value of the argument
+         */
         std::vector<int64_t> insert(instruction_ref arg)
         {
             std::vector<int64_t> result;
@@ -144,16 +147,15 @@ struct parse_slice : op_parser<parse_slice>
             sd.op.axes = axes;
         }
 
-        if(not sd.steps.empty())
+        if(std::any_of(sd.steps.begin(), sd.steps.end(), [](auto s) { return s != 1; }))
         {
             if(sd.op.starts.empty() or sd.op.ends.empty())
-                MIGRAPHX_THROW("PARSE_SLICE: steps and variable starts and ends is not supported");
+                MIGRAPHX_THROW(
+                    "PARSE_SLICE: steps and variable starts and/or ends is not supported");
             if(sd.op.axes.empty())
                 MIGRAPHX_THROW("PARSE_SLICE: steps and variable axes is not supported");
         }
 
-        assert(sd.steps.empty() or sd.steps.size() == sd.op.axes.size());
-
         // If any axes have negative step, prepare to add a "reverse" op
         for(auto i : range(sd.steps.size()))
         {

src/onnx/parse_split.cpp

@@ -68,13 +68,34 @@ struct parse_split : op_parser<parse_split>
         // no split attribute, input is equally divided
         else
         {
-            if((lens[tuned_axis] % info.num_outputs) != 0)
+            std::size_t num_outputs = info.num_outputs;
+            // the num_outputs attribute seems to be redundant since we already have
+            // node_info::num_outputs, but we can still perform an error check
+            if(contains(info.attributes, "num_outputs"))
             {
-                MIGRAPHX_THROW("PARSE_SPLIT: input cannot be equally divided into " +
-                               std::to_string(info.num_outputs) + " splits!");
+                num_outputs =
+                    parser.parse_value(info.attributes.at("num_outputs")).at<std::size_t>();
+                if(num_outputs != info.num_outputs)
+                {
+                    MIGRAPHX_THROW("PARSE_SPLIT: num_outputs attribute " +
+                                   std::to_string(num_outputs) +
+                                   " doesn't match actual number of outputs " +
+                                   std::to_string(info.num_outputs) + "!");
+                }
+            }
+
+            if(lens[tuned_axis] % num_outputs == 0)
+            {
+                std::size_t chunk_size = lens[tuned_axis] / num_outputs;
+                vec_splits.resize(num_outputs, chunk_size);
+            }
+            else
+            {
+                std::size_t chunk_size      = lens[tuned_axis] / num_outputs + 1;
+                std::size_t last_chunk_size = lens[tuned_axis] - chunk_size * (num_outputs - 1);
+                vec_splits.resize(num_outputs - 1, chunk_size);
+                vec_splits.push_back(last_chunk_size);
             }
-            auto dl = lens[tuned_axis] / info.num_outputs;
-            vec_splits.resize(info.num_outputs, dl);
         }
 
         if(std::accumulate(vec_splits.begin(), vec_splits.end(), int64_t(0)) !=

src/py/migraphx_py.cpp

@@ -472,7 +472,8 @@ MIGRAPHX_PYBIND11_MODULE(migraphx, m)
                map_dyn_input_dims,
            bool skip_unknown_operators,
            bool print_program_on_error,
-           int64_t max_loop_iterations) {
+           int64_t max_loop_iterations,
+           int64_t limit_max_iterations) {
             migraphx::onnx_options options;
             options.default_dim_value      = default_dim_value;
             options.default_dyn_dim_value  = default_dyn_dim_value;
@@ -481,6 +482,7 @@ MIGRAPHX_PYBIND11_MODULE(migraphx, m)
             options.skip_unknown_operators = skip_unknown_operators;
             options.print_program_on_error = print_program_on_error;
             options.max_loop_iterations    = max_loop_iterations;
+            options.limit_max_iterations   = limit_max_iterations;
             return migraphx::parse_onnx(filename, options);
         },
         "Parse onnx file",
@@ -492,7 +494,8 @@ MIGRAPHX_PYBIND11_MODULE(migraphx, m)
             std::unordered_map<std::string, std::vector<migraphx::shape::dynamic_dimension>>(),
         py::arg("skip_unknown_operators") = false,
         py::arg("print_program_on_error") = false,
-        py::arg("max_loop_iterations")    = 10);
+        py::arg("max_loop_iterations")    = 10,
+        py::arg("limit_max_iterations")   = std::numeric_limits<uint16_t>::max());
 
     m.def(
         "parse_onnx_buffer",

src/rewrite_quantization.cpp

@@ -47,7 +47,7 @@ void apply_quantizelinear(module& m, instruction_ref ins)
             ins, make_op("convert", {{"target_type", y_scale->get_shape().type()}}), x);
     }
     auto div            = m.insert_instruction(ins, make_op("div"), x, y_scale);
-    auto add_zero_point = m.insert_instruction(ins, make_op("round"), div);
+    auto add_zero_point = m.insert_instruction(ins, make_op("nearbyint"), div);
 
     if(ins->inputs().size() == 3)
     {