Merge branch 'develop' into unet

.github/workflows/ci.yaml

@@ -147,10 +147,14 @@ jobs:
         os:
           - ubuntu-16.04
           - ubuntu-18.04
+          - ubuntu-20.04
         configuration:
           - debug
           - release
           - codecov
+        exclude:
+          - os: ubuntu-16.04
+            configuration: debug
 
     steps:
     - uses: actions/checkout@v2

CMakeLists.txt

@@ -36,7 +36,7 @@ find_package(nlohmann_json 3.8.0 REQUIRED)
 
 include(ROCMSetupVersion)
 
-rocm_setup_version(VERSION 1.2)
+rocm_setup_version(VERSION 1.3)
 set(MIGRAPHX_SO_VERSION ${PROJECT_VERSION_MAJOR}.${PROJECT_VERSION_MINOR})
 
 option( BUILD_SHARED_LIBS "Build as a shared library" ON )

src/CMakeLists.txt

@@ -9,57 +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/argument.cpp

@@ -124,6 +124,14 @@ argument::data_t argument::data_t::from_args(const std::vector<argument>& args)
     return result;
 }
 
+argument argument::copy() const
+{
+    argument result{this->get_shape()};
+    auto* src = this->data();
+    std::copy(src, src + this->get_shape().bytes(), result.data());
+    return result;
+}
+
 argument argument::share() const { return {m_shape, m_data.share()}; }
 
 std::vector<argument> argument::get_sub_objects() const

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/compile_src.cpp

@@ -32,7 +32,7 @@ std::vector<char> src_compiler::compile(const std::vector<src_file>& srcs) const
         }
     }
 
-    params += " -o" + out;
+    params += " -o " + out;
 
     td.execute(compiler, params);
 

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/allocation_model.hpp

@@ -26,6 +26,8 @@ struct allocation_model
     std::string copy() const;
     /// Create an allocation operator for the given shape
     operation allocate(const shape& s) const;
+    /// Create a preallocated operator for the given shape
+    operation preallocate(const shape& s, const std::string& id) const;
 };
 
 #else
@@ -38,6 +40,7 @@ struct allocation_model
  *      std::string name() const;
  *      std::string copy() const;
  *      operation allocate(const shape& s) const;
+ *      operation preallocate(const shape& s,std::string id) const;
  * };
  *
  */
@@ -123,6 +126,12 @@ struct allocation_model
         return (*this).private_detail_te_get_handle().allocate(s);
     }
 
+    operation preallocate(const shape& s, std::string id) const
+    {
+        assert((*this).private_detail_te_handle_mem_var);
+        return (*this).private_detail_te_get_handle().preallocate(s, std::move(id));
+    }
+
     friend bool is_shared(const allocation_model& private_detail_x,
                           const allocation_model& private_detail_y)
     {
@@ -137,9 +146,10 @@ struct allocation_model
         virtual std::shared_ptr<private_detail_te_handle_base_type> clone() const = 0;
         virtual const std::type_info& type() const                                = 0;
 
-        virtual std::string name() const                 = 0;
-        virtual std::string copy() const                 = 0;
-        virtual operation allocate(const shape& s) const = 0;
+        virtual std::string name() const                                    = 0;
+        virtual std::string copy() const                                    = 0;
+        virtual operation allocate(const shape& s) const                    = 0;
+        virtual operation preallocate(const shape& s, std::string id) const = 0;
     };
 
     template <typename PrivateDetailTypeErasedT>
@@ -180,6 +190,12 @@ struct allocation_model
             return private_detail_te_value.allocate(s);
         }
 
+        operation preallocate(const shape& s, std::string id) const override
+        {
+
+            return private_detail_te_value.preallocate(s, std::move(id));
+        }
+
         PrivateDetailTypeErasedT private_detail_te_value;
     };
 

src/include/migraphx/argument.hpp

@@ -60,6 +60,8 @@ struct argument : raw_data<argument>
 
     argument reshape(const shape& s) const;
 
+    argument copy() const;
+
     /// Make copy of the argument that is always sharing the data
     argument share() const;
 

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/include/migraphx/lifetime.hpp

+#ifndef MIGRAPHX_GUARD_MIGRAPHX_LIFETIME_HPP
+#define MIGRAPHX_GUARD_MIGRAPHX_LIFETIME_HPP
+
+#include <migraphx/config.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+enum class lifetime
+{
+    local,
+    global,
+    borrow
+};
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+#endif // MIGRAPHX_GUARD_MIGRAPHX_LIFETIME_HPP

src/include/migraphx/matcher.hpp

@@ -626,7 +626,8 @@ auto tree(M main_op, Ms... ms)
             if(idx != leafs.size())
                 return nullopt;
             // Use explicit captures to workaround ICE on gcc
-            bool found = sequence_c<sizeof...(Ms)>([&ms..., &ctx, &leafs](auto... is) {
+            // Capture by value to workaround compile error on gcc 9
+            bool found = sequence_c<sizeof...(Ms)>([ms..., &ctx, &leafs](auto... is) {
                 return fold(lazy_and{})(ctx.lazy_match(ms, leafs[is])...)();
             });
             if(not found)

src/include/migraphx/op/as_shape.hpp

@@ -8,6 +8,7 @@
 #include <migraphx/literal.hpp>
 #include <migraphx/shape_for_each.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/lifetime.hpp>
 #include <cmath>
 #include <utility>
 
@@ -35,7 +36,7 @@ struct as_shape
     {
         return args.front().reshape(output_shape);
     }
-    bool is_borrowed() const { return true; }
+    lifetime get_lifetime() const { return lifetime::borrow; }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 

src/include/migraphx/op/broadcast.hpp

@@ -6,6 +6,7 @@
 #include <migraphx/argument.hpp>
 #include <migraphx/functional.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/lifetime.hpp>
 #include <cmath>
 #include <utility>
 
@@ -66,7 +67,7 @@ struct broadcast
     {
         return args[0].reshape(output_shape);
     }
-    bool is_borrowed() const { return true; }
+    lifetime get_lifetime() const { return lifetime::borrow; }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 

src/include/migraphx/op/flatten.hpp

@@ -10,6 +10,7 @@
 #include <migraphx/config.hpp>
 #include <migraphx/value.hpp>
 #include <migraphx/op/normalize_attribute.hpp>
+#include <migraphx/lifetime.hpp>
 #include <cmath>
 #include <utility>
 
@@ -50,7 +51,7 @@ struct flatten
     {
         return args[0].reshape(output_shape);
     }
-    bool is_borrowed() const { return true; }
+    lifetime get_lifetime() const { return lifetime::borrow; }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 

src/include/migraphx/op/load.hpp

@@ -6,6 +6,7 @@
 #include <migraphx/functional.hpp>
 #include <migraphx/argument.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/lifetime.hpp>
 #include <cmath>
 #include <utility>
 
@@ -36,7 +37,7 @@ struct load
             MIGRAPHX_THROW("Load access is out of bounds");
         return argument::load(s, args[0].data() + offset);
     }
-    bool is_borrowed() const { return true; }
+    lifetime get_lifetime() const { return lifetime::borrow; }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 
     friend std::ostream& operator<<(std::ostream& os, const load& op)

src/include/migraphx/op/multibroadcast.hpp

@@ -8,6 +8,7 @@
 #include <migraphx/literal.hpp>
 #include <migraphx/shape_for_each.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/lifetime.hpp>
 #include <cmath>
 #include <utility>
 
@@ -68,7 +69,7 @@ struct multibroadcast
     {
         return args[0].reshape(output_shape);
     }
-    bool is_borrowed() const { return true; }
+    lifetime get_lifetime() const { return lifetime::borrow; }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 

src/include/migraphx/op/prefix_scan_op.hpp

@@ -43,7 +43,7 @@ struct prefix_scan_op : op_name<Derived>
 
     argument compute(const shape&, std::vector<argument> args) const
     {
-        argument result = args[0];
+        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();

src/include/migraphx/op/reshape.hpp

@@ -8,6 +8,7 @@
 #include <migraphx/literal.hpp>
 #include <migraphx/shape_for_each.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/lifetime.hpp>
 #include <cmath>
 #include <utility>
 
@@ -71,7 +72,7 @@ struct reshape
         return args[0].reshape(output_shape);
     }
 
-    bool is_borrowed() const { return true; }
+    lifetime get_lifetime() const { return lifetime::borrow; }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };