2 files changed, 109 insertions, 1 deletions
diff --git a/llvm/lib/IR/PassTimingInfo.cpp b/llvm/lib/IR/PassTimingInfo.cpp
index 895cbc1f0ca..8165704a9ec 100644
--- a/llvm/lib/IR/PassTimingInfo.cpp
+++ b/llvm/lib/IR/PassTimingInfo.cpp
@@ -20,6 +20,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/IR/PassInstrumentation.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
@@ -33,6 +34,8 @@
 
 using namespace llvm;
 
+#define DEBUG_TYPE "time-passes"
+
 namespace llvm {
 
 bool TimePassesIsEnabled = false;
@@ -45,7 +48,7 @@ namespace {
 namespace legacy {
 
 //===----------------------------------------------------------------------===//
-// TimingInfo implementation
+// Legacy pass manager's PassTimingInfo implementation
 
 /// Provides an interface for collecting pass timing information.
 ///
@@ -156,4 +159,108 @@ void reportAndResetTimings() {
     legacy::PassTimingInfo::TheTimeInfo->print();
 }
 
+//===----------------------------------------------------------------------===//
+// Pass timing handling for the New Pass Manager
+//===----------------------------------------------------------------------===//
+
+/// Returns the timer for the specified pass invocation of \p PassID.
+/// Each time it creates a new timer.
+Timer &TimePassesHandler::getPassTimer(StringRef PassID) {
+  // Bump counts for each request of the timer.
+  unsigned Count = nextPassID(PassID);
+
+  // Unconditionally appending description with a pass-invocation number.
+  std::string FullDesc = formatv("{0} #{1}", PassID, Count).str();
+
+  PassInvocationID UID{PassID, Count};
+  Timer *T = new Timer(PassID, FullDesc, TG);
+  auto Pair = TimingData.try_emplace(UID, T);
+  assert(Pair.second && "should always create a new timer");
+  return *(Pair.first->second.get());
+}
+
+TimePassesHandler::TimePassesHandler(bool Enabled)
+    : TG("pass", "... Pass execution timing report ..."), Enabled(Enabled) {}
+
+void TimePassesHandler::print() { TG.print(*CreateInfoOutputFile()); }
+
+LLVM_DUMP_METHOD void TimePassesHandler::dump() const {
+  dbgs() << "Dumping timers for " << getTypeName<TimePassesHandler>()
+         << ":\n\tRunning:\n";
+  for (auto &I : TimingData) {
+    const Timer *MyTimer = I.second.get();
+    if (!MyTimer || MyTimer->isRunning())
+      dbgs() << "\tTimer " << MyTimer << " for pass " << I.first.first << "("
+             << I.first.second << ")\n";
+  }
+  dbgs() << "\tTriggered:\n";
+  for (auto &I : TimingData) {
+    const Timer *MyTimer = I.second.get();
+    if (!MyTimer || (MyTimer->hasTriggered() && !MyTimer->isRunning()))
+      dbgs() << "\tTimer " << MyTimer << " for pass " << I.first.first << "("
+             << I.first.second << ")\n";
+  }
+}
+
+void TimePassesHandler::startTimer(StringRef PassID) {
+  Timer &MyTimer = getPassTimer(PassID);
+  TimerStack.push_back(&MyTimer);
+  if (!MyTimer.isRunning())
+    MyTimer.startTimer();
+}
+
+void TimePassesHandler::stopTimer(StringRef PassID) {
+  assert(TimerStack.size() > 0 && "empty stack in popTimer");
+  Timer *MyTimer = TimerStack.pop_back_val();
+  assert(MyTimer && "timer should be present");
+  if (MyTimer->isRunning())
+    MyTimer->stopTimer();
+}
+
+static bool matchPassManager(StringRef PassID) {
+  size_t prefix_pos = PassID.find('<');
+  if (prefix_pos == StringRef::npos)
+    return false;
+  StringRef Prefix = PassID.substr(0, prefix_pos);
+  return Prefix.endswith("PassManager") || Prefix.endswith("PassAdaptor") ||
+         Prefix.endswith("AnalysisManagerProxy");
+}
+
+bool TimePassesHandler::runBeforePass(StringRef PassID, Any IR) {
+  if (matchPassManager(PassID))
+    return true;
+
+  startTimer(PassID);
+
+  LLVM_DEBUG(dbgs() << "after runBeforePass(" << PassID << ")\n");
+  LLVM_DEBUG(dump());
+
+  // we are not going to skip this pass, thus return true.
+  return true;
+}
+
+void TimePassesHandler::runAfterPass(StringRef PassID, Any IR) {
+  if (matchPassManager(PassID))
+    return;
+
+  stopTimer(PassID);
+
+  LLVM_DEBUG(dbgs() << "after runAfterPass(" << PassID << ")\n");
+  LLVM_DEBUG(dump());
+}
+
+void TimePassesHandler::registerCallbacks(PassInstrumentationCallbacks &PIC) {
+  if (!Enabled)
+    return;
+
+  PIC.registerBeforePassCallback(
+      [this](StringRef P, Any IR) { return this->runBeforePass(P, IR); });
+  PIC.registerAfterPassCallback(
+      [this](StringRef P, Any IR) { this->runAfterPass(P, IR); });
+  PIC.registerBeforeAnalysisCallback(
+      [this](StringRef P, Any IR) { this->runBeforePass(P, IR); });
+  PIC.registerAfterAnalysisCallback(
+      [this](StringRef P, Any IR) { this->runAfterPass(P, IR); });
+}
+
 } // namespace llvm
diff --git a/llvm/lib/Passes/StandardInstrumentations.cpp b/llvm/lib/Passes/StandardInstrumentations.cpp
index 6abd39c61ac..aa34584fa12 100644
--- a/llvm/lib/Passes/StandardInstrumentations.cpp
+++ b/llvm/lib/Passes/StandardInstrumentations.cpp
@@ -112,4 +112,5 @@ void StandardInstrumentations::registerCallbacks(
     PIC.registerBeforePassCallback(PrintIR::printBeforePass);
   if (llvm::shouldPrintAfterPass())
     PIC.registerAfterPassCallback(PrintIR::printAfterPass);
+  TimePasses.registerCallbacks(PIC);
 }