Implement callback invoking on the NewPM API

deps/LLVMExtra/lib/NewPM.cpp

@@ -6,6 +6,8 @@
 #include <llvm/Passes/PassBuilder.h>
 #include <llvm/Passes/StandardInstrumentations.h>
 #include <llvm/Support/CBindingWrapping.h>
+#include <llvm/Support/FormatVariadic.h>
+#include <optional>
 
 using namespace llvm;
 
@@ -143,6 +145,167 @@ void LLVMPassBuilderExtensionsSetAAPipeline(LLVMPassBuilderExtensionsRef Extensi
 }
 #endif
 
+static bool checkParametrizedPassName(StringRef Name, StringRef PassName) {
+  if (!Name.consume_front(PassName))
+    return false;
+  // normal pass name w/o parameters == default parameters
+  if (Name.empty())
+    return true;
+#if LLVM_VERSION_MAJOR >= 16
+  return Name.starts_with("<") && Name.ends_with(">");
+#else
+  return Name.startswith("<") && Name.endswith(">");
+#endif
+}
+
+static std::optional<OptimizationLevel> parseOptLevel(StringRef S) {
+  return StringSwitch<std::optional<OptimizationLevel>>(S)
+      .Case("O0", OptimizationLevel::O0)
+      .Case("O1", OptimizationLevel::O1)
+      .Case("O2", OptimizationLevel::O2)
+      .Case("O3", OptimizationLevel::O3)
+      .Case("Os", OptimizationLevel::Os)
+      .Case("Oz", OptimizationLevel::Oz)
+      .Default(std::nullopt);
+}
+
+static Expected<OptimizationLevel> parseOptLevelParam(StringRef S) {
+  std::optional<OptimizationLevel> OptLevel = parseOptLevel(S);
+  if (OptLevel)
+    return *OptLevel;
+  return make_error<StringError>(
+      formatv("invalid optimization level '{}'", S).str(),
+      inconvertibleErrorCode());
+}
+
+template <typename ParametersParseCallableT>
+static auto parsePassParameters(ParametersParseCallableT &&Parser,
+                                StringRef Name, StringRef PassName)
+    -> decltype(Parser(StringRef{})) {
+  using ParametersT = typename decltype(Parser(StringRef{}))::value_type;
+
+  StringRef Params = Name;
+  if (!Params.consume_front(PassName)) {
+    llvm_unreachable(
+        "unable to strip pass name from parametrized pass specification");
+  }
+  if (!Params.empty() &&
+      (!Params.consume_front("<") || !Params.consume_back(">"))) {
+    llvm_unreachable("invalid format for parametrized pass name");
+  }
+
+  Expected<ParametersT> Result = Parser(Params);
+  assert((Result || Result.template errorIsA<StringError>()) &&
+          "Pass parameter parser can only return StringErrors.");
+  return Result;
+}
+
+
+// Register target specific parsing callbacks
+static void registerCallbackParsing(PassBuilder &PB) {
+  PB.registerPipelineParsingCallback(
+      [&](StringRef Name, ModulePassManager &PM,
+             ArrayRef<PassBuilder::PipelineElement>) {
+#define MODULE_CALLBACK(NAME, INVOKE)                                          \
+    if (checkParametrizedPassName(Name, NAME)) {                    \
+      auto L = parsePassParameters(parseOptLevelParam, Name, NAME); \
+      if (!L) {                                                                  \
+        errs() << NAME ": " << toString(L.takeError()) << '\n';                  \
+        return false;                                                            \
+      }                                                                          \
+      PB.INVOKE(PM, L.get());                                                       \
+      return true;                                                               \
+    }
+    #include "callbacks.inc"
+    return false;
+  });
+
+  // Module-level callbacks with LTO phase (use Phase::None for string API)
+  PB.registerPipelineParsingCallback(
+      [&](StringRef Name, ModulePassManager &PM,
+             ArrayRef<PassBuilder::PipelineElement>) {
+#if LLVM_VERSION_MAJOR >= 20
+#define MODULE_LTO_CALLBACK(NAME, INVOKE)                                      \
+    if (checkParametrizedPassName(Name, NAME)) {                    \
+      auto L = parsePassParameters(parseOptLevelParam, Name, NAME); \
+      if (!L) {                                                                  \
+        errs() << NAME ": " << toString(L.takeError()) << '\n';                  \
+        return false;                                                            \
+      }                                                                          \
+      PB.INVOKE(PM, L.get(), ThinOrFullLTOPhase::None);                             \
+      return true;                                                               \
+    }
+    #include "callbacks.inc"
+#else
+#define MODULE_LTO_CALLBACK(NAME, INVOKE)                                      \
+    if (checkParametrizedPassName(Name, NAME)) {                    \
+      auto L = parsePassParameters(parseOptLevelParam, Name, NAME); \
+      if (!L) {                                                                  \
+        errs() << NAME ": " << toString(L.takeError()) << '\n';                  \
+        return false;                                                            \
+      }                                                                          \
+      PB.INVOKE(PM, L.get());                             \
+      return true;                                                               \
+    }
+    #include "callbacks.inc"
+#endif
+    return false;
+  });
+
+  // Function-level callbacks
+  PB.registerPipelineParsingCallback(
+      [&](StringRef Name, FunctionPassManager &PM,
+             ArrayRef<PassBuilder::PipelineElement>) {
+#define FUNCTION_CALLBACK(NAME, INVOKE)                                        \
+    if (checkParametrizedPassName(Name, NAME)) {                    \
+      auto L = parsePassParameters(parseOptLevelParam, Name, NAME); \
+      if (!L) {                                                                  \
+        errs() << NAME ": " << toString(L.takeError()) << '\n';                  \
+        return false;                                                            \
+      }                                                                          \
+      PB.INVOKE(PM, L.get());                                                       \
+      return true;                                                               \
+    }
+    #include "callbacks.inc"
+    return false;
+  });
+
+  // CGSCC-level callbacks
+  PB.registerPipelineParsingCallback(
+      [&](StringRef Name, CGSCCPassManager &PM,
+             ArrayRef<PassBuilder::PipelineElement>) {
+#define CGSCC_CALLBACK(NAME, INVOKE)                                           \
+    if (checkParametrizedPassName(Name, NAME)) {                    \
+      auto L = parsePassParameters(parseOptLevelParam, Name, NAME); \
+      if (!L) {                                                                  \
+        errs() << NAME ": " << toString(L.takeError()) << '\n';                  \
+        return false;                                                            \
+      }                                                                          \
+      PB.INVOKE(PM, L.get());                                                       \
+      return true;                                                               \
+    }
+    #include "callbacks.inc"
+    return false;
+  });
+
+  // Loop-level callbacks
+  PB.registerPipelineParsingCallback(
+      [&](StringRef Name, LoopPassManager &PM,
+             ArrayRef<PassBuilder::PipelineElement>) {
+#define LOOP_CALLBACK(NAME, INVOKE)                                            \
+      if (checkParametrizedPassName(Name, NAME)) {                    \
+        auto L = parsePassParameters(parseOptLevelParam, Name, NAME); \
+        if (!L) {                                                                  \
+          errs() << NAME ": " << toString(L.takeError()) << '\n';                  \
+          return false;                                                            \
+        }                                                                          \
+        PB.INVOKE(PM, L.get());                                                       \
+        return true;                                                               \
+      }
+    #include "callbacks.inc"
+    return false;
+  });
+}
 
 // Vendored API entrypoint
 
@@ -164,7 +327,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,33 @@
+#ifdef MODULE_CALLBACK
+MODULE_CALLBACK("pipeline-start-callbacks", invokePipelineStartEPCallbacks)
+#endif
+#undef MODULE_CALLBACK
+
+// There are some full lto specific ones that are ignored here for now
+#ifdef MODULE_LTO_CALLBACK
+MODULE_LTO_CALLBACK("pipeline-early-simplification-callbacks", invokePipelineEarlySimplificationEPCallbacks)
+MODULE_LTO_CALLBACK("optimizer-early-callbacks", invokeOptimizerEarlyEPCallbacks)
+MODULE_LTO_CALLBACK("optimizer-last-callbacks", invokeOptimizerLastEPCallbacks)
+#endif
+#undef MODULE_LTO_CALLBACK
+
+#ifdef FUNCTION_CALLBACK
+FUNCTION_CALLBACK("peephole-callbacks", invokePeepholeEPCallbacks)
+FUNCTION_CALLBACK("scalar-optimizer-late-callbacks", invokeScalarOptimizerLateEPCallbacks)
+FUNCTION_CALLBACK("vectorizer-start-callbacks", invokeVectorizerStartEPCallbacks)
+#if LLVM_VERSION_MAJOR >= 21
+FUNCTION_CALLBACK("vectorizer-end-callbacks", invokeVectorizerEndEPCallbacks)
+#endif
+#endif
+#undef FUNCTION_CALLBACK
+
+#ifdef CGSCC_CALLBACK
+CGSCC_CALLBACK("cgscc-optimizer-late-callbacks", invokeCGSCCOptimizerLateEPCallbacks)
+#endif
+#undef CGSCC_CALLBACK
+
+#ifdef LOOP_CALLBACK
+LOOP_CALLBACK("late-loop-optimizations-callbacks", invokeLateLoopOptimizationsEPCallbacks)
+LOOP_CALLBACK("loop-optimizer-end-callbacks", invokeLoopOptimizerEndEPCallbacks)
+#endif
+#undef LOOP_CALLBACK

deps/build_local.jl

@@ -47,7 +47,7 @@ LLVM_DIR = joinpath(LLVM.artifact_dir, "lib", "cmake", "llvm")
 
 # build and install
 @info "Building" source_dir scratch_dir build_dir LLVM_DIR
-cmake() do cmake_path
+cmake(adjust_LIBPATH=!Sys.iswindows()) do cmake_path
     config_opts = `-DLLVM_ROOT=$(LLVM_DIR) -DCMAKE_INSTALL_PREFIX=$(scratch_dir)`
     if Sys.iswindows()
         # prevent picking up MSVC

src/newpm.jl

@@ -514,6 +514,29 @@ function InternalizePass(; preserved_gvs::Vector=String[], kwargs...)
 
     "internalize" * kwargs_to_params(kwargs)
 end
+# Module callbacks (not part of general pass sweep)
+export PipelineStartCallbacks, PipelineEarlySimplificationCallbacks,
+       OptimizerEarlyCallbacks, OptimizerLastCallbacks
+function PipelineStartCallbacks(; opt_level=0)
+    opts = Dict{Symbol,Any}()
+    opts[Symbol("O$opt_level")] = true
+    "pipeline-start-callbacks" * kwargs_to_params(opts)
+end
+function PipelineEarlySimplificationCallbacks(; opt_level=0)
+    opts = Dict{Symbol,Any}()
+    opts[Symbol("O$opt_level")] = true
+    "pipeline-early-simplification-callbacks" * kwargs_to_params(opts)
+end
+function OptimizerEarlyCallbacks(; opt_level=0)
+    opts = Dict{Symbol,Any}()
+    opts[Symbol("O$opt_level")] = true
+    "optimizer-early-callbacks" * kwargs_to_params(opts)
+end
+function OptimizerLastCallbacks(; opt_level=0)
+    opts = Dict{Symbol,Any}()
+    opts[Symbol("O$opt_level")] = true
+    "optimizer-last-callbacks" * kwargs_to_params(opts)
+end
 
 # CGSCC passes
 
@@ -526,6 +549,14 @@ end
 @cgscc_pass "inline" InlinerPass
 @cgscc_pass "coro-split" CoroSplitPass
 
+#CGSCC callbacks (not part of general pass sweep)
+export CGSCCOptimizerLateCallbacks
+function CGSCCOptimizerLateCallbacks(; opt_level=0)
+    opts = Dict{Symbol,Any}()
+    opts[Symbol("O$opt_level")] = true
+    "cgscc-optimizer-late-callbacks" * kwargs_to_params(opts)
+end
+
 # function passes
 
 @function_pass "aa-eval" AAEvaluator
@@ -709,6 +740,31 @@ end
 @function_pass "gvn" GVNPass
 @function_pass "print<stack-lifetime>" StackLifetimePrinterPass
 
+# Function pass callbacks (not part of general pass sweep)
+export PeepholeCallbacks, ScalarOptimizerLateCallbacks, VectorizerStartCallbacks
+function PeepholeCallbacks(; opt_level=0)
+    opts = Dict{Symbol,Any}()
+    opts[Symbol("O$opt_level")] = true
+    "peephole-callbacks" * kwargs_to_params(opts)
+end
+function ScalarOptimizerLateCallbacks(; opt_level=0)
+    opts = Dict{Symbol,Any}()
+    opts[Symbol("O$opt_level")] = true
+    "scalar-optimizer-late-callbacks" * kwargs_to_params(opts)
+end
+function VectorizerStartCallbacks(; opt_level=0)
+    opts = Dict{Symbol,Any}()
+    opts[Symbol("O$opt_level")] = true
+    "vectorizer-start-callbacks" * kwargs_to_params(opts)
+end
+@static if version() >= v"21"
+    export VectorizerEndCallbacks
+    function VectorizerEndCallbacks(; opt_level=0)
+        opts = Dict{Symbol,Any}()
+        opts[Symbol("O$opt_level")] = true
+        "vectorizer-end-callbacks" * kwargs_to_params(opts)
+    end
+end
 # loop nest passes
 
 @loop_pass "loop-flatten" LoopFlattenPass
@@ -746,6 +802,19 @@ end
 @loop_pass "licm" LICMPass
 @loop_pass "lnicm" LNICMPass
 
+# Loop Callbacks (not part of general pass sweep)
+export LateLoopOptimizationsCallbacks, LoopOptimizerEndCallbacks
+function LateLoopOptimizationsCallbacks(; opt_level=0)
+    opts = Dict{Symbol,Any}()
+    opts[Symbol("O$opt_level")] = true
+    "late-loop-optimizations-callbacks" * kwargs_to_params(opts)
+end
+function LoopOptimizerEndCallbacks(; opt_level=0)
+    opts = Dict{Symbol,Any}()
+    opts[Symbol("O$opt_level")] = true
+    "loop-optimizer-end-callbacks" * kwargs_to_params(opts)
+end
+
 
 ## alias analyses
 

test/newpm.jl

@@ -195,7 +195,12 @@ end
     end
 
     @testset "loop" begin
-        test_passes("loop", LLVM.loop_passes)
+        # skip opt-level callback pseudo-passes, they require parameters and are provided as functions
+        skip_loop = [
+            "late-loop-optimizations-callbacks",
+            "loop-optimizer-end-callbacks",
+        ]
+        test_passes("loop", LLVM.loop_passes, skip_loop)
     end
 end
 
@@ -409,4 +414,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
+            @test run!("pipeline-start-callbacks<O0>", mod) === nothing
+        end
+        @dispose pb=NewPMPassBuilder() mod=test_module() begin
+            @test run!(PipelineStartCallbacks(opt_level=0), mod) === nothing
+        end
+        # CGSCC callback
+        @dispose pb=NewPMPassBuilder() mod=test_module() begin
+            @test run!("cgscc-optimizer-late-callbacks<O0>", mod) === nothing
+        end
+        @dispose pb=NewPMPassBuilder() mod=test_module() begin
+            @test run!(CGSCCOptimizerLateCallbacks(opt_level=0), mod) === nothing
+        end
+        # function callbacks
+        @dispose pb=NewPMPassBuilder() mod=test_module() begin
+            @test run!("peephole-callbacks<O0>", mod) === nothing
+        end
+        @dispose pb=NewPMPassBuilder() mod=test_module() begin
+            @test run!(PeepholeCallbacks(opt_level=0), mod) === nothing
+        end
+    end
+end
+
 end