Support VectorTransfer splitting on writes also.

[tpopp]

Summary

VectorTransfer split previously only split read xfer ops. This adds
the same logic to write ops. The resulting code involves 2
conditionals for write ops while read ops only needed 1, but the created
ops are built upon the same patterns, so pattern matching/expectations
are all consistent other than in regards to the if/else ops.

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2665

nit: formatting

2804

Early return success(); here and drop the else plz.

mlir/test/Dialect/Vector/vector-transfer-full-partial-split.mlir

207

don't hardcode #map0, either capture it or wildcard it away.

342

Please drop these {{\[}}, you already captured properly (i.e. *without* the %, so you should just be able to write [%[[A]]])

[tpopp]

Respond to comments.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

static Value createScopedSubViewIntersection(VectorTransferOpInterface xferOp,
   int64_t memrefRank = xferOp.getShapedType().getRank();
   // TODO: relax this precondition, will require rank-reducing subviews.
   assert(memrefRank == alloc.getType().cast<MemRefType>().getRank() &&
          "Expected memref rank to match the alloc rank");
   ValueRange leadingIndices =
       xferOp.indices().take_front(xferOp.getLeadingShapedRank());
   SmallVector<OpFoldResult, 4> sizes;
   sizes.append(leadingIndices.begin(), leadingIndices.end());
+  auto isaWrite = isa<vector::TransferWriteOp>(xferOp);
   xferOp.zipResultAndIndexing([&](int64_t resultIdx, int64_t indicesIdx) {
     using MapList = ArrayRef<ArrayRef<AffineExpr>>;
     Value dimMemRef = memref_dim(xferOp.source(), indicesIdx);
     Value dimAlloc = memref_dim(alloc, resultIdx);
     Value index = xferOp.indices()[indicesIdx];
     AffineExpr i, j, k;
     bindDims(xferOp.getContext(), i, j, k);
     SmallVector<AffineMap, 4> maps =
         AffineMap::inferFromExprList(MapList{{i - j, k}});
     // affine_min(%dimMemRef - %index, %dimAlloc)
     Value affineMin = affine_min(index.getType(), maps[0],
                                  ValueRange{dimMemRef, index, dimAlloc});
     sizes.push_back(affineMin);
   });
 
   SmallVector<OpFoldResult, 4> indices = llvm::to_vector<4>(llvm::map_range(
       xferOp.indices(), [](Value idx) -> OpFoldResult { return idx; }));
   return memref_sub_view(
-      xferOp.source(), indices, sizes,
+      isaWrite ? alloc : xferOp.source(), indices, sizes,
       SmallVector<OpFoldResult>(memrefRank, OpBuilder(xferOp).getIndexAttr(1)));
 }
 
 /// Given an `xferOp` for which:
 ///   1. `inBoundsCond` and a `compatibleMemRefType` have been computed.
 ///   2. a memref of single vector `alloc` has been allocated.
 /// Produce IR resembling:
 /// ```
conditionBuilder(
                               zero);
 
         return viewAndIndices;
       },
       &fullPartialIfOp);
   return fullPartialIfOp;
 }
 
+/// Given an `xferOp` for which:
+///   1. `inBoundsCond` and a `compatibleMemRefType` have been computed.
+///   2. a memref of single vector `alloc` has been allocated.
+/// Produce IR resembling:
+/// ```
+///    %1:3 = scf.if (%inBounds) {
+///      memref.cast %A: memref<A...> to compatibleMemRefType
+///      scf.yield %view, ... : compatibleMemRefType, index, index
+///    } else {
+///      %3 = vector.type_cast %extra_alloc :
+///        memref<...> to memref<vector<...>>
+///      %4 = memref.cast %alloc: memref<B...> to compatibleMemRefType
+///      scf.yield %4, ... : compatibleMemRefType, index, index
+///   }
+/// ```
+static ValueRange getLocationToWriteFullVec(vector::TransferWriteOp xferOp,
+                                            TypeRange returnTypes,
+                                            Value inBoundsCond,
+                                            MemRefType compatibleMemRefType,
+                                            Value alloc) {
+  using namespace edsc;
+  using namespace edsc::intrinsics;
+  Value zero = std_constant_index(0);
+  Value memref = xferOp.source();
+  return conditionBuilder(
+      returnTypes, inBoundsCond,
+      [&]() -> scf::ValueVector {
+        Value res = memref;
+        if (compatibleMemRefType != xferOp.getShapedType())
+          res = memref_cast(memref, compatibleMemRefType);
+        scf::ValueVector viewAndIndices{res};
+        viewAndIndices.insert(viewAndIndices.end(), xferOp.indices().begin(),
+                              xferOp.indices().end());
+        return viewAndIndices;
+      },
+      [&]() -> scf::ValueVector {
+        Value casted = memref_cast(alloc, compatibleMemRefType);
+        scf::ValueVector viewAndIndices{casted};
+        viewAndIndices.insert(viewAndIndices.end(), xferOp.getTransferRank(),
+                              zero);
+        return viewAndIndices;
+      });
+}
+
+/// Given an `xferOp` for which:
+///   1. `inBoundsCond` has been computed.
+///   2. a memref of single vector `alloc` has been allocated.
+///   3. it originally wrote to %view
+/// Produce IR resembling:
+/// ```
+///    %notInBounds = xor %inBounds, %true
+///    scf.if (%notInBounds) {
+///      %3 = subview %alloc [...][...][...]
+///      linalg.copy(%3, %view)
+///   }
+/// ```
+static void createScopedFullPartialLinalgCopy(vector::TransferWriteOp xferOp,
+                                              Value inBoundsCond, Value alloc) {
+  using namespace edsc;
+  using namespace edsc::intrinsics;
+  auto &b = ScopedContext::getBuilderRef();
+  auto notInBounds = b.create<XOrOp>(
+      xferOp->getLoc(), inBoundsCond,
+      b.create<::mlir::ConstantIntOp>(xferOp.getLoc(), true, 1));
+
+  conditionBuilder(notInBounds, [&]() {
+    Value memRefSubView = createScopedSubViewIntersection(
+        cast<VectorTransferOpInterface>(xferOp.getOperation()), alloc);
+    linalg_copy(memRefSubView, xferOp.source());
+  });
+}
+
+/// Given an `xferOp` for which:
+///   1. `inBoundsCond` has been computed.
+///   2. a memref of single vector `alloc` has been allocated.
+///   3. it originally wrote to %view
+/// Produce IR resembling:
+/// ```
+///    %notInBounds = xor %inBounds, %true
+///    scf.if (%notInBounds) {
+///      %2 = load %alloc : memref<vector<...>>
+///      vector.transfer_write %2, %view[...] : memref<A...>, vector<...>
+///   }
+/// ```
+static void
+createScopedFullPartialVectorTransferWrite(vector::TransferWriteOp xferOp,
+                                           Value inBoundsCond, Value alloc) {
+  using namespace edsc;
+  using namespace edsc::intrinsics;
+  auto &b = ScopedContext::getBuilderRef();
+  auto notInBounds = b.create<XOrOp>(
+      xferOp->getLoc(), inBoundsCond,
+      b.create<::mlir::ConstantIntOp>(xferOp.getLoc(), true, 1));
+  conditionBuilder(notInBounds, [&]() {
+    BlockAndValueMapping mapping;
+
+    Value load = memref_load(vector_type_cast(
+        MemRefType::get({}, xferOp.vector().getType()), alloc));
+
+    mapping.map(xferOp.vector(), load);
+    b.clone(*xferOp.getOperation(), mapping);
+  });
+}
+
 /// Split a vector.transfer operation into an in-bounds (i.e., no out-of-bounds
 /// masking) fastpath and a slowpath.
+///
+/// For vector.transfer_read:
 /// If `ifOp` is not null and the result is `success, the `ifOp` points to the
 /// newly created conditional upon function return.
 /// To accomodate for the fact that the original vector.transfer indexing may be
 /// arbitrary and the slow path indexes @[0...0] in the temporary buffer, the
 /// scf.if op returns a view and values of type index.
-/// At this time, only vector.transfer_read case is implemented.
 ///
 /// Example (a 2-D vector.transfer_read):
 /// ```
 ///    %1 = vector.transfer_read %0[...], %pad : memref<A...>, vector<...>
 /// ```
 /// is transformed into:
 /// ```
 ///    %1:3 = scf.if (%inBounds) {
 ///      // fastpath, direct cast
 ///      memref.cast %A: memref<A...> to compatibleMemRefType
 ///      scf.yield %view : compatibleMemRefType, index, index
 ///    } else {
 ///      // slowpath, not in-bounds vector.transfer or linalg.copy.
 ///      memref.cast %alloc: memref<B...> to compatibleMemRefType
 ///      scf.yield %4 : compatibleMemRefType, index, index
 //     }
 ///    %0 = vector.transfer_read %1#0[%1#1, %1#2] {in_bounds = [true ... true]}
 /// ```
 /// where `alloc` is a top of the function alloca'ed buffer of one vector.
 ///
+/// For vector.transfer_write:
+/// There are 2 conditional blocks. First a block to decide which memref and
+/// indices to use for an unmasked, inbounds write. Then a conditional block to
+/// further copy a partial buffer into the final result in the slow path case.
+///
+/// Example (a 2-D vector.transfer_write):
+/// ```
+///    vector.transfer_write %arg, %0[...], %pad : memref<A...>, vector<...>
+/// ```
+/// is transformed into:
+/// ```
+///    %1:3 = scf.if (%inBounds) {
+///      memref.cast %A: memref<A...> to compatibleMemRefType
+///      scf.yield %view : compatibleMemRefType, index, index
+///    } else {
+///      memref.cast %alloc: memref<B...> to compatibleMemRefType
+///      scf.yield %4 : compatibleMemRefType, index, index
+///     }
+///    %0 = vector.transfer_write %arg, %1#0[%1#1, %1#2] {in_bounds = [true ...
+///                                                                    true]}
+///    scf.if (%notInBounds) {
+///      // slowpath: not in-bounds vector.transfer or linalg.copy.
+///    }
+/// ```
+/// where `alloc` is a top of the function alloca'ed buffer of one vector.
+///
 /// Preconditions:
 ///  1. `xferOp.permutation_map()` must be a minor identity map
 ///  2. the rank of the `xferOp.source()` and the rank of the `xferOp.vector()`
 ///  must be equal. This will be relaxed in the future but requires
 ///  rank-reducing subviews.
 LogicalResult mlir::vector::splitFullAndPartialTransfer(
     OpBuilder &b, VectorTransferOpInterface xferOp,
     VectorTransformsOptions options, scf::IfOp *ifOp) {
   using namespace edsc;
   using namespace edsc::intrinsics;
 
   if (options.vectorTransferSplit == VectorTransferSplit::None)
     return failure();
 
   SmallVector<bool, 4> bools(xferOp.getTransferRank(), true);
   auto inBoundsAttr = b.getBoolArrayAttr(bools);
   if (options.vectorTransferSplit == VectorTransferSplit::ForceInBounds) {
-    xferOp->setAttr(vector::TransferReadOp::getInBoundsAttrName(),
-                    inBoundsAttr);
+    xferOp->setAttr(xferOp.getInBoundsAttrName(), inBoundsAttr);
     return success();
   }
 
+  // Assert preconditions. Additionally, keep the variables in an inner scope to
+  // ensure they aren't used in the wrong scopes further down.
+  {
     assert(succeeded(splitFullAndPartialTransferPrecondition(xferOp)) &&
            "Expected splitFullAndPartialTransferPrecondition to hold");
+
     auto xferReadOp = dyn_cast<vector::TransferReadOp>(xferOp.getOperation());
+    auto xferWriteOp = dyn_cast<vector::TransferWriteOp>(xferOp.getOperation());
 
-  // TODO: add support for write case.
-  if (!xferReadOp)
+    if (!(xferReadOp || xferWriteOp))
       return failure();
-
-  if (xferReadOp.mask())
+    if (xferWriteOp && xferWriteOp.mask())
+      return failure();
+    if (xferReadOp && xferReadOp.mask())
       return failure();
+  }
 
   OpBuilder::InsertionGuard guard(b);
-  if (Operation *sourceOp = xferOp.source().getDefiningOp())
-    b.setInsertionPointAfter(sourceOp);
-  else
-    b.setInsertionPoint(xferOp);
   ScopedContext scope(b, xferOp.getLoc());
   Value inBoundsCond = createScopedInBoundsCond(
       cast<VectorTransferOpInterface>(xferOp.getOperation()));
   if (!inBoundsCond)
     return failure();
 
   // Top of the function `alloc` for transient storage.
   Value alloc;
   {
     FuncOp funcOp = xferOp->getParentOfType<FuncOp>();
     OpBuilder::InsertionGuard guard(b);
     b.setInsertionPointToStart(&funcOp.getRegion().front());
     auto shape = xferOp.getVectorType().getShape();
     Type elementType = xferOp.getVectorType().getElementType();
     alloc = memref_alloca(MemRefType::get(shape, elementType), ValueRange{},
                           b.getI64IntegerAttr(32));
   }
 
   MemRefType compatibleMemRefType =
       getCastCompatibleMemRefType(xferOp.getShapedType().cast<MemRefType>(),
                                   alloc.getType().cast<MemRefType>());
-
-  // Read case: full fill + partial copy -> in-bounds vector.xfer_read.
   SmallVector<Type, 4> returnTypes(1 + xferOp.getTransferRank(),
                                    b.getIndexType());
   returnTypes[0] = compatibleMemRefType;
+
+  if (auto xferReadOp =
+          dyn_cast<vector::TransferReadOp>(xferOp.getOperation())) {
+    // Read case: full fill + partial copy -> in-bounds vector.xfer_read.
     scf::IfOp fullPartialIfOp =
         options.vectorTransferSplit == VectorTransferSplit::VectorTransfer
             ? createScopedFullPartialVectorTransferRead(
                   xferReadOp, returnTypes, inBoundsCond, compatibleMemRefType,
                   alloc)
             : createScopedFullPartialLinalgCopy(xferReadOp, returnTypes,
                                                 inBoundsCond,
                                                 compatibleMemRefType, alloc);
     if (ifOp)
       *ifOp = fullPartialIfOp;
 
     // Set existing read op to in-bounds, it always reads from a full buffer.
     for (unsigned i = 0, e = returnTypes.size(); i != e; ++i)
       xferReadOp.setOperand(i, fullPartialIfOp.getResult(i));
-  xferOp->setAttr(vector::TransferReadOp::getInBoundsAttrName(), inBoundsAttr);
+
+    xferOp->setAttr(xferOp.getInBoundsAttrName(), inBoundsAttr);
+
+    return success();
+  }
+
+  auto xferWriteOp = cast<vector::TransferWriteOp>(xferOp.getOperation());
+
+  // Decide which location to write the entire vector to.
+  auto memrefAndIndices = getLocationToWriteFullVec(
+      xferWriteOp, returnTypes, inBoundsCond, compatibleMemRefType, alloc);
+
+  // Do an in bounds write to either the output or the extra allocated buffer.
+  // The operation is cloned to prevent deleting information needed for the
+  // later IR creation.
+  BlockAndValueMapping mapping;
+  mapping.map(xferWriteOp.source(), memrefAndIndices.front());
+  mapping.map(xferWriteOp.indices(), memrefAndIndices.drop_front());
+  auto *clone = b.clone(*xferWriteOp, mapping);
+  clone->setAttr(xferWriteOp.getInBoundsAttrName(), inBoundsAttr);
+
+  // Create a potential copy from the allocated buffer to the final output in
+  // the slow path case.
+  if (options.vectorTransferSplit == VectorTransferSplit::VectorTransfer)
+    createScopedFullPartialVectorTransferWrite(xferWriteOp, inBoundsCond,
+                                               alloc);
+  else
+    createScopedFullPartialLinalgCopy(xferWriteOp, inBoundsCond, alloc);
+
+  xferOp->erase();
 
   return success();
 }
 
 LogicalResult mlir::vector::VectorTransferFullPartialRewriter::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   auto xferOp = dyn_cast<VectorTransferOpInterface>(op);
   if (!xferOp || failed(splitFullAndPartialTransferPrecondition(xferOp)) ||
       failed(filter(xferOp)))
     return failure();
   rewriter.startRootUpdate(xferOp);
   if (succeeded(splitFullAndPartialTransfer(rewriter, xferOp, options))) {
     rewriter.finalizeRootUpdate(xferOp);
     return success();
   }
   rewriter.cancelRootUpdate(xferOp);
   return failure();
 }
 
 LogicalResult mlir::vector::PointwiseExtractPattern::matchAndRewrite(
     ExtractMapOp extract, PatternRewriter &rewriter) const {
   Operation *definedOp = extract.vector().getDefiningOp();
   if (!definedOp || definedOp->getNumResults() != 1)
     return failure();
   // TODO: Create an interfaceOp for elementwise operations.
   if (!isa<AddFOp>(definedOp))
     return failure();
   Location loc = extract.getLoc();
   SmallVector<Value, 4> extractOperands;

mlir/test/Dialect/Vector/vector-transfer-full-partial-split.mlir

-// RUN: mlir-opt %s -test-vector-transfer-full-partial-split | FileCheck %s
-// RUN: mlir-opt %s -test-vector-transfer-full-partial-split=use-linalg-copy | FileCheck %s --check-prefix=LINALG
+// RUN: mlir-opt %s -test-vector-transfer-full-partial-split -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-vector-transfer-full-partial-split=use-linalg-copy -split-input-file | FileCheck %s --check-prefix=LINALG
 
 // CHECK-DAG: #[[$map_p4:.*]] = affine_map<()[s0] -> (s0 + 4)>
 // CHECK-DAG: #[[$map_p8:.*]] = affine_map<()[s0] -> (s0 + 8)>
 // CHECK-DAG: #[[$map_2d_stride_1:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 
 // LINALG-DAG: #[[$map_p4:.*]] = affine_map<()[s0] -> (s0 + 4)>
 // LINALG-DAG: #[[$map_p8:.*]] = affine_map<()[s0] -> (s0 + 8)>
 // LINALG-DAG: #[[$map_2d_stride_1:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
func @split_vector_transfer_read_strided_2d(
   //      LINALG: %[[res:.*]] = vector.transfer_read {{.*}} {in_bounds = [true, true]} :
   // LINALG-SAME:   memref<?x8xf32, #[[$map_2d_stride_1]]>, vector<4x8xf32>
   %1 = vector.transfer_read %A[%i, %j], %f0 :
     memref<7x8xf32, offset:?, strides:[?, 1]>, vector<4x8xf32>
 
   // CHECK: return %[[res]] : vector<4x8xf32>
   return %1 : vector<4x8xf32>
 }
+
+// -----
+
+func @split_vector_transfer_write_2d(%V: vector<4x8xf32>, %A: memref<?x8xf32>, %i: index, %j: index) {
+  vector.transfer_write %V, %A[%i, %j] :
+    vector<4x8xf32>, memref<?x8xf32>
+  return
+}
+
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 + 4)>
+// CHECK-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 8)>
+// CHECK:     func @split_vector_transfer_write_2d(
+// CHECK-SAME:                                         %[[VEC:.*]]: vector<4x8xf32>,
+// CHECK-SAME:                                         %[[DEST:.*]]: memref<?x8xf32>,
+// CHECK-SAME:                                         %[[I:.*]]: index,
+// CHECK-SAME:                                         %[[J:.*]]: index) {
+// CHECK-DAG:       %[[C8:.*]] = constant 8 : index
+// CHECK-DAG:       %[[C0:.*]] = constant 0 : index
+// CHECK-DAG:       %[[CT:.*]] = constant true
+// CHECK:           %[[TEMP:.*]] = memref.alloca() {alignment = 32 : i64} : memref<4x8xf32>
+// CHECK:           %[[VAL_8:.*]] = affine.apply #[[MAP0]]()[%[[I]]]
+// CHECK:           %[[DIM0:.*]] = memref.dim %[[DEST]], %[[C0]] : memref<?x8xf32>
+// CHECK:           %[[DIM0_IN:.*]] = cmpi sle, %[[VAL_8]], %[[DIM0]] : index
+// CHECK:           %[[DIM1:.*]] = affine.apply #[[MAP1]]()[%[[J]]]
+// CHECK:           %[[DIM1_IN:.*]] = cmpi sle, %[[DIM1]], %[[C8]] : index
+// CHECK:           %[[IN_BOUNDS:.*]] = and %[[DIM0_IN]], %[[DIM1_IN]] : i1
+// CHECK:           %[[IN_BOUND_DEST:.*]]:3 = scf.if %[[IN_BOUNDS]] ->
+// CHECK-SAME:          (memref<?x8xf32>, index, index) {
+// CHECK:             scf.yield %[[DEST]], %[[I]], %[[J]] : memref<?x8xf32>, index, index
+// CHECK:           } else {
+// CHECK:             %[[VAL_15:.*]] = memref.cast %[[TEMP]]
+// CHECK-SAME:            : memref<4x8xf32> to memref<?x8xf32>
+// CHECK:             scf.yield %[[VAL_15]], %[[C0]], %[[C0]]
+// CHECK-SAME:            : memref<?x8xf32>, index, index
+// CHECK:           }
+// CHECK:           vector.transfer_write %[[VEC]],
+// CHECK-SAME:           %[[IN_BOUND_DEST:.*]]#0[%[[IN_BOUND_DEST]]#1, %[[IN_BOUND_DEST]]#2]
+// CHECK-SAME:           {in_bounds = [true, true]} : vector<4x8xf32>, memref<?x8xf32>
+// CHECK:           %[[OUT_BOUNDS:.*]] = xor %[[IN_BOUNDS]], %[[CT]] : i1
+// CHECK:           scf.if %[[OUT_BOUNDS]] {
+// CHECK:             %[[CASTED:.*]] = vector.type_cast %[[TEMP]]
+// CHECK-SAME:            : memref<4x8xf32> to memref<vector<4x8xf32>>
+// CHECK:             %[[RESULT_COPY:.*]] = memref.load %[[CASTED]][]
+// CHECK-SAME:            : memref<vector<4x8xf32>>
+// CHECK:             vector.transfer_write %[[RESULT_COPY]],
+// CHECK-SAME:            %[[DEST]][%[[I]], %[[J]]]
+// CHECK-SAME:            : vector<4x8xf32>, memref<?x8xf32>
+// CHECK:           }
+// CHECK:           return
+// CHECK:         }
+
+// LINALG-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 + 4)>
+// LINALG-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 8)>
+// LINALG-DAG: #[[MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0 - d1, 4)>
+// LINALG-DAG: #[[MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0 - d1, 8)>
+// LINALG-DAG: #[[MAP4:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 8 + s0 + d1)>
+// LINALG:     func @split_vector_transfer_write_2d(
+// LINALG-SAME:                                         %[[VEC:.*]]: vector<4x8xf32>,
+// LINALG-SAME:                                         %[[DEST:.*]]: memref<?x8xf32>,
+// LINALG-SAME:                                         %[[I:.*]]: index,
+// LINALG-SAME:                                         %[[J:.*]]: index) {
+// LINALG-DAG:       %[[CT:.*]] = constant true
+// LINALG-DAG:       %[[C0:.*]] = constant 0 : index
+// LINALG-DAG:       %[[C4:.*]] = constant 4 : index
+// LINALG-DAG:       %[[C8:.*]] = constant 8 : index
+// LINALG:           %[[TEMP:.*]] = memref.alloca() {alignment = 32 : i64} : memref<4x8xf32>
+// LINALG:           %[[IDX0:.*]] = affine.apply #[[MAP0]]()[%[[I]]]
+// LINALG:           %[[DIM0:.*]] = memref.dim %[[DEST]], %[[C0]] : memref<?x8xf32>
+// LINALG:           %[[DIM0_IN:.*]] = cmpi sle, %[[IDX0]], %[[DIM0]] : index
+// LINALG:           %[[DIM1:.*]] = affine.apply #[[MAP1]]()[%[[J]]]
+// LINALG:           %[[DIM1_IN:.*]] = cmpi sle, %[[DIM1]], %[[C8]] : index
+// LINALG:           %[[IN_BOUNDS:.*]] = and %[[DIM0_IN]], %[[DIM1_IN]] : i1
+// LINALG:           %[[IN_BOUND_DEST:.*]]:3 = scf.if %[[IN_BOUNDS]]
+// LINALG-SAME:          -> (memref<?x8xf32>, index, index) {
+// LINALG:             scf.yield %[[DEST]], %[[I]], %[[J]] : memref<?x8xf32>, index, index
+// LINALG:           } else {
+// LINALG:             %[[VAL_16:.*]] = memref.cast %[[TEMP]] : memref<4x8xf32> to memref<?x8xf32>
+// LINALG:             scf.yield %[[VAL_16]], %[[C0]], %[[C0]] : memref<?x8xf32>, index, index
+// LINALG:           }
+// LINALG:           vector.transfer_write %[[VEC]],
+// LINALG-SAME:          %[[IN_BOUND_DEST:.*]]#0[%[[IN_BOUND_DEST]]#1, %[[IN_BOUND_DEST]]#2]
+// LINALG-SAME:          {in_bounds = [true, true]} : vector<4x8xf32>, memref<?x8xf32>
+// LINALG:           %[[OUT_BOUNDS:.*]] = xor %[[IN_BOUNDS]], %[[CT]] : i1
+// LINALG:           scf.if %[[OUT_BOUNDS]] {
+// LINALG:             %[[VAL_19:.*]] = memref.dim %[[DEST]], %[[C0]] : memref<?x8xf32>
+// LINALG-DAG:         %[[VAL_20:.*]] = affine.min #[[MAP2]](%[[VAL_19]], %[[I]], %[[C4]])
+// LINALG-DAG:         %[[VAL_21:.*]] = affine.min #[[MAP3]](%[[C8]], %[[J]], %[[C8]])
+// LINALG:             %[[VAL_22:.*]] = memref.subview %[[TEMP]]
+// LINALG-SAME:            [%[[I]], %[[J]]] [%[[VAL_20]], %[[VAL_21]]]
+// LINALG-SAME:            [1, 1] : memref<4x8xf32> to memref<?x?xf32, #[[MAP4]]>
+// LINALG:             linalg.copy(%[[VAL_22]], %[[DEST]])
+// LINALG-SAME:            : memref<?x?xf32, #[[MAP4]]>, memref<?x8xf32>
+// LINALG:           }
+// LINALG:           return
+// LINALG:         }
+
+// -----
+
+func @split_vector_transfer_write_strided_2d(
+    %V: vector<4x8xf32>, %A: memref<7x8xf32, offset:?, strides:[?, 1]>,
+    %i: index, %j: index) {
+  vector.transfer_write %V, %A[%i, %j] :
+    vector<4x8xf32>, memref<7x8xf32, offset:?, strides:[?, 1]>
+  return
+}
+
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
+// CHECK-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 4)>
+// CHECK-DAG: #[[MAP2:.*]] = affine_map<()[s0] -> (s0 + 8)>
+// CHECK:   func @split_vector_transfer_write_strided_2d(
+// CHECK-SAME:                                                 %[[VEC:.*]]: vector<4x8xf32>,
+// CHECK-SAME:                                                 %[[DEST:.*]]: memref<7x8xf32, #[[MAP0]]>,
+// CHECK-SAME:                                                 %[[I:.*]]: index,
+// CHECK-SAME:                                                 %[[J:.*]]: index) {
+// CHECK-DAG:       %[[C7:.*]] = constant 7 : index
+// CHECK-DAG:       %[[C8:.*]] = constant 8 : index
+// CHECK-DAG:       %[[C0:.*]] = constant 0 : index
+// CHECK-DAG:       %[[CT:.*]] = constant true
+// CHECK:           %[[TEMP:.*]] = memref.alloca() {alignment = 32 : i64} : memref<4x8xf32>
+// CHECK:           %[[DIM0:.*]] = affine.apply #[[MAP1]]()[%[[I]]]
+// CHECK:           %[[DIM0_IN:.*]] = cmpi sle, %[[DIM0]], %[[C7]] : index
+// CHECK:           %[[DIM1:.*]] = affine.apply #[[MAP2]]()[%[[J]]]
+// CHECK:           %[[DIM1_IN:.*]] = cmpi sle, %[[DIM1]], %[[C8]] : index
+// CHECK:           %[[IN_BOUNDS:.*]] = and %[[DIM0_IN]], %[[DIM1_IN]] : i1
+// CHECK:           %[[IN_BOUND_DEST:.*]]:3 = scf.if %[[IN_BOUNDS]]
+// CHECK-SAME:          -> (memref<?x8xf32, #[[MAP0]]>, index, index) {
+// CHECK:             %[[VAL_15:.*]] = memref.cast %[[DEST]]
+// CHECK-SAME:            : memref<7x8xf32, #[[MAP0]]> to memref<?x8xf32, #[[MAP0]]>
+// CHECK:             scf.yield %[[VAL_15]], %[[I]], %[[J]]
+// CHECK-SAME:            : memref<?x8xf32, #[[MAP0]]>, index, index
+// CHECK:           } else {
+// CHECK:             %[[VAL_16:.*]] = memref.cast %[[TEMP]]
+// CHECK-SAME:            : memref<4x8xf32> to memref<?x8xf32, #[[MAP0]]>
+// CHECK:             scf.yield %[[VAL_16]], %[[C0]], %[[C0]]
+// CHECK-SAME:            : memref<?x8xf32, #[[MAP0]]>, index, index
+// CHECK:           }
+// CHECK:           vector.transfer_write %[[VEC]],
+// CHECK-SAME:          %[[IN_BOUND_DEST:.*]]#0
+// CHECK-SAME:          [%[[IN_BOUND_DEST]]#1, %[[IN_BOUND_DEST]]#2]
+// CHECK-SAME:          {in_bounds = [true, true]} : vector<4x8xf32>, memref<?x8xf32, #[[MAP0]]>
+// CHECK:           %[[OUT_BOUNDS:.*]] = xor %[[IN_BOUNDS]], %[[CT]] : i1
+// CHECK:           scf.if %[[OUT_BOUNDS]] {
+// CHECK:             %[[VAL_19:.*]] = vector.type_cast %[[TEMP]]
+// CHECK-SAME:            : memref<4x8xf32> to memref<vector<4x8xf32>>
+// CHECK:             %[[VAL_20:.*]] = memref.load %[[VAL_19]][]
+// CHECK-SAME:            : memref<vector<4x8xf32>>
+// CHECK:             vector.transfer_write %[[VAL_20]], %[[DEST]][%[[I]], %[[J]]]
+// CHECK-SAME:            : vector<4x8xf32>, memref<7x8xf32, #[[MAP0]]>
+// CHECK:           }
+// CHECK:           return
+// CHECK:         }
+
+// LINALG-DAG: #[[MAP0:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
+// LINALG-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 4)>
+// LINALG-DAG: #[[MAP2:.*]] = affine_map<()[s0] -> (s0 + 8)>
+// LINALG-DAG: #[[MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0 - d1, 4)>
+// LINALG-DAG: #[[MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0 - d1, 8)>
+// LINALG-DAG: #[[MAP5:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 8 + s0 + d1)>
+// LINALG:   func @split_vector_transfer_write_strided_2d(
+// LINALG-SAME:                                                 %[[VEC:.*]]: vector<4x8xf32>,
+// LINALG-SAME:                                                 %[[DEST:.*]]: memref<7x8xf32, #[[MAP0]]>,
+// LINALG-SAME:                                                 %[[I:.*]]: index,
+// LINALG-SAME:                                                 %[[J:.*]]: index) {
+// LINALG-DAG:       %[[C0:.*]] = constant 0 : index
+// LINALG-DAG:       %[[CT:.*]] = constant true
+// LINALG-DAG:       %[[C7:.*]] = constant 7 : index
+// LINALG-DAG:       %[[C4:.*]] = constant 4 : index
+// LINALG-DAG:       %[[C8:.*]] = constant 8 : index
+// LINALG:           %[[TEMP:.*]] = memref.alloca() {alignment = 32 : i64} : memref<4x8xf32>
+// LINALG:           %[[DIM0:.*]] = affine.apply #[[MAP1]]()[%[[I]]]
+// LINALG:           %[[DIM0_IN:.*]] = cmpi sle, %[[DIM0]], %[[C7]] : index
+// LINALG:           %[[DIM1:.*]] = affine.apply #[[MAP2]]()[%[[J]]]
+// LINALG:           %[[DIM1_IN:.*]] = cmpi sle, %[[DIM1]], %[[C8]] : index
+// LINALG:           %[[IN_BOUNDS:.*]] = and %[[DIM0_IN]], %[[DIM1_IN]] : i1
+// LINALG:           %[[IN_BOUND_DEST:.*]]:3 = scf.if %[[IN_BOUNDS]]
+// LINALG-SAME:          -> (memref<?x8xf32, #[[MAP0]]>, index, index) {
+// LINALG:             %[[VAL_16:.*]] = memref.cast %[[DEST]]
+// LINALG-SAME:            : memref<7x8xf32, #[[MAP0]]> to memref<?x8xf32, #[[MAP0]]>
+// LINALG:             scf.yield %[[VAL_16]], %[[I]], %[[J]]
+// LINALG-SAME:            : memref<?x8xf32, #[[MAP0]]>, index, index
+// LINALG:           } else {
+// LINALG:             %[[VAL_17:.*]] = memref.cast %[[TEMP]]
+// LINALG-SAME:            : memref<4x8xf32> to memref<?x8xf32, #[[MAP0]]>
+// LINALG:             scf.yield %[[VAL_17]], %[[C0]], %[[C0]]
+// LINALG-SAME:            : memref<?x8xf32, #[[MAP0]]>, index, index
+// LINALG:           }
+// LINALG:           vector.transfer_write %[[VEC]],
+// LINALG-SAME:          %[[IN_BOUND_DEST:.*]]#0
+// LINALG-SAME:          [%[[IN_BOUND_DEST]]#1, %[[IN_BOUND_DEST]]#2]
+// LINALG-SAME:          {in_bounds = [true, true]}
+// LINALG-SAME:          : vector<4x8xf32>, memref<?x8xf32, #[[MAP0]]>
+// LINALG:           %[[OUT_BOUNDS:.*]] = xor %[[IN_BOUNDS]], %[[CT]] : i1
+// LINALG:           scf.if %[[OUT_BOUNDS]] {
+// LINALG-DAG:         %[[VAL_20:.*]] = affine.min #[[MAP3]](%[[C7]], %[[I]], %[[C4]])
+// LINALG-DAG:         %[[VAL_21:.*]] = affine.min #[[MAP4]](%[[C8]], %[[J]], %[[C8]])
+// LINALG:             %[[VAL_22:.*]] = memref.subview %[[TEMP]]
+// LINALG-SAME:            [%[[I]], %[[J]]] [%[[VAL_20]], %[[VAL_21]]]
+// LINALG-SAME:            [1, 1] : memref<4x8xf32> to memref<?x?xf32, #[[MAP5]]>
+// LINALG:             linalg.copy(%[[VAL_22]], %[[DEST]])
+// LINALG-SAME:            : memref<?x?xf32, #[[MAP5]]>, memref<7x8xf32, #[[MAP0]]>
+// LINALG:           }
+// LINALG:           return
+// LINALG:         }