Implement callback invoking on the NewPM API

Author: gbaraldi

deps/LLVMExtra/lib/NewPM.cpp

@@ -143,6 +143,54 @@ void LLVMPassBuilderExtensionsSetAAPipeline(LLVMPassBuilderExtensionsRef Extensi
 }
 #endif
 
+// Register target specific parsing callbacks
+static void registerCallbackParsing(PassBuilder &PB) {
+    PB.registerPipelineParsingCallback(
+        [&](StringRef Name, FunctionPassManager &PM,
+           ArrayRef<PassBuilder::PipelineElement> InnerPipeline) {
+#define FUNCTION_CALLBACK(NAME, INVOKE) if (Name == NAME) { PB.INVOKE(PM, OptimizationLevel::O2); return true; }
+#include "callbacks.inc"
+#undef FUNCTION_CALLBACK
+            return false;
+        }
+      );
+
+    PB.registerPipelineParsingCallback(
+        [&](StringRef Name, ModulePassManager &PM,
+           ArrayRef<PassBuilder::PipelineElement> InnerPipeline) {
+#define MODULE_CALLBACK(NAME, INVOKE) if (Name == NAME) { PB.INVOKE(PM, OptimizationLevel::O2); return true; }
+#include "callbacks.inc"
+#undef MODULE_CALLBACK
+#if LLVM_VERSION_MAJOR >= 20
+#define MODULE_LTO_CALLBACK(NAME, INVOKE) if (Name == NAME) { PB.INVOKE(PM, OptimizationLevel::O2, ThinOrFullLTOPhase::None); return true; }
+#else
+#define MODULE_LTO_CALLBACK(NAME, INVOKE) if (Name == NAME) { PB.INVOKE(PM, OptimizationLevel::O2); return true; }
+#endif
+#include "callbacks.inc"
+#undef MODULE_LTO_CALLBACK
+            return false;
+        }
+      );
+
+    PB.registerPipelineParsingCallback(
+        [&](StringRef Name, CGSCCPassManager &CGPM,
+           ArrayRef<PassBuilder::PipelineElement> InnerPipeline) {
+#define CGSCC_CALLBACK(NAME, INVOKE) if (Name == NAME) { PB.INVOKE(CGPM, OptimizationLevel::O2); return true; }
+#include "callbacks.inc"
+#undef CGSCC_CALLBACK
+          return false;
+        }
+      );
+    PB.registerPipelineParsingCallback(
+        [&](StringRef Name, LoopPassManager &PM,
+           ArrayRef<PassBuilder::PipelineElement> InnerPipeline) {
+#define LOOP_CALLBACK(NAME, INVOKE) if (Name == NAME) { PB.INVOKE(PM, OptimizationLevel::O2); return true; }
+#include "callbacks.inc"
+#undef LOOP_CALLBACK
+          return false;
+      }
+    );
+}
 
 // Vendored API entrypoint
 
@@ -164,7 +212,7 @@ static LLVMErrorRef runJuliaPasses(Module *Mod, Function *Fun, const char *Passe
     PB.registerPipelineParsingCallback(Callback);
   for (auto &Callback : PassExts->FunctionPipelineParsingCallbacks)
     PB.registerPipelineParsingCallback(Callback);
-
+  registerCallbackParsing(PB); // Parsing for target-specific callbacks
   LoopAnalysisManager LAM;
   FunctionAnalysisManager FAM;
   CGSCCAnalysisManager CGAM;

deps/LLVMExtra/lib/callbacks.inc

@@ -0,0 +1,25 @@
+#ifdef FUNCTION_CALLBACK
+FUNCTION_CALLBACK("PeepholeCallbacks", invokePeepholeEPCallbacks)
+FUNCTION_CALLBACK("ScalarOptimizerLateCallbacks", invokeScalarOptimizerLateEPCallbacks)
+FUNCTION_CALLBACK("VectorizerStartCallbacks", invokeVectorizerStartEPCallbacks)
+#endif
+
+#ifdef LOOP_CALLBACK
+LOOP_CALLBACK("LateLoopOptimizationsCallbacks", invokeLateLoopOptimizationsEPCallbacks)
+LOOP_CALLBACK("LoopOptimizerEndCallbacks", invokeLoopOptimizerEndEPCallbacks)
+#endif
+
+#ifdef MODULE_CALLBACK
+MODULE_CALLBACK("PipelineStartCallbacks", invokePipelineStartEPCallbacks)
+#endif
+
+// There are some full lto specific ones that are ignored here
+#ifdef MODULE_LTO_CALLBACK
+MODULE_LTO_CALLBACK("PipelineEarlySimplificationCallbacks", invokePipelineEarlySimplificationEPCallbacks)
+MODULE_LTO_CALLBACK("OptimizerEarlyCallbacks", invokeOptimizerEarlyEPCallbacks)
+MODULE_LTO_CALLBACK("OptimizerLastCallbacks", invokeOptimizerLastEPCallbacks)
+#endif
+
+#ifdef CGSCC_CALLBACK
+CGSCC_CALLBACK("CGSCCOptimizerLateCallbacks", invokeCGSCCOptimizerLateEPCallbacks)
+#endif

src/newpm.jl

@@ -514,6 +514,11 @@ function InternalizePass(; preserved_gvs::Vector=String[], kwargs...)
 
     "internalize" * kwargs_to_params(kwargs)
 end
+# Module callbacks
+@module_pass "PipelineStartCallbacks" PipelineStartCallbacks
+@module_pass "PipelineEarlySimplificationCallbacks" PipelineEarlySimplificationCallbacks
+@module_pass "OptimizerEarlyCallbacks" OptimizerEarlyCallbacks
+@module_pass "OptimizerLastCallbacks" OptimizerLastCallbacks
 
 # CGSCC passes
 
@@ -526,6 +531,9 @@ end
 @cgscc_pass "inline" InlinerPass
 @cgscc_pass "coro-split" CoroSplitPass
 
+#CGSCC callbacks
+@cgscc_pass "CGSCCOptimizerLateCallbacks" CGSCCOptimizerLateCallbacks
+
 # function passes
 
 @function_pass "aa-eval" AAEvaluator
@@ -709,6 +717,10 @@ end
 @function_pass "gvn" GVNPass
 @function_pass "print<stack-lifetime>" StackLifetimePrinterPass
 
+# Function pass callbacks
+@function_pass "PeepholeCallbacks" PeepholeCallbacks
+@function_pass "ScalarOptimizerLateCallbacks" ScalarOptimizerLateCallbacks
+@function_pass "VectorizerStartCallbacks" VectorizerStartCallbacks
 # loop nest passes
 
 @loop_pass "loop-flatten" LoopFlattenPass
@@ -746,6 +758,10 @@ end
 @loop_pass "licm" LICMPass
 @loop_pass "lnicm" LNICMPass
 
+# Loop Callbacks
+@loop_pass "LateLoopOptimizationsCallbacks" LateLoopOptimizationsCallbacks
+@loop_pass "LoopOptimizerEndCallbacks" LoopOptimizerEndCallbacks
+
 
 ## alias analyses
 

test/newpm.jl

@@ -409,4 +409,31 @@ end
     end
 end
 
+@testset "callbacks" begin
+    # just check that the callbacks can be registered and run without errors
+    @dispose ctx=Context() begin
+        # module callbacks
+        @dispose pb=NewPMPassBuilder() mod=test_module() begin
+            run!("PipelineStartCallbacks", mod) === nothing
+        end
+        @dispose pb=NewPMPassBuilder() mod=test_module() begin
+            run!(PipelineStartCallbacks(), mod) === nothing
+        end
+        # CGSCC callback
+        @dispose pb=NewPMPassBuilder() mod=test_module() begin
+            run!("CGSCCOptimizerLateCallbacks", mod) === nothing
+        end
+        @dispose pb=NewPMPassBuilder() mod=test_module() begin
+            run!(CGSCCOptimizerLateCallbacks(), mod) === nothing
+        end
+        # function callbacks
+        @dispose pb=NewPMPassBuilder() mod=test_module() begin
+            run!("PeepholeCallbacks", mod) === nothing
+        end
+        @dispose pb=NewPMPassBuilder() mod=test_module() begin
+            run!(PeepholeCallbacks(), mod) === nothing
+        end
+    end
+end
+
 end