Initial case

src/include/migraphx/op/reshape.hpp

@@ -48,42 +48,97 @@ struct reshape
     std::string name() const { return "reshape"; }
     shape compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.has(1).standard();
-        auto&& idims = inputs.front().lens();
-        std::vector<std::size_t> rdims(dims.begin(), dims.end());
+        check_shapes{inputs, *this}.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");
-
-        for(std::size_t i = 0; i < dims.size(); i++)
+        auto s0 = inputs[0];
+        if(s0.dynamic())
         {
-            if(dims[i] == 0)
-                rdims[i] = idims[i];
-
-            // since rdims using size_t type, -1 is the max value
-            // is size_t that cause later compuation incorrect
-            if(dims[i] == -1)
-                rdims[i] = 1;
+            auto dyn_dims = s0.dyn_dims();
+            std::size_t not_fixed_index = 0;
+            // track number of fixed elements in input and output
+            std::size_t num_dims_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())
+                {
+                    num_dims_ele *= dims[i];
+                    num_dd_ele *= dyn_dims[i].min;
+                }
+                else
+                {
+                    if(not_fixed_index == 0)
+                    {
+                        not_fixed_index = i;
+                    }
+                    else
+                    {
+                        MIGRAPHX_THROW("Reshape: Only support one non-fixed dynamic_dimension");
+                    }
+                }
+            }
+            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));
+            }
+            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");
+            }
+            // construct output dynamic shape from dims attribute
+            std::vector<shape::dynamic_dimension> output_dyn_dims = {};
+            for(std::size_t i = 0; i < dims.size(); ++i)
+            {
+                if(i == not_fixed_index)
+                {
+                    output_dyn_dims.push_back(dyn_dims[not_fixed_index]);
+                }
+                else
+                {
+                    auto d = static_cast<std::size_t>(dims[i]);
+                    output_dyn_dims.push_back({d, d, 0});
+                }
+            }
+            return {s0.type(), output_dyn_dims};
         }
-
-        if(n_neg_dims > 0)
+        else
         {
-            size_t missing_dim =
-                inputs.front().elements() /
-                std::accumulate(rdims.begin(), rdims.end(), 1, std::multiplies<int64_t>());
-            for(std::size_t i = 0; i < rdims.size(); i++)
+            check_shapes{inputs, *this}.standard();
+            auto&& idims = inputs.front().lens();
+            std::vector<std::size_t> rdims(dims.begin(), dims.end());
+
+            for(std::size_t i = 0; i < dims.size(); i++)
             {
+                if(dims[i] == 0)
+                    rdims[i] = idims[i];
+
+                // since rdims using size_t type, -1 is the max value
+                // is size_t that cause later compuation incorrect
                 if(dims[i] == -1)
-                    rdims[i] = missing_dim;
+                    rdims[i] = 1;
             }
-        }
 
-        shape s{inputs.front().type(), rdims};
-        if(s.elements() != inputs.front().elements())
-            MIGRAPHX_THROW("Reshape: Wrong number of elements for reshape: reshape has " +
-                           std::to_string(s.elements()) + " elements whereas the input has " +
-                           std::to_string(inputs.front().elements()));
-        return s;
+            if(n_neg_dims > 0)
+            {
+                size_t missing_dim =
+                    inputs.front().elements() /
+                    std::accumulate(rdims.begin(), rdims.end(), 1, std::multiplies<int64_t>());
+                for(std::size_t i = 0; i < rdims.size(); i++)
+                {
+                    if(dims[i] == -1)
+                        rdims[i] = missing_dim;
+                }
+            }
+
+            shape s{inputs.front().type(), rdims};
+            if(s.elements() != inputs.front().elements())
+                MIGRAPHX_THROW("Reshape: Wrong number of elements for reshape: reshape has " +
+                               std::to_string(s.elements()) + " elements whereas the input has " +
+                               std::to_string(inputs.front().elements()));
+            return s;
+        }
     }
 
     argument compute(shape output_shape, std::vector<argument> args) const

src/onnx/parse_reshape.cpp

@@ -49,7 +49,7 @@ struct parse_reshape : op_parser<parse_reshape>
         if(args.size() == 2)
         {
             auto s = args[1]->eval();
-            check_arg_empty(s, "Reshape: dynamic shape is not supported");
+            check_arg_empty(s, "Reshape: non-constant shape input is not supported");
             s.visit([&](auto v) { copy(v, std::back_inserter(dims)); });
         }