merge changes from develop branch

src/include/migraphx/op/prefix_scan_op.hpp

@@ -38,18 +38,38 @@ struct prefix_scan_op : op_name<Derived>
     shape normalize_compute_shape(std::vector<shape> inputs) const
     {
         check_shapes{inputs, *this}.has(1);
-        return inputs.at(0);
+        auto s = inputs.front();
+        if(s.broadcasted())
+        {
+            return {s.type(), s.lens()};
+        }
+        else
+        {
+            return s.with_lens(s.lens());
+        }
     }
 
-    argument compute(const shape&, std::vector<argument> args) const
+    argument compute(const shape& output_shape, std::vector<argument> args) const
     {
-        argument result = args[0].copy();
-        auto s          = result.get_shape();
-        auto slice      = shape{s.type(), {s.lens()[axis]}, {s.strides()[axis]}};
-        auto lens       = s.lens();
-        lens[axis]      = 1;
-        auto batch      = shape{s.type(), lens, s.strides()};
-        auto& self      = static_cast<const Derived&>(*this);
+        argument result{output_shape};
+        auto s = args[0].get_shape();
+        if(s == output_shape)
+        {
+            result = args[0].copy();
+        }
+        else
+        {
+            visit_all(result, args[0])([&](auto output, auto input) {
+                par_for(output_shape.elements(),
+                        [&](auto i) { output[output_shape.index(i)] = input[s.index(i)]; });
+            });
+            s = output_shape;
+        }
+        auto slice = shape{s.type(), {s.lens()[axis]}, {s.strides()[axis]}};
+        auto lens  = s.lens();
+        lens[axis] = 1;
+        auto batch = shape{s.type(), lens, s.strides()};
+        auto& self = static_cast<const Derived&>(*this);
         result.visit([&](auto output) {
             using type = decltype(output);
             par_for(batch.elements(), [&](auto i) {

src/include/migraphx/op/reshape.hpp

@@ -75,7 +75,6 @@ struct reshape
         return args[0].reshape(output_shape);
     }
 
-    lifetime get_lifetime() const { return lifetime::borrow; }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 

src/include/migraphx/op/roialign.hpp

@@ -3,6 +3,7 @@
 
 #include <limits>
 #include <migraphx/check_shapes.hpp>
+#include <migraphx/op/common.hpp>
 #include <migraphx/config.hpp>
 #include <migraphx/argument.hpp>
 #include <migraphx/par_for.hpp>
@@ -21,7 +22,7 @@ namespace op {
 struct roialign
 {
     std::string coord_trans_mode = "half_pixel";
-    std::string mode             = "avg";
+    pooling_mode mode            = {pooling_mode::average};
     int64_t output_height        = 1;
     int64_t output_width         = 1;
     int64_t sampling_ratio       = 0;
@@ -241,16 +242,17 @@ struct roialign
                                                            in_dims[0] * in_dims[1]);
                     double output_val;
                     std::tie(output_val, vec_index[c]) =
-                        (mode == "avg") ? this->calc_pooling(offset_bottom_data,
-                                                             bin_grid_size,
-                                                             pre_calc,
-                                                             vec_index[c],
-                                                             avg_pool{})
-                                        : this->calc_pooling(offset_bottom_data,
-                                                             bin_grid_size,
-                                                             pre_calc,
-                                                             vec_index[c],
-                                                             max_pool{});
+                        (mode == migraphx::op::pooling_mode::average)
+                            ? this->calc_pooling(offset_bottom_data,
+                                                 bin_grid_size,
+                                                 pre_calc,
+                                                 vec_index[c],
+                                                 avg_pool{})
+                            : this->calc_pooling(offset_bottom_data,
+                                                 bin_grid_size,
+                                                 pre_calc,
+                                                 vec_index[c],
+                                                 max_pool{});
                     output(n, c, ph, pw) = output_val;
                 });
             });

src/include/migraphx/op/scalar.hpp

@@ -40,7 +40,6 @@ struct scalar
     {
         return args[0].reshape(output_shape);
     }
-    lifetime get_lifetime() const { return lifetime::borrow; }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 

src/include/migraphx/op/squeeze.hpp

@@ -84,7 +84,6 @@ struct squeeze
     {
         return args[0].reshape(output_shape);
     }
-    lifetime get_lifetime() const { return lifetime::borrow; }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 

src/include/migraphx/op/step.hpp

@@ -72,8 +72,6 @@ struct step
         return args[0].reshape(output_shape);
     }
 
-    lifetime get_lifetime() const { return lifetime::borrow; }
-
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 

src/include/migraphx/op/transpose.hpp

@@ -56,7 +56,6 @@ struct transpose
     {
         return args[0].reshape(output_shape);
     }
-    lifetime get_lifetime() const { return lifetime::borrow; }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 

src/include/migraphx/op/unsqueeze.hpp

@@ -81,7 +81,6 @@ struct unsqueeze
     {
         return args[0].reshape(output_shape);
     }
-    lifetime get_lifetime() const { return lifetime::borrow; }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 

src/include/migraphx/operation.hpp

@@ -461,8 +461,8 @@ lifetime get_lifetime_op(const T&)
  *      shape compute_shape(const std::vector<shape>& input) const;
  *      shape compute_shape(const std::vector<shape>& inputs,const std::vector<module_ref>&
  * mod_args) const; argument compute(context& ctx,const shape& output,const std::vector<argument>&
- * input) const; argument compute(const shape& output,const std::vector<argument>& input)
- * const; argument compute(const shape& output,const std::vector<argument>& input,const
+ * input) const; argument compute(const shape& output,const std::vector<argument>& input) const;
+ *      argument compute(const shape& output,const std::vector<argument>& input,const
  * std::vector<module_ref>& module_args,std::function<std::vector<argument>(module_ref&, const
  * std::unordered_map<std::string, argument>&)> run) const; argument compute(context& ctx,const
  * shape& output,const std::vector<argument>& input,const std::vector<module_ref>&

src/include/migraphx/program.hpp

@@ -86,6 +86,9 @@ struct program
                const std::function<void(instruction_ref,
                                         std::unordered_map<instruction_ref, std::string>)>&
                    print_func) const;
+    void print(const std::function<void(instruction_ref ins,
+                                        std::unordered_map<instruction_ref, std::string>)>&
+                   print_func) const;
 
     void print_graph(std::ostream& os, bool brief = false) const;
     void print_cpp(std::ostream& os) const;

src/insert_pad.cpp

@@ -44,7 +44,7 @@ static void update_op(const instruction_ref& input, const instruction_ref& ins,
 static void update_pooling(const instruction_ref& input, const instruction_ref& ins, module& m)
 {
     auto op = any_cast<op::pooling>(ins->get_operator());
-    if(op.mode == "average")
+    if(op.mode == op::pooling_mode::average)
     {
         return;
     }

src/json.cpp

@@ -133,6 +133,12 @@ std::string to_json_string(const value& val)
     return j.dump();
 }
 
+std::string to_pretty_json_string(const value& val, std::size_t indent)
+{
+    json j = val;
+    return j.dump(indent);
+}
+
 migraphx::value from_json_string(const char* str, std::size_t size)
 {
     json j = json::parse(str, str + size);

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

@@ -32,9 +32,20 @@ struct onnx_parser
         instruction_ref add_bias(const std::vector<instruction_ref>& args,
                                  instruction_ref curr_ins,
                                  uint64_t axis) const;
+
         instruction_ref add_broadcastable_binary_op(const std::string& op_name,
                                                     instruction_ref arg0,
                                                     instruction_ref arg1) const;
+
+        instruction_ref add_common_op(const std::string& op_name,
+                                      std::vector<instruction_ref> inputs) const;
+
+        template <class... Ts>
+        instruction_ref add_common_op(const std::string& op_name, Ts... xs) const
+        {
+            return add_common_op(op_name, {xs...});
+        }
+
         instruction_ref add_instruction(const operation& op,
                                         const std::vector<instruction_ref>& args) const;
 

src/onnx/onnx_parser.cpp

@@ -98,7 +98,13 @@ instruction_ref onnx_parser::node_info::add_broadcastable_binary_op(const std::s
                                                                     instruction_ref arg0,
                                                                     instruction_ref arg1) const
 {
-    return add_common_op(*mod, make_op(op_name), {arg0, arg1});
+    return this->add_common_op(op_name, arg0, arg1);
+}
+
+instruction_ref onnx_parser::node_info::add_common_op(const std::string& op_name,
+                                                      std::vector<instruction_ref> inputs) const
+{
+    return migraphx::add_common_op(*mod, make_op(op_name), std::move(inputs));
 }
 
 instruction_ref

src/onnx/padding.cpp

@@ -3,6 +3,7 @@
 #include <migraphx/pad_calc.hpp>
 #include <migraphx/stringutils.hpp>
 #include <migraphx/make_op.hpp>
+#include <migraphx/op/common.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -94,7 +95,7 @@ void tune_padding_size(const value& v,
                        std::vector<int64_t>& s_start)
 {
     // maxpooling or count_include_pad is 1, no change is required.
-    if(v.at("mode").to<std::string>() == "max" or count_include_pad == 1)
+    if(v.at("mode").to<op::pooling_mode>() == op::pooling_mode::max or count_include_pad == 1)
     {
         return;
     }

src/onnx/parse_celu.cpp

+#include <migraphx/onnx/op_parser.hpp>
+#include <migraphx/onnx/checks.hpp>
+#include <migraphx/ranges.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/make_op.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace onnx {
+
+struct parse_celu : op_parser<parse_celu>
+{
+    std::vector<op_desc> operators() const { return {{"Celu"}}; }
+
+    instruction_ref parse(const op_desc&,
+                          const onnx_parser&,
+                          const onnx_parser::node_info& info,
+                          std::vector<instruction_ref> args) const
+    {
+        float alpha = 1.0;
+        if(contains(info.attributes, "alpha"))
+        {
+            alpha = info.attributes.at("alpha").f();
+        }
+        if(float_equal(alpha, 0.0f))
+        {
+            MIGRAPHX_THROW("CELU: alpha is zero (division by zero)");
+        }
+
+        auto input_lens = args[0]->get_shape().lens();
+        auto input_type = args[0]->get_shape().type();
+        if(input_type != migraphx::shape::float_type)
+        {
+            MIGRAPHX_THROW("CELU: input tensor not float type");
+        }
+        auto zero_lit = info.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", input_lens}}),
+            info.add_literal(migraphx::literal{migraphx::shape{input_type}, {0.}}));
+        auto one_lit = info.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", input_lens}}),
+            info.add_literal(migraphx::literal{migraphx::shape{input_type}, {1.}}));
+        auto alpha_lit = info.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", input_lens}}),
+            info.add_literal(migraphx::literal{migraphx::shape{input_type}, {alpha}}));
+        auto linear_part = info.add_instruction(migraphx::make_op("max"), zero_lit, args[0]);
+        auto divi        = info.add_instruction(migraphx::make_op("div"), args[0], alpha_lit);
+        auto expo        = info.add_instruction(migraphx::make_op("exp"), divi);
+        auto sub         = info.add_instruction(migraphx::make_op("sub"), expo, one_lit);
+        auto mul         = info.add_instruction(migraphx::make_op("mul"), alpha_lit, sub);
+        auto exp_part    = info.add_instruction(migraphx::make_op("min"), zero_lit, mul);
+        return info.add_instruction(migraphx::make_op("add"), linear_part, exp_part);
+    }
+};
+
+} // namespace onnx
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/onnx/parse_clip.cpp

@@ -16,7 +16,6 @@ struct parse_clip : op_parser<parse_clip>
                           onnx_parser::node_info info,
                           std::vector<instruction_ref> args) const
     {
-        auto input_lens = args[0]->get_shape().lens();
         instruction_ref min_arg;
         instruction_ref max_arg;
         bool min_used = false;
@@ -45,29 +44,17 @@ struct parse_clip : op_parser<parse_clip>
             max_used      = true;
         }
 
-        if(min_used)
-        {
-            min_arg = info.add_instruction(make_op("multibroadcast", {{"out_lens", input_lens}}),
-                                           min_arg);
-        }
-
-        if(max_used)
-        {
-            max_arg = info.add_instruction(make_op("multibroadcast", {{"out_lens", input_lens}}),
-                                           max_arg);
-        }
-
         if(min_used and max_used)
         {
-            return info.add_instruction(make_op("clip"), args[0], min_arg, max_arg);
+            return info.add_common_op("clip", args[0], min_arg, max_arg);
         }
         else if(max_used)
         {
-            return info.add_instruction(make_op("min"), args[0], max_arg);
+            return info.add_broadcastable_binary_op("min", args[0], max_arg);
         }
         else if(min_used)
         {
-            return info.add_instruction(make_op("max"), args[0], min_arg);
+            return info.add_broadcastable_binary_op("max", args[0], min_arg);
         }
         else
         {

src/onnx/parse_eyelike.cpp

+#include <migraphx/onnx/op_parser.hpp>
+#include <migraphx/onnx/checks.hpp>
+#include <migraphx/ranges.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/make_op.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace onnx {
+
+struct parse_eyelike : op_parser<parse_eyelike>
+{
+    std::vector<op_desc> operators() const { return {{"EyeLike"}}; }
+
+    instruction_ref parse(const op_desc&,
+                          const onnx_parser&,
+                          const onnx_parser::node_info& info,
+                          std::vector<instruction_ref> args) const
+    {
+        auto input_shape = args[0]->get_shape();
+        auto input_lens  = input_shape.lens();
+        if(input_lens.size() != 2)
+        {
+            MIGRAPHX_THROW("EYELIKE: tensor input not of rank 2");
+        }
+        std::ptrdiff_t num_rows = input_lens.front();
+        std::ptrdiff_t num_cols = input_lens.back();
+
+        shape::type_t output_type = args[0]->get_shape().type();
+        if(contains(info.attributes, "dtype"))
+        {
+            output_type = get_type(info.attributes.at("dtype").i());
+        }
+
+        std::ptrdiff_t k = 0;
+        if(contains(info.attributes, "k"))
+        {
+            k = info.attributes.at("k").i();
+        }
+        if(k >= 0)
+        {
+            if(k >= num_cols)
+            {
+                std::ostringstream oss;
+                oss << "EYELIKE: positive k out of bounds, k = " << k << " num_cols = " << num_cols;
+                MIGRAPHX_THROW(oss.str());
+            }
+        }
+        else
+        {
+            if(std::abs(k) >= num_rows)
+            {
+                std::ostringstream oss;
+                oss << "EYELIKE: negative k out of bounds, k = " << k << " num_rows = " << num_cols;
+                MIGRAPHX_THROW(oss.str());
+            }
+        }
+
+        std::vector<char> eyelike_mat(num_rows * num_cols, 0);
+        for(std::ptrdiff_t i = 0; i < num_rows; ++i)
+        {
+            auto idx = i + k;
+            if(idx < num_cols and idx >= 0)
+                eyelike_mat[(num_cols + 1) * i + k] = char{1};
+        }
+        return info.add_literal(
+            migraphx::literal{migraphx::shape{output_type, input_lens}, eyelike_mat});
+    }
+};
+
+} // namespace onnx
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/onnx/parse_lpnormalization.cpp

+#include <migraphx/onnx/op_parser.hpp>
+#include <migraphx/onnx/checks.hpp>
+#include <migraphx/ranges.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/make_op.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace onnx {
+
+//!  Parser for LpNormalization ONNX operator.
+/*!
+  Normalizes a tensor by the L1 or L2 norms along a given axis.
+  Norms that evaluate to 0 are changed to 1 to prevent division by zero.
+*/
+struct parse_lpnormalization : op_parser<parse_lpnormalization>
+{
+    std::vector<op_desc> operators() const { return {{"LpNormalization"}}; }
+
+    instruction_ref parse(const op_desc&,
+                          const onnx_parser&,
+                          const onnx_parser::node_info& info,
+                          std::vector<instruction_ref> args) const
+    {
+        int p = 2;
+        if(contains(info.attributes, "p"))
+        {
+            p = info.attributes.at("p").i();
+        }
+        if(p != 1 and p != 2)
+        {
+            MIGRAPHX_THROW("LPNORMALIZATION: only L1 and L2 norm supported");
+        }
+        auto input              = args.front();
+        auto input_shape        = input->get_shape();
+        const auto& input_lens  = input_shape.lens();
+        auto input_type         = input_shape.type();
+        std::ptrdiff_t num_axes = input_lens.size();
+        std::ptrdiff_t axis     = -1;
+        if(contains(info.attributes, "axis"))
+        {
+            axis = info.attributes.at("axis").i();
+            if(axis < -num_axes or axis >= num_axes)
+            {
+                // handled in normalize_attributes but throwing here might be clearer
+                MIGRAPHX_THROW("LPNORMALIZATION: selected axis out of bounds");
+            }
+        }
+        migraphx::instruction_ref p_val;
+        if(p == 1)
+        {
+            p_val = info.add_instruction(migraphx::make_op("abs"), input);
+        }
+        else
+        {
+            p_val = info.add_instruction(migraphx::make_op("mul"), input, input);
+        }
+
+        // need to check for zeros from lp norm to prevent division by zero
+        // change them to 1 for the element-wise division
+        auto norms =
+            info.add_instruction(migraphx::make_op("reduce_sum", {{"axes", {axis}}}), p_val);
+        if(p == 2)
+        {
+            norms = info.add_instruction(migraphx::make_op("sqrt"), norms);
+        }
+        // broadcast back to initial shape, negative axis option doesn't work with unidirectional
+        norms = info.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", input_lens}}), norms);
+        auto zero_mb = info.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", input_lens}}),
+            info.add_literal(migraphx::literal{migraphx::shape{input_type}, {0.}}));
+        auto one_mb = info.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", input_lens}}),
+            info.add_literal(migraphx::literal{migraphx::shape{input_type}, {1.}}));
+        auto is_zero = info.add_instruction(migraphx::make_op("equal"), norms, zero_mb);
+        auto norms_zeros_to_one =
+            info.add_instruction(migraphx::make_op("where"), is_zero, one_mb, norms);
+        return info.add_instruction(migraphx::make_op("div"), input, norms_zeros_to_one);
+    }
+};
+
+} // namespace onnx
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/onnx/parse_pooling.cpp

@@ -2,6 +2,7 @@
 #include <migraphx/onnx/checks.hpp>
 #include <migraphx/onnx/padding.hpp>
 #include <migraphx/op/pad.hpp>
+#include <migraphx/op/pooling.hpp>
 #include <migraphx/instruction.hpp>
 #include <migraphx/ranges.hpp>
 #include <migraphx/stringutils.hpp>
@@ -27,10 +28,16 @@ struct parse_pooling : op_parser<parse_pooling>
                           std::vector<instruction_ref> args) const
     {
         std::string mode = opd.op_name;
-        operation op     = make_op("pooling", {{"mode", mode}});
-        value values     = op.to_value();
-        auto l0          = args[0];
-        auto in_lens     = l0->get_shape().lens();
+        if(mode != "max" && mode != "average")
+        {
+            MIGRAPHX_THROW("onnx pooling mode must be \"max\" or \"average\"");
+        }
+        operation op = make_op(
+            "pooling",
+            {{"mode", mode == "average" ? op::pooling_mode::average : op::pooling_mode::max}});
+        value values = op.to_value();
+        auto l0      = args[0];
+        auto in_lens = l0->get_shape().lens();
         assert(in_lens.size() > 2);
         auto kdims = in_lens.size() - 2;
 
@@ -72,6 +79,7 @@ struct parse_pooling : op_parser<parse_pooling>
 
         std::vector<int64_t> paddings;
         float pad_val = ((mode == "max") ? std::numeric_limits<float>::lowest() : 0.0f);
+
         if(contains(info.attributes, "pads"))
         {
             values["padding"].clear();