[Mlir-commits] [mlir] [mlir][tensor] Loosen restrictions on folding dynamic reshapes (PR #137963)

Artem Gindinson llvmlistbot at llvm.org
Fri May 9 08:52:35 PDT 2025


https://github.com/AGindinson updated https://github.com/llvm/llvm-project/pull/137963

>From 2ae44808cd573c31331d143d687afcdab6986a72 Mon Sep 17 00:00:00 2001
From: Artem Gindinson <gindinson at roofline.ai>
Date: Mon, 28 Apr 2025 17:24:11 +0000
Subject: [PATCH 1/5] [mlir][tensor] Loosen restrictions on folding dynamic
 reshapes

The main idea behind the change is to allow expand-of-collapse folds
for reshapes like `?x?xk` -> `?` (k>1). The rationale here is that the
expand op must have a coherent index/affine expression specified in its
`output_shape` argument (see example below), and if it doesn't, the IR
has already been invalidated at an earlier stage:
```
%c32 = arith.constant 32 : index
%div = arith.divsi %<some_index>, %c32 : index
%collapsed = tensor.collapse_shape %41#1 [[0], [1, 2], [3, 4]]
	         : tensor<9x?x32x?x32xf32> into tensor<9x?x?xf32>
%affine = affine.apply affine_map<()[s0] -> (s0 * 32)> ()[%div]
%expanded = tensor.expand_shape %collapsed [[0], [1, 2], [3]] output_shape [9, %div, 32, %affine]
		: tensor<9x?x?xf32> into tensor<9x?x32x?xf32>
```

On the above assumption, adjust the routine in
`getReassociationIndicesForCollapse()` to allow dynamic reshapes
beyond just `?x..?x1x1x..x1` -> `?`.

Moreover, the reassociation util was refactored to clearly distinguish
between dynamic and static subshapes. A few known caveats were noted as
a comment; it doesn't seem possible to fold all qualifying dynamic shape
patterns in a deterministic way without looking into affine expressions
simultaneously. That would be difficult to maintain in a single general
utility. Other implementation ideas/larger refactoring could include:
- abandoning the util usage in the `ComposeExpandOfCollapseOp` pattern,
  employing similar logic to `ComposeCollapseOfExpandOp`;
- providing dialect-specific implementations for Linalg/Tensor.

Signed-off-by: Artem Gindinson <gindinson at roofline.ai>
---
 mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp    | 103 ++++++++++--------
 .../Dialect/Linalg/simplify-pack-unpack.mlir  |   4 +-
 mlir/test/Dialect/Tensor/canonicalize.mlir    |  24 +++-
 3 files changed, 79 insertions(+), 52 deletions(-)

diff --git a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
index ed40a080441bc..694783849198a 100644
--- a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
+++ b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
@@ -31,59 +31,70 @@ mlir::getReassociationIndicesForReshape(ShapedType sourceType,
 std::optional<SmallVector<ReassociationIndices>>
 mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
                                          ArrayRef<int64_t> targetShape) {
-  if (sourceShape.size() <= targetShape.size())
+  unsigned numSourceDims = sourceShape.size(),
+           numTargetDims = targetShape.size();
+  if (numSourceDims <= numTargetDims)
     return std::nullopt;
-  unsigned sourceDim = 0;
-  SmallVector<ReassociationIndices> reassociationMap;
-  reassociationMap.reserve(targetShape.size());
-
-  ReassociationIndices currIndices;
-  int64_t prodOfCollapsedDims = 1;
-  while (sourceDim < sourceShape.size()) {
-    unsigned targetDim = reassociationMap.size();
-    // If we have mapped all the target dimensions stop and handle the remaining
-    // tail of size-1 dimensions explicitly.
-    if (targetDim == targetShape.size())
-      break;
+  SmallVector<ReassociationIndices, 4> reassociationMap;
+  reassociationMap.reserve(numTargetDims);
 
+  unsigned sourceDim = 0, targetDim = 0;
+  for (; targetDim < numTargetDims; ++targetDim) {
     int64_t currTargetShape = targetShape[targetDim];
-    while (sourceDim < (sourceShape.size() - 1) &&
-           sourceShape[sourceDim] != ShapedType::kDynamic &&
-           prodOfCollapsedDims * sourceShape[sourceDim] < currTargetShape) {
+    ReassociationIndices currIndices;
+    // 1. Target dimension is dynamic. Source shape should contain at least
+    // one dynamic dimension.
+    if (currTargetShape == ShapedType::kDynamic) {
+      // FIXME: We stop the search with the first dynamic dimension, while in
+      // fact, we can have a valid pattern like 2x?x?x4x8 -> ?x4x8. It becomes
+      // indeterministic altogether when we have neighboring dynamic dimensions
+      // in the target shape. Most of these patterns will be safely rejected,
+      // however we might achieve more correct folds by taking affine
+      // expressions into account, if these can be passed on by the call sites.
+      bool foundDynamic = false;
+      while (sourceDim < numSourceDims) {
+        currIndices.push_back(sourceDim);
+        if (sourceShape[sourceDim++] == ShapedType::kDynamic) {
+          foundDynamic = true;
+          break;
+        }
+      }
+      if (!foundDynamic)
+        return std::nullopt;
+
+      reassociationMap.push_back(currIndices);
+      continue;
+    }
+    // 2. Target dimension is static. The product of dimensions of the expanded
+    // shape should match the collapsed dimension shape.
+    int64_t prodOfCollapsedDims = 1;
+    bool reachedTargetDimSize = false;
+    while (sourceDim < numSourceDims) {
+      // Source shape cannot be dynamic if the target dim is static.
+      if (sourceShape[sourceDim] == ShapedType::kDynamic)
+        return std::nullopt;
       prodOfCollapsedDims *= sourceShape[sourceDim];
-      currIndices.push_back(sourceDim++);
+      if (prodOfCollapsedDims > currTargetShape)
+        break;
+      else if (prodOfCollapsedDims == currTargetShape) {
+        currIndices.push_back(sourceDim++);
+        reachedTargetDimSize = true;
+        break;
+      } else // prodOfCollapsedDims < currTargetShape
+        currIndices.push_back(sourceDim++);
     }
-
-    // If the current expanded dimension is dynamic, then the collapsed
-    // dimensions should also be dynamic and product of all previous unprocessed
-    // dimensions of the expanded shape should be 1.
-    if (sourceShape[sourceDim] == ShapedType::kDynamic &&
-        (currTargetShape != ShapedType::kDynamic || prodOfCollapsedDims != 1))
+    if (!reachedTargetDimSize)
       return std::nullopt;
-
-    // If the collapsed dim is dynamic, the current expanded dim should also
-    // be dynamic.
-    if (currTargetShape == ShapedType::kDynamic &&
-        sourceShape[sourceDim] != ShapedType::kDynamic)
-      return std::nullopt;
-
-    // For static shapes, if the product of dimensions of the expanded shape
-    // should match the collapsed dimension shape.
-    if (prodOfCollapsedDims * sourceShape[sourceDim] != currTargetShape)
-      return std::nullopt;
-
-    currIndices.push_back(sourceDim++);
-    reassociationMap.emplace_back(ReassociationIndices{});
-    std::swap(reassociationMap.back(), currIndices);
-    prodOfCollapsedDims = 1;
+    reassociationMap.push_back(currIndices);
   }
-  // All the dimensions in the target must have been processed.
-  if (reassociationMap.size() != targetShape.size())
-    return std::nullopt;
-  // Process any remaining entries in the source shape. They all need to be
-  // 1 or dynamic.
-  for (; sourceDim < sourceShape.size(); sourceDim++) {
-    if (sourceShape[sourceDim] != ShapedType::kDynamic &&
+  // Now that we've mapped all the target dimensions, process any remaining
+  // entries in the source shape explicitly. Either the last target dimension
+  // is dynamic, or all remaining source entries need to be 1 or dynamic. Same
+  // applies when target shape is empty (can be the case for subshape
+  // reassociations).
+  for (; sourceDim < numSourceDims; sourceDim++) {
+    if ((targetShape.empty() || targetShape.back() != ShapedType::kDynamic) &&
+        sourceShape[sourceDim] != ShapedType::kDynamic &&
         sourceShape[sourceDim] != 1)
       return std::nullopt;
     // The map is empty when the target type is a scalar.
diff --git a/mlir/test/Dialect/Linalg/simplify-pack-unpack.mlir b/mlir/test/Dialect/Linalg/simplify-pack-unpack.mlir
index 51350e5bc8498..6979770154bab 100644
--- a/mlir/test/Dialect/Linalg/simplify-pack-unpack.mlir
+++ b/mlir/test/Dialect/Linalg/simplify-pack-unpack.mlir
@@ -158,8 +158,8 @@ func.func @unpack_to_partial_slice(%arg0: tensor<8x32xf32>) -> tensor<255xf32> {
 // -----
 
 // CHECK-LABEL: func.func @unpack_dynamic
-// CHECK-NOT:     tensor.collapse
-// CHECK:         linalg.unpack
+// CHECK:     tensor.collapse
+// CHECK-NOT:         linalg.unpack
 func.func @unpack_dynamic(%arg0: tensor<?x32xf32>) -> tensor<?xf32> {
   %c32 = arith.constant 32 : index
   %c0 = arith.constant 0 : index
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 85bf6fba52aa4..443f931745557 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -1068,7 +1068,7 @@ func.func @fold_expand_of_collapse(%arg0 : tensor<3x4x4xf32>) -> tensor<3x4x4xf3
 
 // -----
 
-func.func @fold_expand_of_collapse_dynamic(%arg0 : tensor<?x4x?xf32>, %arg1: index, %arg2: index)
+func.func @fold_expand_of_collapse_mixed_subshape(%arg0 : tensor<?x4x?xf32>, %arg1: index, %arg2: index)
     -> tensor<?x4x?xf32> {
   %0 = tensor.collapse_shape %arg0 [[0, 1], [2]]
       : tensor<?x4x?xf32> into tensor<?x?xf32>
@@ -1076,12 +1076,28 @@ func.func @fold_expand_of_collapse_dynamic(%arg0 : tensor<?x4x?xf32>, %arg1: ind
       : tensor<?x?xf32> into tensor<?x4x?xf32>
   return %1 : tensor<?x4x?xf32>
 }
-// CHECK-LABEL: @fold_expand_of_collapse_dynamic
+// CHECK-LABEL: @fold_expand_of_collapse_mixed_subshape
 //   CHECK-NOT:   tensor.{{.*}}_shape
 
 // -----
 
-func.func @no_fold_expand_of_collapse_dynamic(%arg0 : tensor<?x?x?xf32>, %arg1: index, %arg2: index, %arg3: index)
+func.func @fold_expand_of_collapse_mixed_target_subshape(%arg0 : tensor<?x4x?x2xf32>, %arg1: index, %arg2: index)
+    -> tensor<?x4x?xf32> {
+  %0 = tensor.collapse_shape %arg0 [[0, 1], [2, 3]]
+      : tensor<?x4x?x2xf32> into tensor<?x?xf32>
+  %1 = tensor.expand_shape %0 [[0, 1], [2]] output_shape [%arg1, 4, %arg2]
+      : tensor<?x?xf32> into tensor<?x4x?xf32>
+  return %1 : tensor<?x4x?xf32>
+}
+// CHECK-LABEL: @fold_expand_of_collapse_mixed_target_subshape
+//   CHECK-NOT:   tensor.expand_shape
+//       CHECK:   %[[COLLAPSE:.+]] = tensor.collapse_shape %arg0 {{\[}}[0], [1], [2, 3]]
+//  CHECK-SAME:     : tensor<?x4x?x2xf32> into tensor<?x4x?xf32>
+//  CHECK-NEXT:   return %[[COLLAPSE]]
+
+// -----
+
+func.func @no_fold_expand_of_collapse_fully_dynamic(%arg0 : tensor<?x?x?xf32>, %arg1: index, %arg2: index, %arg3: index)
     -> tensor<?x?x?xf32> {
   %0 = tensor.collapse_shape %arg0 [[0, 1], [2]]
       : tensor<?x?x?xf32> into tensor<?x?xf32>
@@ -1089,7 +1105,7 @@ func.func @no_fold_expand_of_collapse_dynamic(%arg0 : tensor<?x?x?xf32>, %arg1:
       : tensor<?x?xf32> into tensor<?x?x?xf32>
   return %1 : tensor<?x?x?xf32>
 }
-// CHECK-LABEL: @no_fold_expand_of_collapse_dynamic
+// CHECK-LABEL: @no_fold_expand_of_collapse_fully_dynamic
 //       CHECK:   tensor.collapse_shape
 //       CHECK:   %[[EXPAND:.+]] = tensor.expand_shape
 //       CHECK:   return %[[EXPAND]]

>From 52ff4e0a1e81e13282975076d730e741a1da1cae Mon Sep 17 00:00:00 2001
From: Artem Gindinson <gindinson at roofline.ai>
Date: Fri, 9 May 2025 15:12:21 +0000
Subject: [PATCH 2/5] [fixup] Algorithm rewrite

Signed-off-by: Artem Gindinson <gindinson at roofline.ai>
---
 mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp | 140 ++++++++++++++-------
 1 file changed, 93 insertions(+), 47 deletions(-)

diff --git a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
index 694783849198a..1cd06a2757363 100644
--- a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
+++ b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
@@ -39,67 +39,113 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
   reassociationMap.reserve(numTargetDims);
 
   unsigned sourceDim = 0, targetDim = 0;
-  for (; targetDim < numTargetDims; ++targetDim) {
-    int64_t currTargetShape = targetShape[targetDim];
-    ReassociationIndices currIndices;
-    // 1. Target dimension is dynamic. Source shape should contain at least
-    // one dynamic dimension.
-    if (currTargetShape == ShapedType::kDynamic) {
-      // FIXME: We stop the search with the first dynamic dimension, while in
-      // fact, we can have a valid pattern like 2x?x?x4x8 -> ?x4x8. It becomes
-      // indeterministic altogether when we have neighboring dynamic dimensions
-      // in the target shape. Most of these patterns will be safely rejected,
-      // however we might achieve more correct folds by taking affine
-      // expressions into account, if these can be passed on by the call sites.
-      bool foundDynamic = false;
-      while (sourceDim < numSourceDims) {
-        currIndices.push_back(sourceDim);
-        if (sourceShape[sourceDim++] == ShapedType::kDynamic) {
-          foundDynamic = true;
-          break;
-        }
-      }
-      if (!foundDynamic)
-        return std::nullopt;
-
-      reassociationMap.push_back(currIndices);
-      continue;
-    }
-    // 2. Target dimension is static. The product of dimensions of the expanded
-    // shape should match the collapsed dimension shape.
+  // Source dimensions iteration logic for static target dimensions.
+  // FIXME: Instead of lambda-capturing this function's source shape index "in
+  // place", consider refactoring this into a separate function.
+  auto collectSourceIndicesForStaticTargetDim =
+      [&](int64_t targetShape,
+          bool mayHaveOffset = false) -> FailureOr<ReassociationIndices> {
+    ReassociationIndices resultIndices;
     int64_t prodOfCollapsedDims = 1;
     bool reachedTargetDimSize = false;
-    while (sourceDim < numSourceDims) {
+    for (; sourceDim < numSourceDims; ++sourceDim) {
       // Source shape cannot be dynamic if the target dim is static.
       if (sourceShape[sourceDim] == ShapedType::kDynamic)
-        return std::nullopt;
+        return failure();
       prodOfCollapsedDims *= sourceShape[sourceDim];
-      if (prodOfCollapsedDims > currTargetShape)
-        break;
-      else if (prodOfCollapsedDims == currTargetShape) {
-        currIndices.push_back(sourceDim++);
+      resultIndices.push_back(sourceDim);
+      if (prodOfCollapsedDims > targetShape && !mayHaveOffset)
+        return failure();
+      while (prodOfCollapsedDims > targetShape) {
+        assert(!resultIndices.empty());
+        auto frontOffsetIdx = resultIndices.begin();
+        prodOfCollapsedDims /= sourceShape[*frontOffsetIdx];
+        resultIndices.erase(frontOffsetIdx);
+      }
+      if (prodOfCollapsedDims == targetShape) {
         reachedTargetDimSize = true;
+        ++sourceDim;
         break;
-      } else // prodOfCollapsedDims < currTargetShape
-        currIndices.push_back(sourceDim++);
+      }
     }
     if (!reachedTargetDimSize)
+      return failure();
+    return resultIndices;
+  };
+  // Source dimensions iteration logic for dynamic target dimensions.
+  // FIXME: Instead of lambda-capturing this function's source shape index "in
+  // place", consider refactoring this into a separate function.
+  auto collectSourceIndicesForDynamicTargetDim =
+      [&](bool allowStaticNonOnes,
+          bool mapConsecutiveDynDims) -> FailureOr<ReassociationIndices> {
+    ReassociationIndices resultIndices;
+    bool foundFirstDynamic = false;
+    while (sourceDim < numSourceDims) {
+      if (sourceShape[sourceDim] == ShapedType::kDynamic) {
+        if (foundFirstDynamic && !mapConsecutiveDynDims)
+          break;
+        foundFirstDynamic |= true;
+      } else {
+        if (foundFirstDynamic)
+          break;
+        else if (sourceShape[sourceDim] > 1 && !allowStaticNonOnes)
+          return failure();
+      }
+      resultIndices.push_back(sourceDim++);
+    }
+    if (!foundFirstDynamic)
+      return failure();
+    return resultIndices;
+  };
+  // Iterate over target shape.
+  bool wasLastDimDynamic = false;
+  for (; targetDim < numTargetDims; ++targetDim) {
+    int64_t currTargetShape = targetShape[targetDim];
+    if (currTargetShape != ShapedType::kDynamic) {
+      unsigned sourceDimAtStart = sourceDim;
+      auto indices = collectSourceIndicesForStaticTargetDim(
+          currTargetShape, /*mayHaveOffset=*/wasLastDimDynamic);
+      if (failed(indices))
+        return std::nullopt;
+      if (wasLastDimDynamic) {
+        assert(!reassociationMap.empty());
+        auto &previousIndices = reassociationMap.back();
+        for (; sourceDimAtStart < indices->front(); ++sourceDimAtStart)
+          previousIndices.push_back(sourceDimAtStart);
+      }
+      reassociationMap.push_back(*indices);
+      wasLastDimDynamic = false;
+      continue;
+    }
+
+    bool isNextDimDynamic = targetDim + 1 < numTargetDims &&
+                            targetShape[targetDim + 1] == ShapedType::kDynamic;
+    auto indices = collectSourceIndicesForDynamicTargetDim(
+        /*allowStaticNonOnes=*/!wasLastDimDynamic,
+        /*mapConsecutiveDynDims=*/!wasLastDimDynamic && !isNextDimDynamic);
+    if (failed(indices))
       return std::nullopt;
-    reassociationMap.push_back(currIndices);
+    reassociationMap.push_back(*indices);
+    wasLastDimDynamic = true;
   }
   // Now that we've mapped all the target dimensions, process any remaining
-  // entries in the source shape explicitly. Either the last target dimension
-  // is dynamic, or all remaining source entries need to be 1 or dynamic. Same
-  // applies when target shape is empty (can be the case for subshape
-  // reassociations).
+  // entries in the source shape explicitly.
   for (; sourceDim < numSourceDims; sourceDim++) {
-    if ((targetShape.empty() || targetShape.back() != ShapedType::kDynamic) &&
-        sourceShape[sourceDim] != ShapedType::kDynamic &&
-        sourceShape[sourceDim] != 1)
+    const bool isOne = sourceShape[sourceDim] == 1,
+               isDynamic = sourceShape[sourceDim] == ShapedType::kDynamic;
+    if (targetShape.empty()) {
+      if (!isOne && !isDynamic)
+        return std::nullopt;
+      continue;
+    }
+    if (wasLastDimDynamic && isDynamic)
+      return std::nullopt;
+    // If the last target dimension is static, only source dimensions of 1 are
+    // acceptable.
+    if (!wasLastDimDynamic && !isOne)
       return std::nullopt;
-    // The map is empty when the target type is a scalar.
-    if (!reassociationMap.empty())
-      reassociationMap.back().push_back(sourceDim);
+    assert(!reassociationMap.empty());
+    reassociationMap.back().push_back(sourceDim);
   }
   return reassociationMap;
 }

>From 1c85a6875fecf52e6714b81bbe9bd2da81178c9e Mon Sep 17 00:00:00 2001
From: Artem Gindinson <gindinson at roofline.ai>
Date: Fri, 9 May 2025 15:12:34 +0000
Subject: [PATCH 3/5] [fixup] Add/expand unit tests

Signed-off-by: Artem Gindinson <gindinson at roofline.ai>
Co-authored-by: Ian Wood <ianwood2024 at u.northwestern.edu>
---
 mlir/test/Dialect/Tensor/canonicalize.mlir    |  15 +++
 mlir/unittests/Dialect/Utils/CMakeLists.txt   |   1 +
 .../Dialect/Utils/ReshapeOpsUtilsTest.cpp     | 125 ++++++++++++++++++
 3 files changed, 141 insertions(+)
 create mode 100644 mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp

diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 443f931745557..035ea850c9102 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -1112,6 +1112,21 @@ func.func @no_fold_expand_of_collapse_fully_dynamic(%arg0 : tensor<?x?x?xf32>, %
 
 // -----
 
+func.func @no_fold_expand_of_collapse_adjacent_dynamic(%arg0 : tensor<?x?x?xf32>, %arg1: index, %arg2: index)
+    -> tensor<?x?xf32> {
+  %0 = tensor.collapse_shape %arg0 [[0, 1, 2]]
+      : tensor<?x?x?xf32> into tensor<?xf32>
+  %1 = tensor.expand_shape %0 [[0, 1]] output_shape [%arg1, %arg2]
+      : tensor<?xf32> into tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+// CHECK-LABEL: @no_fold_expand_of_collapse_adjacent_dynamic
+//       CHECK:   tensor.collapse_shape
+//       CHECK:   %[[EXPAND:.+]] = tensor.expand_shape
+//       CHECK:   return %[[EXPAND]]
+
+// -----
+
 func.func @compose_expand_of_collapse_last_two_dims(%arg0: tensor<?x64x1xf32>) -> tensor<?x384xf32> {
   %collapsed = tensor.collapse_shape %arg0 [[0, 1, 2]] : tensor<?x64x1xf32> into tensor<?xf32>
   %c0 = arith.constant 0 : index
diff --git a/mlir/unittests/Dialect/Utils/CMakeLists.txt b/mlir/unittests/Dialect/Utils/CMakeLists.txt
index 61b9cdcb3b8f3..e921c8bcfb4e5 100644
--- a/mlir/unittests/Dialect/Utils/CMakeLists.txt
+++ b/mlir/unittests/Dialect/Utils/CMakeLists.txt
@@ -1,5 +1,6 @@
 add_mlir_unittest(MLIRDialectUtilsTests
   StructuredOpsUtilsTest.cpp
+  ReshapeOpsUtilsTest.cpp
   IndexingUtilsTest.cpp
 )
 mlir_target_link_libraries(MLIRDialectUtilsTests
diff --git a/mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp b/mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp
new file mode 100644
index 0000000000000..bfcc70150e2ed
--- /dev/null
+++ b/mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp
@@ -0,0 +1,125 @@
+//===- ReshapeOpsUtilsTest.cpp - ReshapeOpsUtils unit tests ---------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
+#include "llvm/ADT/STLExtras.h"
+#include "gtest/gtest.h"
+#include <optional>
+
+using namespace mlir;
+
+/// Helper to make constructing
+/// `std::optional<SmallVector<ReassociationIndices>>` more readable.
+static std::optional<SmallVector<ReassociationIndices>>
+makeOptionalIndices(std::initializer_list<ReassociationIndices> list) {
+  return std::optional<SmallVector<ReassociationIndices>>(list);
+}
+
+TEST(ReassociationIndicesForCollapse, StaticTest) {
+  EXPECT_EQ(getReassociationIndicesForCollapse({10, 20}, {200}),
+            makeOptionalIndices({{0, 1}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse({10, 20, 30}, {10, 600}),
+            makeOptionalIndices({{0}, {1, 2}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse({10, 20, 30}, {200, 30}),
+            makeOptionalIndices({{0, 1}, {2}}));
+}
+
+TEST(ReassociationIndicesForCollapse, StaticTestFailure) {
+  EXPECT_EQ(getReassociationIndicesForCollapse({10, 20}, {10}), std::nullopt);
+  EXPECT_EQ(getReassociationIndicesForCollapse({10, 20}, {10, 20}),
+            std::nullopt);
+  EXPECT_EQ(getReassociationIndicesForCollapse({10, 20, 30}, {200, 300}),
+            std::nullopt);
+  EXPECT_EQ(getReassociationIndicesForCollapse({10, 20, 30}, {1, 10, 20, 30}),
+            std::nullopt);
+}
+
+TEST(ReassociationIndicesForCollapse, StaticTestUnitDims) {
+  EXPECT_EQ(getReassociationIndicesForCollapse({10, 1}, {10}),
+            makeOptionalIndices({{0, 1}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse({1, 20, 30}, {600}),
+            makeOptionalIndices({{0, 1, 2}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse({1, 1, 1}, {1}),
+            makeOptionalIndices({{0, 1, 2}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse({1, 1, 1}, {1, 1}),
+            makeOptionalIndices({{0}, {1, 2}}));
+}
+
+TEST(ReassociationIndicesForCollapse, DynamicTest) {
+  EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic, 1},
+                                               {ShapedType::kDynamic}),
+            makeOptionalIndices({{0, 1}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic, 1, 1},
+                                               {ShapedType::kDynamic}),
+            makeOptionalIndices({{0, 1, 2}}));
+  EXPECT_EQ(
+      getReassociationIndicesForCollapse(
+          {ShapedType::kDynamic, ShapedType::kDynamic}, {ShapedType::kDynamic}),
+      makeOptionalIndices({{0, 1}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse(
+                {1, ShapedType::kDynamic, ShapedType::kDynamic},
+                {1, ShapedType::kDynamic}),
+            makeOptionalIndices({{0}, {1, 2}}));
+
+  EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic, 10},
+                                               {ShapedType::kDynamic}),
+            makeOptionalIndices({{0, 1}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse(
+                {1, ShapedType::kDynamic, ShapedType::kDynamic},
+                {ShapedType::kDynamic}),
+            makeOptionalIndices({{0, 1, 2}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse({10, ShapedType::kDynamic},
+                                               {ShapedType::kDynamic}),
+            makeOptionalIndices({{0, 1}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic, 10, 20},
+                                               {ShapedType::kDynamic, 20}),
+            makeOptionalIndices({{0, 1}, {2}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse({10, ShapedType::kDynamic, 20},
+                                               {ShapedType::kDynamic, 20}),
+            makeOptionalIndices({{0, 1}, {2}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse(
+                {ShapedType::kDynamic, 3, 2, 5, 2}, {ShapedType::kDynamic, 20}),
+            makeOptionalIndices({{0, 1}, {2, 3, 4}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse(
+                {10, ShapedType::kDynamic, 20, ShapedType::kDynamic, 1},
+                {ShapedType::kDynamic, 20, ShapedType::kDynamic}),
+            makeOptionalIndices({{0, 1}, {2}, {3, 4}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse({1, ShapedType::kDynamic, 1},
+                                               {ShapedType::kDynamic}),
+            makeOptionalIndices({{0, 1, 2}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse(
+                {1, ShapedType::kDynamic, ShapedType::kDynamic},
+                {ShapedType::kDynamic, ShapedType::kDynamic}),
+            makeOptionalIndices({{0, 1}, {2}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse(
+                {ShapedType::kDynamic, 1, ShapedType::kDynamic},
+                {ShapedType::kDynamic, ShapedType::kDynamic}),
+            makeOptionalIndices({{0}, {1, 2}}));
+}
+
+TEST(ReassociationIndicesForCollapse, DynamicTestFailure) {
+  EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic, 10, 20},
+                                               {ShapedType::kDynamic, 10}),
+            std::nullopt);
+  EXPECT_EQ(getReassociationIndicesForCollapse(
+                {ShapedType::kDynamic, 10, ShapedType::kDynamic},
+                {ShapedType::kDynamic, ShapedType::kDynamic}),
+            std::nullopt);
+  EXPECT_EQ(getReassociationIndicesForCollapse(
+                {20, ShapedType::kDynamic, 10, ShapedType::kDynamic},
+                {ShapedType::kDynamic, ShapedType::kDynamic}),
+            std::nullopt);
+  EXPECT_EQ(getReassociationIndicesForCollapse(
+                {ShapedType::kDynamic, 5, 3, 2, 2}, {ShapedType::kDynamic, 20}),
+            std::nullopt);
+  EXPECT_EQ(
+      getReassociationIndicesForCollapse(
+          {ShapedType::kDynamic, ShapedType::kDynamic, ShapedType::kDynamic},
+          {ShapedType::kDynamic, ShapedType::kDynamic}),
+      std::nullopt);
+}

>From 0fe986e217e52cc519823a999e83155f1cd2d3ef Mon Sep 17 00:00:00 2001
From: Artem Gindinson <gindinson at roofline.ai>
Date: Fri, 9 May 2025 15:15:08 +0000
Subject: [PATCH 4/5] [fixup] variable renaming

Signed-off-by: Artem Gindinson <gindinson at roofline.ai>
---
 mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp | 39 +++++++++++-----------
 1 file changed, 20 insertions(+), 19 deletions(-)

diff --git a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
index 1cd06a2757363..a6ee21d941e17 100644
--- a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
+++ b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
@@ -38,7 +38,7 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
   SmallVector<ReassociationIndices, 4> reassociationMap;
   reassociationMap.reserve(numTargetDims);
 
-  unsigned sourceDim = 0, targetDim = 0;
+  unsigned sourceDimIdx = 0, targetDimIdx = 0;
   // Source dimensions iteration logic for static target dimensions.
   // FIXME: Instead of lambda-capturing this function's source shape index "in
   // place", consider refactoring this into a separate function.
@@ -48,12 +48,12 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
     ReassociationIndices resultIndices;
     int64_t prodOfCollapsedDims = 1;
     bool reachedTargetDimSize = false;
-    for (; sourceDim < numSourceDims; ++sourceDim) {
+    for (; sourceDimIdx < numSourceDims; ++sourceDimIdx) {
       // Source shape cannot be dynamic if the target dim is static.
-      if (sourceShape[sourceDim] == ShapedType::kDynamic)
+      if (sourceShape[sourceDimIdx] == ShapedType::kDynamic)
         return failure();
-      prodOfCollapsedDims *= sourceShape[sourceDim];
-      resultIndices.push_back(sourceDim);
+      prodOfCollapsedDims *= sourceShape[sourceDimIdx];
+      resultIndices.push_back(sourceDimIdx);
       if (prodOfCollapsedDims > targetShape && !mayHaveOffset)
         return failure();
       while (prodOfCollapsedDims > targetShape) {
@@ -64,7 +64,7 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
       }
       if (prodOfCollapsedDims == targetShape) {
         reachedTargetDimSize = true;
-        ++sourceDim;
+        ++sourceDimIdx;
         break;
       }
     }
@@ -80,18 +80,18 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
           bool mapConsecutiveDynDims) -> FailureOr<ReassociationIndices> {
     ReassociationIndices resultIndices;
     bool foundFirstDynamic = false;
-    while (sourceDim < numSourceDims) {
-      if (sourceShape[sourceDim] == ShapedType::kDynamic) {
+    while (sourceDimIdx < numSourceDims) {
+      if (sourceShape[sourceDimIdx] == ShapedType::kDynamic) {
         if (foundFirstDynamic && !mapConsecutiveDynDims)
           break;
         foundFirstDynamic |= true;
       } else {
         if (foundFirstDynamic)
           break;
-        else if (sourceShape[sourceDim] > 1 && !allowStaticNonOnes)
+        else if (sourceShape[sourceDimIdx] > 1 && !allowStaticNonOnes)
           return failure();
       }
-      resultIndices.push_back(sourceDim++);
+      resultIndices.push_back(sourceDimIdx++);
     }
     if (!foundFirstDynamic)
       return failure();
@@ -99,10 +99,10 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
   };
   // Iterate over target shape.
   bool wasLastDimDynamic = false;
-  for (; targetDim < numTargetDims; ++targetDim) {
-    int64_t currTargetShape = targetShape[targetDim];
+  for (; targetDimIdx < numTargetDims; ++targetDimIdx) {
+    int64_t currTargetShape = targetShape[targetDimIdx];
     if (currTargetShape != ShapedType::kDynamic) {
-      unsigned sourceDimAtStart = sourceDim;
+      unsigned sourceDimAtStart = sourceDimIdx;
       auto indices = collectSourceIndicesForStaticTargetDim(
           currTargetShape, /*mayHaveOffset=*/wasLastDimDynamic);
       if (failed(indices))
@@ -118,8 +118,9 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
       continue;
     }
 
-    bool isNextDimDynamic = targetDim + 1 < numTargetDims &&
-                            targetShape[targetDim + 1] == ShapedType::kDynamic;
+    bool isNextDimDynamic =
+        targetDimIdx + 1 < numTargetDims &&
+        targetShape[targetDimIdx + 1] == ShapedType::kDynamic;
     auto indices = collectSourceIndicesForDynamicTargetDim(
         /*allowStaticNonOnes=*/!wasLastDimDynamic,
         /*mapConsecutiveDynDims=*/!wasLastDimDynamic && !isNextDimDynamic);
@@ -130,9 +131,9 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
   }
   // Now that we've mapped all the target dimensions, process any remaining
   // entries in the source shape explicitly.
-  for (; sourceDim < numSourceDims; sourceDim++) {
-    const bool isOne = sourceShape[sourceDim] == 1,
-               isDynamic = sourceShape[sourceDim] == ShapedType::kDynamic;
+  for (; sourceDimIdx < numSourceDims; sourceDimIdx++) {
+    const bool isOne = sourceShape[sourceDimIdx] == 1,
+               isDynamic = sourceShape[sourceDimIdx] == ShapedType::kDynamic;
     if (targetShape.empty()) {
       if (!isOne && !isDynamic)
         return std::nullopt;
@@ -145,7 +146,7 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
     if (!wasLastDimDynamic && !isOne)
       return std::nullopt;
     assert(!reassociationMap.empty());
-    reassociationMap.back().push_back(sourceDim);
+    reassociationMap.back().push_back(sourceDimIdx);
   }
   return reassociationMap;
 }

>From e3aa2394225bb23cde429cc06bea28df7257070a Mon Sep 17 00:00:00 2001
From: Artem Gindinson <gindinson at roofline.ai>
Date: Fri, 9 May 2025 15:51:59 +0000
Subject: [PATCH 5/5] [fixup] Additional edge-case

Signed-off-by: Artem Gindinson <gindinson at roofline.ai>
---
 mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp           | 10 ++++++++--
 mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp |  9 +++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
index a6ee21d941e17..8c19f20f446da 100644
--- a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
+++ b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
@@ -139,8 +139,14 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
         return std::nullopt;
       continue;
     }
-    if (wasLastDimDynamic && isDynamic)
-      return std::nullopt;
+    // If the last 2 dimensions in the target were dynamic, the tail in the
+    // source shape cannot contain a dynamic value. E.g. ?x?->? is valid,
+    // however ?x?x10x?->?x? would be indeterminate.
+    if (wasLastDimDynamic && numTargetDims > 1 &&
+        targetShape[numTargetDims - 2] == ShapedType::kDynamic) {
+      if (isDynamic)
+        return std::nullopt;
+    }
     // If the last target dimension is static, only source dimensions of 1 are
     // acceptable.
     if (!wasLastDimDynamic && !isOne)
diff --git a/mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp b/mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp
index bfcc70150e2ed..2564866fac493 100644
--- a/mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp
+++ b/mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp
@@ -92,6 +92,10 @@ TEST(ReassociationIndicesForCollapse, DynamicTest) {
   EXPECT_EQ(getReassociationIndicesForCollapse({1, ShapedType::kDynamic, 1},
                                                {ShapedType::kDynamic}),
             makeOptionalIndices({{0, 1, 2}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse(
+                {ShapedType::kDynamic, ShapedType::kDynamic, 1},
+                {ShapedType::kDynamic, ShapedType::kDynamic}),
+            makeOptionalIndices({{0}, {1, 2}}));
   EXPECT_EQ(getReassociationIndicesForCollapse(
                 {1, ShapedType::kDynamic, ShapedType::kDynamic},
                 {ShapedType::kDynamic, ShapedType::kDynamic}),
@@ -122,4 +126,9 @@ TEST(ReassociationIndicesForCollapse, DynamicTestFailure) {
           {ShapedType::kDynamic, ShapedType::kDynamic, ShapedType::kDynamic},
           {ShapedType::kDynamic, ShapedType::kDynamic}),
       std::nullopt);
+  EXPECT_EQ(getReassociationIndicesForCollapse(
+                {ShapedType::kDynamic, ShapedType::kDynamic, 10, 1,
+                 ShapedType::kDynamic},
+                {ShapedType::kDynamic, ShapedType::kDynamic}),
+            std::nullopt);
 }



More information about the Mlir-commits mailing list