[InlineAdvisor] Allow loading advisors as plugins

Adds the ability to load InlineAdvisors as plugins. This allows developing and distributing inlining heuristics outside of tree.

The PluginInlineAdvisorAnalysis class serves as the entry point for dynamic advisors. Plugins must register instances of this class to provide their own InliningAdvisor.

Reviewed By: mtrofin

Differential Revision: https://reviews.llvm.org/D139644

llvm/include/llvm/Analysis/InlineAdvisor.h

@@ -26,7 +26,7 @@ class ImportedFunctionsInliningStatistics;
 class OptimizationRemarkEmitter;
 struct ReplayInlinerSettings;
 
-/// There are 3 scenarios we can use the InlineAdvisor:
+/// There are 4 scenarios we can use the InlineAdvisor:
 /// - Default - use manual heuristics.
 ///
 /// - Release mode, the expected mode for production, day to day deployments.
@@ -39,6 +39,8 @@ struct ReplayInlinerSettings;
 /// requires the full C Tensorflow API library, and evaluates models
 /// dynamically. This mode also permits generating training logs, for offline
 /// training.
+///
+/// - Dynamically load an advisor via a plugin (PluginInlineAdvisorAnalysis)
 enum class InliningAdvisorMode : int { Default, Release, Development };
 
 // Each entry represents an inline driver.
@@ -238,6 +240,79 @@ class DefaultInlineAdvisor : public InlineAdvisor {
   InlineParams Params;
 };
 
+/// Used for dynamically registering InlineAdvisors as plugins
+///
+/// An advisor plugin adds a new advisor at runtime by registering an instance
+/// of PluginInlineAdvisorAnalysis in the current ModuleAnalysisManager.
+/// For example, the following code dynamically registers a
+/// DefaultInlineAdvisor:
+///
+/// namespace {
+///
+/// InlineAdvisor *defaultAdvisorFactory(Module &M, FunctionAnalysisManager
+/// &FAM,
+///                                      InlineParams Params, InlineContext IC)
+///                                      {
+///   return new DefaultInlineAdvisor(M, FAM, Params, IC);
+/// }
+///
+/// struct DefaultDynamicAdvisor : PassInfoMixin<DefaultDynamicAdvisor> {
+///   PreservedAnalyses run(Module &, ModuleAnalysisManager &MAM) {
+///     PluginInlineAdvisorAnalysis PA(defaultAdvisorFactory);
+///     MAM.registerPass([&] { return PA; });
+///     return PreservedAnalyses::all();
+///   }
+/// };
+///
+/// } // namespace
+///
+/// extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
+/// llvmGetPassPluginInfo() {
+///   return {LLVM_PLUGIN_API_VERSION, "DynamicDefaultAdvisor",
+///   LLVM_VERSION_STRING,
+///           [](PassBuilder &PB) {
+///             PB.registerPipelineStartEPCallback(
+///                 [](ModulePassManager &MPM, OptimizationLevel Level) {
+///                   MPM.addPass(DefaultDynamicAdvisor());
+///                 });
+///           }};
+/// }
+///
+/// A plugin must implement an AdvisorFactory and register it with a
+/// PluginInlineAdvisorAnlysis to the provided ModuleanAlysisManager.
+///
+/// If such a plugin has been registered
+/// InlineAdvisorAnalysis::Result::tryCreate will return the dynamically loaded
+/// advisor.
+///
+class PluginInlineAdvisorAnalysis
+    : public AnalysisInfoMixin<PluginInlineAdvisorAnalysis> {
+public:
+  static AnalysisKey Key;
+  static bool HasBeenRegistered;
+
+  typedef InlineAdvisor *(*AdvisorFactory)(Module &M,
+                                           FunctionAnalysisManager &FAM,
+                                           InlineParams Params,
+                                           InlineContext IC);
+
+  PluginInlineAdvisorAnalysis(AdvisorFactory Factory) : Factory(Factory) {
+    HasBeenRegistered = true;
+    assert(Factory != nullptr &&
+           "The plugin advisor factory should not be a null pointer.");
+  }
+
+  struct Result {
+    AdvisorFactory Factory;
+  };
+
+  Result run(Module &M, ModuleAnalysisManager &MAM) { return {Factory}; }
+  Result getResult() { return {Factory}; }
+
+private:
+  AdvisorFactory Factory;
+};
+
 /// The InlineAdvisorAnalysis is a module pass because the InlineAdvisor
 /// needs to capture state right before inlining commences over a module.
 class InlineAdvisorAnalysis : public AnalysisInfoMixin<InlineAdvisorAnalysis> {

llvm/lib/Analysis/InlineAdvisor.cpp

@@ -196,11 +196,18 @@ void InlineAdvice::recordInliningWithCalleeDeleted() {
 }
 
 AnalysisKey InlineAdvisorAnalysis::Key;
+AnalysisKey PluginInlineAdvisorAnalysis::Key;
+bool PluginInlineAdvisorAnalysis::HasBeenRegistered = false;
 
 bool InlineAdvisorAnalysis::Result::tryCreate(
     InlineParams Params, InliningAdvisorMode Mode,
     const ReplayInlinerSettings &ReplaySettings, InlineContext IC) {
   auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+  if (PluginInlineAdvisorAnalysis::HasBeenRegistered) {
+    auto &DA = MAM.getResult<PluginInlineAdvisorAnalysis>(M);
+    Advisor.reset(DA.Factory(M, FAM, Params, IC));
+    return !!Advisor;
+  }
   switch (Mode) {
   case InliningAdvisorMode::Default:
     LLVM_DEBUG(dbgs() << "Using default inliner heuristic.\n");

llvm/unittests/Analysis/CMakeLists.txt

@@ -2,18 +2,12 @@ set(LLVM_LINK_COMPONENTS
   Analysis
   AsmParser
   Core
+  Passes
   Support
   TransformUtils
   )
 
-set(MLGO_TESTS TFUtilsTest.cpp TrainingLoggerTest.cpp)
-if (LLVM_HAVE_TFLITE)
-  LIST(APPEND EXTRA_TESTS ${MLGO_TESTS})
-else()
-  LIST(APPEND LLVM_OPTIONAL_SOURCES ${MLGO_TESTS})
-endif()
-
-add_llvm_unittest_with_input_files(AnalysisTests
+set(ANALYSIS_TEST_SOURCES
   AliasAnalysisTest.cpp
   AliasSetTrackerTest.cpp
   AssumeBundleQueriesTest.cpp
@@ -42,6 +36,7 @@ add_llvm_unittest_with_input_files(AnalysisTests
   MemorySSATest.cpp
   MLModelRunnerTest.cpp
   PhiValuesTest.cpp
+  PluginInlineAdvisorAnalysisTest.cpp
   ProfileSummaryInfoTest.cpp
   ScalarEvolutionTest.cpp
   VectorFunctionABITest.cpp
@@ -53,7 +48,52 @@ add_llvm_unittest_with_input_files(AnalysisTests
   ValueLatticeTest.cpp
   ValueTrackingTest.cpp
   VectorUtilsTest.cpp
-  ${EXTRA_TESTS}
   )
 
- target_link_libraries(AnalysisTests PRIVATE LLVMTestingSupport)
+# The unit tests target does not use InlineAdvisorPlugin.cpp
+# so if not added to LLVM_OPTIONAL_SOURCES, FileCheck will
+# complain about unused file.
+set(LLVM_OPTIONAL_SOURCES InlineAdvisorPlugin.cpp)
+
+set(MLGO_TESTS TFUtilsTest.cpp TrainingLoggerTest.cpp)
+if (LLVM_HAVE_TFLITE)
+  LIST(APPEND ANALYSIS_TEST_SOURCES ${MLGO_TESTS})
+else()
+  LIST(APPEND LLVM_OPTIONAL_SOURCES ${MLGO_TESTS})
+endif()
+
+add_llvm_unittest_with_input_files(AnalysisTests
+  ${ANALYSIS_TEST_SOURCES}
+  )
+
+target_link_libraries(AnalysisTests PRIVATE LLVMTestingSupport)
+
+# The advisor plugin expects to not link against the Analysis, Support and Core 
+# libraries, but expects them to exist in the process loading the plugin. This 
+# doesn't work with DLLs on Windows (where a shared library can't have undefined
+# references), so just skip this testcase on Windows.
+if (NOT WIN32)
+  # On AIX, enable run-time linking to allow symbols from the plugins shared
+  # objects to be properly bound.
+  if(CMAKE_SYSTEM_NAME STREQUAL "AIX")
+    set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -Wl,-brtl")
+  endif()
+
+  # This plugin is built as a stand-alone plugin, so since we don't use the
+  # ANALYSIS_TEST_SOURCES files, we have to add them to LLVM_OPTIONAL_SOURCES
+  # so that FileCheck doesn't complain about unsued files.
+  LIST(APPEND LLVM_OPTIONAL_SOURCES ${ANALYSIS_TEST_SOURCES})
+
+  unset(LLVM_LINK_COMPONENTS)
+  add_llvm_library(InlineAdvisorPlugin MODULE BUILDTREE_ONLY InlineAdvisorPlugin.cpp)
+  # Put PLUGIN next to the unit test executable.
+  set_output_directory(InlineAdvisorPlugin
+      BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_CFG_INTDIR}
+      LIBRARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_CFG_INTDIR}
+      )
+  set_target_properties(InlineAdvisorPlugin PROPERTIES FOLDER "Tests")
+
+  export_executable_symbols_for_plugins(AnalysisTests)
+  add_dependencies(AnalysisTests InlineAdvisorPlugin)
+  set_property(TARGET InlineAdvisorPlugin PROPERTY FOLDER "Tests/UnitTests/AnalysisTests")
+endif()

llvm/unittests/Analysis/InlineAdvisorPlugin.cpp

@@ -0,0 +1,53 @@
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Pass.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Passes/PassPlugin.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "llvm/Analysis/InlineAdvisor.h"
+
+using namespace llvm;
+
+namespace {
+
+InlineAdvisor *DefaultAdvisorFactory(Module &M, FunctionAnalysisManager &FAM,
+                                     InlineParams Params, InlineContext IC) {
+  return new DefaultInlineAdvisor(M, FAM, Params, IC);
+}
+
+struct DefaultDynamicAdvisor : PassInfoMixin<DefaultDynamicAdvisor> {
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
+    PluginInlineAdvisorAnalysis DA(DefaultAdvisorFactory);
+    MAM.registerPass([&] { return DA; });
+    return PreservedAnalyses::all();
+  }
+};
+
+} // namespace
+
+/* New PM Registration */
+llvm::PassPluginLibraryInfo getDefaultDynamicAdvisorPluginInfo() {
+  return {LLVM_PLUGIN_API_VERSION, "DynamicDefaultAdvisor", LLVM_VERSION_STRING,
+          [](PassBuilder &PB) {
+            PB.registerPipelineStartEPCallback(
+                [](ModulePassManager &MPM, OptimizationLevel Level) {
+                  MPM.addPass(DefaultDynamicAdvisor());
+                });
+            PB.registerPipelineParsingCallback(
+                [](StringRef Name, ModulePassManager &PM,
+                   ArrayRef<PassBuilder::PipelineElement> InnerPipeline) {
+                  if (Name == "dynamic-inline-advisor") {
+                    PM.addPass(DefaultDynamicAdvisor());
+                    return true;
+                  }
+                  return false;
+                });
+          }};
+}
+
+extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
+llvmGetPassPluginInfo() {
+  return getDefaultDynamicAdvisorPluginInfo();
+}