Enable reshape on nonstandard shapes

src/include/migraphx/op/reshape.hpp

@@ -96,10 +96,42 @@ 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)
+    {
+        assert(std::distance(dim_start, dim_last) == std::distance(stride_start, 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();
-        auto&& idims = inputs.front().lens();
+        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,11 +157,85 @@ 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 + 1))
+                        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};
+        }
+
         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()));
+        
+        assert(s.bytes() == inputs.front().bytes());
         return s;
     }
 

test/op_shape_test.cpp

@@ -2161,6 +2161,8 @@ TEST_CASE(reshape_shape)
     for(auto&& new_shape :
         std::vector<std::vector<int64_t>>{{8, 3, 2, 2}, {1, 3, -1, -1}, {3, 0}, {3, 2}})
     {
+        std::cout << "input: " << input << std::endl;
+        std::cout << "dims: " << migraphx::to_string_range(new_shape) << std::endl;
         throws_shape(migraphx::make_op("reshape", {{"dims", new_shape}}), input);
     }
 
@@ -2181,6 +2183,34 @@ TEST_CASE(reshape_shape)
     }
 }
 
+TEST_CASE(reshape_nonstandard_unsqueeze)
+{
+    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_nonstandard_squeeze)
+{
+    migraphx::shape input{migraphx::shape::float_type, {2, 16, 16, 1280}, {327680, 16, 1, 256}};
+    std::vector<std::size_t> lens = {2, 256, 1280};
+    std::vector<int64_t> perm     = {0, 2, 1};
+    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_broadcast_squeeze)
+{
+    migraphx::shape input{migraphx::shape::float_type, {2, 16, 16, 1280}, {0, 0, 0, 1}};
+    migraphx::shape output{migraphx::shape::float_type, {2, 256, 1280}, {0, 0, 1}};
+    expect_shape(output, migraphx::make_op("reshape", {{"dims", output.lens()}}), input);
+}
+
 TEST_CASE(reshape_dyn_shape)
 {
     migraphx::shape input{migraphx::shape::float_type, {{1, 4}, {24, 24}, {1, 1}, {1, 1}}};