[Mlir-commits] [mlir] dfb7b3f - [mlir][VectorOps] Fall back to a loop when accessing a vector from a strided memref

[Benjamin Kramer]

Author: Benjamin Kramer
Date: 2020-09-03T16:05:38+02:00
New Revision: dfb7b3fe02c0e84968960f4aef88361e6de10874

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

LOG: [mlir][VectorOps] Fall back to a loop when accessing a vector from a strided memref

The scalar loop is slow but correct.

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

Added: 
    

Modified: 
    mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
    mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
index 267aea90cc9d..3c501f046f07 100644
--- a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
+++ b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
@@ -547,7 +547,15 @@ LogicalResult VectorTransferRewriter<TransferReadOp>::matchAndRewrite(
   using namespace mlir::edsc::op;
 
   TransferReadOp transfer = cast<TransferReadOp>(op);
-  if (transfer.permutation_map().isMinorIdentity()) {
+
+  // Fall back to a loop if the fastest varying stride is not 1 or it is
+  // permuted.
+  int64_t offset;
+  SmallVector<int64_t, 4> strides;
+  auto successStrides =
+      getStridesAndOffset(transfer.getMemRefType(), strides, offset);
+  if (succeeded(successStrides) && strides.back() == 1 &&
+      transfer.permutation_map().isMinorIdentity()) {
     // If > 1D, emit a bunch of loops around 1-D vector transfers.
     if (transfer.getVectorType().getRank() > 1)
       return NDTransferOpHelper<TransferReadOp>(rewriter, transfer, options)
@@ -621,7 +629,15 @@ LogicalResult VectorTransferRewriter<TransferWriteOp>::matchAndRewrite(
   using namespace edsc::op;
 
   TransferWriteOp transfer = cast<TransferWriteOp>(op);
-  if (transfer.permutation_map().isMinorIdentity()) {
+
+  // Fall back to a loop if the fastest varying stride is not 1 or it is
+  // permuted.
+  int64_t offset;
+  SmallVector<int64_t, 4> strides;
+  auto successStrides =
+      getStridesAndOffset(transfer.getMemRefType(), strides, offset);
+  if (succeeded(successStrides) && strides.back() == 1 &&
+      transfer.permutation_map().isMinorIdentity()) {
     // If > 1D, emit a bunch of loops around 1-D vector transfers.
     if (transfer.getVectorType().getRank() > 1)
       return NDTransferOpHelper<TransferWriteOp>(rewriter, transfer, options)

diff  --git a/mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir b/mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir
index b19ea9dde793..986bfe176351 100644
--- a/mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir
+++ b/mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir
@@ -457,3 +457,28 @@ func @transfer_write_minor_identity(%A : vector<3x3xf32>, %B : memref<?x?x?x?xf3
 //       CHECK:      }
 //       CHECK:    }
 //       CHECK:    return
+
+// -----
+
+func @transfer_read_strided(%A : memref<8x4xf32, affine_map<(d0, d1) -> (d0 + d1 * 8)>>) -> vector<4xf32> {
+  %c0 = constant 0 : index
+  %f0 = constant 0.0 : f32
+  %0 = vector.transfer_read %A[%c0, %c0], %f0
+      : memref<8x4xf32, affine_map<(d0, d1) -> (d0 + d1 * 8)>>, vector<4xf32>
+  return %0 : vector<4xf32>
+}
+
+// CHECK-LABEL: transfer_read_strided(
+// CHECK: scf.for
+// CHECK: load
+
+func @transfer_write_strided(%A : vector<4xf32>, %B : memref<8x4xf32, affine_map<(d0, d1) -> (d0 + d1 * 8)>>) {
+  %c0 = constant 0 : index
+  vector.transfer_write %A, %B[%c0, %c0] :
+    vector<4xf32>, memref<8x4xf32, affine_map<(d0, d1) -> (d0 + d1 * 8)>>
+  return
+}
+
+// CHECK-LABEL: transfer_write_strided(
+// CHECK: scf.for
+// CHECK: store