[llvm] Add flags to dump IR to a file before and after LLVM passes (PR #65179)

[Min-Yih Hsu via llvm-commits]

================
@@ -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())
----------------
mshockwave wrote:

nit: `FreqTable.count(PassID)` will suffice

https://github.com/llvm/llvm-project/pull/65179