[Mlir-commits] [mlir] b0c4aaf - Allow only valid vector.shape_cast transitive folding

[Ahmed Taei]

Author: Ahmed Taei
Date: 2021-10-18T07:57:55-07:00
New Revision: b0c4aaff24fb3b55abb3415c14ad776f89c0205d

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

LOG: 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

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/VectorOps.cpp
    mlir/test/Dialect/Vector/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/VectorOps.cpp b/mlir/lib/Dialect/Vector/VectorOps.cpp
index ee87e5b3eec1f..fdec267934bb7 100644
--- a/mlir/lib/Dialect/Vector/VectorOps.cpp
+++ b/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();
   }

diff  --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 6c232cc0f7ab5..cf05308e8129b 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/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>
+}
 
 // -----