[Mlir-commits] [mlir] f82cf74 - [mlir][tensor] Fix `getReassociationForCollapse` for tensor/scalar re… (#144118)

[llvmlistbot at llvm.org]

Author: Artem Gindinson
Date: 2025-06-13T20:03:24+02:00
New Revision: f82cf7442029d3376813db82eca60800e999bfb9

URL: https://github.com/llvm/llvm-project/commit/f82cf7442029d3376813db82eca60800e999bfb9
DIFF: https://github.com/llvm/llvm-project/commit/f82cf7442029d3376813db82eca60800e999bfb9.diff

LOG: [mlir][tensor] Fix `getReassociationForCollapse` for tensor/scalar re… (#144118)

…shapes

Commit 6e5a142 changed the behavior of the function when computing
reassociations between tensors (consisting of unit/dynamic dimensions)
and scalars/0d vectors. The IR representation for such reshapes actually
expects an empty reassociation, like so:
```
func.func @example(%arg0 : tensor<?x?x?xf32>) -> tensor<f32> {
  %0 = tensor.collapse_shape %arg0 [] : tensor<?x?x?xf32> into tensor<f32>
}
```

Restore the original behavior - the routine should resort to reporting
failures when compile time-known non-unit dimensions are part of the
attempted reassociation.

Signed-off-by: Artem Gindinson <gindinson at roofline.ai>

Added: 
    

Modified: 
    mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
    mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
index 3b1fdb69e8ef1..aa566c0086a2f 100644
--- a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
+++ b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
@@ -299,19 +299,17 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
   // this utility).
   if (numSourceDims <= numTargetDims)
     return std::nullopt;
-  // Early handling for scalar target types.
+  // Early handling for scalar target types. We should report an invalid
+  // reassociation for non-unit static dimensions - no chance to collapse these
+  // into a scalar.
   if (numTargetDims == 0) {
-    ReassociationIndices allSourceIndices;
-    allSourceIndices.reserve(numSourceDims);
     for (unsigned sourceDimIdx = 0; sourceDimIdx < numSourceDims;
          ++sourceDimIdx) {
       int64_t sourceSize = sourceShape[sourceDimIdx];
-      // All source dimensions must be unit or dynamic.
       if (sourceSize != 1 && sourceSize != ShapedType::kDynamic)
         return std::nullopt;
-      allSourceIndices.push_back(sourceDimIdx);
     }
-    return SmallVector<ReassociationIndices>{allSourceIndices};
+    return SmallVector<ReassociationIndices>{};
   }
 
   // Collect source ranges by iterating over the target shape left-to-right.

diff  --git a/mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp b/mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp
index db1a87a4de2d5..05f97e875e2dc 100644
--- a/mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp
+++ b/mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp
@@ -23,16 +23,16 @@ makeOptionalIndices(std::initializer_list<ReassociationIndices> list) {
 
 TEST(ReassociationIndicesForCollapse, ScalarTest) {
   EXPECT_EQ(getReassociationIndicesForCollapse({1}, {}),
-            makeOptionalIndices({{0}}));
+            makeOptionalIndices({}));
   EXPECT_EQ(getReassociationIndicesForCollapse({1, 1}, {}),
-            makeOptionalIndices({{0, 1}}));
+            makeOptionalIndices({}));
   EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic}, {}),
-            makeOptionalIndices({{0}}));
+            makeOptionalIndices({}));
   EXPECT_EQ(getReassociationIndicesForCollapse({1, ShapedType::kDynamic,
                                                 ShapedType::kDynamic, 1,
                                                 ShapedType::kDynamic},
                                                {}),
-            makeOptionalIndices({{0, 1, 2, 3, 4}}));
+            makeOptionalIndices({}));
 }
 
 TEST(ReassociationIndicesForCollapse, ScalarTestFailure) {