[mlir][vector] Hoist redundant singleton vector transfer reads

[harsh]

Summary

For singleton transfer reads, we allow hoisting them out
of the enclosing loop if there are no other accesses to
the source memref within the loop.

[ThomasRaoux]

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

127

I think the logic you want is slightly different, you want to check that there are no other use of the memref by ops with side effects (read side-effects are okay).

[harsh]

Updates based on Thomas' comments

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

127

Sounds good. I have updated the patch.

[ThomasRaoux]

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

55

nit: I would rename it to noAliasingUseInLoop

[harsh]

Updated unit test to be more concise

[harsh]

Fix naming of function

[ThomasRaoux]

Looks good!

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

void mlir::linalg::hoistRedundantVectorTransfersOnTensor(func::FuncOp func) {
   IRRewriter rewriter(func->getContext());
   // TODO: walking in some reverse / inside-out order would be more efficient
   // and would capture more cases.
   func.walk([&](scf::ForOp forOp) {
     hoistRedundantSubsetExtractInsert(rewriter, forOp);
   });
 }
 
+static bool noAliasingUseInLoop(vector::TransferReadOp transferRead,
+                                LoopLikeOpInterface loop) {
+  Value source = transferRead.getSource();
+  while (auto subView = source.getDefiningOp<memref::SubViewOp>())
+    source = subView.getSource();
+  llvm::SmallVector<Operation *, 32> users(source.getUsers().begin(),
+                                           source.getUsers().end());
+  llvm::SmallDenseSet<Operation *, 32> processed;
+  while (!users.empty()) {
+    Operation *user = users.pop_back_val();
+    // If the user has already been processed skip.
+    if (!processed.insert(user).second)
+      continue;
+    if (auto subView = dyn_cast<memref::SubViewOp>(user)) {
+      users.append(subView->getUsers().begin(), subView->getUsers().end());
+      continue;
+    }
+    if (isMemoryEffectFree(user) || isa<vector::TransferReadOp>(user))
+      continue;
+    return false;
+  }
+  return true;
+}
+
 void mlir::linalg::hoistRedundantVectorTransfers(func::FuncOp func) {
   bool changed = true;
   while (changed) {
     changed = false;
     // First move loop invariant ops outside of their loop. This needs to be
     // done before as we cannot move ops without interrupting the function walk.
     func.walk(
         [&](LoopLikeOpInterface loopLike) { moveLoopInvariantCode(loopLike); });
func.walk([&](vector::TransferReadOp transferRead) {
         transferWrite = candidateWrite;
       }
 
       // All operands of the TransferRead must be defined outside of the loop.
       for (auto operand : transferRead.getOperands())
         if (!loop.isDefinedOutsideOfLoop(operand))
           return WalkResult::advance();
 
-      // Only hoist transfer_read / transfer_write pairs for now.
-      if (!transferWrite)
+      // Only hoist transfer_read / transfer_write pairs and singleton
+      // transfer_reads for now.
+      if (!transferWrite) {
+        // Make sure there are no other accesses to the memref before
+        // hoisting transfer_read.
+        if (noAliasingUseInLoop(transferRead, loop))
+          loop.moveOutOfLoop(transferRead);
         return WalkResult::advance();
+      }
 
       LLVM_DEBUG(DBGS() << "Candidate: " << *transferWrite.getOperation()
                         << "\n");
 
       // Approximate aliasing by checking that:
       //   1. indices are the same,
       //   2. no other operations in the loop access the same memref except
       //      for transfer_read/transfer_write accessing statically disjoint

mlir/test/Dialect/Linalg/hoisting.mlir

 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["func.func"]} in %arg1
     : (!pdl.operation) -> !pdl.operation
   transform.structured.hoist_redundant_tensor_subsets %0
     : (!pdl.operation) -> !pdl.operation
 }
+
+// -----
+
+// CHECK-LABEL:  func.func @hoist_vector_transfer_read(
+// CHECK-DAG:      %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG:      %[[C128:.+]] = arith.constant 128 : index
+// CHECK-DAG:      %[[C1024:.+]] = arith.constant 1024 : index
+// CHECK-DAG:      %[[CST:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK:          %[[ALLOC:.+]] = memref.alloc() : memref<32x64xf32>
+// CHECK:          %[[ALLOC_0:.+]] = memref.alloc() : memref<32x128xf32>
+// CHECK:          %[[CAST:.+]] = memref.cast %[[ALLOC_0]] : memref<32x128xf32> to memref<32x128xf32, strided<[128, 1],
+// CHECK-SAME:       offset: ?>>
+// CHECK:          %[[D0:.+]] = vector.transfer_read %[[ALLOC]][%[[C0]], %[[C0]]], %[[CST]] {in_bounds = [true, true]} :
+// CHECK-SAME:       memref<32x64xf32>, vector<32x64xf32>
+// CHECK:          scf.for %[[ARG0:.+]] = %[[C0]] to %[[C1024]] step %[[C128]] {
+// CHECK:            %[[D1:.+]] = vector.transfer_read %[[ALLOC_0]][%[[C0]], %[[C0]]], %[[CST]] {in_bounds = [true, true]}
+// CHECK-SAME:         : memref<32x128xf32>, vector<32x128xf32>
+// CHECK:            "some_use"(%[[D0]], %[[D1]], %[[CAST]]) : (vector<32x64xf32>, vector<32x128xf32>, memref<32x128xf32,
+// CHECK-SAME:         strided<[128, 1], offset: ?>>) -> ()
+// CHECK:          }
+// CHECK:          return
+func.func @hoist_vector_transfer_read() {
+  %c0 = arith.constant 0 : index
+  %c128 = arith.constant 128 : index
+  %c1024 = arith.constant 1024 : index
+  %cst_2 = arith.constant 0.000000e+00 : f32
+  %memref0 = memref.alloc() : memref<32x64xf32>
+  %memref2 = memref.alloc() : memref<32x128xf32>
+  %subview2 = memref.subview %memref2[%c0, %c0] [32, 128] [1, 1]: memref<32x128xf32> to memref<32x128xf32, strided<[128, 1], offset: ?>>
+  scf.for %arg0 = %c0 to %c1024 step %c128 {
+    %2 = vector.transfer_read %memref2[%c0, %c0], %cst_2 {in_bounds = [true, true]} : memref<32x128xf32>, vector<32x128xf32>
+    %3 = vector.transfer_read %memref0[%c0, %c0], %cst_2 {in_bounds = [true, true]} : memref<32x64xf32>, vector<32x64xf32>
+    "some_use"(%3, %2, %subview2) : (vector<32x64xf32>, vector<32x128xf32>, memref<32x128xf32, strided<[128, 1], offset: ?>>) -> ()
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["func.func"]} in %arg1
+    : (!pdl.operation) -> !pdl.operation
+  transform.structured.hoist_redundant_vector_transfers %0
+    : (!pdl.operation) -> !pdl.operation
+}