Add capture of "IRUnits" as context for an MLIR Action

[mehdi_amini]

Summary

IRUnit is defined as:

using IRUnit = PointerUnion<Operation *, Region *, Block *, Value>;

The tracing::Action is extended to take an ArrayRef<IRUnit> as context to
describe an Action. It is demonstrated in the "ActionLogging" observer.

[mehdi_amini]

rebase

[Tyker]

The data-flow analysis framework has a concept very similar to IRUnit its called ProgramPoint maybe this should be shared ?

[Mogball]

mlir/include/mlir/IR/Unit.h

18

25

32

44

newline

mlir/lib/IR/AsmPrinter.cpp

681

where does the 2 come from?

[rriddle]

mlir/include/mlir/IR/Unit.h

26

It'd be nice not to document this in the context of actions, I'd rather have something be general.

mlir/lib/Debug/Observers/ActionLogging.cpp

10–13

Some of these could be dropped? Like Region.h?

40–41

Can you add a stream operator and simply this?

mlir/lib/IR/AsmPrinter.cpp

228

It'd be nice to split this change out, it feels separable. I think we could also use this here: https://github.com/llvm/llvm-project/blob/73058e330134d275070d49dba80e5fda50ec015b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp#L519

[mehdi_amini]

In D144814#4163506, @Tyker wrote:

The data-flow analysis framework has a concept very similar to IRUnit its called ProgramPoint maybe this should be shared ?

I'm not sure how to unify it, it feels that somehow the ProgramPoint could inherit from IRUnit but it is declared as a more extensible framework: public PointerUnion<GenericProgramPoint *, Operation *, Value, Block *> so I'm not sure how to unify.

[mehdi_amini]

address comments

mlir/include/mlir/IR/Unit.h

26

Yeah that's a remainder from a time where this was defined inside the Action.h header.

mlir/lib/IR/AsmPrinter.cpp

228

What about updating MLIRServer to use this new flag?

[rriddle]

mlir/include/mlir/IR/OperationSupport.h

867

Missing a rebase?

mlir/include/mlir/IR/Unit.h

34

Should be able to drop most of the llvm:: in this file.

mlir/lib/IR/Unit.cpp

41–45

Can you just do std::distance instead (using block->getIterator())?

54–57

Feels cleaner, but feel free to ignore

[mehdi_amini]

mlir/include/mlir/IR/OperationSupport.h

867

Yeah haven't updated this patch since I wrote the other one :)

[mehdi_amini]

Rebase & address comments

[mehdi_amini]

One more comment addressed

[mehdi_amini]

In D144814#4187700, @mehdi_amini wrote:

In D144814#4163506, @Tyker wrote:

The data-flow analysis framework has a concept very similar to IRUnit its called ProgramPoint maybe this should be shared ?

I'm not sure how to unify it, it feels that somehow the ProgramPoint could inherit from IRUnit but it is declared as a more extensible framework: public PointerUnion<GenericProgramPoint *, Operation *, Value, Block *> so I'm not sure how to unify.

If anyone gets an idea, we'll refactor then, in the meantime I'll push this as-is!

mlir/include/mlir/Debug/Observers/ActionLogging.h

 
 namespace mlir {
 namespace tracing {
 
 /// This class defines an observer that print Actions before and after execution
 /// on the provided stream.
 struct ActionLogger : public ExecutionContext::Observer {
   ActionLogger(raw_ostream &os, bool printActions = true,
-               bool printBreakpoints = true)
-      : os(os), printActions(printActions), printBreakpoints(printBreakpoints) {
-  }
+               bool printBreakpoints = true, bool printIRUnits = true)
+      : os(os), printActions(printActions), printBreakpoints(printBreakpoints),
+        printIRUnits(printIRUnits) {}
 
   void beforeExecute(const ActionActiveStack *action, Breakpoint *breakpoint,
                      bool willExecute) override;
   void afterExecute(const ActionActiveStack *action) override;
 
 private:
   raw_ostream &os;
   bool printActions;
   bool printBreakpoints;
+  bool printIRUnits;
 };
 
 } // namespace tracing
 } // namespace mlir
 
 #endif // MLIR_TRACING_OBSERVERS_ACTIONLOGGING_H

mlir/include/mlir/IR/Action.h

 // allows for external entities to control certain actions taken by the compiler
 // by registering handler functions.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_ACTION_H
 #define MLIR_IR_ACTION_H
 
+#include "mlir/IR/Unit.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/TypeID.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/TypeName.h"
 #include "llvm/Support/raw_ostream.h"
 #include <functional>
public:
   /// Return a string "tag" which intends to uniquely identify this type of
   /// action. For example "pass-application" or "pattern-rewrite".
   virtual StringRef getTag() const = 0;
 
   virtual void print(raw_ostream &os) const {
     os << "Action \"" << getTag() << "\"";
   }
 
+  /// Return the set of IR units that are associated with this action.
+  virtual ArrayRef<IRUnit> getContextIRUnits() const { return irUnits; }
+
 protected:
-  Action(TypeID actionID) : actionID(actionID) {}
+  Action(TypeID actionID, ArrayRef<IRUnit> irUnits)
+      : actionID(actionID), irUnits(irUnits) {}
 
   /// The type of the derived action class, used for `isa`/`dyn_cast`.
   TypeID actionID;
+
+  /// Set of IR units (operations, regions, blocks, values) that are associated
+  /// with this action.
+  ArrayRef<IRUnit> irUnits;
 };
 
 /// CRTP Implementation of an action. This class provides a base class for
 /// implementing specific actions.
 ///  Derived classes are expected to provide the following:
 ///   * static constexpr StringLiteral tag = "...";
 ///     - This method returns a unique string identifier, similar to a command
 ///       line flag or DEBUG_TYPE.
 template <typename Derived>
 class ActionImpl : public Action {
 public:
-  ActionImpl() : Action(TypeID::get<Derived>()) {}
+  ActionImpl(ArrayRef<IRUnit> irUnits = {})
+      : Action(TypeID::get<Derived>(), irUnits) {}
 
   /// Provide classof to allow casting between action types.
   static bool classof(const Action *action) {
     return action->getActionID() == TypeID::get<Derived>();
   }
 
   /// Forward tag access to the derived class.
   StringRef getTag() const final { return Derived::tag; }
 };
 
 } // namespace tracing
 } // namespace mlir
 
 #endif // MLIR_IR_ACTION_H

mlir/include/mlir/IR/MLIRContext.h

 //===- MLIRContext.h - MLIR Global Context Class ----------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_MLIRCONTEXT_H
 #define MLIR_IR_MLIRCONTEXT_H
 
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/TypeID.h"
+#include "llvm/ADT/ArrayRef.h"
 #include <functional>
 #include <memory>
 #include <vector>
 
 namespace llvm {
 class ThreadPool;
 } // namespace llvm
 
void executeAction(function_ref<void()> actionFn,
     if (LLVM_UNLIKELY(hasActionHandler()))
       executeActionInternal(actionFn, action);
     else
       actionFn();
   }
 
   /// Dispatch the provided action to the handler if any, or just execute it.
   template <typename ActionTy, typename... Args>
-  void executeAction(function_ref<void()> actionFn, Args &&...args) {
+  void executeAction(function_ref<void()> actionFn, ArrayRef<IRUnit> irUnits,
+                     Args &&...args) {
     if (LLVM_UNLIKELY(hasActionHandler()))
-      executeActionInternal<ActionTy, Args...>(actionFn,
+      executeActionInternal<ActionTy, Args...>(actionFn, irUnits,
                                                std::forward<Args>(args)...);
     else
       actionFn();
   }
 
 private:
   /// Return true if the given dialect is currently loading.
   bool isDialectLoading(StringRef dialectNamespace);
 
   /// Internal helper for the dispatch method.
   void executeActionInternal(function_ref<void()> actionFn,
                              const tracing::Action &action);
 
   /// Internal helper for the dispatch method. We get here after checking that
   /// there is a handler, for the purpose of keeping this code out-of-line. and
   /// avoid calling the ctor for the Action unnecessarily.
   template <typename ActionTy, typename... Args>
   LLVM_ATTRIBUTE_NOINLINE void
-  executeActionInternal(function_ref<void()> actionFn, Args &&...args) {
-    executeActionInternal(actionFn, ActionTy(std::forward<Args>(args)...));
+  executeActionInternal(function_ref<void()> actionFn, ArrayRef<IRUnit> irUnits,
+                        Args &&...args) {
+    executeActionInternal(actionFn,
+                          ActionTy(irUnits, std::forward<Args>(args)...));
   }
 
   const std::unique_ptr<MLIRContextImpl> impl;
 
   MLIRContext(const MLIRContext &) = delete;
   void operator=(const MLIRContext &) = delete;
 };
 

mlir/include/mlir/IR/OperationSupport.h

public:
   /// 'prettyForm' is set to true, debug information is printed in a more
   /// readable 'pretty' form. Note: The IR generated with 'prettyForm' is not
   /// parsable.
   OpPrintingFlags &enableDebugInfo(bool enable = true, bool prettyForm = false);
 
   /// Always print operations in the generic form.
   OpPrintingFlags &printGenericOpForm();
 
+  /// Skip printing regions.
+  OpPrintingFlags &skipRegions();
+
   /// Do not verify the operation when using custom operation printers.
   OpPrintingFlags &assumeVerified();
 
   /// Use local scope when printing the operation. This allows for using the
   /// printer in a more localized and thread-safe setting, but may not
   /// necessarily be identical to what the IR will look like when dumping
   /// the full module.
   OpPrintingFlags &useLocalScope();
public:
   bool shouldPrintDebugInfo() const;
 
   /// Return if debug information should be printed in the pretty form.
   bool shouldPrintDebugInfoPrettyForm() const;
 
   /// Return if operations should be printed in the generic form.
   bool shouldPrintGenericOpForm() const;
 
+  /// Return if regions should be skipped.
+  bool shouldSkipRegions() const;
+
   /// Return if operation verification should be skipped.
   bool shouldAssumeVerified() const;
 
   /// Return if the printer should use local scope when dumping the IR.
   bool shouldUseLocalScope() const;
 
   /// Return if the printer should print users of values.
   bool shouldPrintValueUsers() const;
 
 private:
   /// Elide large elements attributes if the number of elements is larger than
   /// the upper limit.
   std::optional<int64_t> elementsAttrElementLimit;
 
   /// Print debug information.
   bool printDebugInfoFlag : 1;
   bool printDebugInfoPrettyFormFlag : 1;
 
   /// Print operations in the generic form.
   bool printGenericOpFormFlag : 1;
 
+  /// Always skip Regions.
+  bool skipRegionsFlag : 1;
+
   /// Skip operation verification.
   bool assumeVerifiedFlag : 1;
 
   /// Print operations with numberings local to the current operation.
   bool printLocalScope : 1;
 
   /// Print users of values.
   bool printValueUsersFlag : 1;

mlir/include/mlir/IR/Unit.h

+//===- Unit.h -  IR Unit definition--------------------*- C++ -*-=============//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_UNIT_H
+#define MLIR_IR_UNIT_H
+
+#include "mlir/IR/OperationSupport.h"
+#include "llvm/ADT/PointerUnion.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace llvm {
+class raw_ostream;
+} // namespace llvm
+namespace mlir {
+class Operation;
+class Region;
+class Block;
+class Value;
+
+/// IRUnit is a union of the different types of IR objects that consistute the
+/// IR structure (other than Type and Attribute), that is Operation, Region, and
+/// Block.
+class IRUnit
+    : public llvm::PointerUnion<Operation *, Region *, Block *, Value> {
+public:
+  using PointerUnion::PointerUnion;
+
+  /// Print the IRUnit to the given stream.
+  void print(llvm::raw_ostream &os,
+             OpPrintingFlags flags =
+                 OpPrintingFlags().skipRegions().useLocalScope()) const;
+};
+
+llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const IRUnit &unit);
+
+} // end namespace mlir
+
+#endif // MLIR_IR_UNIT_H

mlir/lib/Debug/Observers/ActionLogging.cpp

 //===- ExecutionContext.cpp - Debug Execution Context Support -------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Debug/Observers/ActionLogging.h"
+#include "mlir/IR/Action.h"
 #include "llvm/Support/Threading.h"
-#include <sstream>
-#include <thread>
 
 using namespace mlir;
 using namespace mlir::tracing;
 
 //===----------------------------------------------------------------------===//
 // ActionLogger
 //===----------------------------------------------------------------------===//
 
 void ActionLogger::beforeExecute(const ActionActiveStack *action,
                                  Breakpoint *breakpoint, bool willExecute) {
   os << "[thread " << llvm::get_threadid() << "] ";
   if (willExecute)
     os << "begins ";
   else
     os << "skipping ";
   if (printBreakpoints) {
     if (breakpoint)
-      os << " (on breakpoint: " << *breakpoint << ") ";
+      os << "(on breakpoint: " << *breakpoint << ") ";
     else
-      os << " (no breakpoint) ";
+      os << "(no breakpoint) ";
   }
   os << "Action ";
   if (printActions)
     action->getAction().print(os);
   else
     os << action->getAction().getTag();
+  if (printIRUnits) {
+    os << " (";
+    interleaveComma(action->getAction().getContextIRUnits(), os);
+    os << ")\n";
+  }
 }
 
 void ActionLogger::afterExecute(const ActionActiveStack *action) {
   os << "[thread " << llvm::get_threadid() << "] completed `"
      << action->getAction().getTag() << "`\n";
 }

mlir/lib/IR/AsmPrinter.cpp

 void mlir::registerAsmPrinterCLOptions() {
   // Make sure that the options struct has been initialized.
   *clOptions;
 }
 
 /// Initialize the printing flags with default supplied by the cl::opts above.
 OpPrintingFlags::OpPrintingFlags()
     : printDebugInfoFlag(false), printDebugInfoPrettyFormFlag(false),
-      printGenericOpFormFlag(false), assumeVerifiedFlag(false),
-      printLocalScope(false), printValueUsersFlag(false) {
+      printGenericOpFormFlag(false), skipRegionsFlag(false),
+      assumeVerifiedFlag(false), printLocalScope(false),
+      printValueUsersFlag(false) {
   // Initialize based upon command line options, if they are available.
   if (!clOptions.isConstructed())
     return;
   if (clOptions->elideElementsAttrIfLarger.getNumOccurrences())
     elementsAttrElementLimit = clOptions->elideElementsAttrIfLarger;
   printDebugInfoFlag = clOptions->printDebugInfoOpt;
   printDebugInfoPrettyFormFlag = clOptions->printPrettyDebugInfoOpt;
   printGenericOpFormFlag = clOptions->printGenericOpFormOpt;

 }
 
 /// Always print operations in the generic form.
 OpPrintingFlags &OpPrintingFlags::printGenericOpForm() {
   printGenericOpFormFlag = true;
   return *this;
 }
 
+/// Always skip Regions.
+OpPrintingFlags &OpPrintingFlags::skipRegions() {
+  skipRegionsFlag = true;
+  return *this;
+}
+
 /// Do not verify the operation when using custom operation printers.
 OpPrintingFlags &OpPrintingFlags::assumeVerified() {
   assumeVerifiedFlag = true;
   return *this;
 }
 
 /// Use local scope when printing the operation. This allows for using the
 /// printer in a more localized and thread-safe setting, but may not necessarily
bool OpPrintingFlags::shouldPrintDebugInfoPrettyForm() const {
   return printDebugInfoPrettyFormFlag;
 }
 
 /// Return if operations should be printed in the generic form.
 bool OpPrintingFlags::shouldPrintGenericOpForm() const {
   return printGenericOpFormFlag;
 }
 
+/// Return if Region should be skipped.
+bool OpPrintingFlags::shouldSkipRegions() const { return skipRegionsFlag; }
+
 /// Return if operation verification should be skipped.
 bool OpPrintingFlags::shouldAssumeVerified() const {
   return assumeVerifiedFlag;
 }
 
 /// Return if the printer should use local scope when dumping the IR.
 bool OpPrintingFlags::shouldUseLocalScope() const { return printLocalScope; }
 
void printCustomOrGenericOp(Operation *op) override {
     // Otherwise print with the generic assembly form.
     printGenericOp(op);
   }
 
 private:
   /// Print the given operation in the generic form.
   void printGenericOp(Operation *op, bool printOpName = true) override {
     // Consider nested operations for aliases.
+    if (!printerFlags.shouldSkipRegions()) {
       for (Region &region : op->getRegions())
         printRegion(region, /*printEntryBlockArgs=*/true,
                     /*printBlockTerminators=*/true);
+    }
 
     // Visit all the types used in the operation.
     for (Type type : op->getOperandTypes())
       printType(type);
     for (Type type : op->getResultTypes())
       printType(type);
 
     // Consider the attributes of the operation for aliases.
private:
   }
 
   /// Print the given region.
   void printRegion(Region &region, bool printEntryBlockArgs,
                    bool printBlockTerminators,
                    bool printEmptyBlock = false) override {
     if (region.empty())
       return;
+    if (printerFlags.shouldSkipRegions()) {
+      os << "{/*skip region*/}";
+      return;
+    }
 
     auto *entryBlock = &region.front();
     print(entryBlock, printEntryBlockArgs, printBlockTerminators);
     for (Block &b : llvm::drop_begin(region, 1))
       print(&b);
   }
 
   void printRegionArgument(BlockArgument arg, ArrayRef<NamedAttribute> argAttrs,
if (op->getNumSuccessors() != 0) {
     interleaveComma(op->getSuccessors(),
                     [&](Block *successor) { printBlockName(successor); });
     os << ']';
   }
 
   // Print regions.
   if (op->getNumRegions() != 0) {
     os << " (";
+    if (!printerFlags.shouldSkipRegions()) {
       interleaveComma(op->getRegions(), [&](Region &region) {
         printRegion(region, /*printEntryBlockArgs=*/true,
                     /*printBlockTerminators=*/true, /*printEmptyBlock=*/true);
       });
+    } else {
+      os << "/*skip " << op->getNumRegions() << " regions*/";
+    }
     os << ')';
   }
 
   auto attrs = op->getAttrs();
   printOptionalAttrDict(attrs);
 
   // Print the type signature of the operation.
   os << " : ";
void OperationPrinter::printSuccessorAndUseList(Block *successor,
   interleaveComma(succOperands,
                   [this](Value operand) { printType(operand.getType()); });
   os << ')';
 }
 
 void OperationPrinter::printRegion(Region &region, bool printEntryBlockArgs,
                                    bool printBlockTerminators,
                                    bool printEmptyBlock) {
+  if (printerFlags.shouldSkipRegions()) {
+    os << "{/*skip region*/}";
+    return;
+  }
   os << "{" << newLine;
   if (!region.empty()) {
     auto restoreDefaultDialect =
         llvm::make_scope_exit([&]() { defaultDialectStack.pop_back(); });
     if (auto iface = dyn_cast<OpAsmOpInterface>(region.getParentOp()))
       defaultDialectStack.push_back(iface.getDefaultDialect());
     else
       defaultDialectStack.push_back("");

mlir/lib/IR/CMakeLists.txt

add_mlir_library(MLIRIR
   PatternMatch.cpp
   Region.cpp
   RegionKindInterface.cpp
   SymbolTable.cpp
   TensorEncoding.cpp
   Types.cpp
   TypeRange.cpp
   TypeUtilities.cpp
+  Unit.cpp
   Value.cpp
   ValueRange.cpp
   Verifier.cpp
   Visitors.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
 

mlir/lib/IR/Unit.cpp

+//===- Unit.cpp - Support for manipulating IR Unit ------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/Unit.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/Region.h"
+#include "llvm/Support/raw_ostream.h"
+#include <sstream>
+
+using namespace mlir;
+
+static void printOp(llvm::raw_ostream &os, Operation *op,
+                    OpPrintingFlags &flags) {
+  if (!op) {
+    os << "<Operation:nullptr>";
+    return;
+  }
+  op->print(os, flags);
+}
+
+static void printRegion(llvm::raw_ostream &os, Region *region,
+                        OpPrintingFlags &flags) {
+  if (!region) {
+    os << "<Region:nullptr>";
+    return;
+  }
+  os << "Region #" << region->getRegionNumber() << " for op ";
+  printOp(os, region->getParentOp(), flags);
+}
+
+static void printBlock(llvm::raw_ostream &os, Block *block,
+                       OpPrintingFlags &flags) {
+  int blockId = 0;
+  Region *region = block->getParent();
+  Block *cur = &region->front();
+  while (cur != block) {
+    cur = cur->getNextNode();
+    blockId++;
+  }
+  os << "Block #" << blockId << " for ";
+  bool shouldSkipRegions = flags.shouldSkipRegions();
+  printRegion(os, region, flags.skipRegions());
+  if (!shouldSkipRegions)
+    block->print(os);
+}
+
+void mlir::IRUnit::print(llvm::raw_ostream &os, OpPrintingFlags flags) const {
+  if (auto *op = this->dyn_cast<Operation *>()) {
+    printOp(os, op, flags);
+    return;
+  }
+  if (auto *region = this->dyn_cast<Region *>()) {
+    printRegion(os, region, flags);
+    return;
+  }
+  if (auto *block = this->dyn_cast<Block *>()) {
+    printBlock(os, block, flags);
+    return;
+  }
+  llvm_unreachable("unknown IRUnit");
+}
+
+llvm::raw_ostream &mlir::operator<<(llvm::raw_ostream &os, const IRUnit &unit) {
+  unit.print(os);
+  return os;
+}

mlir/lib/Pass/Pass.cpp

op->getContext()->executeAction<PassExecutionAction>(
       [&]() {
         // Invoke the virtual runOnOperation method.
         if (auto *adaptor = dyn_cast<OpToOpPassAdaptor>(pass))
           adaptor->runOnOperation(verifyPasses);
         else
           pass->runOnOperation();
         passFailed = pass->passState->irAndPassFailed.getInt();
       },
-      *pass, op);
+      {op}, *pass);
 
   // Invalidate any non preserved analyses.
   am.invalidate(pass->passState->preservedAnalyses);
 
   // When verifyPasses is specified, we run the verifier (unless the pass
   // failed).
   if (!passFailed && verifyPasses) {
     bool runVerifierNow = true;

mlir/lib/Pass/PassDetail.h

 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #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> {
-  PassExecutionAction(const Pass &pass, Operation *op) : pass(pass), op(op) {}
+  using Base = tracing::ActionImpl<PassExecutionAction>;
+  PassExecutionAction(ArrayRef<IRUnit> irUnits, const Pass &pass)
+      : Base(irUnits), pass(pass) {}
   static constexpr StringLiteral tag = "pass-execution-action";
   void print(raw_ostream &os) const override {
-    os << llvm::formatv("`{0}` running \"{1}\" on Operation \"{2}\"\n", tag,
-                        pass.getName(), op->getName());
+    os << "`" << tag << "` "
+       << " running \"" << pass.getName() << "\" on Operation \"";
+    ArrayRef<IRUnit> irUnits = getContextIRUnits();
+    if (irUnits.empty()) {
+      os << "<missing?>";
+    } else {
+      os << irUnits.front().dyn_cast<Operation *>()->getName() << "\"";
+    }
   }
   const Pass &getPass() { return pass; }
   Operation *getOp() { return op; }
 
 public:
   const Pass &pass;
   Operation *op;
 };

mlir/test/Pass/action-logging.mlir

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