[New PM][PassInstrumentation] IR printing support for New Pass Manager

Implementing -print-before-all/-print-after-all/-filter-print-func support
through PassInstrumentation callbacks.

- PrintIR routines implement printing callbacks.

- StandardInstrumentations class provides a central place to manage all
  the "standard" in-tree pass instrumentations. Currently it registers
  PrintIR callbacks.

Reviewers: chandlerc, paquette, philip.pfaffe
Differential Revision: https://reviews.llvm.org/D50923

llvm-svn: 342896

6 files changed, 135 insertions, 12 deletions

diff --git a/llvm/lib/Analysis/CallGraphSCCPass.cpp b/llvm/lib/Analysis/CallGraphSCCPass.cpp
index 3e9226f91de..80d0427529a 100644
--- a/llvm/lib/Analysis/CallGraphSCCPass.cpp
+++ b/llvm/lib/Analysis/CallGraphSCCPass.cpp
@@ -22,6 +22,7 @@
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/IRPrintingPasses.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LegacyPassManagers.h"
diff --git a/llvm/lib/Analysis/LoopInfo.cpp b/llvm/lib/Analysis/LoopInfo.cpp
index 07f72addc81..99ff25a3fd3 100644
--- a/llvm/lib/Analysis/LoopInfo.cpp
+++ b/llvm/lib/Analysis/LoopInfo.cpp
@@ -26,6 +26,7 @@
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/Dominators.h"
+#include "llvm/IR/IRPrintingPasses.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Metadata.h"
diff --git a/llvm/lib/CodeGen/MachineFunctionPrinterPass.cpp b/llvm/lib/CodeGen/MachineFunctionPrinterPass.cpp
index 55d9defced3..9c3288e02db 100644
--- a/llvm/lib/CodeGen/MachineFunctionPrinterPass.cpp
+++ b/llvm/lib/CodeGen/MachineFunctionPrinterPass.cpp
@@ -15,6 +15,7 @@
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/SlotIndexes.h"
+#include "llvm/IR/IRPrintingPasses.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
diff --git a/llvm/lib/IR/LegacyPassManager.cpp b/llvm/lib/IR/LegacyPassManager.cpp
index 3b0c49a4cab..01d14f17bba 100644
--- a/llvm/lib/IR/LegacyPassManager.cpp
+++ b/llvm/lib/IR/LegacyPassManager.cpp
@@ -100,27 +100,31 @@ static cl::list<std::string>
 /// This is a helper to determine whether to print IR before or
 /// after a pass.
 
-static bool ShouldPrintBeforeOrAfterPass(const PassInfo *PI,
+bool llvm::shouldPrintBeforePass() {
+  return PrintBeforeAll || !PrintBefore.empty();
+}
+
+bool llvm::shouldPrintAfterPass() {
+  return PrintAfterAll || !PrintAfter.empty();
+}
+
+static bool ShouldPrintBeforeOrAfterPass(StringRef PassID,
                                          PassOptionList &PassesToPrint) {
   for (auto *PassInf : PassesToPrint) {
     if (PassInf)
-      if (PassInf->getPassArgument() == PI->getPassArgument()) {
+      if (PassInf->getPassArgument() == PassID) {
         return true;
       }
   }
   return false;
 }
 
-/// This is a utility to check whether a pass should have IR dumped
-/// before it.
-static bool ShouldPrintBeforePass(const PassInfo *PI) {
-  return PrintBeforeAll || ShouldPrintBeforeOrAfterPass(PI, PrintBefore);
+bool llvm::shouldPrintBeforePass(StringRef PassID) {
+  return PrintBeforeAll || ShouldPrintBeforeOrAfterPass(PassID, PrintBefore);
 }
 
-/// This is a utility to check whether a pass should have IR dumped
-/// after it.
-static bool ShouldPrintAfterPass(const PassInfo *PI) {
-  return PrintAfterAll || ShouldPrintBeforeOrAfterPass(PI, PrintAfter);
+bool llvm::shouldPrintAfterPass(StringRef PassID) {
+  return PrintAfterAll || ShouldPrintBeforeOrAfterPass(PassID, PrintAfter);
 }
 
 bool llvm::forcePrintModuleIR() { return PrintModuleScope; }
@@ -780,7 +784,7 @@ void PMTopLevelManager::schedulePass(Pass *P) {
     return;
   }
 
-  if (PI && !PI->isAnalysis() && ShouldPrintBeforePass(PI)) {
+  if (PI && !PI->isAnalysis() && shouldPrintBeforePass(PI->getPassArgument())) {
     Pass *PP = P->createPrinterPass(
         dbgs(), ("*** IR Dump Before " + P->getPassName() + " ***").str());
     PP->assignPassManager(activeStack, getTopLevelPassManagerType());
@@ -789,7 +793,7 @@ void PMTopLevelManager::schedulePass(Pass *P) {
   // Add the requested pass to the best available pass manager.
   P->assignPassManager(activeStack, getTopLevelPassManagerType());
 
-  if (PI && !PI->isAnalysis() && ShouldPrintAfterPass(PI)) {
+  if (PI && !PI->isAnalysis() && shouldPrintAfterPass(PI->getPassArgument())) {
     Pass *PP = P->createPrinterPass(
         dbgs(), ("*** IR Dump After " + P->getPassName() + " ***").str());
     PP->assignPassManager(activeStack, getTopLevelPassManagerType());
diff --git a/llvm/lib/Passes/CMakeLists.txt b/llvm/lib/Passes/CMakeLists.txt
index 468fd070a91..d1b873c3fd7 100644
--- a/llvm/lib/Passes/CMakeLists.txt
+++ b/llvm/lib/Passes/CMakeLists.txt
@@ -5,6 +5,7 @@ endif()
 add_llvm_library(LLVMPasses
   PassBuilder.cpp
   PassPlugin.cpp
+  StandardInstrumentations.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/Passes
diff --git a/llvm/lib/Passes/StandardInstrumentations.cpp b/llvm/lib/Passes/StandardInstrumentations.cpp
new file mode 100644
index 00000000000..6abd39c61ac
--- /dev/null
+++ b/llvm/lib/Passes/StandardInstrumentations.cpp
@@ -0,0 +1,115 @@
+//===- 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/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
+//===----------------------------------------------------------------------===//
+
+/// 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) {
+  if (any_isa<const CallGraphSCC *>(IR) ||
+      any_isa<const LazyCallGraph::SCC *>(IR))
+    return;
+
+  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)) {
+    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 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());
+    }
+  }
+  if (M) {
+    dbgs() << Banner << Extra;
+    M->print(dbgs(), nullptr, false);
+  } else {
+    llvm_unreachable("Unknown wrapped IR type");
+  }
+}
+
+bool printBeforePass(StringRef PassID, Any IR) {
+  if (!llvm::shouldPrintBeforePass(PassID))
+    return true;
+
+  if (PassID.startswith("PassManager<") || PassID.contains("PassAdaptor<"))
+    return true;
+
+  SmallString<20> Banner = formatv("*** IR Dump Before {0} ***", PassID);
+  unwrapAndPrint(Banner, IR);
+  return true;
+}
+
+void printAfterPass(StringRef PassID, Any IR) {
+  if (!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;
+}
+} // namespace PrintIR
+} // namespace
+
+void StandardInstrumentations::registerCallbacks(
+    PassInstrumentationCallbacks &PIC) {
+  if (llvm::shouldPrintBeforePass())
+    PIC.registerBeforePassCallback(PrintIR::printBeforePass);
+  if (llvm::shouldPrintAfterPass())
+    PIC.registerAfterPassCallback(PrintIR::printAfterPass);
+}