Index: llvm/include/llvm/Analysis/InlineAdvisor.h =================================================================== --- llvm/include/llvm/Analysis/InlineAdvisor.h +++ llvm/include/llvm/Analysis/InlineAdvisor.h @@ -26,7 +26,7 @@ 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 @@ /// 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 @@ 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 { +/// 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 { +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 { Index: llvm/lib/Analysis/InlineAdvisor.cpp =================================================================== --- llvm/lib/Analysis/InlineAdvisor.cpp +++ llvm/lib/Analysis/InlineAdvisor.cpp @@ -196,11 +196,18 @@ } 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(M).getManager(); + if (PluginInlineAdvisorAnalysis::HasBeenRegistered) { + auto &DA = MAM.getResult(M); + Advisor.reset(DA.Factory(M, FAM, Params, IC)); + return !!Advisor; + } switch (Mode) { case InliningAdvisorMode::Default: LLVM_DEBUG(dbgs() << "Using default inliner heuristic.\n"); Index: llvm/unittests/Analysis/CMakeLists.txt =================================================================== --- llvm/unittests/Analysis/CMakeLists.txt +++ llvm/unittests/Analysis/CMakeLists.txt @@ -2,18 +2,12 @@ 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 @@ MemorySSATest.cpp MLModelRunnerTest.cpp PhiValuesTest.cpp + PluginInlineAdvisorAnalysisTest.cpp ProfileSummaryInfoTest.cpp ScalarEvolutionTest.cpp VectorFunctionABITest.cpp @@ -53,7 +48,58 @@ ValueLatticeTest.cpp ValueTrackingTest.cpp VectorUtilsTest.cpp + ) + +# 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 EXTRA_TESTS ${MLGO_TESTS}) +else() + LIST(APPEND LLVM_OPTIONAL_SOURCES ${MLGO_TESTS}) +endif() + +add_llvm_unittest_with_input_files(AnalysisTests + ${ANALYSIS_TEST_SOURCES} ${EXTRA_TESTS} ) - target_link_libraries(AnalysisTests PRIVATE LLVMTestingSupport) +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) + + # The plugin depends on some of the output files of llvm-tblgen, so make sure + # it is built before the plugin. + add_dependencies(InlineAdvisorPlugin llvm-tblgen) + + set_property(TARGET InlineAdvisorPlugin PROPERTY FOLDER "Tests/UnitTests/AnalysisTests") +endif() Index: llvm/unittests/Analysis/InlineAdvisorPlugin.cpp =================================================================== --- /dev/null +++ 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 { + 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 InnerPipeline) { + if (Name == "dynamic-inline-advisor") { + PM.addPass(DefaultDynamicAdvisor()); + return true; + } + return false; + }); + }}; +} + +extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo +llvmGetPassPluginInfo() { + return getDefaultDynamicAdvisorPluginInfo(); +} Index: llvm/unittests/Analysis/PluginInlineAdvisorAnalysisTest.cpp =================================================================== --- /dev/null +++ llvm/unittests/Analysis/PluginInlineAdvisorAnalysisTest.cpp @@ -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 getAdviceImpl(CallBase &CB) override { + if (CB.getCalledFunction()->getName() == "foo") + return std::make_unique(this, CB, getCallerORE(CB), true); + return std::make_unique(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 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 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