[llvm-jitlink] Move statistics code into a separate file.

Further isolates statistics gathering / reporting code from the rest of llvm-jitlink.

4 files changed, 140 insertions, 118 deletions

diff --git a/llvm/tools/llvm-jitlink/CMakeLists.txt b/llvm/tools/llvm-jitlink/CMakeLists.txt
index b7371e7..a9f8bd1 100644
--- a/llvm/tools/llvm-jitlink/CMakeLists.txt
+++ b/llvm/tools/llvm-jitlink/CMakeLists.txt
@@ -24,6 +24,7 @@ add_llvm_tool(llvm-jitlink
   llvm-jitlink-coff.cpp
   llvm-jitlink-elf.cpp
   llvm-jitlink-macho.cpp
+  llvm-jitlink-statistics.cpp
   )
 
 if(${CMAKE_SYSTEM_NAME} MATCHES "Haiku")
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink-statistics.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink-statistics.cpp
new file mode 100644
index 0000000..1179795
--- /dev/null
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink-statistics.cpp
@@ -0,0 +1,135 @@
+//===-- llvm-jitlink-statistics.cpp -- gathers/reports JIT-linking stats --===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains the code for enabling, gathering and reporting
+// llvm-jitlink statistics.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm-jitlink.h"
+#include "llvm/Support/CommandLine.h"
+
+#define DEBUG_TYPE "llvm_jitlink"
+
+using namespace llvm;
+using namespace llvm::jitlink;
+using namespace llvm::orc;
+
+static cl::opt<bool> ShowPrePruneTotalBlockSize(
+    "pre-prune-total-block-size",
+    cl::desc("Total size of all blocks (including zero-fill) in all "
+             "graphs (pre-pruning)"),
+    cl::init(false));
+
+static cl::opt<bool> ShowPostFixupTotalBlockSize(
+    "post-fixup-total-block-size",
+    cl::desc("Total size of all blocks (including zero-fill) in all "
+             "graphs (post-fixup)"),
+    cl::init(false));
+
+class StatsPlugin : public ObjectLinkingLayer::Plugin {
+public:
+  static void enableIfNeeded(Session &S, bool UsingOrcRuntime) {
+    std::unique_ptr<StatsPlugin> Instance;
+    auto GetStats = [&]() -> StatsPlugin & {
+      if (!Instance)
+        Instance.reset(new StatsPlugin(UsingOrcRuntime));
+      return *Instance;
+    };
+
+    if (ShowPrePruneTotalBlockSize)
+      GetStats().PrePruneTotalBlockSize = 0;
+
+    if (ShowPostFixupTotalBlockSize)
+      GetStats().PostFixupTotalBlockSize = 0;
+
+    if (Instance)
+      S.ObjLayer.addPlugin(std::move(Instance));
+  }
+
+  ~StatsPlugin() { publish(dbgs()); }
+
+  void publish(raw_ostream &OS);
+
+  void modifyPassConfig(MaterializationResponsibility &MR, LinkGraph &G,
+                        PassConfiguration &PassConfig) override {
+    PassConfig.PrePrunePasses.push_back(
+        [this](LinkGraph &G) { return recordPrePruneStats(G); });
+    PassConfig.PostFixupPasses.push_back(
+        [this](LinkGraph &G) { return recordPostFixupStats(G); });
+  }
+
+  Error notifyFailed(MaterializationResponsibility &MR) override {
+    return Error::success();
+  }
+
+  Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override {
+    return Error::success();
+  }
+
+  void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
+                                   ResourceKey SrcKey) override {}
+
+private:
+  StatsPlugin(bool UsingOrcRuntime) : UsingOrcRuntime(UsingOrcRuntime) {}
+  Error recordPrePruneStats(jitlink::LinkGraph &G);
+  Error recordPostFixupStats(jitlink::LinkGraph &G);
+
+  bool UsingOrcRuntime;
+
+  std::mutex M;
+  std::optional<uint64_t> PrePruneTotalBlockSize;
+  std::optional<uint64_t> PostFixupTotalBlockSize;
+  std::optional<DenseMap<size_t, size_t>> EdgeCountDetails;
+};
+
+void StatsPlugin::publish(raw_ostream &OS) {
+
+  if (UsingOrcRuntime)
+    OS << "Note: Session stats include runtime and entry point lookup, but "
+          "not JITDylib initialization/deinitialization.\n";
+
+  OS << "Statistics:\n";
+  if (PrePruneTotalBlockSize)
+    OS << "  Total size of all blocks before pruning: "
+       << *PrePruneTotalBlockSize << "\n";
+
+  if (PostFixupTotalBlockSize)
+    OS << "  Total size of all blocks after fixups: "
+       << *PostFixupTotalBlockSize << "\n";
+}
+
+static uint64_t computeTotalBlockSizes(LinkGraph &G) {
+  uint64_t TotalSize = 0;
+  for (auto *B : G.blocks())
+    TotalSize += B->getSize();
+  return TotalSize;
+}
+
+Error StatsPlugin::recordPrePruneStats(LinkGraph &G) {
+  std::lock_guard<std::mutex> Lock(M);
+
+  if (PrePruneTotalBlockSize)
+    *PrePruneTotalBlockSize += computeTotalBlockSizes(G);
+
+  return Error::success();
+}
+
+Error StatsPlugin::recordPostFixupStats(LinkGraph &G) {
+  std::lock_guard<std::mutex> Lock(M);
+
+  if (PostFixupTotalBlockSize)
+    *PostFixupTotalBlockSize += computeTotalBlockSizes(G);
+  return Error::success();
+}
+
+namespace llvm {
+void enableStatistics(Session &S, bool UsingOrcRuntime) {
+  StatsPlugin::enableIfNeeded(S, UsingOrcRuntime);
+}
+} // namespace llvm
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index eef3aa0..3b8def4 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -181,21 +181,6 @@ static cl::opt<std::string> ShowLinkGraphs(
              "matching that regex after fixups have been applied"),
     cl::Optional, cl::cat(JITLinkCategory));
 
-enum class JITLinkStats {
-  PrePruneTotalBlockSize,
-  PostFixupTotalBlockSize
-};
-
-static cl::list<JITLinkStats> ShowStats(
-    cl::desc("Statistics:"),
-    cl::values(
-        clEnumValN(JITLinkStats::PrePruneTotalBlockSize,
-                   "pre-prune-total-block-size",
-                   "Total size of all blocks in all graphs (pre-pruning)"),
-        clEnumValN(JITLinkStats::PostFixupTotalBlockSize,
-                   "post-fixup-total-block-size",
-                   "Tatal size of all blocks in all graphs (post-fixup)")));
-
 static cl::opt<bool> ShowTimes("show-times",
                                cl::desc("Show times for llvm-jitlink phases"),
                                cl::init(false), cl::cat(JITLinkCategory));
@@ -1926,89 +1911,6 @@ static Error addSelfRelocations(LinkGraph &G) {
   return Error::success();
 }
 
-LLVMJITLinkStatistics::LLVMJITLinkStatistics(Session &S) {
-
-  class StatisticsPlugin : public ObjectLinkingLayer::Plugin {
-  public:
-    StatisticsPlugin(LLVMJITLinkStatistics &Stats) : Stats(Stats) {}
-
-    void modifyPassConfig(MaterializationResponsibility &MR, LinkGraph &G,
-                          PassConfiguration &PassConfig) override {
-      PassConfig.PrePrunePasses.push_back(
-          [this](LinkGraph &G) { return Stats.recordPrePruneStats(G); });
-      PassConfig.PostFixupPasses.push_back(
-          [this](LinkGraph &G) { return Stats.recordPostFixupStats(G); });
-    }
-
-    Error notifyFailed(MaterializationResponsibility &MR) override {
-      return Error::success();
-    }
-
-    Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override {
-      return Error::success();
-    }
-
-    void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
-                                     ResourceKey SrcKey) override {}
-
-  private:
-    LLVMJITLinkStatistics &Stats;
-  };
-
-  S.ObjLayer.addPlugin(std::make_unique<StatisticsPlugin>(*this));
-
-  // Walk command line option and enable requested stats.
-  for (auto &Stat : ShowStats) {
-    switch (Stat) {
-    case JITLinkStats::PrePruneTotalBlockSize:
-      PrePruneTotalBlockSize = 0;
-      break;
-    case JITLinkStats::PostFixupTotalBlockSize:
-      PostFixupTotalBlockSize = 0;
-      break;
-    }
-  }
-}
-
-void LLVMJITLinkStatistics::print(raw_ostream &OS) {
-
-  if (!OrcRuntime.empty())
-    OS << "Note: Session stats include runtime and entry point lookup, but "
-          "not JITDylib initialization/deinitialization.\n";
-
-  OS << "Statistics:\n";
-  if (PrePruneTotalBlockSize)
-    OS << "  Total size of all blocks before pruning: "
-       << *PrePruneTotalBlockSize << "\n";
-
-  if (PostFixupTotalBlockSize)
-    OS << "  Total size of all blocks after fixups: "
-       << *PostFixupTotalBlockSize << "\n";
-}
-
-static uint64_t computeTotalBlockSizes(LinkGraph &G) {
-  uint64_t TotalSize = 0;
-  for (auto *B : G.blocks())
-    TotalSize += B->getSize();
-  return TotalSize;
-}
-
-Error LLVMJITLinkStatistics::recordPrePruneStats(LinkGraph &G) {
-  std::lock_guard<std::mutex> Lock(M);
-
-  if (PrePruneTotalBlockSize)
-    *PrePruneTotalBlockSize += computeTotalBlockSizes(G);
-  return Error::success();
-}
-
-Error LLVMJITLinkStatistics::recordPostFixupStats(LinkGraph &G) {
-  std::lock_guard<std::mutex> Lock(M);
-
-  if (PostFixupTotalBlockSize)
-    *PostFixupTotalBlockSize += computeTotalBlockSizes(G);
-  return Error::success();
-}
-
 static Expected<ExecutorSymbolDef> getMainEntryPoint(Session &S) {
   return S.ES.lookup(S.JDSearchOrder, S.ES.intern(EntryPointName));
 }
@@ -2100,9 +2002,7 @@ int main(int argc, char *argv[]) {
 
   auto S = ExitOnErr(Session::Create(std::move(TT), std::move(Features)));
 
-  std::unique_ptr<LLVMJITLinkStatistics> Stats;
-  if (!ShowStats.empty())
-    Stats = std::make_unique<LLVMJITLinkStatistics>(*S);
+  enableStatistics(*S, !OrcRuntime.empty());
 
   {
     TimeRegion TR(Timers ? &Timers->LoadObjectsTimer : nullptr);
@@ -2129,9 +2029,6 @@ int main(int argc, char *argv[]) {
   if (ShowAddrs)
     S->dumpSessionInfo(outs());
 
-  if (Stats)
-    Stats->print(outs());
-
   if (!EntryPoint) {
     if (Timers)
       Timers->JITLinkTG.printAll(errs());
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.h b/llvm/tools/llvm-jitlink/llvm-jitlink.h
index a019bde..9389be4 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.h
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.h
@@ -88,20 +88,6 @@ private:
   Session(std::unique_ptr<orc::ExecutorProcessControl> EPC, Error &Err);
 };
 
-class LLVMJITLinkStatistics {
-public:
-  LLVMJITLinkStatistics(Session &S);
-  void print(raw_ostream &OS);
-
-private:
-  Error recordPrePruneStats(jitlink::LinkGraph &G);
-  Error recordPostFixupStats(jitlink::LinkGraph &G);
-
-  std::mutex M;
-  std::optional<uint64_t> PrePruneTotalBlockSize;
-  std::optional<uint64_t> PostFixupTotalBlockSize;
-};
-
 /// Record symbols, GOT entries, stubs, and sections for ELF file.
 Error registerELFGraphInfo(Session &S, jitlink::LinkGraph &G);
 
@@ -111,6 +97,9 @@ Error registerMachOGraphInfo(Session &S, jitlink::LinkGraph &G);
 /// Record symbols, GOT entries, stubs, and sections for COFF file.
 Error registerCOFFGraphInfo(Session &S, jitlink::LinkGraph &G);
 
+/// Adds a statistics gathering plugin if any stats options are used.
+void enableStatistics(Session &S, bool UsingOrcRuntime);
+
 } // end namespace llvm
 
 #endif // LLVM_TOOLS_LLVM_JITLINK_LLVM_JITLINK_H