Merge branch 'operand-alias'

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

@@ -280,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
@@ -343,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
@@ -424,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
@@ -471,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
@@ -509,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
@@ -560,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
@@ -702,6 +708,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
 {
@@ -744,6 +751,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
@@ -765,6 +773,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
@@ -817,6 +826,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/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/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

test/include/basic_ops.hpp

@@ -79,6 +79,7 @@ struct pass_op
             return {};
         return inputs.front();
     }
+    int output_alias(const std::vector<migraph::shape>&) const { return 0; }
 };
 
 struct pass_standard_op
@@ -103,6 +104,7 @@ struct pass_standard_op
             return {};
         return inputs.front();
     }
+    int output_alias(const std::vector<migraph::shape>&) const { return 0; }
 };
 
 struct nop

test/output_alias.cpp

+#include <migraph/program.hpp>
+#include <migraph/instruction.hpp>
+#include <test.hpp>
+#include <basic_ops.hpp>
+
+void simple_alias()
+{
+    migraph::program p;
+    auto l  = p.add_literal(1);
+    auto p1 = p.add_instruction(pass_op{}, l);
+    EXPECT(bool{migraph::instruction::get_output_alias(l) == l});
+    EXPECT(bool{migraph::instruction::get_output_alias(p1) == l});
+}
+
+void cascade_alias()
+{
+    migraph::program p;
+    auto l  = p.add_literal(1);
+    auto p1 = p.add_instruction(pass_op{}, l);
+    auto p2 = p.add_instruction(pass_op{}, p1);
+    auto p3 = p.add_instruction(pass_op{}, p2);
+    EXPECT(bool{migraph::instruction::get_output_alias(l) == l});
+    EXPECT(bool{migraph::instruction::get_output_alias(p1) == l});
+    EXPECT(bool{migraph::instruction::get_output_alias(p2) == l});
+    EXPECT(bool{migraph::instruction::get_output_alias(p3) == l});
+}
+
+void no_alias()
+{
+    migraph::program p;
+    auto x   = p.add_literal(1);
+    auto y   = p.add_literal(2);
+    auto sum = p.add_instruction(sum_op{}, x, y);
+    EXPECT(bool{migraph::instruction::get_output_alias(sum) == sum});
+}
+
+int main()
+{
+    simple_alias();
+    cascade_alias();
+    no_alias();
+}

tools/include/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,10 +111,34 @@ 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);
+}
+
 <%
  interface(
      'operation',
      virtual('name', returns = 'std::string', const = True),
+     virtual('output_alias',
+             returns = 'int',
+             input   = 'const std::vector<shape>&',
+             const   = True,
+             default = 'output_alias_op'),
      virtual('compute_shape', returns = 'shape', input = 'const std::vector<shape>&', const = True),
      virtual('compute',
              returns = 'argument',