[StandardInstrumentation] Add YieldCallback to NewPM

[annamthomas]

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.

[github-actions]

✅ With the latest revision this PR passed the Python code formatter.

[github-actions]

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff 05dc5d927bc5f0eeadd277a5e31b04399415f9d0 b611ac9c299dcb4ff0ee9a312751375230fead10 -- llvm/include/llvm/Passes/StandardInstrumentations.h llvm/lib/Passes/StandardInstrumentations.cpp

View the diff from clang-format here.

diff --git a/llvm/include/llvm/Passes/StandardInstrumentations.h b/llvm/include/llvm/Passes/StandardInstrumentations.h
index 258602d09b..d47878056b 100644
--- a/llvm/include/llvm/Passes/StandardInstrumentations.h
+++ b/llvm/include/llvm/Passes/StandardInstrumentations.h
@@ -451,7 +451,6 @@ public:
   void registerCallbacks(PassInstrumentationCallbacks &PIC);
 };
 
-
 class YieldInstrumentation {
   LLVMContext &Context;
 
diff --git a/llvm/lib/Passes/StandardInstrumentations.cpp b/llvm/lib/Passes/StandardInstrumentations.cpp
index 563b1821bf..5e03eaed38 100644
--- a/llvm/lib/Passes/StandardInstrumentations.cpp
+++ b/llvm/lib/Passes/StandardInstrumentations.cpp
@@ -2346,8 +2346,7 @@ void DotCfgChangeReporter::registerCallbacks(
 StandardInstrumentations::StandardInstrumentations(
     LLVMContext &Context, bool DebugLogging, bool VerifyEach,
     PrintPassOptions PrintPassOpts)
-    : PrintPass(DebugLogging, PrintPassOpts),
-      OptNone(DebugLogging),
+    : PrintPass(DebugLogging, PrintPassOpts), OptNone(DebugLogging),
       OptPassGate(Context),
       PrintChangedIR(PrintChanged == ChangePrinter::Verbose),
       PrintChangedDiff(PrintChanged == ChangePrinter::DiffVerbose ||

[aeubanks]

LLVMContext::yield isn't a great API. it claims that it doesn't need to be called, plus you can only have one callback. could you instead add your own pass instrumentation downstream that does what you want it to do? pass instrumentation is more principled in when it'll be called. and with the work on the new PM for codegen, hopefully you can retire your yield callback and always use pass instrumentation at some point

[annamthomas]

@aeubanks I'm not sure I understand. If end users of LLVM have a yield callback implemented, this implementation can be like "interrupt a compilation that is going on since it consumes too much memory" or "interrupt a compilation if it takes too much compile time and figure out why". We can add as many "pathological checks" as needed to decide what to do with a compile.

For example, in our case, we have some checks to do memory consumption instead of relying on mallinfo/mallinfo2 (which is buggy in older versions) to decide if a compilation is consuming too much memory.

When you tie LLVM to something like a JIT (our use case), having a yield callback gives more functionality to drop compilations which have pathological behaviour and even figure out what these pathological behaviour is, so that we crash with artifacts we need. Is the problem with "yield" that it is open ended? My hope was that once codegen moves to NewPM, we will also have these yield callback after every pass (rather than after function/module pass as it is now).

[aeubanks]

I'm saying that your use case makes sense, but I'd move away specifically from LLVMContext::yield and just directly use your own custom pass instrumentation (which you can register when creating a PassBuilder, e.g. like this).

Basically instead of this PR, have your downstream add a pass instrumentation that does the JIT'y thing you want (aka call the same thing your current yield callback calls). That way you get more guarantees on when it's called (after every pass), rather than LLVMContext::yield's "we might never call this" note.

[annamthomas]

makes sense. thank you.