From 3a781efe9704593edb641dcb170c0b35cd93fa1c Mon Sep 17 00:00:00 2001 From: gbaraldi Date: Fri, 5 Sep 2025 14:12:24 -0300 Subject: [PATCH] Implement callback invoking on the NewPM API --- deps/LLVMExtra/lib/NewPM.cpp | 50 +++++++++++++++++++++++++++++++- deps/LLVMExtra/lib/callbacks.inc | 25 ++++++++++++++++ src/newpm.jl | 16 ++++++++++ test/newpm.jl | 27 +++++++++++++++++ 4 files changed, 117 insertions(+), 1 deletion(-) create mode 100644 deps/LLVMExtra/lib/callbacks.inc diff --git a/deps/LLVMExtra/lib/NewPM.cpp b/deps/LLVMExtra/lib/NewPM.cpp index ccd65adb..47321f47 100644 --- a/deps/LLVMExtra/lib/NewPM.cpp +++ b/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 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 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 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 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; diff --git a/deps/LLVMExtra/lib/callbacks.inc b/deps/LLVMExtra/lib/callbacks.inc new file mode 100644 index 00000000..1369e306 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/newpm.jl b/src/newpm.jl index 1d96894a..422d8774 100644 --- a/src/newpm.jl +++ b/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" 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 diff --git a/test/newpm.jl b/test/newpm.jl index abcc4aa0..68ae97c8 100644 --- a/test/newpm.jl +++ b/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