Add tracing for pattern application in a ApplyPatternAction

[mehdi_amini]

rebase

[Mogball]

LGTM but River mentioned there were some perf issues. Do have a quick capture of the before vs. after time to run a pattern-heavy workload?

[mehdi_amini]

In D144816#4172601, @Mogball wrote:

LGTM but River mentioned there were some perf issues. Do have a quick capture of the before vs. after time to run a pattern-heavy workload?

I didn't benchmark but dumped the assembly and exchanged with River, he was fine with the current state. We'll pay the price of the capture only (here setting up two pointers on the stack).

[mehdi_amini]

Update with one extra no-inline indirection

[rriddle]

mlir/lib/Rewrite/PatternApplicator.cpp

188–192

Seeing this pattern multiple times almost makes me think executeAction should have a return value, but that's a weak thought.

[mehdi_amini]

mlir/lib/Rewrite/PatternApplicator.cpp

188–192

Do you have some C++ dark-magic ideas about this? It's not clear to me to see how to build some return value with unknown type ahead of time there (unless you mean we'd just say we'll forward a bool and that's it.

[mehdi_amini]

mlir/lib/Rewrite/PatternApplicator.cpp

188–192

Still interested in this if you have thoughts!

mlir/include/mlir/Rewrite/PatternApplicator.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_REWRITE_PATTERNAPPLICATOR_H
 #define MLIR_REWRITE_PATTERNAPPLICATOR_H
 
 #include "mlir/Rewrite/FrozenRewritePatternSet.h"
 
+#include "mlir/IR/Action.h"
+
 namespace mlir {
 class PatternRewriter;
 
 namespace detail {
 class PDLByteCodeMutableState;
 } // namespace detail
 
+/// This is the type of Action that is dispatched when a pattern is applied.
+/// It captures the pattern to apply on top of the usual context.
+class ApplyPatternAction : public tracing::ActionImpl<ApplyPatternAction> {
+public:
+  using Base = tracing::ActionImpl<ApplyPatternAction>;
+  ApplyPatternAction(ArrayRef<IRUnit> irUnits, const Pattern &pattern)
+      : Base(irUnits), pattern(pattern) {}
+  static constexpr StringLiteral tag = "apply-pattern-action";
+  static constexpr StringLiteral desc =
+      "Encapsulate the application of rewrite patterns";
+
+  void print(raw_ostream &os) const override {
+    os << "`" << tag << "`\n"
+       << " pattern: " << pattern.getDebugName() << '\n';
+  }
+
+private:
+  const Pattern &pattern;
+};
+
 /// This class manages the application of a group of rewrite patterns, with a
 /// user-provided cost model.
 class PatternApplicator {
 public:
   /// The cost model dynamically assigns a PatternBenefit to a particular
   /// pattern. Users can query contained patterns and pass analysis results to
   /// applyCostModel. Patterns to be discarded should have a benefit of
   /// `impossibleToMatch`.

mlir/lib/Rewrite/PatternApplicator.cpp

do {
 
     // Check that the pattern can be applied.
     if (canApply && !canApply(*bestPattern))
       continue;
 
     // Try to match and rewrite this pattern. The patterns are sorted by
     // benefit, so if we match we can immediately rewrite. For PDL patterns, the
     // match has already been performed, we just need to rewrite.
+    bool matched = false;
+    op->getContext()->executeAction<ApplyPatternAction>(
+        [&]() {
           rewriter.setInsertionPoint(op);
 #ifndef NDEBUG
-    // Operation `op` may be invalidated after applying the rewrite pattern.
+          // Operation `op` may be invalidated after applying the rewrite
+          // pattern.
           Operation *dumpRootOp = getDumpRootOp(op);
 #endif
           if (pdlMatch) {
-      result = bytecode->rewrite(rewriter, *pdlMatch, *mutableByteCodeState);
+            result =
+                bytecode->rewrite(rewriter, *pdlMatch, *mutableByteCodeState);
           } else {
             LLVM_DEBUG(llvm::dbgs() << "Trying to match \""
                                     << bestPattern->getDebugName() << "\"\n");
 
-      const auto *pattern = static_cast<const RewritePattern *>(bestPattern);
+            const auto *pattern =
+                static_cast<const RewritePattern *>(bestPattern);
             result = pattern->matchAndRewrite(op, rewriter);
 
-      LLVM_DEBUG(llvm::dbgs() << "\"" << bestPattern->getDebugName()
-                              << "\" result " << succeeded(result) << "\n");
+            LLVM_DEBUG(llvm::dbgs()
+                       << "\"" << bestPattern->getDebugName() << "\" result "
+                       << succeeded(result) << "\n");
           }
 
           // Process the result of the pattern application.
           if (succeeded(result) && onSuccess && failed(onSuccess(*bestPattern)))
             result = failure();
           if (succeeded(result)) {
             LLVM_DEBUG(logSucessfulPatternApplication(dumpRootOp));
-      break;
+            matched = true;
+            return;
           }
 
           // Perform any necessary cleanups.
           if (onFailure)
             onFailure(*bestPattern);
+        },
+        {op}, *bestPattern);
+    if (matched)
+      break;
   } while (true);
 
   if (mutableByteCodeState)
     mutableByteCodeState->cleanupAfterMatchAndRewrite();
   return result;
 }