[mlir][vector] Move transferOp on tensor opt to folder/canonicalization

[ThomasRaoux]

Summary

Move the existing optimization for transfer op on tensor to folder and canonicalization. This handles the write after write case and read after write and also add write after read case.

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorOps.cpp

2525

I'd suggest to rename foldRAW so it becomes very clear what type of folding this is and how it relates to the others.

2530

I'd suggest to rename checkSameValueRAW (or a better name but with RAW in it) so it becomes very clear what type of folding this is and how it relates to the others.

2770

I'd suggest to rename foldWAR so it becomes very clear what type of folding this is and how it relates to the others.

2779

I'd suggest hoisting that in a helper checkSameValueWAR (or a better name but with RAW in it) so it becomes very clear what type of folding this is and how it relates to the others.

2838

I'd suggest to rename foldRAW so it becomes very clear what type of folding this is and how it relates to the others.

2841

Why does this need to be a canonicalization pattern and not a simple fold ?
At some point in space time, https://reviews.llvm.org/D100586 had it working as a folder.

2848

I'd suggest to rename checkSameValueWAW (or a better name but with WAW in it) so it becomes very clear what type of folding this is and how it relates to the others.

[ThomasRaoux]

Renaming based on review comments

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorOps.cpp

2841

If we want to support looking more than one level up in the WaW case we need to modify a Transfer_write that isn't the root like in the example here:

%w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]}
  : vector<1x4xf32>, tensor<4x4xf32>
%w1 = vector.transfer_write %v0, %arg0[%c2, %c0] {in_bounds = [true, true]}
  : vector<1x4xf32>, tensor<4x4xf32>
%w2 = vector.transfer_write %v1, %w1[%c1, %c0] {in_bounds = [true, true]}
  : vector<1x4xf32>, tensor<4x4xf32>

In this case we need to update %w1.

It could be implemented as a folding starting from %w0 considering only the case where the write has a single use and going through the SSA chain in the other direction. Do you think this is better?

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorOps.cpp

2841

I had missed that is more powerful than what you could get with a folder and hop-1 use-def.
This is great as is, thanks!

mlir/include/mlir/Dialect/Vector/VectorOps.td

OpBuilder<(ins "Value":$vector, "Value":$source, "ValueRange":$indices,
       "AffineMapAttr":$permutationMap, "ArrayAttr":$inBounds)>,
     OpBuilder<(ins "Value":$vector, "Value":$source, "ValueRange":$indices,
       "AffineMap":$permutationMap, "Value":$mask, "ArrayAttr":$inBounds)>,
     OpBuilder<(ins "Value":$vector, "Value":$source, "ValueRange":$indices,
       "AffineMap":$permutationMap, "ArrayAttr":$inBounds)>,
   ];
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 def Vector_LoadOp : Vector_Op<"load"> {
   let summary = "reads an n-D slice of memory into an n-D vector";
   let description = [{
     The 'vector.load' operation reads an n-D slice of memory into an n-D
     vector. It takes a 'base' memref, an index for each memref dimension and a
     result vector type as arguments. It returns a value of the result vector

mlir/include/mlir/Dialect/Vector/VectorUtils.h

 class Location;
 class OpBuilder;
 class Operation;
 class ShapedType;
 class Value;
 class VectorType;
 class VectorTransferOpInterface;
 
+namespace vector {
+class TransferWriteOp;
+class TransferReadOp;
+} // namespace vector
+
 /// Return the number of elements of basis, `0` if empty.
 int64_t computeMaxLinearIndex(ArrayRef<int64_t> basis);
 
 /// Given a shape with sizes greater than 0 along all dimensions,
 /// return the distance, in number of elements, between a slice in a dimension
 /// and the next slice in the same dimension.
 ///   e.g. shape[3, 4, 5] -> linearization_basis[20, 5, 1]
 SmallVector<int64_t, 8> computeStrides(ArrayRef<int64_t> shape);
bool isDisjointTransferSet(VectorTransferOpInterface transferA,
                            VectorTransferOpInterface transferB);
 
 /// Same behavior as `isDisjointTransferSet` but doesn't require the operations
 /// to have the same tensor/memref. This allows comparing operations accessing
 /// different tensors.
 bool isDisjointTransferIndices(VectorTransferOpInterface transferA,
                                VectorTransferOpInterface transferB);
 
+/// Return true if the transfer_write fully writes the data accessed by the
+/// transfer_read.
+bool checkSameValueRAW(vector::TransferWriteOp defWrite,
+                       vector::TransferReadOp read);
+
+/// Return true if the write op fully over-write the priorWrite transfer_write
+/// op.
+bool checkSameValueWAW(vector::TransferWriteOp write,
+                       vector::TransferWriteOp priorWrite);
+
 namespace matcher {
 
 /// Matches vector.transfer_read, vector.transfer_write and ops that return a
 /// vector type that is a multiple of the sub-vector type. This allows passing
 /// over other smaller vector types in the function and avoids interfering with
 /// operations on those.
 /// This is a first approximation, it can easily be extended in the future.
 /// TODO: this could all be much simpler if we added a bit that a vector type to
 /// mark that a vector is a strict super-vector but it still does not warrant
 /// adding even 1 extra bit in the IR for now.
 bool operatesOnSuperVectorsOf(Operation &op, VectorType subVectorType);
 
 } // end namespace matcher
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_VECTOR_VECTORUTILS_H_

mlir/lib/Dialect/Vector/VectorOps.cpp

if (!changed)
     return failure();
   // OpBuilder is only used as a helper to build an I64ArrayAttr.
   OpBuilder b(op.getContext());
   op->setAttr(TransferOp::getInBoundsAttrName(),
               b.getBoolArrayAttr(newInBounds));
   return success();
 }
 
+///  ```
+///  %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]}
+///    : vector<1x4xf32>, tensor<4x4xf32>
+///  %0 = vector.transfer_read %w0[%c1, %c0], %cf0 {in_bounds = [true, true]}
+///    : tensor<4x4xf32>, vector<1x4xf32>
+///  ```
+///  -> Folds into
+///  ```
+///  %v0
+///  ```
+static Value foldRAW(TransferReadOp readOp) {
+  if (!readOp.getShapedType().isa<RankedTensorType>())
+    return {};
+  auto defWrite = readOp.source().getDefiningOp<vector::TransferWriteOp>();
+  while (defWrite) {
+    if (checkSameValueRAW(defWrite, readOp))
+      return defWrite.vector();
+    if (!isDisjointTransferIndices(
+            cast<VectorTransferOpInterface>(defWrite.getOperation()),
+            cast<VectorTransferOpInterface>(readOp.getOperation())))
+      break;
+    defWrite = defWrite.source().getDefiningOp<vector::TransferWriteOp>();
+  }
+  return {};
+}
+
 OpFoldResult TransferReadOp::fold(ArrayRef<Attribute>) {
+  if (Value vec = foldRAW(*this))
+    return vec;
   /// transfer_read(memrefcast) -> transfer_read
   if (succeeded(foldTransferInBoundsAttribute(*this)))
     return getResult();
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
   if (succeeded(foldTensorCast(*this)))
     return getResult();
   return OpFoldResult();
static LogicalResult foldReadInitWrite(TransferWriteOp write,
   if (llvm::any_of(read.indices(), isNotConstantZero) ||
       llvm::any_of(write.indices(), isNotConstantZero))
     return failure();
   // Success.
   results.push_back(read.source());
   return success();
 }
 
+static bool checkSameValueWAR(vector::TransferReadOp read,
+                              vector::TransferWriteOp write) {
+  return read.source() == write.source() && read.indices() == write.indices() &&
+         read.permutation_map() == write.permutation_map() &&
+         read.getVectorType() == write.getVectorType() && !read.mask() &&
+         !write.mask();
+}
+/// Fold transfer_write write after read:
+/// ```
+///    %t0 = ...
+///    %v = vector.transfer_read %t0[%c0...] :
+///      tensor<static_sizesxf32>, vector<static_sizesxf32>
+///    %t1 = vector.transfer_write %v, %t0[%c0...] :
+///      vector<static_sizesxf32>, tensor<static_sizesxf32>
+/// ```
+///
+/// into:
+///
+/// ```
+///    %t0
+/// ```
+static LogicalResult foldWAR(TransferWriteOp write,
+                             SmallVectorImpl<OpFoldResult> &results) {
+  if (!write.source().getType().isa<RankedTensorType>())
+    return failure();
+  auto read = write.vector().getDefiningOp<vector::TransferReadOp>();
+  if (!read)
+    return failure();
+
+  if (!checkSameValueWAR(read, write))
+    return failure();
+  results.push_back(read.source());
+  return success();
+}
+
 LogicalResult TransferWriteOp::fold(ArrayRef<Attribute> operands,
                                     SmallVectorImpl<OpFoldResult> &results) {
   if (succeeded(foldReadInitWrite(*this, operands, results)))
     return success();
+  if (succeeded(foldWAR(*this, results)))
+    return success();
   if (succeeded(foldTransferInBoundsAttribute(*this)))
     return success();
   return foldMemRefCast(*this);
 }
 
 Optional<SmallVector<int64_t, 4>> TransferWriteOp::getShapeForUnroll() {
   return llvm::to_vector<4>(getVectorType().getShape());
 }
 
 void TransferWriteOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
   if (getShapedType().isa<MemRefType>())
     effects.emplace_back(MemoryEffects::Write::get(), source(),
                          SideEffects::DefaultResource::get());
 }
 
+namespace {
+/// Remove dead transfer write from the SSA chain so that it an be eliminated by
+/// DCE
+/// ```
+///  %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]}
+///    : vector<1x4xf32>, tensor<4x4xf32>
+///  %w1 = vector.transfer_write %v0, %w0[%c2, %c0] {in_bounds = [true, true]}
+///    : vector<1x4xf32>, tensor<4x4xf32>
+///  %w2 = vector.transfer_write %v1, %w1[%c1, %c0] {in_bounds = [true, true]}
+///    : vector<1x4xf32>, tensor<4x4xf32>
+/// ```
+///
+/// into:
+///
+/// ```
+///  %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]}
+///    : vector<1x4xf32>, tensor<4x4xf32>
+///  %w1 = vector.transfer_write %v0, %arg0[%c2, %c0] {in_bounds = [true, true]}
+///    : vector<1x4xf32>, tensor<4x4xf32>
+///  %w2 = vector.transfer_write %v1, %w1[%c1, %c0] {in_bounds = [true, true]}
+///    : vector<1x4xf32>, tensor<4x4xf32>
+/// ```
+///
+/// `%w0 = vector.transfer_write` op will be removed by DCE if it doesn't have
+/// any other uses.
+class foldWAW final : public OpRewritePattern<TransferWriteOp> {
+public:
+  using OpRewritePattern<TransferWriteOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(TransferWriteOp writeOp,
+                                PatternRewriter &rewriter) const override {
+    if (!writeOp.getShapedType().isa<RankedTensorType>())
+      return failure();
+    vector::TransferWriteOp writeToModify = writeOp;
+
+    auto defWrite = writeOp.source().getDefiningOp<vector::TransferWriteOp>();
+    while (defWrite) {
+      if (checkSameValueWAW(writeOp, defWrite)) {
+        writeToModify.sourceMutable().assign(defWrite.source());
+        return success();
+      }
+      if (!isDisjointTransferIndices(
+              cast<VectorTransferOpInterface>(defWrite.getOperation()),
+              cast<VectorTransferOpInterface>(writeOp.getOperation())))
+        break;
+      // If the previous write op doesn't have any other use we an safely look
+      // at the previous store to see if it can be removed.
+      if (!defWrite->hasOneUse())
+        break;
+      writeToModify = defWrite;
+      defWrite = defWrite.source().getDefiningOp<vector::TransferWriteOp>();
+    }
+    return failure();
+  }
+};
+} // namespace
+
+void TransferWriteOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                                  MLIRContext *context) {
+  results.add<foldWAW>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // LoadOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verifyLoadStoreMemRefLayout(Operation *op,
                                                  MemRefType memRefTy) {
   auto affineMaps = memRefTy.getAffineMaps();
   if (!affineMaps.empty())

mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp

 /// an ancestor of the op.
 static Operation *findAncestorOpInRegion(Region *region, Operation *op) {
   for (; op != nullptr && op->getParentRegion() != region;
        op = op->getParentOp())
     ;
   return op;
 }
 
-/// Return true if the transfer_write fully writes the data accessed by the
-/// transfer_read.
-static bool transferEncompasses(vector::TransferWriteOp defWrite,
-                                vector::TransferReadOp read) {
-  return !defWrite.hasOutOfBoundsDim() &&
-         defWrite.indices() == read.indices() &&
-         defWrite.getVectorType() == read.getVectorType() &&
-         defWrite.permutation_map() == read.permutation_map();
-}
-
-/// Return true if the write op fully over-write the priorWrite transfer_write
-/// op.
-static bool transferEncompasses(vector::TransferWriteOp write,
-                                vector::TransferWriteOp priorWrite) {
-  return priorWrite.indices() == write.indices() &&
-         priorWrite.getVectorType() == write.getVectorType() &&
-         priorWrite.permutation_map() == write.permutation_map();
-}
-
 namespace {
 
 class TransferOptimization {
 public:
   TransferOptimization(FuncOp func) : dominators(func), postDominators(func) {}
   void deadStoreOp(vector::TransferWriteOp);
-  void deadStoreOpTensor(vector::TransferWriteOp);
   void storeToLoadForwarding(vector::TransferReadOp);
-  void storeToLoadForwardingTensor(vector::TransferReadOp);
   void removeDeadOp() {
     for (Operation *op : opToErase)
       op->erase();
     opToErase.clear();
   }
 
 private:
   bool isReachable(Operation *start, Operation *dest);
LLVM_DEBUG(DBGS() << "Candidate for dead store: " << *write.getOperation()
                     << "\n");
   llvm::SmallVector<Operation *, 8> reads;
   Operation *firstOverwriteCandidate = nullptr;
   for (auto *user : write.source().getUsers()) {
     if (user == write.getOperation())
       continue;
     if (auto nextWrite = dyn_cast<vector::TransferWriteOp>(user)) {
       // Check candidate that can override the store.
-      if (transferEncompasses(nextWrite, write) &&
+      if (checkSameValueWAW(nextWrite, write) &&
           postDominators.postDominates(nextWrite, write)) {
         if (firstOverwriteCandidate == nullptr ||
             postDominators.postDominates(firstOverwriteCandidate, nextWrite))
           firstOverwriteCandidate = nextWrite;
         else
           assert(
               postDominators.postDominates(nextWrite, firstOverwriteCandidate));
       }
if (isa<vector::TransferReadOp>(user))
       continue;
     if (auto write = dyn_cast<vector::TransferWriteOp>(user)) {
       // If there is a write, but we can prove that it is disjoint we can ignore
       // the write.
       if (isDisjointTransferSet(
               cast<VectorTransferOpInterface>(write.getOperation()),
               cast<VectorTransferOpInterface>(read.getOperation())))
         continue;
-      if (dominators.dominates(write, read) &&
-          transferEncompasses(write, read)) {
+      if (dominators.dominates(write, read) && checkSameValueRAW(write, read)) {
         if (lastwrite == nullptr || dominators.dominates(lastwrite, write))
           lastwrite = write;
         else
           assert(dominators.dominates(write, lastwrite));
         continue;
       }
     }
     blockingWrites.push_back(user);
void TransferOptimization::storeToLoadForwarding(vector::TransferReadOp read) {
   }
 
   LLVM_DEBUG(DBGS() << "Forward value from " << *lastwrite.getOperation()
                     << " to: " << *read.getOperation() << "\n");
   read.replaceAllUsesWith(lastwrite.vector());
   opToErase.push_back(read.getOperation());
 }
 
-/// Walk up the SSA links, if any write gets fully overwritten we can skip it.
-/// If it has no more uses it becomes dead.
-void TransferOptimization::deadStoreOpTensor(vector::TransferWriteOp write) {
-  auto defWrite = write.source().getDefiningOp<vector::TransferWriteOp>();
-  while (defWrite) {
-    if (transferEncompasses(write, defWrite)) {
-      write.sourceMutable().assign(defWrite.source());
-      if (defWrite->use_empty())
-        opToErase.push_back(defWrite.getOperation());
-      return;
-    }
-    if (!isDisjointTransferIndices(
-            cast<VectorTransferOpInterface>(defWrite.getOperation()),
-            cast<VectorTransferOpInterface>(write.getOperation())))
-      break;
-    defWrite = defWrite.source().getDefiningOp<vector::TransferWriteOp>();
-  }
-}
-
-/// Walk up the SSA links, if any write fully match the written vector we can
-/// replace the read by the vector. The read becomes dead and can be removed.
-void TransferOptimization::storeToLoadForwardingTensor(
-    vector::TransferReadOp read) {
-  auto defWrite = read.source().getDefiningOp<vector::TransferWriteOp>();
-  while (defWrite) {
-    if (transferEncompasses(defWrite, read)) {
-      read.replaceAllUsesWith(defWrite.vector());
-      opToErase.push_back(read.getOperation());
-      return;
-    }
-    if (!isDisjointTransferIndices(
-            cast<VectorTransferOpInterface>(defWrite.getOperation()),
-            cast<VectorTransferOpInterface>(read.getOperation())))
-      break;
-    defWrite = defWrite.source().getDefiningOp<vector::TransferWriteOp>();
-  }
-}
-
 } // namespace
 
 void mlir::vector::transferOpflowOpt(FuncOp func) {
   TransferOptimization opt(func);
   // Run store to load forwarding first since it can expose more dead store
   // opportunity.
   func.walk([&](vector::TransferReadOp read) {
     if (read.getShapedType().isa<MemRefType>())
       opt.storeToLoadForwarding(read);
-    else
-      opt.storeToLoadForwardingTensor(read);
   });
   opt.removeDeadOp();
   func.walk([&](vector::TransferWriteOp write) {
     if (write.getShapedType().isa<MemRefType>())
       opt.deadStoreOp(write);
-    else
-      opt.deadStoreOpTensor(write);
   });
   opt.removeDeadOp();
 }

mlir/lib/Dialect/Vector/VectorUtils.cpp

 }
 
 bool mlir::isDisjointTransferSet(VectorTransferOpInterface transferA,
                                  VectorTransferOpInterface transferB) {
   if (transferA.source() != transferB.source())
     return false;
   return isDisjointTransferIndices(transferA, transferB);
 }
+
+bool mlir::checkSameValueRAW(vector::TransferWriteOp defWrite,
+                             vector::TransferReadOp read) {
+  return !defWrite.hasOutOfBoundsDim() && !defWrite.mask() && !read.mask() &&
+         defWrite.indices() == read.indices() &&
+         defWrite.getVectorType() == read.getVectorType() &&
+         defWrite.permutation_map() == read.permutation_map();
+}
+
+bool mlir::checkSameValueWAW(vector::TransferWriteOp write,
+                             vector::TransferWriteOp priorWrite) {
+  return priorWrite.indices() == write.indices() &&
+         priorWrite.mask() == write.mask() &&
+         priorWrite.getVectorType() == write.getVectorType() &&
+         priorWrite.permutation_map() == write.permutation_map();
+}

mlir/test/Dialect/Vector/canonicalize.mlir

func @transfer_folding_1(%t0: tensor<2x3x4xf32>, %t1: tensor<2x3x4xf32>)
   // CHECK-NEXT: some_op_that_may_have_side_effects
   %t3 = "some_op_that_may_have_side_effects"() : () -> (tensor<2x3x4xf32>)
   %r2 = vector.transfer_write %v, %t0[%c0, %c0, %c0] {in_bounds = [true, true, true]} :
     vector<2x3x4xf32>, tensor<2x3x4xf32>
 
   // CHECK-NEXT: return %[[T0]], %[[T0]], %[[T0]]
   return %r0, %r1, %r2: tensor<2x3x4xf32>, tensor<2x3x4xf32>, tensor<2x3x4xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @store_after_load_tensor
+//  CHECK-SAME: (%[[ARG:.*]]: tensor<4x4xf32>)
+//   CHECK-NOT:   vector.transfer_read
+//   CHECK-NOT:   vector.transfer_write
+//       CHECK:   return %[[ARG]] : tensor<4x4xf32>
+func @store_after_load_tensor(%arg0 : tensor<4x4xf32>) -> tensor<4x4xf32> {
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %cf0 = constant 0.0 : f32
+  %0 = vector.transfer_read %arg0[%c1, %c0], %cf0 :
+    tensor<4x4xf32>, vector<1x4xf32>
+  %w0 = vector.transfer_write %0, %arg0[%c1, %c0] :
+    vector<1x4xf32>, tensor<4x4xf32>
+  return %w0 : tensor<4x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @store_after_load_tensor_negative
+//       CHECK:   vector.transfer_read
+//       CHECK:   vector.transfer_write
+//       CHECK:   return
+func @store_after_load_tensor_negative(%arg0 : tensor<4x4xf32>) -> tensor<4x4xf32> {
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %cf0 = constant 0.0 : f32
+  %0 = vector.transfer_read %arg0[%c1, %c0], %cf0 :
+    tensor<4x4xf32>, vector<1x4xf32>
+  %w0 = vector.transfer_write %0, %arg0[%c0, %c0] :
+    vector<1x4xf32>, tensor<4x4xf32>
+  return %w0 : tensor<4x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @store_to_load_tensor
+//  CHECK-SAME: (%[[ARG:.*]]: tensor<4x4xf32>, %[[V0:.*]]: vector<1x4xf32>, %[[V1:.*]]: vector<1x4xf32>)
+//   CHECK-NOT:   vector.transfer_write
+//   CHECK-NOT:   vector.transfer_read
+//       CHECK:   return %[[V0]] : vector<1x4xf32>
+func @store_to_load_tensor(%arg0 : tensor<4x4xf32>,
+  %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>) -> vector<1x4xf32> {
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c0 = constant 0 : index
+  %cf0 = constant 0.0 : f32
+  %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]} :
+    vector<1x4xf32>, tensor<4x4xf32>
+  %w1 = vector.transfer_write %v1, %w0[%c2, %c0] {in_bounds = [true, true]} :
+    vector<1x4xf32>, tensor<4x4xf32>
+  %0 = vector.transfer_read %w1[%c1, %c0], %cf0 {in_bounds = [true, true]} :
+    tensor<4x4xf32>, vector<1x4xf32>
+  return %0 : vector<1x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @store_to_load_negative_tensor
+//       CHECK:   vector.transfer_write
+//       CHECK:   vector.transfer_write
+//       CHECK:   %[[V:.*]] = vector.transfer_read
+//       CHECK:   return %[[V]] : vector<1x4xf32>
+func @store_to_load_negative_tensor(%arg0 : tensor<4x4xf32>,
+  %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>, %i : index) -> vector<1x4xf32> {
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c0 = constant 0 : index
+  %cf0 = constant 0.0 : f32
+  %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]} :
+    vector<1x4xf32>, tensor<4x4xf32>
+  %w1 = vector.transfer_write %v0, %w0[%i, %i] {in_bounds = [true, true]} :
+    vector<1x4xf32>, tensor<4x4xf32>
+  %0 = vector.transfer_read %w1[%c1, %c0], %cf0 {in_bounds = [true, true]} :
+    tensor<4x4xf32>, vector<1x4xf32>
+  return %0 : vector<1x4xf32>
+}
+
+// -----
+
+
+// CHECK-LABEL: func @dead_store_tensor
+//   CHECK-DAG:      %[[C0:.*]] = constant 0 : index
+//   CHECK-DAG:      %[[C1:.*]] = constant 1 : index
+//   CHECK-DAG:      %[[C2:.*]] = constant 2 : index
+//   CHECK-NOT:   vector.transfer_write {{.*}}, {{.*}}[%[[C1]], %[[C0]]
+//       CHECK:   vector.transfer_write {{.*}}, {{.*}}[%[[C2]], %[[C0]]
+//       CHECK:   %[[VTW:.*]] = vector.transfer_write {{.*}}, {{.*}}[%[[C1]], %[[C0]]
+//       CHECK:   return %[[VTW]] : tensor<4x4xf32>
+func @dead_store_tensor(%arg0 : tensor<4x4xf32>,
+  %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>, %i : index) -> tensor<4x4xf32> {
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c0 = constant 0 : index
+  %cf0 = constant 0.0 : f32
+  %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]} :
+    vector<1x4xf32>, tensor<4x4xf32>
+  %w1 = vector.transfer_write %v0, %w0[%c2, %c0] {in_bounds = [true, true]} :
+    vector<1x4xf32>, tensor<4x4xf32>
+  %w2 = vector.transfer_write %v1, %w1[%c1, %c0] {in_bounds = [true, true]} :
+    vector<1x4xf32>, tensor<4x4xf32>
+  return %w2 : tensor<4x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @dead_store_tensor_negative
+//   CHECK-DAG:      %[[C0:.*]] = constant 0 : index
+//   CHECK-DAG:      %[[C1:.*]] = constant 1 : index
+//       CHECK:   vector.transfer_write
+//       CHECK:   vector.transfer_write
+//       CHECK:   vector.transfer_read
+//       CHECK:   %[[VTW:.*]] = vector.transfer_write {{.*}}, {{.*}}[%[[C1]], %[[C0]]]
+//       CHECK:   return %[[VTW]] : tensor<4x4xf32>
+func @dead_store_tensor_negative(%arg0 : tensor<4x4xf32>,
+  %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>, %i : index) -> tensor<4x4xf32> {
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c0 = constant 0 : index
+  %cf0 = constant 0.0 : f32
+  %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]} :
+    vector<1x4xf32>, tensor<4x4xf32>
+  %w1 = vector.transfer_write %v0, %w0[%c2, %c0] {in_bounds = [true, true]} :
+    vector<1x4xf32>, tensor<4x4xf32>
+  %0 = vector.transfer_read %w1[%i, %i], %cf0 {in_bounds = [true, true]} :
+    tensor<4x4xf32>, vector<1x4xf32>
+  %x = addf %0, %0 : vector<1x4xf32>
+  %w2 = vector.transfer_write %x, %w0[%c1, %c0] {in_bounds = [true, true]} :
+    vector<1x4xf32>, tensor<4x4xf32>
+  return %w2 : tensor<4x4xf32>
+}

mlir/test/Dialect/Vector/vector-transfer-unroll.mlir

 //  CHECK-NEXT:   %[[VTR2:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C0]]], %{{.*}} : tensor<4x4xf32>, vector<2x2xf32>
 //  CHECK-NEXT:   %[[VTR3:.*]] = vector.transfer_read {{.*}}[%[[C2]], %[[C2]]], %{{.*}} : tensor<4x4xf32>, vector<2x2xf32>
 //  CHECK-NEXT:   %[[VTW0:.*]] = vector.transfer_write %[[VTR0]], {{.*}}[%[[C0]], %[[C0]]] {{.*}} : vector<2x2xf32>, tensor<4x4xf32>
 //  CHECK-NEXT:   %[[VTW1:.*]] = vector.transfer_write %[[VTR1]], %[[VTW0]][%[[C0]], %[[C2]]] {{.*}} : vector<2x2xf32>, tensor<4x4xf32>
 //  CHECK-NEXT:   %[[VTW2:.*]] = vector.transfer_write %[[VTR2]], %[[VTW1]][%[[C2]], %[[C0]]] {{.*}} : vector<2x2xf32>, tensor<4x4xf32>
 //  CHECK-NEXT:   %[[VTW3:.*]] = vector.transfer_write %[[VTR3]], %[[VTW2]][%[[C2]], %[[C2]]] {{.*}} : vector<2x2xf32>, tensor<4x4xf32>
 //  CHECK-NEXT:   return %[[VTW3]] : tensor<4x4xf32>
 
-func @transfer_readwrite_unroll_tensor(%arg0 : tensor<4x4xf32>) ->
+func @transfer_readwrite_unroll_tensor(%arg0 : tensor<4x4xf32>, %arg1 : tensor<4x4xf32>) ->
   tensor<4x4xf32> {
   %c0 = constant 0 : index
   %cf0 = constant 0.0 : f32
   %0 = vector.transfer_read %arg0[%c0, %c0], %cf0 : tensor<4x4xf32>, vector<4x4xf32>
-  %r = vector.transfer_write %0, %arg0[%c0, %c0] : vector<4x4xf32>, tensor<4x4xf32>
+  %r = vector.transfer_write %0, %arg1[%c0, %c0] : vector<4x4xf32>, tensor<4x4xf32>
   return %r: tensor<4x4xf32>
 }

mlir/test/Dialect/Vector/vector-transferop-opt.mlir

scf.if %arg1 {
         vector<1x4xf32>, memref<4x4xf32>
     }
     vector.transfer_write %v0, %arg2[%c1, %c0] {in_bounds = [true, true]} :
       vector<1x4xf32>, memref<4x4xf32>
   }
   return
 }
 
-// CHECK-LABEL: func @forward_dead_store_tensor
-//   CHECK-NOT:   vector.transfer_write
-//   CHECK-NOT:   vector.transfer_read
-//       CHECK:   scf.for
-//       CHECK:   }
-//       CHECK:   %[[VTW:.*]] = vector.transfer_write
-//       CHECK:   return %[[VTW]] : tensor<4x4xf32>
-func @forward_dead_store_tensor(%arg0: i1, %arg1 : tensor<4x4xf32>,
-  %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>, %i : index) -> tensor<4x4xf32> {
-  %c1 = constant 1 : index
-  %c4 = constant 4 : index
-  %c0 = constant 0 : index
-  %cf0 = constant 0.0 : f32
-  %w0 = vector.transfer_write %v0, %arg1[%c1, %c0] {in_bounds = [true, true]} :
-    vector<1x4xf32>, tensor<4x4xf32>
-  %0 = vector.transfer_read %w0[%c1, %c0], %cf0 {in_bounds = [true, true]} :
-    tensor<4x4xf32>, vector<1x4xf32>
-  %x = scf.for %i0 = %c0 to %c4 step %c1 iter_args(%acc = %0)
-    -> (vector<1x4xf32>) {
-    %1 = addf %acc, %acc : vector<1x4xf32>
-    scf.yield %1 : vector<1x4xf32>
-  }
-  %w1 = vector.transfer_write %x, %w0[%c1, %c0] {in_bounds = [true, true]} :
-    vector<1x4xf32>, tensor<4x4xf32>
-  return %w1 : tensor<4x4xf32>
-}
-
-// CHECK-LABEL: func @forward_dead_store_negative_tensor
-//       CHECK:   vector.transfer_write
-//       CHECK:   vector.transfer_read
-//       CHECK:   scf.for
-//       CHECK:   }
-//       CHECK:   %[[VTW:.*]] = vector.transfer_write
-//       CHECK:   return %[[VTW]] : tensor<4x4xf32>
-func @forward_dead_store_negative_tensor(%arg0: i1, %arg1 : tensor<4x4xf32>,
-  %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>, %i : index) -> tensor<4x4xf32> {
-  %c1 = constant 1 : index
-  %c4 = constant 4 : index
-  %c0 = constant 0 : index
-  %cf0 = constant 0.0 : f32
-  %w0 = vector.transfer_write %v0, %arg1[%c1, %i] {in_bounds = [true, true]} :
-    vector<1x4xf32>, tensor<4x4xf32>
-  %0 = vector.transfer_read %w0[%c1, %c0], %cf0 {in_bounds = [true, true]} :
-    tensor<4x4xf32>, vector<1x4xf32>
-  %x = scf.for %i0 = %c0 to %c4 step %c1 iter_args(%acc = %0)
-    -> (vector<1x4xf32>) {
-    %1 = addf %acc, %acc : vector<1x4xf32>
-    scf.yield %1 : vector<1x4xf32>
-  }
-  %w1 = vector.transfer_write %x, %w0[%c1, %c0] {in_bounds = [true, true]} :
-    vector<1x4xf32>, tensor<4x4xf32>
-  return %w1 : tensor<4x4xf32>
-}