[PassBuilder] Add callback invoking to PassBuilder string API

llvm/lib/Passes/PassBuilder.cpp

@@ -188,6 +188,7 @@
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Regex.h"
@@ -479,6 +480,26 @@ class RequireAllMachineFunctionPropertiesPass
 
 } // namespace
 
+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());
+}
+
 PassBuilder::PassBuilder(TargetMachine *TM, PipelineTuningOptions PTO,
                          std::optional<PGOOptions> PGOOpt,
                          PassInstrumentationCallbacks *PIC)
@@ -528,6 +549,101 @@ PassBuilder::PassBuilder(TargetMachine *TM, PipelineTuningOptions PTO,
 #include "llvm/Passes/MachinePassRegistry.def"
     });
   }
+
+  // Module-level callbacks without LTO phase
+  registerPipelineParsingCallback(
+      [this](StringRef Name, ModulePassManager &PM,
+             ArrayRef<PassBuilder::PipelineElement>) {
+#define MODULE_CALLBACK(NAME, INVOKE)                                          \
+  if (PassBuilder::checkParametrizedPassName(Name, NAME)) {                    \
+    auto L = PassBuilder::parsePassParameters(parseOptLevelParam, Name, NAME); \
+    if (!L) {                                                                  \
+      errs() << NAME ": " << toString(L.takeError()) << '\n';                  \
+      return false;                                                            \
+    }                                                                          \
+    INVOKE(PM, L.get());                                                       \
+    return true;                                                               \
+  }
+#include "PassRegistry.def"
+#undef MODULE_CALLBACK
+        return false;
+      });
+
+  // Module-level callbacks with LTO phase (use Phase::None for string API)
+  registerPipelineParsingCallback(
+      [this](StringRef Name, ModulePassManager &PM,
+             ArrayRef<PassBuilder::PipelineElement>) {
+#define MODULE_LTO_CALLBACK(NAME, INVOKE)                                      \
+  if (PassBuilder::checkParametrizedPassName(Name, NAME)) {                    \
+    auto L = PassBuilder::parsePassParameters(parseOptLevelParam, Name, NAME); \
+    if (!L) {                                                                  \
+      errs() << NAME ": " << toString(L.takeError()) << '\n';                  \
+      return false;                                                            \
+    }                                                                          \
+    INVOKE(PM, L.get(), ThinOrFullLTOPhase::None);                             \
+    return true;                                                               \
+  }
+#include "PassRegistry.def"
+#undef MODULE_LTO_CALLBACK
+        return false;
+      });
+
+  // Function-level callbacks
+  registerPipelineParsingCallback(
+      [this](StringRef Name, FunctionPassManager &PM,
+             ArrayRef<PassBuilder::PipelineElement>) {
+#define FUNCTION_CALLBACK(NAME, INVOKE)                                        \
+  if (PassBuilder::checkParametrizedPassName(Name, NAME)) {                    \
+    auto L = PassBuilder::parsePassParameters(parseOptLevelParam, Name, NAME); \
+    if (!L) {                                                                  \
+      errs() << NAME ": " << toString(L.takeError()) << '\n';                  \
+      return false;                                                            \
+    }                                                                          \
+    INVOKE(PM, L.get());                                                       \
+    return true;                                                               \
+  }
+#include "PassRegistry.def"
+#undef FUNCTION_CALLBACK
+        return false;
+      });
+
+  // CGSCC-level callbacks
+  registerPipelineParsingCallback(
+      [this](StringRef Name, CGSCCPassManager &PM,
+             ArrayRef<PassBuilder::PipelineElement>) {
+#define CGSCC_CALLBACK(NAME, INVOKE)                                           \
+  if (PassBuilder::checkParametrizedPassName(Name, NAME)) {                    \
+    auto L = PassBuilder::parsePassParameters(parseOptLevelParam, Name, NAME); \
+    if (!L) {                                                                  \
+      errs() << NAME ": " << toString(L.takeError()) << '\n';                  \
+      return false;                                                            \
+    }                                                                          \
+    INVOKE(PM, L.get());                                                       \
+    return true;                                                               \
+  }
+#include "PassRegistry.def"
+#undef CGSCC_CALLBACK
+        return false;
+      });
+
+  // Loop-level callbacks
+  registerPipelineParsingCallback(
+      [this](StringRef Name, LoopPassManager &PM,
+             ArrayRef<PassBuilder::PipelineElement>) {
+#define LOOP_CALLBACK(NAME, INVOKE)                                            \
+  if (PassBuilder::checkParametrizedPassName(Name, NAME)) {                    \
+    auto L = PassBuilder::parsePassParameters(parseOptLevelParam, Name, NAME); \
+    if (!L) {                                                                  \
+      errs() << NAME ": " << toString(L.takeError()) << '\n';                  \
+      return false;                                                            \
+    }                                                                          \
+    INVOKE(PM, L.get());                                                       \
+    return true;                                                               \
+  }
+#include "PassRegistry.def"
+#undef LOOP_CALLBACK
+        return false;
+      });
 }
 
 void PassBuilder::registerModuleAnalyses(ModuleAnalysisManager &MAM) {
@@ -613,26 +729,6 @@ static std::optional<int> parseDevirtPassName(StringRef Name) {
   return Count;
 }
 
-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());
-}
-
 Expected<bool> PassBuilder::parseSinglePassOption(StringRef Params,
                                                   StringRef OptionName,
                                                   StringRef PassName) {

llvm/lib/Passes/PassRegistry.def

@@ -725,7 +725,7 @@ FUNCTION_PASS_WITH_PARAMS(
       return ExpandFpPass(TM, OL);
     },
     parseExpandFpOptions, "opt-level=N")
-    
+
 #undef FUNCTION_PASS_WITH_PARAMS
 
 #ifndef LOOPNEST_PASS
@@ -804,3 +804,31 @@ LOOP_PASS_WITH_PARAMS(
     },
     parseLoopUnswitchOptions, "nontrivial;no-nontrivial;trivial;no-trivial")
 #undef LOOP_PASS_WITH_PARAMS
+
+#ifdef MODULE_CALLBACK
+MODULE_CALLBACK("PipelineStartCallbacks", invokePipelineStartEPCallbacks)
+#endif
+
+// There are some full lto specific ones that are ignored here for now
+#ifdef MODULE_LTO_CALLBACK
+MODULE_LTO_CALLBACK("PipelineEarlySimplificationCallbacks", invokePipelineEarlySimplificationEPCallbacks)
+MODULE_LTO_CALLBACK("OptimizerEarlyCallbacks", invokeOptimizerEarlyEPCallbacks)
+MODULE_LTO_CALLBACK("OptimizerLastCallbacks", invokeOptimizerLastEPCallbacks)
+#endif
+
+#ifdef FUNCTION_CALLBACK
+FUNCTION_CALLBACK("PeepholeCallbacks", invokePeepholeEPCallbacks)
+FUNCTION_CALLBACK("ScalarOptimizerLateCallbacks", invokeScalarOptimizerLateEPCallbacks)
+FUNCTION_CALLBACK("VectorizerStartCallbacks", invokeVectorizerStartEPCallbacks)
+FUNCTION_CALLBACK("VectorizerEndCallbacks", invokeVectorizerEndEPCallbacks)
+#endif
+
+#ifdef LOOP_CALLBACK
+LOOP_CALLBACK("LateLoopOptimizationsCallbacks", invokeLateLoopOptimizationsEPCallbacks)
+LOOP_CALLBACK("LoopOptimizerEndCallbacks", invokeLoopOptimizerEndEPCallbacks)
+#endif
+
+#ifdef CGSCC_CALLBACK
+CGSCC_CALLBACK("CGSCCOptimizerLateCallbacks", invokeCGSCCOptimizerLateEPCallbacks)
+#endif
+

llvm/test/Other/pipeline-callbacks-string-api.ll

@@ -0,0 +1,52 @@
+; RUN: opt -disable-output -print-pipeline-passes \
+; RUN:   -passes-ep-vectorizer-start='no-op-function' \
+; RUN:   -passes='function(VectorizerStartCallbacks<O3>)' < %s 2>&1 | FileCheck %s --check-prefix=VECSTART
+; RUN: opt -disable-output -print-pipeline-passes \
+; RUN:   -passes-ep-peephole='no-op-function' \
+; RUN:   -passes='PeepholeCallbacks<Os>' < %s 2>&1 | FileCheck %s --check-prefix=PEEP
+; RUN: opt -disable-output -print-pipeline-passes \
+; RUN:   -passes-ep-pipeline-start='no-op-module' \
+; RUN:   -passes='PipelineStartCallbacks<O1>' < %s 2>&1 | FileCheck %s --check-prefix=MODSTART
+; RUN: opt -disable-output -print-pipeline-passes \
+; RUN:   -passes-ep-pipeline-early-simplification='no-op-module' \
+; RUN:   -passes='PipelineEarlySimplificationCallbacks<O2>' < %s 2>&1 | FileCheck %s --check-prefix=LTOEARLY
+; RUN: opt -disable-output -print-pipeline-passes \
+; RUN:   -passes-ep-optimizer-early='no-op-module' \
+; RUN:   -passes='OptimizerEarlyCallbacks<O2>' < %s 2>&1 | FileCheck %s --check-prefix=OPTEARLY
+; RUN: opt -disable-output -print-pipeline-passes \
+; RUN:   -passes-ep-optimizer-last='no-op-module' \
+; RUN:   -passes='OptimizerLastCallbacks<O2>' < %s 2>&1 | FileCheck %s --check-prefix=OPTLAST
+; RUN: opt -disable-output -print-pipeline-passes \
+; RUN:   -passes-ep-scalar-optimizer-late='no-op-function' \
+; RUN:   -passes='function(ScalarOptimizerLateCallbacks<O2>)' < %s 2>&1 | FileCheck %s --check-prefix=SCALATE
+; RUN: opt -disable-output -print-pipeline-passes \
+; RUN:   -passes-ep-vectorizer-end='no-op-function' \
+; RUN:   -passes='function(VectorizerEndCallbacks<O3>)' < %s 2>&1 | FileCheck %s --check-prefix=VECEND
+; RUN: opt -disable-output -print-pipeline-passes \
+; RUN:   -passes-ep-late-loop-optimizations='no-op-loop' \
+; RUN:   -passes='loop(LateLoopOptimizationsCallbacks<O2>)' < %s 2>&1 | FileCheck %s --check-prefix=LATELOOP
+; RUN: opt -disable-output -print-pipeline-passes \
+; RUN:   -passes-ep-loop-optimizer-end='no-op-loop' \
+; RUN:   -passes='loop(LoopOptimizerEndCallbacks<O2>)' < %s 2>&1 | FileCheck %s --check-prefix=LOOPOPTEND
+; RUN: opt -disable-output -print-pipeline-passes \
+; RUN:   -passes-ep-cgscc-optimizer-late='no-op-cgscc' \
+; RUN:   -passes='cgscc(CGSCCOptimizerLateCallbacks<O2>)' < %s 2>&1 | FileCheck %s --check-prefix=CGSCCTLATE
+; RUN: not opt -disable-output -passes='VectorizerStartCallbacks<foo>' < %s 2>&1 | FileCheck %s --check-prefix=INVALID
+
+; VECSTART: no-op-function
+; PEEP: no-op-function
+; MODSTART: no-op-module
+; LTOEARLY: no-op-module
+; OPTEARLY: no-op-module
+; OPTLAST: no-op-module
+; SCALATE: no-op-function
+; VECEND: no-op-function
+; LATELOOP: no-op-loop
+; LOOPOPTEND: no-op-loop
+; CGSCCTLATE: no-op-cgscc
+; INVALID: invalid optimization level 'foo'
+
+define void @f() {
+entry:
+  ret void
+}