[mlir] Fix potential crash in hoistRedundantVectorTransfers

[ThomasRaoux]

Summary

Operations cannot be moved in the middle of function walk.

[mravishankar]

Looks fine to me, but I dont have much context on this. Stamping anyway.

[nicolasvasilache]

Seems like you'd rather want to record the ops during the walk and then process them so that you don't have the issue?

The current change seems intrusive as the reordering may have non-trivial implications?

Putting a blocker for now.

[ThomasRaoux]

In D107856#2998315, @nicolasvasilache wrote:

Seems like you'd rather want to record the ops during the walk and then process them so that you don't have the issue?

The current change seems intrusive as the reordering may have non-trivial implications?

Putting a blocker for now.

I can record the ops by breaking down moveLoopInvariantCode so that it just returns the ops to hoist without doing it but I don't fully understand how this is better?

Is your concern that the hoisting of transfer op may expose more loop invariant code which could create more opportunity? If that the case we can just move the code doing the loop invariant code motion within the while loop and this should handle all the cases.

[nicolasvasilache]

In D107856#3002773, @ThomasRaoux wrote:

In D107856#2998315, @nicolasvasilache wrote:

Seems like you'd rather want to record the ops during the walk and then process them so that you don't have the issue?

The current change seems intrusive as the reordering may have non-trivial implications?

Putting a blocker for now.

I can record the ops by breaking down moveLoopInvariantCode so that it just returns the ops to hoist without doing it but I don't fully understand how this is better?

Is your concern that the hoisting of transfer op may expose more loop invariant code which could create more opportunity? If that the case we can just move the code doing the loop invariant code motion within the while loop and this should handle all the cases.

Yes, the latter is my concern, moving the code a few lines below under the while but outside of the walk seems good to me.
I don't have data that this will be better, just an intuition.

[ThomasRaoux]

Address review comment

[ThomasRaoux]

In D107856#3003247, @nicolasvasilache wrote:

In D107856#3002773, @ThomasRaoux wrote:

In D107856#2998315, @nicolasvasilache wrote:

Seems like you'd rather want to record the ops during the walk and then process them so that you don't have the issue?

The current change seems intrusive as the reordering may have non-trivial implications?

Putting a blocker for now.

I can record the ops by breaking down moveLoopInvariantCode so that it just returns the ops to hoist without doing it but I don't fully understand how this is better?

Is your concern that the hoisting of transfer op may expose more loop invariant code which could create more opportunity? If that the case we can just move the code doing the loop invariant code motion within the while loop and this should handle all the cases.

Yes, the latter is my concern, moving the code a few lines below under the while but outside of the walk seems good to me.
I don't have data that this will be better, just an intuition.

Sounds good, I moved the code into the while loop.

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

if (changed) {
       RewritePatternSet patterns(func->getContext());
       scf::ForOp::getCanonicalizationPatterns(patterns, func->getContext());
       (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
     }
   }
 }
 
 void mlir::linalg::hoistRedundantVectorTransfers(FuncOp func) {
+  // First move loop invariant ops outside of their loop. This needs to be done
+  // before as we cannot move ops without interputing the function walk.
+  func.walk([&](LoopLikeOpInterface loopLike) {
+    if (failed(moveLoopInvariantCode(loopLike)))
+      llvm_unreachable("Unexpected failure to move invariant code out of loop");
+  });
+
   bool changed = true;
   while (changed) {
     changed = false;
 
     func.walk([&](vector::TransferReadOp transferRead) {
       if (!transferRead.getShapedType().isa<MemRefType>())
         return WalkResult::advance();
 
       LLVM_DEBUG(DBGS() << "Candidate for hoisting: "
                         << *transferRead.getOperation() << "\n");
       auto loop = dyn_cast<scf::ForOp>(transferRead->getParentOp());
       LLVM_DEBUG(DBGS() << "Parent op: " << *transferRead->getParentOp()
                         << "\n");
       if (!loop)
         return WalkResult::advance();
 
-      if (failed(moveLoopInvariantCode(
-              cast<LoopLikeOpInterface>(loop.getOperation()))))
-        llvm_unreachable(
-            "Unexpected failure to move invariant code out of loop");
-
       LLVM_DEBUG(DBGS() << "Candidate read: " << *transferRead.getOperation()
                         << "\n");
 
       SetVector<Operation *> forwardSlice;
       getForwardSlice(transferRead.getOperation(), &forwardSlice);
 
       // Look for the last TransferWriteOp in the forwardSlice of
       // `transferRead` that operates on the same memref.