[PassManager] Add pretty stack frames (#96078)

In NewPM pass managers, add a "pretty stack frame" that tells you which
pass crashed while running which function.

For example `opt -O3 -passes-ep-peephole=trigger-crash-function test.ll`
will print something like this:

```
Stack dump:
0.	Program arguments: build/bin/opt -S -O3 -passes-ep-peephole=trigger-crash-function test.ll
1.	Running pass "function<eager-inv>(mem2reg,instcombine<max-iterations=1;no-use-loop-info;no-verify-fixpoint>,trigger-crash-function,simplifycfg<bonus-inst-threshold=1;no-forward-switch-cond;switch-range-to-icmp;no-switch-to-lookup;keep-loops;no-hoist-common-insts;no-sink-common-insts;speculate-blocks;simplify-cond-branch>)" on module "test.ll"
2.	Running pass "trigger-crash-function" on function "fshl_concat_i8_i8"
```

While the crashing pass is usually evident from the stack trace, this
also shows which function triggered the crash, as well as the pipeline
string for the pass (including options).

Similar functionality existed in the LegacyPM.

8 files changed, 95 insertions, 11 deletions

diff --git a/llvm/include/llvm/IR/PassInstrumentation.h b/llvm/include/llvm/IR/PassInstrumentation.h
index f2eb8a9..9fcc2d5 100644
--- a/llvm/include/llvm/IR/PassInstrumentation.h
+++ b/llvm/include/llvm/IR/PassInstrumentation.h
@@ -332,6 +332,13 @@ public:
     if (Callbacks)
       Callbacks->BeforeNonSkippedPassCallbacks.pop_back();
   }
+
+  /// Get the pass name for a given pass class name.
+  StringRef getPassNameForClassName(StringRef ClassName) const {
+    if (Callbacks)
+      return Callbacks->getPassNameForClassName(ClassName);
+    return {};
+  }
 };
 
 bool isSpecialPass(StringRef PassID, const std::vector<StringRef> &Specials);
diff --git a/llvm/include/llvm/IR/PassManager.h b/llvm/include/llvm/IR/PassManager.h
index 23a05c0..226eba7 100644
--- a/llvm/include/llvm/IR/PassManager.h
+++ b/llvm/include/llvm/IR/PassManager.h
@@ -224,11 +224,21 @@ protected:
   std::vector<std::unique_ptr<PassConceptT>> Passes;
 };
 
+template <typename IRUnitT>
+void printIRUnitNameForStackTrace(raw_ostream &OS, const IRUnitT &IR);
+
+template <>
+void printIRUnitNameForStackTrace<Module>(raw_ostream &OS, const Module &IR);
+
 extern template class PassManager<Module>;
 
 /// Convenience typedef for a pass manager over modules.
 using ModulePassManager = PassManager<Module>;
 
+template <>
+void printIRUnitNameForStackTrace<Function>(raw_ostream &OS,
+                                            const Function &IR);
+
 extern template class PassManager<Function>;
 
 /// Convenience typedef for a pass manager over functions.
diff --git a/llvm/include/llvm/IR/PassManagerImpl.h b/llvm/include/llvm/IR/PassManagerImpl.h
index e8a9576..91f5774 100644
--- a/llvm/include/llvm/IR/PassManagerImpl.h
+++ b/llvm/include/llvm/IR/PassManagerImpl.h
@@ -15,9 +15,10 @@
 #ifndef LLVM_IR_PASSMANAGERIMPL_H
 #define LLVM_IR_PASSMANAGERIMPL_H
 
-#include "llvm/Support/CommandLine.h"
 #include "llvm/IR/PassInstrumentation.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/PrettyStackTrace.h"
 
 extern llvm::cl::opt<bool> UseNewDbgInfoFormat;
 
@@ -26,6 +27,28 @@ namespace llvm {
 template <typename IRUnitT, typename AnalysisManagerT, typename... ExtraArgTs>
 PreservedAnalyses PassManager<IRUnitT, AnalysisManagerT, ExtraArgTs...>::run(
     IRUnitT &IR, AnalysisManagerT &AM, ExtraArgTs... ExtraArgs) {
+  class StackTraceEntry : public PrettyStackTraceEntry {
+    const PassInstrumentation &PI;
+    PassConceptT &Pass;
+    IRUnitT &IR;
+
+  public:
+    explicit StackTraceEntry(const PassInstrumentation &PI, PassConceptT &Pass,
+                             IRUnitT &IR)
+        : PI(PI), Pass(Pass), IR(IR) {}
+
+    void print(raw_ostream &OS) const override {
+      OS << "Running pass \"";
+      Pass.printPipeline(OS, [this](StringRef ClassName) {
+        auto PassName = PI.getPassNameForClassName(ClassName);
+        return PassName.empty() ? ClassName : PassName;
+      });
+      OS << "\" on ";
+      printIRUnitNameForStackTrace(OS, IR);
+      OS << "\n";
+    }
+  };
+
   PreservedAnalyses PA = PreservedAnalyses::all();
 
   // Request PassInstrumentation from analysis manager, will use it to run
@@ -47,6 +70,7 @@ PreservedAnalyses PassManager<IRUnitT, AnalysisManagerT, ExtraArgTs...>::run(
     if (!PI.runBeforePass<IRUnitT>(*Pass, IR))
       continue;
 
+    StackTraceEntry Entry(PI, *Pass, IR);
     PreservedAnalyses PassPA = Pass->run(IR, AM, ExtraArgs...);
 
     // Update the analysis manager as each pass runs and potentially
diff --git a/llvm/lib/IR/PassManager.cpp b/llvm/lib/IR/PassManager.cpp
index cbddf3d..64a6eff 100644
--- a/llvm/lib/IR/PassManager.cpp
+++ b/llvm/lib/IR/PassManager.cpp
@@ -144,6 +144,18 @@ PreservedAnalyses ModuleToFunctionPassAdaptor::run(Module &M,
   return PA;
 }
 
+template <>
+void llvm::printIRUnitNameForStackTrace<Module>(raw_ostream &OS,
+                                                const Module &IR) {
+  OS << "module \"" << IR.getName() << "\"";
+}
+
+template <>
+void llvm::printIRUnitNameForStackTrace<Function>(raw_ostream &OS,
+                                                  const Function &IR) {
+  OS << "function \"" << IR.getName() << "\"";
+}
+
 AnalysisSetKey CFGAnalyses::SetKey;
 
 AnalysisSetKey PreservedAnalyses::AllAnalysesKey;
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index a87d0b5..8d2a64d 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -326,15 +326,25 @@ AnalysisKey NoOpLoopAnalysis::Key;
 
 namespace {
 
-// A pass for testing -print-on-crash.
+// Passes for testing crashes.
 // DO NOT USE THIS EXCEPT FOR TESTING!
-class TriggerCrashPass : public PassInfoMixin<TriggerCrashPass> {
+class TriggerCrashModulePass : public PassInfoMixin<TriggerCrashModulePass> {
 public:
   PreservedAnalyses run(Module &, ModuleAnalysisManager &) {
     abort();
     return PreservedAnalyses::all();
   }
-  static StringRef name() { return "TriggerCrashPass"; }
+  static StringRef name() { return "TriggerCrashModulePass"; }
+};
+
+class TriggerCrashFunctionPass
+    : public PassInfoMixin<TriggerCrashFunctionPass> {
+public:
+  PreservedAnalyses run(Function &, FunctionAnalysisManager &) {
+    abort();
+    return PreservedAnalyses::all();
+  }
+  static StringRef name() { return "TriggerCrashFunctionPass"; }
 };
 
 // A pass for testing message reporting of -verify-each failures.
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 60c5177..6c1b04e 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -137,7 +137,7 @@ MODULE_PASS("strip-debug-declare", StripDebugDeclarePass())
 MODULE_PASS("strip-nondebug", StripNonDebugSymbolsPass())
 MODULE_PASS("strip-nonlinetable-debuginfo", StripNonLineTableDebugInfoPass())
 MODULE_PASS("synthetic-counts-propagation", SyntheticCountsPropagation())
-MODULE_PASS("trigger-crash", TriggerCrashPass())
+MODULE_PASS("trigger-crash-module", TriggerCrashModulePass())
 MODULE_PASS("trigger-verifier-error", TriggerVerifierErrorPass())
 MODULE_PASS("tsan-module", ModuleThreadSanitizerPass())
 MODULE_PASS("verify", VerifierPass())
@@ -459,6 +459,7 @@ FUNCTION_PASS("structurizecfg", StructurizeCFGPass())
 FUNCTION_PASS("tailcallelim", TailCallElimPass())
 FUNCTION_PASS("tlshoist", TLSVariableHoistPass())
 FUNCTION_PASS("transform-warning", WarnMissedTransformationsPass())
+FUNCTION_PASS("trigger-crash-function", TriggerCrashFunctionPass())
 FUNCTION_PASS("trigger-verifier-error", TriggerVerifierErrorPass())
 FUNCTION_PASS("tsan", ThreadSanitizerPass())
 FUNCTION_PASS("typepromotion", TypePromotionPass(TM))
diff --git a/llvm/test/Other/crash-stack-trace.ll b/llvm/test/Other/crash-stack-trace.ll
new file mode 100644
index 0000000..29e43fe
--- /dev/null
+++ b/llvm/test/Other/crash-stack-trace.ll
@@ -0,0 +1,20 @@
+; REQUIRES: asserts
+
+; RUN: not --crash opt -passes=trigger-crash-module %s -disable-output 2>&1 | \
+; RUN: FileCheck %s --check-prefix=CHECK-MODULE
+
+; CHECK-MODULE:      Stack dump:
+; CHECK-MODULE-NEXT: 0. Program arguments:
+; CHECK-MODULE-NEXT: 1. Running pass "trigger-crash-module" on module "{{.*}}crash-stack-trace.ll"
+
+; RUN: not --crash opt -passes='sroa,trigger-crash-function' %s -disable-output 2>&1 | \
+; RUN: FileCheck %s --check-prefix=CHECK-FUNCTION
+
+; CHECK-FUNCTION:      Stack dump:
+; CHECK-FUNCTION-NEXT: 0. Program arguments:
+; CHECK-FUNCTION-NEXT: 1. Running pass "function(sroa<modify-cfg>,trigger-crash-function)" on module "{{.*}}crash-stack-trace.ll"
+; CHECK-FUNCTION-NEXT: 2. Running pass "trigger-crash-function" on function "foo"
+
+define void @foo() {
+  ret void
+}
diff --git a/llvm/test/Other/print-on-crash.ll b/llvm/test/Other/print-on-crash.ll
index f9bf8f0..565da0c 100644
--- a/llvm/test/Other/print-on-crash.ll
+++ b/llvm/test/Other/print-on-crash.ll
@@ -1,9 +1,9 @@
 ; A test that the hidden option -print-on-crash properly sets a signal handler
-; which gets called when a pass crashes.  The trigger-crash pass asserts.
+; which gets called when a pass crashes.  The trigger-crash-module pass asserts.
 
-; RUN: not --crash opt -print-on-crash -passes=trigger-crash < %s 2>&1 | FileCheck %s --check-prefix=CHECK_SIMPLE
+; RUN: not --crash opt -print-on-crash -passes=trigger-crash-module < %s 2>&1 | FileCheck %s --check-prefix=CHECK_SIMPLE
 
-; RUN: not --crash opt -print-on-crash-path=%t -passes=trigger-crash < %s
+; RUN: not --crash opt -print-on-crash-path=%t -passes=trigger-crash-module < %s
 ; RUN: FileCheck %s --check-prefix=CHECK_SIMPLE --input-file=%t
 
 ; A test that the signal handler set by the  hidden option -print-on-crash
@@ -11,11 +11,11 @@
 
 ; RUN: opt -disable-output -print-on-crash -passes="default<O2>" < %s 2>&1 | FileCheck %s --check-prefix=CHECK_NO_CRASH --allow-empty
 
-; RUN: not --crash opt -print-on-crash -print-module-scope -passes=trigger-crash < %s 2>&1 | FileCheck %s --check-prefix=CHECK_MODULE
+; RUN: not --crash opt -print-on-crash -print-module-scope -passes=trigger-crash-module < %s 2>&1 | FileCheck %s --check-prefix=CHECK_MODULE
 
-; RUN: not --crash opt -print-on-crash -print-module-scope -passes=trigger-crash -filter-passes=trigger-crash < %s 2>&1 | FileCheck %s --check-prefix=CHECK_MODULE
+; RUN: not --crash opt -print-on-crash -print-module-scope -passes=trigger-crash-module -filter-passes=trigger-crash-module < %s 2>&1 | FileCheck %s --check-prefix=CHECK_MODULE
 
-; RUN: not --crash opt -print-on-crash -print-module-scope -passes=trigger-crash -filter-passes=blah < %s 2>&1 | FileCheck %s --check-prefix=CHECK_FILTERED
+; RUN: not --crash opt -print-on-crash -print-module-scope -passes=trigger-crash-module -filter-passes=blah < %s 2>&1 | FileCheck %s --check-prefix=CHECK_FILTERED
 
 ; CHECK_SIMPLE: *** Dump of IR Before Last Pass {{.*}} Started ***
 ; CHECK_SIMPLE: @main