[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
This commit is contained in:
ibricchi 2022-12-16 14:11:52 -08:00 committed by Mircea Trofin
parent 7d2c1150d3
commit a00aaf2b13
5 changed files with 484 additions and 11 deletions

View File

@ -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 @@ private:
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> {

View File

@ -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");

View File

@ -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()

View File

@ -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();
}

View File

@ -0,0 +1,298 @@
#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 {
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);
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
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);
ASSERT_TRUE(true);
}
};
StringRef TestIRS[] = {
// 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
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
R"(
define void @f1() {
call void @f2()
ret void
}
define void @f2() {
call void @f3()
ret void
}
define void @f3() {
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
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
br label %loop_cond
loop_cond:
%i_val = load i32, i32* %i
%cmp = icmp slt i32 %i_val, %n
br i1 %cmp, label %loop_body, label %loop_end
loop_body:
%curr_val = load i32, i32* %curr
%last_val = load i32, i32* %last
%add = add i32 %curr_val, %last_val
store i32 %add, i32* %last
store i32 %curr_val, i32* %curr
%i_val2 = load i32, i32* %i
%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){
%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)
%sub2 = sub i32 %n, 2
%call2 = call i32 @fib_rec(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)
%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) {
#if !defined(LLVM_ENABLE_PLUGINS)
// Disable the test if plugins are disabled.
return;
#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");
}
}
}
}
} // namespace llvm