[mlir][Transforms] GreedyPatternRewriteDriver: Worklist randomizer

[springerm]

Summary

Instead of always taking the last op from the worklist, take a random one. For testing/debugging purposes only. This feature can be used to ensure that lowering pipelines work correctly regardless of the order in which ops are processed by the GreedyPatternRewriteDriver.

The randomizer can be enabled by setting a numeric MLIR_GREEDY_REWRITE_RANDOMIZER_SEED option.

Note: When enabled, 27 tests are currently failing. Partly because FileCheck tests are looking for exact IR.

Discussion: https://discourse.llvm.org/t/discussion-fuzzing-pattern-application/67911

[kuhar]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

360–361

I'm don't have experience with random number in LLVM, but I wonder if we should use some centralized instance like https://llvm.org/doxygen/classllvm_1_1RandomNumberGenerator.html?

519

nit: I think llvm::contains would work here

[springerm]

rebase

[kuhar]

As a general comment, IMO we should call it 'randomizer'/'randomized' throughout the patch. IMHO fuzzing has a more specific meaning.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

239

Doesn't this modify the global rand state? Could we use something stateful from <random> instead?

249

nit: either for (int i = pos, e = ....size(); i < e; ++i) or we could use llvm::enumerate(list)

[mehdi_amini]

LGTM on the principle, but @kuhar made good comments to address.

Also can you document this flag (and the one you already introduced) in the debugging guide?

[springerm]

address comments

[kuhar]

LGTM. I'd love to see this as a runtime option in the future -- not sure how difficult that is to plumb through though.

[mehdi_amini]

LG!

This is the doc I was referring to earlier: https://mlir.llvm.org/getting_started/Debugging/

mlir/include/mlir/Config/mlir-config.h.cmake

 #define MLIR_CONFIG_H
 
 /* Enable expensive checks to detect invalid pattern API usage. Failed checks
    manifest as fatal errors or invalid memory accesses (e.g., accessing
    deallocated memory) that cause a crash. Running with ASAN is recommended for
    easier debugging. */
 #cmakedefine01 MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
 
+/* If set, greedy pattern application is randomized: ops on the worklist are
+   chosen at random. For testing/debugging purposes only. This feature can be
+   used to ensure that lowering pipelines work correctly regardless of the order
+   in which ops are processed by the GreedyPatternRewriteDriver. This flag is
+   numeric seed that is passed to the random number generator. */
+#cmakedefine MLIR_GREEDY_REWRITE_RANDOMIZER_SEED ${MLIR_GREEDY_REWRITE_RANDOMIZER_SEED}
+
 #endif

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "llvm/Support/raw_ostream.h"
 
+#ifdef MLIR_GREEDY_REWRITE_RANDOMIZER_SEED
+#include <random>
+#endif // MLIR_GREEDY_REWRITE_RANDOMIZER_SEED
+
 using namespace mlir;
 
 #define DEBUG_TYPE "greedy-rewriter"
 
 namespace {
 
 //===----------------------------------------------------------------------===//
 // Debugging Infrastructure
public:
   Operation *pop();
 
   /// Remove an operation from the worklist.
   void remove(Operation *op);
 
   /// Reverse the worklist.
   void reverse();
 
-private:
+protected:
   /// 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::reverse() {
   std::reverse(list.begin(), list.end());
   for (size_t i = 0, e = list.size(); i != e; ++i)
     map[list[i]] = i;
 }
 
+#ifdef MLIR_GREEDY_REWRITE_RANDOMIZER_SEED
+/// A worklist that pops elements at a random position. This worklist is for
+/// testing/debugging purposes only. It can be used to ensure that lowering
+/// pipelines work correctly regardless of the order in which ops are processed
+/// by the GreedyPatternRewriteDriver.
+class RandomizedWorklist : public Worklist {
+public:
+  RandomizedWorklist() : Worklist() {
+    generator.seed(MLIR_GREEDY_REWRITE_RANDOMIZER_SEED);
+  }
+
+  /// Pop a random non-empty op from the worklist.
+  Operation *pop() {
+    Operation *op = nullptr;
+    do {
+      assert(!list.empty() && "cannot pop from empty worklist");
+      int64_t pos = generator() % list.size();
+      op = list[pos];
+      list.erase(list.begin() + pos);
+      for (int64_t i = pos, e = list.size(); i < e; ++i)
+        map[list[i]] = i;
+      map.erase(op);
+    } while (!op);
+    return op;
+  }
+
+private:
+  std::minstd_rand0 generator;
+};
+#endif // MLIR_GREEDY_REWRITE_RANDOMIZER_SEED
+
 //===----------------------------------------------------------------------===//
 // 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
protected:
   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 (re)visited.
+#ifdef MLIR_GREEDY_REWRITE_RANDOMIZER_SEED
+  RandomizedWorklist worklist;
+#else
   Worklist worklist;
+#endif // MLIR_GREEDY_REWRITE_RANDOMIZER_SEED
 
   /// 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
#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
 {
   // Apply a simple cost model based solely on pattern benefit.
   matcher.applyDefaultCostModel();
 
void GreedyPatternRewriteDriver::addToWorklist(Operation *op) {
   } while ((op = region->getParentOp()));
 }
 
 void GreedyPatternRewriteDriver::addSingleOpToWorklist(Operation *op) {
   if (config.strictMode == GreedyRewriteStrictness::AnyOp ||
       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() << "** Insert  : '" << op->getName() << "'(" << op