Add an action for each iteration of the GreedyPatternRewriteDriver

[jpienaar]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

419

int64_t ?

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 //===- GreedyPatternRewriteDriver.cpp - A greedy 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements mlir::applyPatternsAndFoldGreedily.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/IR/Action.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Rewrite/PatternApplicator.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/ScopeExit.h"
RegionPatternRewriteDriver::RegionPatternRewriteDriver(
     const GreedyRewriteConfig &config, Region &region)
     : GreedyPatternRewriteDriver(ctx, patterns, config), region(region) {
   // Populate strict mode ops.
   if (config.strictMode != GreedyRewriteStrictness::AnyOp) {
     region.walk([&](Operation *op) { strictModeFilteredOps.insert(op); });
   }
 }
 
+namespace {
+class GreedyPatternRewriteIteration
+    : public tracing::ActionImpl<GreedyPatternRewriteIteration> {
+public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GreedyPatternRewriteIteration)
+  GreedyPatternRewriteIteration(ArrayRef<IRUnit> units, int64_t iteration)
+      : tracing::ActionImpl<GreedyPatternRewriteIteration>(units),
+        iteration(iteration) {}
+  static constexpr StringLiteral tag = "GreedyPatternRewriteIteration";
+  void print(raw_ostream &os) const override {
+    os << "GreedyPatternRewriteIteration(" << iteration << ")";
+  }
+
+private:
+  int64_t iteration = 0;
+};
+} // namespace
+
 LogicalResult RegionPatternRewriteDriver::simplify() && {
   auto insertKnownConstant = [&](Operation *op) {
     // Check for existing constants when populating the worklist. This avoids
     // accidentally reversing the constant order during processing.
     Attribute constValue;
     if (matchPattern(op, m_Constant(&constValue)))
       if (!folder.insertKnownConstant(op, constValue))
         return true;
     return false;
   };
 
   bool changed = false;
   int64_t iteration = 0;
+  MLIRContext *ctx = getContext();
   do {
     // Check if the iteration limit was reached.
     if (iteration++ >= config.maxIterations &&
         config.maxIterations != GreedyRewriteConfig::kNoLimit)
       break;
 
     worklist.clear();
     worklistMap.clear();
if (!config.useTopDownTraversal) {
 
       // Reverse the list so our pop-back loop processes them in-order.
       std::reverse(worklist.begin(), worklist.end());
       // Remember the reverse index.
       for (size_t i = 0, e = worklist.size(); i != e; ++i)
         worklistMap[worklist[i]] = i;
     }
 
+    ctx->executeAction<GreedyPatternRewriteIteration>(
+        [&] {
           changed = processWorklist();
 
-    // After applying patterns, make sure that the CFG of each of the regions
-    // is kept up to date.
+          // After applying patterns, make sure that the CFG of each of the
+          // regions is kept up to date.
           if (config.enableRegionSimplification)
             changed |= succeeded(simplifyRegions(*this, region));
+        },
+        {&region}, iteration);
   } while (changed);
 
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return success(!changed);
 }
 
 LogicalResult
 mlir::applyPatternsAndFoldGreedily(Region &region,

mlir/test/Pass/action-logging-filter.mlir

 
 func.func @c() {
     return
 }
 
 ////////////////////////////////////
 /// 1. All actions should be logged.
 
-// RUN: mlir-opt %s --log-actions-to=- -pass-pipeline="builtin.module(func.func(canonicalize))" -o %t --mlir-disable-threading | FileCheck %s
+// RUN: mlir-opt %s --log-actions-to=- -pass-pipeline="builtin.module(func.func(test-stats-pass))" -o %t --mlir-disable-threading | FileCheck %s
 // Specify the current file as filter, expect to see all actions.
-// RUN: mlir-opt %s --log-mlir-actions-filter=%s --log-actions-to=- -pass-pipeline="builtin.module(func.func(canonicalize))" -o %t --mlir-disable-threading | FileCheck %s
+// RUN: mlir-opt %s --log-mlir-actions-filter=%s --log-actions-to=- -pass-pipeline="builtin.module(func.func(test-stats-pass))" -o %t --mlir-disable-threading | FileCheck %s
 
-// CHECK: [thread {{.*}}] begins (no breakpoint) Action `pass-execution`  running `Canonicalizer` on Operation `func.func` (func.func @a() {...}
+// CHECK: [thread {{.*}}] begins (no breakpoint) Action `pass-execution`  running `{{.*}}TestStatisticPass` on Operation `func.func` (func.func @a() {...}
 // CHECK-NEXT: [thread {{.*}}] completed `pass-execution`
-// CHECK-NEXT: [thread {{.*}}] begins (no breakpoint) Action `pass-execution`  running `Canonicalizer` on Operation `func.func` (func.func @b() {...}
+// CHECK-NEXT: [thread {{.*}}] begins (no breakpoint) Action `pass-execution`  running `{{.*}}TestStatisticPass` on Operation `func.func` (func.func @b() {...}
 // CHECK-NEXT: [thread {{.*}}] completed `pass-execution`
-// CHECK-NEXT: [thread {{.*}}] begins (no breakpoint) Action `pass-execution`  running `Canonicalizer` on Operation `func.func` (func.func @c() {...}
+// CHECK-NEXT: [thread {{.*}}] begins (no breakpoint) Action `pass-execution`  running `{{.*}}TestStatisticPass` on Operation `func.func` (func.func @c() {...}
 // CHECK-NEXT: [thread {{.*}}] completed `pass-execution`
 
 ////////////////////////////////////
 /// 2. No match
 
 // Specify a non-existing file as filter, expect to see no actions.
-// RUN: mlir-opt %s --log-mlir-actions-filter=foo.mlir --log-actions-to=- -pass-pipeline="builtin.module(func.func(canonicalize))" -o %t --mlir-disable-threading | FileCheck %s --check-prefix=CHECK-NONE --allow-empty
+// RUN: mlir-opt %s --log-mlir-actions-filter=foo.mlir --log-actions-to=- -pass-pipeline="builtin.module(func.func(test-stats-pass))" -o %t --mlir-disable-threading | FileCheck %s --check-prefix=CHECK-NONE --allow-empty
 // Filter on a non-matching line, expect to see no actions.
-// RUN: mlir-opt %s --log-mlir-actions-filter=%s:1 --log-actions-to=- -pass-pipeline="builtin.module(func.func(canonicalize))" -o %t --mlir-disable-threading | FileCheck %s --check-prefix=CHECK-NONE --allow-empty
+// RUN: mlir-opt %s --log-mlir-actions-filter=%s:1 --log-actions-to=- -pass-pipeline="builtin.module(func.func(test-stats-pass))" -o %t --mlir-disable-threading | FileCheck %s --check-prefix=CHECK-NONE --allow-empty
 
 // Invalid Filter
 // CHECK-NONE-NOT: Canonicalizer
 
 ////////////////////////////////////
 /// 3. Matching filters
 
 // Filter the second function only
-// RUN: mlir-opt %s --log-mlir-actions-filter=%s:8 --log-actions-to=- -pass-pipeline="builtin.module(func.func(canonicalize))" -o %t --mlir-disable-threading | FileCheck %s --check-prefix=CHECK-SECOND
+// RUN: mlir-opt %s --log-mlir-actions-filter=%s:8 --log-actions-to=- -pass-pipeline="builtin.module(func.func(test-stats-pass))" -o %t --mlir-disable-threading | FileCheck %s --check-prefix=CHECK-SECOND
 
 // CHECK-SECOND-NOT: @a
 // CHECK-SECOND-NOT: @c
-// CHECK-SECOND: [thread {{.*}}] begins (no breakpoint) Action `pass-execution`  running `Canonicalizer` on Operation `func.func` (func.func @b() {...}
+// CHECK-SECOND: [thread {{.*}}] begins (no breakpoint) Action `pass-execution`  running `{{.*}}TestStatisticPass` on Operation `func.func` (func.func @b() {...}
 // CHECK-SECOND-NEXT: [thread {{.*}}] completed `pass-execution`
 
 // Filter the first and third functions
-// RUN: mlir-opt %s --log-mlir-actions-filter=%s:4,%s:12 --log-actions-to=- -pass-pipeline="builtin.module(func.func(canonicalize))" -o %t --mlir-disable-threading | FileCheck %s  --check-prefix=CHECK-FIRST-THIRD
+// RUN: mlir-opt %s --log-mlir-actions-filter=%s:4,%s:12 --log-actions-to=- -pass-pipeline="builtin.module(func.func(test-stats-pass))" -o %t --mlir-disable-threading | FileCheck %s  --check-prefix=CHECK-FIRST-THIRD
 
 // CHECK-FIRST-THIRD-NOT: Canonicalizer
-// CHECK-FIRST-THIRD: [thread {{.*}}] begins (no breakpoint) Action `pass-execution`  running `Canonicalizer` on Operation `func.func` (func.func @a() {...}
+// CHECK-FIRST-THIRD: [thread {{.*}}] begins (no breakpoint) Action `pass-execution`  running `{{.*}}TestStatisticPass` on Operation `func.func` (func.func @a() {...}
 // CHECK-FIRST-THIRD-NEXT: [thread {{.*}}] completed `pass-execution`
-// CHECK-FIRST-THIRD-NEXT: [thread {{.*}}] begins (no breakpoint) Action `pass-execution`  running `Canonicalizer` on Operation `func.func` (func.func @c() {...}
+// CHECK-FIRST-THIRD-NEXT: [thread {{.*}}] begins (no breakpoint) Action `pass-execution`  running `{{.*}}TestStatisticPass` on Operation `func.func` (func.func @c() {...}
 // CHECK-FIRST-THIRD-NEXT: [thread {{.*}}] completed `pass-execution`
 // CHECK-FIRST-THIRD-NOT: Canonicalizer

mlir/test/Pass/action-logging.mlir

-// RUN: mlir-opt %s --log-actions-to=- -canonicalize -test-module-pass | FileCheck %s
+// RUN: mlir-opt %s --log-actions-to=- -test-stats-pass -test-module-pass | FileCheck %s
 
-// CHECK: [thread {{.*}}] begins (no breakpoint) Action `pass-execution` running `Canonicalizer` on Operation `builtin.module` (module {...}
+// CHECK: [thread {{.*}}] begins (no breakpoint) Action `pass-execution` running `{{.*}}TestStatisticPass` on Operation `builtin.module` (module {...})`
 // CHECK-NEXT: [thread {{.*}}] completed `pass-execution`
 // CHECK-NEXT: [thread {{.*}}] begins (no breakpoint) Action `pass-execution`  running `{{.*}}TestModulePass` on Operation `builtin.module` (module {...}
 // CHECK-NEXT: [thread {{.*}}] completed `pass-execution`
 // CHECK-NOT: Action