[IR] Reimplement time tracing of NewPassManager by PassInstrumentation framework

[dongjunduo]

Summary

This patch is a enhancement of D74516.

The previous implementation of time tracing in NewPassManager is direct but messive.

The key codes are like the demo below:

 /// Runs the function pass across every function in the module.
 PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM,
                       LazyCallGraph &CG, CGSCCUpdateResult &UR) {
     /// ...
     PreservedAnalyses PassPA;
     {
       TimeTraceScope TimeScope(Pass.name());
       PassPA = Pass.run(F, FAM);
     }
     /// ...
}

It can be bothered to judge where should we add the tracing codes by hands.

With the PassInstrumentation framework, we can easily add Before/After callback
functions to add time tracing codes.

Thus, the implementation of that can be clearer than before by this mechanism.

Further more, we need not to add time tracing codes by hands when new types of PassManager are added.

Through this improved method, the time tracing results of pass execution time now are approximately equal to the previous one.

Demo:

{   // ... (Previous)
    "dur": 31,
    "name": "Total AnnotationRemarksPass",
    // ...
},
{
    // ... (Now)
    "dur": 50,
    "name": "Total AnnotationRemarksPass",
    // ...
},
{
    // ... (Previous)
    "dur": 840,
    "name": "Total AlwaysInlinerPass",
    // ...
},
{
    // ... (Now)
    "dur": 725,
    "name": "Total AlwaysInlinerPass",
    // ...
}

[myhsu]

I don't think this is Clang-specific, could you update the title?

[jamieschmeiser]

llvm/lib/IR/PassTimeProfiling.cpp

42 ↗

(On Diff #452973)

This looks like it is copied with changes from Passes/StandardInstructions.cpp. You should almost never replicate code for several reasons. 1) If there is a bug in the code, it probably won't be fixed in all the copies; 2) code bloat 3) It is confusing because it is unexpected. If there are only minor changes, it is difficult to recognize the changes. I see that the loop case is different. Perhaps a parameter to control what happens for a loop? Why not put the new functionality in Passes/StandardInstructions.cpp?

[jamieschmeiser]

Considering that getIRName is only used for identifying the IR, do you even need the change?

[dongjunduo]

Move class PassTimeProfiling to StandardInstrumentation.cpp

[jamieschmeiser]

You want the timing extensions to be the last ones called before the pass and the first ones at the end of the pass so that the overhead of the other callbacks is not included in the timings. Does PassInstrumentationCallbacks have the ability to specify that the pass be registered as the first callback? If not, add a bool parameter that allows it to be placed at the front with a default of false, and then put all timing extensions at the end, with the parameter being set to true for the registering of the atEnd callback (to put them as the first ones called). Also, add a comment that the timing ones should be last for this reason.

llvm/lib/Passes/StandardInstrumentations.cpp

1165

Rather than testing if the option is set before and after each pass, test in registerCallBacks and only register when getTimeTraceProfilerInstance() returns a pointer.

[dongjunduo]

Adjust the location of TimeProfiler's start and end functions

[myhsu]

llvm/include/llvm/Passes/StandardInstrumentations.h

409

Could we use /// here to get better doxygen appearances?

413

nit: can we move private members here (or using struct instead)

llvm/lib/Passes/StandardInstrumentations.cpp

1151

nit: personally I would simply rename the local PIC rather than spelling out this

1175

ditto the one above

[dongjunduo]

Restyle the callback registration and related comment

[dongjunduo]

llvm/include/llvm/Passes/StandardInstrumentations.h

413

nit: can we move private members here (or using struct instead)

413

That class is designed following the class PrintIRInstrumentation(ref)

I think it's best to keep the class structure consistent. : )

llvm/lib/Passes/StandardInstrumentations.cpp

1151

dtto the one above.

[dongjunduo]

llvm/lib/Passes/StandardInstrumentations.cpp

1151

@jamieschmeiser
The callback functions of TimeProfiler can be registered if and only if "-ftime-trace" is given. : )

[dongjunduo]

Restyle codes

[jamieschmeiser]

llvm/include/llvm/Passes/StandardInstrumentations.h

422

Why are you saving PIC as a member? You do not refer to it other than initializing it. Please remove it.

[jamieschmeiser]

llvm/lib/Passes/StandardInstrumentations.cpp

1151

1153

This is unnecessary (see comment in header)

[dongjunduo]

Remove the useless code

[jamieschmeiser]

llvm/lib/Passes/StandardInstrumentations.cpp

2093

This comment should be expanded to indicate why it must be the last instrumentation (ie, so that it is the last instrumentation before the pass and that this also ensures that the after routines are truly put at the front so that other instrumentations are not included in the timing).

[dongjunduo]

Add more comments

[jamieschmeiser]

A couple of nits left: small edits to comments and the change of "...() == nullptr" to "!...()" and it is good to go.
LGTM.

llvm/lib/Passes/StandardInstrumentations.cpp

2095

2096

[dongjunduo]

Restyle codes and comments

[jamieschmeiser]

As discussed, change your code to use the name of the loop or call getIRName and add a FIXME comment to resolve the 2 pieces of code. https://reviews.llvm.org/D133587 has been created to fix the generation of the loop names and fix the lit tests.

[jamieschmeiser]

I have delivered https://reviews.llvm.org/D133587 so you should be good to go without the FIXME or changes.

[aeubanks]

nice cleanup, one nit

llvm/lib/Passes/StandardInstrumentations.cpp

1159

I'd rename the param to InsertAtFront and here it should be /*InsertAtFront=*/ true. We do this for params where we pass constants to make it harder to mess up if we add more bool params in the future, and just for reading clarity.

[aeubanks]

(sorry to post here, not sure where else to post this in regards to the gsoc project I just learned about) somebody had brought up the idea of noting which passes don't make any changes to the IR. if we could somehow filter by passes that haven't made any changes that could be insightful

[jamieschmeiser]

@aeubanks Not sure what you mean... something like print-changed=no-change? That would different from this work but we can continue the discussion.

[jamieschmeiser]

or do you mean noting which passes don't make changes and calculating the time spent in those passes? That would combine this work with the print-changed type filters...

[aeubanks]

In D131960#3781325, @jamieschmeiser wrote:

@aeubanks Not sure what you mean... something like print-changed=no-change? That would different from this work but we can continue the discussion.

Nothing to do with printing, I meant somehow augmenting TimeTrace filter out passes that actually transform the IR (since they're actually doing something) and only report passes that made no change (wasting compile time). That could be useful for investigating if there's any way to avoid running passes.

and yeah I'm not asking for any work here, just a random thought if you have time

[jamieschmeiser]

@aeubanks An interesting idea. I'll give it some thought... it should be doable by combining this with the logic in the change printers but I'd be concerned that the detection of the changes would become a hotspot itself. Perhaps a double run, the first to determine which passes make changes and then a recompile to time those passes that are then known to make changes.

[aeubanks]

In D131960#3781458, @jamieschmeiser wrote:

@aeubanks An interesting idea. I'll give it some thought... it should be doable by combining this with the logic in the change printers but I'd be concerned that the detection of the changes would become a hotspot itself. Perhaps a double run, the first to determine which passes make changes and then a recompile to time those passes that are then known to make changes.

Ideally we'd have some sort of hashing mechanism to hash functions rather than relying on printing the IR. But printing seems like a fine first step.
If we can augment the time trace to have some way to mark the last data point with some flag, we can stop the timer, determine if IR changed, then mark the data point accordingly and we shouldn't need two runs. Two runs seems like too much hassle.

[jamieschmeiser]

I agree. We could time the pass, figure out whether it changed anything (outside of the timing) and then decide how to save the timing info for the pass.

[aeubanks]

this actually slightly improved compile times, probably by not calculating pass names when we aren't using time tracing. thanks again!

https://llvm-compile-time-tracker.com/compare.php?from=69d9bb2aade39c76de9d6da9e2383377852e4d73&to=6975ab71260c79ddc7a616814678913e67ea417c&stat=instructions

llvm/include/llvm/IR/PassInstrumentation.h

void registerBeforeSkippedPassCallback(CallableT C) {
     BeforeSkippedPassCallbacks.emplace_back(std::move(C));
   }
 
   template <typename CallableT>
   void registerBeforeNonSkippedPassCallback(CallableT C) {
     BeforeNonSkippedPassCallbacks.emplace_back(std::move(C));
   }
 
-  template <typename CallableT> void registerAfterPassCallback(CallableT C) {
+  template <typename CallableT>
+  void registerAfterPassCallback(CallableT C, bool ToFront = false) {
+    if (ToFront)
+      AfterPassCallbacks.insert(AfterPassCallbacks.begin(), std::move(C));
+    else
       AfterPassCallbacks.emplace_back(std::move(C));
   }
 
   template <typename CallableT>
-  void registerAfterPassInvalidatedCallback(CallableT C) {
+  void registerAfterPassInvalidatedCallback(CallableT C, bool ToFront = false) {
+    if (ToFront)
+      AfterPassInvalidatedCallbacks.insert(
+          AfterPassInvalidatedCallbacks.begin(), std::move(C));
+    else
       AfterPassInvalidatedCallbacks.emplace_back(std::move(C));
   }
 
   template <typename CallableT>
   void registerBeforeAnalysisCallback(CallableT C) {
     BeforeAnalysisCallbacks.emplace_back(std::move(C));
   }
 
   template <typename CallableT>
-  void registerAfterAnalysisCallback(CallableT C) {
+  void registerAfterAnalysisCallback(CallableT C, bool ToFront = false) {
+    if (ToFront)
+      AfterAnalysisCallbacks.insert(AfterAnalysisCallbacks.begin(),
+                                    std::move(C));
+    else
       AfterAnalysisCallbacks.emplace_back(std::move(C));
   }
 
   template <typename CallableT>
   void registerAnalysisInvalidatedCallback(CallableT C) {
     AnalysisInvalidatedCallbacks.emplace_back(std::move(C));
   }
 
   template <typename CallableT>

llvm/include/llvm/IR/PassManager.h

PreservedAnalyses run(IRUnitT &IR, AnalysisManagerT &AM,
 
     for (auto &Pass : Passes) {
       // Check the PassInstrumentation's BeforePass callbacks before running the
       // pass, skip its execution completely if asked to (callback returns
       // false).
       if (!PI.runBeforePass<IRUnitT>(*Pass, IR))
         continue;
 
-      PreservedAnalyses PassPA;
-      {
-        TimeTraceScope TimeScope(Pass->name(), IR.getName());
-        PassPA = Pass->run(IR, AM, ExtraArgs...);
-      }
+      PreservedAnalyses PassPA = Pass->run(IR, AM, ExtraArgs...);
 
       // Call onto PassInstrumentation's AfterPass callbacks immediately after
       // running the pass.
       PI.runAfterPass<IRUnitT>(*Pass, IR, PassPA);
 
       // Update the analysis manager as each pass runs and potentially
       // invalidates analyses.
       AM.invalidate(IR, PassPA);

llvm/include/llvm/Passes/StandardInstrumentations.h

 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/OptBisect.h"
 #include "llvm/IR/PassTimingInfo.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/TimeProfiler.h"
 #include "llvm/Transforms/IPO/SampleProfileProbe.h"
 
 #include <string>
 #include <utility>
 
 namespace llvm {
 
 class Module;

 class VerifyInstrumentation {
   bool DebugLogging;
 
 public:
   VerifyInstrumentation(bool DebugLogging) : DebugLogging(DebugLogging) {}
   void registerCallbacks(PassInstrumentationCallbacks &PIC);
 };
 
+/// This class implements --time-trace functionality for new pass manager.
+/// It provides the pass-instrumentation callbacks that measure the pass
+/// execution time. They collect time tracing info by TimeProfiler.
+class TimeProfilingPassesHandler {
+public:
+  TimeProfilingPassesHandler();
+  // We intend this to be unique per-compilation, thus no copies.
+  TimeProfilingPassesHandler(const TimeProfilingPassesHandler &) = delete;
+  void operator=(const TimeProfilingPassesHandler &) = delete;
+
+  void registerCallbacks(PassInstrumentationCallbacks &PIC);
+
+private:
+  // Implementation of pass instrumentation callbacks.
+  void runBeforePass(StringRef PassID, Any IR);
+  void runAfterPass();
+};
+
 // Class that holds transitions between basic blocks.  The transitions
 // are contained in a map of values to names of basic blocks.
 class DCData {
 public:
   // Fill the map with the transitions from basic block \p B.
   DCData(const BasicBlock &B);
 
   // Return an iterator to the names of the successor blocks.

 };
 
 /// This class provides an interface to register all the standard pass
 /// instrumentations and manages their state (if any).
 class StandardInstrumentations {
   PrintIRInstrumentation PrintIR;
   PrintPassInstrumentation PrintPass;
   TimePassesHandler TimePasses;
+  TimeProfilingPassesHandler TimeProfilingPasses;
   OptNoneInstrumentation OptNone;
   OptBisectInstrumentation OptBisect;
   PreservedCFGCheckerInstrumentation PreservedCFGChecker;
   IRChangedPrinter PrintChangedIR;
   PseudoProbeVerifier PseudoProbeVerification;
   InLineChangePrinter PrintChangedDiff;
   DotCfgChangeReporter WebsiteChangeReporter;
   PrintCrashIRInstrumentation PrintCrashIR;

llvm/include/llvm/Transforms/Scalar/LoopPassManager.h

Optional<PreservedAnalyses> LoopPassManager::runSinglePass(
   // Get the loop in case of Loop pass and outermost loop in case of LoopNest
   // pass which is to be passed to BeforePass and AfterPass call backs.
   const Loop &L = getLoopFromIR(IR);
   // Check the PassInstrumentation's BeforePass callbacks before running the
   // pass, skip its execution completely if asked to (callback returns false).
   if (!PI.runBeforePass<Loop>(*Pass, L))
     return None;
 
-  PreservedAnalyses PA;
-  {
-    TimeTraceScope TimeScope(Pass->name(), IR.getName());
-    PA = Pass->run(IR, AM, AR, U);
-  }
+  PreservedAnalyses PA = Pass->run(IR, AM, AR, U);
 
   // do not pass deleted Loop into the instrumentation
   if (U.skipCurrentLoop())
     PI.runAfterPassInvalidated<IRUnitT>(*Pass, PA);
   else
     PI.runAfterPass<Loop>(*Pass, L, PA);
   return PA;
 }

llvm/lib/Analysis/CGSCCPassManager.cpp

FunctionAnalysisManager &FAM =
       AM.getCachedResult<FunctionAnalysisManagerCGSCCProxy>(*C)->getManager();
 
   for (auto &Pass : Passes) {
     // Check the PassInstrumentation's BeforePass callbacks before running the
     // pass, skip its execution completely if asked to (callback returns false).
     if (!PI.runBeforePass(*Pass, *C))
       continue;
 
-    PreservedAnalyses PassPA;
-    {
-      TimeTraceScope TimeScope(Pass->name());
-      PassPA = Pass->run(*C, AM, G, UR);
-    }
+    PreservedAnalyses PassPA = Pass->run(*C, AM, G, UR);
 
     if (UR.InvalidatedSCCs.count(C))
       PI.runAfterPassInvalidated<LazyCallGraph::SCC>(*Pass, PassPA);
     else
       PI.runAfterPass<LazyCallGraph::SCC>(*Pass, *C, PassPA);
 
     // Update the SCC if necessary.
     C = UR.UpdatedC ? UR.UpdatedC : C;
do {
           UR.UpdatedC = nullptr;
 
           // Check the PassInstrumentation's BeforePass callbacks before
           // running the pass, skip its execution completely if asked to
           // (callback returns false).
           if (!PI.runBeforePass<LazyCallGraph::SCC>(*Pass, *C))
             continue;
 
-          PreservedAnalyses PassPA;
-          {
-            TimeTraceScope TimeScope(Pass->name());
-            PassPA = Pass->run(*C, CGAM, CG, UR);
-          }
+          PreservedAnalyses PassPA = Pass->run(*C, CGAM, CG, UR);
 
           if (UR.InvalidatedSCCs.count(C))
             PI.runAfterPassInvalidated<LazyCallGraph::SCC>(*Pass, PassPA);
           else
             PI.runAfterPass<LazyCallGraph::SCC>(*Pass, *C, PassPA);
 
           // Update the SCC and RefSCC if necessary.
           C = UR.UpdatedC ? UR.UpdatedC : C;
for (LazyCallGraph::Node *N : Nodes) {
 
     if (NoRerun && FAM.getCachedResult<ShouldNotRunFunctionPassesAnalysis>(F))
       continue;
 
     PassInstrumentation PI = FAM.getResult<PassInstrumentationAnalysis>(F);
     if (!PI.runBeforePass<Function>(*Pass, F))
       continue;
 
-    PreservedAnalyses PassPA;
-    {
-      TimeTraceScope TimeScope(Pass->name());
-      PassPA = Pass->run(F, FAM);
-    }
-
+    PreservedAnalyses PassPA = Pass->run(F, FAM);
     PI.runAfterPass<Function>(*Pass, F, PassPA);
 
     // We know that the function pass couldn't have invalidated any other
     // function's analyses (that's the contract of a function pass), so
     // directly handle the function analysis manager's invalidation here.
     FAM.invalidate(F, EagerlyInvalidate ? PreservedAnalyses::none() : PassPA);
     if (NoRerun)
       (void)FAM.getResult<ShouldNotRunFunctionPassesAnalysis>(F);

llvm/lib/IR/PassManager.cpp

if (F.isDeclaration())
       continue;
 
     // Check the PassInstrumentation's BeforePass callbacks before running the
     // pass, skip its execution completely if asked to (callback returns
     // false).
     if (!PI.runBeforePass<Function>(*Pass, F))
       continue;
 
-    PreservedAnalyses PassPA;
-    {
-      TimeTraceScope TimeScope(Pass->name(), F.getName());
-      PassPA = Pass->run(F, FAM);
-    }
-
+    PreservedAnalyses PassPA = Pass->run(F, FAM);
     PI.runAfterPass(*Pass, F, PassPA);
 
     // We know that the function pass couldn't have invalidated any other
     // function's analyses (that's the contract of a function pass), so
     // directly handle the function analysis manager's invalidation here.
     FAM.invalidate(F, EagerlyInvalidate ? PreservedAnalyses::none() : PassPA);
 
     // Then intersect the preserved set so that invalidation of module

llvm/lib/Passes/StandardInstrumentations.cpp

 void InLineChangePrinter::registerCallbacks(PassInstrumentationCallbacks &PIC) {
   if (PrintChanged == ChangePrinter::DiffVerbose ||
       PrintChanged == ChangePrinter::DiffQuiet ||
       PrintChanged == ChangePrinter::ColourDiffVerbose ||
       PrintChanged == ChangePrinter::ColourDiffQuiet)
     TextChangeReporter<IRDataT<EmptyData>>::registerRequiredCallbacks(PIC);
 }
 
+TimeProfilingPassesHandler::TimeProfilingPassesHandler() {}
+
+void TimeProfilingPassesHandler::registerCallbacks(
+    PassInstrumentationCallbacks &PIC) {
+  if (!getTimeTraceProfilerInstance())
+    return;
+  PIC.registerBeforeNonSkippedPassCallback(
+      [this](StringRef P, Any IR) { this->runBeforePass(P, IR); });
+  PIC.registerAfterPassCallback(
+      [this](StringRef P, Any IR, const PreservedAnalyses &) {
+        this->runAfterPass();
+      },
+      true);
+  PIC.registerAfterPassInvalidatedCallback(
+      [this](StringRef P, const PreservedAnalyses &) { this->runAfterPass(); },
+      true);
+  PIC.registerBeforeAnalysisCallback(
+      [this](StringRef P, Any IR) { this->runBeforePass(P, IR); });
+  PIC.registerAfterAnalysisCallback(
+      [this](StringRef P, Any IR) { this->runAfterPass(); }, true);
+}
+
+void TimeProfilingPassesHandler::runBeforePass(StringRef PassID, Any IR) {
+  timeTraceProfilerBegin(PassID, getIRName(IR));
+}
+
+void TimeProfilingPassesHandler::runAfterPass() { timeTraceProfilerEnd(); }
+
 namespace {
 
 class DisplayNode;
 class DotCfgDiffDisplayGraph;
 
 // Base class for a node or edge in the dot-cfg-changes graph.
 class DisplayElement {
 public:
   // Is this in before, after, or both?
if (FAM)
     PreservedCFGChecker.registerCallbacks(PIC, *FAM);
   PrintChangedIR.registerCallbacks(PIC);
   PseudoProbeVerification.registerCallbacks(PIC);
   if (VerifyEach)
     Verify.registerCallbacks(PIC);
   PrintChangedDiff.registerCallbacks(PIC);
   WebsiteChangeReporter.registerCallbacks(PIC);
   PrintCrashIR.registerCallbacks(PIC);
+  // TimeProfiling records the pass running time cost.
+  // Its 'BeforePassCallback' can be appended at the tail of all the
+  // BeforeCallbacks by calling `registerCallbacks` in the end.
+  // Its 'AfterPassCallback' is put at the front of all the
+  // AfterCallbacks by its `registerCallbacks`. This is necessary
+  // to ensure that other callbacks are not included in the timings.
+  TimeProfilingPasses.registerCallbacks(PIC);
 }
 
 template class ChangeReporter<std::string>;
 template class TextChangeReporter<std::string>;
 
 template class BlockDataT<EmptyData>;
 template class FuncDataT<EmptyData>;
 template class IRDataT<EmptyData>;
 template class ChangeReporter<IRDataT<EmptyData>>;
 template class TextChangeReporter<IRDataT<EmptyData>>;
 template class IRComparer<EmptyData>;
 
 } // namespace llvm

llvm/lib/Transforms/Scalar/LoopPassManager.cpp

#ifndef NDEBUG
     Updater.ParentL = L->getParentLoop();
 #endif
     // Check the PassInstrumentation's BeforePass callbacks before running the
     // pass, skip its execution completely if asked to (callback returns
     // false).
     if (!PI.runBeforePass<Loop>(*Pass, *L))
       continue;
 
-    PreservedAnalyses PassPA;
-    {
-      TimeTraceScope TimeScope(Pass->name());
-      PassPA = Pass->run(*L, LAM, LAR, Updater);
-    }
+    PreservedAnalyses PassPA = Pass->run(*L, LAM, LAR, Updater);
 
     // Do not pass deleted Loop into the instrumentation.
     if (Updater.skipCurrentLoop())
       PI.runAfterPassInvalidated<Loop>(*Pass, PassPA);
     else
       PI.runAfterPass<Loop>(*Pass, *L, PassPA);
 
     if (LAR.MSSA && !PassPA.getChecker<MemorySSAAnalysis>().preserved())