[MLIR] Factor pass timing out into a dedicated timing manager

[fabianschuiki]

Summary

This factors out the pass timing code into a separate TimingManager that can be plugged into the PassManager from the outside. Users are able to provide their own implementation of this manager, and use it to time additional code paths outside of the pass manager. Also allows for multiple PassManagers to run and contribute to a single timing report.

More specifically, moves most of the existing infrastructure in Pass/PassTiming.cpp into a new Support/Timing.cpp file and adds a public interface in Support/Timing.h. The PassTiming instrumentation becomes a wrapper around the new timing infrastructure which adapts the instrumentation callbacks to the new timers.

This is a concrete suggestion towards https://llvm.discourse.group/t/add-entries-to-pass-timing-from-outside-the-passmanager/3115.

[mehdi_amini]

I support making the timing infrastructure usable and pluggable outside of the pass manager! Thanks for working on this, I looked into it at some point but didn't need it badly enough to finish it.

However I'm not sure about this approach at the moment: it seems that you're trying to keep the PM in charge of all the timing, and using it as a "timing manager" for things that aren't related to pass management.
I would think that it'd be better modeled by extracting the timing manager as a separate entities usable outside of the PM, and have the PM report to "timing manager" (through injection or similar).

[rriddle]

In D100647#2695523, @mehdi_amini wrote:

I support making the timing infrastructure usable and pluggable outside of the pass manager! Thanks for working on this, I looked into it at some point but didn't need it badly enough to finish it.

However I'm not sure about this approach at the moment: it seems that you're trying to keep the PM in charge of all the timing, and using it as a "timing manager" for things that aren't related to pass management.
I would think that it'd be better modeled by extracting the timing manager as a separate entities usable outside of the PM, and have the PM report to "timing manager" (through injection or similar).

+1 to what Mehdi says here. We want to have general infra related to timing (that ideally users can inject their own implementations of), with the pass manager simply hooking into that.

[lattner]

I agree with Mehdi and River. it seems like the passmgr should be a client of the more general shared thing.

I'd also recommend changing the top level flyer:
... Pass execution timing report ... to just ... Execution time report ...

[fabianschuiki]

Thanks for the great feedback! I also had a more general timing manager in mind that we could inject into the PM, but opted for this simpler change to get a discussion rolling. Let me factor out the timing interface into some abstract base class and the current implementation of time keeping into a concrete subclass. That would allow us to default to an MLIR-provided TimingManager class as a convenience (and to not break existing code maybe?), but still allow the user to provide a custom implementation if desired.

[fabianschuiki]

• Moved most of the timing infrastructure in Pass/PassTiming.cpp into a new Support/Timing.cpp file
• Added public interfaces for the timing stuff in Support/Timing.h
• Made the pass manager use the timers provided by the new Timing.h
• Added a new TimingManager that holds the root timer and prints timing reports (similar to the role the PassTiming instrumentation had previously)
• Added TimingManagerOptions and corresponding registerTimingManagerCLOptions and applyTimingManagerCLOptions functions

One problem I ran into when merging async executions of the same pipeline was identifying whether two executed passes were actually "the same" from the user's perspective (i.e. just cloned for the sake of threading). Just using the pass pointer as unique identifier would cause each thread to show up as a separate set of passes. An obvious solution would have been to use pass.getTypeID().getAsOpaquePointer(), but if the user intentionally added the same pass multiple times they would all be collapsed into one. I ended up adding a passPrototype field to Pass itself to track the "pass identity" through cloning inside the PM. (The original pass timing had the same problem and opted for merging the passes in a pipeline by position.)

I'm not sure how far we want to preserve compatibility with respect to CL options. The patch as it is now preserves the old --pass-timing and --pass-timing-display options of the PM, as well as the PM's enableTiming(config) function. If no external timing manager is provided, the PM will just create its own temporary one internally, report there, and print results as before. The new --timing and --timing-display can still be used and if present take precedence over --pass-timing. I'm also happy to throw out the old --pass-timing options, or rename the new ones to --pass-timing instead.

Finally I also added reporting for the "remaining" time not covered by dedicated timers, as suggested by @lattner. It is currently only shown for the top of the hierarchy. It's easy to add for all nested timers that have children, but that adds *a lot* of entries to the pipeline.

---

The output from the CIRCT project's firtool with the new timing manager looks as follows:

===-------------------------------------------------------------------------===
                        ... Execution timing report ...
===-------------------------------------------------------------------------===
  Total Execution Time: 0.0067 seconds

  ----User Time----  ----Wall Time----  ----Name----
    0.0010 ( 13.4%)    0.0010 ( 14.9%)  Parser
    0.0012 ( 15.7%)    0.0012 ( 17.5%)  'firrtl.circuit' Pipeline
    0.0004 (  6.0%)    0.0004 (  6.6%)    LowerFIRRTLTypes
    0.0006 (  7.7%)    0.0006 (  8.5%)    'firrtl.module' Pipeline
    0.0002 (  2.1%)    0.0002 (  2.3%)      CSE
    0.0000 (  0.1%)    0.0000 (  0.1%)        (A) DominanceInfo
    0.0004 (  5.3%)    0.0004 (  6.0%)      Canonicalizer
    0.0014 ( 19.3%)    0.0014 ( 21.5%)  LowerFIRRTLToRTL
    0.0002 (  2.2%)    0.0002 (  2.5%)  RTLMemSimImpl
    0.0011 ( 14.1%)    0.0004 (  5.5%)  'rtl.module' Pipeline
    0.0003 (  4.5%)    0.0002 (  2.3%)    RTLCleanup
    0.0002 (  3.2%)    0.0001 (  1.4%)    CSE
    0.0000 (  0.2%)    0.0000 (  0.1%)      (A) DominanceInfo
    0.0003 (  4.3%)    0.0001 (  1.6%)    Canonicalizer
    0.0005 (  6.5%)    0.0005 (  7.3%)  RTLLegalizeNames
    0.0002 (  2.2%)    0.0002 (  2.4%)    (A) circt::sv::LegalNamesAnalysis
    0.0008 ( 10.6%)    0.0008 ( 11.8%)  Output
    0.0008 ( 11.2%)    0.0008 ( 12.5%)  Rest
    0.0075 (100.0%)    0.0067 (100.0%)  Total

Of these Parser and Output are timers run before and after the PM, and Rest indicates the time that is not covered by the top-level timers.

[rriddle]

mlir/include/mlir/Pass/Pass.h

154

You can't rely on this meaning that the pass is a threading sibling with the original pass.

mlir/include/mlir/Support/Timing.h

113

nit: Just spell out auto here.

121

This display mode looks specific to the current pass manager implementation, I'm not sure this is generally applicable to the other potential timer backends. Couple of related general questions:

• Do we need a general TimingConfig anymore?

This was necessary before because of how the pass manager hid the internal timing instrumentation, but I'm not sure that it is really necessary anymore. How much of it would be shared across different potential timer backends?

• Can you rename TimingManagerBase base to just TimingManager?

I don't think the default implementation should be called TimingManager, it should likely be called something else(not sure the name though). This would establish that TimingManager (as you currently have it) is just an implementation and not the interface.

• Can you add comment blocks? I find these to help separate different concepts within a source file (e.g. https://github.com/llvm/llvm-project/blob/a7f33d720b74a92031958fce0e209794d296393d/mlir/include/mlir/Pass/Pass.h#L310)

• Do you have an idea of what supporting another timing backend would look like?

Not that you have to do the work, but when I've thought about this before ideally I would want to be able to support something like (1) or (2) without needing a new API/interface. Ideally these would all just be implementations that a user can choose from and pass the general infra to use for timer reporting.

(1): https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/Support/TimeProfiler.h
(2): https://github.com/wolfpld/tracy

mlir/lib/Support/Timing.cpp

134

The virtual isn't necessary when you have override.

139

nit: Drop the else after a return.

171

Drop the trivial braces here.

[lattner]

This is looking really great Fabian!

mlir/include/mlir/Pass/PassManager.h

306

Unless this is a widely used typedef that could occur in downstream code, I'd recommend just eliminating the PassTimingConfig name, updating clients to use TimingConfig directly.

mlir/include/mlir/Support/Timing.h

26

Please follow clang-tidy's advice here and add a virtual dtor. Also, please declare one method (e.g. the dtor) out of line in a .cpp file to give the vtable a home in a .o file.

161

I agree with River's comment that it would be great to eliminate this class.

mlir/lib/Support/Timing.cpp

107

As clang tidy suggests please declare this as a class. The difference matters for MSVC

[fabianschuiki]

mlir/include/mlir/Pass/Pass.h

154

Renamed this to threadingSibling and moved the tracking of copies from Pass::clone() into the OpPassManager constructor. Since the OpPassManager is only duplicated in OpToOpPassAdaptor::runOnOperationAsyncImpl, that should reliably identify clones that were made for the sake of threading.

[fabianschuiki]

Thanks for the great feedback @rriddle and @lattner! Updated the diff with your suggestions and fixed the check-mlir run.

• @rriddle's question about how easily we can spin up additional implementations for stuff like the LLVM TimeProfiler or tracy got me thinking. The previous design would require you to subclass both TimerBase and TimingManagerBase to hook into the infrastructure. Since you had to hand out TimerBase instances, you had to allocate them somewhere in memory, which was particularly bad if the external timing API you were wrapping was giving you IDs or pointers to its own data structures. This revision now simplifies the design: TimerBase has been replaced by a Timer that is not intended to be subclassed, but packages a pointer to a timing manager and an opaque ID/handle identifying the timer. TimingManagerBase is now called TimingManager and has overridable callback functions for timer start, stop, and nesting which Timer forwards to. That means there is one thing for implementations to subclass (TimingManager), and one thing for library code to interact with (Timer, and by extension TimingScope). Makes things much easier to use. Also, wrappers around external timing libs can just pass around the pointers they get from those libs as opaque handles in Timer.

• Dropped all the TimingConfig business and made the TimingDisplayMode a characteristic of the default timing manager implementation.

• DefaultTimingManager is maybe not the best of names; suggestions are very welcome! Maybe ConsoleTimingManager, BasicTimingManager, SimpleTimingManager?

• The redundancy between the timing options in PassManagerOptions and the new ones in TimingManagerOptions felt awkward. I dropped the timing options from PM altogether in favor of dealing with timing config completely outside of the PM. The pass manager can own a TimingManager and there is now a applyDefaultTimingPassManagerCLOptions() function which generates a dedicated timing manager for the PM with some options in case the user does not want to deal with timing themselves. The intended usage though is through PassManager::enableTiming(TimingManager).

• I renamed the --pass-timing and --pass-timing-display options to --timing and --timing-display, respectively, to reflect that the scope of these has broadened beyond pass timing.

• I updated MlirOptMain and the pass-timing.mlir test to use the new timing infrastructure and options, and to act as an example on how to use them.

[lattner]

This looks really nice Fabian!

[rriddle]

How do you envision this being used deep within the compiler? Should we have a static API that is accessible from anywhere? Should we have something on the context? What are your thoughts here?

Also, please fix the clang tidy issues.

mlir/include/mlir/Support/Timing.h

160

nit: Please make bool conversion operators explicit.

[rriddle]

(Apologies for the delay in review, will take a pass over tomorrow)

[fabianschuiki]

Fixed clang-tidy issues and made bool conversion explicit. Sorry about those, having a somewhat hard time to get my local clang-tidy to reproduce the set of upstream checks.

Regarding usage within the compiler: Having to thread TimingScopes through everywhere sounds pretty unergonomic if you just want to time a few things within the compiler. Ideally we would have a per-thread stack of active timing scopes somewhere (Maybe in the context? As a thread-local static?). Pointers to the active TimingScopes would then be kept on that stack, and you could just ask for the currently-active one to nest into. That way you could get at whatever scope within the parent code you're implicitly executing in, without needing to explicitly push scope references down the call chain. Maybe context.getTimingScope(id, nameBuilder) and context.getTimer(id, nameBuilder)?

[mehdi_amini]

In general, I have concerns with any design that would single global data structure for book-keeping. I would look into hooking the storage onto the MLIRContext as needed.

[rriddle]

In D100647#2735576, @fabianschuiki wrote:

Fixed clang-tidy issues and made bool conversion explicit. Sorry about those, having a somewhat hard time to get my local clang-tidy to reproduce the set of upstream checks.

Regarding usage within the compiler: Having to thread TimingScopes through everywhere sounds pretty unergonomic if you just want to time a few things within the compiler. Ideally we would have a per-thread stack of active timing scopes somewhere (Maybe in the context? As a thread-local static?). Pointers to the active TimingScopes would then be kept on that stack, and you could just ask for the currently-active one to nest into. That way you could get at whatever scope within the parent code you're implicitly executing in, without needing to explicitly push scope references down the call chain. Maybe context.getTimingScope(id, nameBuilder) and context.getTimer(id, nameBuilder)?

I would not add timing API to the context itself, but instead hook in a reference to a timing manager. Mehdi touches a bit on the general opinion that MLIR takes about these in terms of static vs non-static, and returning a manager object is inline with the other things hooked into the context. One important piece here though, is that it should generally be zero-cost if there is no timing manager.

mlir/include/mlir/Pass/PassManager.h

350

Why can't this be handled by the timing instrumentation? I don't think we need these fields here.

mlir/include/mlir/Support/Timing.h

17

This is an indicator that the dependencies are off. Support/ shouldn't depend on anything in IR/.

333

nit: I don't think the llvm:: here is necessary if you include LLVM.h

mlir/lib/Pass/PassTiming.cpp

95

nit: Drop the trivial braces here.

mlir/lib/Support/MlirOptMain.cpp

51 ↗

(On Diff #342652)

I'm not sure about this type of integration. This seems to always enforce timing, which isn't something that I would expect. If we start integrating timing more deeply into the compiler. I would expect it to do nothing unless explicitly enabled.

[fabianschuiki]

I think having a reference to a TimingManager in the context would be great. That would allow us to encode timing enable/disable through the presence/absence of a such a reference (right now it's an explicit field in the TM). For the sake of keeping code that uses the timing API streamlined, we should probably have the getTimingManager() return a NullTimingManager reference that always hands out disabled Timers. Since the latter are just a wrapper around a pointer, and their functions all return immediately if disabled, the runtime cost in code should be equivalent to having manual "is timing enabled" checks.

Do you guys want to tackle the integration into MLIRContext already as part of this diff, or push that off into a separate one?

mlir/include/mlir/Pass/PassManager.h

350

Yeah good point, I'll move that into the timing instrumentation.

mlir/include/mlir/Support/Timing.h

17

Good point. Will check for some other/dedicated mechanism to map a string to a opaque pointer identifying the string content. Since we have a TimingManager now, that could just as well go in there.

mlir/lib/Support/MlirOptMain.cpp

51 ↗

(On Diff #342652)

The timing is disabled in the manager upon construction. The call to applyDefaultTimingManagerCLOptions(tm) on the next line enables timing if the options are present on the command line. If the timing remains disabled, getRootScope() below returns a disabled TimingScope that has a fast path to do nothing (just a pointer comparison).

With some integration into MLIRContext and a stack of running timers, you could remove the timing variable below and instead have things like:

DefaultTimingManager tm;
applyDefaultTimingManagerCLOptions(tm);
context->setTimingManager(tm); // or some ownership transfer

auto parserTiming = context->getTimingManager()->addTimingScope("Parser");

parserTiming.stop()

auto outputTiming = context->getTimingManager()->addTimingScope("Output");

[rriddle]

In D100647#2738046, @fabianschuiki wrote:

I think having a reference to a TimingManager in the context would be great. That would allow us to encode timing enable/disable through the presence/absence of a such a reference (right now it's an explicit field in the TM). For the sake of keeping code that uses the timing API streamlined, we should probably have the getTimingManager() return a NullTimingManager reference that always hands out disabled Timers. Since the latter are just a wrapper around a pointer, and their functions all return immediately if disabled, the runtime cost in code should be equivalent to having manual "is timing enabled" checks.

Do you guys want to tackle the integration into MLIRContext already as part of this diff, or push that off into a separate one?

I'm okay with shifting that a follow up, so that it is easier to review the diff.

Can you also draft up some documentation for the timing infra? with examples on users would interact with it, and how it could be used to inject custom timers? I'm fine with that being a follow up as well.

This looks really great! I'm looking forward to being able to use it.

mlir/include/mlir/Support/Timing.h

87–88

/// Implementations override this method to provide access to the top-level timer.

This method isn't virtual, what do you mean by the comment above?

mlir/lib/Support/MlirOptMain.cpp

51 ↗

(On Diff #342652)

Ah, I completely missed that. Thanks for pointing it out! I like what you have here.

[fabianschuiki]

mlir/include/mlir/Support/Timing.h

87–88

Oops, my bad -- that's a leftover from when rootTimer() below wasn't a thing.

[fabianschuiki]

• Moved the owned timing manager and scope into PassTiming, dropping those fields from PassManager.
• Added a TimingIdentifier to remove the dependency on IR/Identifier.h. This new identifier is a version of Identifier/MLIRContext but stripped-down to cover only the functionality needed by the TM.
• Other cleanup

[rriddle]

Looks good to me with comments resolved.

(Also should clarify when I said docs, I was thinking something in mlir/docs/)

mlir/include/mlir/Support/Timing.h

80

demarcate

231

Can you change all of the uses of std::function<> && to function_ref? std::function generally ends up allocating memory, which we don't really want to do. Given that we aren't storing the function, we only really need a reference anyways.

403

Is this something we could make non-virtual? I think it would be nice if we could avoid an indirect call in the case where timing is disabled.

mlir/lib/Support/Timing.cpp

34

This shouldn't be necessary, StringRef is exported in the mlir namespace (see LLVM.h)

543

Can you prefix these with mlir-? I'm always hesitant about adding very commonly named things to the global command line namespace.

[fabianschuiki]

mlir/include/mlir/Support/Timing.h

403

The way you use the TM at the moment is roughly:

DefaultTimingManager tm;
TimingScope ts = tm.getRootScope();
// all further timing calls go through `ts`

If timing is disabled, ts is marked as diabled through a nullptr in its tm field, which causes all subsequent interaction to short-circuit into no-ops. So in its current implementation, if timing is disabled there is just one virtual function call for the duration of the program -- that should not be visible in practice.

I do agree that it would be nicer if the setup-related interaction with the TM would be non-virtual. Especially if we start to keep per-thread stacks of active timing scopes, that should all be directly accessible. We could promote setEnabled and isEnabled into the TimingManager base class such that it can just never ask for the rootTimer if disabled. Pretty sure that all TM impls can get behind the concept of enabling/disabling.

Do you want to tackle this in this diff, or push that off to a follow-up diff alongside tracking a TM in the context?

mlir/lib/Support/Timing.cpp

543

Yeah that's a good idea. --timing did seem awfully brief.

[fabianschuiki]

Applied suggestions by @rriddle.

[rriddle]

mlir/include/mlir/Support/Timing.h

403

Followup is alright with me, thanks!

[fabianschuiki]

Awesome, thanks! What are the next steps (sorry, first time Phabricator user)?

[lattner]

git pull --rebase
ninja check-mlir # sanity check
git push

if needed:
git pull --rebase
git push

[fabianschuiki]

Thanks @lattner! Wasn't sure whether the ready-to-land is the green light to push to the repo myself. Done!

[lattner]

Thanks again Fabian. It would also be good to make sure this gets mentioned in MLIR biweekly

[mehdi_amini]

In D100647#2754448, @lattner wrote:

Thanks again Fabian. It would also be good to make sure this gets mentioned in MLIR biweekly

Yes please! :)

(here is the link: https://llvm.discourse.group/t/work-in-progress-next-mlir-news-33rd-edition-5-1-5-14-2021/3417 )

[fabianschuiki]

In D100647#2755459, @mehdi_amini wrote:

In D100647#2754448, @lattner wrote:

Thanks again Fabian. It would also be good to make sure this gets mentioned in MLIR biweekly

Yes please! :)

(here is the link: https://llvm.discourse.group/t/work-in-progress-next-mlir-news-33rd-edition-5-1-5-14-2021/3417 )

Done! :-)

mlir/include/mlir/Pass/Pass.h

private:
     /// Hide some of the details of llvm::Statistic that we don't use.
     using llvm::Statistic::getDebugType;
   };
 
   /// Returns the main statistics for this pass instance.
   ArrayRef<Statistic *> getStatistics() const { return statistics; }
   MutableArrayRef<Statistic *> getStatistics() { return statistics; }
 
+  /// Returns the prototype of this pass.
+  ///
+  /// If this pass was created through `clonePass()`, returns the original pass
+  /// it was cloned from. This is useful for diagnostic purposes to distinguish
+  /// passes that were replicated internally in the pass manager (e.g., for
+  /// threading purposes) from passes purposefully instantiated by the user.
+  /// Used to collapse passes in timing statistics.
+  const Pass *getPrototype() const {
+    return passPrototype ? passPrototype : this;
+  }
+
 protected:
   explicit Pass(TypeID passID, Optional<StringRef> opName = llvm::None)
       : passID(passID), opName(opName) {}
   Pass(const Pass &other) : Pass(other.passID, other.opName) {}
 
   /// Returns the current pass state.
   detail::PassExecutionState &getPassState() {
     assert(passState && "pass state was never initialized");
protected:
   /// The provided operation must be the current one or one nested below.
   LogicalResult runPipeline(OpPassManager &pipeline, Operation *op) {
     return passState->pipelineExecutor(pipeline, op);
   }
 
   /// A clone method to create a copy of this pass.
   std::unique_ptr<Pass> clone() const {
     auto newInst = clonePass();
+    newInst->passPrototype = getPrototype();
     newInst->copyOptionValuesFrom(this);
     return newInst;
   }
 
   /// Return the current operation being transformed.
   Operation *getOperation() {
     return getPassState().irAndPassFailed.getPointer();
   }
private:
   Optional<detail::PassExecutionState> passState;
 
   /// The set of statistics held by this pass.
   std::vector<Statistic *> statistics;
 
   /// The pass options registered to this pass instance.
   detail::PassOptions passOptions;
 
+  /// A pointer to the pass this pass was cloned from, or `nullptr` if the pass
+  /// is an original.
+  const Pass *passPrototype = nullptr;
+
   /// Allow access to 'clone'.
   friend class OpPassManager;
 
   /// Allow access to 'passState'.
   friend detail::OpToOpPassAdaptor;
 
   /// Allow access to 'passOptions'.
   friend class PassInfo;

mlir/include/mlir/Pass/PassManager.h

 //===- PassManager.h - Pass Management Interface ----------------*- 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_PASSMANAGER_H
 #define MLIR_PASS_PASSMANAGER_H
 
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/Support/LogicalResult.h"
+#include "mlir/Support/Timing.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include <functional>
 #include <vector>
 
private:
   /// Allow access.
   friend detail::OpPassManagerImpl;
 };
 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//
 
-/// An enum describing the different display modes for the information within
-/// the pass manager.
-enum class PassDisplayMode {
-  // In this mode the results are displayed in a list sorted by total,
-  // with each pass/analysis instance aggregated into one unique result.
-  List,
-
-  // In this mode the results are displayed in a nested pipeline view that
-  // mirrors the internal pass pipeline that is being executed in the pass
-  // manager.
-  Pipeline,
-};
+using PassDisplayMode = TimingDisplayMode;
 
 /// The main pass manager and pipeline builder.
 class PassManager : public OpPassManager {
 public:
   /// Create a new pass manager under the given context with a specific nesting
   /// style. The created pass manager can schedule operations that match
   /// `operationName`.
   PassManager(MLIRContext *ctx, Nesting nesting = Nesting::Explicit,
void enableIRPrinting(
           [](Pass *, Operation *) { return true; },
       bool printModuleScope = true, bool printAfterOnlyOnChange = true,
       raw_ostream &out = llvm::errs(),
       OpPrintingFlags opPrintingFlags = OpPrintingFlags());
 
   //===--------------------------------------------------------------------===//
   // Pass Timing
 
-  /// A configuration struct provided to the pass timing feature.
-  class PassTimingConfig {
-  public:
-    using PrintCallbackFn = function_ref<void(raw_ostream &)>;
+  using PassTimingConfig = TimingConfig;
 
-    /// Initialize the configuration.
-    /// * 'displayMode' switch between list or pipeline display (see the
-    /// `PassDisplayMode` enum documentation).
-    explicit PassTimingConfig(
-        PassDisplayMode displayMode = PassDisplayMode::Pipeline)
-        : displayMode(displayMode) {}
-
-    virtual ~PassTimingConfig();
-
-    /// A hook that may be overridden by a derived config to control the
-    /// printing. The callback is supplied by the framework and the config is
-    /// responsible to call it back with a stream for the output.
-    virtual void printTiming(PrintCallbackFn printCallback);
-
-    /// Return the `PassDisplayMode` this config was created with.
-    PassDisplayMode getDisplayMode() { return displayMode; }
-
-  private:
-    PassDisplayMode displayMode;
-  };
+  /// Add an instrumentation to time the execution of passes and the computation
+  /// of analyses. This function creates a temporary timing manager with the
+  /// given optional `config`. An execution time report will be printed as soon
+  /// as this `PassManager` is destroyed.
+  ///
+  /// Note: Timing should be enabled after all other instrumentations to avoid
+  /// any potential "ghost" timing from other instrumentations being
+  /// unintentionally included in the timing results.
+  void enableTiming(std::unique_ptr<TimingConfig> config = nullptr);
 
   /// Add an instrumentation to time the execution of passes and the computation
-  /// of analyses.
+  /// of analyses. Timing will be reported by nesting timers into the provided
+  /// `timingScope`.
+  ///
   /// Note: Timing should be enabled after all other instrumentations to avoid
   /// any potential "ghost" timing from other instrumentations being
   /// unintentionally included in the timing results.
-  void enableTiming(std::unique_ptr<PassTimingConfig> config = nullptr);
+  void enableTiming(TimingScope &timingScope);
 
   /// Prompts the pass manager to print the statistics collected for each of the
   /// held passes after each call to 'run'.
   void
   enableStatistics(PassDisplayMode displayMode = PassDisplayMode::Pipeline);
 
 private:
   /// Dump the statistics of the passes within this pass manager.
   void dumpStatistics();
 
   /// Run the pass manager with crash recover enabled.
   LogicalResult runWithCrashRecovery(Operation *op, AnalysisManager am);
   /// Run the given passes with crash recover enabled.
   LogicalResult
   runWithCrashRecovery(MutableArrayRef<std::unique_ptr<Pass>> passes,
                        Operation *op, AnalysisManager am);
 
   /// Context this PassManager was initialized with.
   MLIRContext *context;
 
   /// Flag that specifies if pass statistics should be dumped.
   Optional<PassDisplayMode> passStatisticsMode;
 
+  /// Timing manager owned by this pass manager (in case timing is enabled by
+  /// the user without providing an external timing manager). This *must* appear
+  /// before the instrumentor to ensure the timing manager is destroyed *after*
+  /// the instrumentor, since the latter may hold timer pointers into the
+  /// former.
+  std::unique_ptr<TimingManager> ownedTimingManager;
+
+  /// See `ownedTimingManager`.
+  TimingScope ownedTimingScope;
+
   /// A manager for pass instrumentations.
   std::unique_ptr<PassInstrumentor> instrumentor;
 
   /// An optional factory to use when generating a crash reproducer if valid.
   ReproducerStreamFactory crashReproducerStreamFactory;
 
   /// A hash key used to detect when reinitialization is necessary.
   llvm::hash_code initializationKey;

mlir/include/mlir/Support/Timing.h

+//===- Timing.h - Execution time measurement facilities ---------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Facilities to measure and provide statistics on execution time.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_SUPPORT_TIMING_H
+#define MLIR_SUPPORT_TIMING_H
+
+#include "mlir/IR/Identifier.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <functional>
+#include <memory>
+#include <string>
+
+namespace mlir {
+
+class TimerBase {
+public:
+  /// Start the timer.
+  virtual void start() = 0;
+
+  /// Stop the timer.
+  virtual void stop() = 0;
+
+  /// Hide the timer in timing reports and directly show its children. Optional.
+  /// May be unsupported by implementations.
+  virtual void setTransparent(bool transparent) {}
+
+  /// Create a child timer nested within this one. Multiple calls to this
+  /// function with the same unique identifier `id` will return the same child
+  /// timer.
+  ///
+  /// This function can be called from other threads, as long as this timer
+  /// outlives any uses of the child timer on the other thread.
+  virtual TimerBase *nest(const void *id,
+                          std::function<std::string()> &&nameBuilder) = 0;
+
+  /// See above.
+  TimerBase *nest(Identifier name) {
+    return nest(name.getAsOpaquePointer(), [=]() { return name.str(); });
+  }
+
+  /// See above.
+  TimerBase *nest(StringRef name, MLIRContext *context) {
+    return nest(Identifier::get(name, context));
+  }
+};
+
+/// An RAII-style wrapper around a timer that ensures the timer is properly
+/// started and stopped.
+class TimingScope {
+  TimerBase *timer;
+
+public:
+  TimingScope() : timer(nullptr) {}
+  TimingScope(TimerBase *timer) : timer(timer) {
+    if (timer)
+      timer->start();
+  }
+  TimingScope(TimingScope &&other) : timer(other.release()) {}
+  ~TimingScope() { stop(); }
+
+  TimingScope &operator=(TimingScope &&other) {
+    stop();
+    timer = other.release();
+    return *this;
+  }
+
+  // Disable copying of the `TimingScope`.
+  TimingScope(const TimingScope &) = delete;
+  TimingScope &operator=(const TimingScope &) = delete;
+
+  /// Manually stop the timer early.
+  void stop() {
+    if (timer)
+      timer->stop();
+    timer = nullptr;
+  }
+
+  /// Hide the timer in timing reports and directly show its children.
+  void setTransparent(bool transparent) {
+    if (timer)
+      timer->setTransparent(transparent);
+  }
+
+  /// Create a nested timing scope.
+  ///
+  /// This returns a new `TimingScope` with a timer nested within the current
+  /// scope. In this fashion, the time in this scope may be further subdivided
+  /// in a more fine-grained fashion.
+  template <typename... Args>
+  TimingScope nest(Args... args) {
+    return TimingScope(timer ? timer->nest(std::forward<Args>(args)...)
+                             : nullptr);
+  }
+
+  /// Implicit call to `bool` to check if the timing scope actually contains a
+  /// valid timer.
+  operator bool() const { return timer != nullptr; }
+
+private:
+  /// Release control of the timer in this scope, without stopping it.
+  TimerBase *release() {
+    auto t = timer;
+    timer = nullptr;
+    return t;
+  }
+};
+
+/// An enum describing the different display modes for the information within
+/// the timing manager.
+enum class TimingDisplayMode {
+  /// In this mode the results are displayed in a list sorted by total time,
+  /// with timers aggregated into one unique result per timer name.
+  List,
+
+  // In this mode the results are displayed in a nested pipeline view that
+  // resembles the internal pass pipeline that is being executed in a pass
+  // manager.
+  Pipeline,
+};
+
+/// A configuration struct provided to the timing facilities.
+class TimingConfig {
+public:
+  using PrintCallbackFn = llvm::function_ref<void(llvm::raw_ostream &)>;
+
+  /// Initialize the configuration.
+  /// * 'displayMode' switches between list or pipeline display (see the
+  /// `TimingDisplayMode` enum documentation).
+  explicit TimingConfig(
+      TimingDisplayMode displayMode = TimingDisplayMode::Pipeline)
+      : displayMode(displayMode) {}
+
+  virtual ~TimingConfig();
+
+  /// A hook that may be overridden by a derived config to control printing of
+  /// timing results. The callback is supplied by the framework and the config
+  /// is responsible to call it back with a stream for the output. The default
+  /// implementation uses `llvm::CreateInfoOutputFile()` as stream.
+  virtual void printTiming(PrintCallbackFn printCallback);
+
+  /// Return the `TimingDisplayMode` this config was created with.
+  TimingDisplayMode getDisplayMode() { return displayMode; }
+
+private:
+  TimingDisplayMode displayMode;
+};
+
+/// An abstract interface to an object that can be used for execution time
+/// measurement.
+class TimingManagerBase {
+public:
+  /// Get the root timer of this timing manager. The returned timer must be
+  /// started and stopped manually. Execution time can be measured by nesting
+  /// timers within the root timer and starting/stopping them as appropriate.
+  /// Use this function only if you need access to the timer itself. Otherwise
+  /// consider the more convenient `getRootScope()` which offers an RAII-style
+  /// wrapper around the timer. This function returns a `nullptr` if timing is
+  /// disabled.
+  virtual TimerBase *getRootTimer() = 0;
+
+  /// Get the root timer of this timing manager wrapped in a `TimingScope` for
+  /// convenience. Automatically starts the timer and stops it as soon as the
+  /// `TimingScope` goes out of scope (is destroyed).
+  virtual TimingScope getRootScope() { return TimingScope(getRootTimer()); }
+
+  /// Enable execution time sampling.
+  virtual void enableTiming(std::unique_ptr<TimingConfig> config = nullptr) = 0;
+};
+
+/// A manager for execution time measurement and report printing.
+///
+/// This is MLIR's default implementation of a `TimingManagerBase`. Prints an
+/// execution time report upon destruction, or manually through `print()`. Use
+/// a `TimingConfig` to control the formatting of the report and the stream to
+/// which it is printed. Use `getRootScope()` for the most convenient way to
+/// measure execution time. Library code should accept a `TimingManagerBase` or
+/// `TimerBase` rather than this concrete implementation to allow the user to
+/// substitute their own implementation.
+class TimingManager : public TimingManagerBase {
+public:
+  TimingManager();
+  TimingManager(TimingManager &&rhs);
+  ~TimingManager();
+
+  // Disable copying of the `TimingManager`.
+  TimingManager(const TimingManager &rhs) = delete;
+  TimingManager &operator=(const TimingManager &rhs) = delete;
+
+  virtual TimerBase *getRootTimer() override;
+
+  /// Enable execution time sampling.
+  virtual void
+  enableTiming(std::unique_ptr<TimingConfig> config = nullptr) override;
+
+  /// Print and clear the timing results. Only call this when there are no more
+  /// references to nested timers around, as printing post-processes and clears
+  /// the timers.
+  void print();
+
+  /// Clear the timing results. Only call this when there are no more references
+  /// to nested timers around, as clearing invalidates them.
+  void clear();
+
+  /// Debug print the timer data structures to an output stream.
+  void dumpTimers(llvm::raw_ostream &os = llvm::errs());
+
+  /// Debug print the timers as a list. Only call this when there are no more
+  /// references to nested timers around.
+  void dumpAsList(llvm::raw_ostream &os = llvm::errs());
+
+  /// Debug print the timers as a pipeline. Only call this when there are no
+  /// more references to nested timers around.
+  void dumpAsPipeline(llvm::raw_ostream &os = llvm::errs());
+
+private:
+  class Timer;
+  std::unique_ptr<Timer> rootTimer;
+  std::unique_ptr<TimingConfig> config;
+};
+
+/// Register a set of useful command-line options that can be used to configure
+/// a timing manager. The values of these options can be applied via the
+/// `applyTimingManagerCLOptions` method below.
+void registerTimingManagerCLOptions();
+
+/// Apply any values provided to the timing manager options that were registered
+/// with 'registerTimingManagerOptions'.
+void applyTimingManagerCLOptions(TimingManagerBase &tm);
+
+} // namespace mlir
+
+#endif // MLIR_SUPPORT_TIMING_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 "PassDetail.h"
 #include "mlir/Pass/PassManager.h"
-#include "llvm/ADT/MapVector.h"
-#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/Statistic.h"
-#include "llvm/Support/Format.h"
-#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Threading.h"
+
 #include <chrono>
 
 using namespace mlir;
 using namespace mlir::detail;
 
-constexpr StringLiteral kPassTimingDescription =
-    "... Pass execution timing report ...";
+//===----------------------------------------------------------------------===//
+// PassTiming
+//===----------------------------------------------------------------------===//
 
 namespace {
-/// Simple record class to record timing information.
-struct TimeRecord {
-  TimeRecord(double wall = 0.0, double user = 0.0) : wall(wall), user(user) {}
-
-  TimeRecord &operator+=(const TimeRecord &other) {
-    wall += other.wall;
-    user += other.user;
-    return *this;
-  }
-
-  /// Print the current time record to 'os', with a breakdown showing
-  /// contributions to the give 'total' time record.
-  void print(raw_ostream &os, const TimeRecord &total) {
-    if (total.user != total.wall)
-      os << llvm::format("  %7.4f (%5.1f%%)  ", user,
-                         100.0 * user / total.user);
-    os << llvm::format("  %7.4f (%5.1f%%)  ", wall, 100.0 * wall / total.wall);
-  }
-
-  double wall, user;
-};
-
-/// An enumeration of the different types of timers.
-enum class TimerKind {
-  /// This timer represents an ordered collection of pass timers, corresponding
-  /// to a pass pipeline.
-  Pipeline,
-
-  /// This timer represents a collection of pipeline timers.
-  PipelineCollection,
-
-  /// This timer represents an analysis or pass timer.
-  PassOrAnalysis
-};
-
-struct Timer {
-  explicit Timer(std::string &&name, TimerKind kind)
-      : name(std::move(name)), kind(kind) {}
-
-  /// Start the timer.
-  void start() { startTime = std::chrono::system_clock::now(); }
-
-  /// Stop the timer.
-  void stop() {
-    auto newTime = std::chrono::system_clock::now() - startTime;
-    wallTime += newTime;
-    userTime += newTime;
-  }
-
-  /// Get or create a child timer with the provided name and id.
-  Timer *getChildTimer(const void *id, TimerKind kind,
-                       std::function<std::string()> &&nameBuilder) {
-    auto &child = children[id];
-    if (!child)
-      child = std::make_unique<Timer>(nameBuilder(), kind);
-    return child.get();
-  }
-
-  /// Returns the total time for this timer in seconds.
-  TimeRecord getTotalTime() {
-    // If this is a pass or analysis timer, use the recorded time directly.
-    if (kind == TimerKind::PassOrAnalysis) {
-      return TimeRecord(
-          std::chrono::duration_cast<std::chrono::duration<double>>(wallTime)
-              .count(),
-          std::chrono::duration_cast<std::chrono::duration<double>>(userTime)
-              .count());
-    }
-
-    // Otherwise, accumulate the timing from each of the children.
-    TimeRecord totalTime;
-    for (auto &child : children)
-      totalTime += child.second->getTotalTime();
-    return totalTime;
-  }
-
-  /// A map of unique identifiers to child timers.
-  using ChildrenMap = llvm::MapVector<const void *, std::unique_ptr<Timer>>;
-
-  /// Merge the timing data from 'other' into this timer.
-  void merge(Timer &&other) {
-    if (wallTime < other.wallTime)
-      wallTime = other.wallTime;
-    userTime += other.userTime;
-    mergeChildren(std::move(other.children));
-  }
-
-  /// Merge the timer children in 'otherChildren' with the children of this
-  /// timer.
-  void mergeChildren(ChildrenMap &&otherChildren) {
-    // Check for an empty children list.
-    if (children.empty()) {
-      children = std::move(otherChildren);
-      return;
-    }
-
-    // Pipeline merges are handled separately as the children are merged
-    // lexicographically.
-    if (kind == TimerKind::Pipeline) {
-      assert(children.size() == otherChildren.size() &&
-             "pipeline merge requires the same number of children");
-      for (auto it : llvm::zip(children, otherChildren))
-        std::get<0>(it).second->merge(std::move(*std::get<1>(it).second));
-      return;
-    }
-
-    // Otherwise, we merge children based upon their timer key.
-    for (auto &otherChild : otherChildren)
-      mergeChild(std::move(otherChild));
-  }
-
-  /// Merge in the given child timer and id into this timer.
-  void mergeChild(ChildrenMap::value_type &&childIt) {
-    auto &child = children[childIt.first];
-    if (!child)
-      child = std::move(childIt.second);
-    else
-      child->merge(std::move(*childIt.second));
-  }
-
-  /// Raw timing information.
-  std::chrono::time_point<std::chrono::system_clock> startTime;
-  std::chrono::nanoseconds wallTime = std::chrono::nanoseconds(0);
-  std::chrono::nanoseconds userTime = std::chrono::nanoseconds(0);
-
-  /// A map of unique identifiers to child timers.
-  ChildrenMap children;
-
-  /// A descriptive name for this timer.
-  std::string name;
-
-  /// The type of timer this instance represents.
-  TimerKind kind;
-};
-
 struct PassTiming : public PassInstrumentation {
-  PassTiming(std::unique_ptr<PassManager::PassTimingConfig> config)
-      : config(std::move(config)) {}
-  ~PassTiming() override { print(); }
+  PassTiming(TimingScope &timingScope) : rootScope(timingScope) {}
 
-  /// Setup the instrumentation hooks.
-  void runBeforePipeline(Identifier name,
-                         const PipelineParentInfo &parentInfo) override;
-  void runAfterPipeline(Identifier name,
-                        const PipelineParentInfo &parentInfo) override;
-  void runBeforePass(Pass *pass, Operation *) override { startPassTimer(pass); }
-  void runAfterPass(Pass *pass, Operation *) override;
-  void runAfterPassFailed(Pass *pass, Operation *op) override {
-    runAfterPass(pass, op);
-  }
-  void runBeforeAnalysis(StringRef name, TypeID id, Operation *) override {
-    startAnalysisTimer(name, id);
-  }
-  void runAfterAnalysis(StringRef, TypeID, Operation *) override;
-
-  /// Print and clear the timing results.
-  void print();
-
-  /// Start a new timer for the given pass.
-  void startPassTimer(Pass *pass);
+  /// 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 root timing scope into which timing is reported.
+  TimingScope &rootScope;
+
+  //===--------------------------------------------------------------------===//
+  // Pipeline
+  //===--------------------------------------------------------------------===//
 
-  /// Start a new timer for the given analysis.
-  void startAnalysisTimer(StringRef name, TypeID id);
-
-  /// Pop the last active timer for the current thread.
-  Timer *popLastActiveTimer() {
+  void runBeforePipeline(Identifier name,
+                         const PipelineParentInfo &parentInfo) override {
     auto tid = llvm::get_threadid();
     auto &activeTimers = activeThreadTimers[tid];
-    assert(!activeTimers.empty() && "expected active timer");
-    return activeTimers.pop_back_val();
-  }
-
-  /// Print the timing result in list mode.
-  void printResultsAsList(raw_ostream &os, Timer *root, TimeRecord totalTime);
-
-  /// Print the timing result in pipeline mode.
-  void printResultsAsPipeline(raw_ostream &os, Timer *root,
-                              TimeRecord totalTime);
 
-  /// Returns a timer for the provided identifier and name.
-  Timer *getTimer(const void *id, TimerKind kind,
-                  std::function<std::string()> &&nameBuilder) {
-    auto tid = llvm::get_threadid();
-
-    // If there is no active timer then add to the root timer.
-    auto &activeTimers = activeThreadTimers[tid];
-    Timer *parentTimer;
+    TimingScope *parentScope;
     if (activeTimers.empty()) {
-      auto &rootTimer = rootTimers[tid];
-      if (!rootTimer)
-        rootTimer = std::make_unique<Timer>("root", TimerKind::Pipeline);
-      parentTimer = rootTimer.get();
+      auto it = parentTimerIndices.find(parentInfo);
+      if (it != parentTimerIndices.end())
+        parentScope =
+            &activeThreadTimers[parentInfo.parentThreadID][it->second];
+      else
+        parentScope = &rootScope;
     } else {
-      // Otherwise, add this to the active timer.
-      parentTimer = activeTimers.back();
+      parentScope = &activeTimers.back();
     }
-
-    auto timer = parentTimer->getChildTimer(id, kind, std::move(nameBuilder));
-    activeTimers.push_back(timer);
-    return timer;
+    activeTimers.push_back(parentScope->nest(name.getAsOpaquePointer(), [name] {
+      return ("'" + name.strref() + "' Pipeline").str();
+    }));
   }
 
-  /// The root top level timers for each thread.
-  DenseMap<uint64_t, std::unique_ptr<Timer>> rootTimers;
-
-  /// A stack of the currently active pass timers per thread.
-  DenseMap<uint64_t, SmallVector<Timer *, 4>> activeThreadTimers;
-
-  /// The configuration object to use when printing the timing results.
-  std::unique_ptr<PassManager::PassTimingConfig> config;
-
-  /// A mapping of pipeline timers that need to be merged into the parent
-  /// collection. The timers are mapped to the parent info to merge into.
-  DenseMap<PipelineParentInfo, SmallVector<Timer::ChildrenMap::value_type, 4>>
-      pipelinesToMerge;
-};
-} // end anonymous namespace
-
-void PassTiming::runBeforePipeline(Identifier name,
-                                   const PipelineParentInfo &parentInfo) {
-  // We don't actually want to time the pipelines, they gather their total
-  // from their held passes.
-  getTimer(name.getAsOpaquePointer(), TimerKind::Pipeline,
-           [&] { return ("'" + name.strref() + "' Pipeline").str(); });
+  void runAfterPipeline(Identifier, const PipelineParentInfo &) override {
+    auto &activeTimers = activeThreadTimers[llvm::get_threadid()];
+    assert(!activeTimers.empty() && "expected active timer");
+    activeTimers.pop_back();
   }
 
-void PassTiming::runAfterPipeline(Identifier name,
-                                  const PipelineParentInfo &parentInfo) {
-  // Pop the timer for the pipeline.
+  //===--------------------------------------------------------------------===//
+  // Pass
+  //===--------------------------------------------------------------------===//
+
+  void runBeforePass(Pass *pass, Operation *) override {
     auto tid = llvm::get_threadid();
     auto &activeTimers = activeThreadTimers[tid];
-  assert(!activeTimers.empty() && "expected active timer");
-  activeTimers.pop_back();
+    auto &parentScope = activeTimers.empty() ? rootScope : activeTimers.back();
 
-  // If the current thread is the same as the parent, there is nothing left to
-  // do.
-  if (tid == parentInfo.parentThreadID)
-    return;
-
-  // Otherwise, mark the pipeline timer for merging into the correct parent
-  // thread.
-  assert(activeTimers.empty() && "expected parent timer to be root");
-  auto *parentTimer = rootTimers[tid].get();
-  assert(parentTimer->children.size() == 1 &&
-         parentTimer->children.count(name.getAsOpaquePointer()) &&
-         "expected a single pipeline timer");
-  pipelinesToMerge[parentInfo].push_back(
-      std::move(*parentTimer->children.begin()));
-  rootTimers.erase(tid);
-}
-
-/// Start a new timer for the given pass.
-void PassTiming::startPassTimer(Pass *pass) {
-  auto kind = isa<OpToOpPassAdaptor>(pass) ? TimerKind::PipelineCollection
-                                           : TimerKind::PassOrAnalysis;
-  Timer *timer = getTimer(pass, kind, [pass]() -> std::string {
-    if (auto *adaptor = dyn_cast<OpToOpPassAdaptor>(pass))
+    if (auto *adaptor = dyn_cast<OpToOpPassAdaptor>(pass)) {
+      parentTimerIndices[{tid, pass}] = activeTimers.size();
+      auto scope = parentScope.nest(pass->getPrototype(), [adaptor]() {
         return adaptor->getAdaptorName();
-    return std::string(pass->getName());
       });
-
-  // We don't actually want to time the adaptor passes, they gather their total
-  // from their held passes.
-  if (!isa<OpToOpPassAdaptor>(pass))
-    timer->start();
-}
-
-/// Start a new timer for the given analysis.
-void PassTiming::startAnalysisTimer(StringRef name, TypeID id) {
-  Timer *timer = getTimer(id.getAsOpaquePointer(), TimerKind::PassOrAnalysis,
-                          [name] { return "(A) " + name.str(); });
-  timer->start();
-}
-
-/// Stop a pass timer.
-void PassTiming::runAfterPass(Pass *pass, Operation *) {
-  Timer *timer = popLastActiveTimer();
-
-  // Check to see if we need to merge in the timing data for the pipelines
-  // running on other threads.
-  auto toMerge = pipelinesToMerge.find({llvm::get_threadid(), pass});
-  if (toMerge != pipelinesToMerge.end()) {
-    for (auto &it : toMerge->second)
-      timer->mergeChild(std::move(it));
-    pipelinesToMerge.erase(toMerge);
+      scope.setTransparent(adaptor->getPassManagers().size() <= 1);
+      activeTimers.push_back(std::move(scope));
+    } else {
+      activeTimers.push_back(parentScope.nest(pass->getPrototype(), [pass]() {
+        return std::string(pass->getName());
+      }));
     }
-
-  timer->stop();
   }
 
-/// Stop a timer.
-void PassTiming::runAfterAnalysis(StringRef, TypeID, Operation *) {
-  popLastActiveTimer()->stop();
+  void runAfterPass(Pass *pass, Operation *) override {
+    auto tid = llvm::get_threadid();
+    if (isa<OpToOpPassAdaptor>(pass)) {
+      parentTimerIndices.erase({tid, pass});
     }
-
-/// Utility to print the timer heading information.
-static void printTimerHeader(raw_ostream &os, TimeRecord total) {
-  os << "===" << std::string(73, '-') << "===\n";
-  // Figure out how many spaces to description name.
-  unsigned padding = (80 - kPassTimingDescription.size()) / 2;
-  os.indent(padding) << kPassTimingDescription << '\n';
-  os << "===" << std::string(73, '-') << "===\n";
-
-  // Print the total time followed by the section headers.
-  os << llvm::format("  Total Execution Time: %5.4f seconds\n\n", total.wall);
-  if (total.user != total.wall)
-    os << "   ---User Time---";
-  os << "   ---Wall Time---  --- Name ---\n";
+    auto &activeTimers = activeThreadTimers[tid];
+    assert(!activeTimers.empty() && "expected active timer");
+    activeTimers.pop_back();
   }
 
-/// Utility to print a single line entry in the timer output.
-static void printTimeEntry(raw_ostream &os, unsigned indent, StringRef name,
-                           TimeRecord time, TimeRecord totalTime) {
-  time.print(os, totalTime);
-  os.indent(indent) << name << "\n";
+  void runAfterPassFailed(Pass *pass, Operation *op) override {
+    runAfterPass(pass, op);
   }
 
-/// Print out the current timing information.
-void PassTiming::print() {
-  // Don't print anything if there is no timing data.
-  if (rootTimers.empty())
-    return;
-
-  assert(rootTimers.size() == 1 && "expected one remaining root timer");
+  //===--------------------------------------------------------------------===//
+  // Analysis
+  //===--------------------------------------------------------------------===//
 
-  auto printCallback = [&](raw_ostream &os) {
-    auto &rootTimer = rootTimers.begin()->second;
-    // Print the timer header.
-    TimeRecord totalTime = rootTimer->getTotalTime();
-    printTimerHeader(os, totalTime);
-    // Defer to a specialized printer for each display mode.
-    switch (config->getDisplayMode()) {
-    case PassDisplayMode::List:
-      printResultsAsList(os, rootTimer.get(), totalTime);
-      break;
-    case PassDisplayMode::Pipeline:
-      printResultsAsPipeline(os, rootTimer.get(), totalTime);
-      break;
-    }
-    printTimeEntry(os, 0, "Total", totalTime, totalTime);
-    os.flush();
-
-    // Reset root timers.
-    rootTimers.clear();
-    activeThreadTimers.clear();
-  };
-
-  config->printTiming(printCallback);
+  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(); }));
   }
 
-// The default implementation for printTiming uses
-// `llvm::CreateInfoOutputFile()` as stream, it can be overridden by clients
-// to customize the output.
-void PassManager::PassTimingConfig::printTiming(PrintCallbackFn printCallback) {
-  printCallback(*llvm::CreateInfoOutputFile());
-}
-
-/// Print the timing result in list mode.
-void PassTiming::printResultsAsList(raw_ostream &os, Timer *root,
-                                    TimeRecord totalTime) {
-  llvm::StringMap<TimeRecord> mergedTimings;
-
-  std::function<void(Timer *)> addTimer = [&](Timer *timer) {
-    // Only add timing information for passes and analyses.
-    if (timer->kind == TimerKind::PassOrAnalysis)
-      mergedTimings[timer->name] += timer->getTotalTime();
-    for (auto &children : timer->children)
-      addTimer(children.second.get());
-  };
-
-  // Add each of the top level timers.
-  for (auto &topLevelTimer : root->children)
-    addTimer(topLevelTimer.second.get());
-
-  // Sort the timing information by wall time.
-  std::vector<std::pair<StringRef, TimeRecord>> timerNameAndTime;
-  for (auto &it : mergedTimings)
-    timerNameAndTime.emplace_back(it.first(), it.second);
-  llvm::array_pod_sort(timerNameAndTime.begin(), timerNameAndTime.end(),
-                       [](const std::pair<StringRef, TimeRecord> *lhs,
-                          const std::pair<StringRef, TimeRecord> *rhs) {
-                         return llvm::array_pod_sort_comparator<double>(
-                             &rhs->second.wall, &lhs->second.wall);
-                       });
-
-  // Print the timing information sequentially.
-  for (auto &timeData : timerNameAndTime)
-    printTimeEntry(os, 0, timeData.first, timeData.second, totalTime);
-}
-
-/// Print the timing result in pipeline mode.
-void PassTiming::printResultsAsPipeline(raw_ostream &os, Timer *root,
-                                        TimeRecord totalTime) {
-  std::function<void(unsigned, Timer *)> printTimer = [&](unsigned indent,
-                                                          Timer *timer) {
-    // If this is a timer for a pipeline collection and the collection only has
-    // one pipeline child, then only print the child.
-    if (timer->kind == TimerKind::PipelineCollection &&
-        timer->children.size() == 1)
-      return printTimer(indent, timer->children.begin()->second.get());
-
-    printTimeEntry(os, indent, timer->name, timer->getTotalTime(), totalTime);
-
-    // If this timer is a pipeline, then print the children in-order.
-    if (timer->kind == TimerKind::Pipeline) {
-      for (auto &child : timer->children)
-        printTimer(indent + 2, child.second.get());
-      return;
-    }
-
-    // Otherwise, sort the children by name to give a deterministic ordering
-    // when emitting the time.
-    SmallVector<Timer *, 4> children;
-    children.reserve(timer->children.size());
-    for (auto &child : timer->children)
-      children.push_back(child.second.get());
-    llvm::array_pod_sort(children.begin(), children.end(),
-                         [](Timer *const *lhs, Timer *const *rhs) {
-                           return (*lhs)->name.compare((*rhs)->name);
-                         });
-    for (auto &child : children)
-      printTimer(indent + 2, child);
-  };
-
-  // Print each of the top level timers.
-  for (auto &topLevelTimer : root->children)
-    printTimer(0, topLevelTimer.second.get());
-}
-
-// Out-of-line as key function.
-PassManager::PassTimingConfig::~PassTimingConfig() {}
+  void runAfterAnalysis(StringRef, TypeID, Operation *) override {
+    auto &activeTimers = activeThreadTimers[llvm::get_threadid()];
+    assert(!activeTimers.empty() && "expected active timer");
+    activeTimers.pop_back();
+} // namespace
 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//
 
 /// Add an instrumentation to time the execution of passes and the computation
 /// of analyses.
 void PassManager::enableTiming(std::unique_ptr<PassTimingConfig> config) {
-  // Check if pass timing is already enabled.
-  if (passTiming)
+  if (passTiming || ownedTimingManager)
+    return;
+  ownedTimingManager = std::make_unique<TimingManager>();
+  ownedTimingManager->enableTiming(std::move(config));
+  ownedTimingScope = ownedTimingManager->getRootScope();
+  enableTiming(ownedTimingScope);
+}
+
+/// Add an instrumentation to time the execution of passes and the computation
+/// of analyses.
+void PassManager::enableTiming(TimingScope &timingScope) {
+  if (passTiming || !timingScope)
     return;
-  if (!config)
-    config = std::make_unique<PassManager::PassTimingConfig>();
-  addInstrumentation(std::make_unique<PassTiming>(std::move(config)));
+  addInstrumentation(std::make_unique<PassTiming>(timingScope));
   passTiming = true;
 }

mlir/lib/Support/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   DebugCounter.cpp
   FileUtilities.cpp
   IndentedOstream.cpp
   MlirOptMain.cpp
   StorageUniquer.cpp
+  Timing.cpp
   ToolUtilities.cpp
 )
 
 add_mlir_library(MLIRSupport
   DebugCounter.cpp
   FileUtilities.cpp
   StorageUniquer.cpp
+  Timing.cpp
   ToolUtilities.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Support
 
   LINK_LIBS PUBLIC
   ${LLVM_PTHREAD_LIB})
 

mlir/lib/Support/Timing.cpp

+//===- Timing.cpp - Execution time measurement facilities -----------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Facilities to measure and provide statistics on execution time.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Support/Timing.h"
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Format.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/ManagedStatic.h"
+#include "llvm/Support/Threading.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <atomic>
+#include <chrono>
+
+using namespace mlir;
+using llvm::StringRef;
+
+constexpr llvm::StringLiteral kTimingDescription =
+    "... Execution time report ...";
+
+//===----------------------------------------------------------------------===//
+// TimingConfig
+//===----------------------------------------------------------------------===//
+
+// Out-of-line as key function.
+TimingConfig::~TimingConfig() {}
+
+// The default implementation of the printing hook.
+void TimingConfig::printTiming(PrintCallbackFn printCallback) {
+  printCallback(*llvm::CreateInfoOutputFile());
+}
+
+//===----------------------------------------------------------------------===//
+// Timer
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// Simple record class to record timing information.
+struct TimeRecord {
+  TimeRecord(double wall = 0.0, double user = 0.0) : wall(wall), user(user) {}
+
+  TimeRecord &operator+=(const TimeRecord &other) {
+    wall += other.wall;
+    user += other.user;
+    return *this;
+  }
+
+  TimeRecord &operator-=(const TimeRecord &other) {
+    wall -= other.wall;
+    user -= other.user;
+    return *this;
+  }
+
+  /// Print the current time record to 'os', with a breakdown showing
+  /// contributions to the give 'total' time record.
+  void print(llvm::raw_ostream &os, const TimeRecord &total) {
+    if (total.user != total.wall)
+      os << llvm::format("  %8.4f (%5.1f%%)", user, 100.0 * user / total.user);
+    os << llvm::format("  %8.4f (%5.1f%%)  ", wall, 100.0 * wall / total.wall);
+  }
+
+  double wall, user;
+};
+
+} // namespace
+
+/// Utility to print a single line entry in the timer output.
+static void printTimeEntry(llvm::raw_ostream &os, unsigned indent,
+                           StringRef name, TimeRecord time, TimeRecord total) {
+  time.print(os, total);
+  os.indent(indent) << name << "\n";
+}
+
+/// Utility to print the timer heading information.
+static void printTimeHeader(llvm::raw_ostream &os, TimeRecord total) {
+  // Figure out how many spaces to description name.
+  unsigned padding = (80 - kTimingDescription.size()) / 2;
+  os << "===" << std::string(73, '-') << "===\n";
+  os.indent(padding) << kTimingDescription << '\n';
+  os << "===" << std::string(73, '-') << "===\n";
+
+  // Print the total time followed by the section headers.
+  os << llvm::format("  Total Execution Time: %.4f seconds\n\n", total.wall);
+  if (total.user != total.wall)
+    os << "  ----User Time----";
+  os << "  ----Wall Time----  ----Name----\n";
+}
+
+/// A timer used to sample execution time.
+///
+/// Separately tracks wall time and user time to account for parallel threads of
+/// execution. Timers are intended to be started and stopped multiple times.
+/// Each start and stop will add to the timer's wall and user time.
+struct TimingManager::Timer : public TimerBase {
+  using ChildrenMap = llvm::MapVector<const void *, std::unique_ptr<Timer>>;
+  using AsyncChildrenMap = llvm::DenseMap<uint64_t, ChildrenMap>;
+
+  Timer(std::string &&name) : threadId(llvm::get_threadid()), name(name) {}
+
+  /// Start the timer.
+  virtual void start() override {
+    startTime = std::chrono::system_clock::now();
+  }
+
+  /// Stop the timer.
+  virtual void stop() override {
+    auto newTime = std::chrono::system_clock::now() - startTime;
+    wallTime += newTime;
+    userTime += newTime;
+  }
+
+  /// Hide the timer in timing reports and directly show its children.
+  virtual void setTransparent(bool x) override { transparent = x; }
+
+  /// Create a child timer nested within this one. Multiple calls to this
+  /// function with the same unique identifier `id` will return the same child
+  /// timer.
+  ///
+  /// This function can be called from other threads, as long as this timer
+  /// outlives any uses of the child timer on the other thread.
+  virtual TimerBase *nest(const void *id,
+                          std::function<std::string()> &&nameBuilder) override {
+    auto tid = llvm::get_threadid();
+    if (tid == threadId) {
+      return nestTail(children[id], std::move(nameBuilder));
+    } else {
+      std::unique_lock<std::mutex> lock(asyncMutex);
+      return nestTail(asyncChildren[tid][id], std::move(nameBuilder));
+    }
+  }
+
+  /// Tail-called from `nest()`.
+  TimerBase *nestTail(std::unique_ptr<Timer> &child,
+                      std::function<std::string()> &&nameBuilder) {
+    if (!child)
+      child = std::make_unique<Timer>(nameBuilder());
+    return child.get();
+  }
+
+  /// Finalize this timer and all its children.
+  ///
+  /// If this timer has async children, which happens if `nest()` was called
+  /// from another thread, this function merges the async childr timers into the
+  /// main list of child timers.
+  ///
+  /// Caution: Call this function only after all nested timers running on other
+  /// threads no longer need their timers!
+  void finalize() {
+    addAsyncUserTime();
+    mergeAsyncChildren();
+  }
+
+  /// Add the user time of all async children to this timer's user time. This is
+  /// necessary since the user time already contains all regular child timers,
+  /// but not the asynchronous ones (by the nesting nature of the timers).
+  std::chrono::nanoseconds addAsyncUserTime() {
+    auto added = std::chrono::nanoseconds(0);
+    for (auto &child : children) {
+      added += child.second->addAsyncUserTime();
+    }
+    for (auto &thread : asyncChildren) {
+      for (auto &child : thread.second) {
+        child.second->addAsyncUserTime();
+        added += child.second->userTime;
+      }
+    }
+    userTime += added;
+    return added;
+  }
+
+  /// Ensure that this timer and recursively all its children have their async
+  /// children folded into the main map of children.
+  void mergeAsyncChildren() {
+    for (auto &child : children)
+      child.second->mergeAsyncChildren();
+    mergeChildren(std::move(asyncChildren));
+    assert(asyncChildren.empty());
+  }
+
+  /// Merge multiple child timers into this timer.
+  ///
+  /// Children in `other` are added as children to this timer, or, if this timer
+  /// already contains a child with the corresponding unique identifier, are
+  /// merged into the existing child.
+  void mergeChildren(ChildrenMap &&other) {
+    if (children.empty()) {
+      children = std::move(other);
+      for (auto &child : other)
+        child.second->mergeAsyncChildren();
+    } else {
+      for (auto &child : other)
+        mergeChild(child.first, std::move(child.second));
+      other.clear();
+    }
+  }
+
+  /// See above.
+  void mergeChildren(AsyncChildrenMap &&other) {
+    for (auto &thread : other) {
+      mergeChildren(std::move(thread.second));
+      assert(thread.second.empty());
+    }
+    other.clear();
+  }
+
+  /// Merge a child timer into this timer for a given unique identifier.
+  ///
+  /// Moves all child and async child timers of `other` into this timer's child
+  /// for the given unique identifier.
+  void mergeChild(const void *id, std::unique_ptr<Timer> &&other) {
+    auto &into = children[id];
+    if (!into) {
+      into = std::move(other);
+      into->mergeAsyncChildren();
+    } else {
+      into->wallTime = std::max(into->wallTime, other->wallTime);
+      into->userTime += other->userTime;
+      into->mergeChildren(std::move(other->children));
+      into->mergeChildren(std::move(other->asyncChildren));
+      other.reset();
+    }
+  }
+
+  /// Dump a human-readable tree representation of the timer and its children.
+  /// This is useful for debugging the timing mechanisms and structure of the
+  /// timers.
+  void dump(llvm::raw_ostream &os, unsigned indent = 0,
+            unsigned markThreadId = 0) {
+    auto time = getTimeRecord();
+    os << std::string(indent * 2, ' ') << name << " [" << threadId << "]"
+       << llvm::format("  %7.4f / %7.4f", time.user, time.wall);
+    if (threadId != markThreadId && markThreadId != 0)
+      os << " (*)";
+    os << "\n";
+    for (auto &child : children)
+      child.second->dump(os, indent + 1, threadId);
+    for (auto &thread : asyncChildren)
+      for (auto &child : thread.second)
+        child.second->dump(os, indent + 1, threadId);
+  }
+
+  /// Returns the time for this timer in seconds.
+  TimeRecord getTimeRecord() {
+    return TimeRecord(
+        std::chrono::duration_cast<std::chrono::duration<double>>(wallTime)
+            .count(),
+        std::chrono::duration_cast<std::chrono::duration<double>>(userTime)
+            .count());
+  }
+
+  /// Print the timing result in list mode.
+  void printAsList(llvm::raw_ostream &os, TimeRecord total) {
+    // Flatten the leaf timers in the tree and merge them by name.
+    llvm::StringMap<TimeRecord> mergedTimers;
+    std::function<void(Timer *)> addTimer = [&](Timer *timer) {
+      if (timer->children.empty())
+        mergedTimers[timer->name] += timer->getTimeRecord();
+      else
+        for (auto &children : timer->children)
+          addTimer(children.second.get());
+    };
+    addTimer(this);
+
+    // Sort the timing information by wall time.
+    std::vector<std::pair<StringRef, TimeRecord>> timerNameAndTime;
+    for (auto &it : mergedTimers)
+      timerNameAndTime.emplace_back(it.first(), it.second);
+    llvm::array_pod_sort(timerNameAndTime.begin(), timerNameAndTime.end(),
+                         [](const std::pair<StringRef, TimeRecord> *lhs,
+                            const std::pair<StringRef, TimeRecord> *rhs) {
+                           return llvm::array_pod_sort_comparator<double>(
+                               &rhs->second.wall, &lhs->second.wall);
+                         });
+
+    // Print the timing information sequentially.
+    for (auto &timeData : timerNameAndTime)
+      printTimeEntry(os, 0, timeData.first, timeData.second, total);
+  }
+
+  /// Print the timing result in pipeline mode.
+  void printAsPipeline(llvm::raw_ostream &os, TimeRecord total,
+                       unsigned indent = 0) {
+    unsigned childIndent = indent;
+    if (!transparent) {
+      printTimeEntry(os, indent, name, getTimeRecord(), total);
+      childIndent += 2;
+    }
+    for (auto &child : children) {
+      child.second->printAsPipeline(os, total, childIndent);
+    }
+  }
+
+  /// Print out the current timing information.
+  void print(llvm::raw_ostream &os, TimingDisplayMode displayMode) {
+    // Print the banner.
+    auto total = getTimeRecord();
+    printTimeHeader(os, total);
+
+    // Defer to a specialized printer for each display mode.
+    switch (displayMode) {
+    case TimingDisplayMode::List:
+      printAsList(os, total);
+      break;
+    case TimingDisplayMode::Pipeline:
+      printAsPipeline(os, total);
+      break;
+    default:
+      llvm_unreachable("unhandled display mode");
+    }
+
+    // Print the top-level time not accounted for by child timers, and the
+    // total.
+    auto rest = total;
+    for (auto &child : children)
+      rest -= child.second->getTimeRecord();
+    printTimeEntry(os, 0, "Rest", rest, total);
+    printTimeEntry(os, 0, "Total", total, total);
+    os.flush();
+  }
+
+  /// The last time instant at which the timer was started.
+  std::chrono::time_point<std::chrono::system_clock> startTime;
+
+  /// Accumulated wall time. If multiple threads of execution are merged into
+  /// this timer, the wall time will hold the maximum wall time of each thread
+  /// of execution.
+  std::chrono::nanoseconds wallTime = std::chrono::nanoseconds(0);
+
+  /// Accumulated user time. If multiple threads of execution are merged into
+  /// this timer, each thread's user time is added here.
+  std::chrono::nanoseconds userTime = std::chrono::nanoseconds(0);
+
+  /// The thread on which this timer is running.
+  uint64_t threadId;
+
+  /// A descriptive name for this timer.
+  std::string name;
+
+  /// Whether omit this timer from reports and directly show its children.
+  bool transparent = false;
+
+  /// Child timers on the same thread the timer itself. We keep at most one
+  /// timer per unique identifier.
+  ChildrenMap children;
+
+  /// Child timers on other threads. We keep at most one timer per unique
+  /// identifier.
+  AsyncChildrenMap asyncChildren;
+
+  /// Mutex for the async children.
+  std::mutex asyncMutex;
+};
+
+//===----------------------------------------------------------------------===//
+// TimingManager
+//===----------------------------------------------------------------------===//
+
+TimingManager::TimingManager() {
+  clear(); // initializes the root timer
+}
+
+TimingManager::~TimingManager() { print(); }
+
+TimerBase *TimingManager::getRootTimer() {
+  return config ? rootTimer.get() : nullptr;
+}
+
+/// Enable execution time sampling.
+void TimingManager::enableTiming(std::unique_ptr<TimingConfig> newConfig) {
+  if (!newConfig)
+    newConfig = std::make_unique<TimingConfig>();
+  config = std::move(newConfig);
+}
+
+/// Print and clear the timing results.
+void TimingManager::print() {
+  if (config) {
+    rootTimer->finalize();
+    config->printTiming([&](llvm::raw_ostream &os) {
+      rootTimer->print(os, config->getDisplayMode());
+    });
+  }
+  clear();
+}
+
+/// Clear the timing results.
+void TimingManager::clear() {
+  rootTimer = std::make_unique<Timer>("root");
+  rootTimer->transparent = true;
+}
+
+/// Debug print the timer data structures to an output stream.
+void TimingManager::dumpTimers(llvm::raw_ostream &os) { rootTimer->dump(os); }
+
+/// Debug print the timers as a list.
+void TimingManager::dumpAsList(llvm::raw_ostream &os) {
+  rootTimer->finalize();
+  rootTimer->print(os, TimingDisplayMode::List);
+}
+
+/// Debug print the timers as a pipeline.
+void TimingManager::dumpAsPipeline(llvm::raw_ostream &os) {
+  rootTimer->finalize();
+  rootTimer->print(os, TimingDisplayMode::Pipeline);
+}
+
+//===----------------------------------------------------------------------===//
+// TimingManagerOptions
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct TimingManagerOptions {
+  llvm::cl::opt<bool> timing{"timing",
+                             llvm::cl::desc("Display execution times")};
+  llvm::cl::opt<TimingDisplayMode> timingDisplayMode{
+      "timing-display", llvm::cl::desc("Display method for timing data"),
+      llvm::cl::init(TimingDisplayMode::Pipeline),
+      llvm::cl::values(
+          clEnumValN(TimingDisplayMode::List, "list",
+                     "display the results in a list sorted by total time"),
+          clEnumValN(TimingDisplayMode::Pipeline, "pipeline",
+                     "display the results with a nested pipeline view"))};
+};
+} // end anonymous namespace
+
+static llvm::ManagedStatic<TimingManagerOptions> options;
+
+void mlir::registerTimingManagerCLOptions() {
+  // Make sure that the options struct has been constructed.
+  *options;
+}
+
+void mlir::applyTimingManagerCLOptions(TimingManagerBase &tm) {
+  if (!options.isConstructed())
+    return;
+  if (options->timing)
+    tm.enableTiming(std::make_unique<TimingConfig>(options->timingDisplayMode));
+}