[llvm] [llvm] Optimize print-on-crash serialization (PR #205682)

[Miguel A. Arroyo via llvm-commits]

https://github.com/mayanez updated https://github.com/llvm/llvm-project/pull/205682

>From f7134dfda4c89348464c47ce058700f3d987dba6 Mon Sep 17 00:00:00 2001
From: Miguel Arroyo <Miguel.Arroyo at rockstargames.com>
Date: Wed, 24 Jun 2026 14:20:08 -0700
Subject: [PATCH] [llvm] Optimize print-on-crash serialization

---
 .../llvm/Passes/StandardInstrumentations.h    |  8 +-
 llvm/lib/Passes/StandardInstrumentations.cpp  | 88 +++++++++++++++++--
 llvm/test/Other/print-on-crash.ll             |  4 +
 3 files changed, 93 insertions(+), 7 deletions(-)

diff --git a/llvm/include/llvm/Passes/StandardInstrumentations.h b/llvm/include/llvm/Passes/StandardInstrumentations.h
index 4ee5ab2554868..f107e757f96b4 100644
--- a/llvm/include/llvm/Passes/StandardInstrumentations.h
+++ b/llvm/include/llvm/Passes/StandardInstrumentations.h
@@ -579,13 +579,17 @@ class LLVM_ABI DotCfgChangeReporter : public ChangeReporter<IRDataT<DCData>> {
 class PrintCrashIRInstrumentation {
 public:
   PrintCrashIRInstrumentation()
-      : SavedIR("*** Dump of IR Before Last Pass Unknown ***") {}
+      : SavedString("*** Dump of IR Before Last Pass Unknown ***") {}
   LLVM_ABI ~PrintCrashIRInstrumentation();
   LLVM_ABI void registerCallbacks(PassInstrumentationCallbacks &PIC);
   LLVM_ABI void reportCrashIR();
+  LLVM_ABI void printToStream(raw_ostream &OS);
 
 protected:
-  std::string SavedIR;
+  std::string SavedString;
+
+  std::unique_ptr<Module> SavedModule;
+  std::vector<llvm::Function *> SavedFunctions;
 
 private:
   // The crash reporter that will report on a crash.
diff --git a/llvm/lib/Passes/StandardInstrumentations.cpp b/llvm/lib/Passes/StandardInstrumentations.cpp
index 075cd56402ca9..4ce78bbd64bfe 100644
--- a/llvm/lib/Passes/StandardInstrumentations.cpp
+++ b/llvm/lib/Passes/StandardInstrumentations.cpp
@@ -42,6 +42,7 @@
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Support/xxhash.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 #include <utility>
 #include <vector>
 
@@ -331,6 +332,51 @@ void unwrapAndPrint(raw_ostream &OS, Any IR) {
   llvm_unreachable("Unknown wrapped IR type");
 }
 
+std::optional<std::pair<std::unique_ptr<Module>, std::vector<Function *>>>
+unwrapAndSaveClone(Any IR) {
+  std::vector<Function *> ValuesToSave;
+
+  if (!shouldPrintIR(IR))
+    return std::nullopt;
+
+  auto *OrigM = unwrapModule(IR);
+  auto M = CloneModule(*OrigM);
+
+  if (forcePrintModuleIR()) {
+    return {{std::move(M), ValuesToSave}};
+  }
+
+  if (const auto *_ = unwrapIR<Module>(IR)) {
+    return {{std::move(M), ValuesToSave}};
+  }
+
+  if (const auto *F = unwrapIR<Function>(IR)) {
+    auto *SaveF = M->getFunction(F->getName());
+    ValuesToSave.push_back(SaveF);
+    return {{std::move(M), ValuesToSave}};
+  }
+
+  if (auto *C = unwrapIR<LazyCallGraph::SCC>(IR)) {
+    for (LazyCallGraph::Node &N : *C) {
+      Function &F = N.getFunction();
+      if (!F.isDeclaration() && isFunctionInPrintList(F.getName())) {
+        ValuesToSave.push_back(&F);
+      }
+    }
+    return {{std::move(M), ValuesToSave}};
+  }
+
+  if (const auto *_ = unwrapIR<Loop>(IR)) {
+    return std::nullopt;
+  }
+
+  if (const auto *_ = unwrapIR<MachineFunction>(IR)) {
+    return std::nullopt;
+  }
+
+  llvm_unreachable("Unknown wrapped IR type");
+}
+
 // Return true when this is a pass for which changes should be ignored
 bool isIgnored(StringRef PassID) {
   return isSpecialPass(PassID,
@@ -2452,15 +2498,34 @@ StandardInstrumentations::StandardInstrumentations(
 PrintCrashIRInstrumentation *PrintCrashIRInstrumentation::CrashReporter =
     nullptr;
 
+void PrintCrashIRInstrumentation::printToStream(raw_ostream &OS) {
+  // We always print SavedString.
+  OS << SavedString;
+
+  if (SavedModule) {
+    // We avoid pushing all functions to ValuesToSave
+    // to avoid unnecessary memory pressure. We treat
+    // empty to mean all functions.
+    if (SavedFunctions.empty()) {
+      printIR(OS, SavedModule.get());
+    } else {
+      for (auto *GV : SavedFunctions) {
+        printIR(OS, cast<Function>(GV));
+      }
+    }
+  }
+}
+
 void PrintCrashIRInstrumentation::reportCrashIR() {
   if (!PrintOnCrashPath.empty()) {
     std::error_code EC;
     raw_fd_ostream Out(PrintOnCrashPath, EC);
     if (EC)
       report_fatal_error(errorCodeToError(EC));
-    Out << SavedIR;
+
+    printToStream(Out);
   } else {
-    dbgs() << SavedIR;
+    printToStream(dbgs());
   }
 }
 
@@ -2494,8 +2559,10 @@ void PrintCrashIRInstrumentation::registerCallbacks(
 
   PIC.registerBeforeNonSkippedPassCallback(
       [&PIC, this](StringRef PassID, Any IR) {
-        SavedIR.clear();
-        raw_string_ostream OS(SavedIR);
+        SavedModule.reset();
+        SavedString.clear();
+        raw_string_ostream OS(SavedString);
+
         OS << formatv("; *** Dump of {0}IR Before Last Pass {1}",
                       llvm::forcePrintModuleIR() ? "Module " : "", PassID);
         if (!isInteresting(IR, PassID, PIC.getPassNameForClassName(PassID))) {
@@ -2503,7 +2570,18 @@ void PrintCrashIRInstrumentation::registerCallbacks(
           return;
         }
         OS << " Started ***\n";
-        unwrapAndPrint(OS, IR);
+
+        // Cloning a Module is significantly faster than serializing to a string
+        // in most situations.
+        auto UseClonePair = unwrapAndSaveClone(IR);
+
+        if (UseClonePair.has_value()) {
+          SavedModule = std::move(UseClonePair->first);
+          SavedFunctions = UseClonePair->second;
+        } else {
+          // We only rely on SavedString for MachineFunctions & Loops
+          unwrapAndPrint(OS, IR);
+        }
       });
 }
 
diff --git a/llvm/test/Other/print-on-crash.ll b/llvm/test/Other/print-on-crash.ll
index f1e3414fb6352..8ea53b0884372 100644
--- a/llvm/test/Other/print-on-crash.ll
+++ b/llvm/test/Other/print-on-crash.ll
@@ -17,6 +17,8 @@
 
 ; 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
 
+; RUN: not --crash opt -print-on-crash -passes=trigger-crash-function < %s 2>&1 | FileCheck %s --check-prefix=CHECK_FUNCTION
+
 ; CHECK_SIMPLE: ; *** Dump of IR Before Last Pass {{.*}} Started ***
 ; CHECK_SIMPLE: @main
 ; CHECK_SIMPLE: entry:
@@ -24,6 +26,8 @@
 ; CHECK_MODULE: *** Dump of Module IR Before Last Pass {{.*}} Started ***
 ; CHECK_MODULE: ; ModuleID = {{.*}}
 ; CHECK_FILTERED: *** Dump of Module IR Before Last Pass {{.*}} Filtered Out ***
+; CHECK_FUNCTION: *** Dump of IR Before Last Pass {{.*}} Started ***
+; CHECK_FUNCTION: define i32 @main()
 
 define i32 @main() {
 entry: