Compute reshape for nonstandard shape

src/include/migraphx/op/reshape.hpp

@@ -96,10 +96,34 @@ struct reshape
         return {s0.type(), output_dyn_dims};
     }
 
+    template<class Iterator>
+    static auto compute_end_dim(Iterator start, Iterator last, std::size_t dim)
+    {
+        std::size_t x = 1;
+        auto it = std::find_if(start, last, [&](auto i) {
+            x *= i;
+            return x >= dim;
+        });
+        if (x != dim)
+            return start;
+        return it;
+    }
+
+    template<class DimIterator, class StrideIterator>
+    static auto can_strides_merge(DimIterator dim_start, DimIterator dim_last, StrideIterator stride_start, StrideIterator stride_last)
+    {
+        auto cstride = *std::prev(stride_last);
+        return std::equal(std::make_reverse_iterator(dim_last), std::make_reverse_iterator(dim_start+1), std::make_reverse_iterator(stride_last-1), std::make_reverse_iterator(stride_start), [&](auto dim, auto stride) {
+            cstride *= dim;
+            return stride == cstride;
+        });
+    }
+
     shape static_compute_shape(std::vector<shape> inputs, std::size_t n_neg_dims) const
     {
-        check_shapes{inputs, *this}.standard();
+        check_shapes{inputs, *this}.has(1);
         auto&& idims = inputs.front().lens();
+        auto&& istrides = inputs.front().strides();
         std::vector<std::size_t> rdims(dims.begin(), dims.end());
 
         for(std::size_t i = 0; i < dims.size(); i++)
@@ -125,7 +149,81 @@ struct reshape
             }
         }
 
-        shape s{inputs.front().type(), rdims};
+        shape s;
+        if (inputs.front().standard())
+        {
+            s = shape{inputs.front().type(), rdims};
+        }
+        else
+        {
+            std::vector<std::size_t> rstrides;
+            std::size_t i = 0;
+            std::size_t r = 0;
+            while(i < idims.size() and r < rdims.size())
+            {
+                auto idim = idims[i];
+                auto rdim = rdims[r];
+                if (rdim == idim)
+                {
+                    rstrides.push_back(istrides[i]);
+                }
+                // squeeze
+                else if(rdim > idim)
+                {
+                    auto start = idims.begin()+i;
+                    auto it = compute_end_dim(start, idims.end(), rdim);
+                    if (it == start)
+                        break;
+                    auto n = it - start;
+                    if ((i+n) > istrides.size())
+                        break;
+                    if (not can_strides_merge(start, it+1, istrides.begin()+i, istrides.begin()+i+n))
+                        break;
+                    i += n;
+                    rstrides.push_back(istrides[i]);
+                }
+                // unsqueeze
+                else if(rdim < idim)
+                {
+                    auto start = rdims.begin()+i;
+                    auto it = compute_end_dim(start, rdims.end(), idim);
+                    if (it == start)
+                        break;
+                    auto n = it - start;
+                    if ((r+n) > rdims.size())
+                        break;
+                    auto stride = istrides[i] * idim;
+                    std::for_each(start, it+1, [&](auto dim) {
+                        stride /= dim;
+                        rstrides.push_back(stride);
+                    });
+                    r += n;
+                }
+                i++;
+                r++;
+            }
+
+            // Handle trailing 1s
+            if (rstrides.size() < rdims.size() and not rstrides.empty())
+            {
+                auto stride = rstrides.back();
+                for(auto d:range(rdims.begin()+rstrides.size(), rdims.end()))
+                {
+                    if (d != 1)
+                        break;
+                    rstrides.push_back(stride);
+                }
+            }
+
+            if (rdims.size() != rstrides.size())
+                MIGRAPHX_THROW("Reshape on axis that is not standard");
+
+            s = shape{inputs.front().type(), rdims, rstrides};
+
+        }
+
+        assert(s.bytes() == inputs.front().bytes());
+
         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 " +

test/op_shape_test.cpp

@@ -2129,6 +2129,15 @@ TEST_CASE(reshape_shape)
     }
 }
 
+TEST_CASE(reshape_nonstandard_unsqeeze)
+{
+    migraphx::shape input{migraphx::shape::float_type, {4, 24, 1, 1, 1}, {1, 4, 1, 1, 1}};
+    std::vector<std::size_t> lens = {4, 1, 3, 4, 2};
+    std::vector<int64_t> perm = {4, 0, 1, 2, 3};
+    migraphx::shape output = migraphx::shape::from_permutation(migraphx::shape::float_type, lens, migraphx::invert_permutation(perm));
+    expect_shape(output, migraphx::make_op("reshape", {{"dims", lens}}), input);
+}
+
 TEST_CASE(reshape_dyn_shape)
 {
     migraphx::shape input{migraphx::shape::float_type,