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

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) const;
+  bool shouldDumpAfterPass(StringRef PassID) const;
+
+  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 = StringMap<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,181 @@ 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) const {
+  if (shouldDumpBeforeAll())
+    return true;
+
+  StringRef PassName = PIC->getPassNameForClassName(PassID);
+  return is_contained(dumpBeforePasses(), PassName);
+}
+
+bool DumpIRInstrumentation::shouldDumpAfterPass(StringRef PassID) const {
+  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.count(PassID))
+    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,
+                    sys::path::relative_path(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 +2464,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);