[StandardInstrumentations] Support -print-after-pass-number option

llvm/include/llvm/Passes/StandardInstrumentations.h

@@ -64,8 +64,11 @@ class PrintIRInstrumentation {
 
   bool shouldPrintBeforePass(StringRef PassID);
   bool shouldPrintAfterPass(StringRef PassID);
+  bool shouldPrintBeforeCurrentPassNumber();
+  bool shouldPrintAfterCurrentPassNumber();
   bool shouldPrintPassNumbers();
-  bool shouldPrintBeforePassNumber();
+  bool shouldPrintBeforeSomePassNumber();
+  bool shouldPrintAfterSomePassNumber();
 
   void pushPassRunDescriptor(StringRef PassID, Any IR,
                              std::string &DumpIRFilename);

llvm/lib/Passes/StandardInstrumentations.cpp

@@ -123,6 +123,11 @@ static cl::opt<unsigned> PrintBeforePassNumber(
     cl::desc("Print IR before the pass with this number as "
              "reported by print-pass-numbers"));
 
+static cl::opt<unsigned>
+    PrintAfterPassNumber("print-after-pass-number", cl::init(0), cl::Hidden,
+                         cl::desc("Print IR after the pass with this number as "
+                                  "reported by print-pass-numbers"));
+
 static cl::opt<std::string> IRDumpDirectory(
     "ir-dump-directory",
     cl::desc("If specified, IR printed using the "
@@ -799,7 +804,9 @@ void PrintIRInstrumentation::printBeforePass(StringRef PassID, Any IR) {
 
   std::string DumpIRFilename;
   if (!IRDumpDirectory.empty() &&
-      (shouldPrintBeforePass(PassID) || shouldPrintAfterPass(PassID)))
+      (shouldPrintBeforePass(PassID) || shouldPrintAfterPass(PassID) ||
+       shouldPrintBeforeCurrentPassNumber() ||
+       shouldPrintAfterCurrentPassNumber()))
     DumpIRFilename = fetchDumpFilename(PassID, IR);
 
   // Saving Module for AfterPassInvalidated operations.
@@ -818,12 +825,15 @@ void PrintIRInstrumentation::printBeforePass(StringRef PassID, Any IR) {
     dbgs() << " Running pass " << CurrentPassNumber << " " << PassID
            << " on " << getIRName(IR) << "\n";
 
-  if (!shouldPrintBeforePass(PassID))
+  if (shouldPrintAfterCurrentPassNumber())
+    pushPassRunDescriptor(PassID, IR, DumpIRFilename);
+
+  if (!shouldPrintBeforePass(PassID) && !shouldPrintBeforeCurrentPassNumber())
     return;
 
   auto WriteIRToStream = [&](raw_ostream &Stream) {
     Stream << "; *** IR Dump Before ";
-    if (shouldPrintBeforePassNumber())
+    if (shouldPrintBeforeSomePassNumber())
       Stream << CurrentPassNumber << "-";
     Stream << PassID << " on " << getIRName(IR) << " ***\n";
     unwrapAndPrint(Stream, IR);
@@ -843,18 +853,21 @@ void PrintIRInstrumentation::printAfterPass(StringRef PassID, Any IR) {
   if (isIgnored(PassID))
     return;
 
-  if (!shouldPrintAfterPass(PassID))
+  if (!shouldPrintAfterPass(PassID) && !shouldPrintAfterCurrentPassNumber())
     return;
 
   auto [M, DumpIRFilename, IRName, StoredPassID] = popPassRunDescriptor(PassID);
   assert(StoredPassID == PassID && "mismatched PassID");
 
-  if (!shouldPrintIR(IR) || !shouldPrintAfterPass(PassID))
+  if (!shouldPrintIR(IR) ||
+      (!shouldPrintAfterPass(PassID) && !shouldPrintAfterCurrentPassNumber()))
     return;
 
   auto WriteIRToStream = [&](raw_ostream &Stream, const StringRef IRName) {
-    Stream << "; *** IR Dump " << StringRef(formatv("After {0}", PassID))
-           << " on " << IRName << " ***\n";
+    Stream << "; *** IR Dump After ";
+    if (shouldPrintAfterSomePassNumber())
+      Stream << CurrentPassNumber << "-";
+    Stream << StringRef(formatv("{0}", PassID)) << " on " << IRName << " ***\n";
     unwrapAndPrint(Stream, IR);
   };
 
@@ -876,14 +889,15 @@ void PrintIRInstrumentation::printAfterPassInvalidated(StringRef PassID) {
   if (isIgnored(PassID))
     return;
 
-  if (!shouldPrintAfterPass(PassID))
+  if (!shouldPrintAfterPass(PassID) && !shouldPrintAfterCurrentPassNumber())
     return;
 
   auto [M, DumpIRFilename, IRName, StoredPassID] = popPassRunDescriptor(PassID);
   assert(StoredPassID == PassID && "mismatched PassID");
   // Additional filtering (e.g. -filter-print-func) can lead to module
   // printing being skipped.
-  if (!M || !shouldPrintAfterPass(PassID))
+  if (!M ||
+      (!shouldPrintAfterPass(PassID) && !shouldPrintAfterCurrentPassNumber()))
     return;
 
   auto WriteIRToStream = [&](raw_ostream &Stream, const Module *M,
@@ -913,10 +927,6 @@ bool PrintIRInstrumentation::shouldPrintBeforePass(StringRef PassID) {
   if (shouldPrintBeforeAll())
     return true;
 
-  if (shouldPrintBeforePassNumber() &&
-      CurrentPassNumber == PrintBeforePassNumber)
-    return true;
-
   StringRef PassName = PIC->getPassNameForClassName(PassID);
   return is_contained(printBeforePasses(), PassName);
 }
@@ -929,26 +939,42 @@ bool PrintIRInstrumentation::shouldPrintAfterPass(StringRef PassID) {
   return is_contained(printAfterPasses(), PassName);
 }
 
+bool PrintIRInstrumentation::shouldPrintBeforeCurrentPassNumber() {
+  return shouldPrintBeforeSomePassNumber() &&
+         (CurrentPassNumber == PrintBeforePassNumber);
+}
+
+bool PrintIRInstrumentation::shouldPrintAfterCurrentPassNumber() {
+  return shouldPrintAfterSomePassNumber() &&
+         (CurrentPassNumber == PrintAfterPassNumber);
+}
+
 bool PrintIRInstrumentation::shouldPrintPassNumbers() {
   return PrintPassNumbers;
 }
 
-bool PrintIRInstrumentation::shouldPrintBeforePassNumber() {
+bool PrintIRInstrumentation::shouldPrintBeforeSomePassNumber() {
   return PrintBeforePassNumber > 0;
 }
 
+bool PrintIRInstrumentation::shouldPrintAfterSomePassNumber() {
+  return PrintAfterPassNumber > 0;
+}
+
 void PrintIRInstrumentation::registerCallbacks(
     PassInstrumentationCallbacks &PIC) {
   this->PIC = &PIC;
 
   // BeforePass callback is not just for printing, it also saves a Module
-  // for later use in AfterPassInvalidated.
-  if (shouldPrintPassNumbers() || shouldPrintBeforePassNumber() ||
-      shouldPrintBeforeSomePass() || shouldPrintAfterSomePass())
+  // for later use in AfterPassInvalidated and keeps tracks of the
+  // CurrentPassNumber.
+  if (shouldPrintPassNumbers() || shouldPrintBeforeSomePassNumber() ||
+      shouldPrintAfterSomePassNumber() || shouldPrintBeforeSomePass() ||
+      shouldPrintAfterSomePass())
     PIC.registerBeforeNonSkippedPassCallback(
         [this](StringRef P, Any IR) { this->printBeforePass(P, IR); });
 
-  if (shouldPrintAfterSomePass()) {
+  if (shouldPrintAfterSomePass() || shouldPrintAfterSomePassNumber()) {
     PIC.registerAfterPassCallback(
         [this](StringRef P, Any IR, const PreservedAnalyses &) {
           this->printAfterPass(P, IR);

llvm/test/Other/print-at-pass-number.ll

@@ -1,9 +1,13 @@
 ; RUN: opt -passes="loop(indvars,loop-deletion,loop-unroll-full)" -print-pass-numbers -S -o /dev/null %s 2>&1 | FileCheck %s --check-prefix=NUMBER
+; RUN: opt -passes="loop(indvars,loop-deletion,loop-unroll-full)" -print-pass-numbers -filter-print-funcs=baz -S -o /dev/null %s 2>&1 | FileCheck %s --check-prefix=NUMBER-FILTERED
 ; RUN: opt -passes="loop(indvars,loop-deletion,loop-unroll-full)" -print-module-scope -print-before-pass-number=3 -S -o /dev/null %s 2>&1 | FileCheck %s --check-prefix=BEFORE
+; RUN: opt -passes="loop(indvars,loop-deletion,loop-unroll-full)" -print-module-scope -print-after-pass-number=2 -S -o /dev/null %s 2>&1 | FileCheck %s --check-prefix=AFTER
 
 define i32 @bar(i32 %arg) {
 ; BEFORE: *** IR Dump Before 3-IndVarSimplifyPass on bb1 ***
 ; BEFORE: define i32 @bar(i32 %arg) {
+; AFTER:  *** IR Dump After 2-LCSSAPass on bar ***
+; AFTER:  define i32 @bar(i32 %arg) {
 
 bb:
   br label %bb1
@@ -20,8 +24,18 @@ bb4:                                              ; preds = %bb1
   ret i32 %add3
 }
 
+define i32 @baz(i32 %arg) {
+  ret i32 0;
+}
+
 ; NUMBER:  Running pass 1 LoopSimplifyPass on bar
 ; NUMBER-NEXT: Running pass 2 LCSSAPass on bar
 ; NUMBER-NEXT: Running pass 3 IndVarSimplifyPass on bb1
 ; NUMBER-NEXT: Running pass 4 LoopDeletionPass on bb1
+; NUMBER-NEXT: Running pass 5 LoopSimplifyPass on baz
+; NUMBER-NEXT: Running pass 6 LCSSAPass on baz
 ; NUMBER-NOT: Running pass
+
+; NUMBER-FILTERED: Running pass 1 LoopSimplifyPass on baz
+; NUMBER-FILTERED-NEXT: Running pass 2 LCSSAPass on baz
+