[mlir][Pass] Move PassExecutionAction to Pass.h, NFC.

mlir/include/mlir/Pass/Pass.h

@@ -9,6 +9,7 @@
 #ifndef MLIR_PASS_PASS_H
 #define MLIR_PASS_PASS_H
 
+#include "mlir/IR/Action.h"
 #include "mlir/Pass/AnalysisManager.h"
 #include "mlir/Pass/PassRegistry.h"
 #include "mlir/Support/LogicalResult.h"
@@ -457,6 +458,52 @@ class PassWrapper : public BaseT {
   }
 };
 
+/// This class encapsulates the "action" of executing a single pass. This allows
+/// a user of the Action infrastructure to query information about an action in
+/// (for example) a breakpoint context. You could use it like this:
+///
+///  auto onBreakpoint = [&](const ActionActiveStack *backtrace) {
+///    if (auto passExec = dyn_cast<PassExecutionAction>(anAction))
+///      record(passExec.getPass());
+///    return ExecutionContext::Apply;
+///  };
+///  ExecutionContext exeCtx(onBreakpoint);
+///
+class PassExecutionAction : public tracing::ActionImpl<PassExecutionAction> {
+  using Base = tracing::ActionImpl<PassExecutionAction>;
+
+public:
+  /// Define a TypeID for this PassExecutionAction.
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PassExecutionAction)
+  /// Construct a PassExecutionAction. This is called by the OpToOpPassAdaptor
+  /// when it calls `executeAction`.
+  PassExecutionAction(ArrayRef<IRUnit> irUnits, const Pass &pass);
+
+  /// The tag required by ActionImpl to identify this action.
+  static constexpr StringLiteral tag = "pass-execution";
+
+  /// Print a textual version of this action to `os`.
+  void print(raw_ostream &os) const override;
+
+  /// Get the pass that will be executed by this action. This is not a class of
+  /// passes, or all instances of a pass kind, this is a single pass.
+  const Pass &getPass() const { return pass; }
+
+  /// Get the operation that is the base of this pass. For example, an
+  /// OperationPass<ModuleOp> would return a ModuleOp.
+  Operation *getOp() const;
+
+public:
+  /// Reference to the pass being run. Notice that this will *not* extend the
+  /// lifetime of the pass, and so this class is therefore unsafe to keep past
+  /// the lifetime of the `executeAction` call.
+  const Pass &pass;
+
+  /// The base op for this pass. For an OperationPass<ModuleOp>, we would have a
+  /// ModuleOp here.
+  Operation *op;
+};
+
 } // namespace mlir
 
 #endif // MLIR_PASS_PASS_H

mlir/lib/Pass/Pass.cpp

@@ -36,11 +36,21 @@ using namespace mlir::detail;
 // PassExecutionAction
 //===----------------------------------------------------------------------===//
 
+PassExecutionAction::PassExecutionAction(ArrayRef<IRUnit> irUnits,
+                                         const Pass &pass)
+    : Base(irUnits), pass(pass) {}
+
 void PassExecutionAction::print(raw_ostream &os) const {
   os << llvm::formatv("`{0}` running `{1}` on Operation `{2}`", tag,
                       pass.getName(), getOp()->getName());
 }
 
+Operation *PassExecutionAction::getOp() const {
+  ArrayRef<IRUnit> irUnits = getContextIRUnits();
+  return irUnits.empty() ? nullptr
+                         : llvm::dyn_cast_if_present<Operation *>(irUnits[0]);
+}
+
 //===----------------------------------------------------------------------===//
 // Pass
 //===----------------------------------------------------------------------===//

mlir/lib/Pass/PassDetail.h

@@ -15,26 +15,6 @@
 #include "llvm/Support/FormatVariadic.h"
 
 namespace mlir {
-/// Encapsulate the "action" of executing a single pass, used for the MLIR
-/// tracing infrastructure.
-struct PassExecutionAction : public tracing::ActionImpl<PassExecutionAction> {
-  using Base = tracing::ActionImpl<PassExecutionAction>;
-  PassExecutionAction(ArrayRef<IRUnit> irUnits, const Pass &pass)
-      : Base(irUnits), pass(pass) {}
-  static constexpr StringLiteral tag = "pass-execution";
-  void print(raw_ostream &os) const override;
-  const Pass &getPass() const { return pass; }
-  Operation *getOp() const {
-    ArrayRef<IRUnit> irUnits = getContextIRUnits();
-    return irUnits.empty() ? nullptr
-                           : llvm::dyn_cast_if_present<Operation *>(irUnits[0]);
-  }
-
-public:
-  const Pass &pass;
-  Operation *op;
-};
-
 namespace detail {
 
 //===----------------------------------------------------------------------===//

mlir/unittests/Pass/CMakeLists.txt

@@ -5,5 +5,6 @@ add_mlir_unittest(MLIRPassTests
 )
 target_link_libraries(MLIRPassTests
   PRIVATE
+  MLIRDebug
   MLIRFuncDialect
   MLIRPass)

mlir/unittests/Pass/PassManagerTest.cpp

@@ -7,6 +7,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Pass/PassManager.h"
+#include "mlir/Debug/BreakpointManagers/TagBreakpointManager.h"
+#include "mlir/Debug/ExecutionContext.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -86,6 +88,104 @@ TEST(PassManagerTest, OpSpecificAnalysis) {
   }
 }
 
+/// Simple pass to annotate a func::FuncOp with a single attribute `didProcess`.
+struct AddAttrFunctionPass
+    : public PassWrapper<AddAttrFunctionPass, OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddAttrFunctionPass)
+
+  void runOnOperation() override {
+    func::FuncOp op = getOperation();
+    Builder builder(op->getParentOfType<ModuleOp>());
+    if (op->hasAttr("didProcess"))
+      op->setAttr("didProcessAgain", builder.getUnitAttr());
+
+    // We always want to set this one.
+    op->setAttr("didProcess", builder.getUnitAttr());
+  }
+};
+
+/// Simple pass to annotate a func::FuncOp with a single attribute
+/// `didProcess2`.
+struct AddSecondAttrFunctionPass
+    : public PassWrapper<AddSecondAttrFunctionPass,
+                         OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddSecondAttrFunctionPass)
+
+  void runOnOperation() override {
+    func::FuncOp op = getOperation();
+    Builder builder(op->getParentOfType<ModuleOp>());
+    op->setAttr("didProcess2", builder.getUnitAttr());
+  }
+};
+
+TEST(PassManagerTest, ExecutionAction) {
+  MLIRContext context;
+  context.loadDialect<func::FuncDialect>();
+  Builder builder(&context);
+
+  // Create a module with 2 functions.
+  OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
+  auto f =
+      func::FuncOp::create(builder.getUnknownLoc(), "process_me_once",
+                           builder.getFunctionType(std::nullopt, std::nullopt));
+  f.setPrivate();
+  module->push_back(f);
+
+  // Instantiate our passes.
+  auto pm = PassManager::on<ModuleOp>(&context);
+  auto pass = std::make_unique<AddAttrFunctionPass>();
+  auto *passPtr = pass.get();
+  pm.addNestedPass<func::FuncOp>(std::move(pass));
+  pm.addNestedPass<func::FuncOp>(std::make_unique<AddSecondAttrFunctionPass>());
+  // Duplicate the first pass to ensure that we *only* run the *first* pass, not
+  // all instances of this pass kind. Notice that this pass (and the test as a
+  // whole) are built to ensure that we can run just a single pass out of a
+  // pipeline that may contain duplicates.
+  pm.addNestedPass<func::FuncOp>(std::make_unique<AddAttrFunctionPass>());
+
+  // Use the action manager to only hit the first pass, not the second one.
+  auto onBreakpoint = [&](const tracing::ActionActiveStack *backtrace)
+      -> tracing::ExecutionContext::Control {
+    // Not a PassExecutionAction, apply the action.
+    auto *passExec = dyn_cast<PassExecutionAction>(&backtrace->getAction());
+    if (!passExec)
+      return tracing::ExecutionContext::Next;
+
+    // If this isn't a function, apply the action.
+    if (!isa<func::FuncOp>(passExec->getOp()))
+      return tracing::ExecutionContext::Next;
+
+    // Only apply the first function pass. Not all instances of the first pass,
+    // only the first pass.
+    if (passExec->getPass().getThreadingSiblingOrThis() == passPtr)
+      return tracing::ExecutionContext::Next;
+
+    // Do not apply any other passes in the pass manager.
+    return tracing::ExecutionContext::Skip;
+  };
+
+  // Set up our breakpoint manager.
+  tracing::TagBreakpointManager simpleManager;
+  tracing::ExecutionContext executionCtx(onBreakpoint);
+  executionCtx.addBreakpointManager(&simpleManager);
+  simpleManager.addBreakpoint(PassExecutionAction::tag);
+
+  // Register the execution context in the MLIRContext.
+  context.registerActionHandler(executionCtx);
+
+  // Run the pass manager, expecting our handler to be called.
+  LogicalResult result = pm.run(module.get());
+  EXPECT_TRUE(succeeded(result));
+
+  // Verify that each function got annotated with `didProcess` and *not*
+  // `didProcess2`.
+  for (func::FuncOp func : module->getOps<func::FuncOp>()) {
+    ASSERT_TRUE(func->getDiscardableAttr("didProcess"));
+    ASSERT_FALSE(func->getDiscardableAttr("didProcess2"));
+    ASSERT_FALSE(func->getDiscardableAttr("didProcessAgain"));
+  }
+}
+
 namespace {
 struct InvalidPass : Pass {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InvalidPass)