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

[fedor.sergeev]

Summary

-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.

[fedor.sergeev]

oops, this needs a small update

[fedor.sergeev]

stack handling fixed (AfterPass now pops stack as well)

[fedor.sergeev]

minor cleanup

[philip.pfaffe]

What exactly does this change solve?

[fedor.sergeev]

In D55278#1321579, @philip.pfaffe wrote:

What exactly does this change solve?

updated the summary.

[fedor.sergeev]

ping?

[philip.pfaffe]

This generally looks fine, some code comments inline.

lib/Passes/StandardInstrumentations.cpp

34

This should be an Optional.

35

Why this intermiate? Why not just return?

169

Avoid AAA

[philip.pfaffe]

lib/Passes/StandardInstrumentations.cpp

35

*intermediate

[fedor.sergeev]

addressing comments

[philip.pfaffe]

lib/Passes/StandardInstrumentations.cpp

88

Is this truly more readable than UnwrappedModule->first, UnwrappedModule->second?

[fedor.sergeev]

fixing Optional use

[fedor.sergeev]

lib/Passes/StandardInstrumentations.cpp

88

Definitely not, thanks for suggestion.

[philip.pfaffe]

Thanks, LGTM!

include/llvm/Passes/StandardInstrumentations.h

 /// This header defines a class that provides bookkeeping for all standard
 /// (i.e in-tree) pass instrumentations.
 ///
 //===----------------------------------------------------------------------===//
 
 #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:
   StandardInstrumentations() = default;
 
   void registerCallbacks(PassInstrumentationCallbacks &PIC);
 };
 } // namespace llvm
 
 #endif

lib/Passes/StandardInstrumentations.cpp

 #include "llvm/IR/PassInstrumentation.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
 
 namespace {
-namespace PrintIR {
 
-//===----------------------------------------------------------------------===//
-// IR-printing instrumentation
-//===----------------------------------------------------------------------===//
-
-/// 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"};
+/// Extracting Module out of \p IR unit. Also fills a textual description
+/// of \p IR for use in header when printing.
+std::pair<const Module *, std::string> unwrapModule(Any IR) {
+  std::string Extra;
   const Module *M = nullptr;
   if (any_isa<const Module *>(IR)) {
     M = any_cast<const Module *>(IR);
     assert(M && "module should be valid for printing");
   } else if (any_isa<const Function *>(IR)) {
     const Function *F = any_cast<const Function *>(IR);
     assert(F && "function should be valid for printing");
     if (!llvm::isFunctionInPrintList(F->getName()))
-      return;
-    if (!llvm::forcePrintModuleIR()) {
-      dbgs() << Banner << Extra << static_cast<const Value &>(*F);
-      return;
-    }
+      return {nullptr, std::string()};
     M = F->getParent();
-    Extra = formatv(" (function: {0})\n", F->getName());
+    Extra = formatv(" (function: {0})", F->getName());
   } else if (any_isa<const LazyCallGraph::SCC *>(IR)) {
     const LazyCallGraph::SCC *C = any_cast<const LazyCallGraph::SCC *>(IR);
     assert(C);
-    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;
       }
     }
     if (!M)
-      return;
-    Extra = formatv(" (for scc: {0})\n", C->getName());
+      return {nullptr, std::string()};
+    Extra = formatv(" (scc: {0})", C->getName());
   } else if (any_isa<const Loop *>(IR)) {
     const Loop *L = any_cast<const Loop *>(IR);
     assert(L && "Loop should be valid for printing");
     const Function *F = L->getHeader()->getParent();
     if (!isFunctionInPrintList(F->getName()))
-      return;
-    if (!llvm::forcePrintModuleIR()) {
-      llvm::printLoop(const_cast<Loop &>(*L), dbgs(), Banner);
-      return;
-    }
+      return {nullptr, std::string()};
     M = F->getParent();
     {
       std::string LoopName;
       raw_string_ostream ss(LoopName);
       L->getHeader()->printAsOperand(ss, false);
-      Extra = formatv(" (loop: {0})\n", ss.str());
+      Extra = formatv(" (loop: {0})", ss.str());
     }
   }
-  if (M) {
-    dbgs() << Banner << Extra;
-    M->print(dbgs(), nullptr, false);
-  } else {
-    llvm_unreachable("Unknown wrapped IR type");
-  }
+  assert(M && "module should be valid for printing");
+  return {M, std::move(Extra)};
+}
+
+void printIR(const Module *M, StringRef Banner, StringRef Extra = StringRef()) {
+  dbgs() << Banner << Extra << "\n";
+}
+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());
+    }
+  }
+}
+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) {
+    const Module *M;
+    std::string Extra;
+    std::tie(M, Extra) = unwrapModule(IR);
+    if (M)
+      printIR(M, Banner, Extra);
+    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;
+  }
+
+} // namespace
+
+PrintIRInstrumentation::~PrintIRInstrumentation() {
+  assert(ModuleDescStack.empty() && "ModuleDescStack is not empty at exit");
 }
 
-bool printBeforePass(StringRef PassID, Any IR) {
-  if (!llvm::shouldPrintBeforePass(PassID))
-    return true;
+void PrintIRInstrumentation::pushModuleDesc(StringRef PassID, Any IR) {
+  assert(StoreModuleDesc);
+  const Module *M;
+  std::string Extra;
+  std::tie(M, Extra) = unwrapModule(IR);
+  ModuleDescStack.emplace_back(M, Extra, PassID);
+}
+
+PrintIRInstrumentation::PrintModuleDesc
+PrintIRInstrumentation::popModuleDesc(StringRef PassID) {
+  assert(!ModuleDescStack.empty() && "empty ModuleDescStack");
+  auto 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);
+  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);
 }

test/Other/loop-deletion-printer.ll

 ; Make sure that Loop which was invalidated by loop-deletion
 ; does not lead to problems for -print-after-all and is just skipped.
 ;
 ; RUN: opt < %s -disable-output \
 ; 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:
   br label %loop
 loop:
   %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop ]
   %iv.next = add i32 %iv, 1
   %check = icmp ult i32 %iv.next, 3
   br i1 %check, label %loop, label %exit
 exit:
   ret void
 }
 

test/Other/scc-deleted-printer.ll

 ; RUN: opt < %s 2>&1 -disable-output \
 ; RUN:     -passes=inline -print-before-all -print-after-all | FileCheck %s -check-prefix=INL
 ; RUN: opt < %s 2>&1 -disable-output \
 ; RUN:     -passes=inline -print-before-all -print-after-all -print-module-scope | FileCheck %s -check-prefix=INL-MOD
 
 ; INL: IR Dump Before {{InlinerPass .*scc: .tester, foo}}
 ; INL-NOT: IR Dump After {{InlinerPass}}
 ; INL: IR Dump Before {{InlinerPass .*scc: .tester}}
 ; 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}}
 
 
 define void @tester() noinline {
   call void @foo()
   ret void
 }
 
 define internal void @foo() alwaysinline {
   call void @tester()
   ret void
 }