[mlir][Pass] Integrate pass timing hooks into Pass

[rriddle]

Summary

This allows for passes to nest additional timing scopes by
using the new createTimingScope method. This will allow
for pass authors to inject more valuable timing information
into their passes when necessary, which is valuable when
debugging expensive passes, or just for tracking important
passes in general. Timing scopes are a no-op when timing
is not enabled.

[mehdi_amini]

That is awesome! I was thinking about rebasing the Timing on top of the Action framework this after I publish it (soonish). Can we table this for a couple of weeks so we can check how we converge this nicely?

[mehdi_amini]

My current thought is that this is a great direction to be able to use timing for more than passes, the ability to delimitate arbitrary sections of code to trace is a common need (showcased by people using Tracy for example).

Something that seems not orthogonal right now in the way timing is setup is that this is very much anchored to the Pass infrastructure, but other than historical reason I am not sure why?
Seems like this could be decoupled from the pass infra, and attached to the Context instead for example. That would give the possibility to add timers anywhere there is an IR construct: if I'm in a utility method and I want to instrument it temporarily I likely don't want to thread a callback to createTimingScope through multiple layers of API so that I can ultimately pass it in from all the passes that transitively use this utility.
Have you already thought about that?

mlir/docs/PassManagement.md

 Instrumentations added to the PassManager are run in a stack like fashion, i.e.
 the last instrumentation to execute a `runBefore*` hook will be the first to
 execute the respective `runAfter*` hook. The hooks of a `PassInstrumentation`
 class are guaranteed to be executed in a thread-safe fashion, so additional
 synchronization is not necessary. Below in an example instrumentation that
 counts the number of times the `DominanceInfo` analysis is computed:
 
 ```c++
-struct DominanceCounterInstrumentation : public PassInstrumentation {
+struct DominanceCounterInstrumentation
+  : public PassInstrumentation<DominanceCounterInstrumentation> {
   /// The cumulative count of how many times dominance has been calculated.
   unsigned &count;
 
   DominanceCounterInstrumentation(unsigned &count) : count(count) {}
   void runAfterAnalysis(llvm::StringRef, TypeID id, Operation *) override {
     if (id == TypeID::get<DominanceInfo>())
       ++count;
   }

mlir/include/mlir/Pass/Pass.h

 //===- Pass.h - Base classes for compiler passes ----------------*- 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_PASS_PASS_H
 #define MLIR_PASS_PASS_H
 
 #include "mlir/Pass/AnalysisManager.h"
 #include "mlir/Pass/PassRegistry.h"
 #include "mlir/Support/LogicalResult.h"
+#include "mlir/Support/Timing.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/Statistic.h"
 #include <optional>
 
 namespace mlir {
 namespace detail {
 class OpToOpPassAdaptor;
 struct OpPassManagerImpl;
protected:
 
   /// Schedule an arbitrary pass pipeline on the provided operation.
   /// This can be invoke any time in a pass to dynamic schedule more passes.
   /// The provided operation must be the current one or one nested below.
   LogicalResult runPipeline(OpPassManager &pipeline, Operation *op) {
     return passState->pipelineExecutor(pipeline, op);
   }
 
+  /// Create a nested timing scope. This returns a new `TimingScope` with a
+  /// timer nested within the current pass scope. In this fashion, the time in
+  /// this scope may be further subdivided in a more fine-grained fashion.
+  template <typename... Args>
+  TimingScope createTimingScope(Args... args) {
+    if (auto *pi = getAnalysisManager().getPassInstrumentor())
+      return pi->getCurrentTimer().nest(std::forward<Args>(args)...);
+    return TimingScope();
+  }
+
   /// A clone method to create a copy of this pass.
   std::unique_ptr<Pass> clone() const {
     auto newInst = clonePass();
     newInst->copyOptionValuesFrom(this);
     return newInst;
   }
 
   /// Return the current operation being transformed.

mlir/include/mlir/Pass/PassInstrumentation.h

 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/TypeID.h"
 #include <optional>
 
 namespace mlir {
 class OperationName;
 class Operation;
 class Pass;
+class Timer;
 
 namespace detail {
 struct PassInstrumentorImpl;
 } // namespace detail
 
-/// PassInstrumentation provides several entry points into the pass manager
-/// infrastructure. Instrumentations should be added directly to a PassManager
-/// before running a pipeline.
-class PassInstrumentation {
+/// PassInstrumentationBase represents the opaque instrumentation base class,
+/// with several entry points into the pass manager infrastructure.
+/// Instrumentations should be added directly to a PassManager before running a
+/// pipeline.
+class PassInstrumentationBase {
 public:
   /// This struct represents information related to the parent pass of pipeline.
   /// It includes information that allows for effectively linking pipelines that
   /// run on different threads.
   struct PipelineParentInfo {
     /// The thread of the parent pass that the current pipeline was spawned
     /// from. Note: This is acquired from llvm::get_threadid().
     uint64_t parentThreadID;
 
     /// The pass that spawned this pipeline.
     Pass *parentPass;
   };
 
-  virtual ~PassInstrumentation() = 0;
+  virtual ~PassInstrumentationBase() = 0;
 
   /// A callback to run before a pass pipeline is executed. This function takes
   /// the name of the operation type being operated on, or std::nullopt if the
   /// pipeline is op-agnostic, and information related to the parent that
   /// spawned this pipeline.
   virtual void runBeforePipeline(std::optional<OperationName> name,
                                  const PipelineParentInfo &parentInfo);
 
public:
   /// name of the analysis to be computed, its TypeID, as well as the
   /// current operation being analyzed.
   virtual void runBeforeAnalysis(StringRef name, TypeID id, Operation *op) {}
 
   /// A callback to run before an analysis is computed. This function takes the
   /// name of the analysis that was computed, its TypeID, as well as the
   /// current operation being analyzed.
   virtual void runAfterAnalysis(StringRef name, TypeID id, Operation *op) {}
+
+  /// Return the TypeID of the derived instrumentation.
+  TypeID getTypeID() const { return typeID; }
+
+protected:
+  PassInstrumentationBase(TypeID typeID) : typeID(typeID) {}
+
+private:
+  /// The TypeID of the derived class.
+  TypeID typeID;
+};
+
+/// This class provides a CRTP base class to simplify defining
+/// instrumentations.
+template <typename T>
+class PassInstrumentation : public PassInstrumentationBase {
+public:
+  static bool classof(const PassInstrumentationBase *pi) {
+    return pi->getTypeID() == TypeID::get<T>();
+  }
+
+protected:
+  PassInstrumentation() : PassInstrumentationBase(TypeID::get<T>()) {}
 };
 
 /// This class holds a collection of PassInstrumentation objects, and invokes
 /// their respective call backs.
 class PassInstrumentor {
 public:
   PassInstrumentor();
   PassInstrumentor(PassInstrumentor &&) = delete;
   PassInstrumentor(const PassInstrumentor &) = delete;
   ~PassInstrumentor();
 
-  /// See PassInstrumentation::runBeforePipeline for details.
-  void
-  runBeforePipeline(std::optional<OperationName> name,
-                    const PassInstrumentation::PipelineParentInfo &parentInfo);
-
-  /// See PassInstrumentation::runAfterPipeline for details.
-  void
-  runAfterPipeline(std::optional<OperationName> name,
-                   const PassInstrumentation::PipelineParentInfo &parentInfo);
+  /// Return the currently active timer for the current thread. This can be used
+  /// to nest additional timers under the current instrumentation scope.
+  Timer getCurrentTimer();
+
+  /// See PassInstrumentationBase::runBeforePipeline for details.
+  void runBeforePipeline(
+      std::optional<OperationName> name,
+      const PassInstrumentationBase::PipelineParentInfo &parentInfo);
+
+  /// See PassInstrumentationBase::runAfterPipeline for details.
+  void runAfterPipeline(
+      std::optional<OperationName> name,
+      const PassInstrumentationBase::PipelineParentInfo &parentInfo);
 
-  /// See PassInstrumentation::runBeforePass for details.
+  /// See PassInstrumentationBase::runBeforePass for details.
   void runBeforePass(Pass *pass, Operation *op);
 
-  /// See PassInstrumentation::runAfterPass for details.
+  /// See PassInstrumentationBase::runAfterPass for details.
   void runAfterPass(Pass *pass, Operation *op);
 
-  /// See PassInstrumentation::runAfterPassFailed for details.
+  /// See PassInstrumentationBase::runAfterPassFailed for details.
   void runAfterPassFailed(Pass *pass, Operation *op);
 
-  /// See PassInstrumentation::runBeforeAnalysis for details.
+  /// See PassInstrumentationBase::runBeforeAnalysis for details.
   void runBeforeAnalysis(StringRef name, TypeID id, Operation *op);
 
-  /// See PassInstrumentation::runAfterAnalysis for details.
+  /// See PassInstrumentationBase::runAfterAnalysis for details.
   void runAfterAnalysis(StringRef name, TypeID id, Operation *op);
 
   /// Add the given instrumentation to the collection.
-  void addInstrumentation(std::unique_ptr<PassInstrumentation> pi);
+  void addInstrumentation(std::unique_ptr<PassInstrumentationBase> pi);
 
 private:
   std::unique_ptr<detail::PassInstrumentorImpl> impl;
 };
 
 } // namespace mlir
 
 namespace llvm {
 template <>
-struct DenseMapInfo<mlir::PassInstrumentation::PipelineParentInfo> {
-  using T = mlir::PassInstrumentation::PipelineParentInfo;
+struct DenseMapInfo<mlir::PassInstrumentationBase::PipelineParentInfo> {
+  using T = mlir::PassInstrumentationBase::PipelineParentInfo;
   using PairInfo = DenseMapInfo<std::pair<uint64_t, void *>>;
 
   static T getEmptyKey() {
     auto pair = PairInfo::getEmptyKey();
     return {pair.first, reinterpret_cast<mlir::Pass *>(pair.second)};
   }
   static T getTombstoneKey() {
     auto pair = PairInfo::getTombstoneKey();

mlir/include/mlir/Pass/PassManager.h

 class Any;
 } // namespace llvm
 
 namespace mlir {
 class AnalysisManager;
 class MLIRContext;
 class Operation;
 class Pass;
-class PassInstrumentation;
+class PassInstrumentationBase;
 class PassInstrumentor;
 
 namespace detail {
 struct OpPassManagerImpl;
 class OpToOpPassAdaptor;
 class PassCrashReproducerGenerator;
 struct PassExecutionState;
 } // namespace detail
public:
   /// Runs the verifier after each individual pass.
   void enableVerifier(bool enabled = true);
 
   //===--------------------------------------------------------------------===//
   // Instrumentations
   //===--------------------------------------------------------------------===//
 
   /// Add the provided instrumentation to the pass manager.
-  void addInstrumentation(std::unique_ptr<PassInstrumentation> pi);
+  void addInstrumentation(std::unique_ptr<PassInstrumentationBase> pi);
 
   //===--------------------------------------------------------------------===//
   // IR Printing
 
   /// A configuration struct provided to the IR printer instrumentation.
   class IRPrinterConfig {
   public:
     using PrintCallbackFn = function_ref<void(raw_ostream &)>;

mlir/include/mlir/Support/Timing.h

public:
   template <typename... Args>
   TimingScope nest(Args... args) {
     return TimingScope(std::move(timer.nest(std::forward<Args>(args)...)));
   }
 
   /// Hide the timer in timing reports and directly show its children.
   void hide() { timer.hide(); }
 
+  /// Return the underlying timer managed by this scope.
+  Timer getTimer() { return timer; }
+
 private:
   /// The wrapped timer.
   Timer timer;
 };
 
 //===----------------------------------------------------------------------===//
 // DefaultTimingManager
 //===----------------------------------------------------------------------===//

mlir/lib/Pass/IRPrinting.cpp

 using namespace mlir;
 using namespace mlir::detail;
 
 namespace {
 //===----------------------------------------------------------------------===//
 // IRPrinter
 //===----------------------------------------------------------------------===//
 
-class IRPrinterInstrumentation : public PassInstrumentation {
+class IRPrinterInstrumentation
+    : public PassInstrumentation<IRPrinterInstrumentation> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(IRPrinterInstrumentation)
   IRPrinterInstrumentation(std::unique_ptr<PassManager::IRPrinterConfig> config)
       : config(std::move(config)) {}
 
 private:
   /// Instrumentation hooks.
   void runBeforePass(Pass *pass, Operation *op) override;
   void runAfterPass(Pass *pass, Operation *op) override;
   void runAfterPassFailed(Pass *pass, Operation *op) override;

mlir/lib/Pass/Pass.cpp

 //===- Pass.cpp - Pass infrastructure implementation ----------------------===//
 //
 // 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 common pass infrastructure.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Pass/Pass.h"
 #include "PassDetail.h"
+#include "PassTiming.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/Threading.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Support/FileUtilities.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
return op->emitOpError()
            << "trying to schedule a pass on an unregistered operation";
   if (!opInfo->hasTrait<OpTrait::IsIsolatedFromAbove>())
     return op->emitOpError() << "trying to schedule a pass on an operation not "
                                 "marked as 'IsolatedFromAbove'";
 
   // Initialize the pass state with a callback for the pass to dynamically
   // execute a pipeline on the currently visited operation.
   PassInstrumentor *pi = am.getPassInstrumentor();
-  PassInstrumentation::PipelineParentInfo parentInfo = {llvm::get_threadid(),
-                                                        pass};
+  PassInstrumentationBase::PipelineParentInfo parentInfo = {
+      llvm::get_threadid(), pass};
   auto dynamicPipelineCallback = [&](OpPassManager &pipeline,
                                      Operation *root) -> LogicalResult {
     if (!op->isAncestor(root))
       return root->emitOpError()
              << "Trying to schedule a dynamic pipeline on an "
                 "operation that isn't "
                 "nested under the current operation the pass is processing";
     assert(
#endif
   // Return if the pass signaled a failure.
   return failure(passFailed);
 }
 
 /// Run the given operation and analysis manager on a provided op pass manager.
 LogicalResult OpToOpPassAdaptor::runPipeline(
     OpPassManager &pm, Operation *op, AnalysisManager am, bool verifyPasses,
     unsigned parentInitGeneration, PassInstrumentor *instrumentor,
-    const PassInstrumentation::PipelineParentInfo *parentInfo) {
+    const PassInstrumentationBase::PipelineParentInfo *parentInfo) {
   assert((!instrumentor || parentInfo) &&
          "expected parent info if instrumentor is provided");
   auto scopeExit = llvm::make_scope_exit([&] {
     // Clear out any computed operation analyses. These analyses won't be used
     // any more in this pipeline, and this helps reduce the current working set
     // of memory. If preserving these analyses becomes important in the future
     // we can re-evaluate this.
     am.clear();
if (getContext().isMultithreadingEnabled())
     runOnOperationAsyncImpl(verifyPasses);
   else
     runOnOperationImpl(verifyPasses);
 }
 
 /// Run this pass adaptor synchronously.
 void OpToOpPassAdaptor::runOnOperationImpl(bool verifyPasses) {
   auto am = getAnalysisManager();
-  PassInstrumentation::PipelineParentInfo parentInfo = {llvm::get_threadid(),
-                                                        this};
+  PassInstrumentationBase::PipelineParentInfo parentInfo = {
+      llvm::get_threadid(), this};
   auto *instrumentor = am.getPassInstrumentor();
   for (auto &region : getOperation()->getRegions()) {
     for (auto &block : region) {
       for (auto &op : block) {
         auto *mgr = findPassManagerFor(mgrs, op.getName(), *op.getContext());
         if (!mgr)
           continue;
 
for (Operation &op : region.getOps()) {
 
       // If this operation can be scheduled, add it to the list.
       if (pmIdxIt.first->second)
         opInfos.emplace_back(*pmIdxIt.first->second, &op, am.nest(&op));
     }
   }
 
   // Get the current thread for this adaptor.
-  PassInstrumentation::PipelineParentInfo parentInfo = {llvm::get_threadid(),
-                                                        this};
+  PassInstrumentationBase::PipelineParentInfo parentInfo = {
+      llvm::get_threadid(), this};
   auto *instrumentor = am.getPassInstrumentor();
 
   // An atomic failure variable for the async executors.
   std::vector<std::atomic<bool>> activePMs(asyncExecutors.size());
   std::fill(activePMs.begin(), activePMs.end(), false);
   auto processFn = [&](OpPMInfo &opInfo) {
     // Find an executor for this operation.
     auto it = llvm::find_if(activePMs, [](std::atomic<bool> &isActive) {
LogicalResult PassManager::run(Operation *op) {
 
   // Dump all of the pass statistics if necessary.
   if (passStatisticsMode)
     dumpStatistics();
   return result;
 }
 
 /// Add the provided instrumentation to the pass manager.
-void PassManager::addInstrumentation(std::unique_ptr<PassInstrumentation> pi) {
+void PassManager::addInstrumentation(
+    std::unique_ptr<PassInstrumentationBase> pi) {
   if (!instrumentor)
     instrumentor = std::make_unique<PassInstrumentor>();
 
   instrumentor->addInstrumentation(std::move(pi));
 }
 
 LogicalResult PassManager::runPasses(Operation *op, AnalysisManager am) {
   return OpToOpPassAdaptor::runPipeline(*this, op, am, verifyPasses,
while (!mapsToInvalidate.empty()) {
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 // PassInstrumentation
 //===----------------------------------------------------------------------===//
 
-PassInstrumentation::~PassInstrumentation() = default;
+PassInstrumentationBase::~PassInstrumentationBase() = default;
 
-void PassInstrumentation::runBeforePipeline(
+void PassInstrumentationBase::runBeforePipeline(
     std::optional<OperationName> name, const PipelineParentInfo &parentInfo) {}
 
-void PassInstrumentation::runAfterPipeline(
+void PassInstrumentationBase::runAfterPipeline(
     std::optional<OperationName> name, const PipelineParentInfo &parentInfo) {}
 
 //===----------------------------------------------------------------------===//
 // PassInstrumentor
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace detail {
 struct PassInstrumentorImpl {
   /// Mutex to keep instrumentation access thread-safe.
   llvm::sys::SmartMutex<true> mutex;
 
   /// Set of registered instrumentations.
-  std::vector<std::unique_ptr<PassInstrumentation>> instrumentations;
+  std::vector<std::unique_ptr<PassInstrumentationBase>> instrumentations;
 };
 } // namespace detail
 } // namespace mlir
 
 PassInstrumentor::PassInstrumentor() : impl(new PassInstrumentorImpl()) {}
 PassInstrumentor::~PassInstrumentor() = default;
 
-/// See PassInstrumentation::runBeforePipeline for details.
+Timer PassInstrumentor::getCurrentTimer() {
+  llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+  for (auto &instr : impl->instrumentations)
+    if (auto *timing = dyn_cast<PassTiming>(instr.get()))
+      return timing->getActiveTimer();
+  return Timer();
+}
+
+/// See PassInstrumentationBase::runBeforePipeline for details.
 void PassInstrumentor::runBeforePipeline(
     std::optional<OperationName> name,
-    const PassInstrumentation::PipelineParentInfo &parentInfo) {
+    const PassInstrumentationBase::PipelineParentInfo &parentInfo) {
   llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
   for (auto &instr : impl->instrumentations)
     instr->runBeforePipeline(name, parentInfo);
 }
 
-/// See PassInstrumentation::runAfterPipeline for details.
+/// See PassInstrumentationBase::runAfterPipeline for details.
 void PassInstrumentor::runAfterPipeline(
     std::optional<OperationName> name,
-    const PassInstrumentation::PipelineParentInfo &parentInfo) {
+    const PassInstrumentationBase::PipelineParentInfo &parentInfo) {
   llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
   for (auto &instr : llvm::reverse(impl->instrumentations))
     instr->runAfterPipeline(name, parentInfo);
 }
 
-/// See PassInstrumentation::runBeforePass for details.
+/// See PassInstrumentationBase::runBeforePass for details.
 void PassInstrumentor::runBeforePass(Pass *pass, Operation *op) {
   llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
   for (auto &instr : impl->instrumentations)
     instr->runBeforePass(pass, op);
 }
 
-/// See PassInstrumentation::runAfterPass for details.
+/// See PassInstrumentationBase::runAfterPass for details.
 void PassInstrumentor::runAfterPass(Pass *pass, Operation *op) {
   llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
   for (auto &instr : llvm::reverse(impl->instrumentations))
     instr->runAfterPass(pass, op);
 }
 
-/// See PassInstrumentation::runAfterPassFailed for details.
+/// See PassInstrumentationBase::runAfterPassFailed for details.
 void PassInstrumentor::runAfterPassFailed(Pass *pass, Operation *op) {
   llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
   for (auto &instr : llvm::reverse(impl->instrumentations))
     instr->runAfterPassFailed(pass, op);
 }
 
-/// See PassInstrumentation::runBeforeAnalysis for details.
+/// See PassInstrumentationBase::runBeforeAnalysis for details.
 void PassInstrumentor::runBeforeAnalysis(StringRef name, TypeID id,
                                          Operation *op) {
   llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
   for (auto &instr : impl->instrumentations)
     instr->runBeforeAnalysis(name, id, op);
 }
 
-/// See PassInstrumentation::runAfterAnalysis for details.
+/// See PassInstrumentationBase::runAfterAnalysis for details.
 void PassInstrumentor::runAfterAnalysis(StringRef name, TypeID id,
                                         Operation *op) {
   llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
   for (auto &instr : llvm::reverse(impl->instrumentations))
     instr->runAfterAnalysis(name, id, op);
 }
 
 /// Add the given instrumentation to the collection.
 void PassInstrumentor::addInstrumentation(
-    std::unique_ptr<PassInstrumentation> pi) {
+    std::unique_ptr<PassInstrumentationBase> pi) {
   llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
   impl->instrumentations.emplace_back(std::move(pi));
 }

mlir/lib/Pass/PassCrashRecovery.cpp

void PassCrashReproducerGenerator::removeLastReproducerFor(Pass *pass,
   }
 }
 
 //===----------------------------------------------------------------------===//
 // CrashReproducerInstrumentation
 //===----------------------------------------------------------------------===//
 
 namespace {
-struct CrashReproducerInstrumentation : public PassInstrumentation {
+struct CrashReproducerInstrumentation
+    : public PassInstrumentation<CrashReproducerInstrumentation> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CrashReproducerInstrumentation)
+
   CrashReproducerInstrumentation(PassCrashReproducerGenerator &generator)
       : generator(generator) {}
   ~CrashReproducerInstrumentation() override = default;
 
   void runBeforePass(Pass *pass, Operation *op) override {
     if (!isa<OpToOpPassAdaptor>(pass))
       generator.prepareReproducerFor(pass, op);
   }

mlir/lib/Pass/PassDetail.h

private:
 
   /// Run the given operation and analysis manager on a provided op pass
   /// manager. `parentInitGeneration` is the initialization generation of the
   /// parent pass manager, and is used to initialize any dynamic pass pipelines
   /// run by the given passes.
   static LogicalResult runPipeline(
       OpPassManager &pm, Operation *op, AnalysisManager am, bool verifyPasses,
       unsigned parentInitGeneration, PassInstrumentor *instrumentor = nullptr,
-      const PassInstrumentation::PipelineParentInfo *parentInfo = nullptr);
+      const PassInstrumentationBase::PipelineParentInfo *parentInfo = nullptr);
 
   /// A set of adaptors to run.
   SmallVector<OpPassManager, 1> mgrs;
 
   /// A set of executors, cloned from the main executor, that run asynchronously
   /// on different threads. This is used when threading is enabled.
   SmallVector<SmallVector<OpPassManager, 1>, 8> asyncExecutors;
 

mlir/lib/Pass/PassTiming.h

+//===- PassTiming.h - MLIR Pass Timing --------------------------*- 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 LIB_MLIR_PASS_PASSTIMING_H_
+#define LIB_MLIR_PASS_PASSTIMING_H_
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+
+namespace mlir::detail {
+/// This class implements a pass instrumentation that records the time spent in
+/// each pass and pipeline.
+class PassTiming : public PassInstrumentation<PassTiming> {
+public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PassTiming)
+
+  PassTiming(TimingScope &timingScope) : rootScope(timingScope) {}
+  PassTiming(std::unique_ptr<TimingManager> tm)
+      : ownedTimingManager(std::move(tm)),
+        ownedTimingScope(ownedTimingManager->getRootScope()),
+        rootScope(ownedTimingScope) {}
+
+  /// Return the currently active timer for the current thread.
+  Timer getActiveTimer();
+
+  //===--------------------------------------------------------------------===//
+  // Pipeline
+  //===--------------------------------------------------------------------===//
+
+  void runBeforePipeline(std::optional<OperationName> name,
+                         const PipelineParentInfo &parentInfo) override;
+
+  void runAfterPipeline(std::optional<OperationName>,
+                        const PipelineParentInfo &) override;
+
+  //===--------------------------------------------------------------------===//
+  // Pass
+  //===--------------------------------------------------------------------===//
+
+  void runBeforePass(Pass *pass, Operation *) override;
+  void runAfterPass(Pass *pass, Operation *) override;
+  void runAfterPassFailed(Pass *pass, Operation *op) override {
+    runAfterPass(pass, op);
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Analysis
+  //===--------------------------------------------------------------------===//
+
+  void runBeforeAnalysis(StringRef name, TypeID id, Operation *) override;
+  void runAfterAnalysis(StringRef, TypeID, Operation *) override;
+
+private:
+  /// If a pass can spawn additional work on other threads, it records the
+  /// index to its currently active timer here. Passes that run on a
+  /// newly-forked thread will check this list to find the active timer of the
+  /// parent thread into which the new thread should be nested.
+  DenseMap<PipelineParentInfo, unsigned> parentTimerIndices;
+
+  /// A stack of the currently active timing scopes per thread.
+  DenseMap<uint64_t, SmallVector<TimingScope, 4>> activeThreadTimers;
+
+  /// The timing manager owned by this instrumentation (in case timing was
+  /// enabled by the user on the pass manager without providing an external
+  /// timing manager). This *must* appear before the `ownedTimingScope` to
+  /// ensure the timing manager is destroyed *after* the scope, since the latter
+  /// may hold a timer that points into the former.
+  std::unique_ptr<TimingManager> ownedTimingManager;
+  TimingScope ownedTimingScope;
+
+  /// The root timing scope into which timing is reported.
+  TimingScope &rootScope;
+};
+} // namespace mlir::detail
+
+#endif // LIB_MLIR_PASS_PASSTIMING_H_

mlir/lib/Pass/PassTiming.cpp

 //===- PassTiming.cpp -----------------------------------------------------===//
 //
 // 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 "PassTiming.h"
 #include "PassDetail.h"
 #include "mlir/Pass/PassManager.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Threading.h"
 
 #include <chrono>
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
 // PassTiming
 //===----------------------------------------------------------------------===//
 
-namespace {
-struct PassTiming : public PassInstrumentation {
-  PassTiming(TimingScope &timingScope) : rootScope(timingScope) {}
-  PassTiming(std::unique_ptr<TimingManager> tm)
-      : ownedTimingManager(std::move(tm)),
-        ownedTimingScope(ownedTimingManager->getRootScope()),
-        rootScope(ownedTimingScope) {}
-
-  /// If a pass can spawn additional work on other threads, it records the
-  /// index to its currently active timer here. Passes that run on a
-  /// newly-forked thread will check this list to find the active timer of the
-  /// parent thread into which the new thread should be nested.
-  DenseMap<PipelineParentInfo, unsigned> parentTimerIndices;
-
-  /// A stack of the currently active timing scopes per thread.
-  DenseMap<uint64_t, SmallVector<TimingScope, 4>> activeThreadTimers;
-
-  /// The timing manager owned by this instrumentation (in case timing was
-  /// enabled by the user on the pass manager without providing an external
-  /// timing manager). This *must* appear before the `ownedTimingScope` to
-  /// ensure the timing manager is destroyed *after* the scope, since the latter
-  /// may hold a timer that points into the former.
-  std::unique_ptr<TimingManager> ownedTimingManager;
-  TimingScope ownedTimingScope;
-
-  /// The root timing scope into which timing is reported.
-  TimingScope &rootScope;
+Timer PassTiming::getActiveTimer() {
+  auto tid = llvm::get_threadid();
+  auto &activeTimers = activeThreadTimers[tid];
+  if (activeTimers.empty())
+    return Timer();
+  return activeTimers.back().getTimer();
+}
 
-  //===--------------------------------------------------------------------===//
-  // Pipeline
-  //===--------------------------------------------------------------------===//
 
-  void runBeforePipeline(std::optional<OperationName> name,
-                         const PipelineParentInfo &parentInfo) override {
-    auto tid = llvm::get_threadid();
-    auto &activeTimers = activeThreadTimers[tid];
+void PassTiming::runBeforePipeline(std::optional<OperationName> name,
+                                   const PipelineParentInfo &parentInfo) {
 
-    // Find the parent scope, either using the parent info or the root scope
-    // (e.g. in the case of the top-level pipeline).
-    TimingScope *parentScope;
-    auto it = parentTimerIndices.find(parentInfo);
-    if (it != parentTimerIndices.end())
-      parentScope = &activeThreadTimers[parentInfo.parentThreadID][it->second];
-    else
-      parentScope = &rootScope;
 
-    // Use nullptr to anchor op-agnostic pipelines, otherwise use the name of
-    // the operation.
-    const void *timerId = name ? name->getAsOpaquePointer() : nullptr;
-    activeTimers.push_back(parentScope->nest(timerId, [name] {
-      return ("'" + (name ? name->getStringRef() : "any") + "' Pipeline").str();
-    }));
-  }
 
-  void runAfterPipeline(std::optional<OperationName>,
-                        const PipelineParentInfo &) override {
-    auto &activeTimers = activeThreadTimers[llvm::get_threadid()];
-    assert(!activeTimers.empty() && "expected active timer");
-    activeTimers.pop_back();
-  }
+void PassTiming::runAfterPipeline(std::optional<OperationName>,
+                                  const PipelineParentInfo &) {
 
-  //===--------------------------------------------------------------------===//
-  // Pass
-  //===--------------------------------------------------------------------===//
 
-  void runBeforePass(Pass *pass, Operation *) override {
-    auto tid = llvm::get_threadid();
-    auto &activeTimers = activeThreadTimers[tid];
-    auto &parentScope = activeTimers.empty() ? rootScope : activeTimers.back();
+void PassTiming::runBeforePass(Pass *pass, Operation *) {
 
-    if (auto *adaptor = dyn_cast<OpToOpPassAdaptor>(pass)) {
-      parentTimerIndices[{tid, pass}] = activeTimers.size();
-      auto scope =
-          parentScope.nest(pass->getThreadingSiblingOrThis(),
-                           [adaptor]() { return adaptor->getAdaptorName(); });
-      if (adaptor->getPassManagers().size() <= 1)
-        scope.hide();
-      activeTimers.push_back(std::move(scope));
-    } else {
-      activeTimers.push_back(
-          parentScope.nest(pass->getThreadingSiblingOrThis(),
-                           [pass]() { return std::string(pass->getName()); }));
-    }
-  }
 
-  void runAfterPass(Pass *pass, Operation *) override {
-    auto tid = llvm::get_threadid();
-    if (isa<OpToOpPassAdaptor>(pass))
-      parentTimerIndices.erase({tid, pass});
-    auto &activeTimers = activeThreadTimers[tid];
-    assert(!activeTimers.empty() && "expected active timer");
-    activeTimers.pop_back();
-  }
+void PassTiming::runAfterPass(Pass *pass, Operation *) {
 
-  void runAfterPassFailed(Pass *pass, Operation *op) override {
-    runAfterPass(pass, op);
-  }
-
-  //===--------------------------------------------------------------------===//
-  // Analysis
-  //===--------------------------------------------------------------------===//
 
-  void runBeforeAnalysis(StringRef name, TypeID id, Operation *) override {
-    auto tid = llvm::get_threadid();
-    auto &activeTimers = activeThreadTimers[tid];
-    auto &parentScope = activeTimers.empty() ? rootScope : activeTimers.back();
-    activeTimers.push_back(parentScope.nest(
-        id.getAsOpaquePointer(), [name] { return "(A) " + name.str(); }));
-  }
+void PassTiming::runBeforeAnalysis(StringRef name, TypeID id, Operation *) {
 
-  void runAfterAnalysis(StringRef, TypeID, Operation *) override {
-    auto &activeTimers = activeThreadTimers[llvm::get_threadid()];
-    assert(!activeTimers.empty() && "expected active timer");
-    activeTimers.pop_back();
-  }
-};
-} // namespace
+void PassTiming::runAfterAnalysis(StringRef, TypeID, Operation *) {
 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//
 
 /// Add an instrumentation to time the execution of passes and the computation
 /// of analyses.
 void PassManager::enableTiming(TimingScope &timingScope) {
   if (!timingScope)
     return;
   addInstrumentation(std::make_unique<PassTiming>(timingScope));
 }
 
 /// Add an instrumentation to time the execution of passes and the computation
 /// of analyses.
 void PassManager::enableTiming(std::unique_ptr<TimingManager> tm) {
+  // If the timing manager isn't active, no need to add it.
   if (!tm->getRootTimer())
-    return; // no need to keep the timing manager around if it's disabled
+    return;
   addInstrumentation(std::make_unique<PassTiming>(std::move(tm)));
 }
 
 /// Add an instrumentation to time the execution of passes and the computation
 /// of analyses.
 void PassManager::enableTiming() {
   auto tm = std::make_unique<DefaultTimingManager>();
   tm->setEnabled(true);
   enableTiming(std::move(tm));
 }

mlir/test/Pass/pass-timing.mlir

 // RUN: mlir-opt %s -mlir-disable-threading=true -verify-each=true -pass-pipeline='builtin.module(func.func(cse,canonicalize,cse))' -mlir-timing -mlir-timing-display=list 2>&1 | FileCheck -check-prefix=LIST %s
-// RUN: mlir-opt %s -mlir-disable-threading=true -verify-each=true -pass-pipeline='builtin.module(func.func(cse,canonicalize,cse))' -mlir-timing -mlir-timing-display=tree 2>&1 | FileCheck -check-prefix=PIPELINE %s
+// RUN: mlir-opt %s -mlir-disable-threading=true -verify-each=true -pass-pipeline='builtin.module(func.func(cse,canonicalize,cse,test-timingscope-pass))' -mlir-timing -mlir-timing-display=tree 2>&1 | FileCheck -check-prefix=PIPELINE %s
 // RUN: mlir-opt %s -mlir-disable-threading=false -verify-each=true -pass-pipeline='builtin.module(func.func(cse,canonicalize,cse))' -mlir-timing -mlir-timing-display=list 2>&1 | FileCheck -check-prefix=MT_LIST %s
 // RUN: mlir-opt %s -mlir-disable-threading=false -verify-each=true -pass-pipeline='builtin.module(func.func(cse,canonicalize,cse))' -mlir-timing -mlir-timing-display=tree 2>&1 | FileCheck -check-prefix=MT_PIPELINE %s
 // RUN: mlir-opt %s -mlir-disable-threading=true -verify-each=false -test-pm-nested-pipeline -mlir-timing -mlir-timing-display=tree 2>&1 | FileCheck -check-prefix=NESTED_PIPELINE %s
 
 // LIST: Execution time report
 // LIST: Total Execution Time:
 // LIST: Name
 // LIST-DAG: Canonicalizer
 // LIST-DAG: CSE
 // LIST-DAG: DominanceInfo
 // LIST: Total
 
 // PIPELINE: Execution time report
 // PIPELINE: Total Execution Time:
 // PIPELINE: Name
 // PIPELINE-NEXT: Parser
 // PIPELINE-NEXT: 'func.func' Pipeline
 // PIPELINE-NEXT:   CSE
 // PIPELINE-NEXT:     (A) DominanceInfo
 // PIPELINE-NEXT:   Canonicalizer
 // PIPELINE-NEXT:   CSE
 // PIPELINE-NEXT:     (A) DominanceInfo
+// PIPELINE-NEXT:   TestTimingScopePass
+// PIPELINE-NEXT:     Test Timing Scope
 // PIPELINE-NEXT: Output
 // PIPELINE-NEXT: Rest
 // PIPELINE-NEXT: Total
 
 // MT_LIST: Execution time report
 // MT_LIST: Total Execution Time:
 // MT_LIST: Name
 // MT_LIST-DAG: Canonicalizer

mlir/test/lib/Pass/TestPassManager.cpp

Statistic opCountDuplicate{this, "num-ops2",
                              "Number of operations counted one more time"};
   Statistic opCount{this, "num-ops", "Number of operations counted"};
 
   void runOnOperation() final {
     getOperation()->walk([&](Operation *) { ++opCount; });
     getOperation()->walk([&](Operation *) { ++opCountDuplicate; });
   }
 };
+
+/// A test pass that adds a nested timing scope.
+struct TestTimingScopePass
+    : public PassWrapper<TestTimingScopePass, OperationPass<>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTimingScopePass)
+
+  TestTimingScopePass() = default;
+  TestTimingScopePass(const TestStatisticPass &) : PassWrapper() {}
+  StringRef getArgument() const final { return "test-timingscope-pass"; }
+  StringRef getDescription() const final { return "Test pass timing scopes"; }
+
+  void runOnOperation() final { createTimingScope("Test Timing Scope"); }
+};
 } // namespace
 
 static void testNestedPipeline(OpPassManager &pm) {
   // Nest a module pipeline that contains:
   /// A module pass.
   auto &modulePM = pm.nest<ModuleOp>();
   modulePM.addPass(std::make_unique<TestModulePass>());
   /// A nested function pass.
void registerPassManagerTestPass() {
 
   PassRegistration<TestCrashRecoveryPass>();
   PassRegistration<TestFailurePass>();
   PassRegistration<TestInvalidIRPass>();
   PassRegistration<TestInvalidParentPass>();
 
   PassRegistration<TestStatisticPass>();
 
+  PassRegistration<TestTimingScopePass>();
+
   PassPipelineRegistration<>("test-pm-nested-pipeline",
                              "Test a nested pipeline in the pass manager",
                              testNestedPipeline);
   PassPipelineRegistration<>("test-textual-pm-nested-pipeline",
                              "Test a nested pipeline in the pass manager",
                              testNestedPipelineTextual);
 
   PassPipelineRegistration<TestOptionsPass::Options>
       registerOptionsPassPipeline(
           "test-options-pass-pipeline",
           "Parses options using pass pipeline registration",
           [](OpPassManager &pm, const TestOptionsPass::Options &options) {
             pm.addPass(std::make_unique<TestOptionsPass>(options));
           });
 }
 } // namespace mlir