[NewPM][CodeGen] Support printing machine functions from StandardInstrumentation

llvm/lib/Passes/StandardInstrumentations.cpp

 ///
 //===----------------------------------------------------------------------===//
 
 #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/CodeGen/MIRPrinter.h"
+#include "llvm/CodeGen/MachineFunction.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"
 
Optional<std::pair<const Module *, std::string>> unwrapModule(Any IR) {
   if (any_isa<const Function *>(IR)) {
     const Function *F = any_cast<const Function *>(IR);
     if (!llvm::isFunctionInPrintList(F->getName()))
       return None;
     const Module *M = F->getParent();
     return std::make_pair(M, formatv(" (function: {0})", F->getName()).str());
   }
 
+  if (any_isa<const MachineFunction *>(IR)) {
+    const MachineFunction *MF = any_cast<const MachineFunction *>(IR);
+    if (!llvm::isFunctionInPrintList(MF->getName()))
+      return None;
+    const Module *M = MF->getFunction().getParent();
+    return std::make_pair(
+        M, formatv(" (machine function: {0})", MF->getName()).str());
+  }
+
   if (any_isa<const LazyCallGraph::SCC *>(IR)) {
     const LazyCallGraph::SCC *C = any_cast<const LazyCallGraph::SCC *>(IR);
     for (const LazyCallGraph::Node &N : *C) {
       const Function &F = N.getFunction();
       if (!F.isDeclaration() && isFunctionInPrintList(F.getName())) {
         const Module *M = F.getParent();
         return std::make_pair(M, formatv(" (scc: {0})", C->getName()).str());
       }

 }
 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(), std::string(Banner));
 }
 
+void printIR(const MachineFunction *MF, StringRef Banner,
+             StringRef Extra = StringRef()) {
+  if (!llvm::isFunctionInPrintList(MF->getName()))
+    return;
+  dbgs() << Banner << Extra << "\n";
+  printMIR(dbgs(), *MF);
+}
 /// 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;
   }
void unwrapAndPrint(Any IR, StringRef Banner, bool ForceModule = false) {
   }
 
   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 (any_isa<const MachineFunction *>(IR)) {
+    const MachineFunction *MF = any_cast<const MachineFunction *>(IR);
+    assert(MF && "machine function should be valid for printing");
+    printIR(MF, Banner);
+    return;
+  }
+
   llvm_unreachable("Unknown wrapped IR type");
 }
 
 } // namespace
 
 PrintIRInstrumentation::~PrintIRInstrumentation() {
   assert(ModuleDescStack.empty() && "ModuleDescStack is not empty at exit");
 }