[llvm] [StandardInstrumentation] Add YieldCallback to NewPM (PR #86561)

[via llvm-commits]

https://github.com/annamthomas created https://github.com/llvm/llvm-project/pull/86561

Currently, this is done in Legacy Pass Manager which is used only in code gen (and previously when we were using legacy PM in opt). When we switched to new PM, we lost this functionality.

Tested this with our downstream yield callback implementation to interrupt compilation when compile takes more than X seconds.

>From b611ac9c299dcb4ff0ee9a312751375230fead10 Mon Sep 17 00:00:00 2001
From: Anna Thomas <anna at azul.com>
Date: Mon, 25 Mar 2024 14:31:29 -0400
Subject: [PATCH] [StandardInstrumentation] Add YieldCallback to NewPM

Currently, this is done in Legacy Pass Manager which is used only in
code gen (and previously when we were using legacy PM in opt). When we
switched to new PM, we lost this functionality.

Tested this with our downstream yield callback implementation to
interrupt compilation when compile takes more than X seconds.
---
 llvm/include/llvm/Passes/StandardInstrumentations.h | 10 ++++++++++
 llvm/lib/Passes/StandardInstrumentations.cpp        | 13 ++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/Passes/StandardInstrumentations.h b/llvm/include/llvm/Passes/StandardInstrumentations.h
index 8c6a44876d5459..258602d09b4fc1 100644
--- a/llvm/include/llvm/Passes/StandardInstrumentations.h
+++ b/llvm/include/llvm/Passes/StandardInstrumentations.h
@@ -451,6 +451,15 @@ class VerifyInstrumentation {
   void registerCallbacks(PassInstrumentationCallbacks &PIC);
 };
 
+
+class YieldInstrumentation {
+  LLVMContext &Context;
+
+public:
+  YieldInstrumentation(LLVMContext &Context) : Context(Context) {}
+  void registerCallbacks(PassInstrumentationCallbacks &PIC);
+};
+
 /// This class implements --time-trace functionality for new pass manager.
 /// It provides the pass-instrumentation callbacks that measure the pass
 /// execution time. They collect time tracing info by TimeProfiler.
@@ -580,6 +589,7 @@ class StandardInstrumentations {
   PrintCrashIRInstrumentation PrintCrashIR;
   IRChangedTester ChangeTester;
   VerifyInstrumentation Verify;
+  YieldInstrumentation YieldAfterPass;
 
   bool VerifyEach;
 
diff --git a/llvm/lib/Passes/StandardInstrumentations.cpp b/llvm/lib/Passes/StandardInstrumentations.cpp
index 697988b3fc7c0b..563b1821bf382a 100644
--- a/llvm/lib/Passes/StandardInstrumentations.cpp
+++ b/llvm/lib/Passes/StandardInstrumentations.cpp
@@ -2355,7 +2355,7 @@ StandardInstrumentations::StandardInstrumentations(
                        PrintChanged == ChangePrinter::ColourDiffVerbose ||
                            PrintChanged == ChangePrinter::ColourDiffQuiet),
       WebsiteChangeReporter(PrintChanged == ChangePrinter::DotCfgVerbose),
-      Verify(DebugLogging), VerifyEach(VerifyEach) {}
+      Verify(DebugLogging), YieldAfterPass(Context), VerifyEach(VerifyEach) {}
 
 PrintCrashIRInstrumentation *PrintCrashIRInstrumentation::CrashReporter =
     nullptr;
@@ -2415,6 +2415,16 @@ void PrintCrashIRInstrumentation::registerCallbacks(
       });
 }
 
+void YieldInstrumentation::registerCallbacks(
+    PassInstrumentationCallbacks &PIC) {
+  PIC.registerAfterPassCallback(
+      [this](StringRef P, Any IR, const PreservedAnalyses &PassPA) {
+        if (isIgnored(P))
+          return;
+        Context.yield();
+      });
+}
+
 void StandardInstrumentations::registerCallbacks(
     PassInstrumentationCallbacks &PIC, ModuleAnalysisManager *MAM) {
   PrintIR.registerCallbacks(PIC);
@@ -2432,6 +2442,7 @@ void StandardInstrumentations::registerCallbacks(
   PrintCrashIR.registerCallbacks(PIC);
   if (MAM)
     PreservedCFGChecker.registerCallbacks(PIC, *MAM);
+  YieldAfterPass.registerCallbacks(PIC);
 
   // TimeProfiling records the pass running time cost.
   // Its 'BeforePassCallback' can be appended at the tail of all the