Add implicit type conversion for broadcastable operators (#875)

* Add implicit type conversion for broadcastable operators

* Formatting

* Fix warnings

src/CMakeLists.txt

@@ -9,58 +9,59 @@ add_library(migraphx
     analyze_streams.cpp
     argument.cpp
     auto_contiguous.cpp
-    eliminate_common_subexpression.cpp
-    decompose.cpp
-    propagate_constant.cpp
+    common.cpp
     compile_src.cpp
+    convert_to_json.cpp
     cpp_generator.cpp
     dead_code_elimination.cpp
+    decompose.cpp
     dom_info.cpp
     dynamic_loader.cpp
     eliminate_allocation.cpp
-    eliminate_contiguous.cpp
+    eliminate_common_subexpression.cpp
     eliminate_concat.cpp
+    eliminate_contiguous.cpp
     eliminate_data_type.cpp
     eliminate_identity.cpp
     eliminate_pad.cpp
-    insert_pad.cpp
-    file_buffer.cpp
-    rewrite_batchnorm.cpp
-    rewrite_rnn.cpp
-    rewrite_pooling.cpp
     env.cpp
+    file_buffer.cpp
     generate.cpp
     inline_module.cpp
+    insert_pad.cpp
     instruction.cpp
+    json.cpp
     load_save.cpp
     make_op.cpp
+    module.cpp
     msgpack.cpp
+    normalize_attributes.cpp
+    normalize_ops.cpp
     operation.cpp
+    opt/memory_coloring.cpp
+    opt/memory_coloring_impl.cpp
+    pass_manager.cpp
     permutation.cpp
     preallocate_param.cpp
     process.cpp
     program.cpp
-    module.cpp
+    propagate_constant.cpp
     quantization.cpp
     reduce_dims.cpp
+    register_op.cpp
+    register_target.cpp
     remap.cpp
-    shape.cpp
+    rewrite_batchnorm.cpp
+    rewrite_pooling.cpp
+    rewrite_rnn.cpp
     schedule.cpp
     serialize.cpp
-    pass_manager.cpp
-    register_op.cpp
-    register_target.cpp
+    shape.cpp
     simplify_algebra.cpp
     simplify_reshapes.cpp
     tmp_dir.cpp
     value.cpp
     verify_args.cpp
-    json.cpp
-    convert_to_json.cpp
-    opt/memory_coloring.cpp
-    opt/memory_coloring_impl.cpp
-    normalize_attributes.cpp
-    normalize_ops.cpp
 )
 configure_file(version.h.in include/migraphx/version.h)
 rocm_set_soversion(migraphx ${MIGRAPHX_SO_VERSION})

src/common.cpp

+#include <migraphx/common.hpp>
+#include <migraphx/module.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/algorithm.hpp>
+#include <migraphx/stringutils.hpp>
+#include <migraphx/instruction.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+// Example:
+// s0 = (3,2,4,5) and s1 = (2,1,1)
+//
+// In this case we need to broadcast (:,1,1) portion of
+// s1 plus broadcast the 1st dimension of s1
+// giving output_lens = (3,2,4,5)
+//
+// Another example:
+// s0 = (3,2,1,5) and s1 = (2,7,5)
+// In this case we need to broadcast the (:,:,1:,:) axis
+// of s0 plus the 1st dimension of s1 giving
+// output_lens = (3,2,7,5)
+std::vector<std::size_t> compute_broadcasted_lens(std::vector<std::size_t> s0,
+                                                  std::vector<std::size_t> s1)
+{
+    if(s0 == s1)
+        return s0;
+    if(s0.size() > s1.size())
+        s0.swap(s1);
+
+    std::vector<std::size_t> out_lens(s1);
+    auto offset = s1.size() - s0.size();
+    std::transform(
+        s0.begin(), s0.end(), s1.begin() + offset, out_lens.begin() + offset, [&](auto a, auto b) {
+            if(a != b and a != 1 and b != 1)
+            {
+                MIGRAPHX_THROW("COMPUTE_BROADCASTLEN: shape {" + to_string_range(s0) + "} and {" +
+                               to_string_range(s1) + "} mismatch!");
+            }
+            return std::max(a, b);
+        });
+
+    return out_lens;
+}
+
+std::vector<std::size_t> compute_common_lens(const std::vector<shape>& shapes)
+{
+    assert(not shapes.empty());
+    return transform_accumulate(shapes.begin() + 1,
+                                shapes.end(),
+                                shapes.front().lens(),
+                                &compute_broadcasted_lens,
+                                [](auto s) { return s.lens(); });
+}
+
+shape::type_t compute_common_type(shape::type_t t1, shape::type_t t2)
+{
+    if(t1 == t2)
+        return t1;
+    shape::type_t result;
+    shape::visit(t1, [&](auto x) {
+        shape::visit(t2, [&](auto y) {
+            // Workaround broken warning on gcc 5
+            (void)x;
+            (void)y;
+            using type = std::common_type_t<decltype(x()), decltype(y())>;
+            result     = shape::get_type<type>{};
+        });
+    });
+    return result;
+}
+
+shape::type_t compute_common_types(const std::vector<shape>& shapes)
+{
+    assert(not shapes.empty());
+    return transform_accumulate(
+        shapes.begin() + 1, shapes.end(), shapes.front().type(), &compute_common_type, [&](auto s) {
+            return s.type();
+        });
+}
+
+shape common_shape(const std::vector<shape>& shapes)
+{
+    if(shapes.empty())
+        return {};
+    return {compute_common_types(shapes), compute_common_lens(shapes)};
+}
+
+instruction_ref insert_common_op(module& m,
+                                 instruction_ref ins,
+                                 const operation& op,
+                                 std::vector<instruction_ref> inputs)
+{
+    auto common = common_shape(to_shapes(inputs));
+    std::transform(inputs.begin(), inputs.end(), inputs.begin(), [&](auto input) {
+        if(input->get_shape().lens() != common.lens())
+        {
+            input = m.insert_instruction(
+                ins, make_op("multibroadcast", {{"output_lens", common.lens()}}), input);
+        }
+        if(input->get_shape().type() != common.type())
+        {
+            input = m.insert_instruction(
+                ins, make_op("convert", {{"target_type", common.type()}}), input);
+        }
+        return input;
+    });
+    return m.insert_instruction(ins, op, inputs);
+}
+
+instruction_ref add_common_op(module& m, const operation& op, std::vector<instruction_ref> inputs)
+{
+    return insert_common_op(m, m.end(), op, std::move(inputs));
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/include/migraphx/algorithm.hpp

@@ -2,6 +2,7 @@
 #define MIGRAPHX_GUARD_RTGLIB_ALGORITHM_HPP
 
 #include <algorithm>
+#include <numeric>
 #include <migraphx/config.hpp>
 
 namespace migraphx {
@@ -21,6 +22,13 @@ void transform_if(Iterator start, Iterator last, Output out, Predicate pred, F f
     }
 }
 
+template <class Iterator, class T, class BinaryOp, class UnaryOp>
+T transform_accumulate(Iterator first, Iterator last, T init, BinaryOp binop, UnaryOp unaryop)
+{
+    return std::inner_product(
+        first, last, first, init, binop, [&](auto&& x, auto&&) { return unaryop(x); });
+}
+
 template <class Iterator, class Output, class Predicate>
 void group_by(Iterator start, Iterator last, Output out, Predicate pred)
 {

src/include/migraphx/common.hpp

+#ifndef MIGRAPHX_GUARD_MIGRAPHX_COMMON_HPP
+#define MIGRAPHX_GUARD_MIGRAPHX_COMMON_HPP
+
+#include <migraphx/config.hpp>
+#include <migraphx/shape.hpp>
+#include <migraphx/instruction_ref.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+struct module;
+struct operation;
+
+std::vector<std::size_t> compute_broadcasted_lens(std::vector<std::size_t> s0,
+                                                  std::vector<std::size_t> s1);
+shape common_shape(const std::vector<shape>& shapes);
+
+instruction_ref insert_common_op(module& m,
+                                 instruction_ref ins,
+                                 const operation& op,
+                                 std::vector<instruction_ref> inputs);
+instruction_ref add_common_op(module& m, const operation& op, std::vector<instruction_ref> inputs);
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+#endif // MIGRAPHX_GUARD_MIGRAPHX_COMMON_HPP

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

@@ -84,9 +84,6 @@ struct onnx_parser
 
 shape::type_t get_type(int dtype);
 
-std::vector<std::size_t> compute_broadcasted_lens(std::vector<std::size_t> s0,
-                                                  std::vector<std::size_t> s1);
-
 } // namespace onnx
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/onnx/onnx_parser.cpp

@@ -6,6 +6,7 @@
 #include <migraphx/ranges.hpp>
 #include <migraphx/instruction.hpp>
 #include <migraphx/pad_calc.hpp>
+#include <migraphx/common.hpp>
 #include <migraphx/type_traits.hpp>
 #include <migraphx/float_equal.hpp>
 #include <migraphx/file_buffer.hpp>
@@ -91,66 +92,11 @@ instruction_ref onnx_parser::node_info::add_bias(const std::vector<instruction_r
     return curr_ins;
 }
 
-std::vector<std::size_t> compute_broadcasted_lens(std::vector<std::size_t> s0,
-                                                  std::vector<std::size_t> s1)
-{
-    // Example:
-    // s0 = (3,2,4,5) and s1 = (2,1,1)
-    //
-    // In this case we need to broadcast (:,1,1) portion of
-    // s1 plus broadcast the 1st dimension of s1
-    // giving output_lens = (3,2,4,5)
-    //
-    // Another example:
-    // s0 = (3,2,1,5) and s1 = (2,7,5)
-    // In this case we need to broadcast the (:,:,1:,:) axis
-    // of s0 plus the 1st dimension of s1 giving
-    // output_lens = (3,2,7,5)
-    if(s0.size() > s1.size())
-    {
-        s0.swap(s1);
-    }
-
-    std::vector<std::size_t> out_lens(s1);
-    auto offset = s1.size() - s0.size();
-    std::transform(
-        s0.begin(), s0.end(), s1.begin() + offset, out_lens.begin() + offset, [&](auto a, auto b) {
-            if(a != b and a != 1 and b != 1)
-            {
-                MIGRAPHX_THROW("COMPUTE_BROADCASTLEN: shape {" + to_string_range(s0) + "} and {" +
-                               to_string_range(s1) + "} mismatch!");
-            }
-            return std::max(a, b);
-        });
-
-    return out_lens;
-}
-
 instruction_ref onnx_parser::node_info::add_broadcastable_binary_op(const std::string& op_name,
                                                                     instruction_ref arg0,
                                                                     instruction_ref arg1) const
 {
-    if(arg0->get_shape().lens() != arg1->get_shape().lens())
-    {
-        // Get lengths for both arguments
-        auto s0       = arg0->get_shape().lens();
-        auto s1       = arg1->get_shape().lens();
-        auto out_lens = compute_broadcasted_lens(s0, s1);
-
-        auto l0 = arg0;
-        if(arg0->get_shape().lens() != out_lens)
-            l0 = add_instruction(make_op("multibroadcast", {{"output_lens", out_lens}}), arg0);
-
-        auto l1 = arg1;
-        if(arg1->get_shape().lens() != out_lens)
-            l1 = add_instruction(make_op("multibroadcast", {{"output_lens", out_lens}}), arg1);
-
-        return add_instruction(make_op(op_name), l0, l1);
-    }
-    else
-    {
-        return add_instruction(make_op(op_name), {arg0, arg1});
-    }
+    return add_common_op(*mod, make_op(op_name), {arg0, arg1});
 }
 
 instruction_ref

src/onnx/parse_expand.cpp

@@ -2,6 +2,7 @@
 #include <migraphx/onnx/checks.hpp>
 #include <migraphx/ranges.hpp>
 #include <migraphx/instruction.hpp>
+#include <migraphx/common.hpp>
 #include <migraphx/make_op.hpp>
 
 namespace migraphx {

src/onnx/parse_matmul.cpp

 #include <migraphx/onnx/op_parser.hpp>
 #include <migraphx/ranges.hpp>
 #include <migraphx/instruction.hpp>
+#include <migraphx/common.hpp>
 #include <migraphx/make_op.hpp>
 
 namespace migraphx {

src/onnx/parse_where.cpp

 #include <migraphx/onnx/op_parser.hpp>
 #include <migraphx/ranges.hpp>
 #include <migraphx/instruction.hpp>
+#include <migraphx/common.hpp>
 #include <migraphx/make_op.hpp>
 
 namespace migraphx {

src/tf/tf_parser.cpp

@@ -17,6 +17,7 @@
 #include <migraphx/instruction.hpp>
 #include <migraphx/config.hpp>
 #include <migraphx/tf.hpp>
+#include <migraphx/common.hpp>
 #include <migraphx/make_op.hpp>
 
 #include <migraphx/tf/tf_parser.hpp>
@@ -74,66 +75,11 @@ instruction_ref tf_parser::node_info::make_contiguous(instruction_ref ins) const
         return mm->add_instruction(make_op("contiguous"), ins);
 }
 
-std::vector<std::size_t> compute_broadcasted_lens(std::vector<std::size_t> s0,
-                                                  std::vector<std::size_t> s1)
-{
-    // Example:
-    // s0 = (3,2,4,5) and s1 = (2,1,1)
-    //
-    // In this case we need to broadcast (:,1,1) portion of
-    // s1 plus broadcast the 1st dimension of s1
-    // giving output_lens = (3,2,4,5)
-    //
-    // Another example:
-    // s0 = (3,2,1,5) and s1 = (2,7,5)
-    // In this case we need to broadcast the (:,:,1:,:) axis
-    // of s0 plus the 1st dimension of s1 giving
-    // output_lens = (3,2,7,5)
-    if(s0.size() > s1.size())
-    {
-        s0.swap(s1);
-    }
-
-    std::vector<std::size_t> out_lens(s1);
-    auto offset = s1.size() - s0.size();
-    std::transform(
-        s0.begin(), s0.end(), s1.begin() + offset, out_lens.begin() + offset, [&](auto a, auto b) {
-            if(a != b and a != 1 and b != 1)
-            {
-                MIGRAPHX_THROW("COMPUTE_BROADCASTLEN: shape {" + to_string_range(s0) + "} and {" +
-                               to_string_range(s1) + "} mismatch!");
-            }
-            return std::max(a, b);
-        });
-
-    return out_lens;
-}
-
 instruction_ref tf_parser::node_info::add_broadcastable_binary_op(const std::string& op_name,
                                                                   instruction_ref arg0,
                                                                   instruction_ref arg1) const
 {
-    if(arg0->get_shape().lens() != arg1->get_shape().lens())
-    {
-        // Get lengths for both arguments
-        auto s0       = arg0->get_shape().lens();
-        auto s1       = arg1->get_shape().lens();
-        auto out_lens = compute_broadcasted_lens(s0, s1);
-
-        auto l0 = arg0;
-        if(arg0->get_shape().lens() != out_lens)
-            l0 = add_instruction(make_op("multibroadcast", {{"output_lens", out_lens}}), arg0);
-
-        auto l1 = arg1;
-        if(arg1->get_shape().lens() != out_lens)
-            l1 = add_instruction(make_op("multibroadcast", {{"output_lens", out_lens}}), arg1);
-
-        return add_instruction(make_op(op_name), l0, l1);
-    }
-    else
-    {
-        return add_instruction(make_op(op_name), {arg0, arg1});
-    }
+    return add_common_op(*mm, make_op(op_name), {arg0, arg1});
 }
 
 int64_t tf_parser::parse_axis(const int64_t dim, const size_t num_dims) const