Add tests, fix bug

src/include/migraphx/op/reshape.hpp

@@ -28,6 +28,7 @@
 #include <migraphx/argument.hpp>
 #include <migraphx/config.hpp>
 #include <migraphx/value.hpp>
+#include <migraphx/dyn_output.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -48,18 +49,18 @@ struct reshape
     std::string name() const { return "reshape"; }
     shape compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.has(1);
+        check_shapes{inputs, *this, true}.has(1);
         auto n_neg_dims = std::count(dims.begin(), dims.end(), -1);
         if(n_neg_dims > 1)
             MIGRAPHX_THROW("Reshape: Dimensions for reshape can only have one -1 dim");
         auto s0 = inputs[0];
         if(s0.dynamic())
         {
-            auto dyn_dims = s0.dyn_dims();
-            std::size_t not_fixed_index = 0;
+            auto dyn_dims               = s0.dyn_dims();
+            int not_fixed_index = -1;
             // track number of fixed elements in input and output
             std::size_t num_dims_ele = 1;
-            std::size_t num_dd_ele = 1;
+            std::size_t num_dd_ele   = 1;
             for(std::size_t i = 0; i < dyn_dims.size(); ++i)
             {
                 if(dyn_dims[i].is_fixed())
@@ -69,7 +70,7 @@ struct reshape
                 }
                 else
                 {
-                    if(not_fixed_index == 0)
+                    if(not_fixed_index == -1)
                     {
                         not_fixed_index = i;
                     }
@@ -79,13 +80,16 @@ struct reshape
                     }
                 }
             }
-            if (num_dims_ele != num_dd_ele)
+            if(num_dims_ele != num_dd_ele)
             {
-                MIGRAPHX_THROW("Reshape: Number of fixed elements must match. Input: " + std::to_string(num_dd_ele) + " Output: " + std::to_string(num_dims_ele));
+                MIGRAPHX_THROW("Reshape: Number of fixed elements must match. Input: " +
+                               std::to_string(num_dd_ele) +
+                               " Output: " + std::to_string(num_dims_ele));
             }
             if(dims[not_fixed_index] != 0 and dims[not_fixed_index] != -1)
             {
-                MIGRAPHX_THROW("Reshape: Non-fixed dynamic_dimension doesn't match with 0 or -1 output dimension");
+                MIGRAPHX_THROW("Reshape: Non-fixed dynamic_dimension doesn't match with 0 or -1 "
+                               "output dimension");
             }
             // construct output dynamic shape from dims attribute
             std::vector<shape::dynamic_dimension> output_dyn_dims = {};
@@ -141,9 +145,9 @@ struct reshape
         }
     }
 
-    argument compute(shape output_shape, std::vector<argument> args) const
+    argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
     {
-        return args[0].reshape(output_shape);
+        return args[0].reshape(dyn_out.computed_shape);
     }
 
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }

test/op_shape_test.cpp

@@ -1984,6 +1984,51 @@ TEST_CASE(reshape_shape)
     }
 }
 
+TEST_CASE(reshape_dyn_shape)
+{
+    migraphx::shape input{migraphx::shape::float_type, {{1, 4, 0}, {24, 24, 0}, {1, 1, 0}, {1, 1, 0}}};
+    for(auto&& new_shape :
+        std::vector<std::vector<int64_t>>{{-1, 1, 1, 24}, {0, 8, 3, 1}, {-1, 3, 4, 2}, {0, 2, 4, 3}})
+    {
+        std::vector<migraphx::shape::dynamic_dimension> out_dyn_dims{};
+        for(std::size_t i = 0; i < new_shape.size(); ++i)
+        {
+            if(new_shape[i] == 0 or new_shape[i] == -1)
+            {
+                out_dyn_dims.push_back(input.dyn_dims().at(i));
+            }
+            else
+            {
+                auto d = static_cast<std::size_t>(new_shape[i]);
+                out_dyn_dims.push_back({d, d, 0});
+            }
+        }
+        migraphx::shape output{migraphx::shape::float_type, out_dyn_dims};
+        expect_shape(output, migraphx::make_op("reshape", {{"dims", new_shape}}), input);
+    }
+}
+
+TEST_CASE(reshape_multiple_non_fixed_error)
+{
+    migraphx::shape input{migraphx::shape::float_type, {{1, 4, 0}, {24, 24, 0}, {10, 20, 0}, {1, 1, 0}}};
+    std::vector<int64_t> new_shape = {0, 1, 0, 24};
+    throws_shape(migraphx::make_op("reshape", {{"dims", new_shape}}), input);
+}
+
+TEST_CASE(reshape_fixed_ele_not_matching_error)
+{
+    migraphx::shape input{migraphx::shape::float_type, {{1, 4, 0}, {24, 24, 0}, {10, 10, 0}, {1, 1, 0}}};
+    std::vector<int64_t> new_shape = {0, 1, 5, 24};
+    throws_shape(migraphx::make_op("reshape", {{"dims", new_shape}}), input);
+}
+
+TEST_CASE(reshape_non_fixed_not_matching_error)
+{
+    migraphx::shape input{migraphx::shape::float_type, {{1, 4, 0}, {24, 24, 0}, {1, 1, 0}, {1, 1, 0}}};
+    std::vector<int64_t> new_shape = {2, 1, 1, 24};
+    throws_shape(migraphx::make_op("reshape", {{"dims", new_shape}}), input);
+}
+
 TEST_CASE(rnn)
 {
     {

test/ref_ops_test.cpp

@@ -6045,47 +6045,76 @@ TEST_CASE(relu_dyn_test)
     EXPECT(migraphx::verify_range(results_vector, gold));
 }
 
-TEST_CASE(reshape_test)
+TEST_CASE(reshape_test0)
 {
     migraphx::shape a_shape{migraphx::shape::float_type, {24, 1, 1, 1}};
     std::vector<float> data(24);
     std::iota(data.begin(), data.end(), -3);
-    {
-        migraphx::program p;
-        auto* mm                       = p.get_main_module();
-        auto l                         = mm->add_literal(migraphx::literal{a_shape, data});
-        std::vector<int64_t> new_shape = {8, 3, 1, 1};
-        mm->add_instruction(migraphx::make_op("reshape", {{"dims", new_shape}}), l);
-        p.compile(migraphx::ref::target{});
-        auto result = p.eval({}).back();
-        std::vector<float> results_vector(3);
-        result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
-        EXPECT(migraphx::verify_range(results_vector, data));
-    }
-    {
-        migraphx::program p;
-        auto* mm                       = p.get_main_module();
-        auto l                         = mm->add_literal(migraphx::literal{a_shape, data});
-        std::vector<int64_t> new_shape = {1, 3, 4, 2};
-        mm->add_instruction(migraphx::make_op("reshape", {{"dims", new_shape}}), l);
-        p.compile(migraphx::ref::target{});
-        auto result = p.eval({}).back();
-        std::vector<float> results_vector(3);
-        result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
-        EXPECT(migraphx::verify_range(results_vector, data));
-    }
-    {
-        migraphx::program p;
-        auto* mm                       = p.get_main_module();
-        auto l                         = mm->add_literal(migraphx::literal{a_shape, data});
-        std::vector<int64_t> new_shape = {1, 3, 4, 2};
-        mm->add_instruction(migraphx::make_op("reshape", {{"dims", new_shape}}), l);
-        p.compile(migraphx::ref::target{});
-        auto result = p.eval({}).back();
-        std::vector<float> results_vector(3);
-        result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
-        EXPECT(migraphx::verify_range(results_vector, data));
-    }
+    migraphx::program p;
+    auto* mm                       = p.get_main_module();
+    auto l                         = mm->add_literal(migraphx::literal{a_shape, data});
+    std::vector<int64_t> new_shape = {8, 3, 1, 1};
+    mm->add_instruction(migraphx::make_op("reshape", {{"dims", new_shape}}), l);
+    p.compile(migraphx::ref::target{});
+    auto result = p.eval({}).back();
+    std::vector<float> results_vector{};
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    EXPECT(migraphx::verify_range(results_vector, data));
+}
+
+TEST_CASE(reshape_test1)
+{
+    migraphx::shape a_shape{migraphx::shape::float_type, {24, 1, 1, 1}};
+    std::vector<float> data(24);
+    std::iota(data.begin(), data.end(), -3);
+    migraphx::program p;
+    auto* mm                       = p.get_main_module();
+    auto l                         = mm->add_literal(migraphx::literal{a_shape, data});
+    std::vector<int64_t> new_shape = {1, 3, 4, 2};
+    mm->add_instruction(migraphx::make_op("reshape", {{"dims", new_shape}}), l);
+    p.compile(migraphx::ref::target{});
+    auto result = p.eval({}).back();
+    std::vector<float> results_vector{};
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    EXPECT(migraphx::verify_range(results_vector, data));
+}
+
+TEST_CASE(reshape_test2)
+{
+    migraphx::shape a_shape{migraphx::shape::float_type, {24, 1, 1, 1}};
+    std::vector<float> data(24);
+    std::iota(data.begin(), data.end(), -3);
+    migraphx::program p;
+    auto* mm                       = p.get_main_module();
+    auto l                         = mm->add_literal(migraphx::literal{a_shape, data});
+    std::vector<int64_t> new_shape = {1, 2, 3, 4};
+    mm->add_instruction(migraphx::make_op("reshape", {{"dims", new_shape}}), l);
+    p.compile(migraphx::ref::target{});
+    auto result = p.eval({}).back();
+    std::vector<float> results_vector{};
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    EXPECT(migraphx::verify_range(results_vector, data));
+}
+
+TEST_CASE(reshape_dyn_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape s{migraphx::shape::float_type, {{1, 4, 0}, {24, 24, 0}, {1, 1, 0}, {1, 1, 0}}};
+    std::vector<int64_t> new_shape = {0, 8, 3, 1};
+    auto input = mm->add_parameter("X", s);
+    mm->add_instruction(migraphx::make_op("reshape", {{"dims", new_shape}}), input);
+    p.compile(migraphx::ref::target{});
+    
+    std::vector<float> data(48);
+    std::iota(data.begin(), data.end(), -3);
+    migraphx::parameter_map params;
+    migraphx::shape input_fixed_shape{migraphx::shape::float_type, {2, 24, 1, 1}};
+    params["X"] = migraphx::argument(input_fixed_shape, data.data());
+    auto result = p.eval(params).back();
+    std::vector<float> results_vector{};
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    EXPECT(migraphx::verify_range(results_vector, data));
 }
 
 TEST_CASE(reverse_test_axis0)