[NFC] Add raw_ostream parameter to printIR routines

[jamieschmeiser]

Summary

This is a non-functional-change to generalize the printIR routines so that
the output can be saved and manipulated rather than being directly output
to dbgs(). This is a prerequisite change for many upcoming changes that
allow new ways of examining changes made to the IR in the new pass manager.
These changes will be presented at the developer's conference in Oct, 2020
(see https://hotcrp.llvm.org/usllvm2020/paper/29 for more details).

[aeubanks]

lgtm

llvm/lib/Passes/StandardInstrumentations.cpp

if (any_isa<const Loop *>(IR)) {
     raw_string_ostream ss(LoopName);
     L->getHeader()->printAsOperand(ss, false);
     return std::make_pair(M, formatv(" (loop: {0})", ss.str()).str());
   }
 
   llvm_unreachable("Unknown IR unit");
 }
 
-void printIR(const Function *F, StringRef Banner, StringRef Extra = StringRef(),
-             bool Brief = false) {
+void printIR(raw_ostream &OS, const Function *F, StringRef Banner,
+             StringRef Extra = StringRef(), bool Brief = false) {
   if (Brief) {
-    dbgs() << F->getName() << '\n';
+    OS << F->getName() << '\n';
     return;
   }
 
   if (!llvm::isFunctionInPrintList(F->getName()))
     return;
-  dbgs() << Banner << Extra << "\n" << static_cast<const Value &>(*F);
+  OS << Banner << Extra << "\n" << static_cast<const Value &>(*F);
 }
 
-void printIR(const Module *M, StringRef Banner, StringRef Extra = StringRef(),
-             bool Brief = false) {
+void printIR(raw_ostream &OS, const Module *M, StringRef Banner,
+             StringRef Extra = StringRef(), bool Brief = false) {
   if (Brief) {
-    dbgs() << M->getName() << '\n';
+    OS << M->getName() << '\n';
     return;
   }
 
   if (llvm::isFunctionInPrintList("*") || llvm::forcePrintModuleIR()) {
-    dbgs() << Banner << Extra << "\n";
-    M->print(dbgs(), nullptr, false);
+    OS << Banner << Extra << "\n";
+    M->print(OS, nullptr, false);
   } else {
     for (const auto &F : M->functions()) {
-      printIR(&F, Banner, Extra);
+      printIR(OS, &F, Banner, Extra);
     }
   }
 }
 
-void printIR(const LazyCallGraph::SCC *C, StringRef Banner,
+void printIR(raw_ostream &OS, const LazyCallGraph::SCC *C, StringRef Banner,
              StringRef Extra = StringRef(), bool Brief = false) {
   if (Brief) {
-    dbgs() << *C << '\n';
+    OS << *C << '\n';
     return;
   }
 
   bool BannerPrinted = false;
   for (const LazyCallGraph::Node &N : *C) {
     const Function &F = N.getFunction();
     if (!F.isDeclaration() && llvm::isFunctionInPrintList(F.getName())) {
       if (!BannerPrinted) {
-        dbgs() << Banner << Extra << "\n";
+        OS << Banner << Extra << "\n";
         BannerPrinted = true;
       }
-      F.print(dbgs());
+      F.print(OS);
     }
   }
 }
 
-void printIR(const Loop *L, StringRef Banner, bool Brief = false) {
+void printIR(raw_ostream &OS, const Loop *L, StringRef Banner,
+             bool Brief = false) {
   if (Brief) {
-    dbgs() << *L;
+    OS << *L;
     return;
   }
 
   const Function *F = L->getHeader()->getParent();
   if (!llvm::isFunctionInPrintList(F->getName()))
     return;
-  llvm::printLoop(const_cast<Loop &>(*L), dbgs(), std::string(Banner));
+  llvm::printLoop(const_cast<Loop &>(*L), OS, std::string(Banner));
 }
 
 /// Generic IR-printing helper that unpacks a pointer to IRUnit wrapped into
 /// llvm::Any and does actual print job.
-void unwrapAndPrint(Any IR, StringRef Banner, bool ForceModule = false,
-                    bool Brief = false) {
+void unwrapAndPrint(raw_ostream &OS, Any IR, StringRef Banner,
+                    bool ForceModule = false, bool Brief = false) {
   if (ForceModule) {
     if (auto UnwrappedModule = unwrapModule(IR))
-      printIR(UnwrappedModule->first, Banner, UnwrappedModule->second);
+      printIR(OS, UnwrappedModule->first, Banner, UnwrappedModule->second);
     return;
   }
 
   if (any_isa<const Module *>(IR)) {
     const Module *M = any_cast<const Module *>(IR);
     assert(M && "module should be valid for printing");
-    printIR(M, Banner, "", Brief);
+    printIR(OS, M, Banner, "", Brief);
     return;
   }
 
   if (any_isa<const Function *>(IR)) {
     const Function *F = any_cast<const Function *>(IR);
     assert(F && "function should be valid for printing");
-    printIR(F, Banner, "", Brief);
+    printIR(OS, F, Banner, "", Brief);
     return;
   }
 
   if (any_isa<const LazyCallGraph::SCC *>(IR)) {
     const LazyCallGraph::SCC *C = any_cast<const LazyCallGraph::SCC *>(IR);
     assert(C && "scc should be valid for printing");
     std::string Extra = std::string(formatv(" (scc: {0})", C->getName()));
-    printIR(C, Banner, Extra, Brief);
+    printIR(OS, C, Banner, Extra, Brief);
     return;
   }
 
   if (any_isa<const Loop *>(IR)) {
     const Loop *L = any_cast<const Loop *>(IR);
     assert(L && "Loop should be valid for printing");
-    printIR(L, Banner, Brief);
+    printIR(OS, L, Banner, Brief);
     return;
   }
   llvm_unreachable("Unknown wrapped IR type");
 }
 
 } // namespace
 
 PrintIRInstrumentation::~PrintIRInstrumentation() {
void PrintIRInstrumentation::printBeforePass(StringRef PassID, Any IR) {
   // for all print operations that has not happen yet.
   if (StoreModuleDesc && llvm::shouldPrintAfterPass(PassID))
     pushModuleDesc(PassID, IR);
 
   if (!llvm::shouldPrintBeforePass(PassID))
     return;
 
   SmallString<20> Banner = formatv("*** IR Dump Before {0} ***", PassID);
-  unwrapAndPrint(IR, Banner, llvm::forcePrintModuleIR());
+  unwrapAndPrint(dbgs(), IR, Banner, llvm::forcePrintModuleIR());
   return;
 }
 
 void PrintIRInstrumentation::printAfterPass(StringRef PassID, Any IR) {
   if (PassID.startswith("PassManager<") || PassID.contains("PassAdaptor<"))
     return;
 
   if (!llvm::shouldPrintAfterPass(PassID))
     return;
 
   if (StoreModuleDesc)
     popModuleDesc(PassID);
 
   SmallString<20> Banner = formatv("*** IR Dump After {0} ***", PassID);
-  unwrapAndPrint(IR, Banner, llvm::forcePrintModuleIR());
+  unwrapAndPrint(dbgs(), IR, Banner, llvm::forcePrintModuleIR());
 }
 
 void PrintIRInstrumentation::printAfterPassInvalidated(StringRef PassID) {
   if (!StoreModuleDesc || !llvm::shouldPrintAfterPass(PassID))
     return;
 
   if (PassID.startswith("PassManager<") || PassID.contains("PassAdaptor<"))
     return;
 
   const Module *M;
   std::string Extra;
   StringRef StoredPassID;
   std::tie(M, Extra, StoredPassID) = popModuleDesc(PassID);
   // Additional filtering (e.g. -filter-print-func) can lead to module
   // printing being skipped.
   if (!M)
     return;
 
   SmallString<20> Banner =
       formatv("*** IR Dump After {0} *** invalidated: ", PassID);
-  printIR(M, Banner, Extra);
+  printIR(dbgs(), M, Banner, Extra);
 }
 
 void PrintIRInstrumentation::registerCallbacks(
     PassInstrumentationCallbacks &PIC) {
   // BeforePass callback is not just for printing, it also saves a Module
   // for later use in AfterPassInvalidated.
   StoreModuleDesc = llvm::forcePrintModuleIR() && llvm::shouldPrintAfterPass();
   if (llvm::shouldPrintBeforePass() || StoreModuleDesc)
if (!DebugPMVerbose)
     SpecialPasses.emplace_back("PassAdaptor");
 
   PIC.registerBeforeSkippedPassCallback(
       [SpecialPasses](StringRef PassID, Any IR) {
         assert(!isSpecialPass(PassID, SpecialPasses) &&
                "Unexpectedly skipping special pass");
 
         dbgs() << "Skipping pass: " << PassID << " on ";
-        unwrapAndPrint(IR, "", false, true);
+        unwrapAndPrint(dbgs(), IR, "", false, true);
       });
 
   PIC.registerBeforeNonSkippedPassCallback(
       [SpecialPasses](StringRef PassID, Any IR) {
         if (isSpecialPass(PassID, SpecialPasses))
           return;
 
         dbgs() << "Running pass: " << PassID << " on ";
-        unwrapAndPrint(IR, "", false, true);
+        unwrapAndPrint(dbgs(), IR, "", false, true);
       });
 
   PIC.registerBeforeAnalysisCallback([](StringRef PassID, Any IR) {
     dbgs() << "Running analysis: " << PassID << " on ";
-    unwrapAndPrint(IR, "", false, true);
+    unwrapAndPrint(dbgs(), IR, "", false, true);
   });
 }
 
 void StandardInstrumentations::registerCallbacks(
     PassInstrumentationCallbacks &PIC) {
   PrintIR.registerCallbacks(PIC);
   PrintPass.registerCallbacks(PIC);
   TimePasses.registerCallbacks(PIC);
   OptNone.registerCallbacks(PIC);
 }