Merge branch 'rand_uniform' into multinomial_parse_merge_random

src/include/migraphx/op/random_seed.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_OPERATORS_RANDOM_SEED_HPP
+#define MIGRAPHX_GUARD_OPERATORS_RANDOM_SEED_HPP
+
+#include <migraphx/check_shapes.hpp>
+#include <migraphx/argument.hpp>
+#include <random>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace op {
+
+/**
+ *    Generates a random seed for the use of random number generators.  Generating the seed
+ * at runtime guarantees there will be a different random sequence on every execution.
+ * This operation has no inputs or attributes, and outputs an unsigned integer tensor with
+ * a single value.
+ */
+struct random_seed
+{
+    shape::type_t dtype = shape::type_t::uint64_type;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.dtype, "dtype"));
+    }
+
+    std::string name() const { return "random_seed"; }
+    shape compute_shape(const std::vector<shape>& inputs) const
+    {
+        check_shapes{inputs, *this}.has(0);
+        return shape{dtype};
+    }
+
+    argument compute(const shape& output_shape, const std::vector<argument>&) const
+    {
+        argument result(output_shape);
+
+        result.visit([&](auto output) { output.front() = std::random_device{}(); });
+        return result;
+    }
+};
+
+} // namespace op
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/op/random_uniform.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.
+ */
+
+/**
+ * Random Uniform distribution operator.  Given a shape, populate it with random
+ * values.  Calls to random_uniform using the same randomization seed will
+ * always generate the same pseudo-random sequence.  Seed can
+ * be given as a runtime argument containing a single value, or a compile-time
+ * attribute.
+ *
+ *      Inputs:   (1) randomization seed (any type is allowed)
+ *                (2) the shape of the set to be populated.
+ *
+ *
+ *      Attributes:  none
+ *
+ *      Output:   Same shape.
+ *
+ */
+#ifndef MIGRAPHX_GUARD_OPERATORS_RANDOM_UNIFORM_HPP
+#define MIGRAPHX_GUARD_OPERATORS_RANDOM_UNIFORM_HPP
+
+#include <migraphx/check_shapes.hpp>
+#include <migraphx/argument.hpp>
+#include <random>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace op {
+
+/**
+ * random_uniform populates the passed shape with random numbers, in a uniform
+ * distribution.  Range for floating-point data types is (0, 1);
+ * for integer types it is [0, <max value for the type>]
+ *
+ *   Input 1:  seed
+ *   Input 2:  output shape
+ */
+struct random_uniform
+{
+    // The random_uniform operation needs the random number generator seed
+    // to be passed as a runtime input.
+
+    std::string name() const { return "random_uniform"; }
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs, *this, true}.has(2);
+
+        return inputs.at(1);
+    }
+
+    argument compute(const shape&, std::vector<argument> args) const
+    {
+        // Output goes into the passed buffer, not the shape output.
+        auto result = args[1];
+
+        uint64_t local_seed = args[0].at<uint64_t>(0);
+        std::mt19937 gen(local_seed);
+
+        result.visit([&](auto output) {
+            using type = typename decltype(output)::value_type;
+            if constexpr(std::is_integral<type>{})
+            {
+                // default range for all integer types is (0,
+                // std::uniform_int_distribution<type>::max()).
+                // Todo:  enable different ranges
+                std::uniform_int_distribution<type> dis;
+                std::generate(output.begin(), output.end(), [&] { return dis(gen); });
+            }
+            else
+            {
+                // default real distribution type is double with range (0, 1);
+                std::uniform_real_distribution<> dis;
+                std::generate(output.begin(), output.end(), [&] { return dis(gen); });
+            }
+        });
+        return result;
+    }
+
+    std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 1; }
+};
+
+} // namespace op
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/op/slice.hpp

@@ -27,19 +27,34 @@
 #include <migraphx/check_shapes.hpp>
 #include <migraphx/argument.hpp>
 #include <migraphx/config.hpp>
-#include <migraphx/dyn_output.hpp>
 #include <migraphx/value.hpp>
+#include <migraphx/dyn_output.hpp>
 #include <migraphx/op/normalize_attribute.hpp>
+#include <migraphx/normalize_attributes.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 namespace op {
 
+/**
+ * Slice operator that accepts variable axes, starts and ends.
+ *
+ * Attributes:
+ * axes: constant axes to slice over (optional)
+ * starts: constant slice starting indices (optional)
+ * ends: constant slice ending indices (optional)
+ *
+ * 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_axes: axes to slice over (optional, static shape)
+ */
 struct slice
 {
-    std::vector<int64_t> axes;
-    std::vector<int64_t> starts;
-    std::vector<int64_t> ends;
+    std::vector<int64_t> axes{};
+    std::vector<int64_t> starts{};
+    std::vector<int64_t> ends{};
 
     template <class Self, class F>
     static auto reflect(Self& self, F f)
@@ -48,8 +63,8 @@ struct slice
     }
 
     /**
-     * Ensure that attribute vectors axes, starts, and ends are all the same size and values are in
-     * limits.
+     * Ensure that attribute vectors axes, starts, and ends are all the same size and values are
+     * within limits.
      */
     value attributes() const
     {
@@ -70,6 +85,90 @@ struct slice
 
     std::string name() const { return "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
+     * object's axes attribute. Assumes in_starts and in_ends are normalized; in_axes are valid.
+     */
+    template <class A, class B>
+    std::vector<std::size_t>
+    lens_calc(const std::vector<std::size_t>& lengths, A in_starts, A in_ends, B in_axes) const
+    {
+        auto new_lens = lengths;
+        for(std::size_t i = 0; i < in_axes.size(); ++i)
+        {
+            auto axis      = in_axes[i];
+            new_lens[axis] = in_ends[i] - in_starts[i];
+        }
+        return new_lens;
+    }
+
+    shape normalize_compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs, *this, true}.has(1, 3, 4);
+        auto input_shape = inputs[0];
+        if(inputs.size() == 1)
+        {
+            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();
+               }))
+            {
+                MIGRAPHX_THROW("SLICE: slicing is not allowed on non-fixed dynamic input axis ");
+            }
+            if(input_shape.dynamic())
+            {
+                return shape{t,
+                             lens_calc(input_shape.min_lens(), starts, ends, axes),
+                             lens_calc(input_shape.max_lens(), starts, ends, axes),
+                             {}};
+            }
+            else
+            {
+                return shape{
+                    t, lens_calc(input_shape.lens(), starts, ends, axes), input_shape.strides()};
+            }
+        }
+        else
+        {
+            // 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(inputs[1].lens().at(0) != axes.size())
+                {
+                    MIGRAPHX_THROW("SLICE: inputs starts and ends do not have the same dimension "
+                                   "as the axes attribute");
+                }
+                std::for_each(axes.cbegin(), axes.cend(), [&](const auto& axis) {
+                    dds.at(axis) = {0, dds.at(axis).max};
+                });
+            }
+            else
+            {
+                // if axes is an input, then all the output dimensions could be 0 to the max value
+                std::transform(dds.begin(), dds.end(), dds.begin(), [](auto dd) {
+                    return shape::dynamic_dimension{0, dd.max};
+                });
+            }
+            return shape{input_shape.type(), dds};
+        }
+    }
+
+    /**
+     * Calculates the starting offset for the sliced tensor.
+     * Used in compute when only data input and all other information are in the attributes.
+     *
+     * \param s static input shape
+     */
     auto compute_offset(const shape& s) const
     {
         const std::vector<std::size_t>& lens    = s.lens();
@@ -90,80 +189,131 @@ struct slice
                 offset += starts[axis] * strides[axis];
             }
         }
-        return offset;
+        return offset * s.type_size();
     }
 
-    shape normalize_compute_shape(std::vector<shape> inputs) const
+    /**
+     * Calculates the starting offset for the sliced tensor (for aliasing).
+     * Used when the starts and/or the axes are inputs.
+     *
+     * \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
     {
-        check_shapes{inputs, *this, true}.has(1);
-        auto input_shape = inputs[0];
-        auto t           = input_shape.type();
-
-        // TODO:  When support for dynamic shapes is added to normalize_attributes,
-        //  remove this restriction.
-        if(input_shape.dynamic() and std::any_of(axes.begin(), axes.end(), [&](auto axis) {
-               return not input_shape.dyn_dims()[axis].is_fixed();
-           }))
+        auto ret = 0;
+        for(std::size_t i = 0; i < ax_vec.size(); ++i)
         {
-            MIGRAPHX_THROW("SLICE: slicing is not allowed on non-fixed dynamic input axis ");
+            auto axis = ax_vec[i];
+            ret += input_starts[i] * s.strides().at(axis);
         }
+        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.
+     */
+    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
+    {
+        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}};
+    }
 
-        // 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 optimals if possible.
-        std::vector<std::size_t> new_mins;
-        std::vector<std::size_t> old_lens;
-        std::vector<std::size_t> old_strides;
-        // Doesn't handle optimals
-        if(input_shape.dynamic())
+    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())
         {
-            old_lens = input_shape.max_lens();
-            new_mins = input_shape.min_lens();
+        case 1: {
+            std::size_t offset = compute_offset(input_shape);
+            return {dyn_out.computed_shape, [=] { return input.data() + offset; }};
         }
-        else
-        {
-            old_lens = input_shape.lens();
-            // For static shape (including during eval step after a dynamic input) the strides are
-            // indexed into the pre-slice array, so they are larger than the apparent size of the
-            // resulting shape.
-            old_strides = input_shape.strides();
+        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; }};
         }
-
-        std::vector<std::size_t> new_lens = old_lens;
-        for(std::size_t i = 0; i < axes.size(); i++)
-        {
-            auto axis            = axes[i];
-            size_t sliced_length = ends[i] - starts[i];
-            // A Numpy indexing convention: a slice size larger than the actual dimension
-            // is legal and the "ends" value is clipped to the axis size
-            new_lens[axis] = std::min(new_lens[axis], sliced_length);
-            if(input_shape.dynamic())
-            {
-                // TODO: when non-fixed shape slicing is allowed, this will be different than
-                // sliced_length, making use of TBD start/end values.
-                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]);
-            }
+        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()};
+                });
+            return {calc_shape, [=] { return input.data() + offset; }};
         }
-        if(input_shape.dynamic())
-        {
-            return shape{t, new_mins, new_lens, {}};
+        default: {
+            // Should never get here; covering in case some code change occurs
+            MIGRAPHX_THROW("SLICE: invalid number of inputs");
         }
-        else
-        {
-            return shape{t, new_lens, old_strides};
         }
     }
 
-    argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
-    {
-        auto input = args[0];
-
-        auto offset = compute_offset(input.get_shape()) * dyn_out.computed_shape.type_size();
-        return {dyn_out.computed_shape, [=] { return input.data() + offset; }};
-    }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 

src/include/migraphx/permutation.hpp

@@ -66,6 +66,10 @@ MIGRAPHX_EXPORT std::vector<int64_t> invert_permutation(const std::vector<int64_
 MIGRAPHX_EXPORT std::vector<int64_t> find_permutation(const shape& s);
 MIGRAPHX_EXPORT std::vector<int64_t> find_permutation(const std::vector<shape>& shapes);
 
+/// Normalize the shapes so the order of dimensions will be in the order it is
+/// in memory as much as possible.
+MIGRAPHX_EXPORT std::vector<shape> normalize_permutation(const std::vector<shape>& shapes);
+
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx
 

src/instruction.cpp

@@ -64,10 +64,7 @@ void instruction::replace(const shape& r)
         result = r;
         for(auto&& ins : output)
         {
-            if(ins->name() == "@return")
-                continue;
-
-            assert(ins->name().front() != '@');
+            assert(ins->name() == "@return" or ins->name().front() != '@');
             ins->recompute_shape();
         }
     }
@@ -122,10 +119,6 @@ bool instruction::valid() const
     {
         computed = result;
     }
-    else if(op.name() == "@return")
-    {
-        computed = {};
-    }
     else
     {
         try
@@ -145,6 +138,7 @@ bool instruction::valid() const
 }
 
 shape instruction::get_shape() const { return result; }
+
 const literal& instruction::get_literal() const
 {
     assert(op.name() == "@literal");
@@ -395,7 +389,7 @@ void instruction::print(std::ostream& os,
     if(not ins->module_inputs().empty())
     {
         std::string delim = ", [";
-        for(auto&& mod_arg : ins->module_inputs())
+        for(const const_module_ref& mod_arg : ins->module_inputs())
         {
             os << delim << mod_arg->name();
             delim = ", ";

src/memory_coloring.cpp

@@ -23,9 +23,9 @@
  */
 #include <migraphx/memory_coloring.hpp>
 #include <migraphx/module.hpp>
-#include <migraphx/operators.hpp>
 #include <migraphx/instruction.hpp>
 #include <migraphx/iterator_for.hpp>
+#include <migraphx/make_op.hpp>
 #include <migraphx/functional.hpp>
 #include <migraphx/algorithm.hpp>
 #include <migraphx/ranges.hpp>
@@ -382,7 +382,8 @@ void memory_coloring::apply(module& m) const
         auto s             = ins->get_shape();
         std::size_t offset = seg.first * alignment;
         assert(offset < n);
-        m.replace_instruction(ins, op::load{s, offset}, mem);
+        m.replace_instruction(
+            ins, make_op("load", {{"shape", to_value(s)}, {"offset", offset}}), mem);
     }
 
     // Replace zero allocation
@@ -391,7 +392,8 @@ void memory_coloring::apply(module& m) const
         if(ins->name() != allocation_op)
             continue;
         assert(ins->get_shape().bytes() == 0);
-        m.replace_instruction(ins, op::load{ins->get_shape(), 0}, mem);
+        m.replace_instruction(
+            ins, make_op("load", {{"shape", to_value(ins->get_shape())}, {"offset", 0}}), mem);
     }
 
     // Remove scratch parameter if its not used

src/module.cpp

@@ -460,11 +460,11 @@ instruction_ref module::add_parameter(std::string name, shape s)
 
 instruction_ref module::add_return(std::vector<instruction_ref> args)
 {
-    impl->push_back({builtin::returns{}, {}, std::move(args)});
+    shape instr_shape = compute_shape(builtin::returns{}, args);
+    impl->push_back({builtin::returns{}, instr_shape, std::move(args)});
     auto result = std::prev(impl->instructions.end());
     instruction::backreference(result);
     assert(result->valid(begin()));
-
     return result;
 }
 
@@ -873,12 +873,11 @@ module::print_py(std::ostream& os,
             if(ins->name() == "@literal")
             {
                 os << mname << ".add_literal(";
-                bool use_abs = false;
-                ins->get_literal().visit([&](auto v) {
-                    use_abs = std::none_of(v.begin(), v.end(), [](auto x) { return x < 0; });
-                });
+                const bool use_abs = false;
                 // Disable abs for now
-                use_abs = false;
+                // ins->get_literal().visit([&](auto v) {
+                //     use_abs = std::none_of(v.begin(), v.end(), [](auto x) { return x < 0; });
+                // });
                 if(use_abs)
                     os << "migraphx.abs_literal(";
                 os << "migraphx.generate_argument(";
@@ -1011,9 +1010,17 @@ std::vector<module_ref> module::get_sub_modules(bool shallow) const
 
 module& module::sort()
 {
+    auto implicit_deps = calc_implicit_deps();
     fix([&](auto self, auto ins) {
         this->move_instruction(ins, this->begin());
-        for(auto child : ins->inputs())
+        auto ins_inputs = ins->inputs();
+        if(implicit_deps.find(ins) != implicit_deps.end())
+        {
+            auto ins_implict_inputs = implicit_deps.at(ins);
+            ins_inputs.insert(
+                ins_inputs.end(), ins_implict_inputs.begin(), ins_implict_inputs.end());
+        }
+        for(auto child : ins_inputs)
         {
             if(not contains(this->impl->instructions, child))
             {

src/normalize_attributes.cpp

@@ -49,6 +49,10 @@ auto tune_attribute(const std::vector<int64_t>& vec,
                     Message m)
 {
     std::vector<int64_t> result(vec);
+    if(result.empty())
+    {
+        return result;
+    };
     int64_t n_rank                                 = input_shape.ndim();
     std::vector<op::normalize_attribute> vec_attrs = val.to_vector<op::normalize_attribute>();
     if(contains(vec_attrs, op::normalize_attribute::use_output))
@@ -251,5 +255,22 @@ bool normalize_attributes(operation& op, const shape& input_shape)
     return tuned;
 }
 
+std::vector<int64_t> normalize_axes(const std::vector<int64_t>& axes,
+                                    const shape& input_shape,
+                                    const value& attr_val,
+                                    const std::string& prefix)
+{
+    return tune_attribute(axes, {}, attr_val, input_shape, [&] { return prefix; });
+}
+
+std::vector<int64_t> normalize_indices(const std::vector<int64_t>& indices,
+                                       const std::vector<int64_t>& axes,
+                                       const shape& input_shape,
+                                       const value& attr_val,
+                                       const std::string& prefix)
+{
+    return tune_attribute(indices, axes, attr_val, input_shape, [&] { return prefix; });
+}
+
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

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

@@ -117,6 +117,7 @@ struct onnx_parser
     parse_graph(module* mod, const onnx::GraphProto& graph, bool inlining = false);
     literal parse_value(const onnx::AttributeProto& attr) const;
     literal parse_tensor(const onnx::TensorProto& t) const;
+    shape parse_type(const onnx::TypeProto& t) const;
     shape parse_type(const onnx::TypeProto& t, const std::vector<std::size_t>& input_dims) const;
 };
 

src/onnx/onnx_parser.cpp

@@ -357,10 +357,9 @@ parse_inputs(const onnx_parser& parser,
             }
 
             shape s;
-            std::vector<std::size_t> dims;
             if(parser.map_input_dims.count(name) > 0)
             {
-                dims = parser.map_input_dims.at(name);
+                std::vector<std::size_t> dims = parser.map_input_dims.at(name);
                 s    = parser.parse_type(input.type(), dims);
             }
             else if(parser.map_dyn_input_dims.count(name) > 0)
@@ -370,7 +369,7 @@ parse_inputs(const onnx_parser& parser,
             }
             else
             {
-                s = parser.parse_type(input.type(), dims);
+                s = parser.parse_type(input.type());
             }
             mod_insts[name] = mod->add_parameter(name, s);
         }
@@ -553,14 +552,9 @@ literal onnx_parser::parse_tensor(const onnx::TensorProto& t) const
     }
     MIGRAPHX_THROW("PARSE_TENSOR: Invalid tensor type");
 }
-shape onnx_parser::parse_type(const onnx::TypeProto& t,
-                              const std::vector<std::size_t>& input_dims) const
+shape onnx_parser::parse_type(const onnx::TypeProto& t) const
 {
     shape::type_t shape_type = get_type(t.tensor_type().elem_type());
-    if(not input_dims.empty())
-    {
-        return {shape_type, input_dims};
-    }
 
     std::vector<shape::dynamic_dimension> dynamic_dims;
     auto&& tensor_dims = t.tensor_type().shape().dim();
@@ -590,6 +584,15 @@ shape onnx_parser::parse_type(const onnx::TypeProto& t,
     return shape_from_dyn_dims(shape_type, dynamic_dims);
 }
 
+shape onnx_parser::parse_type(const onnx::TypeProto& t,
+                              const std::vector<std::size_t>& input_dims) const
+{
+    shape::type_t shape_type = get_type(t.tensor_type().elem_type());
+    if(input_dims.empty())
+        return {shape_type};
+    return {shape_type, input_dims};
+}
+
 shape::type_t get_type(int dtype)
 {
     switch(dtype)

src/onnx/parse_constant_of_shape.cpp

@@ -55,9 +55,6 @@ struct parse_constant_of_shape : op_parser<parse_constant_of_shape>
             l_val = literal({shape::float_type, {1}, {0}}, {0.0f});
         }
 
-        // input is empty, output is a scalar
-        auto type = l_val.get_shape().type();
-
         if(args.empty())
         {
             MIGRAPHX_THROW("ConstantOfShape : must have 1 input!");
@@ -65,6 +62,8 @@ struct parse_constant_of_shape : op_parser<parse_constant_of_shape>
         else
         {
             migraphx::shape s;
+            // input is empty, output is a scalar
+            auto type = l_val.get_shape().type();
             // empty input tensor, output is a scalar
             if(args[0]->get_shape().elements() == 0)
             {

src/onnx/parse_randomuniform_ops.cpp

@@ -96,7 +96,7 @@ struct parse_randomuniform_ops : op_parser<parse_randomuniform_ops>
         if(contains(info.attributes, "seed"))
             gen.seed(info.attributes.at("seed").f());
 
-        std::uniform_real_distribution<> d(high, low);
+        std::uniform_real_distribution<> d(low, high);
         std::vector<double> rand_vals(out_shape.elements());
         std::generate(rand_vals.begin(), rand_vals.end(), [&]() { return d(gen); });
 

src/onnx/parse_slice.cpp

@@ -34,16 +34,65 @@ namespace onnx {
 
 struct parse_slice : op_parser<parse_slice>
 {
+
     std::vector<op_desc> operators() const { return {{"Slice"}}; }
 
+    struct slice_desc
+    {
+        op::slice op;
+        std::vector<instruction_ref> op_args;
+        std::vector<int64_t> steps;
+        std::vector<int64_t> raxes;
+
+        void always_insert(instruction_ref arg) { op_args.insert(op_args.begin(), arg); }
+
+        std::vector<int64_t> insert(instruction_ref arg)
+        {
+            std::vector<int64_t> result;
+            migraphx::argument arg_value = arg->eval();
+            if(arg_value.empty())
+            {
+                op_args.insert(op_args.begin(), arg);
+            }
+            else
+            {
+                arg_value.visit([&](auto s) { result.assign(s.begin(), s.end()); });
+            }
+            return result;
+        }
+    };
+
     instruction_ref parse(const op_desc& /*opd*/,
                           const onnx_parser& parser,
-                          onnx_parser::node_info info,
-                          std::vector<instruction_ref> args) const
+                          const onnx_parser::node_info& info,
+                          const std::vector<instruction_ref>& args) const
     {
-        op::slice op;
+        auto sd  = construct_slice_desc(parser, info, args);
+        auto ins = info.add_instruction(sd.op, sd.op_args);
+        if(not sd.raxes.empty())
+        {
+            ins = info.add_instruction(make_op("reverse", {{"axes", sd.raxes}}), ins);
+        }
+        // If any steps are other than default 1, add a "steps" op
+        if(std::any_of(sd.steps.begin(), sd.steps.end(), [](auto s) { return std::abs(s) != 1; }))
+        {
+            std::vector<int64_t> nsteps;
+            std::transform(sd.steps.begin(),
+                           sd.steps.end(),
+                           std::back_inserter(nsteps),
+                           [](auto s) { return std::abs(s); });
+            return ins = info.add_instruction(
+                       make_op("step", {{"axes", sd.op.axes}, {"steps", nsteps}}), ins);
+        }
+        else
+            return ins;
+    }
 
-        std::vector<int64_t> steps;
+    slice_desc construct_slice_desc(const onnx_parser& parser,
+                                    onnx_parser::node_info info,
+                                    std::vector<instruction_ref> args) const
+    {
+        slice_desc sd;
 
         // slice can have up to 5 inputs, we first check the 5th one
         // to decide whether MIGRAPHX can handle this slice.
@@ -51,89 +100,73 @@ struct parse_slice : op_parser<parse_slice>
         {
             migraphx::argument step_arg = args.back()->eval();
             check_arg_empty(step_arg, "PARSE_SLICE: cannot handle variable steps for slice");
-            step_arg.visit([&](auto s) { steps.assign(s.begin(), s.end()); });
+            step_arg.visit([&](auto s) { sd.steps.assign(s.begin(), s.end()); });
         }
 
         if(args.size() >= 4)
         {
-            migraphx::argument axes_arg = args.at(3)->eval();
-            check_arg_empty(axes_arg, "PARSE_SLICE: cannot handle variable axes for slice");
-            axes_arg.visit([&](auto s) { op.axes.assign(s.begin(), s.end()); });
+            sd.op.axes = sd.insert(args.at(3));
         }
         else if(contains(info.attributes, "axes"))
         {
             literal s = parser.parse_value(info.attributes.at("axes"));
-            s.visit([&](auto v) { copy(v, std::back_inserter(op.axes)); });
+            s.visit([&](auto v) { copy(v, std::back_inserter(sd.op.axes)); });
         }
 
         if(args.size() >= 3)
         {
-            migraphx::argument end_arg = args.at(2)->eval();
-            check_arg_empty(end_arg, "PARSE_SLICE: cannot handle variable ends for slice");
-            end_arg.visit([&](auto s) { op.ends.assign(s.begin(), s.end()); });
+            sd.op.ends = sd.insert(args.at(2));
         }
         else if(contains(info.attributes, "ends"))
         {
             literal s = parser.parse_value(info.attributes.at("ends"));
-            s.visit([&](auto v) { copy(v, std::back_inserter(op.ends)); });
+            s.visit([&](auto v) { copy(v, std::back_inserter(sd.op.ends)); });
         }
 
         if(args.size() >= 2)
         {
-            migraphx::argument start_arg = args.at(1)->eval();
-            check_arg_empty(start_arg, "PARSE_SLICE: cannot handle variable starts for slice");
-            start_arg.visit([&](auto s) { op.starts.assign(s.begin(), s.end()); });
+            sd.op.starts = sd.insert(args.at(1));
         }
         else if(contains(info.attributes, "starts"))
         {
             literal s = parser.parse_value(info.attributes.at("starts"));
-            s.visit([&](auto v) { copy(v, std::back_inserter(op.starts)); });
+            s.visit([&](auto v) { copy(v, std::back_inserter(sd.op.starts)); });
         }
 
+        // data input argument
+        sd.always_insert(args.at(0));
+
         // If axes arg is not given, the default is all of them.
-        if(op.axes.empty())
+        if(sd.op.axes.empty() and sd.op_args.size() < 3)
         {
             std::vector<int64_t> axes(args[0]->get_shape().ndim());
             std::iota(axes.begin(), axes.end(), int64_t{0});
-            op.axes = axes;
+            sd.op.axes = axes;
         }
 
-        std::vector<int64_t> raxes;
+        if(not sd.steps.empty())
+        {
+            if(sd.op.starts.empty() or sd.op.ends.empty())
+                MIGRAPHX_THROW("PARSE_SLICE: steps and variable starts and ends is not supported");
+            if(sd.op.axes.empty())
+                MIGRAPHX_THROW("PARSE_SLICE: steps and variable axes is not supported");
+        }
 
-        assert(steps.empty() or steps.size() == op.axes.size());
-        assert(op.axes.size() == op.starts.size());
-        assert(op.axes.size() == op.ends.size());
+        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(steps.size()))
+        for(auto i : range(sd.steps.size()))
         {
-            if(steps[i] >= 0)
+            if(sd.steps[i] >= 0)
                 continue;
-            op.starts[i] += 1;
-            if(op.starts[i] == 0)
-                op.starts[i] = INT_MAX;
-            op.ends[i] += 1;
-            raxes.push_back(op.axes[i]);
-            std::swap(op.starts[i], op.ends[i]);
-        }
-
-        auto ins = info.add_instruction(op, args[0]);
-        if(not raxes.empty())
-        {
-            ins = info.add_instruction(make_op("reverse", {{"axes", raxes}}), ins);
+            sd.op.starts[i] += 1;
+            if(sd.op.starts[i] == 0)
+                sd.op.starts[i] = INT_MAX;
+            sd.op.ends[i] += 1;
+            sd.raxes.push_back(sd.op.axes[i]);
+            std::swap(sd.op.starts[i], sd.op.ends[i]);
         }
-        // If any steps are other than default 1, add a "steps" op
-        if(std::any_of(steps.begin(), steps.end(), [](auto s) { return std::abs(s) != 1; }))
-        {
-            std::vector<int64_t> nsteps;
-            std::transform(steps.begin(), steps.end(), std::back_inserter(nsteps), [](auto s) {
-                return std::abs(s);
-            });
-            return ins = info.add_instruction(
-                       make_op("step", {{"axes", op.axes}, {"steps", nsteps}}), ins);
-        }
-        else
-            return ins;
+        return sd;
     }
 };
 

src/permutation.cpp

@@ -74,5 +74,15 @@ std::vector<int64_t> find_permutation(const std::vector<shape>& shapes)
     return it->first;
 }
 
+std::vector<shape> normalize_permutation(const std::vector<shape>& shapes)
+{
+    auto result = shapes;
+    auto perm   = find_permutation(shapes);
+    std::transform(result.begin(), result.end(), result.begin(), [&](auto s) {
+        return reorder_shape(s, perm);
+    });
+    return result;
+}
+
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/program.cpp

@@ -40,13 +40,14 @@
 #include <migraphx/make_op.hpp>
 #include <migraphx/marker.hpp>
 #include <migraphx/supported_segments.hpp>
+
 #include <iostream>
+#include <queue>
 #include <sstream>
 #include <algorithm>
 #include <set>
 #include <unordered_map>
 #include <utility>
-
 #include <unordered_set>
 #include <map>
 #include <cassert>
@@ -222,7 +223,7 @@ void program::compile(const std::vector<target>& targets, std::vector<compile_op
     // Gather all the target roots
     std::unordered_multimap<std::size_t, module_ref> roots;
     auto mods = this->get_modules();
-    for(auto* mod : mods)
+    for(const auto* mod : mods)
     {
         for(const auto& ins : *mod)
         {
@@ -547,7 +548,7 @@ std::vector<argument> program::eval(parameter_map params, execution_environment
             ins_out[x] = ss.str();
         });
         ret = generic_eval(*this, contexts, std::move(params), [&](instruction_ref ins, auto f) {
-            auto& ctx = contexts[ins->get_target_id()];
+            const auto& ctx = contexts[ins->get_target_id()];
             ctx.finish();
             std::cout << "Run instruction: " << ins_out.at(ins) << std::endl;
             timer t{};
@@ -727,7 +728,7 @@ static void mod_from_val(module_ref mod,
                                std::back_inserter(module_inputs),
                                [&](const value& i) { return map_mods.at(i.to<std::string>()); });
 
-                for(auto& smod : module_inputs)
+                for(const auto& smod : module_inputs)
                 {
                     mod_from_val(smod, v, instructions, map_mods);
                 }
@@ -1185,17 +1186,25 @@ void program::remove_unused_modules()
     std::vector<module*> unused;
     generic_get_unused_modules(
         impl->modules, generic_get_modules(this->get_main_module()), std::back_inserter(unused));
-    for(auto* m : unused)
+    for(const auto* m : unused)
         this->remove_module(m->name());
 }
 
 program& program::sort()
 {
-    for(auto& pp : this->impl->modules)
+    std::queue<migraphx::module_ref> mqueue;
+    mqueue.push(get_main_module());
+    while(not mqueue.empty())
     {
-        pp.second.sort();
+        module_ref current_mod = mqueue.front();
+        current_mod->sort();
+        mqueue.pop();
+        auto child_mods = current_mod->get_sub_modules(true);
+        for(auto& sub_mod : child_mods)
+        {
+            mqueue.push(sub_mod);
+        }
     }
-
     return *this;
 }
 

src/py/CMakeLists.txt

@@ -23,14 +23,25 @@
 #####################################################################################
 
 option(MIGRAPHX_ENABLE_PYTHON "Enable python bindings" ON)
+add_library(migraphx_py py_loader.cpp)
+migraphx_generate_export_header(migraphx_py)
+target_include_directories(migraphx_py PRIVATE include)
+target_link_libraries(migraphx_py PUBLIC migraphx)
+rocm_install_targets(TARGETS migraphx_py INCLUDE include)
 if(MIGRAPHX_ENABLE_PYTHON)
     include(PythonModules)
 
-    add_custom_target(migraphx_py)
 
     foreach(PYTHON_VERSION ${PYTHON_VERSIONS})
-        py_add_module(migraphx_py_${PYTHON_VERSION} migraphx_py.cpp PYTHON_VERSION ${PYTHON_VERSION} PYTHON_MODULE migraphx)
-        target_link_libraries(migraphx_py_${PYTHON_VERSION} PRIVATE migraphx migraphx_tf migraphx_onnx migraphx_all_targets)
+        py_add_module(migraphx_pybind_${PYTHON_VERSION} migraphx_py.cpp PYTHON_VERSION ${PYTHON_VERSION} PYTHON_MODULE migraphx)
+        target_link_libraries(migraphx_pybind_${PYTHON_VERSION} PRIVATE migraphx migraphx_tf migraphx_onnx migraphx_all_targets)
+        rocm_install_targets(TARGETS migraphx_pybind_${PYTHON_VERSION})
+        add_dependencies(migraphx_py migraphx_pybind_${PYTHON_VERSION})
+        
+        add_library(migraphx_py_${PYTHON_VERSION} py.cpp)
+        target_include_directories(migraphx_py_${PYTHON_VERSION} PRIVATE include)
+        target_link_libraries(migraphx_py_${PYTHON_VERSION} PUBLIC migraphx)
+        target_link_libraries(migraphx_py_${PYTHON_VERSION} PRIVATE pybind11::pybind11 python${PYTHON_VERSION}::runtime)
         rocm_install_targets(TARGETS migraphx_py_${PYTHON_VERSION})
         add_dependencies(migraphx_py migraphx_py_${PYTHON_VERSION})
     endforeach()

src/py/include/migraphx/py.hpp

+/*
+ * 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.
+ */
+#ifndef MIGRAPHX_GUARD_MIGRAPHX_PY_HPP
+#define MIGRAPHX_GUARD_MIGRAPHX_PY_HPP
+
+#include <migraphx/config.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/py/export.h>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+MIGRAPHX_PY_EXPORT program load_py(const std::string& filename);
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+#endif // MIGRAPHX_GUARD_MIGRAPHX_PY_HPP

src/py/py.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/config.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/dynamic_loader.hpp>
+#include <migraphx/file_buffer.hpp>
+#include <pybind11/embed.h>
+
+namespace py = pybind11;
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+#ifdef __clang__
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wreturn-type-c-linkage"
+#endif
+// extern "C" is used to disable name mangling, but the function will still be called from C++
+extern "C" program migraphx_load_py(const std::string& filename);
+#ifdef __clang__
+#pragma clang diagnostic pop
+#endif
+
+const std::string& python_path()
+{
+    static const auto path = dynamic_loader::path(&migraphx_load_py).parent_path().string();
+    return path;
+}
+
+static py::dict run_file(const std::string& file)
+{
+    py::object scope = py::module_::import("__main__").attr("__dict__");
+    std::string buffer;
+    buffer.append("import sys\n");
+    buffer.append("sys.path.insert(0, '" + python_path() + "')\n");
+    buffer.append("import migraphx\n");
+    buffer.append(read_string(file));
+    py::exec(buffer, scope);
+    return scope.cast<py::dict>();
+}
+
+extern "C" program migraphx_load_py(const std::string& filename)
+{
+    py::scoped_interpreter guard{};
+    py::dict vars = run_file(filename);
+    auto it       = std::find_if(vars.begin(), vars.end(), [](const auto& p) {
+        return py::isinstance<migraphx::program>(p.second);
+    });
+    if(it == vars.end())
+        MIGRAPHX_THROW("No program variable found");
+    return it->second.cast<migraphx::program>();
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/py/py_loader.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/py.hpp>
+#include <migraphx/dynamic_loader.hpp>
+#include <migraphx/process.hpp>
+#include <migraphx/ranges.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+static std::vector<fs::path> find_available_python_versions()
+{
+    std::vector<fs::path> result;
+    auto path = dynamic_loader::path(&load_py).parent_path();
+    for(const auto& entry : fs::directory_iterator{path})
+    {
+        auto p = entry.path();
+        if(not fs::is_regular_file(p))
+            continue;
+        if(not contains(p.stem().string(), "migraphx_py_"))
+            continue;
+        result.push_back(p);
+    }
+    std::sort(result.begin(), result.end(), std::greater<>{});
+    return result;
+}
+
+static dynamic_loader load_py_lib()
+{
+    auto libs = find_available_python_versions();
+    for(const auto& lib : libs)
+    {
+        auto result = dynamic_loader::try_load(lib);
+        if(result.has_value())
+            return *result;
+    }
+    MIGRAPHX_THROW("Cant find a viable version of python");
+}
+
+static dynamic_loader py_lib()
+{
+    static dynamic_loader lib = load_py_lib();
+    return lib;
+}
+
+MIGRAPHX_PY_EXPORT program load_py(const std::string& filename)
+{
+    static auto f = py_lib().get_function<program(const std::string&)>("migraphx_load_py");
+    return f(filename);
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/rewrite_quantization.cpp

@@ -28,6 +28,7 @@
 #include <migraphx/tune_axis.hpp>
 #include <migraphx/program.hpp>
 #include <migraphx/shape.hpp>
+#include <migraphx/common.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -61,13 +62,10 @@ void apply_quantizelinear(module& m, instruction_ref ins)
         max_quant = qt.max();
         min_quant = qt.min();
     });
-    auto s = add_zero_point->get_shape();
-    std::vector<int> min_data(s.elements(), min_quant);
-    std::vector<int> max_data(s.elements(), max_quant);
-    auto min_arg = m.add_literal(literal(s, min_data));
-    auto max_arg = m.add_literal(literal(s, max_data));
-
-    auto saturate = m.insert_instruction(ins, make_op("clip"), add_zero_point, min_arg, max_arg);
+    auto s        = add_zero_point->get_shape();
+    auto min_arg  = m.add_literal(literal{shape{s.type()}, {min_quant}});
+    auto max_arg  = m.add_literal(literal{shape{s.type()}, {max_quant}});
+    auto saturate = insert_common_op(m, ins, make_op("clip"), {add_zero_point, min_arg, max_arg});
     m.replace_instruction(
         ins, make_op("convert", {{"target_type", ins->get_shape().type()}}), saturate);
 }