[NewPM] -print-module-scope -print-after now prints module even after invalidated Loop/SCC

-print-after IR printing generally can not print the IR unit (Loop or SCC)
which has just been invalidated by the pass. However, when working in -print-module-scope
mode even if Loop was invalidated there is still a valid module that we can print.

Since we can not access invalidated IR unit from AfterPassInvalidated instrumentation
point we can remember the module to be printed *before* pass. This change introduces
BeforePass instrumentation that stores all the information required for module printing
into the stack and then after pass (in AfterPassInvalidated) just print whatever
has been placed on stack.

Reviewed By: philip.pfaffe
Differential Revision: https://reviews.llvm.org/D55278

llvm-svn: 349896

4 files changed, 214 insertions, 72 deletions

diff --git a/llvm/include/llvm/Passes/StandardInstrumentations.h b/llvm/include/llvm/Passes/StandardInstrumentations.h
index b6a73b0f487..8c6f5e1e22f 100644
--- a/llvm/include/llvm/Passes/StandardInstrumentations.h
+++ b/llvm/include/llvm/Passes/StandardInstrumentations.h
@@ -16,14 +16,48 @@
 #ifndef LLVM_PASSES_STANDARDINSTRUMENTATIONS_H
 #define LLVM_PASSES_STANDARDINSTRUMENTATIONS_H
 
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/PassInstrumentation.h"
 #include "llvm/IR/PassTimingInfo.h"
 
+#include <string>
+#include <utility>
+
 namespace llvm {
 
+class Module;
+
+/// Instrumentation to print IR before/after passes.
+///
+/// Needs state to be able to print module after pass that invalidates IR unit
+/// (typically Loop or SCC).
+class PrintIRInstrumentation {
+public:
+  PrintIRInstrumentation() = default;
+  ~PrintIRInstrumentation();
+
+  void registerCallbacks(PassInstrumentationCallbacks &PIC);
+
+private:
+  bool printBeforePass(StringRef PassID, Any IR);
+  void printAfterPass(StringRef PassID, Any IR);
+  void printAfterPassInvalidated(StringRef PassID);
+
+  using PrintModuleDesc = std::tuple<const Module *, std::string, StringRef>;
+
+  void pushModuleDesc(StringRef PassID, Any IR);
+  PrintModuleDesc popModuleDesc(StringRef PassID);
+
+  /// Stack of Module description, enough to print the module after a given
+  /// pass.
+  SmallVector<PrintModuleDesc, 2> ModuleDescStack;
+  bool StoreModuleDesc = false;
+};
+
 /// This class provides an interface to register all the standard pass
 /// instrumentations and manages their state (if any).
 class StandardInstrumentations {
+  PrintIRInstrumentation PrintIR;
   TimePassesHandler TimePasses;
 
 public:
diff --git a/llvm/lib/Passes/StandardInstrumentations.cpp b/llvm/lib/Passes/StandardInstrumentations.cpp
index 765ffe62662..a1dfc39d472 100644
--- a/llvm/lib/Passes/StandardInstrumentations.cpp
+++ b/llvm/lib/Passes/StandardInstrumentations.cpp
@@ -14,6 +14,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Passes/StandardInstrumentations.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/Analysis/LoopInfo.h"
@@ -28,112 +29,215 @@
 using namespace llvm;
 
 namespace {
-namespace PrintIR {
 
-//===----------------------------------------------------------------------===//
-// IR-printing instrumentation
-//===----------------------------------------------------------------------===//
+/// Extracting Module out of \p IR unit. Also fills a textual description
+/// of \p IR for use in header when printing.
+Optional<std::pair<const Module *, std::string>> unwrapModule(Any IR) {
+  if (any_isa<const Module *>(IR))
+    return std::make_pair(any_cast<const Module *>(IR), std::string());
 
-/// Generic IR-printing helper that unpacks a pointer to IRUnit wrapped into
-/// llvm::Any and does actual print job.
-void unwrapAndPrint(StringRef Banner, Any IR) {
-  SmallString<40> Extra{"\n"};
-  const Module *M = nullptr;
-  if (any_isa<const Module *>(IR)) {
-    M = any_cast<const Module *>(IR);
-  } else if (any_isa<const Function *>(IR)) {
+  if (any_isa<const Function *>(IR)) {
     const Function *F = any_cast<const Function *>(IR);
     if (!llvm::isFunctionInPrintList(F->getName()))
-      return;
-    if (!llvm::forcePrintModuleIR()) {
-      dbgs() << Banner << Extra << static_cast<const Value &>(*F);
-      return;
-    }
-    M = F->getParent();
-    Extra = formatv(" (function: {0})\n", F->getName());
-  } else if (any_isa<const LazyCallGraph::SCC *>(IR)) {
+      return None;
+    const Module *M = F->getParent();
+    return std::make_pair(M, formatv(" (function: {0})", F->getName()).str());
+  }
+
+  if (any_isa<const LazyCallGraph::SCC *>(IR)) {
     const LazyCallGraph::SCC *C = any_cast<const LazyCallGraph::SCC *>(IR);
-    if (!llvm::forcePrintModuleIR()) {
-      Extra = formatv(" (scc: {0})\n", C->getName());
-      bool BannerPrinted = false;
-      for (const LazyCallGraph::Node &N : *C) {
-        const Function &F = N.getFunction();
-        if (!F.isDeclaration() && isFunctionInPrintList(F.getName())) {
-          if (!BannerPrinted) {
-            dbgs() << Banner << Extra;
-            BannerPrinted = true;
-          }
-          F.print(dbgs());
-        }
-      }
-      return;
-    }
     for (const LazyCallGraph::Node &N : *C) {
       const Function &F = N.getFunction();
       if (!F.isDeclaration() && isFunctionInPrintList(F.getName())) {
-        M = F.getParent();
-        break;
+        const Module *M = F.getParent();
+        return std::make_pair(M, formatv(" (scc: {0})", C->getName()).str());
       }
     }
-    if (!M)
-      return;
-    Extra = formatv(" (for scc: {0})\n", C->getName());
-  } else if (any_isa<const Loop *>(IR)) {
+    return None;
+  }
+
+  if (any_isa<const Loop *>(IR)) {
     const Loop *L = any_cast<const Loop *>(IR);
     const Function *F = L->getHeader()->getParent();
     if (!isFunctionInPrintList(F->getName()))
-      return;
-    if (!llvm::forcePrintModuleIR()) {
-      llvm::printLoop(const_cast<Loop &>(*L), dbgs(), Banner);
-      return;
-    }
-    M = F->getParent();
-    {
-      std::string LoopName;
-      raw_string_ostream ss(LoopName);
-      L->getHeader()->printAsOperand(ss, false);
-      Extra = formatv(" (loop: {0})\n", ss.str());
+      return None;
+    const Module *M = F->getParent();
+    std::string LoopName;
+    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 Module *M, StringRef Banner, StringRef Extra = StringRef()) {
+  dbgs() << Banner << Extra << "\n";
+  M->print(dbgs(), nullptr, false);
+}
+void printIR(const Function *F, StringRef Banner,
+             StringRef Extra = StringRef()) {
+  if (!llvm::isFunctionInPrintList(F->getName()))
+    return;
+  dbgs() << Banner << Extra << "\n" << static_cast<const Value &>(*F);
+}
+void printIR(const LazyCallGraph::SCC *C, StringRef Banner,
+             StringRef Extra = StringRef()) {
+  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";
+        BannerPrinted = true;
+      }
+      F.print(dbgs());
     }
   }
-  if (M) {
-    dbgs() << Banner << Extra;
-    M->print(dbgs(), nullptr, false);
-  } else {
-    llvm_unreachable("Unknown wrapped IR type");
+}
+void printIR(const Loop *L, StringRef Banner) {
+  const Function *F = L->getHeader()->getParent();
+  if (!llvm::isFunctionInPrintList(F->getName()))
+    return;
+  llvm::printLoop(const_cast<Loop &>(*L), dbgs(), 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) {
+  if (ForceModule) {
+    if (auto UnwrappedModule = unwrapModule(IR))
+      printIR(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);
+    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);
+    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 = formatv(" (scc: {0})", C->getName());
+    printIR(C, Banner, Extra);
+    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);
+    return;
+  }
+  llvm_unreachable("Unknown wrapped IR type");
 }
 
-bool printBeforePass(StringRef PassID, Any IR) {
-  if (!llvm::shouldPrintBeforePass(PassID))
-    return true;
+} // namespace
 
+PrintIRInstrumentation::~PrintIRInstrumentation() {
+  assert(ModuleDescStack.empty() && "ModuleDescStack is not empty at exit");
+}
+
+void PrintIRInstrumentation::pushModuleDesc(StringRef PassID, Any IR) {
+  assert(StoreModuleDesc);
+  const Module *M = nullptr;
+  std::string Extra;
+  if (auto UnwrappedModule = unwrapModule(IR))
+    std::tie(M, Extra) = UnwrappedModule.getValue();
+  ModuleDescStack.emplace_back(M, Extra, PassID);
+}
+
+PrintIRInstrumentation::PrintModuleDesc
+PrintIRInstrumentation::popModuleDesc(StringRef PassID) {
+  assert(!ModuleDescStack.empty() && "empty ModuleDescStack");
+  PrintModuleDesc ModuleDesc = ModuleDescStack.pop_back_val();
+  assert(std::get<2>(ModuleDesc).equals(PassID) && "malformed ModuleDescStack");
+  return ModuleDesc;
+}
+
+bool PrintIRInstrumentation::printBeforePass(StringRef PassID, Any IR) {
   if (PassID.startswith("PassManager<") || PassID.contains("PassAdaptor<"))
     return true;
 
+  // Saving Module for AfterPassInvalidated operations.
+  // Note: here we rely on a fact that we do not change modules while
+  // traversing the pipeline, so the latest captured module is good
+  // for all print operations that has not happen yet.
+  if (StoreModuleDesc && llvm::shouldPrintAfterPass(PassID))
+    pushModuleDesc(PassID, IR);
+
+  if (!llvm::shouldPrintBeforePass(PassID))
+    return true;
+
   SmallString<20> Banner = formatv("*** IR Dump Before {0} ***", PassID);
-  unwrapAndPrint(Banner, IR);
+  unwrapAndPrint(IR, Banner, llvm::forcePrintModuleIR());
   return true;
 }
 
-void printAfterPass(StringRef PassID, Any IR) {
+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());
+}
+
+void PrintIRInstrumentation::printAfterPassInvalidated(StringRef PassID) {
+  if (!StoreModuleDesc || !llvm::shouldPrintAfterPass(PassID))
+    return;
+
   if (PassID.startswith("PassManager<") || PassID.contains("PassAdaptor<"))
     return;
 
-  SmallString<20> Banner = formatv("*** IR Dump After {0} ***", PassID);
-  unwrapAndPrint(Banner, IR);
-  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);
+}
+
+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)
+    PIC.registerBeforePassCallback(
+        [this](StringRef P, Any IR) { return this->printBeforePass(P, IR); });
+
+  if (llvm::shouldPrintAfterPass()) {
+    PIC.registerAfterPassCallback(
+        [this](StringRef P, Any IR) { this->printAfterPass(P, IR); });
+    PIC.registerAfterPassInvalidatedCallback(
+        [this](StringRef P) { this->printAfterPassInvalidated(P); });
+  }
 }
-} // namespace PrintIR
-} // namespace
 
 void StandardInstrumentations::registerCallbacks(
     PassInstrumentationCallbacks &PIC) {
-  if (llvm::shouldPrintBeforePass())
-    PIC.registerBeforePassCallback(PrintIR::printBeforePass);
-  if (llvm::shouldPrintAfterPass())
-    PIC.registerAfterPassCallback(PrintIR::printAfterPass);
+  PrintIR.registerCallbacks(PIC);
   TimePasses.registerCallbacks(PIC);
 }
diff --git a/llvm/test/Other/loop-deletion-printer.ll b/llvm/test/Other/loop-deletion-printer.ll
index d344568506c..6cb29ae374f 100644
--- a/llvm/test/Other/loop-deletion-printer.ll
+++ b/llvm/test/Other/loop-deletion-printer.ll
@@ -5,10 +5,14 @@
 ; RUN:     -passes=loop-instsimplify -print-after-all  2>&1 | FileCheck %s -check-prefix=SIMPLIFY
 ; RUN: opt < %s -disable-output \
 ; RUN:     -passes=loop-deletion,loop-instsimplify -print-after-all  2>&1 | FileCheck %s -check-prefix=DELETED
+; RUN: opt < %s -disable-output \
+; RUN:     -passes=loop-deletion,loop-instsimplify -print-after-all -print-module-scope  2>&1 | FileCheck %s -check-prefix=DELETED-BUT-PRINTED
 ;
 ; SIMPLIFY: IR Dump {{.*}} LoopInstSimplifyPass
 ; DELETED-NOT: IR Dump {{.*}}LoopInstSimplifyPass
 ; DELETED-NOT: IR Dump {{.*}}LoopDeletionPass
+; DELETED-BUT-PRINTED: IR Dump {{.*}}LoopDeletionPass {{.*invalidated:}}
+; DELETED-BUT-PRINTED-NOT: IR Dump {{.*}}LoopInstSimplifyPass
 
 define void @deleteme() {
 entry:
diff --git a/llvm/test/Other/scc-deleted-printer.ll b/llvm/test/Other/scc-deleted-printer.ll
index e6323f0274a..f7b97c3d30e 100644
--- a/llvm/test/Other/scc-deleted-printer.ll
+++ b/llvm/test/Other/scc-deleted-printer.ll
@@ -9,7 +9,7 @@
 ; INL: IR Dump After {{InlinerPass .*scc: .tester}}
 
 ; INL-MOD: IR Dump Before {{InlinerPass .*scc: .tester, foo}}
-; INL-MOD-NOT: IR Dump After {{InlinerPass}}
+; INL-MOD: IR Dump After {{InlinerPass .*invalidated: .*scc: .tester, foo}}
 ; INL-MOD: IR Dump Before {{InlinerPass .*scc: .tester}}
 ; INL-MOD: IR Dump After {{InlinerPass .*scc: .tester}}