Merge

src/include/migraphx/op/argmax.hpp

@@ -62,7 +62,7 @@ struct argmax
         if(s0.dynamic())
         {
             auto dyn_dims  = s0.dyn_dims();
-            dyn_dims[axis] = {1, 1, 0};
+            dyn_dims[axis] = {1, 1};
             return {shape::int64_type, dyn_dims};
         }
         else

src/include/migraphx/op/concat.hpp

@@ -134,7 +134,7 @@ struct concat
             }
 
             auto new_dims  = inputs[0].dyn_dims();
-            new_dims[axis] = migraphx::shape::dynamic_dimension{new_min, new_max, 0};
+            new_dims[axis] = migraphx::shape::dynamic_dimension{new_min, new_max};
             return {inputs[0].type(), new_dims};
         }
         else

src/include/migraphx/op/contiguous.hpp

@@ -48,7 +48,7 @@ struct contiguous
     {
         check_shapes{inputs, *this, true}.has(1);
         auto s0 = inputs.front();
-        if(s0.dynamic() or s0.standard())
+        if(s0.dynamic())
         {
             return s0;
         }

src/include/migraphx/op/convolution.hpp

@@ -24,9 +24,12 @@
 #ifndef MIGRAPHX_GUARD_OPERATORS_CONVOLUTION_HPP
 #define MIGRAPHX_GUARD_OPERATORS_CONVOLUTION_HPP
 
+#include <migraphx/argument.hpp>
 #include <migraphx/op/common.hpp>
 #include <migraphx/check_shapes.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/convolution.hpp>
+#include <migraphx/pad_calc.hpp>
 #include <migraphx/value.hpp>
 #include <cmath>
 #include <utility>
@@ -35,6 +38,10 @@ namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 namespace op {
 
+/**
+ * Convolution operator. Does not support optimal dimensions for spatial dimensions. Returns empty
+ * optimals.
+ */
 struct convolution
 {
     std::vector<std::size_t> padding  = {0, 0};
@@ -145,7 +152,7 @@ struct convolution
             else
             {
                 auto l = input_shape.lens().at(0);
-                output_dyn_dims.push_back({l, l, 0});
+                output_dyn_dims.push_back({l, l});
             }
         };
 
@@ -162,25 +169,30 @@ struct convolution
                 if(x_shape.dynamic())
                 {
                     auto x = x_shape.dyn_dims()[i + 2];
-                    output_dyn_dims.push_back(shape::dynamic_dimension{
-                        ceil_div(x.min, s), ceil_div(x.max, s), ceil_div(x.opt, s)});
+                    std::set<std::size_t> optimals{};
+                    std::transform(x.optimals.begin(),
+                                   x.optimals.end(),
+                                   std::inserter(optimals, optimals.begin()),
+                                   [&](auto o) { return ceil_div(o, s); });
+                    output_dyn_dims.push_back(
+                        shape::dynamic_dimension{ceil_div(x.min, s), ceil_div(x.max, s), optimals});
                 }
                 else
                 {
                     auto od = ceil_div(x_shape.lens()[i + 2], s);
-                    output_dyn_dims.push_back(shape::dynamic_dimension{od, od, 0});
+                    output_dyn_dims.push_back(shape::dynamic_dimension{od, od});
                 }
             }
         }
         else
         {
+            // Does not compute for optimals
             auto min_spatial_dims = calc_conv_lens(x_shape.min_lens(), w_shape.max_lens());
             auto max_spatial_dims = calc_conv_lens(x_shape.max_lens(), w_shape.min_lens());
-            auto opt_spatial_dims = calc_conv_lens(x_shape.opt_lens(), w_shape.opt_lens());
             for(size_t i = 0; i < num_spatial_dims; ++i)
             {
-                output_dyn_dims.push_back(shape::dynamic_dimension{
-                    min_spatial_dims[i], max_spatial_dims[i], opt_spatial_dims[i]});
+                output_dyn_dims.push_back(
+                    shape::dynamic_dimension{min_spatial_dims[i], max_spatial_dims[i], {}});
             }
         }
         return shape{x_shape.type(), output_dyn_dims};
@@ -201,6 +213,37 @@ struct convolution
         check_attribute_size();
         return stride.size();
     }
+
+    argument compute(shape output_shape, std::vector<argument> args) const
+    {
+        std::vector<std::size_t> new_padding;
+        if(padding_mode != op::padding_mode_t::default_)
+        {
+            auto input_lens   = args[0].get_shape().lens();
+            auto weights_lens = args[1].get_shape().lens();
+            new_padding =
+                padding_mode == op::same_upper
+                    ? calc_dyn_auto_pad(input_lens, weights_lens, stride, dilation, true)
+                    : calc_dyn_auto_pad(input_lens, weights_lens, stride, dilation, false);
+            output_shape = compute_padded_shape(
+                args[0].get_shape(), args[1].get_shape(), new_padding, stride, dilation);
+        }
+        else
+        {
+            new_padding = padding;
+            if(output_shape.dynamic())
+            {
+                output_shape =
+                    normalize_compute_shape({args.at(0).get_shape(), args.at(1).get_shape()});
+            }
+        }
+
+        argument result{output_shape};
+        visit_all(result, args[0], args[1])([&](auto output, auto input, auto weights) {
+            migraphx::convolution(output, input, weights, new_padding, stride, group);
+        });
+        return result;
+    }
 };
 
 } // namespace op

src/include/migraphx/op/dequantizelinear.hpp

@@ -40,7 +40,11 @@ struct dequantizelinear
     std::string name() const { return "dequantizelinear"; }
     shape compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.same_dims();
+        check_shapes{inputs, *this}.same_dims().has(2, 3);
+        if(inputs.size() == 3 and inputs[0].type() != inputs[2].type())
+        {
+            MIGRAPHX_THROW("DEQUANTIZELINEAR: Zero point and input should be the same type.");
+        }
         return {inputs[1].type(), inputs[0].lens(), inputs[0].strides()};
     }
 

src/include/migraphx/op/flatten.hpp

@@ -29,6 +29,7 @@
 #include <migraphx/config.hpp>
 #include <migraphx/value.hpp>
 #include <migraphx/op/normalize_attribute.hpp>
+#include <migraphx/dyn_output.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -59,27 +60,22 @@ struct flatten
         auto s = inputs[0];
         if(s.dynamic())
         {
+            // Doesn't handle optimals
             auto min_lens = s.min_lens();
             auto max_lens = s.max_lens();
-            auto opt_lens = s.opt_lens();
             // If any of the opt values is 0, output opt will be 0
             shape::dynamic_dimension x = {
                 std::accumulate(
                     min_lens.begin(), min_lens.begin() + axis, std::size_t{1}, std::multiplies<>{}),
                 std::accumulate(
                     max_lens.begin(), max_lens.begin() + axis, std::size_t{1}, std::multiplies<>{}),
-                std::accumulate(opt_lens.begin(),
-                                opt_lens.begin() + axis,
-                                std::size_t{1},
-                                std::multiplies<>{})};
+                {}};
             shape::dynamic_dimension y = {
                 std::accumulate(
                     min_lens.begin() + axis, min_lens.end(), std::size_t{1}, std::multiplies<>{}),
                 std::accumulate(
                     max_lens.begin() + axis, max_lens.end(), std::size_t{1}, std::multiplies<>{}),
-                std::accumulate(
-                    opt_lens.begin() + axis, opt_lens.end(), std::size_t{1}, std::multiplies<>{}),
-            };
+                {}};
             return {s.type(), {x, y}};
         }
         else

src/include/migraphx/op/gathernd.hpp

@@ -121,7 +121,7 @@ struct gathernd
 
             // A rank 0 output is a scalar
             if(output_ndim == 0)
-                return shape(data_shape.type(), {shape::dynamic_dimension({1, 1, 0})});
+                return shape(data_shape.type(), {shape::dynamic_dimension({1, 1})});
 
             // Part of the output shape comes from indices tensor, part from data tensor
             std::vector<shape::dynamic_dimension> output_dims(output_ndim);

src/include/migraphx/op/nonmaxsuppression.hpp

@@ -119,8 +119,8 @@ struct nonmaxsuppression
                 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});
+            out_lens.push_back({0, max_num_boxes});
+            out_lens.push_back({3, 3});
             return {shape::int64_type, out_lens};
         }
         else

src/include/migraphx/op/pooling.hpp

@@ -89,25 +89,17 @@ struct pooling
         std::vector<std::size_t> output_lens{};
         for(size_t i = 0; i < kdims; ++i)
         {
-            if(input_lens[i + 2] == 0)
-            {
-                // handle opt = 0
-                output_lens.push_back(0);
-            }
-            else
-            {
-                std::size_t padding_factor = 2 * padding[i];
-                if(padding.size() == 2 * kdims)
-                    padding_factor = padding[i] + padding[i + kdims];
-                assert(input_lens[i + 2] + padding_factor >= lengths[i]);
-                std::size_t dim_size = input_lens[i + 2] + padding_factor - lengths[i];
-                std::size_t len =
-                    (ceil_mode)
-                        ? dim_size / stride[i] + static_cast<std::size_t>((dim_size % stride[i] !=
-                                                                           0)) // ceil uint divide
-                        : dim_size / stride[i];                                // floor divide
-                output_lens.push_back(len + 1);
-            }
+            std::size_t padding_factor = 2 * padding[i];
+            if(padding.size() == 2 * kdims)
+                padding_factor = padding[i] + padding[i + kdims];
+            assert(input_lens[i + 2] + padding_factor >= lengths[i]);
+            std::size_t dim_size = input_lens[i + 2] + padding_factor - lengths[i];
+            std::size_t len =
+                (ceil_mode)
+                    ? dim_size / stride[i] +
+                          static_cast<std::size_t>((dim_size % stride[i] != 0)) // ceil uint divide
+                    : dim_size / stride[i];                                     // floor divide
+            output_lens.push_back(len + 1);
         }
         return output_lens;
     }
@@ -134,19 +126,19 @@ struct pooling
             {
                 for(size_t i = 0; i < kdims; ++i)
                 {
-                    output_dyn_dims.push_back(shape::dynamic_dimension{1, 1, 1});
+                    output_dyn_dims.push_back(shape::dynamic_dimension{1, 1});
                 }
                 return {input.type(), output_dyn_dims};
             }
             else
             {
+                // does not compute for optimals
                 auto min_spatial_dims = calc_spatial_dim_out(input.min_lens(), kdims);
                 auto max_spatial_dims = calc_spatial_dim_out(input.max_lens(), kdims);
-                auto opt_spatial_dims = calc_spatial_dim_out(input.opt_lens(), kdims);
                 for(size_t i = 0; i < kdims; ++i)
                 {
-                    output_dyn_dims.push_back(shape::dynamic_dimension{
-                        min_spatial_dims[i], max_spatial_dims[i], opt_spatial_dims[i]});
+                    output_dyn_dims.push_back(
+                        shape::dynamic_dimension{min_spatial_dims[i], max_spatial_dims[i], {}});
                 }
                 return {input.type(), output_dyn_dims};
             }

src/include/migraphx/op/quant_convolution.hpp

@@ -25,8 +25,10 @@
 #define MIGRAPHX_GUARD_OPERATORS_QUANT_CONVOLUTION_HPP
 
 #include <migraphx/op/common.hpp>
+#include <migraphx/argument.hpp>
 #include <migraphx/check_shapes.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/convolution.hpp>
 #include <migraphx/value.hpp>
 #include <cmath>
 #include <utility>
@@ -114,6 +116,17 @@ struct quant_convolution
         check_attribute_size();
         return stride.size();
     }
+
+    argument compute(shape output_shape, std::vector<argument> args) const
+    {
+        argument result{output_shape};
+        result.visit([&](auto output) {
+            visit_all(args[0], args[1])([&](auto input, auto weights) {
+                migraphx::convolution(output, input, weights, padding, stride, group);
+            });
+        });
+        return result;
+    }
 };
 
 } // namespace op

src/include/migraphx/op/quantizelinear.hpp

@@ -40,7 +40,11 @@ struct quantizelinear
     std::string name() const { return "quantizelinear"; }
     shape compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.same_dims();
+        check_shapes{inputs, *this}.same_dims().has(2, 3);
+        if(inputs[0].type() != inputs[1].type())
+        {
+            MIGRAPHX_THROW("QUANTIZELINEAR: Scales and input must be the same type");
+        }
         if(inputs.size() == 3)
         {
             return {inputs[2].type(), inputs[0].lens(), inputs[0].strides()};
@@ -61,17 +65,15 @@ struct quantizelinear
 
         argument result{output_shape};
         visit_all(result, y_zero_point)([&](auto output, auto zero_pts) {
-            x.visit([&](auto input) {
-                y_scale.visit([&](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])) +
-                                            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));
-                    });
+            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])) +
+                                        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));
                 });
             });
         });

src/include/migraphx/op/reduce_op.hpp

@@ -91,7 +91,7 @@ struct reduce_op : op_name<Derived>
     {
         value normalize;
         normalize["axes"] = value::array{normalize_attribute::include_min};
-        return {{"normalize_axes", normalize}};
+        return {{"normalize_axes", normalize}, {"reduce", true}};
     }
 
     std::vector<int64_t> tune_axes(std::size_t n_dim) const
@@ -123,9 +123,7 @@ struct reduce_op : op_name<Derived>
             auto tuned_axes      = tune_axes(output_dyn_dims.size());
             for(const auto& axis : tuned_axes)
             {
-                // At the time of writing, there's no functional difference between
-                // optimum of 0 (no opt) or 1.
-                output_dyn_dims[axis] = {1, 1, 0};
+                output_dyn_dims[axis] = {1, 1};
             }
 
             return shape{s.type(), output_dyn_dims};

src/include/migraphx/op/slice.hpp

@@ -111,16 +111,15 @@ struct slice
         // For a static shape, old_lens will be adjusted to a new size
         // for those axes that are sliced.
         // For dynamic shape, the adjusted old_lens become the new max values,
-        // while updating the old mins and opts if possible.
+        // while updating the old mins and optimals if possible.
         std::vector<std::size_t> new_mins;
-        std::vector<std::size_t> new_opts;
         std::vector<std::size_t> old_lens;
         std::vector<std::size_t> old_strides;
+        // Doesn't handle optimals
         if(input_shape.dynamic())
         {
             old_lens = input_shape.max_lens();
             new_mins = input_shape.min_lens();
-            new_opts = input_shape.opt_lens();
         }
         else
         {
@@ -146,17 +145,11 @@ struct slice
                 std::size_t sliced_min_length = ends[i] - starts[i];
                 // if the slice size is smaller than maxes but larger than mins
                 new_mins[axis] = std::min(sliced_min_length, new_mins[axis]);
-
-                auto sliced_opt_length = ends[i] - starts[i];
-                if(new_opts[axis] != 0)
-                    new_opts[axis] = sliced_opt_length;
-                if(new_opts[axis] < new_mins[axis] or new_opts[axis] > new_lens[axis])
-                    new_opts[axis] = 0;
             }
         }
         if(input_shape.dynamic())
         {
-            return shape{t, new_mins, new_lens, new_opts};
+            return shape{t, new_mins, new_lens, {}};
         }
         else
         {

src/include/migraphx/op/unsqueeze.hpp

@@ -81,7 +81,7 @@ struct unsqueeze
             {
                 if(std::find(axes.begin(), axes.end(), i) != axes.end())
                 {
-                    dyn_dims.push_back({1, 1, 0});
+                    dyn_dims.push_back({1, 1});
                 }
                 else
                 {

src/include/migraphx/pass_manager.hpp

@@ -39,6 +39,7 @@ struct module_pass_manager
     virtual module& get_module()                           = 0;
     virtual module* create_module(const std::string& name) = 0;
     virtual module* get_common_parent()                    = 0;
+    virtual module* get_root_module()                      = 0;
     virtual void run_pass(const pass& p)                   = 0;
 
     protected:

src/include/migraphx/promote_literals.hpp

+/*
+ * 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.
+ */
+#ifndef MIGRAPHX_GUARD_RTGLIB_PROMOTE_LITERALS_HPP
+#define MIGRAPHX_GUARD_RTGLIB_PROMOTE_LITERALS_HPP
+
+#include <string>
+#include <migraphx/pass_manager.hpp>
+#include <migraphx/config.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+/**
+ * Replace literals in submodules with literals in the root module.
+ * Intended to allow for reuse of the literals between submodules.
+ */
+struct promote_literals
+{
+    std::string name() const { return "promote_literals"; }
+    void apply(module_pass_manager&) const;
+};
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/serialize.hpp

@@ -188,7 +188,8 @@ auto from_value_impl(rank<3>, const value& v, T& x)
 }
 
 template <class T>
-auto from_value_impl(rank<4>, const value& v, T& x) -> decltype(x.insert(*x.begin()), void())
+auto from_value_impl(rank<4>, const value& v, T& x)
+    -> decltype(x.insert(*x.begin()), std::declval<typename T::mapped_type>(), void())
 {
     x.clear();
     for(auto&& e : v)

src/include/migraphx/shape.hpp

@@ -29,10 +29,12 @@
 #include <ostream>
 #include <numeric>
 #include <memory>
+#include <set>
 
 #include <migraphx/functional.hpp>
 #include <migraphx/errors.hpp>
 #include <migraphx/half.hpp>
+#include <migraphx/serialize.hpp>
 #include <migraphx/config.hpp>
 
 namespace migraphx {
@@ -87,12 +89,12 @@ struct shape
     {
         std::size_t min = 0;
         std::size_t max = 0;
-        std::size_t opt = 0;
+        std::set<std::size_t> optimals{};
 
         template <class Self, class F>
         static auto reflect(Self& self, F f)
         {
-            return pack(f(self.min, "min"), f(self.max, "max"), f(self.opt, "opt"));
+            return pack(f(self.min, "min"), f(self.max, "max"), f(self.optimals, "optimals"));
         }
 
         bool is_fixed() const;
@@ -132,11 +134,12 @@ struct shape
 
     shape(type_t t, std::vector<dynamic_dimension> dims);
 
-    // Construct a dynamic shape from three sets of lengths (of the same rank)
+    // Construct a dynamic shape from vectors of mins, maxes, and optimals.
+    // optimals_list is a vector of optimals that corresponds to each min and max.
     shape(type_t t,
           std::vector<std::size_t> mins,
           std::vector<std::size_t> maxes,
-          std::vector<std::size_t> opts);
+          std::vector<std::set<std::size_t>> optimals_list);
 
     template <class Range>
     shape(type_t t, const Range& l) : shape(t, std::vector<std::size_t>(l.begin(), l.end()))
@@ -186,21 +189,21 @@ struct shape
 
     /*!
      * Minimum lengths for dynamic shape.
-     * lens() for fixed shape.
+     * lens() for static shape.
      */
     std::vector<std::size_t> min_lens() const;
 
     /*!
      * Maximum lengths for dynamic shape.
-     * lens() for fixed shape.
+     * lens() for static shape.
      */
     std::vector<std::size_t> max_lens() const;
 
     /*!
      * Optimum lengths for dynamic shape.
-     * lens() for fixed shape.
+     * Empty for static shape.
      */
-    std::vector<std::size_t> opt_lens() const;
+    std::vector<std::set<std::size_t>> opt_lens() const;
 
     /// Map multiple indices to space index
     std::size_t index(std::initializer_list<std::size_t> l) const;
@@ -253,9 +256,12 @@ struct shape
 
     shape with_type(type_t t) const;
 
-    // convert the shape to an equivalent dynamic shape
+    // convert the shape to an equivalent dynamic shape with empty optimals
     shape to_dynamic() const;
 
+    // convert the shape to a static one setting any non-fixed dynamic_dimensions to x
+    shape to_static(std::size_t x) const;
+
     friend bool operator==(const shape& x, const shape& y);
     friend bool operator!=(const shape& x, const shape& y);
     friend std::ostream& operator<<(std::ostream& os, const shape& x);

src/include/migraphx/split_single_dyn_dim.hpp

+/*
+ * 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.
+ */
+#ifndef MIGRAPHX_GUARD_RTGLIB_SPLIT_SINGLE_DYN_DIM_HPP
+#define MIGRAPHX_GUARD_RTGLIB_SPLIT_SINGLE_DYN_DIM_HPP
+
+#include <string>
+#include <migraphx/pass_manager.hpp>
+#include <migraphx/instruction_ref.hpp>
+#include <migraphx/config.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+/**
+ * Split dynamic dimension over submodules if exactly one dimension in the parameter list is
+ * dynamic.
+ */
+struct split_single_dyn_dim
+{
+    std::string name() const { return "split_single_dyn_dim"; }
+    void apply(module_pass_manager&) const;
+};
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/module.cpp

@@ -595,6 +595,14 @@ std::vector<shape> module::get_output_shapes() const
     }
 }
 
+std::vector<instruction_ref> module::get_returns() const
+{
+    auto last = std::prev(this->end());
+    if(last->name() == "@return")
+        return last->inputs();
+    return {last};
+}
+
 instruction_ref module::validate() const
 {
     return std::find_if(