[InlineOrder] Plugin Inline Order

Adds the ability to load a plugin to control the inline order.
This allows developing and distributing inlining heuristics
outside of tree. And together with the inline advisor plugins
allows for fine grained control of the inliner.

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

Reviewed By: kazu

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

llvm/include/llvm/Analysis/InlineOrder.h

@@ -32,7 +32,52 @@ template <typename T> class InlineOrder {
 };
 
 std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>>
-getInlineOrder(FunctionAnalysisManager &FAM, const InlineParams &Params);
+getDefaultInlineOrder(FunctionAnalysisManager &FAM, const InlineParams &Params,
+                      ModuleAnalysisManager &MAM, Module &M);
+
+std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>>
+getInlineOrder(FunctionAnalysisManager &FAM, const InlineParams &Params,
+               ModuleAnalysisManager &MAM, Module &M);
+
+/// Used for dynamically loading instances of InlineOrder as plugins
+///
+/// Plugins must implement an InlineOrderFactory, for an example refer to:
+/// llvm/unittests/Analysis/InlineOrderPlugin/InlineOrderPlugin.cpp
+///
+/// If a PluginInlineOrderAnalysis has been registered with the
+/// current ModuleAnalysisManager, llvm::getInlineOrder returns an
+/// InlineOrder created by the PluginInlineOrderAnalysis' Factory.
+///
+class PluginInlineOrderAnalysis
+    : public AnalysisInfoMixin<PluginInlineOrderAnalysis> {
+public:
+  static AnalysisKey Key;
+
+  typedef std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>> (
+      *InlineOrderFactory)(FunctionAnalysisManager &FAM,
+                           const InlineParams &Params,
+                           ModuleAnalysisManager &MAM, Module &M);
+
+  PluginInlineOrderAnalysis(InlineOrderFactory Factory) : Factory(Factory) {
+    HasBeenRegistered = true;
+    assert(Factory != nullptr &&
+           "The plugin inline order factory should not be a null pointer.");
+  }
+
+  struct Result {
+    InlineOrderFactory Factory;
+  };
+
+  Result run(Module &, ModuleAnalysisManager &) { return {Factory}; }
+  Result getResult() { return {Factory}; }
+
+  static bool isRegistered() { return HasBeenRegistered; }
+  static void unregister() { HasBeenRegistered = false; }
+
+private:
+  static bool HasBeenRegistered;
+  InlineOrderFactory Factory;
+};
 
 } // namespace llvm
 #endif // LLVM_ANALYSIS_INLINEORDER_H

llvm/lib/Analysis/InlineOrder.cpp

@@ -33,8 +33,7 @@ static cl::opt<InlinePriorityMode> UseInlinePriority(
                           "Use inline cost priority."),
                clEnumValN(InlinePriorityMode::CostBenefit, "cost-benefit",
                           "Use cost-benefit ratio."),
-               clEnumValN(InlinePriorityMode::ML, "ml",
-                          "Use ML.")));
+               clEnumValN(InlinePriorityMode::ML, "ml", "Use ML.")));
 
 static cl::opt<int> ModuleInlinerTopPriorityThreshold(
     "moudle-inliner-top-priority-threshold", cl::Hidden, cl::init(0),
@@ -281,8 +280,13 @@ class PriorityInlineOrder : public InlineOrder<std::pair<CallBase *, int>> {
 
 } // namespace
 
+AnalysisKey llvm::PluginInlineOrderAnalysis::Key;
+bool llvm::PluginInlineOrderAnalysis::HasBeenRegistered;
+
 std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>>
-llvm::getInlineOrder(FunctionAnalysisManager &FAM, const InlineParams &Params) {
+llvm::getDefaultInlineOrder(FunctionAnalysisManager &FAM,
+                            const InlineParams &Params,
+                            ModuleAnalysisManager &MAM, Module &M) {
   switch (UseInlinePriority) {
   case InlinePriorityMode::Size:
     LLVM_DEBUG(dbgs() << "    Current used priority: Size priority ---- \n");
@@ -295,11 +299,22 @@ llvm::getInlineOrder(FunctionAnalysisManager &FAM, const InlineParams &Params) {
   case InlinePriorityMode::CostBenefit:
     LLVM_DEBUG(
         dbgs() << "    Current used priority: cost-benefit priority ---- \n");
-    return std::make_unique<PriorityInlineOrder<CostBenefitPriority>>(FAM, Params);
+    return std::make_unique<PriorityInlineOrder<CostBenefitPriority>>(FAM,
+                                                                      Params);
   case InlinePriorityMode::ML:
-    LLVM_DEBUG(
-        dbgs() << "    Current used priority: ML priority ---- \n");
+    LLVM_DEBUG(dbgs() << "    Current used priority: ML priority ---- \n");
     return std::make_unique<PriorityInlineOrder<MLPriority>>(FAM, Params);
   }
   return nullptr;
 }
+
+std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>>
+llvm::getInlineOrder(FunctionAnalysisManager &FAM, const InlineParams &Params,
+                     ModuleAnalysisManager &MAM, Module &M) {
+  if (llvm::PluginInlineOrderAnalysis::isRegistered()) {
+    LLVM_DEBUG(dbgs() << "    Current used priority: plugin ---- \n");
+    return MAM.getResult<PluginInlineOrderAnalysis>(M).Factory(FAM, Params, MAM,
+                                                               M);
+  }
+  return getDefaultInlineOrder(FAM, Params, MAM, M);
+}

llvm/lib/Transforms/IPO/ModuleInliner.cpp

@@ -138,7 +138,7 @@ PreservedAnalyses ModuleInlinerPass::run(Module &M,
   //
   // TODO: Here is a huge amount duplicate code between the module inliner and
   // the SCC inliner, which need some refactoring.
-  auto Calls = getInlineOrder(FAM, Params);
+  auto Calls = getInlineOrder(FAM, Params, MAM, M);
   assert(Calls != nullptr && "Expected an initialized InlineOrder");
 
   // Populate the initial list of calls in this module.

llvm/unittests/Analysis/CMakeLists.txt

@@ -39,6 +39,7 @@ set(ANALYSIS_TEST_SOURCES
   MLModelRunnerTest.cpp
   PhiValuesTest.cpp
   PluginInlineAdvisorAnalysisTest.cpp
+  PluginInlineOrderAnalysisTest.cpp
   ProfileSummaryInfoTest.cpp
   ScalarEvolutionTest.cpp
   VectorFunctionABITest.cpp
@@ -72,4 +73,10 @@ if(CMAKE_SYSTEM_NAME STREQUAL "AIX")
   set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -Wl,-brtl")
 endif()
 
+# Export symbols from the plugins shared objects.
+if(NOT WIN32)
+  export_executable_symbols_for_plugins(AnalysisTests)
+endif()
+
 add_subdirectory(InlineAdvisorPlugin)
+add_subdirectory(InlineOrderPlugin)

llvm/unittests/Analysis/InlineAdvisorPlugin/CMakeLists.txt

@@ -14,7 +14,6 @@ if (NOT WIN32)
       )
   set_target_properties(InlineAdvisorPlugin PROPERTIES FOLDER "Tests")
 
-  export_executable_symbols_for_plugins(AnalysisTests)
   # The plugin depends on some of the output files of intrinsics_gen, so make sure
   # it is built before the plugin.
   add_dependencies(InlineAdvisorPlugin intrinsics_gen)

llvm/unittests/Analysis/InlineOrderPlugin/CMakeLists.txt

@@ -0,0 +1,20 @@
+# The order 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)
+  unset(LLVM_LINK_COMPONENTS)
+  add_llvm_library(InlineOrderPlugin MODULE BUILDTREE_ONLY InlineOrderPlugin.cpp)
+  # Put PLUGIN next to the unit test executable.
+  set_output_directory(InlineOrderPlugin
+      BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_CFG_INTDIR}/../
+      LIBRARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_CFG_INTDIR}/../
+      )
+  set_target_properties(InlineOrderPlugin PROPERTIES FOLDER "Tests")
+
+  # The plugin depends on some of the output files of intrinsics_gen, so make sure
+  # it is built before the plugin.
+  add_dependencies(InlineOrderPlugin intrinsics_gen)
+  add_dependencies(AnalysisTests InlineOrderPlugin)
+  set_property(TARGET InlineOrderPlugin PROPERTY FOLDER "Tests/UnitTests/AnalysisTests")
+endif()

llvm/unittests/Analysis/InlineOrderPlugin/InlineOrderPlugin.cpp

@@ -0,0 +1,70 @@
+#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/InlineOrder.h"
+
+using namespace llvm;
+
+namespace {
+
+class NoFooInlineOrder : public InlineOrder<std::pair<CallBase *, int>> {
+public:
+  NoFooInlineOrder(FunctionAnalysisManager &FAM, const InlineParams &Params,
+                   ModuleAnalysisManager &MAM, Module &M) {
+    DefaultInlineOrder = getDefaultInlineOrder(FAM, Params, MAM, M);
+  }
+  size_t size() override { return DefaultInlineOrder->size(); }
+  void push(const std::pair<CallBase *, int> &Elt) override {
+    // We ignore calles named "foo"
+    if (Elt.first->getCalledFunction()->getName() == "foo") {
+      DefaultInlineOrder->push(Elt);
+    }
+  }
+  std::pair<CallBase *, int> pop() override {
+    return DefaultInlineOrder->pop();
+  }
+  void erase_if(function_ref<bool(std::pair<CallBase *, int>)> Pred) override {
+    DefaultInlineOrder->erase_if(Pred);
+  }
+
+private:
+  std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>> DefaultInlineOrder;
+};
+
+std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>>
+NoFooInlineOrderFactory(FunctionAnalysisManager &FAM,
+                        const InlineParams &Params, ModuleAnalysisManager &MAM,
+                        Module &M) {
+  return std::make_unique<NoFooInlineOrder>(FAM, Params, MAM, M);
+}
+
+} // namespace
+
+/* New PM Registration */
+llvm::PassPluginLibraryInfo getDefaultDynamicInlineOrderPluginInfo() {
+  return {LLVM_PLUGIN_API_VERSION, "DynamicDefaultInlineOrder",
+          LLVM_VERSION_STRING, [](PassBuilder &PB) {
+            // We use the PassBuilder's callback mechanism
+            // to register our Analysis: this will register
+            // our PluginInlineOrderAnalysis instance with
+            // the ModuleAnalysisManager
+            PB.registerAnalysisRegistrationCallback(
+                [](ModuleAnalysisManager &MAM) {
+                  MAM.registerPass([] {
+                    // defaultInlineOrderFactory will be
+                    // used to create an InlineOrder
+                    return PluginInlineOrderAnalysis(NoFooInlineOrderFactory);
+                  });
+                });
+          }};
+}
+
+extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
+llvmGetPassPluginInfo() {
+  return getDefaultDynamicInlineOrderPluginInfo();
+}

llvm/unittests/Analysis/PluginInlineAdvisorAnalysisTest.cpp

@@ -10,6 +10,8 @@
 
 namespace llvm {
 
+namespace {
+
 void anchor() {}
 
 static std::string libPath(const std::string Name = "InlineAdvisorPlugin") {
@@ -84,6 +86,12 @@ struct CompilerInstance {
                                   ThinOrFullLTOPhase::None));
   }
 
+  ~CompilerInstance() {
+    // Reset the static variable that tracks if the plugin has been registered.
+    // This is needed to allow the test to run multiple times.
+    PluginInlineAdvisorAnalysis::HasBeenRegistered = false;
+  }
+
   std::string output;
   std::unique_ptr<Module> outputM;
 
@@ -256,6 +264,8 @@ define i32 @fib_check(){
 }
   )"};
 
+} // namespace
+
 // check that loading a plugin works
 // the plugin being loaded acts identically to the default inliner
 TEST(PluginInlineAdvisorTest, PluginLoad) {