Handle more reshapes with the shape transform descriptor (#3527)

ROCm · Oct 21, 2024 · c1d2d5a · c1d2d5a
1 parent fdd28ca
commit c1d2d5a
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 0 deletions.
diff --git a/src/include/migraphx/shape_transform_descriptor.hpp b/src/include/migraphx/shape_transform_descriptor.hpp
@@ -76,6 +76,7 @@ struct MIGRAPHX_EXPORT shape_transform_descriptor
 
     bool apply(const std::vector<operation>& ops);
     bool apply_reshape(const std::vector<std::size_t>& rdims);
+    bool apply_reshape_impl(const std::vector<std::size_t>& rdims);
     bool apply_transpose(const std::vector<std::int64_t>& permutation);
     bool apply_broadcast(const std::vector<std::size_t>& out_lens,
                          optional<std::size_t> axis = nullopt);

diff --git a/src/shape_transform_descriptor.cpp b/src/shape_transform_descriptor.cpp
@@ -143,8 +143,22 @@ bool shape_transform_descriptor::apply(const std::vector<operation>& ops)
     return true;
 }
 bool shape_transform_descriptor::apply_reshape(const std::vector<std::size_t>& rdims)
+{
+    std::vector<std::size_t> idims;
+    transform(get_all_subdimensions(dimensions),
+              std::back_inserter(idims),
+              std::mem_fn(&dimension::sub::len));
+    auto cdims = common_dims::compute(idims, rdims).dims;
+    if(not cdims.empty() and not apply_reshape_impl(cdims))
+        return false;
+    return apply_reshape_impl(rdims);
+}
+bool shape_transform_descriptor::apply_reshape_impl(const std::vector<std::size_t>& rdims)
 {
     assert(migraphx::elements(rdims) == this->elements());
+    if(migraphx::equal(
+           dimensions, rdims, [](const dimension& d, std::size_t rdim) { return d.len() == rdim; }))
+        return true;
     std::vector<dimension> new_dims;
     auto subs     = get_all_subdimensions(dimensions);
     std::size_t i = 0;

diff --git a/test/shape_transform_descriptor.cpp b/test/shape_transform_descriptor.cpp
@@ -134,6 +134,15 @@ TEST_CASE(record_reshape_split)
     EXPECT(get_all_axes(desc) == all_axes{d_axes{{0}}, d_axes{{1, 0}}, d_axes{{1, 1}}});
 }
 
+TEST_CASE(record_reshape_merge_split)
+{
+    auto desc = make_descriptor({3, 10, 16}, make_op("reshape", {{"dims", {3, 40, 2, 2}}}));
+    EXPECT(get_final_lens(desc) == final_lens{3, 40, 2, 2});
+    EXPECT(get_all_lens(desc) == all_lens{{3}, {10, 4}, {2}, {2}});
+    EXPECT(get_all_axes(desc) ==
+           all_axes{d_axes{{0}}, d_axes{{1}, {2, 0}}, d_axes{{2, 1}}, d_axes{{2, 2}}});
+}
+
 TEST_CASE(record_squeeze_trailing_1s)
 {
     auto desc = make_descriptor({3, 4, 4, 1, 1}, make_op("reshape", {{"dims", {3, 4, 4}}}));