[Mlir-commits] [mlir] 8a5cb93 - [mlir][vector] Check minor identity map in FoldExtractSliceIntoTransferRead

[Lei Zhang]

Author: Lei Zhang
Date: 2022-09-14T17:24:02-04:00
New Revision: 8a5cb939e7069eb8be8d46f6db5b598dc1954218

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

LOG: [mlir][vector] Check minor identity map in FoldExtractSliceIntoTransferRead

vecotr.transfer_read ops with minor identity indexing map is rank
reducing, with implicit leading unit dimensions. This should be
a natural extension to support in addition to full identity indexing
maps.

Reviewed By: ThomasRaoux

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 37725cf3c90bf..65cb37737f329 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -3294,7 +3294,7 @@ struct FoldExtractSliceIntoTransferRead
       return failure();
     if (xferOp.hasOutOfBoundsDim())
       return failure();
-    if (!xferOp.getPermutationMap().isIdentity())
+    if (!xferOp.getPermutationMap().isMinorIdentity())
       return failure();
     if (xferOp.getMask())
       return failure();

diff  --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index ac5b857938581..86ee934ad25c6 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -1169,6 +1169,24 @@ func.func @transfer_read_of_extract_slice(%t : tensor<?x?xf32>, %s1 : index, %s2
 
 // -----
 
+// CHECK-LABEL: func @transfer_read_of_extract_slice(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index
+//   CHECK-DAG:   %[[c4:.*]] = arith.constant 4 : index
+//   CHECK-DAG:   %[[c8:.*]] = arith.constant 8 : index
+//       CHECK:   %[[add:.*]] = arith.addi %[[s1]], %[[c4]]
+//       CHECK:   %[[r:.*]] = vector.transfer_read %[[t]][%[[c8]], %[[add]]], %{{.*}} {in_bounds = [true]} : tensor<?x?xf32>, vector<6xf32>
+//       CHECK:   return %[[r]]
+func.func @transfer_read_of_extract_slice(%t : tensor<?x?xf32>, %s1 : index, %s2 : index) -> vector<6xf32> {
+  %c3 = arith.constant 3 : index
+  %c4 = arith.constant 4 : index
+  %cst = arith.constant 0.0 : f32
+  %0 = tensor.extract_slice %t[5, %s1] [10, %s2] [1, 1] : tensor<?x?xf32> to tensor<10x?xf32>
+  %1 = vector.transfer_read %0[%c3, %c4], %cst {in_bounds = [true]} : tensor<10x?xf32>, vector<6xf32>
+  return %1 : vector<6xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @transfer_read_of_extract_slice_rank_reducing(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index
 //   CHECK-DAG:   %[[c3:.*]] = arith.constant 3 : index