[Mlir-commits] [mlir] [mlir] Add fast walk-based pattern rewrite driver (PR #113825)

[Jakub Kuderski]

================
@@ -0,0 +1,86 @@
+//===- WalkPatternRewriteDriver.cpp - A fast walk-based rewriter ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Implements mlir::walkAndApplyPatterns.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/WalkPatternRewriteDriver.h"
+
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Verifier.h"
+#include "mlir/IR/Visitors.h"
+#include "mlir/Rewrite/PatternApplicator.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/ErrorHandling.h"
+
+#define DEBUG_TYPE "walk-rewriter"
+
+namespace mlir {
+
+namespace {
+// Forwarding listener to guard against unsupported erasures. Because we use
+// walk-based pattern application, erasing the op from the *next* iteration
+// (e.g., a user of the visited op) is not valid.
+struct ErasedOpsListener final : RewriterBase::ForwardingListener {
+  using RewriterBase::ForwardingListener::ForwardingListener;
+
+  void notifyOperationErased(Operation *op) override {
+    if (op != visitedOp)
+      llvm::report_fatal_error("unsupported op erased in WalkPatternRewriter; "
+                               "erasure is only supported for matched ops");
+
+    ForwardingListener::notifyOperationErased(op);
+  }
+
+  Operation *visitedOp = nullptr;
+};
+} // namespace
+
+void walkAndApplyPatterns(Operation *op,
+                          const FrozenRewritePatternSet &patterns,
+                          RewriterBase::Listener *listener) {
+#if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+  if (failed(verify(op)))
+    llvm::report_fatal_error("walk pattern rewriter input IR failed to verify");
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+
+  PatternRewriter rewriter(op->getContext());
+  ErasedOpsListener erasedListener(listener);
+#if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+  rewriter.setListener(&erasedListener);
+#else
+  (void)erasedListener;
+  rewriter.setListener(listener);
+#endif
+
+  PatternApplicator applicator(patterns);
+  applicator.applyDefaultCostModel();
+
+  op->walk([&](Operation *visitedOp) {
----------------
kuhar wrote:

> Would it make sense to let users specify a traversal order object?

I had that in an earlier implementation but decided to drop it to keep things simple. We can add it later on if there's a concrete use case.

https://github.com/llvm/llvm-project/pull/113825