[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

40

This should be an Optional.

41

Why this intermiate? Why not just return?

164

Avoid AAA

[philip.pfaffe]

lib/Passes/StandardInstrumentations.cpp

41

*intermediate

[fedor.sergeev]

addressing comments

[philip.pfaffe]

lib/Passes/StandardInstrumentations.cpp

80

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

[fedor.sergeev]

fixing Optional use

[fedor.sergeev]

lib/Passes/StandardInstrumentations.cpp

80

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

 //===- Standard pass instrumentations handling ----------------*- C++ -*--===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 /// \file
 ///
 /// This file defines IR-printing pass instrumentation callbacks as well as
 /// StandardInstrumentations class that manages standard pass instrumentations.
 ///
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRPrintingPasses.h"
 #include "llvm/IR/Module.h"
 #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
-//===----------------------------------------------------------------------===//
+/// 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());
+    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 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()))
+  dbgs() << Banner << Extra << "\n" << static_cast<const Value &>(*F);
+void printIR(const LazyCallGraph::SCC *C, StringRef Banner,
+             StringRef Extra = StringRef()) {
+  bool BannerPrinted = false;
+    if (!F.isDeclaration() && llvm::isFunctionInPrintList(F.getName())) {
+      if (!BannerPrinted) {
+        dbgs() << Banner << Extra << "\n";
+        BannerPrinted = true;
       }
+      F.print(dbgs());
     }
-    if (!M)
-      return;
-    Extra = formatv(" (for scc: {0})\n", C->getName());
-  } else 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());
-    }
+  }
+}
+void printIR(const Loop *L, StringRef Banner) {
+  if (!llvm::isFunctionInPrintList(F->getName()))
+}
+
+/// 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);
+
+  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;
   }
-  if (M) {
-    dbgs() << Banner << Extra;
-    M->print(dbgs(), nullptr, false);
-  } else {
-    llvm_unreachable("Unknown wrapped IR type");
-  }
+
+} // 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 = 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);
+  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
 }