Added tests for eliminate_concat

src/eliminate_concat.cpp

@@ -56,8 +56,9 @@ void eliminate_concat::apply(program& p) const
             for(auto x : allocations)
             {
                 migraph::op::load op{x->get_shape(), offset};
+                //migraph::op::load op{x->get_shape(), 0};
                 p.replace_instruction(x, op, {super});
-                offset += x->get_shape().elements();
+                offset += x->get_shape().bytes();
             }
             std::vector<instruction_ref> args = {super};
             std::copy(ins->inputs().begin(), ins->inputs().end() - 1, std::back_inserter(args));

src/opt/memory_coloring_impl.cpp

@@ -193,6 +193,7 @@ void memory_coloring_impl::register_operand_alias()
     operand_alias["transpose"]         = 0;
     operand_alias["flatten"]           = 0;
     operand_alias["broadcast"]         = 0;
+    operand_alias["identity"]          = 0;
     operand_alias["reshape"]           = 0;
     operand_alias["pass"]              = 0;
 }

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

+#ifndef MIGRAPH_GUARD_RTGLIB_CONCAT_GPU_OPT_HPP
+#define MIGRAPH_GUARD_RTGLIB_CONCAT_GPU_OPT_HPP
+
+#include <migraph/gpu/concat.hpp>
+
+namespace migraph {
+namespace gpu {
+
+struct concat_gpu_optimization
+{
+    std::string name() const { return "gpu::concat"; }
+    std::string allocate() const { return "hip::allocate"; }
+    migraph::op::concat get_concat(const migraph::operation& op) const
+    {
+        return migraph::any_cast<migraph::gpu::hip_concat>(op).op;
+    }
+};
+
+} // namespace gpu
+
+} // namespace migraph
+
+#endif

src/targets/gpu/target.cpp

@@ -15,6 +15,8 @@
 #include <migraph/eliminate_contiguous.hpp>
 #include <migraph/common_subexpression_elimination.hpp>
 #include <migraph/fwd_conv_batchnorm_rewrite.hpp>
+#include <migraph/eliminate_concat.hpp>
+#include <migraph/gpu/concat_gpu_opt.hpp>
 
 namespace migraph {
 namespace gpu {
@@ -38,6 +40,8 @@ std::vector<pass> target::get_passes(migraph::context& gctx) const
         simplify_reshapes{},
         dead_code_elimination{},
         lowering{ctx},
+        eliminate_concat{concat_gpu_optimization{}},
+        dead_code_elimination{},
         eliminate_contiguous{},
         dead_code_elimination{},
         fuse_ops{&ctx},

test/eliminate_concat_test.cpp

@@ -106,10 +106,10 @@ void basic()
             migraph::op::load{migraph::shape{migraph::shape::float_type, {1, 2, 8, 8}}, 0}, {a1});
         auto p1 = p.add_instruction(fred_op{}, l1);
         auto l2 = p.add_instruction(
-            migraph::op::load{migraph::shape{migraph::shape::float_type, {1, 3, 8, 8}}, 128}, {a1});
+            migraph::op::load{migraph::shape{migraph::shape::float_type, {1, 3, 8, 8}}, 512}, {a1});
         auto p2 = p.add_instruction(fred_op{}, l2);
         auto l3 = p.add_instruction(
-            migraph::op::load{migraph::shape{migraph::shape::float_type, {1, 5, 8, 8}}, 320}, {a1});
+            migraph::op::load{migraph::shape{migraph::shape::float_type, {1, 5, 8, 8}}, 1280}, {a1});
         auto p3 = p.add_instruction(fred_op{}, l3);
         auto i1 = p.add_instruction(migraph::op::identity{}, {a1, p1, p2, p3});
         return p;

test/gpu/miopen.cpp

@@ -611,6 +611,73 @@ struct test_concat2
     }
 };
 
+struct test_concat_relu
+{
+    migraph::program create_program() const
+    {
+        migraph::program p;
+        std::size_t axis = 0;
+        migraph::shape s0{migraph::shape::float_type, {2, 2}};
+        migraph::shape s1{migraph::shape::float_type, {3, 2}};
+        migraph::shape s2{migraph::shape::float_type, {1, 2}};
+        auto l0 = p.add_parameter("x", s0);
+        auto l1 = p.add_parameter("y", s1);
+        auto l2 = p.add_parameter("z", s2);
+        auto r0 = p.add_instruction(migraph::op::activation{"relu"}, l0);
+        auto r1 = p.add_instruction(migraph::op::activation{"relu"}, l1);
+        auto r2 = p.add_instruction(migraph::op::activation{"relu"}, l2);
+        auto c0 = p.add_instruction(migraph::op::concat{axis}, r0, r1, r2);
+        p.add_instruction(migraph::op::activation{"relu"}, c0);
+        return p;
+    }
+};
+
+void manual_identity()
+{
+    migraph::program p;
+    std::vector<float> data0 = {0, 1, 2, 3};
+    migraph::shape s0{migraph::shape::float_type, {2, 2}};
+    auto l0 = p.add_literal(migraph::literal{s0, data0});
+    p.add_instruction(migraph::op::identity{}, l0);
+    p.compile(migraph::gpu::target{});
+    migraph::program::parameter_map m;
+    for(auto&& x : p.get_parameter_shapes())
+    {
+        m[x.first] = migraph::gpu::to_gpu(migraph::generate_argument(x.second));
+    }
+    auto result = migraph::gpu::from_gpu(p.eval(m));
+    std::cout << result << std::endl;
+}
+
+void manual_test_concat_relu()
+{
+    migraph::program p;
+    std::size_t axis = 0;
+    std::vector<float> data0 = {0, 1, 2, 3};
+    std::vector<float> data1 = {4, 5, 6, 7, 8, 9};
+    std::vector<float> data2 = {10, 11};
+    migraph::shape s0{migraph::shape::float_type, {2, 2}};
+    migraph::shape s1{migraph::shape::float_type, {3, 2}};
+    migraph::shape s2{migraph::shape::float_type, {1, 2}};
+    auto l0 = p.add_literal(migraph::literal{s0, data0});
+    auto l1 = p.add_literal(migraph::literal{s1, data1});
+    auto l2 = p.add_literal(migraph::literal{s2, data2});
+    auto r0 = p.add_instruction(migraph::op::activation{"relu"}, l0);
+    auto r1 = p.add_instruction(migraph::op::activation{"relu"}, l1);
+    auto r2 = p.add_instruction(migraph::op::activation{"relu"}, l2);
+    auto c0 = p.add_instruction(migraph::op::concat{axis}, r0, r1, r2);
+    p.add_instruction(migraph::op::activation{"relu"}, c0);
+
+    p.compile(migraph::gpu::target{});
+    migraph::program::parameter_map m;
+    for(auto&& x : p.get_parameter_shapes())
+    {
+        m[x.first] = migraph::gpu::to_gpu(migraph::generate_argument(x.second));
+    }
+    auto result = migraph::gpu::from_gpu(p.eval(m));
+    std::cout << result << std::endl;
+}
+
 struct test_conv_bn_relu_pooling2
 {
     static migraph::instruction_ref
@@ -651,6 +718,7 @@ int main()
 {
     verify_program<test_concat>();
     verify_program<test_concat2>();
+    verify_program<test_concat_relu>();
     verify_program<test_add>();
     verify_program<test_triadd>();
     verify_program<test_triadd2>();