fix bug

Author: Max191

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

@@ -104,11 +104,6 @@ static OpFoldResult foldReshapeOp(ReshapeOpTy reshapeOp,
   if (srcType != resultType)
     return nullptr;
 
-  // If the reshapes are expanding and then collapsing, the ops can be folded
-  // despite multiple dynamic dimensions.
-  if (srcType.getRank() < reshapeSrcOp.getResultType().getRank())
-    return reshapeSrcOp.getSrc();
-  // Otherwise, only 1 dynamic dimension is allowed.
   if (srcType == resultType &&
       llvm::count_if(srcType.getShape(), ShapedType::isDynamic) < 2) {
     return reshapeSrcOp.getSrc();
@@ -124,6 +119,10 @@ static OpFoldResult foldReshapeOp(ReshapeOpTy reshapeOp,
   auto inverseReassociations = reshapeSrcOp.getReassociationIndices();
   if (reassociations != inverseReassociations)
     return nullptr;
+  // If the reshapes are expanding and then collapsing, the ops can be folded
+  // despite multiple dynamic dimensions.
+  if (srcType.getRank() < reshapeSrcOp.getResultType().getRank())
+    return reshapeSrcOp.getSrc();
   ArrayRef<int64_t> expandedSrcShape = srcType.getShape();
   ArrayRef<int64_t> expandedResultShape = resultType.getShape();
   if (llvm::none_of(reassociations, [&](auto reInd) {

mlir/test/Dialect/Tensor/canonicalize.mlir

@@ -1169,6 +1169,21 @@ func.func @fold_collapse_of_expand_fully_dynamic(%arg0 : tensor<?x?xf32>, %arg1:
 
 // -----
 
+func.func @no_fold_parallel_collapse_of_expand_dynamic(%arg0 : tensor<?x?x?xf32>, %arg1: index, %arg2: index, %arg3: index, %arg4: index)
+    -> tensor<?x?x?xf32> {
+  %0 = tensor.expand_shape %arg0 [[0, 1], [2], [3]] output_shape [%arg1, %arg2, %arg3, %arg4]
+      : tensor<?x?x?xf32> into tensor<?x?x?x?xf32>
+  %1 = tensor.collapse_shape %0 [[0], [1], [2, 3]]
+      : tensor<?x?x?x?xf32> into tensor<?x?x?xf32>
+  return %1 : tensor<?x?x?xf32>
+}
+// CHECK-LABEL: @no_fold_parallel_collapse_of_expand_dynamic
+//       CHECK:   tensor.expand_shape
+//       CHECK:   %[[COLLAPSE:.+]] = tensor.collapse_shape
+//       CHECK:   return %[[COLLAPSE]]
+
+// -----
+
 func.func @fold_expand_of_collapse(%arg0 : tensor<3x4x4xf32>) -> tensor<3x4x4xf32> {
   %0 = tensor.collapse_shape %arg0 [[0, 1], [2]]
       : tensor<3x4x4xf32> into tensor<12x4xf32>