[mlir][Transforms][NFC] GreedyPatternRewriteDriver: Add worklist class

[springerm]

Summary

Encapsulate all worklist-related functionality in a separate Worklist class. This makes the remaining code more readable and allows for custom worklist implementations (e.g., a randomized worklist for fuzzing pattern application: D142447).

[rriddle]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

58–62

Seems like this could be a trivial llvm::any_of.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 
 #define DEBUG_TYPE "greedy-rewriter"
 
+namespace {
+
 //===----------------------------------------------------------------------===//
 // Debugging Infrastructure
 //===----------------------------------------------------------------------===//
 
-namespace {
 #if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
 /// A helper struct that stores finger prints of ops in order to detect broken
 /// RewritePatterns. A rewrite pattern is broken if it modifies IR without
 /// using the rewriter API or if it returns an inconsistent return value.
 struct DebugFingerPrints : public RewriterBase::ForwardingListener {
   DebugFingerPrints(RewriterBase::Listener *driver)
       : RewriterBase::ForwardingListener(driver) {}
 
   /// Compute finger prints of the given op and its nested ops.
   void computeFingerPrints(Operation *topLevel) {
     this->topLevel = topLevel;
     this->topLevelFingerPrint.emplace(topLevel);
     topLevel->walk([&](Operation *op) { fingerprints.try_emplace(op, op); });
   }
 
   /// Clear all finger prints.
   void clear() {
     topLevel = nullptr;
     topLevelFingerPrint.reset();
     fingerprints.clear();
   }
 
   void notifyRewriteSuccess() {
     // Pattern application success => IR must have changed.
     OperationFingerPrint afterFingerPrint(topLevel);
     if (*topLevelFingerPrint == afterFingerPrint) {
       // Note: Run "mlir-opt -debug" to see which pattern is broken.
       llvm::report_fatal_error(
           "pattern returned success but IR did not change");
     }
     for (const auto &it : fingerprints) {
protected:
   Operation *topLevel = nullptr;
 
   /// Finger print of the top-level operation.
   std::optional<OperationFingerPrint> topLevelFingerPrint;
 };
 #endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
 
 //===----------------------------------------------------------------------===//
+// Worklist
+//===----------------------------------------------------------------------===//
+
+/// A LIFO worklist of operations with efficient removal and set semantics.
+///
+/// This class maintains a vector of operations and a mapping of operations to
+/// positions in the vector, so that operations can be removed efficiently at
+/// random. When an operation is removed, it is replaced with nullptr. Such
+/// nullptr are skipped when pop'ing elements.
+class Worklist {
+public:
+  Worklist();
+
+  /// Clear the worklist.
+  void clear();
+
+  /// Return whether the worklist is empty.
+  bool empty() const;
+
+  /// Push an operation to the end of the worklist, unless the operation is
+  /// already on the worklist.
+  void push(Operation *op);
+
+  /// Pop the an operation from the end of the worklist. Only allowed on
+  /// non-empty worklists.
+  Operation *pop();
+
+  /// Remove an operation from the worklist.
+  void remove(Operation *op);
+
+  /// Reverse the worklist.
+  void reverse();
+
+private:
+  /// The worklist of operations.
+  std::vector<Operation *> list;
+
+  /// A mapping of operations to positions in `list`.
+  DenseMap<Operation *, unsigned> map;
+};
+
+Worklist::Worklist() { list.reserve(64); }
+
+void Worklist::clear() {
+  list.clear();
+  map.clear();
+}
+
+bool Worklist::empty() const {
+  // Skip all nullptr.
+  return !llvm::any_of(list,
+                       [](Operation *op) { return static_cast<bool>(op); });
+}
+
+void Worklist::push(Operation *op) {
+  assert(op && "cannot push nullptr to worklist");
+  // Check to see if the worklist already contains this op.
+  if (map.count(op))
+    return;
+  map[op] = list.size();
+  list.push_back(op);
+}
+
+Operation *Worklist::pop() {
+  assert(!empty() && "cannot pop from empty worklist");
+  // Skip and remove all trailing nullptr.
+  while (!list.back())
+    list.pop_back();
+  Operation *op = list.back();
+  list.pop_back();
+  map.erase(op);
+  // Cleanup: Remove all trailing nullptr.
+  while (!list.empty() && !list.back())
+    list.pop_back();
+  return op;
+}
+
+void Worklist::remove(Operation *op) {
+  assert(op && "cannot remove nullptr from worklist");
+  auto it = map.find(op);
+  if (it != map.end()) {
+    assert(list[it->second] == op && "malformed worklist data structure");
+    list[it->second] = nullptr;
+    map.erase(it);
+  }
+}
+
+void Worklist::reverse() {
+  std::reverse(list.begin(), list.end());
+  for (size_t i = 0, e = list.size(); i != e; ++i)
+    map[list[i]] = i;
+}
+
+//===----------------------------------------------------------------------===//
 // GreedyPatternRewriteDriver
 //===----------------------------------------------------------------------===//
 
 /// This is a worklist-driven driver for the PatternMatcher, which repeatedly
 /// applies the locally optimal patterns.
 ///
 /// This abstract class manages the worklist and contains helper methods for
 /// rewriting ops on the worklist. Derived classes specify how ops are added
protected:
   /// worklist as needed: New users are added enqueued.
   void notifyOperationReplaced(Operation *op, ValueRange replacement) override;
 
   /// Process ops until the worklist is empty or `config.maxNumRewrites` is
   /// reached. Return `true` if any IR was changed.
   bool processWorklist();
 
   /// The worklist for this transformation keeps track of the operations that
-  /// need to be revisited, plus their index in the worklist.  This allows us to
-  /// efficiently remove operations from the worklist when they are erased, even
-  /// if they aren't the root of a pattern.
-  std::vector<Operation *> worklist;
-  DenseMap<Operation *, unsigned> worklistMap;
+  /// need to be (re)visited.
+  Worklist worklist;
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 
   /// Configuration information for how to simplify.
   const GreedyRewriteConfig config;
 
   /// The list of ops we are restricting our rewrites to. These include the
   /// supplied set of ops as well as new ops created while rewriting those ops
   /// depending on `strictMode`. This set is not maintained when
   /// `config.strictMode` is GreedyRewriteStrictness::AnyOp.
   llvm::SmallDenseSet<Operation *, 4> strictModeFilteredOps;
 
 private:
   /// Look over the provided operands for any defining operations that should
   /// be re-added to the worklist. This function should be called when an
   /// operation is modified or removed, as it may trigger further
   /// simplifications.
   void addOperandsToWorklist(ValueRange operands);
 
-  /// Pop the next operation from the worklist.
-  Operation *popFromWorklist();
-
   /// Notify the driver that the given block was created.
   void notifyBlockCreated(Block *block) override;
 
   /// For debugging only: Notify the driver of a pattern match failure.
   LogicalResult
   notifyMatchFailure(Location loc,
                      function_ref<void(Diagnostic &)> reasonCallback) override;
 
-  /// If the specified operation is in the worklist, remove it.
-  void removeFromWorklist(Operation *op);
-
 #ifndef NDEBUG
   /// A logger used to emit information during the application process.
   llvm::ScopedPrinter logger{llvm::dbgs()};
 #endif
 
   /// The low-level pattern applicator.
   PatternApplicator matcher;
 
 #if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
   DebugFingerPrints debugFingerPrints;
 #endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
 };
 } // namespace
 
 GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
     MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
     const GreedyRewriteConfig &config)
     : PatternRewriter(ctx), folder(ctx, this), config(config), matcher(patterns)
 #if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
       // clang-format off
       , debugFingerPrints(this)
 // clang-format on
 #endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
 {
-  worklist.reserve(64);
-
   // Apply a simple cost model based solely on pattern benefit.
   matcher.applyDefaultCostModel();
 
   // Set up listener.
 #if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
   // Send IR notifications to the debug handler. This handler will then forward
   // all notifications to this GreedyPatternRewriteDriver.
   setListener(&debugFingerPrints);
#ifndef NDEBUG
   };
 #endif
 
   bool changed = false;
   int64_t numRewrites = 0;
   while (!worklist.empty() &&
          (numRewrites < config.maxNumRewrites ||
           config.maxNumRewrites == GreedyRewriteConfig::kNoLimit)) {
-    auto *op = popFromWorklist();
-
-    // Nulls get added to the worklist when operations are removed, ignore
-    // them.
-    if (op == nullptr)
-      continue;
+    auto *op = worklist.pop();
 
     LLVM_DEBUG({
       logger.getOStream() << "\n";
       logger.startLine() << logLineComment;
       logger.startLine() << "Processing operation : '" << op->getName() << "'("
                          << op << ") {\n";
       logger.indent();
 
do {
     }
     if (region == nullptr)
       return;
   } while ((op = region->getParentOp()));
 }
 
 void GreedyPatternRewriteDriver::addSingleOpToWorklist(Operation *op) {
   if (config.strictMode == GreedyRewriteStrictness::AnyOp ||
-      strictModeFilteredOps.contains(op)) {
-    // Check to see if the worklist already contains this op.
-    if (worklistMap.count(op))
-      return;
-
-    worklistMap[op] = worklist.size();
-    worklist.push_back(op);
-  }
-}
-
-Operation *GreedyPatternRewriteDriver::popFromWorklist() {
-  auto *op = worklist.back();
-  worklist.pop_back();
-
-  // This operation is no longer in the worklist, keep worklistMap up to date.
-  if (op)
-    worklistMap.erase(op);
-  return op;
-}
-
-void GreedyPatternRewriteDriver::removeFromWorklist(Operation *op) {
-  auto it = worklistMap.find(op);
-  if (it != worklistMap.end()) {
-    assert(worklist[it->second] == op && "malformed worklist data structure");
-    worklist[it->second] = nullptr;
-    worklistMap.erase(it);
-  }
+      strictModeFilteredOps.contains(op))
+    worklist.push(op);
 }
 
 void GreedyPatternRewriteDriver::notifyBlockCreated(Block *block) {
   if (config.listener)
     config.listener->notifyBlockCreated(block);
 }
 
 void GreedyPatternRewriteDriver::notifyOperationInserted(Operation *op) {
LLVM_DEBUG({
     logger.startLine() << "** Erase   : '" << op->getName() << "'(" << op
                        << ")\n";
   });
   if (config.listener)
     config.listener->notifyOperationRemoved(op);
 
   addOperandsToWorklist(op->getOperands());
   op->walk([this](Operation *operation) {
-    removeFromWorklist(operation);
+    worklist.remove(operation);
     folder.notifyRemoval(operation);
   });
 
   if (config.strictMode != GreedyRewriteStrictness::AnyOp)
     strictModeFilteredOps.erase(op);
 }
 
 void GreedyPatternRewriteDriver::notifyOperationReplaced(
LogicalResult RegionPatternRewriteDriver::simplify() && {
   MLIRContext *ctx = getContext();
   do {
     // Check if the iteration limit was reached.
     if (iteration++ >= config.maxIterations &&
         config.maxIterations != GreedyRewriteConfig::kNoLimit)
       break;
 
     worklist.clear();
-    worklistMap.clear();
 
     if (!config.useTopDownTraversal) {
       // Add operations to the worklist in postorder.
       region.walk([&](Operation *op) {
         if (!insertKnownConstant(op))
           addToWorklist(op);
       });
     } else {
       // Add all nested operations to the worklist in preorder.
       region.walk<WalkOrder::PreOrder>([&](Operation *op) {
         if (!insertKnownConstant(op)) {
           addToWorklist(op);
           return WalkResult::advance();
         }
         return WalkResult::skip();
       });
 
       // Reverse the list so our pop-back loop processes them in-order.
-      std::reverse(worklist.begin(), worklist.end());
-      // Remember the reverse index.
-      for (size_t i = 0, e = worklist.size(); i != e; ++i)
-        worklistMap[worklist[i]] = i;
+      worklist.reverse();
     }
 
     ctx->executeAction<GreedyPatternRewriteIteration>(
         [&] {
           changed = processWorklist();
 
           // After applying patterns, make sure that the CFG of each of the
           // regions is kept up to date.