[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"
@@ -480,6 +481,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)
@@ -529,6 +550,96 @@ 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"
+        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"
+        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"
+        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"
+        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"
+        return false;
+      });
 }
 
 void PassBuilder::registerModuleAnalyses(ModuleAnalysisManager &MAM) {
@@ -614,26 +725,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

@@ -726,7 +726,7 @@ FUNCTION_PASS_WITH_PARAMS(
       return ExpandFpPass(TM, OL);
     },
     parseExpandFpOptions, "O0;O1;O2;O3")
-    
+
 #undef FUNCTION_PASS_WITH_PARAMS
 
 #ifndef LOOPNEST_PASS
@@ -805,3 +805,35 @@ LOOP_PASS_WITH_PARAMS(
     },
     parseLoopUnswitchOptions, "nontrivial;no-nontrivial;trivial;no-trivial")
 #undef LOOP_PASS_WITH_PARAMS
+
+#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)
+FUNCTION_CALLBACK("vectorizer-end-callbacks", invokeVectorizerEndEPCallbacks)
+#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

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(vectorizer-start-callbacks<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='peephole-callbacks<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='pipeline-start-callbacks<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='pipeline-early-simplification-callbacks<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='optimizer-early-callbacks<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='optimizer-last-callbacks<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(scalar-optimizer-late-callbacks<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(vectorizer-end-callbacks<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(late-loop-optimizations-callbacks<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(loop-optimizer-end-callbacks<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(cgscc-optimizer-late-callbacks<O2>)' < %s 2>&1 | FileCheck %s --check-prefix=CGSCCTLATE
+; RUN: not opt -disable-output -passes='vectorizer-start-callbacks<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
+}