[mlir][vector] Support transfer op on tensor optimizations

[ThomasRaoux]

Summary

Support store to load forwarding and dead store transformations for transfer op on tensor.

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp

224

How about having a transferEncompasses(A, B) and implementing it for the equality case for now?

232

xxxIndices

237

else loop forever ?

I am not sure I understand the case disjunction ?
Seems like things would break in the partial overlap case?

245

same Q

252

same question here.

[ThomasRaoux]

Fix bug caught in review and add extra test. Address more review feedback.

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp

224

I moved this logic and did the same for the read case.

232

Fixed, this is depend on the hoisting review, I'm updating this patch as well.

237

This is a bug indeed. Thanks for catching this. I fixed it and added a test to handle that case.

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.hasMaskedDim() && 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 (write.indices() == nextWrite.indices() &&
-          write.getVectorType() == nextWrite.getVectorType() &&
-          write.permutation_map() == write.permutation_map() &&
+      if (transferEncompasses(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) && !write.hasMaskedDim() &&
-          write.indices() == read.indices() &&
-          write.getVectorType() == read.getVectorType() &&
-          write.permutation_map() == read.permutation_map()) {
+      if (dominators.dominates(write, read) &&
+          transferEncompasses(write, read)) {
         if (lastwrite == nullptr || dominators.dominates(lastwrite, write))
           lastwrite = write;
         else
           assert(dominators.dominates(write, lastwrite));
         continue;
       }
     }
     blockingWrites.push_back(user);
for (Operation *write : blockingWrites) {
     // TODO: if the store and read have the same ancestor we could recurse in
     // the region to know if the read is reachable with more precision.
     if (writeAncestor == nullptr || !isReachable(writeAncestor, readAncestor))
       continue;
     if (!postDominators.postDominates(lastwrite, write)) {
       LLVM_DEBUG(DBGS() << "Fail to do write to read forwarding due to op: "
                         << *write << "\n");
       return;
     }
   }
 
   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) { opt.storeToLoadForwarding(read); });
+  func.walk([&](vector::TransferReadOp read) {
+    if (read.getShapedType().isa<MemRefType>())
+      opt.storeToLoadForwarding(read);
+    else
+      opt.storeToLoadForwardingTensor(read);
+  });
   opt.removeDeadOp();
-  func.walk([&](vector::TransferWriteOp write) { opt.deadStoreOp(write); });
+  func.walk([&](vector::TransferWriteOp write) {
+    if (write.getShapedType().isa<MemRefType>())
+      opt.deadStoreOp(write);
+    else
+      opt.deadStoreOpTensor(write);
+  });
   opt.removeDeadOp();
 }

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

 // RUN: mlir-opt %s -test-vector-transferop-opt | FileCheck %s
 
 // CHECK-LABEL: func @forward_dead_store
 //   CHECK-NOT:   vector.transfer_write
 //   CHECK-NOT:   vector.transfer_read
 //       CHECK:   scf.for
 //       CHECK:   }
 //       CHECK:   vector.transfer_write
 //       CHECK:   return
 func @forward_dead_store(%arg0: i1, %arg1 : memref<4x4xf32>,
   %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>, %i : index) {
   %c1 = constant 1 : index
   %c4 = constant 4 : index
   %c0 = constant 0 : index
   %cf0 = constant 0.0 : f32
-  vector.transfer_write %v0, %arg1[%c1, %c0] {masked = [false, false]} : 
+  vector.transfer_write %v0, %arg1[%c1, %c0] {masked = [false, false]} :
     vector<1x4xf32>, memref<4x4xf32>
-  %0 = vector.transfer_read %arg1[%c1, %c0], %cf0 {masked = [false, false]} : 
+  %0 = vector.transfer_read %arg1[%c1, %c0], %cf0 {masked = [false, false]} :
     memref<4x4xf32>, vector<1x4xf32>
-  %x = scf.for %i0 = %c0 to %c4 step %c1 iter_args(%acc = %0) 
+  %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>
   }
-  vector.transfer_write %x, %arg1[%c1, %c0] {masked = [false, false]} : 
+  vector.transfer_write %x, %arg1[%c1, %c0] {masked = [false, false]} :
     vector<1x4xf32>, memref<4x4xf32>
   return
 }
 
 // CHECK-LABEL: func @forward_nested
 //       CHECK:   vector.transfer_write
 //       CHECK:   vector.transfer_write
 //       CHECK:   scf.if

 //       CHECK:   scf.if
 //   CHECK-NOT:     vector.transfer_write
 //       CHECK:   }
 //       CHECK:   vector.transfer_write
 //   CHECK-NOT:   vector.transfer_write
 //       CHECK:   vector.transfer_read
 //       CHECK:   return
 func @dead_store_region(%arg0: i1, %arg1 : memref<4x4xf32>,
-  %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>, %i : index) 
+  %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>, %i : index)
   -> (vector<1x4xf32>) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %cf0 = constant 0.0 : f32
   vector.transfer_write %v0, %arg1[%c1, %c0] {masked = [false, false]} :
     vector<1x4xf32>, memref<4x4xf32>
   %x = scf.if %arg0 -> (vector<1x4xf32>) {
     scf.yield %v1 : vector<1x4xf32>
scf.if %arg1 {
         vector<1x4xf32>, memref<4x4xf32>
     }
     vector.transfer_write %v0, %arg2[%c1, %c0] {masked = [false, false]} :
       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] {masked = [false, false]} :
+    vector<1x4xf32>, tensor<4x4xf32>
+  %0 = vector.transfer_read %w0[%c1, %c0], %cf0 {masked = [false, false]} :
+    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] {masked = [false, false]} :
+    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] {masked = [false, false]} :
+    vector<1x4xf32>, tensor<4x4xf32>
+  %0 = vector.transfer_read %w0[%c1, %c0], %cf0 {masked = [false, false]} :
+    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] {masked = [false, false]} :
+    vector<1x4xf32>, tensor<4x4xf32>
+  return %w1 : tensor<4x4xf32>
+}