[InlineOrder] Plugin Inline Order

[IBricchi]

Summary

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.

[IBricchi]

This patch builds off https://reviews.llvm.org/D139644 to add plugin support for inline order. This should allow for full fine grained control of the inliner using plugins.

[IBricchi]

rebase

[IBricchi]

format

[IBricchi]

Implement fix for AIX cmake issue from https://reviews.llvm.org/D140559

[IBricchi]

ping

[kazu]

I think I have a high-level understanding, but I haven't understood the implementation yet.

llvm/include/llvm/Analysis/InlineOrder.h

40–43

Sorry, I am slow to understand this. Is the following correct understanding?

PluginInlineOrderAnalysis is always there, regardless of whether you have a plugin. If you actually have a plugin, we can use Factory in PluginInlineOrderAnalysis to construct a custom inline order.

Now, why are we using an analysis here? Is that because it's a good hook to use?

45–80

Could we document the plugin as inline comments in llvm/unittests/Analysis/InlineOrderPlugin/InlineOrderPlugin.cpp or something? I am afraid that the big code block comments might diverge from what the actual code should look like. If you are OK with this idea, you might want to put a pointer to the example file here.

83

ModuleAnalysisManager?

llvm/unittests/Analysis/InlineOrderPlugin/InlineOrderPlugin.cpp

61–62

I am not familiar with this. Would you please explain? The new pass would load the plugin or something?

llvm/unittests/Analysis/PluginInlineAdvisorAnalysisTest.cpp

90–91

Could you capitalize the first letter of the sentence and put a period at the end. Easier to parse that way.

llvm/unittests/Analysis/PluginInlineOrderAnalysisTest.cpp

7

How are we using this?

19

Could we drop static inside the anonymous namespace?

64–65

Could you capitalize the first letter of the sentence and put a period at the end. Easier to parse that way.

[IBricchi]

[IBricchi]

[IBricchi]

Thanks for the review! I updated the patch to reflect the comments.

llvm/include/llvm/Analysis/InlineOrder.h

40–43

An instance of PluginInlineOrderAnalysis only exists if it has been registered, in a normal compilation it won't exist. If it's registered then indeed it's Factory will be used to create a custom order
We use an analysis here because we need the PassBuilder's ability to dynamically load passes as plugins (else we'd need to implement our own custom mechanism).

llvm/unittests/Analysis/InlineOrderPlugin/InlineOrderPlugin.cpp

61–62

In the initial commit there was an unnecessary level of indirection, it has been fixed now. The plugin should just register a PluginInlineOrderAnalysis with the ModuleAnalysisManager.

[kazu]

Thank you for explaining how things work.

llvm/include/llvm/Analysis/InlineOrder.h

52

Could we make this variable private and provide accessors like so?

static bool isRegistered() { return HasBeenRegistered; }
static void unregister() { HasBeenRegistered = false; }

I am suggesting this because I don't want anyone to set HasBeenRegistered to true except in the constructor here.

llvm/unittests/Analysis/InlineOrderPlugin/InlineOrderPlugin.cpp

20–24

Could we say something like this?

DefaultInlineOrder = getDefaultInlineOrder(FAM, Params, MAM, M);

To do this, we would need to rename getInlineOrder to getDefaultInlineOrder, turning getInlineOrder into a mere selector:

std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>>
llvm::getDefaultInlineOrder(FunctionAnalysisManager &FAM, const InlineParams &Params,
                           ModuleAnalysisManager &MAM, Module &M) {
  // The original function body goes here.
}

std::unique_ptr<InlineOrder<std::pair<CallBase *, int>>>
llvm::getInlineOrder(FunctionAnalysisManager &FAM, const InlineParams &Params,
                     ModuleAnalysisManager &MAM, Module &M) {
  if (llvm::PluginInlineOrderAnalysis::HasBeenRegistered) {
    LLVM_DEBUG(dbgs() << "    Current used priority: plugin ---- \n");
    return MAM.getResult<PluginInlineOrderAnalysis>(M).Factory(FAM, Params, MAM,
                                                               M);
  }

  return getDefaultInlineOrder(FAM, Params, MAM, M);
}

28

May I suggest the following?

// Ignore callees named "foo".

45

May I suggest getNoFooInlineOrder or noFooInlineOrderFactory? At least, I would drop default here because it's not the default inline order.

llvm/unittests/Analysis/PluginInlineOrderAnalysisTest.cpp

239–240

nit: Could we say named instead of called to avoid confusion when discussing inlining?

[IBricchi]

Thanks for the suggestions

[kazu]

Thank you for updating the patch. LGTM.

[IBricchi]

Ignore tests on windows due to limited pluign suport

[IBricchi]

I don't have commit access, could you commit it for me? ibricchi <[e-mail]>

[Jake-Egan]

Hi, we're seeing a CMake error on AIX after this change. Can you take a look please?

https://lab.llvm.org/buildbot/#/builders/214/builds/6366/steps/4/logs/stdio

[IBricchi]

I think the issue was that:

export_executable_symbols_for_plugins(AnalysisTests)

Was being called twice, once in each of the plugin cmake files. I've moved it into the main Analysys CMakeFile and I think this should fix the issue we were seeing, but I don't have an AIX machine to test it on. Would you mind checking @Jake-Egan?

[Jake-Egan]

In D140637#4196543, @IBricchi wrote:

I think the issue was that:

export_executable_symbols_for_plugins(AnalysisTests)

Was being called twice, once in each of the plugin cmake files. I've moved it into the main Analysys CMakeFile and I think this should fix the issue we were seeing, but I don't have an AIX machine to test it on. Would you mind checking @Jake-Egan?

It resolves the CMake error - thanks for the fix!

[IBricchi]

Glad to hear that! Would you mind pushing the fix? I don't have commit access, thanks

[tmatheson]

This test fails when I build and run with the following:

cmake ../llvm -GNinja -DCMAKE_BUILD_TYPE=Release -DLLVM_EXPORT_SYMBOLS_FOR_PLUGINS=ON -DLLVM_ENABLE_PLUGINS=ON -DBUILD_SHARED_LIBS=OFF -DLLVM_ENABLE_PROJECTS=clang
ninja unittests/Analysis/AnalysisTests
./unittests/Analysis/AnalysisTests --gtest_filter=PluginInlineOrderTest.NoInlineFoo

I get the following failure:

$ ./unittests/Analysis/AnalysisTests --gtest_filter=PluginInlineOrderTest.NoInlineFoo
Note: Google Test filter = PluginInlineOrderTest.NoInlineFoo
[==========] Running 1 test from 1 test suite.
[----------] Global test environment set-up.
[----------] 1 test from PluginInlineOrderTest
[ RUN      ] PluginInlineOrderTest.NoInlineFoo
/ssd/upstream/llvm/unittests/Analysis/PluginInlineOrderAnalysisTest.cpp:48: Failure
Value of: !!Plugin
  Actual: false
Expected: true
Plugin path: /ssd/upstream/build2/unittests/Analysis/InlineOrderPlugin.so
[  FAILED  ] PluginInlineOrderTest.NoInlineFoo (4 ms)
[----------] 1 test from PluginInlineOrderTest (4 ms total)

[----------] Global test environment tear-down
[==========] 1 test from 1 test suite ran. (4 ms total)
[  PASSED  ] 0 tests.
[  FAILED  ] 1 test, listed below:
[  FAILED  ] PluginInlineOrderTest.NoInlineFoo

 1 FAILED TEST

The error returned by CI.setupPlugin() is InlineOrderPlugin.so: undefined symbol: _ZN4llvm13AllAnalysesOnINS_6ModuleEE6SetKeyE".

This symbol is llvm::AllAnalysesOn<llvm::Module>::SetKey defined in LLVMCore. It is present in the unit test executable but is a local symbol:

$ readelf -s -W ./unittests/Analysis/AnalysisTests | rg _ZN4llvm13AllAnalysesOnINS_6ModuleEE6SetKeyE
 41264: 0000000001c79738     8 OBJECT  LOCAL  DEFAULT   29 _ZN4llvm13AllAnalysesOnINS_6ModuleEE6SetKeyE

Therefore it is not resolved by the dlopen.

[ttheodor]

This comment was removed by ttheodor.

llvm/include/llvm/Analysis/InlineOrder.h

public:
   virtual T pop() = 0;
 
   virtual void erase_if(function_ref<bool(T)> Pred) = 0;
 
   bool empty() { return !size(); }
 };
 
 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

static cl::opt<InlinePriorityMode> UseInlinePriority(
     "inline-priority-mode", cl::init(InlinePriorityMode::Size), cl::Hidden,
     cl::desc("Choose the priority mode to use in module inline"),
     cl::values(clEnumValN(InlinePriorityMode::Size, "size",
                           "Use callee size priority."),
                clEnumValN(InlinePriorityMode::Cost, "cost",
                           "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),
     cl::desc("The cost threshold for call sites that get inlined without the "
              "cost-benefit analysis"));
 
 namespace {
 
private:
   DenseMap<CallBase *, int> InlineHistoryMap;
   DenseMap<const CallBase *, PriorityT> Priorities;
   FunctionAnalysisManager &FAM;
   const InlineParams &Params;
 };
 
 } // 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");
     return std::make_unique<PriorityInlineOrder<SizePriority>>(FAM, Params);
 
   case InlinePriorityMode::Cost:
     LLVM_DEBUG(dbgs() << "    Current used priority: Cost priority ---- \n");
     return std::make_unique<PriorityInlineOrder<CostPriority>>(FAM, 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);
+}
 No newline at end of file

llvm/lib/Transforms/IPO/ModuleInliner.cpp

PreservedAnalyses ModuleInlinerPass::run(Module &M,
   // the entire Module. With this module inliner, the inline order is not
   // limited to bottom-up order. More globally scope inline order is enabled.
   // Also, the inline deferral logic become unnecessary in this module inliner.
   // It is possible to use other priority heuristics, e.g. profile-based
   // heuristic.
   //
   // 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.
   for (Function &F : M) {
     auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
     // We want to generally process call sites top-down in order for
     // simplifications stemming from replacing the call with the returned value
     // after inlining to be visible to subsequent inlining decisions.

llvm/unittests/Analysis/CMakeLists.txt

set(ANALYSIS_TEST_SOURCES
   LoopInfoTest.cpp
   LoopNestTest.cpp
   MemoryBuiltinsTest.cpp
   MemoryProfileInfoTest.cpp
   MemorySSATest.cpp
   MLModelRunnerTest.cpp
   PhiValuesTest.cpp
   PluginInlineAdvisorAnalysisTest.cpp
+  PluginInlineOrderAnalysisTest.cpp
   ProfileSummaryInfoTest.cpp
   ScalarEvolutionTest.cpp
   VectorFunctionABITest.cpp
   SparsePropagation.cpp
   TargetLibraryInfoTest.cpp
   TensorSpecTest.cpp
   TBAATest.cpp
   UnrollAnalyzerTest.cpp

 target_link_libraries(AnalysisTests PRIVATE LLVMTestingSupport)
 
 # 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()
 
+# 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

 # 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)
   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)
   # 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)
   add_dependencies(AnalysisTests InlineAdvisorPlugin)
   set_property(TARGET InlineAdvisorPlugin PROPERTY FOLDER "Tests/UnitTests/AnalysisTests")
 endif()

llvm/unittests/Analysis/InlineOrderPlugin/CMakeLists.txt

-# The advisor plugin expects to not link against the Analysis, Support and Core 
+# 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(InlineAdvisorPlugin MODULE BUILDTREE_ONLY
-    InlineAdvisorPlugin.cpp
-    )
+  add_llvm_library(InlineOrderPlugin MODULE BUILDTREE_ONLY InlineOrderPlugin.cpp)
   # Put PLUGIN next to the unit test executable.
-  set_output_directory(InlineAdvisorPlugin
+  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(InlineAdvisorPlugin PROPERTIES FOLDER "Tests")
+  set_target_properties(InlineOrderPlugin 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)
-  add_dependencies(AnalysisTests InlineAdvisorPlugin)
-  set_property(TARGET InlineAdvisorPlugin PROPERTY FOLDER "Tests/UnitTests/AnalysisTests")
+  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

+#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

 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/Config/config.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
 
 namespace llvm {
 
+namespace {
+
 void anchor() {}
 
 static std::string libPath(const std::string Name = "InlineAdvisorPlugin") {
   const auto &Argvs = testing::internal::GetArgvs();
   const char *Argv0 =
       Argvs.size() > 0 ? Argvs[0].c_str() : "PluginInlineAdvisorAnalysisTest";
   void *Ptr = (void *)(intptr_t)anchor;
   std::string Path = sys::fs::getMainExecutable(Argv0, Ptr);
CompilerInstance() {
     PB.registerCGSCCAnalyses(CGAM);
     PB.registerFunctionAnalyses(FAM);
     PB.registerLoopAnalyses(LAM);
     PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
     MPM.addPass(ModuleInlinerPass(IP, InliningAdvisorMode::Default,
                                   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;
 
   // run with the default inliner
   auto run_default(StringRef IR) {
     PluginInlineAdvisorAnalysis::HasBeenRegistered = false;
     outputM = parseAssemblyString(IR, Error, Ctx);
     MPM.run(*outputM, MAM);
if_end:
     store i32 %add, i32* %i
     br label %loop_cond
   loop_end:
     %correct_val = load i32, i32* %correct
     ret i32 %correct_val
 }
   )"};
 
+} // namespace
+
 // check that loading a plugin works
 // the plugin being loaded acts identically to the default inliner
 TEST(PluginInlineAdvisorTest, PluginLoad) {
 #if !defined(LLVM_ENABLE_PLUGINS)
   // Skip the test if plugins are disabled.
   GTEST_SKIP();
 #endif
   CompilerInstance CI{};

llvm/unittests/Analysis/PluginInlineOrderAnalysisTest.cpp

 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/Config/config.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
 
+#include "llvm/Analysis/InlineOrder.h"
+
 namespace llvm {
 
+namespace {
+
 void anchor() {}
 
-static std::string libPath(const std::string Name = "InlineAdvisorPlugin") {
+std::string libPath(const std::string Name = "InlineOrderPlugin") {
   const auto &Argvs = testing::internal::GetArgvs();
   const char *Argv0 =
-      Argvs.size() > 0 ? Argvs[0].c_str() : "PluginInlineAdvisorAnalysisTest";
+      Argvs.size() > 0 ? Argvs[0].c_str() : "PluginInlineOrderAnalysisTest";
   void *Ptr = (void *)(intptr_t)anchor;
   std::string Path = sys::fs::getMainExecutable(Argv0, Ptr);
   llvm::SmallString<256> Buf{sys::path::parent_path(Path)};
   sys::path::append(Buf, (Name + LLVM_PLUGIN_EXT).c_str());
   return std::string(Buf.str());
 }
 
-// Example of a custom InlineAdvisor that only inlines calls to functions called
-// "foo".
-class FooOnlyInlineAdvisor : public InlineAdvisor {
-public:
-  FooOnlyInlineAdvisor(Module &M, FunctionAnalysisManager &FAM,
-                       InlineParams Params, InlineContext IC)
-      : InlineAdvisor(M, FAM, IC) {}
-
-  std::unique_ptr<InlineAdvice> getAdviceImpl(CallBase &CB) override {
-    if (CB.getCalledFunction()->getName() == "foo")
-      return std::make_unique<InlineAdvice>(this, CB, getCallerORE(CB), true);
-    return std::make_unique<InlineAdvice>(this, CB, getCallerORE(CB), false);
-  }
-};
-
-static InlineAdvisor *fooOnlyFactory(Module &M, FunctionAnalysisManager &FAM,
-                                     InlineParams Params, InlineContext IC) {
-  return new FooOnlyInlineAdvisor(M, FAM, Params, IC);
-}
-
 struct CompilerInstance {
   LLVMContext Ctx;
   ModulePassManager MPM;
   InlineParams IP;
 
   PassBuilder PB;
   LoopAnalysisManager LAM;
   FunctionAnalysisManager FAM;
   CGSCCAnalysisManager CGAM;
   ModuleAnalysisManager MAM;
 
   SMDiagnostic Error;
 
-  // connect the plugin to our compiler instance
+  // Connect the plugin to our compiler instance.
   void setupPlugin() {
     auto PluginPath = libPath();
     ASSERT_NE("", PluginPath);
     Expected<PassPlugin> Plugin = PassPlugin::Load(PluginPath);
     ASSERT_TRUE(!!Plugin) << "Plugin path: " << PluginPath;
     Plugin->registerPassBuilderCallbacks(PB);
-    ASSERT_THAT_ERROR(PB.parsePassPipeline(MPM, "dynamic-inline-advisor"),
-                      Succeeded());
-  }
-
-  // connect the FooOnlyInlineAdvisor to our compiler instance
-  void setupFooOnly() {
-    MAM.registerPass(
-        [&] { return PluginInlineAdvisorAnalysis(fooOnlyFactory); });
   }
 
   CompilerInstance() {
     IP = getInlineParams(3, 0);
     PB.registerModuleAnalyses(MAM);
     PB.registerCGSCCAnalyses(CGAM);
     PB.registerFunctionAnalyses(FAM);
     PB.registerLoopAnalyses(LAM);
     PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
     MPM.addPass(ModuleInlinerPass(IP, InliningAdvisorMode::Default,
                                   ThinOrFullLTOPhase::None));
   }
 
-  std::string output;
-  std::unique_ptr<Module> outputM;
-
-  // run with the default inliner
-  auto run_default(StringRef IR) {
-    PluginInlineAdvisorAnalysis::HasBeenRegistered = false;
-    outputM = parseAssemblyString(IR, Error, Ctx);
-    MPM.run(*outputM, MAM);
-    ASSERT_TRUE(outputM);
-    output.clear();
-    raw_string_ostream o_stream{output};
-    outputM->print(o_stream, nullptr);
-    ASSERT_TRUE(true);
-  }
-
-  // run with the dnamic inliner
-  auto run_dynamic(StringRef IR) {
-    // note typically the constructor for the DynamicInlineAdvisorAnalysis
-    // will automatically set this to true, we controll it here only to
-    // altenate between the default and dynamic inliner in our test
-    PluginInlineAdvisorAnalysis::HasBeenRegistered = true;
-    outputM = parseAssemblyString(IR, Error, Ctx);
-    MPM.run(*outputM, MAM);
-    ASSERT_TRUE(outputM);
-    output.clear();
-    raw_string_ostream o_stream{output};
-    outputM->print(o_stream, nullptr);
+  ~CompilerInstance() {
+    // Reset the static variable that tracks if the plugin has been registered.
+    // This is needed to allow the test to run multiple times.
+    PluginInlineOrderAnalysis::unregister();
+  }
+
+  std::string Output;
+  std::unique_ptr<Module> OutputM;
+
+  // Run with the dynamic inline order.
+  auto run(StringRef IR) {
+    OutputM = parseAssemblyString(IR, Error, Ctx);
+    MPM.run(*OutputM, MAM);
+    ASSERT_TRUE(OutputM);
+    Output.clear();
+    raw_string_ostream OStream{Output};
+    OutputM->print(OStream, nullptr);
     ASSERT_TRUE(true);
   }
 };
 
 StringRef TestIRS[] = {
-    // Simple 3 function inline case
+    // Simple 3 function inline case.
     R"(
 define void @f1() {
   call void @foo()
   ret void
 }
 define void @foo() {
   call void @f3()
   ret void
 }
 define void @f3() {
   ret void
 }
   )",
-    // Test that has 5 functions of which 2 are recursive
+    // Test that has 5 functions of which 2 are recursive.
     R"(
 define void @f1() {
   call void @foo()
   ret void
 }
 define void @f2() {
   call void @foo()
   ret void
 }
 define void @foo() {
   call void @f4()
   call void @f5()
   ret void
 }
 define void @f4() {
   ret void
 }
 define void @f5() {
   call void @foo()
   ret void
 }
   )",
-    // test with 2 mutually recursive functions and 1 function with a loop
+    // Test with 2 mutually recursive functions and 1 function with a loop.
     R"(
 define void @f1() {
   call void @f2()
   ret void
 }
 define void @f2() {
-  call void @f3()
+  call void @foo()
   ret void
 }
-define void @f3() {
+define void @foo() {
   call void @f1()
   ret void
 }
 define void @f4() {
   br label %loop
 loop:
   call void @f5()
   br label %loop
 }
 define void @f5() {
   ret void
 }
   )",
-    // test that has a function that computes fibonacci in a loop, one in a
-    // recurisve manner, and one that calls both and compares them
+    // Test that has a function that computes fibonacci in a loop, one in a
+    // recursive manner, and one that calls both and compares them.
     R"(
 define i32 @fib_loop(i32 %n){
     %curr = alloca i32
     %last = alloca i32
     %i = alloca i32
     store i32 1, i32* %curr
     store i32 1, i32* %last
     store i32 2, i32* %i
loop_body:
     %add2 = add i32 %i_val2, 1
     store i32 %add2, i32* %i
     br label %loop_cond
   loop_end:
     %curr_val3 = load i32, i32* %curr
     ret i32 %curr_val3
 }
 
-define i32 @fib_rec(i32 %n){
+define i32 @foo(i32 %n){
     %cmp = icmp eq i32 %n, 0
     %cmp2 = icmp eq i32 %n, 1
     %or = or i1 %cmp, %cmp2
     br i1 %or, label %if_true, label %if_false
   if_true:
     ret i32 1
   if_false:
     %sub = sub i32 %n, 1
-    %call = call i32 @fib_rec(i32 %sub)
+    %call = call i32 @foo(i32 %sub)
     %sub2 = sub i32 %n, 2
-    %call2 = call i32 @fib_rec(i32 %sub2)
+    %call2 = call i32 @foo(i32 %sub2)
     %add = add i32 %call, %call2
     ret i32 %add
 }
 
 define i32 @fib_check(){
     %correct = alloca i32
     %i = alloca i32
     store i32 1, i32* %correct
     store i32 0, i32* %i
     br label %loop_cond
   loop_cond:
     %i_val = load i32, i32* %i
     %cmp = icmp slt i32 %i_val, 10
     br i1 %cmp, label %loop_body, label %loop_end
   loop_body:
     %i_val2 = load i32, i32* %i
     %call = call i32 @fib_loop(i32 %i_val2)
     %i_val3 = load i32, i32* %i
-    %call2 = call i32 @fib_rec(i32 %i_val3)
+    %call2 = call i32 @foo(i32 %i_val3)
     %cmp2 = icmp ne i32 %call, %call2
     br i1 %cmp2, label %if_true, label %if_false
   if_true:
     store i32 0, i32* %correct
     br label %if_end
   if_false:
     br label %if_end
   if_end:
     %i_val4 = load i32, i32* %i
     %add = add i32 %i_val4, 1
     store i32 %add, i32* %i
     br label %loop_cond
   loop_end:
     %correct_val = load i32, i32* %correct
     ret i32 %correct_val
 }
   )"};
 
-// check that loading a plugin works
-// the plugin being loaded acts identically to the default inliner
-TEST(PluginInlineAdvisorTest, PluginLoad) {
+} // namespace
+
+// Check that the behaviour of a custom inline order is correct.
+// The custom order drops any functions named "foo" so all tests
+// should contain at least one function named foo.
+TEST(PluginInlineOrderTest, NoInlineFoo) {
 #if !defined(LLVM_ENABLE_PLUGINS)
   // Skip the test if plugins are disabled.
   GTEST_SKIP();
 #endif
   CompilerInstance CI{};
   CI.setupPlugin();
 
   for (StringRef IR : TestIRS) {
-    CI.run_default(IR);
-    std::string default_output = CI.output;
-    CI.run_dynamic(IR);
-    std::string dynamic_output = CI.output;
-    ASSERT_EQ(default_output, dynamic_output);
-  }
-}
-
-// check that the behaviour of a custom inliner is correct
-// the custom inliner inlines all functions that are not named "foo"
-// this testdoes not require plugins to be enabled
-TEST(PluginInlineAdvisorTest, CustomAdvisor) {
-  CompilerInstance CI{};
-  CI.setupFooOnly();
-
-  for (StringRef IR : TestIRS) {
-    CI.run_dynamic(IR);
-    CallGraph CGraph = CallGraph(*CI.outputM);
-    for (auto &node : CGraph) {
-      for (auto &edge : *node.second) {
-        if (!edge.first)
-          continue;
-        ASSERT_NE(edge.second->getFunction()->getName(), "foo");
+    bool FoundFoo = false;
+    CI.run(IR);
+    CallGraph CGraph = CallGraph(*CI.OutputM);
+    for (auto &Node : CGraph) {
+      for (auto &Edge : *Node.second) {
+        FoundFoo |= Edge.second->getFunction()->getName() == "foo";
       }
     }
+    ASSERT_TRUE(FoundFoo);
   }
 }
 
 } // namespace llvm