Allow only valid vector.shape_cast transitive folding

When folding A->B->C => A->C only accept A->C that is valid shape cast

Reviewed By: ThomasRaoux, nicolasvasilache

Differential Revision: https://reviews.llvm.org/D111473

Author: asaadaldien

mlir/lib/Dialect/Vector/VectorOps.cpp

@@ -3634,6 +3634,20 @@ OpFoldResult ShapeCastOp::fold(ArrayRef<Attribute> operands) {
   if (auto otherOp = source().getDefiningOp<ShapeCastOp>()) {
     if (result().getType() == otherOp.source().getType())
       return otherOp.source();
+
+    // Only allows valid transitive folding.
+    VectorType srcType = otherOp.source().getType().cast<VectorType>();
+    VectorType resultType = getResult().getType().cast<VectorType>();
+    if (srcType.getRank() < resultType.getRank()) {
+      if (!isValidShapeCast(srcType.getShape(), resultType.getShape()))
+        return {};
+    } else if (srcType.getRank() > resultType.getRank()) {
+      if (!isValidShapeCast(resultType.getShape(), srcType.getShape()))
+        return {};
+    } else {
+      return {};
+    }
+
     setOperand(otherOp.source());
     return getResult();
   }

mlir/test/Dialect/Vector/canonicalize.mlir

@@ -542,6 +542,17 @@ func @fold_extract_shapecast_negative(%arg0 : vector<16xf32>,
   return %r : vector<4x2xf32>
 }
 
+// -----
+
+// CHECK-LABEL: dont_fold_expand_collapse
+//       CHECK:   %[[A:.*]] = vector.shape_cast %{{.*}} : vector<1x1x64xf32> to vector<1x1x8x8xf32>
+//       CHECK:   %[[B:.*]] = vector.shape_cast %{{.*}} : vector<1x1x8x8xf32> to vector<8x8xf32>
+//       CHECK:   return %[[B]] : vector<8x8xf32>
+func @dont_fold_expand_collapse(%arg0: vector<1x1x64xf32>) -> vector<8x8xf32> {
+    %0 = vector.shape_cast %arg0 : vector<1x1x64xf32> to vector<1x1x8x8xf32>
+    %1 = vector.shape_cast %0 : vector<1x1x8x8xf32> to vector<8x8xf32>
+    return %1 : vector<8x8xf32>
+}
 
 // -----