Fix bug in eliminate_concat with negative axis (#576)

* Fix bug in eliminate_concat with negative axis

* Formatting

* Fix unused parameter

* Formatting
Co-authored-by: Shucai Xiao <shucai.xiao@amd.com>

src/driver/main.cpp

@@ -132,7 +132,7 @@ struct compiler
         ap(offload_copy,
            {"--enable-offload-copy"},
            ap.help("Enable implicit offload copying"),
-           ap.set_value(false));
+           ap.set_value(true));
         ap(quantize, {"--fp16"}, ap.help("Quantize for fp16"), ap.set_value(q_fp16));
         ap(quantize, {"--int8"}, ap.help("Quantize for int8"), ap.set_value(q_int8));
         ap(fill1, {"--fill1"}, ap.help("Fill parameter with 1s"), ap.append());
@@ -226,9 +226,14 @@ struct verify : command<verify>
     double tolerance     = 80;
     bool per_instruction = false;
     bool reduce          = false;
+    bool offload_copy    = false;
     void parse(argument_parser& ap)
     {
         l.parse(ap);
+        ap(offload_copy,
+           {"--enable-offload-copy"},
+           ap.help("Enable implicit offload copying"),
+           ap.set_value(true));
         ap(tolerance, {"--tolerance"}, ap.help("Tolerance for errors"));
         ap(per_instruction,
            {"-i", "--per-instruction"},
@@ -242,17 +247,20 @@ struct verify : command<verify>
         auto p = l.load();
         std::cout << p << std::endl;
 
+        compile_options options;
+        options.offload_copy = offload_copy;
+
         if(per_instruction)
         {
-            verify_instructions(p, tolerance);
+            verify_instructions(p, options, tolerance);
         }
         else if(reduce)
         {
-            verify_reduced_program(p, tolerance);
+            verify_reduced_program(p, options, tolerance);
         }
         else
         {
-            verify_program(l.file, p, tolerance);
+            verify_program(l.file, p, options, tolerance);
         }
     }
 };

src/driver/verify.cpp

@@ -4,6 +4,7 @@
 #include <migraphx/generate.hpp>
 #include <migraphx/verify_args.hpp>
 #include <migraphx/instruction.hpp>
+#include <migraphx/compile_options.hpp>
 
 #ifdef HAVE_GPU
 #include <migraphx/gpu/target.hpp>
@@ -33,34 +34,39 @@ std::vector<argument> run_cpu(program p)
     return out;
 }
 
-std::vector<argument> run_gpu(program p)
+std::vector<argument> run_gpu(program p, const compile_options& options)
 {
 #ifdef HAVE_GPU
-    p.compile(gpu::target{});
+    p.compile(gpu::target{}, options);
 
     program::parameter_map m;
     for(auto&& x : p.get_parameter_shapes())
     {
-        m[x.first] = gpu::to_gpu(generate_argument(x.second, get_hash(x.first)));
+        auto arg   = generate_argument(x.second, get_hash(x.first));
+        m[x.first] = options.offload_copy ? arg : gpu::to_gpu(arg);
     }
     auto gpu_out = p.eval(m);
     std::vector<argument> output(gpu_out.size());
     std::cout << p << std::endl;
     std::transform(gpu_out.begin(), gpu_out.end(), output.begin(), [&](auto& argu) {
-        return gpu::from_gpu(argu);
+        return options.offload_copy ? argu : gpu::from_gpu(argu);
     });
     return output;
 
 #else
     (void)p;
+    (void)options;
     MIGRAPHX_THROW("Gpu unsupported!");
 #endif
 }
 
-void verify_program(const std::string& name, const program& p, double tolerance)
+void verify_program(const std::string& name,
+                    const program& p,
+                    compile_options options,
+                    double tolerance)
 {
     auto x = run_cpu(p);
-    auto y = run_gpu(p);
+    auto y = run_gpu(p, options);
 
     std::size_t output_num = x.size();
     for(std::size_t i = 0; i < output_num; ++i)
@@ -71,7 +77,7 @@ void verify_program(const std::string& name, const program& p, double tolerance)
     // std::cout << "gpu: " << y << std::endl;
 }
 
-void verify_instructions(const program& prog, double tolerance)
+void verify_instructions(const program& prog, compile_options options, double tolerance)
 {
     for(auto&& ins : prog)
     {
@@ -99,7 +105,7 @@ void verify_instructions(const program& prog, double tolerance)
         {
             std::cout << "Verify: " << ins.name() << std::endl;
             std::cout << p << std::endl;
-            verify_program(ins.name(), p, tolerance);
+            verify_program(ins.name(), p, options, tolerance);
         }
         catch(...)
         {
@@ -109,21 +115,21 @@ void verify_instructions(const program& prog, double tolerance)
     }
 }
 
-void verify_reduced(program p, int n, double tolerance)
+void verify_reduced(program p, int n, compile_options options, double tolerance)
 {
     auto last = std::prev(p.end(), n + 1);
     p.remove_instructions(last, p.end());
     std::cout << "Verify: " << std::endl;
     std::cout << p << std::endl;
-    verify_program(std::to_string(n), p, tolerance);
+    verify_program(std::to_string(n), p, options, tolerance);
 }
 
-void verify_reduced_program(const program& p, double tolerance)
+void verify_reduced_program(const program& p, compile_options options, double tolerance)
 {
     auto n = std::distance(p.begin(), p.end());
     for(std::size_t i = 0; i < n; i++)
     {
-        verify_reduced(p, i, tolerance);
+        verify_reduced(p, i, options, tolerance);
     }
 }
 

src/driver/verify.hpp

@@ -9,9 +9,16 @@ inline namespace MIGRAPHX_INLINE_NS {
 
 std::vector<argument> run_cpu(program p);
 std::vector<argument> run_gpu(program p);
-void verify_program(const std::string& name, const program& p, double tolerance = 100);
-void verify_instructions(const program& prog, double tolerance = 80);
-void verify_reduced_program(const program& p, double tolerance = 80);
+void verify_program(const std::string& name,
+                    const program& p,
+                    compile_options options = compile_options{},
+                    double tolerance        = 100);
+void verify_instructions(const program& prog,
+                         compile_options options = compile_options{},
+                         double tolerance        = 80);
+void verify_reduced_program(const program& p,
+                            compile_options options = compile_options{},
+                            double tolerance        = 80);
 
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace driver

src/eliminate_concat.cpp

@@ -36,7 +36,7 @@ void eliminate_concat::apply(program& p) const
         std::size_t axis_index =
             (concat_op.axis < 0) ? (concat_op.axis + lens.size()) : concat_op.axis;
         if(axis_index == 0 ||
-           std::all_of(lens.begin(), lens.begin() + concat_op.axis, [](auto x) { return x == 1; }))
+           std::all_of(lens.begin(), lens.begin() + axis_index, [](auto x) { return x == 1; }))
         {
             // Last input should be an allocation
             auto last = ins->inputs().back();

src/targets/gpu/hip.cpp

@@ -47,7 +47,7 @@ hip_ptr allocate_gpu(std::size_t sz, bool host = false)
         if(host)
             MIGRAPHX_THROW("Gpu allocation failed: " + hip_error(status));
         else
-            allocate_gpu(sz, true);
+            return allocate_gpu(sz, true);
     }
     return hip_ptr{result};
 }

test/eliminate_concat_test.cpp

@@ -21,6 +21,7 @@ struct concat
     std::string name() const { return "eliminate_concat::concat"; }
     migraphx::shape compute_shape(std::vector<migraphx::shape> inputs) const
     {
+        inputs.pop_back();
         return op.compute_shape(std::move(inputs));
     }
     migraphx::argument compute(migraphx::context&,
@@ -132,6 +133,90 @@ TEST_CASE(simple)
     EXPECT(p1 == p2);
 }
 
+TEST_CASE(negative_axis1)
+{
+    auto create_test_program = [] {
+        migraphx::program p;
+        auto a1          = p.add_instruction(allocate{create_shape(2, 2)});
+        auto p1          = p.add_instruction(simple_op{}, a1);
+        auto a2          = p.add_instruction(allocate{create_shape(2, 2)});
+        auto p2          = p.add_instruction(simple_op{}, a2);
+        std::size_t axis = -1;
+        auto a3          = p.add_instruction(allocate{create_shape(4, 2)});
+        p.add_instruction(concat(axis), p1, p2, a3);
+        return p;
+    };
+    auto create_control_program = create_test_program;
+
+    auto p1 = create_test_program();
+    auto p2 = create_control_program();
+    run_pass(p1);
+
+    EXPECT(p1 == p2);
+}
+
+TEST_CASE(negative_axis2)
+{
+    auto create_test_program = [] {
+        migraphx::program p;
+        auto a1          = p.add_instruction(allocate{create_shape(2, 2)});
+        auto p1          = p.add_instruction(simple_op{}, a1);
+        auto a2          = p.add_instruction(allocate{create_shape(2, 2)});
+        auto p2          = p.add_instruction(simple_op{}, a2);
+        std::size_t axis = -2;
+        auto a3          = p.add_instruction(allocate{create_shape(4, 2)});
+        p.add_instruction(concat(axis), p1, p2, a3);
+        return p;
+    };
+    auto create_control_program = [] {
+        migraphx::program p;
+        auto a1 = p.add_instruction(allocate{create_shape(4, 2)});
+        auto l1 = p.add_instruction(load{create_shape(2, 2), 0}, a1);
+        auto p1 = p.add_instruction(simple_op{}, l1);
+        auto l2 = p.add_instruction(load{create_shape(2, 2), 16}, a1);
+        auto p2 = p.add_instruction(simple_op{}, l2);
+        p.add_instruction(identity{}, a1, p1, p2);
+        return p;
+    };
+
+    auto p1 = create_test_program();
+    auto p2 = create_control_program();
+    run_pass(p1);
+
+    EXPECT(p1 == p2);
+}
+
+TEST_CASE(negative_axis3)
+{
+    auto create_test_program = [] {
+        migraphx::program p;
+        auto a1          = p.add_instruction(allocate{create_shape(1, 2, 2)});
+        auto p1          = p.add_instruction(simple_op{}, a1);
+        auto a2          = p.add_instruction(allocate{create_shape(1, 2, 2)});
+        auto p2          = p.add_instruction(simple_op{}, a2);
+        std::size_t axis = -2;
+        auto a3          = p.add_instruction(allocate{create_shape(1, 4, 2)});
+        p.add_instruction(concat(axis), p1, p2, a3);
+        return p;
+    };
+    auto create_control_program = [] {
+        migraphx::program p;
+        auto a1 = p.add_instruction(allocate{create_shape(1, 4, 2)});
+        auto l1 = p.add_instruction(load{create_shape(1, 2, 2), 0}, a1);
+        auto p1 = p.add_instruction(simple_op{}, l1);
+        auto l2 = p.add_instruction(load{create_shape(1, 2, 2), 16}, a1);
+        auto p2 = p.add_instruction(simple_op{}, l2);
+        p.add_instruction(identity{}, a1, p1, p2);
+        return p;
+    };
+
+    auto p1 = create_test_program();
+    auto p2 = create_control_program();
+    run_pass(p1);
+
+    EXPECT(p1 == p2);
+}
+
 TEST_CASE(reversed)
 {
     auto create_test_program = [] {