[mlir][Rewrite] Add match failure/success diagnostics to `PatternApplicator`users/matthias-springer/pattern_applicator_diagnostics

5 files changed, 71 insertions, 0 deletions

diff --git a/mlir/include/mlir/Rewrite/PatternApplicator.h b/mlir/include/mlir/Rewrite/PatternApplicator.h
index f7871f819a27..0fb00154a67b 100644
--- a/mlir/include/mlir/Rewrite/PatternApplicator.h
+++ b/mlir/include/mlir/Rewrite/PatternApplicator.h
@@ -96,6 +96,9 @@ private:
   std::unique_ptr<detail::PDLByteCodeMutableState> mutableByteCodeState;
 };
 
+/// Register command-line options for the pattern applicator.
+void registerPatternApplicatorCLOptions();
+
 } // namespace mlir
 
 #endif // MLIR_REWRITE_PATTERNAPPLICATOR_H
diff --git a/mlir/lib/Rewrite/PatternApplicator.cpp b/mlir/lib/Rewrite/PatternApplicator.cpp
index e1b56fd6efda..1d64dc2df170 100644
--- a/mlir/lib/Rewrite/PatternApplicator.cpp
+++ b/mlir/lib/Rewrite/PatternApplicator.cpp
@@ -13,7 +13,10 @@
 
 #include "mlir/Rewrite/PatternApplicator.h"
 #include "ByteCode.h"
+#include "mlir/IR/Diagnostics.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/DebugLog.h"
+#include "llvm/Support/ManagedStatic.h"
 
 #ifndef NDEBUG
 #include "llvm/ADT/ScopeExit.h"
@@ -24,6 +27,41 @@
 using namespace mlir;
 using namespace mlir::detail;
 
+namespace {
+enum PatternMatchingRemarkMode {
+  /// Do not emit any diagnostic remarks for pattern matching.
+  None,
+  /// Emit a remark for successful pattern matching.
+  MatchSuccess,
+  /// Emit a remark for failed pattern matching.
+  MatchFailure,
+  /// Emit a remark for both successful and failed pattern matching.
+  MatchSuccessAndFailure,
+};
+struct PatternApplicatorOptions {
+  llvm::cl::opt<PatternMatchingRemarkMode> mode{
+      "mlir-emit-pattern-match-diagnostics",
+      llvm::cl::desc("Emit diagnostic remarks for pattern matching"),
+      llvm::cl::init(PatternMatchingRemarkMode::None), // default value
+      llvm::cl::values(
+          clEnumValN(PatternMatchingRemarkMode::None, "none", "no remarks"),
+          clEnumValN(PatternMatchingRemarkMode::MatchSuccess, "match-success",
+                     "pattern match success"),
+          clEnumValN(PatternMatchingRemarkMode::MatchFailure, "match-failure",
+                     "pattern match failure"),
+          clEnumValN(PatternMatchingRemarkMode::MatchSuccessAndFailure,
+                     "match-success-and-failure",
+                     "pattern match success and failure"))};
+};
+} // namespace
+
+static llvm::ManagedStatic<PatternApplicatorOptions> clOptions;
+
+void mlir::registerPatternApplicatorCLOptions() {
+  // Make sure that the options struct has been initialized.
+  *clOptions;
+}
+
 PatternApplicator::PatternApplicator(
     const FrozenRewritePatternSet &frozenPatternList)
     : frozenPatternList(frozenPatternList) {
@@ -220,9 +258,26 @@ LogicalResult PatternApplicator::matchAndRewrite(
             llvm::scope_exit resetListenerCallback(
                 [&] { rewriter.setListener(oldListener); });
 #endif
+            Location loc = op->getLoc();
             result = pattern->matchAndRewrite(op, rewriter);
             LDBG() << " -> matchAndRewrite "
                    << (succeeded(result) ? "successful" : "failed");
+            bool shouldEmitMatchSuccessRemark =
+                clOptions->mode == PatternMatchingRemarkMode::MatchSuccess ||
+                clOptions->mode ==
+                    PatternMatchingRemarkMode::MatchSuccessAndFailure;
+            bool shouldEmitMatchFailureRemark =
+                clOptions->mode == PatternMatchingRemarkMode::MatchFailure ||
+                clOptions->mode ==
+                    PatternMatchingRemarkMode::MatchSuccessAndFailure;
+            if (succeeded(result) && shouldEmitMatchSuccessRemark) {
+              mlir::emitRemark(loc)
+                  << "pattern match success: " << pattern->getDebugName();
+            }
+            if (failed(result) && shouldEmitMatchFailureRemark) {
+              mlir::emitRemark(loc)
+                  << "pattern match failure: " << pattern->getDebugName();
+            }
           }
 
           // Process the result of the pattern application.
diff --git a/mlir/lib/Tools/mlir-opt/CMakeLists.txt b/mlir/lib/Tools/mlir-opt/CMakeLists.txt
index 858c9c1f97f9..8af4750d3149 100644
--- a/mlir/lib/Tools/mlir-opt/CMakeLists.txt
+++ b/mlir/lib/Tools/mlir-opt/CMakeLists.txt
@@ -11,6 +11,7 @@ add_mlir_library(MLIROptLib
   MLIRPass
   MLIRParser
   MLIRPluginsLib
+  MLIRRewrite
   MLIRSupport
   MLIRIRDL
   MLIRRemarkStreamer
diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
index 560ef6effd2f..4d03ae66dbea 100644
--- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
+++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
@@ -28,6 +28,7 @@
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Pass/PassRegistry.h"
 #include "mlir/Remark/RemarkStreamer.h"
+#include "mlir/Rewrite/PatternApplicator.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/Timing.h"
 #include "mlir/Support/ToolUtilities.h"
@@ -688,6 +689,7 @@ std::string mlir::registerCLIOptions(llvm::StringRef toolName,
   registerMLIRContextCLOptions();
   registerPassManagerCLOptions();
   registerDefaultTimingManagerCLOptions();
+  registerPatternApplicatorCLOptions();
   tracing::DebugCounter::registerCLOptions();
 
   // Build the list of dialects as a header for the --help message.
diff --git a/mlir/test/Rewrite/test-pattern-applicator-diagnostics.mlir b/mlir/test/Rewrite/test-pattern-applicator-diagnostics.mlir
new file mode 100644
index 000000000000..771e1b5373d8
--- /dev/null
+++ b/mlir/test/Rewrite/test-pattern-applicator-diagnostics.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-opt %s -canonicalize -mlir-emit-pattern-match-diagnostics="match-success" -verify-diagnostics
+
+func.func @tensor_empty_canonicalization() -> tensor<?xf32> {
+  %c4 = arith.constant 4 : index
+  // Do not match "(anonymous namespace)::", as it may render differently
+  // depending on the C++ compiler.
+  // expected-remark-re @+1 {{pattern match success: {{.*}}ReplaceEmptyTensorStaticShapeDims}}
+  %r = tensor.empty(%c4) : tensor<?xf32>
+  return %r : tensor<?xf32>
+}