[opt] Add --save-stats option (llvm#167304)

This patch adds a Clang-compatible --save-stats option to opt, to
provide an easy to use way to save LLVM statistics files when working
with opt on the middle end.

This is a follow up on the addition to `llc`:
llvm#163967

Like on Clang, one can specify --save-stats, --save-stats=cwd, and
--save-stats=obj with the same semantics and JSON format. The
pre-existing --stats option is not affected.

The implementation extracts the flag and its methods into the common
`CodeGen/CommandFlags` as `LLVM_ABI`, using a new registration class to
conservatively enable opt-in rather than let all tools take it. Its only
needed for llc and opt for now. Then it refactors llc and adds support
for opt.

(cherry-pick 35ffe10)

Author: tomershafir

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

@@ -298,6 +298,7 @@ Changes to the LLVM tools
 * 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

@@ -1,5 +1,8 @@
 ; 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

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

@@ -66,6 +66,7 @@
 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
@@ -203,20 +204,6 @@ static cl::opt<std::string> RemarksFormat(
     cl::desc("The format used for serializing remarks (default: YAML)"),
     cl::value_desc("format"), cl::init("yaml"));
 
-enum SaveStatsMode { None, Cwd, Obj };
-
-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);
-
 static cl::opt<bool> EnableNewPassManager(
     "enable-new-pm", cl::desc("Enable the new pass manager"), cl::init(false));
 
@@ -391,45 +378,6 @@ static void verifyCASOptions(const Triple &TheTriple) {
 }
 // END MCCAS
 
-static int MaybeEnableStats() {
-  if (SaveStats == SaveStatsMode::None)
-    return 0;
-
-  llvm::EnableStatistics(false);
-  return 0;
-}
-
-static int MaybeSaveStats(std::string &&OutputFilename) {
-  if (SaveStats == SaveStatsMode::None)
-    return 0;
-
-  SmallString<128> StatsFilename;
-  if (SaveStats == SaveStatsMode::Obj) {
-    StatsFilename = OutputFilename;
-    llvm::sys::path::remove_filename(StatsFilename);
-  } else {
-    assert(SaveStats == 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(), "llc")
-        << "Unable to open statistics file: " << EC.message() << "\n";
-    return 1;
-  }
-
-  llvm::PrintStatisticsJSON(*StatsOS);
-  return 0;
-}
-
 // main - Entry point for the llc compiler.
 //
 int main(int argc, char **argv) {
@@ -507,8 +455,7 @@ int main(int argc, char **argv) {
     reportError(std::move(E), RemarksFilename);
   std::unique_ptr<ToolOutputFile> RemarksFile = std::move(*RemarksFileOrErr);
 
-  if (int RetVal = MaybeEnableStats())
-    return RetVal;
+  codegen::MaybeEnableStatistics();
   std::string OutputFilename;
 
   if (InputLanguage != "" && InputLanguage != "ir" && InputLanguage != "mir")
@@ -523,7 +470,7 @@ int main(int argc, char **argv) {
   if (RemarksFile)
     RemarksFile->keep();
 
-  return MaybeSaveStats(std::move(OutputFilename));
+  return codegen::MaybeSaveStatistics(OutputFilename, "llc");
 }
 
 static bool addPass(PassManagerBase &PM, const char *argv0,

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");
 }