Merge branch 'master' into remove_concat

src/include/migraph/instruction.hpp

@@ -69,6 +69,8 @@ struct instruction
     static void
     replace(instruction_ref ins, operation o, const shape& r, std::vector<instruction_ref> args);
 
+    static instruction_ref get_output_alias(instruction_ref ins);
+
     private:
     // internal
     void replace(operation o, const shape& r, std::vector<instruction_ref> args);

src/include/migraph/operation.hpp

@@ -43,6 +43,9 @@ struct operation
      * the same the `output` shape.
      */
     argument compute(context& ctx, const shape& output, const std::vector<argument>& input) const;
+    /// An optional method to return which argument the output will alias. If
+    /// there is no aliased output then -1 can be returned.
+    int output_alias(const std::vector<shape>& input) const;
     /// An optional stream operator to print the operation. When this is not
     /// implemented, it will just print the operation's name.
     friend std::ostream& operator<<(std::ostream& os, const operation& op);
@@ -108,12 +111,32 @@ compute_op(const T& x, context& ctx, const shape& output_shape, const std::vecto
     return compute_op(rank<1>{}, x, ctx, output_shape, input);
 }
 
+template <class T>
+int output_alias_op(rank<0>, const T&, const std::vector<shape>&)
+{
+    return -1;
+}
+
+template <class T>
+auto output_alias_op(rank<1>, const T& x, const std::vector<shape>& shapes)
+    -> decltype(x.output_alias(shapes))
+{
+    return x.output_alias(shapes);
+}
+
+template <class T>
+int output_alias_op(const T& x, const std::vector<shape>& shapes)
+{
+    return output_alias_op(rank<1>{}, x, shapes);
+}
+
 /*
  * Type-erased interface for:
  *
  * struct operation
  * {
  *      std::string name() const;
+ *      int output_alias(const std::vector<shape>& input) const;
  *      shape compute_shape(const std::vector<shape>& input) const;
  *      argument compute(context& ctx,const shape& output,const std::vector<argument>& input) const;
  *     friend std::ostream & operator<<(std::ostream & os,const operation & op) ;
@@ -185,6 +208,12 @@ struct operation
         return (*this).private_detail_te_get_handle().name();
     }
 
+    int output_alias(const std::vector<shape>& input) const
+    {
+        assert((*this).private_detail_te_handle_mem_var);
+        return (*this).private_detail_te_get_handle().output_alias(input);
+    }
+
     shape compute_shape(const std::vector<shape>& input) const
     {
         assert((*this).private_detail_te_handle_mem_var);
@@ -217,6 +246,7 @@ struct operation
         virtual const std::type_info& type() const                                = 0;
 
         virtual std::string name() const                                   = 0;
+        virtual int output_alias(const std::vector<shape>& input) const    = 0;
         virtual shape compute_shape(const std::vector<shape>& input) const = 0;
         virtual argument
         compute(context& ctx, const shape& output, const std::vector<argument>& input) const = 0;
@@ -254,8 +284,15 @@ struct operation
 
         std::string name() const override { return private_detail_te_value.name(); }
 
+        int output_alias(const std::vector<shape>& input) const override
+        {
+
+            return output_alias_op(private_detail_te_value, input);
+        }
+
         shape compute_shape(const std::vector<shape>& input) const override
         {
+
             return private_detail_te_value.compute_shape(input);
         }
 

src/include/migraph/operators.hpp

@@ -223,22 +223,6 @@ struct pooling
     }
 };
 
-struct activation
-{
-    std::string mode;
-    std::string name() const { return "activation"; }
-    shape compute_shape(std::vector<shape> inputs) const
-    {
-        check_shapes{inputs, *this}.has(1);
-        return inputs.front();
-    }
-    friend std::ostream& operator<<(std::ostream& os, const activation& op)
-    {
-        os << op.name() << ":" << op.mode;
-        return os;
-    }
-};
-
 struct leaky_relu
 {
     std::string name() const { return "leaky_relu"; }
@@ -296,6 +280,7 @@ struct transpose
     {
         return {std::move(output_shape), std::move(args.front().data)};
     }
+    int output_alias(const std::vector<shape>&) const { return 0; }
 };
 
 struct contiguous
@@ -359,6 +344,7 @@ struct concat
         new_lens[axis] = new_dim_axis;
         return {type, new_lens};
     }
+    int output_alias(const std::vector<shape>&) const { return 0; }
 };
 
 struct slice
@@ -440,6 +426,7 @@ struct slice
         auto offset = compute_offset(input.get_shape()) * output_shape.type_size();
         return {std::move(output_shape), [=] { return input.data() + offset; }};
     }
+    int output_alias(const std::vector<shape>&) const { return 0; }
 };
 
 struct squeeze
@@ -487,6 +474,7 @@ struct squeeze
     {
         return {std::move(output_shape), std::move(args.front().data)};
     }
+    int output_alias(const std::vector<shape>&) const { return 0; }
 };
 
 struct unsqueeze
@@ -525,6 +513,7 @@ struct unsqueeze
     {
         return {std::move(output_shape), std::move(args.front().data)};
     }
+    int output_alias(const std::vector<shape>&) const { return 0; }
 };
 
 struct reshape
@@ -576,6 +565,7 @@ struct reshape
     {
         return {std::move(output_shape), std::move(args.front().data)};
     }
+    int output_alias(const std::vector<shape>&) const { return 0; }
 };
 
 struct dot
@@ -678,6 +668,11 @@ struct neg : unary
     std::string name() const { return "neg"; }
 };
 
+struct relu : unary
+{
+    std::string name() const { return "relu"; }
+};
+
 struct softmax
 {
     std::string name() const { return "softmax"; }
@@ -718,6 +713,7 @@ struct flatten
     {
         return {std::move(output_shape), std::move(args.front().data)};
     }
+    int output_alias(const std::vector<shape>&) const { return 0; }
 };
 struct broadcast
 {
@@ -760,6 +756,7 @@ struct broadcast
     {
         return {std::move(output_shape), std::move(args.at(0).data)};
     }
+    int output_alias(const std::vector<shape>&) const { return 0; }
 };
 
 struct scalar
@@ -781,6 +778,7 @@ struct scalar
     {
         return {std::move(output_shape), std::move(args.at(0).data)};
     }
+    int output_alias(const std::vector<shape>&) const { return 0; }
 };
 
 struct binary
@@ -833,6 +831,7 @@ struct load
     {
         return {s, args[0].data() + offset};
     }
+    int output_alias(const std::vector<shape>&) const { return 0; }
 };
 
 struct outline

src/instruction.cpp

@@ -161,12 +161,25 @@ void instruction::replace_argument(instruction_ref old, instruction_ref new_ins)
     old->remove_output(*this);
 }
 
-shape compute_shape(const operation& op, const std::vector<instruction_ref>& args)
+std::vector<shape> compute_shapes(const std::vector<instruction_ref>& args)
 {
     std::vector<shape> shapes(args.size());
     std::transform(
         args.begin(), args.end(), shapes.begin(), [](instruction_ref i) { return i->get_shape(); });
-    return op.compute_shape(shapes);
+    return shapes;
+}
+
+instruction_ref instruction::get_output_alias(instruction_ref ins)
+{
+    auto i = ins->get_operator().output_alias(compute_shapes(ins->inputs()));
+    if(i < 0)
+        return ins;
+    return get_output_alias(ins->inputs().at(i));
+}
+
+shape compute_shape(const operation& op, const std::vector<instruction_ref>& args)
+{
+    return op.compute_shape(compute_shapes(args));
 }
 
 } // namespace migraph

src/onnx/onnx.cpp

@@ -52,7 +52,7 @@ struct onnx_parser
         add_generic_op("Div", op::div{});
         add_generic_op("MatMul", op::dot{});
         add_generic_op("Mul", op::mul{});
-        add_generic_op("Relu", op::activation{"relu"});
+        add_generic_op("Relu", op::relu{});
         add_generic_op("Sub", op::sub{});
         add_generic_op("Sum", op::add{});
 

src/program.cpp

@@ -281,7 +281,7 @@ void program::compile(const target& t, tracer trace)
 {
     assert(this->validate() == impl->instructions.end());
     this->impl->ctx = t.get_context();
-    if(not trace.enabled() and enabled(MIGRAPH_TRACE_COMPILE{}))
+    if(not trace.enabled() or enabled(MIGRAPH_TRACE_COMPILE{}))
         trace = tracer{std::cout};
     trace(*this);
     trace();

src/targets/cpu/lowering.cpp

@@ -606,6 +606,7 @@ struct cpu_apply
         apply_map["sin"]        = simple_op<cpu_unary<sin_op>>();
         apply_map["cos"]        = simple_op<cpu_unary<cos_op>>();
         apply_map["tan"]        = simple_op<cpu_unary<tan_op>>();
+        apply_map["relu"]       = simple_op<cpu_unary<relu_op>>();
         apply_map["add"]        = simple_op<cpu_binary<add_op>>();
         apply_map["sub"]        = simple_op<cpu_binary<sub_op>>();
         apply_map["mul"]        = simple_op<cpu_binary<mul_op>>();
@@ -619,11 +620,7 @@ struct cpu_apply
         init();
         for(auto it : iterator_for(*prog))
         {
-            if(it->name() == "activation")
-            {
-                apply_activation(it);
-            }
-            else if(it->name() == "pooling")
+            if(it->name() == "pooling")
             {
                 apply_pooling(it);
             }
@@ -647,13 +644,6 @@ struct cpu_apply
         prog->replace_instruction(ins, T{op}, ins->inputs());
     }
 
-    void apply_activation(instruction_ref ins)
-    {
-        auto&& op = any_cast<op::activation>(ins->get_operator());
-        if(op.mode == "relu")
-            prog->replace_instruction(ins, cpu_unary<relu_op>{}, ins->inputs());
-    }
-
     void apply_pooling(instruction_ref ins)
     {
         auto&& op = any_cast<op::pooling>(ins->get_operator());

src/targets/gpu/include/migraph/gpu/add.hpp

@@ -25,6 +25,7 @@ struct hip_add
     std::string name() const { return "gpu::add"; }
     shape compute_shape(const std::vector<shape>& inputs) const;
     argument compute(context&, const shape&, const std::vector<argument>& args) const;
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
 };
 
 struct miopen_add
@@ -33,6 +34,7 @@ struct miopen_add
     shape compute_shape(const std::vector<shape>& inputs) const;
     argument
     compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
 };
 
 } // namespace gpu

src/targets/gpu/include/migraph/gpu/batchnorm.hpp

@@ -27,6 +27,7 @@ struct miopen_batch_norm_inference
     shape compute_shape(const std::vector<shape>& inputs) const;
     argument
     compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
 };
 
 } // namespace gpu

src/targets/gpu/include/migraph/gpu/concat.hpp

@@ -28,6 +28,7 @@ struct hip_concat
     shape compute_shape(std::vector<shape> inputs) const;
     argument
     compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
 };
 
 } // namespace gpu

src/targets/gpu/include/migraph/gpu/contiguous.hpp

@@ -26,6 +26,7 @@ struct miopen_contiguous
     std::string name() const { return "gpu::contiguous"; }
     shape compute_shape(const std::vector<shape>& inputs) const;
     argument compute(context&, shape output_shape, const std::vector<argument>& args) const;
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
 };
 
 } // namespace gpu

src/targets/gpu/include/migraph/gpu/convolution.hpp

@@ -38,6 +38,7 @@ struct miopen_convolution
     argument
     compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
     shape compile(context& ctx, const shape& output_shape, std::vector<instruction_ref> inputs);
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
 };
 
 } // namespace gpu

src/targets/gpu/include/migraph/gpu/gemm.hpp

@@ -27,6 +27,7 @@ struct miopen_gemm
     shape compute_shape(const std::vector<shape>& inputs) const;
     argument
     compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
 };
 
 } // namespace gpu

src/targets/gpu/include/migraph/gpu/hip.hpp

@@ -67,6 +67,7 @@ struct hip_write
     {
         return to_gpu(args.front());
     }
+    int output_alias(const std::vector<shape>&) const { return 0; }
 };
 
 struct hip_copy
@@ -82,6 +83,7 @@ struct hip_copy
         copy_to_gpu(args[0], args[1]);
         return args[1];
     }
+    int output_alias(const std::vector<shape>&) const { return 1; }
 };
 } // namespace gpu
 } // namespace migraph

src/targets/gpu/include/migraph/gpu/leaky_relu.hpp

@@ -27,6 +27,7 @@ struct miopen_leaky_relu
     shape compute_shape(const std::vector<shape>& inputs) const;
     argument
     compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
 };
 
 } // namespace gpu

src/targets/gpu/include/migraph/gpu/mul.hpp

@@ -25,6 +25,7 @@ struct hip_mul
     std::string name() const { return "gpu::mul"; }
     shape compute_shape(const std::vector<shape>& inputs) const;
     argument compute(context&, const shape&, const std::vector<argument>& args) const;
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
 };
 
 } // namespace gpu

src/targets/gpu/include/migraph/gpu/pooling.hpp

@@ -29,6 +29,7 @@ struct miopen_pooling
     shape compute_shape(const std::vector<shape>& inputs) const;
     argument
     compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
 };
 
 } // namespace gpu

src/targets/gpu/include/migraph/gpu/relu.hpp

@@ -27,6 +27,7 @@ struct miopen_relu
     shape compute_shape(const std::vector<shape>& inputs) const;
     argument
     compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
 };
 
 } // namespace gpu

src/targets/gpu/include/migraph/gpu/softmax.hpp

@@ -27,6 +27,7 @@ struct miopen_softmax
     shape compute_shape(const std::vector<shape>& inputs) const;
     argument
     compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
 };
 
 } // namespace gpu

src/targets/gpu/lowering.cpp

@@ -50,9 +50,9 @@ struct miopen_apply
             {
                 check_shape(s, apply_convolution(it));
             }
-            else if(it->name() == "activation")
+            else if(it->name() == "relu")
             {
-                check_shape(s, apply_activation(it));
+                check_shape(s, apply_relu(it));
             }
             else if(it->name() == "leaky_relu")
             {
@@ -131,17 +131,13 @@ struct miopen_apply
             ins, miopen_pooling{op, std::move(pd)}, ins->inputs().at(0), output);
     }
 
-    instruction_ref apply_activation(instruction_ref ins)
+    instruction_ref apply_relu(instruction_ref ins)
     {
-        auto&& op = any_cast<op::activation>(ins->get_operator());
-        auto ad   = make_relu();
-        if(op.mode == "relu")
-        {
-            auto output = insert_allocation(ins, ins->get_shape());
-            return prog->replace_instruction(
-                ins, miopen_relu{std::move(ad)}, ins->inputs().at(0), output);
-        }
-        return ins;
+        auto ad = make_relu();
+
+        auto output = insert_allocation(ins, ins->get_shape());
+        return prog->replace_instruction(
+            ins, miopen_relu{std::move(ad)}, ins->inputs().at(0), output);
     }
 
     instruction_ref apply_leaky_relu(instruction_ref ins)