diff --git a/mlir/include/mlir/Config/mlir-config.h.cmake b/mlir/include/mlir/Config/mlir-config.h.cmake --- 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_FUZZER_SEED ${MLIR_GREEDY_REWRITE_FUZZER_SEED} + #endif diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp --- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp +++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp @@ -165,7 +165,7 @@ /// Reverse the worklist. void reverse(); -private: +protected: /// The worklist of operations. std::vector list; @@ -225,6 +225,32 @@ map[list[i]] = i; } +#ifdef MLIR_GREEDY_REWRITE_FUZZER_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 FuzzedWorklist : public Worklist { +public: + FuzzedWorklist() : Worklist() { std::srand(MLIR_GREEDY_REWRITE_FUZZER_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 = std::rand() % list.size(); + op = list[pos]; + list.erase(list.begin() + pos); + for (int64_t i = pos; i < static_cast(list.size()); ++i) + map[list[i]] = i; + map.erase(op); + } while (!op); + return op; + } +}; +#endif // MLIR_GREEDY_REWRITE_FUZZER_SEED + //===----------------------------------------------------------------------===// // GreedyPatternRewriteDriver //===----------------------------------------------------------------------===// @@ -272,7 +298,11 @@ /// The worklist for this transformation keeps track of the operations that /// need to be (re)visited. +#ifdef MLIR_GREEDY_REWRITE_FUZZER_SEED + FuzzedWorklist worklist; +#else Worklist worklist; +#endif // MLIR_GREEDY_REWRITE_FUZZER_SEED /// Non-pattern based folder for operations. OperationFolder folder;