[Mlir-commits] [mlir] ce954e1 - [mlir][Transforms] GreedyPatternRewriteDriver: Worklist randomizer

[Matthias Springer]

Author: Matthias Springer
Date: 2023-05-31T09:38:34+02:00
New Revision: ce954e1cda5c9b55325903d51285cd742152a0c3

URL: https://github.com/llvm/llvm-project/commit/ce954e1cda5c9b55325903d51285cd742152a0c3
DIFF: https://github.com/llvm/llvm-project/commit/ce954e1cda5c9b55325903d51285cd742152a0c3.diff

LOG: [mlir][Transforms] GreedyPatternRewriteDriver: Worklist randomizer

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

Differential Revision: https://reviews.llvm.org/D142447

Added: 
    

Modified: 
    mlir/include/mlir/Config/mlir-config.h.cmake
    mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Config/mlir-config.h.cmake b/mlir/include/mlir/Config/mlir-config.h.cmake
index 2bcc9bf9f6b09..efa77b2e5ce5d 100644
--- a/mlir/include/mlir/Config/mlir-config.h.cmake
+++ b/mlir/include/mlir/Config/mlir-config.h.cmake
@@ -19,4 +19,11 @@
    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

diff  --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 050f18c8677b7..2a39cccfc580d 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -27,6 +27,10 @@
 #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"
@@ -165,7 +169,7 @@ class Worklist {
   /// Reverse the worklist.
   void reverse();
 
-private:
+protected:
   /// The worklist of operations.
   std::vector<Operation *> list;
 
@@ -225,6 +229,37 @@ void Worklist::reverse() {
     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
 //===----------------------------------------------------------------------===//
@@ -272,7 +307,11 @@ class GreedyPatternRewriteDriver : public PatternRewriter,
 
   /// 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;