Add flags to dump IR to a file before and after LLVM passes

Summary:

LLVM offers -print-after and -print-before flags that allow you
to print IR to stderr before and after any pass you want. This can
be useful for debugging LLVM optimization issue, but is far too
noisy for large builds.

This patch adds analogous options -dump-after and -dump-before that
dump the IR to appropriately named files. In addition, it also
introduces flags -dump-after-all, -dump-before-all and -ir-dump-directory
to control where the files are written to.

Test Plan:

Included LIT tests:
```
ninja check-llvm
```

Author: Nuri Amari

llvm/include/llvm/IR/PrintPasses.h

@@ -48,6 +48,27 @@ bool shouldPrintAfterAll();
 std::vector<std::string> printBeforePasses();
 std::vector<std::string> printAfterPasses();
 
+// Returns true if dumping IR to a file before/after some pass is enabled
+// wether all passes or a specific pass.
+bool shouldDumpBeforeSomePass();
+bool shouldDumpAfterSomePass();
+
+// Returns true if we should dump IR to a file before/after a specific pass. The
+// argument should be the pass ID, e.g. "instcombine"
+bool shouldDumpBeforePass(StringRef PassID);
+bool shouldDumpAfterPass(StringRef PassID);
+
+// Returns true if we should dump IR to a file before/after all passes.
+bool shouldDumpBeforeAll();
+bool shouldDumpAfterAll();
+
+// The list of passes to dump IR to a file before/after, if we only want
+// to print before/after specific passes.
+std::vector<std::string> dumpBeforePasses();
+std::vector<std::string> dumpAfterPasses();
+
+StringRef irInstrumentationDumpDirectory();
+
 // Returns true if we should always print the entire module.
 bool forcePrintModuleIR();
 

llvm/include/llvm/Passes/StandardInstrumentations.h

@@ -69,6 +69,62 @@ class PrintIRInstrumentation {
   unsigned CurrentPassNumber = 0;
 };
 
+class DumpIRInstrumentation {
+public:
+  void registerCallbacks(PassInstrumentationCallbacks &PIC);
+
+private:
+  void dumpBeforePass(StringRef PassID, Any IR);
+  void dumpAfterPass(StringRef PassID, Any IR);
+
+  bool shouldDumpBeforePass(StringRef PassID);
+  bool shouldDumpAfterPass(StringRef PassID);
+
+  PassInstrumentationCallbacks *PIC;
+
+  // The module currently being processed in the pipeline.
+  Module const *CurrentModule = nullptr;
+
+  void pushPass(StringRef PassID, Any IR);
+  void popPass(StringRef PassID);
+
+  SmallString<16> InstrumentationDumpDirectory;
+  StringRef fetchInstrumentationDumpDirectory();
+
+  SmallString<16> fetchCurrentInstrumentationDumpFile(StringRef Suffix);
+
+  // A table to store how many times a given pass has run at the current "nested
+  // level"
+  using PassRunsFrequencyTableT = DenseMap<StringRef, unsigned>;
+  // A stack each frame of which (aside from the very first) represents a pass
+  // being run on some unit of IR. The larger, the stack grows, the smaller the
+  // unit of IR. For example, we would first push a module pass, then for each
+  // function pass in that module pass, we would push a frame and so on. This
+  // information is used to craft the output path for this logging.
+  //
+  // Each frame contains a map to track how many times a given subpass runs. For
+  // example, to keep track of how many times a function pass Foo runs within a
+  // module pass Bar. The first frame of the stack represents the module being
+  // processed rather than any particular pass. This is to create a frequency
+  // table to track module level pass run counts without having to special case
+  // that logic.
+  //
+  // When a change in the module being processed is detected, this first frame
+  // is updated accordingly.
+
+  struct PipelineStateStackFrame {
+    StringRef PassID;
+    PassRunsFrequencyTableT FreqTable;
+    unsigned int PassCount;
+
+    PipelineStateStackFrame(StringRef PassID) : PassID{PassID}, PassCount{0} {}
+  };
+
+  using PipelineStateStackT = SmallVector<PipelineStateStackFrame>;
+
+  PipelineStateStackT PipelineStateStack;
+};
+
 class OptNoneInstrumentation {
 public:
   OptNoneInstrumentation(bool DebugLogging) : DebugLogging(DebugLogging) {}
@@ -555,6 +611,7 @@ class PrintCrashIRInstrumentation {
 /// instrumentations and manages their state (if any).
 class StandardInstrumentations {
   PrintIRInstrumentation PrintIR;
+  DumpIRInstrumentation DumpIR;
   PrintPassInstrumentation PrintPass;
   TimePassesHandler TimePasses;
   TimeProfilingPassesHandler TimeProfilingPasses;

llvm/lib/IR/PrintPasses.cpp

@@ -33,6 +33,31 @@ static cl::opt<bool> PrintAfterAll("print-after-all",
                                    llvm::cl::desc("Print IR after each pass"),
                                    cl::init(false), cl::Hidden);
 
+static cl::list<std::string>
+    DumpBefore("dump-before",
+               llvm::cl::desc("Dump IR to a file before specified passes"),
+               cl::CommaSeparated, cl::Hidden);
+
+static cl::list<std::string>
+    DumpAfter("dump-after",
+              llvm::cl::desc("Dump IR to a file after specified passes"),
+              cl::CommaSeparated, cl::Hidden);
+
+static cl::opt<bool>
+    DumpBeforeAll("dump-before-all",
+                  llvm::cl::desc("Dump IR to a file before each pass"),
+                  cl::init(false), cl::Hidden);
+static cl::opt<bool>
+    DumpAfterAll("dump-after-all",
+                 llvm::cl::desc("Dump IR to a file after each pass"),
+                 cl::init(false), cl::Hidden);
+
+static cl::opt<std::string> DumpDirectory(
+    "ir-dump-directory",
+    llvm::cl::desc("A directory to dump IR log files into before and after "
+                   "passes as specified using -dump-before / -dump-after"),
+    cl::init(""), cl::Hidden, cl::value_desc("filename"));
+
 // Print out the IR after passes, similar to -print-after-all except that it
 // only prints the IR after passes that change the IR. Those passes that do not
 // make changes to the IR are reported as not making any changes. In addition,
@@ -139,6 +164,35 @@ std::vector<std::string> llvm::printAfterPasses() {
   return std::vector<std::string>(PrintAfter);
 }
 
+bool llvm::shouldDumpBeforeSomePass() {
+  return DumpBeforeAll || !DumpBefore.empty();
+}
+
+bool llvm::shouldDumpAfterSomePass() {
+  return DumpAfterAll || !DumpAfter.empty();
+}
+
+bool llvm::shouldDumpBeforeAll() { return DumpBeforeAll; }
+
+bool llvm::shouldDumpAfterAll() { return DumpAfterAll; }
+
+bool llvm::shouldDumpBeforePass(StringRef PassID) {
+  return DumpBeforeAll || llvm::is_contained(DumpBefore, PassID);
+}
+
+bool llvm::shouldDumpAfterPass(StringRef PassID) {
+  return DumpAfterAll || llvm::is_contained(DumpAfter, PassID);
+}
+std::vector<std::string> llvm::dumpBeforePasses() {
+  return std::vector<std::string>(DumpBefore);
+}
+
+std::vector<std::string> llvm::dumpAfterPasses() {
+  return std::vector<std::string>(DumpAfter);
+}
+
+StringRef llvm::irInstrumentationDumpDirectory() { return DumpDirectory; }
+
 bool llvm::forcePrintModuleIR() { return PrintModuleScope; }
 
 bool llvm::isPassInPrintList(StringRef PassName) {

llvm/lib/Passes/PassBuilder.cpp

@@ -386,7 +386,8 @@ AnalysisKey NoOpLoopAnalysis::Key;
 /// We currently only use this for --print-before/after.
 bool shouldPopulateClassToPassNames() {
   return PrintPipelinePasses || !printBeforePasses().empty() ||
-         !printAfterPasses().empty() || !isFilterPassesEmpty();
+         !printAfterPasses().empty() || !isFilterPassesEmpty() ||
+         !dumpAfterPasses().empty() || !dumpBeforePasses().empty();
 }
 
 // A pass for testing -print-on-crash.

llvm/lib/Passes/StandardInstrumentations.cpp

@@ -33,6 +33,7 @@
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/GraphWriter.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/Signals.h"
@@ -830,6 +831,182 @@ void PrintIRInstrumentation::registerCallbacks(
   }
 }
 
+void DumpIRInstrumentation::registerCallbacks(
+    PassInstrumentationCallbacks &PIC) {
+
+  if (!(shouldDumpBeforeSomePass() || shouldDumpAfterSomePass()))
+    return;
+
+  this->PIC = &PIC;
+
+  PIC.registerBeforeNonSkippedPassCallback(
+      [this](StringRef P, Any IR) { this->pushPass(P, IR); });
+
+  if (shouldDumpBeforeSomePass())
+    PIC.registerBeforeNonSkippedPassCallback(
+        [this](StringRef P, Any IR) { this->dumpBeforePass(P, IR); });
+
+  if (shouldDumpAfterSomePass()) {
+    PIC.registerAfterPassCallback(
+        [this](StringRef P, Any IR, const PreservedAnalyses &) {
+          this->dumpAfterPass(P, IR);
+        });
+  }
+
+  // It is important the the "popPass" callback fires after the dumpAfterPass
+  // callback
+  PIC.registerAfterPassCallback(
+      [this](StringRef P, Any IR, const PreservedAnalyses &) {
+        this->popPass(P);
+      });
+}
+
+void DumpIRInstrumentation::dumpBeforePass(StringRef PassID, Any IR) {
+  if (isIgnored(PassID))
+    return;
+
+  if (!shouldDumpBeforePass(PassID)) {
+    return;
+  }
+
+  SmallString<16> OutputPath =
+      fetchCurrentInstrumentationDumpFile("-before.ll");
+  std::error_code EC;
+  llvm::raw_fd_ostream OS(OutputPath, EC, llvm::sys::fs::CD_CreateAlways);
+  if (EC)
+    report_fatal_error(Twine("Failed to open ") + OutputPath +
+                       " to support dump-before: " + EC.message());
+
+  OS << "*** IR Dump Before " << PassID << " on " << getIRName(IR) << " ***\n";
+  unwrapAndPrint(OS, IR);
+}
+
+void DumpIRInstrumentation::dumpAfterPass(StringRef PassID, Any IR) {
+  if (isIgnored(PassID))
+    return;
+
+  if (!shouldDumpAfterPass(PassID))
+    return;
+
+  SmallString<16> OutputPath = fetchCurrentInstrumentationDumpFile("-after.ll");
+  std::error_code EC;
+  llvm::raw_fd_ostream OS(OutputPath, EC, llvm::sys::fs::CD_CreateAlways);
+  if (EC)
+    report_fatal_error(Twine("Failed to open ") + OutputPath +
+                       " to support -dump-after: " + EC.message());
+
+  OS << "*** IR Dump After " << PassID << " on " << getIRName(IR) << " ***\n";
+  unwrapAndPrint(OS, IR);
+}
+
+bool DumpIRInstrumentation::shouldDumpBeforePass(StringRef PassID) {
+  if (shouldDumpBeforeAll())
+    return true;
+
+  StringRef PassName = PIC->getPassNameForClassName(PassID);
+  return is_contained(dumpBeforePasses(), PassName);
+}
+
+bool DumpIRInstrumentation::shouldDumpAfterPass(StringRef PassID) {
+  if (shouldDumpAfterAll())
+    return true;
+
+  StringRef PassName = PIC->getPassNameForClassName(PassID);
+  return is_contained(dumpAfterPasses(), PassName);
+}
+
+void DumpIRInstrumentation::pushPass(StringRef PassID, Any IR) {
+  const Module *M = unwrapModule(IR);
+  if (CurrentModule != M) {
+    // If currentModule is nullptr, or is not equal to M, we are starting to
+    // process a new module.
+
+    // The first frame of the stack should maintain a frequency table
+    // for module level passes.
+    PipelineStateStack.clear();
+    PipelineStateStack.push_back(PipelineStateStackFrame(M->getName()));
+
+    CurrentModule = M;
+  }
+
+  PassRunsFrequencyTableT &FreqTable = PipelineStateStack.back().FreqTable;
+  if (FreqTable.find(PassID) == FreqTable.end())
+    FreqTable[PassID] = 0;
+
+  PipelineStateStack.push_back(PipelineStateStackFrame(PassID));
+}
+
+void DumpIRInstrumentation::popPass(StringRef PassID) {
+  PipelineStateStack.pop_back();
+  assert(!PipelineStateStack.empty());
+
+  PassRunsFrequencyTableT &FreqTable = PipelineStateStack.back().FreqTable;
+  assert(FreqTable.find(PassID) != FreqTable.end());
+  FreqTable[PassID]++;
+  PipelineStateStack.back().PassCount++;
+}
+
+StringRef DumpIRInstrumentation::fetchInstrumentationDumpDirectory() {
+  if (!InstrumentationDumpDirectory.empty())
+    return InstrumentationDumpDirectory;
+
+  if (!irInstrumentationDumpDirectory().empty())
+    return irInstrumentationDumpDirectory();
+
+  std::error_code EC =
+      sys::fs::createUniqueDirectory("dumped-ir", InstrumentationDumpDirectory);
+  if (EC)
+    report_fatal_error(
+        Twine("Failed to create unique directory for IR dumping: ") +
+        EC.message());
+
+  return InstrumentationDumpDirectory;
+}
+
+SmallString<16>
+DumpIRInstrumentation::fetchCurrentInstrumentationDumpFile(StringRef Suffix) {
+  SmallString<16> OutputPath;
+  sys::path::append(OutputPath, fetchInstrumentationDumpDirectory());
+  assert(CurrentModule);
+  sys::path::append(OutputPath, CurrentModule->getName());
+  auto *StateStackIt = PipelineStateStack.begin();
+  // Skip over the first frame in the stack which represents the module being
+  // processed
+  for (++StateStackIt; StateStackIt != PipelineStateStack.end();
+       ++StateStackIt) {
+    SmallString<8> PathComponentName;
+    // Check the previous frame's pass count to see how many passes of
+    // any kind have run on this "nesting level".
+    unsigned int PassCount = (StateStackIt - 1)->PassCount;
+    PathComponentName += std::to_string(PassCount);
+    PathComponentName += ".";
+    // Check the previous frame's frequency table to see how many times
+    // this pass has run at this "nesting level".
+    PassRunsFrequencyTableT &FreqTable = (StateStackIt - 1)->FreqTable;
+    StringRef FramePassID = StateStackIt->PassID;
+    PathComponentName += FramePassID;
+    PathComponentName += ".";
+    assert(FreqTable.find(FramePassID) != FreqTable.end() &&
+           "DumpIRInstrumentation pass frequency table missing entry");
+    // Make sure the uint is converted to a character and not interpretted as
+    // one
+    PathComponentName += std::to_string(FreqTable[FramePassID]);
+    sys::path::append(OutputPath, PathComponentName);
+  }
+  OutputPath += Suffix;
+
+  // Make sure the directory we wish to write our log file into exists.
+  StringRef ParentDirectory = sys::path::parent_path(OutputPath);
+  if (!ParentDirectory.empty()) {
+    if (auto EC = llvm::sys::fs::create_directories(ParentDirectory)) {
+      report_fatal_error(Twine("Failed to create directory '") +
+                         ParentDirectory + "' :" + EC.message());
+    }
+  }
+
+  return OutputPath;
+}
+
 void OptNoneInstrumentation::registerCallbacks(
     PassInstrumentationCallbacks &PIC) {
   PIC.registerShouldRunOptionalPassCallback(
@@ -2288,6 +2465,7 @@ void PrintCrashIRInstrumentation::registerCallbacks(
 void StandardInstrumentations::registerCallbacks(
     PassInstrumentationCallbacks &PIC, ModuleAnalysisManager *MAM) {
   PrintIR.registerCallbacks(PIC);
+  DumpIR.registerCallbacks(PIC);
   PrintPass.registerCallbacks(PIC);
   TimePasses.registerCallbacks(PIC);
   OptNone.registerCallbacks(PIC);