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

[bkramer]

Summary

The scalar loop is slow but correct.

[nicolasvasilache]

Note I am also making some changes to the scalar path for other correctness purposes.
There may be something interesting to do down the line.

[bkramer]

Also handle transfer_write

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

 
 /// Performs the rewrite.
 template <>
 LogicalResult VectorTransferRewriter<TransferReadOp>::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   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)
           .doReplace();
     // If 1-D this is now handled by the target-specific lowering.
     if (transfer.getVectorType().getRank() == 1)
       return failure();
   }

 /// TODO: implement alternatives to clipping.
 /// TODO: support non-data-parallel operations.
 template <>
 LogicalResult VectorTransferRewriter<TransferWriteOp>::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   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)
           .doReplace();
     // If 1-D this is now handled by the target-specific lowering.
     if (transfer.getVectorType().getRank() == 1)
       return failure();
   }

mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir

 //       CHECK:   affine.for %[[arg2:.*]] = 0 to 3 {
 //       CHECK:      %[[cmp:.*]] = cmpi "slt", %[[arg2]], %[[d]] : index
 //       CHECK:      scf.if %[[cmp]] {
 //       CHECK:        %[[tmp:.*]] = load %[[m]][%[[arg2]]] : memref<3xvector<3xf32>>
 //       CHECK:        vector.transfer_write %[[tmp]], %[[B]][%[[c0]], %[[c0]], %[[arg2]], %[[c0]]] : vector<3xf32>, memref<?x?x?x?xf32>
 //       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