[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

 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.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 (defWrite.indices() == write.indices() &&
+        defWrite.getVectorType() == write.getVectorType() &&
+        defWrite.permutation_map() == write.permutation_map()) {
+      write.sourceMutable().assign(defWrite.source());
+      if (defWrite->use_empty())
+        opToErase.push_back(defWrite.getOperation());
+      return;
+    }
+    if (isDisjointTransferIndicies(
+            cast<VectorTransferOpInterface>(defWrite.getOperation()),
+            cast<VectorTransferOpInterface>(write.getOperation())))
+      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 (!defWrite.hasMaskedDim() && defWrite.indices() == read.indices() &&
+        defWrite.getVectorType() == read.getVectorType() &&
+        defWrite.permutation_map() == read.permutation_map()) {
+      read.replaceAllUsesWith(defWrite.vector());
+      opToErase.push_back(read.getOperation());
+      return;
+    }
+    if (isDisjointTransferIndicies(
+            cast<VectorTransferOpInterface>(defWrite.getOperation()),
+            cast<VectorTransferOpInterface>(read.getOperation())))
+      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

   %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:   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

   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>
+}