Extend TimeTrace to LLVM's new pass manager

[amonshiz]

Summary

With the addition of the LLD time tracing it made sense to include coverage
for LLVM's various passes. Doing so ensures that ThinLTO is also covered
with a time trace.

Before:

After:

[modocache]

Friendly ping for reviews on this one. I've found this to be a great tool for debugging the passes running during ThinLTO.

[anton-afanasyev]

I've faced with objections from @philip.pfaffe trying to implement the same: D62067. Since that time, trace profiler has supported multiple threads (D71059), so I do believe this can now correctly work. I would wait for review from new pm guys.

[modocache]

Friendly ping! New PM guys, maybe @chandlerc or @lebedev.ri could review?

[lebedev.ri]

I'm very much not the right reviewer here, sorry!

[rnk]

lgtm!

Sorry for the delay, I have been busy. I hadn't opened this up until now, but it seems straightforward, I don't think we need additional thorough review from local area code owners.

[rnk]

Based on perusing https://reviews.llvm.org/p/amonshiz/, I assume you need someone to push this. I'll go ahead and do that.

[amonshiz]

thank you @rnk for committing this for me!

[fedor.sergeev]

I'm sorry for the late comment, but this should have really been done through the PassInstrumentation framework.
It looks like a classic use-case.

[fedor.sergeev]

As an explanation why PassInstrumentation is a better choice - you do not need to manually extend all the individual pass managers,
so you never miss one. This particular change missed at least two of them: RepeatPass and DevirtSCCRepeatedPass.

[anton-afanasyev]

In D74516#1913475, @fedor.sergeev wrote:

I'm sorry for the late comment, but this should have really been done through the PassInstrumentation framework.
It looks like a classic use-case.

AFAII, this couldn't be done using register[Before/After]PassCallback() (like here llvm/lib/IR/PassTimingInfo.cpp) since profiler timer have to be assuredly finished if started, and that should be done directly by inserting [start/stop]TimeProfiler() in run[Before/After/Pass/Invalidated]().

llvm/include/llvm/Analysis/CGSCCPassManager.h

for (LazyCallGraph::Node *N : Nodes) {
         continue;
 
       Function &F = N->getFunction();
 
       PassInstrumentation PI = FAM.getResult<PassInstrumentationAnalysis>(F);
       if (!PI.runBeforePass<Function>(Pass, F))
         continue;
 
-      PreservedAnalyses PassPA = Pass.run(F, FAM);
+      PreservedAnalyses PassPA;
+      {
+        TimeTraceScope TimeScope(Pass.name());
+        PassPA = Pass.run(F, FAM);
+      }
 
       PI.runAfterPass<Function>(Pass, F);
 
       // 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, PassPA);
 
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 = Pass.run(*C, CGAM, CG, UR);
+          PreservedAnalyses PassPA;
+          {
+            TimeTraceScope TimeScope(Pass.name());
+            PassPA = Pass.run(*C, CGAM, CG, UR);
+          }
 
           if (UR.InvalidatedSCCs.count(C))
             PI.runAfterPassInvalidated<LazyCallGraph::SCC>(Pass);
           else
             PI.runAfterPass<LazyCallGraph::SCC>(Pass, *C);
 
           // Update the SCC and RefSCC if necessary.
           C = UR.UpdatedC ? UR.UpdatedC : C;

llvm/include/llvm/IR/PassManager.h

 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassInstrumentation.h"
 #include "llvm/IR/PassManagerInternal.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/TimeProfiler.h"
 #include "llvm/Support/TypeName.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <cassert>
 #include <cstring>
 #include <iterator>
 #include <list>
 #include <memory>
for (unsigned Idx = 0, Size = Passes.size(); Idx != Size; ++Idx) {
                << "\n";
 
       // Check the PassInstrumentation's BeforePass callbacks before running the
       // pass, skip its execution completely if asked to (callback returns
       // false).
       if (!PI.runBeforePass<IRUnitT>(*P, IR))
         continue;
 
-      PreservedAnalyses PassPA = P->run(IR, AM, ExtraArgs...);
+      PreservedAnalyses PassPA;
+      {
+        TimeTraceScope TimeScope(P->name(), IR.getName());
+        PassPA = P->run(IR, AM, ExtraArgs...);
+      }
 
       // Call onto PassInstrumentation's AfterPass callbacks immediately after
       // running the pass.
       PI.runAfterPass<IRUnitT>(*P, IR);
 
       // Update the analysis manager as each pass runs and potentially
       // invalidates analyses.
       AM.invalidate(IR, PassPA);
for (Function &F : M) {
       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 = Pass.run(F, FAM);
+
+      PreservedAnalyses PassPA;
+      {
+        TimeTraceScope TimeScope(Pass.name(), F.getName());
+        PassPA = Pass.run(F, FAM);
+      }
 
       PI.runAfterPass(Pass, F);
 
       // 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, PassPA);
 

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

#ifndef NDEBUG
       assert(L->isRecursivelyLCSSAForm(LAR.DT, LI) &&
              "Loops must remain in LCSSA form!");
 #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 = Pass.run(*L, LAM, LAR, Updater);
+
+      PreservedAnalyses PassPA;
+      {
+        TimeTraceScope TimeScope(Pass.name());
+        PassPA = Pass.run(*L, LAM, LAR, Updater);
+      }
 
       // Do not pass deleted Loop into the instrumentation.
       if (Updater.skipCurrentLoop())
         PI.runAfterPassInvalidated<Loop>(Pass);
       else
         PI.runAfterPass<Loop>(Pass, *L);
 
       // FIXME: We should verify the set of analyses relevant to Loop passes

llvm/lib/Analysis/CGSCCPassManager.cpp

 #include "llvm/IR/Constant.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/PassManagerImpl.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Support/TimeProfiler.h"
 #include <algorithm>
 #include <cassert>
 #include <iterator>
 
 #define DEBUG_TYPE "cgscc"
 
 using namespace llvm;
 
for (auto &Pass : Passes) {
     if (DebugLogging)
       dbgs() << "Running pass: " << Pass->name() << " on " << *C << "\n";
 
     // 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 = Pass->run(*C, AM, G, UR);
+    PreservedAnalyses PassPA;
+    {
+      TimeTraceScope TimeScope(Pass->name());
+      PassPA = Pass->run(*C, AM, G, UR);
+    }
 
     if (UR.InvalidatedSCCs.count(C))
       PI.runAfterPassInvalidated<LazyCallGraph::SCC>(*Pass);
     else
       PI.runAfterPass<LazyCallGraph::SCC>(*Pass, *C);
 
     // Update the SCC if necessary.
     C = UR.UpdatedC ? UR.UpdatedC : C;

llvm/lib/Transforms/Scalar/LoopPassManager.cpp

 //===- LoopPassManager.cpp - Loop pass management -------------------------===//
 //
 // 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 "llvm/Support/TimeProfiler.h"
 #include "llvm/Transforms/Scalar/LoopPassManager.h"
 #include "llvm/Analysis/LoopInfo.h"
 
 using namespace llvm;
 
 // Explicit template instantiations and specialization defininitions for core
 // template typedefs.
 namespace llvm {
for (auto &Pass : Passes) {
     if (DebugLogging)
       dbgs() << "Running pass: " << Pass->name() << " on " << L;
 
     // 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 = Pass->run(L, AM, AR, U);
+    PreservedAnalyses PassPA;
+    {
+      TimeTraceScope TimeScope(Pass->name(), L.getName());
+      PassPA = Pass->run(L, AM, AR, U);
+    }
 
     // do not pass deleted Loop into the instrumentation
     if (U.skipCurrentLoop())
       PI.runAfterPassInvalidated<Loop>(*Pass);
     else
       PI.runAfterPass<Loop>(*Pass, L);
 
     // If the loop was deleted, abort the run and return to the outer walk.

llvm/test/Other/new-pm-time-trace.ll

+; RUN: opt -time-trace -time-trace-file %t.json \
+; RUN:     -disable-output -disable-verify \
+; RUN:     -passes='default<O3>' %s
+; RUN: cat %t.json \
+; RUN:  | %python -c 'import json, sys; json.dump(json.loads(sys.stdin.read()), sys.stdout, sort_keys=True, indent=2)' \
+; RUN:  | FileCheck %s
+
+; CHECK: "name": "Total FunctionToLoopPassAdaptor
+
+define void @foo(i1 %x, i8* %p1, i8* %p2) {
+entry:
+  store i8 42, i8* %p1
+  br i1 %x, label %loop, label %exit
+
+loop:
+  %tmp1 = load i8, i8* %p2
+  br label %loop
+
+exit:
+  ret void
+}
+
+declare void @bar()
+

llvm/tools/opt/opt.cpp

static cl::opt<bool> DiscardValueNames(
     cl::desc("Discard names from Value (other than GlobalValue)."),
     cl::init(false), cl::Hidden);
 
 static cl::opt<bool> Coroutines(
   "enable-coroutines",
   cl::desc("Enable coroutine passes."),
   cl::init(false), cl::Hidden);
 
+static cl::opt<bool> TimeTrace(
+    "time-trace",
+    cl::desc("Record time trace"));
+
+static cl::opt<unsigned> TimeTraceGranularity(
+    "time-trace-granularity",
+    cl::desc("Minimum time granularity (in microseconds) traced by time profiler"),
+    cl::init(500), cl::Hidden);
+
+static cl::opt<std::string>
+    TimeTraceFile("time-trace-file",
+                    cl::desc("Specify time trace file destination"),
+                    cl::value_desc("filename"));
+
 static cl::opt<bool> RemarksWithHotness(
     "pass-remarks-with-hotness",
     cl::desc("With PGO, include profile count in optimization remarks"),
     cl::Hidden);
 
 static cl::opt<unsigned>
     RemarksHotnessThreshold("pass-remarks-hotness-threshold",
                             cl::desc("Minimum profile count required for "
for (const auto &Entry : Map) {
     DebugifyStatistics Stats = Entry.second;
 
     OS << Pass << ',' << Stats.NumDbgValuesMissing << ','
        << Stats.NumDbgLocsMissing << ',' << Stats.getMissingValueRatio() << ','
        << Stats.getEmptyLocationRatio() << '\n';
   }
 }
 
+struct TimeTracerRAII {
+  TimeTracerRAII(StringRef ProgramName) {
+    if (TimeTrace)
+      timeTraceProfilerInitialize(TimeTraceGranularity, ProgramName);
+  }
+  ~TimeTracerRAII() {
+    if (TimeTrace) {
+      if (auto E = timeTraceProfilerWrite(TimeTraceFile, OutputFilename)) {
+        handleAllErrors(std::move(E), [&](const StringError &SE) {
+          errs() << SE.getMessage() << "\n";
+        });
+        return;
+      }
+      timeTraceProfilerCleanup();
+    }
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // main for opt
 //
 int main(int argc, char **argv) {
   InitLLVM X(argc, argv);
 
   // Enable debug stream buffering.
   EnableDebugBuffering = true;
#endif
   cl::ParseCommandLineOptions(argc, argv,
     "llvm .bc -> .bc modular optimizer and analysis printer\n");
 
   if (AnalyzeOnly && NoOutput) {
     errs() << argv[0] << ": analyze mode conflicts with no-output mode.\n";
     return 1;
   }
 
+  TimeTracerRAII TimeTracer(argv[0]);
+
   SMDiagnostic Err;
 
   Context.setDiscardValueNames(DiscardValueNames);
   if (!DisableDITypeMap)
     Context.enableDebugTypeODRUniquing();
 
   Expected<std::unique_ptr<ToolOutputFile>> RemarksFileOrErr =
       setupLLVMOptimizationRemarks(Context, RemarksFilename, RemarksPasses,