Merge pull request swiftlang#11901 from swiftlang/cherrypick-save-stats

Cherrypick save stats

Author: tomershafir

llvm/docs/CommandGuide/llc.rst

@@ -136,6 +136,12 @@ End-user Options
 
  Print statistics recorded by code-generation passes.
 
+.. option:: --save-stats, --save-stats=cwd, --save-stats=obj
+
+ Save LLVM statistics to a file in the current directory
+ (:option:`--save-stats`/"--save-stats=cwd") or the directory
+ of the output file ("--save-stats=obj") in JSON format.
+
 .. option:: --time-passes
 
  Record the amount of time needed for each pass and print a report to standard

llvm/docs/CommandGuide/opt.rst

@@ -70,6 +70,12 @@ OPTIONS
 
  Print statistics.
 
+.. option:: --save-stats, --save-stats=cwd, --save-stats=obj
+
+ Save LLVM statistics to a file in the current directory
+ (:option:`--save-stats`/"--save-stats=cwd") or the directory
+ of the output file ("--save-stats=obj") in JSON format.
+
 .. option:: -time-passes
 
  Record the amount of time needed for each pass and print it to standard

llvm/docs/ReleaseNotes.md

@@ -297,6 +297,8 @@ Changes to the LLVM tools
   ([#47468](https://github.com/llvm/llvm-project/issues/47468))
 * llvm-addr2line now supports a `+` prefix when specifying an address.
 * Support for `SHT_LLVM_BB_ADDR_MAP` versions 0 and 1 has been dropped.
+* Add `--save-stats` option to `llc` to save LLVM statistics to a file. Compatible with the Clang option.
+* Add `--save-stats` option to `opt` to save LLVM statistics to a file. Compatible with the Clang option.
 
 Changes to LLDB
 ---------------------------------

llvm/include/llvm/CodeGen/CommandFlags.h

@@ -156,12 +156,22 @@ LLVM_ABI bool getJMCInstrument();
 
 LLVM_ABI bool getXCOFFReadOnlyPointers();
 
+enum SaveStatsMode { None, Cwd, Obj };
+
+LLVM_ABI SaveStatsMode getSaveStats();
+
 /// Create this object with static storage to register codegen-related command
 /// line options.
 struct RegisterCodeGenFlags {
   LLVM_ABI RegisterCodeGenFlags();
 };
 
+/// Tools that support stats saving should create this object with static
+/// storage to register the --save-stats command line option.
+struct RegisterSaveStatsFlag {
+  LLVM_ABI RegisterSaveStatsFlag();
+};
+
 LLVM_ABI bool getEnableBBAddrMap();
 
 LLVM_ABI llvm::BasicBlockSection
@@ -205,6 +215,17 @@ LLVM_ABI Expected<std::unique_ptr<TargetMachine>> createTargetMachineForTriple(
     StringRef TargetTriple,
     CodeGenOptLevel OptLevel = CodeGenOptLevel::Default);
 
+/// Conditionally enables the collection of LLVM statistics during the tool run,
+/// based on the value of the flag. Must be called before the tool run to
+/// actually collect data.
+LLVM_ABI void MaybeEnableStatistics();
+
+/// Conditionally saves the collected LLVM statistics to the received output
+/// file, based on the value of the flag. Should be called after the tool run,
+/// and must follow a call to `MaybeEnableStatistics()` to actually have data to
+/// write.
+LLVM_ABI int MaybeSaveStatistics(StringRef OutputFilename, StringRef ToolName);
+
 } // namespace codegen
 } // namespace llvm
 

llvm/lib/CodeGen/CommandFlags.cpp

@@ -13,33 +13,43 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/CommandFlags.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Module.h"
 #include "llvm/MC/MCTargetOptionsCommandFlags.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/WithColor.h"
+#include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/TargetParser/Host.h"
 #include "llvm/TargetParser/SubtargetFeature.h"
 #include "llvm/TargetParser/Triple.h"
+#include <cassert>
+#include <memory>
 #include <optional>
+#include <system_error>
 
 using namespace llvm;
 
 #define CGOPT(TY, NAME)                                                        \
   static cl::opt<TY> *NAME##View;                                              \
   TY codegen::get##NAME() {                                                    \
-    assert(NAME##View && "RegisterCodeGenFlags not created.");                 \
+    assert(NAME##View && "Flag not registered.");                              \
     return *NAME##View;                                                        \
   }
 
 #define CGLIST(TY, NAME)                                                       \
   static cl::list<TY> *NAME##View;                                             \
   std::vector<TY> codegen::get##NAME() {                                       \
-    assert(NAME##View && "RegisterCodeGenFlags not created.");                 \
+    assert(NAME##View && "Flag not registered.");                              \
     return *NAME##View;                                                        \
   }
 
@@ -111,13 +121,14 @@ CGOPT(bool, DebugStrictDwarf)
 CGOPT(unsigned, AlignLoops)
 CGOPT(bool, JMCInstrument)
 CGOPT(bool, XCOFFReadOnlyPointers)
+CGOPT(codegen::SaveStatsMode, SaveStats)
 
-codegen::RegisterCodeGenFlags::RegisterCodeGenFlags() {
 #define CGBINDOPT(NAME)                                                        \
   do {                                                                         \
     NAME##View = std::addressof(NAME);                                         \
   } while (0)
 
+codegen::RegisterCodeGenFlags::RegisterCodeGenFlags() {
   static cl::opt<std::string> MArch(
       "march", cl::desc("Architecture to generate code for (see --version)"));
   CGBINDOPT(MArch);
@@ -523,11 +534,25 @@ codegen::RegisterCodeGenFlags::RegisterCodeGenFlags() {
       cl::init(false));
   CGBINDOPT(DisableIntegratedAS);
 
-#undef CGBINDOPT
-
   mc::RegisterMCTargetOptionsFlags();
 }
 
+codegen::RegisterSaveStatsFlag::RegisterSaveStatsFlag() {
+  static cl::opt<SaveStatsMode> SaveStats(
+      "save-stats",
+      cl::desc(
+          "Save LLVM statistics to a file in the current directory"
+          "(`-save-stats`/`-save-stats=cwd`) or the directory of the output"
+          "file (`-save-stats=obj`). (default: cwd)"),
+      cl::values(clEnumValN(SaveStatsMode::Cwd, "cwd",
+                            "Save to the current working directory"),
+                 clEnumValN(SaveStatsMode::Cwd, "", ""),
+                 clEnumValN(SaveStatsMode::Obj, "obj",
+                            "Save to the output file directory")),
+      cl::init(SaveStatsMode::None), cl::ValueOptional);
+  CGBINDOPT(SaveStats);
+}
+
 llvm::BasicBlockSection
 codegen::getBBSectionsMode(llvm::TargetOptions &Options) {
   if (getBBSections() == "all")
@@ -774,3 +799,42 @@ codegen::createTargetMachineForTriple(StringRef TargetTriple,
                                  TargetTriple);
   return std::unique_ptr<TargetMachine>(Target);
 }
+
+void codegen::MaybeEnableStatistics() {
+  if (getSaveStats() == SaveStatsMode::None)
+    return;
+
+  llvm::EnableStatistics(false);
+}
+
+int codegen::MaybeSaveStatistics(StringRef OutputFilename, StringRef ToolName) {
+  auto SaveStatsValue = getSaveStats();
+  if (SaveStatsValue == codegen::SaveStatsMode::None)
+    return 0;
+
+  SmallString<128> StatsFilename;
+  if (SaveStatsValue == codegen::SaveStatsMode::Obj) {
+    StatsFilename = OutputFilename;
+    llvm::sys::path::remove_filename(StatsFilename);
+  } else {
+    assert(SaveStatsValue == codegen::SaveStatsMode::Cwd &&
+           "Should have been a valid --save-stats value");
+  }
+
+  auto BaseName = llvm::sys::path::filename(OutputFilename);
+  llvm::sys::path::append(StatsFilename, BaseName);
+  llvm::sys::path::replace_extension(StatsFilename, "stats");
+
+  auto FileFlags = llvm::sys::fs::OF_TextWithCRLF;
+  std::error_code EC;
+  auto StatsOS =
+      std::make_unique<llvm::raw_fd_ostream>(StatsFilename, EC, FileFlags);
+  if (EC) {
+    WithColor::error(errs(), ToolName)
+        << "Unable to open statistics file: " << EC.message() << "\n";
+    return 1;
+  }
+
+  llvm::PrintStatisticsJSON(*StatsOS);
+  return 0;
+}

llvm/test/tools/llc/save-stats.ll

@@ -0,0 +1,18 @@
+; REQUIRES: asserts
+
+; Ensure the test runs in a temp directory. See https://github.com/llvm/llvm-project/pull/167403#event-20848739526
+; RUN: rm -rf %t.dir && mkdir -p %t.dir && cd %t.dir
+
+; RUN: llc --save-stats=obj -o %t.s %s && cat %t.stats | FileCheck %s
+; RUN: llc --save-stats=cwd -o %t.s %s && cat %{t:stem}.tmp.stats | FileCheck %s
+; RUN: llc --save-stats -o %t.s %s && cat %{t:stem}.tmp.stats | FileCheck %s
+; RUN: not llc --save-stats=invalid -o %t.s %s 2>&1 | FileCheck %s --check-prefix=INVALID_ARG
+
+; CHECK: {
+; CHECK: "asm-printer.EmittedInsts":
+; CHECK: }
+
+; INVALID_ARG: {{.*}}llc{{.*}}: for the --save-stats option: Cannot find option named 'invalid'!
+define i32 @func() {
+  ret i32 0
+}

llvm/test/tools/opt/save-stats.ll

@@ -0,0 +1,18 @@
+; REQUIRES: asserts
+
+; Ensure the test runs in a temp directory. See https://github.com/llvm/llvm-project/pull/167403#event-20848739526
+; RUN: rm -rf %t.dir && mkdir -p %t.dir && cd %t.dir
+
+; RUN: opt -S -passes=instcombine --save-stats=obj -o %t.ll %s && cat %t.stats | FileCheck %s
+; RUN: opt -S -passes=instcombine --save-stats=cwd -o %t.ll %s && cat %{t:stem}.tmp.stats | FileCheck %s
+; RUN: opt -S -passes=instcombine --save-stats -o %t.ll %s && cat %{t:stem}.tmp.stats | FileCheck %s
+; RUN: not opt -S --save-stats=invalid -o %t.ll %s 2>&1 | FileCheck %s --check-prefix=INVALID_ARG
+
+; CHECK: {
+; CHECK: "instcombine.NumWorklistIterations":
+; CHECK: }
+
+; INVALID_ARG: {{.*}}opt{{.*}}: for the --save-stats option: Cannot find option named 'invalid'!
+define i32 @func() {
+  ret i32 0
+}

llvm/tools/llc/llc.cpp

@@ -15,6 +15,7 @@
 #include "NewPMDriver.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/CAS/ObjectStore.h"
 #include "llvm/CodeGen/CommandFlags.h"
@@ -45,6 +46,7 @@
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/FormattedStream.h"
 #include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/Path.h"
 #include "llvm/Support/PluginLoader.h"
 #include "llvm/Support/Process.h"
 #include "llvm/Support/SourceMgr.h"
@@ -58,11 +60,13 @@
 #include "llvm/TargetParser/SubtargetFeature.h"
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Utils/Cloning.h"
+#include <cassert>
 #include <memory>
 #include <optional>
 using namespace llvm;
 
 static codegen::RegisterCodeGenFlags CGF;
+static codegen::RegisterSaveStatsFlag SSF;
 
 // General options for llc.  Other pass-specific options are specified
 // within the corresponding llc passes, and target-specific options
@@ -262,7 +266,8 @@ static cl::opt<RunPassOption, true, cl::parser<std::string>> RunPass(
     cl::desc("Run compiler only for specified passes (comma separated list)"),
     cl::value_desc("pass-name"), cl::location(RunPassOpt));
 
-static int compileModule(char **, LLVMContext &);
+static int compileModule(char **argv, LLVMContext &Context,
+                         std::string &OutputFilename);
 
 [[noreturn]] static void reportError(Twine Msg, StringRef Filename = "") {
   SmallString<256> Prefix;
@@ -282,9 +287,7 @@ static int compileModule(char **, LLVMContext &);
   llvm_unreachable("reportError() should not return");
 }
 
-static std::unique_ptr<ToolOutputFile> GetOutputStream(const char *TargetName,
-                                                       Triple::OSType OS,
-                                                       const char *ProgName) {
+static std::unique_ptr<ToolOutputFile> GetOutputStream(Triple::OSType OS) {
   // If we don't yet have an output filename, make one.
   if (OutputFilename.empty()) {
     if (InputFilename == "-")
@@ -450,18 +453,22 @@ int main(int argc, char **argv) {
     reportError(std::move(E), RemarksFilename);
   std::unique_ptr<ToolOutputFile> RemarksFile = std::move(*RemarksFileOrErr);
 
+  codegen::MaybeEnableStatistics();
+  std::string OutputFilename;
+
   if (InputLanguage != "" && InputLanguage != "ir" && InputLanguage != "mir")
     reportError("input language must be '', 'IR' or 'MIR'");
 
   // Compile the module TimeCompilations times to give better compile time
   // metrics.
   for (unsigned I = TimeCompilations; I; --I)
-    if (int RetVal = compileModule(argv, Context))
+    if (int RetVal = compileModule(argv, Context, OutputFilename))
       return RetVal;
 
   if (RemarksFile)
     RemarksFile->keep();
-  return 0;
+
+  return codegen::MaybeSaveStatistics(OutputFilename, "llc");
 }
 
 static bool addPass(PassManagerBase &PM, const char *argv0,
@@ -493,7 +500,8 @@ static bool addPass(PassManagerBase &PM, const char *argv0,
   return false;
 }
 
-static int compileModule(char **argv, LLVMContext &Context) {
+static int compileModule(char **argv, LLVMContext &Context,
+                         std::string &OutputFilename) {
   // Load the module to be compiled...
   SMDiagnostic Err;
   std::unique_ptr<Module> M;
@@ -671,13 +679,16 @@ static int compileModule(char **argv, LLVMContext &Context) {
     Target->Options.FloatABIType = codegen::getFloatABIForCalls();
 
   // Figure out where we are going to send the output.
-  std::unique_ptr<ToolOutputFile> Out =
-      GetOutputStream(TheTarget->getName(), TheTriple.getOS(), argv[0]);
-  if (!Out) return 1;
+  std::unique_ptr<ToolOutputFile> Out = GetOutputStream(TheTriple.getOS());
+  if (!Out)
+    return 1;
 
   // Ensure the filename is passed down to CodeViewDebug.
   Target->Options.ObjectFilenameForDebug = Out->outputFilename();
 
+  // Return a copy of the output filename via the output param
+  OutputFilename = Out->outputFilename();
+
   // Tell target that this tool is not necessarily used with argument ABI
   // compliance (i.e. narrow integer argument extensions).
   Target->Options.VerifyArgABICompliance = 0;

llvm/tools/opt/optdriver.cpp

@@ -64,6 +64,7 @@ using namespace llvm;
 using namespace opt_tool;
 
 static codegen::RegisterCodeGenFlags CFG;
+static codegen::RegisterSaveStatsFlag SSF;
 
 // The OptimizationList is automatically populated with registered Passes by the
 // PassNameParser.
@@ -511,6 +512,8 @@ extern "C" int optMain(
   }
   std::unique_ptr<ToolOutputFile> RemarksFile = std::move(*RemarksFileOrErr);
 
+  codegen::MaybeEnableStatistics();
+
   // Load the input module...
   auto SetDataLayout = [&](StringRef IRTriple,
                            StringRef IRLayout) -> std::optional<std::string> {
@@ -741,14 +744,14 @@ extern "C" int optMain(
     // The user has asked to use the new pass manager and provided a pipeline
     // string. Hand off the rest of the functionality to the new code for that
     // layer.
-    return runPassPipeline(
-               argv[0], *M, TM.get(), &TLII, Out.get(), ThinLinkOut.get(),
-               RemarksFile.get(), Pipeline, PluginList, PassBuilderCallbacks,
-               OK, VK, PreserveAssemblyUseListOrder,
-               PreserveBitcodeUseListOrder, EmitSummaryIndex, EmitModuleHash,
-               EnableDebugify, VerifyDebugInfoPreserve, UnifiedLTO)
-               ? 0
-               : 1;
+    if (!runPassPipeline(
+            argv[0], *M, TM.get(), &TLII, Out.get(), ThinLinkOut.get(),
+            RemarksFile.get(), Pipeline, PluginList, PassBuilderCallbacks, OK,
+            VK, PreserveAssemblyUseListOrder, PreserveBitcodeUseListOrder,
+            EmitSummaryIndex, EmitModuleHash, EnableDebugify,
+            VerifyDebugInfoPreserve, UnifiedLTO))
+      return 1;
+    return codegen::MaybeSaveStatistics(OutputFilename, "opt");
   }
 
   if (OptLevelO0 || OptLevelO1 || OptLevelO2 || OptLevelOs || OptLevelOz ||
@@ -924,5 +927,5 @@ extern "C" int optMain(
   if (ThinLinkOut)
     ThinLinkOut->keep();
 
-  return 0;
+  return codegen::MaybeSaveStatistics(OutputFilename, "opt");
 }