Allow only valid vector.shape_cast transitive folding

[asaadaldien]

Summary

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

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorOps.cpp

3645

What if srcType.getShape() == resultType.getRank()? This case would assert right now.

[asaadaldien]

mlir/lib/Dialect/Vector/VectorOps.cpp

3645

if its the same rank, in order to be valid it will have to be the same type. so we want be executing this isValidShapeCast

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorOps.cpp

3645

But isValidShapeCast has assert(rankA < rankB);
https://github.com/llvm/llvm-project/blob/e356027016c6365b3d8924f54c33e2c63d931492/mlir/lib/Dialect/Vector/VectorOps.cpp#L3465

so we would just assert in this case?

[asaadaldien]

Assert src rank != result rank

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorOps.cpp

3642

How do we know that? Can't we have something like:

%0 = vector.shape_cast %arg0 : vector<4x2x2xf32> to vector<8x2xf32>
%1 = vector.shape_cast %0 : vector<8x2xf32> to vector<2x4x2xf32>

[asaadaldien]

don't fold A -> B -> C when rank(A) == rank(C) && type(A) != type(C)

[asaadaldien]

mlir/lib/Dialect/Vector/VectorOps.cpp

3642

right, it will be always invalid.

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorOps.cpp

3646

nit: remove braces for simple single line if:
https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

3649–3651

Can you add a test for this case?

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorOps.cpp

3642

Can you put this if/else logic under isValidShapeCast(typea, typeb) ?

[asaadaldien]

Rebase...

mlir/lib/Dialect/Vector/VectorOps.cpp

OpFoldResult ShapeCastOp::fold(ArrayRef<Attribute> operands) {
   // Nop shape cast.
   if (source().getType() == result().getType())
     return source();
 
   // Canceling shape casts.
   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();
   }
   return {};
 }
 
 namespace {
 // Pattern to rewrite a ShapeCast(splat ConstantOp) -> ConstantOp.
 class ShapeCastConstantFolder final : public OpRewritePattern<ShapeCastOp> {

mlir/test/Dialect/Vector/canonicalize.mlir

 //       CHECK:   return %[[R]] : vector<4x2xf32>
 func @fold_extract_shapecast_negative(%arg0 : vector<16xf32>,
                              %arg1 : vector<8x4x2xf32>) -> vector<4x2xf32> {
   %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<2x4x2xf32>
   %r = vector.extract %0[1] : vector<2x4x2xf32>
   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>
+}
 
 // -----
 
 // CHECK-LABEL: fold_vector_transfers
 func @fold_vector_transfers(%A: memref<?x8xf32>) -> (vector<4x8xf32>, vector<4x9xf32>) {
   %c0 = arith.constant 0 : index
   %f0 = arith.constant 0.0 : f32