[mlir][IR] Add Iterator template option to IR walkers

[springerm]

Summary

This allows users to specify a top-down or bottom-up traversal of the IR, in addition to the already existing WalkOrder.

Certain transformations work better with a forward traversal. E.g., when cloning a piece of IR, operations should be cloned top-down so that all uses are defined when creating an op.

Certain transformations work better with a reverse traversal. E.g., when erasing a piece of IR, operations should be erased bottom-up to avoid erasing operations that still have users.

[springerm]

rebase

[springerm]

update

[springerm]

mlir/lib/IR/Visitors.cpp

42 ↗

(On Diff #498317)

Note: Explicit instantiations would not be needed if detail::walk would be defined in the header file.

[springerm]

simplify

[springerm]

update

[Mogball]

This change overall makes sense to me

mlir/lib/IR/Visitors.cpp

31 ↗

(On Diff #498420)

Can you implement this with if constexpr instead? Walking is very performance sensitive and it's better not to regress existing code.

[springerm]

address comments

[springerm]

mlir/lib/IR/Visitors.cpp

31 ↗

(On Diff #498420)

I was able to simplify TopDownIterator::makeRange and make it a constexpr function that simplify passes through the argument as a reference. Together with Iterator being a template argument, there should be no difference in performance in case of a top-down traversal.

[rriddle]

mlir/include/mlir/IR/Visitors.h

65–82 ↗

(On Diff #499475)

Why not just name them "ForwardIterator" and "ReverseIterator" instead of TopDown/BottomUp (which aren't correct, given the hierarchal structure of the IR).

[springerm]

address comments

mlir/lib/IR/PatternMatch.cpp

 void RewriterBase::replaceOp(Operation *op, ValueRange newValues) {
   // Notify the rewriter subclass that we're about to replace this root.
   notifyRootReplaced(op, newValues);
 
   assert(op->getNumResults() == newValues.size() &&
          "incorrect # of replacement values");
   op->replaceAllUsesWith(newValues);
 
-  notifyOperationRemoved(op);
+  op->walk<WalkOrder::PostOrder, IteratorOrder::Reverse>(
+      [&](Operation *op) { notifyOperationRemoved(op); });
   op->erase();
 }
 
 /// This method erases an operation that is known to have no uses. The uses of
 /// the given operation *must* be known to be dead.
 void RewriterBase::eraseOp(Operation *op) {
   assert(op->use_empty() && "expected 'op' to have no uses");
-  notifyOperationRemoved(op);
+  op->walk<WalkOrder::PostOrder, IteratorOrder::Reverse>(
+      [&](Operation *op) { notifyOperationRemoved(op); });
   op->erase();
 }
 
 void RewriterBase::eraseBlock(Block *block) {
   for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block))) {
     assert(op.use_empty() && "expected 'op' to have no uses");
     eraseOp(&op);
   }

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

LLVM_DEBUG({
             logger.startLine(),
             OpPrintingFlags().printGenericOpForm().elideLargeElementsAttrs());
         logger.getOStream() << "\n\n";
       }
     });
 
     // If the operation is trivially dead - remove it.
     if (isOpTriviallyDead(op)) {
-      notifyOperationRemoved(op);
-      op->erase();
+      eraseOp(op);
       changed = true;
 
       LLVM_DEBUG(logResultWithLine("success", "operation is trivially dead"));
       continue;
     }
 
     // Collects all the operands and result uses of the given `op` into work
     // list. Also remove `op` and nested ops from worklist.
     originalOperands.assign(op->operand_begin(), op->operand_end());
     auto preReplaceAction = [&](Operation *op) {
       // Add the operands to the worklist for visitation.
       addOperandsToWorklist(originalOperands);
 
       // Add all the users of the result to the worklist so we make sure
       // to revisit them.
       for (auto result : op->getResults())
         for (auto *userOp : result.getUsers())
           addToWorklist(userOp);
 
-      notifyOperationRemoved(op);
+      op->walk<WalkOrder::PostOrder, IteratorOrder::Reverse>(
+          [&](Operation *op) { notifyOperationRemoved(op); });
     };
 
     // Add the given operation to the worklist.
     auto collectOps = [this](Operation *op) { addToWorklist(op); };
 
     // Try to fold this op.
     bool inPlaceUpdate;
     if ((succeeded(folder.tryToFold(op, collectOps, preReplaceAction,

 
 void GreedyPatternRewriteDriver::notifyOperationRemoved(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Erase   : '" << op->getName() << "'(" << op
                        << ")\n";
   });
 
   addOperandsToWorklist(op->getOperands());
-  op->walk([this](Operation *operation) {
-    removeFromWorklist(operation);
-    folder.notifyRemoval(operation);
-  });
+  removeFromWorklist(op);
+  folder.notifyRemoval(op);
 
   if (config.strictMode != GreedyRewriteStrictness::AnyOp)
     strictModeFilteredOps.erase(op);
 }
 
 void GreedyPatternRewriteDriver::notifyRootReplaced(Operation *op,
                                                     ValueRange replacement) {
   LLVM_DEBUG({