Make MLIR Pass Timing output configurable through injection

[mehdi_amini]

Summary

This makes it possible for the client to control where the pass timings will
be printed.

[rriddle]

mlir/include/mlir/Pass/PassManager.h

235

Do we need to pass the display mode back? Is there a situation where it would differ from the config?

[mehdi_amini]

Remove displayMode from callback arguments

[mehdi_amini]

(actually do it)

[rriddle]

mlir/lib/Pass/PassTiming.cpp

235

Is this actually optional?

mlir/include/mlir/Pass/PassManager.h

public:
   void enableIRPrinting(
       std::function<bool(Pass *, Operation *)> shouldPrintBeforePass,
       std::function<bool(Pass *, Operation *)> shouldPrintAfterPass,
       bool printModuleScope, bool printAfterOnlyOnChange, raw_ostream &out);
 
   //===--------------------------------------------------------------------===//
   // Pass Timing
 
+  /// A configuration struct provided to the pass timing feature.
+  class PassTimingConfig {
+  public:
+    using PrintCallbackFn = function_ref<void(raw_ostream &)>;
+
+    /// 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.
   /// 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(PassDisplayMode displayMode = PassDisplayMode::Pipeline);
+  void enableTiming(std::unique_ptr<PassTimingConfig> config = nullptr);
 
   /// 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.

mlir/lib/Pass/PassManagerOptions.cpp

void PassManagerOptions::addPrinterInstrumentation(PassManager &pm) {
   // Otherwise, add the IR printing instrumentation.
   pm.enableIRPrinting(shouldPrintBeforePass, shouldPrintAfterPass,
                       printModuleScope, printAfterChange, llvm::errs());
 }
 
 /// Add a pass timing instrumentation if enabled by 'pass-timing' flags.
 void PassManagerOptions::addTimingInstrumentation(PassManager &pm) {
   if (passTiming)
-    pm.enableTiming(passTimingDisplayMode);
+    pm.enableTiming(
+        std::make_unique<PassManager::PassTimingConfig>(passTimingDisplayMode));
 }
 
 void mlir::registerPassManagerCLOptions() {
   // Make sure that the options struct has been constructed.
   *options;
 }
 
 void mlir::applyPassManagerCLOptions(PassManager &pm) {

mlir/lib/Pass/PassTiming.cpp

struct Timer {
   /// A descriptive name for this timer.
   std::string name;
 
   /// The type of timer this instance represents.
   TimerKind kind;
 };
 
 struct PassTiming : public PassInstrumentation {
-  PassTiming(PassDisplayMode displayMode) : displayMode(displayMode) {}
+  PassTiming(std::unique_ptr<PassManager::PassTimingConfig> config)
+      : config(std::move(config)) {}
   ~PassTiming() override { print(); }
 
   /// Setup the instrumentation hooks.
   void runBeforePipeline(const OperationName &name,
                          const PipelineParentInfo &parentInfo) override;
   void runAfterPipeline(const OperationName &name,
                         const PipelineParentInfo &parentInfo) override;
   void runBeforePass(Pass *pass, Operation *) override { startPassTimer(pass); }
struct PassTiming : public PassInstrumentation {
   }
 
   /// 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 display mode to use when printing the timing results.
-  PassDisplayMode displayMode;
+  /// The (optional) 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
 

 
 /// 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");
-  auto &rootTimer = rootTimers.begin()->second;
-  auto os = llvm::CreateInfoOutputFile();
 
+  auto printCallback = [&](raw_ostream &os) {
+    auto &rootTimer = rootTimers.begin()->second;
     // Print the timer header.
     TimeRecord totalTime = rootTimer->getTotalTime();
-  printTimerHeader(*os, totalTime);
-
+    printTimerHeader(os, totalTime);
     // Defer to a specialized printer for each display mode.
-  switch (displayMode) {
+    switch (config->getDisplayMode()) {
     case PassDisplayMode::List:
-    printResultsAsList(*os, rootTimer.get(), totalTime);
+      printResultsAsList(os, rootTimer.get(), totalTime);
       break;
     case PassDisplayMode::Pipeline:
-    printResultsAsPipeline(*os, rootTimer.get(), totalTime);
+      printResultsAsPipeline(os, rootTimer.get(), totalTime);
       break;
     }
-  printTimeEntry(*os, 0, "Total", totalTime, totalTime);
-  os->flush();
+    printTimeEntry(os, 0, "Total", totalTime, totalTime);
+    os.flush();
 
     // Reset root timers.
     rootTimers.clear();
     activeThreadTimers.clear();
+  };
+
+  config->printTiming(printCallback);
+}
+
+// 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) {
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() {}
+
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//
 
 /// Add an instrumentation to time the execution of passes and the computation
 /// of analyses.
-void PassManager::enableTiming(PassDisplayMode displayMode) {
+void PassManager::enableTiming(std::unique_ptr<PassTimingConfig> config) {
   // Check if pass timing is already enabled.
   if (passTiming)
     return;
-  addInstrumentation(std::make_unique<PassTiming>(displayMode));
+  if (!config)
+    config = std::make_unique<PassManager::PassTimingConfig>();
+  addInstrumentation(std::make_unique<PassTiming>(std::move(config)));
   passTiming = true;
 }