diff --git a/llvm/CMakeLists.txt b/llvm/CMakeLists.txt --- a/llvm/CMakeLists.txt +++ b/llvm/CMakeLists.txt @@ -962,6 +962,25 @@ llvm_replace_compiler_option(CMAKE_CXX_FLAGS_RELEASE "-O3" "-O2") endif() +# For up-to-date instructions for installing the Tensorflow dependency, refer to +# the bot setup script: https://github.com/google/ml-compiler-opt/blob/master/buildbot/buildbot_init.sh +# Specifically, assuming python3 is installed: +# python3 -m pip install --upgrade pip && python3 -m pip install --user tf_nightly==2.3.0.dev20200528 +# Then set TENSORFLOW_AOT_PATH to the package install - usually it's ~/.local/lib/python3.7/site-packages/tensorflow +# +set(TENSORFLOW_AOT_PATH "" CACHE PATH "Path to TensorFlow pip install dir") + +if (NOT TENSORFLOW_AOT_PATH STREQUAL "") + set(LLVM_HAVE_TF_AOT "ON" CACHE BOOL "Tensorflow AOT available") + set(TENSORFLOW_AOT_COMPILER + "${TENSORFLOW_AOT_PATH}/../../../../bin/saved_model_cli" + CACHE PATH "Path to the Tensorflow AOT compiler") + add_definitions("-DLLVM_HAVE_TF_AOT") + include_directories(${TENSORFLOW_AOT_PATH}/include) + add_subdirectory(${TENSORFLOW_AOT_PATH}/xla_aot_runtime_src + ${CMAKE_ARCHIVE_OUTPUT_DIRECTORY}/tf_runtime) +endif() + # Put this before tblgen. Else we have a circular dependence. add_subdirectory(lib/Demangle) add_subdirectory(lib/Support) diff --git a/llvm/cmake/modules/TensorFlowCompile.cmake b/llvm/cmake/modules/TensorFlowCompile.cmake new file mode 100644 --- /dev/null +++ b/llvm/cmake/modules/TensorFlowCompile.cmake @@ -0,0 +1,38 @@ +# Run the tensorflow compiler (saved_model_cli) on the saved model in the +# ${model} directory, looking for the ${tag_set} tag set, and the SignatureDef +# ${signature_def_key}. +# Produce a pair of files called ${fname}.h and ${fname}.o in the +# ${CMAKE_CURRENT_BINARY_DIR}. The generated header will define a C++ class +# called ${cpp_class} - which may be a namespace-qualified class name. +function(tfcompile model tag_set signature_def_key fname cpp_class) + if (IS_ABSOLUTE ${model}) + set(LLVM_ML_MODELS_ABSOLUTE ${model}) + else() + set(LLVM_ML_MODELS_ABSOLUTE + ${CMAKE_CURRENT_SOURCE_DIR}/${model}) + endif() + + set(prefix ${CMAKE_CURRENT_BINARY_DIR}/${fname}) + set(obj_file ${prefix}.o) + set(hdr_file ${prefix}.h) + add_custom_command(OUTPUT ${obj_file} ${hdr_file} + COMMAND "XLA_FLAGS=\"--xla_cpu_multi_thread_eigen=false\"" ${TENSORFLOW_AOT_COMPILER} aot_compile_cpu + --dir ${LLVM_ML_MODELS_ABSOLUTE} + --tag_set ${tag_set} + --signature_def_key ${signature_def_key} + --output_prefix ${prefix} + --cpp_class ${cpp_class} + --target_triple ${LLVM_HOST_TRIPLE} + ) + + # Aggregate the objects so that results of different tfcompile calls may be + # grouped into one target. + set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file} PARENT_SCOPE) + set_source_files_properties(${obj_file} PROPERTIES + GENERATED 1 EXTERNAL_OBJECT 1) + + set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file} PARENT_SCOPE) + set_source_files_properties(${hdr_file} PROPERTIES + GENERATED 1) + +endfunction() diff --git a/llvm/include/llvm/Analysis/InlineAdvisor.h b/llvm/include/llvm/Analysis/InlineAdvisor.h --- a/llvm/include/llvm/Analysis/InlineAdvisor.h +++ b/llvm/include/llvm/Analysis/InlineAdvisor.h @@ -203,6 +203,11 @@ Result run(Module &M, ModuleAnalysisManager &MAM) { return Result(M, MAM); } }; +#ifdef LLVM_HAVE_TF_AOT +std::unique_ptr +getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM); +#endif + // Default (manual policy) decision making helper APIs. Shared with the legacy // pass manager inliner. diff --git a/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h b/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h new file mode 100644 --- /dev/null +++ b/llvm/include/llvm/Analysis/InlineModelFeatureMaps.h @@ -0,0 +1,70 @@ +//===- InlineModelFeatureMaps.h - common model runner defs ------*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// + +#ifndef LLVM_ANALYSIS_INLINEMODELFEATUREMAPS_H +#define LLVM_ANALYSIS_INLINEMODELFEATUREMAPS_H + +#include +#include +#include + +namespace llvm { + +// List of features. Each feature is defined through a triple: +// - the name of an enum member, which will be the feature index +// - a textual name, used for Tensorflow model binding (so it needs to match the +// names used by the Tensorflow model) +// - a documentation description. Currently, that is not used anywhere +// programmatically, and serves as workaround to inability of inserting comments +// in macros. +#define INLINE_FEATURE_ITERATOR(M) \ + M(CalleeBasicBlockCount, "callee_basic_block_count", \ + "number of basic blocks of the callee") \ + M(CallSiteHeight, "callsite_height", \ + "position of the call site in the original call graph - measured from " \ + "the farthest SCC") \ + M(NodeCount, "node_count", \ + "total current number of defined functions in the module") \ + M(NrCtantParams, "nr_ctant_params", \ + "number of parameters in the call site that are constants") \ + M(CostEstimate, "cost_estimate", "total cost estimate (threshold - free)") \ + M(EdgeCount, "edge_count", \ + "number of module-internal users of the caller, +1 if the caller is " \ + "exposed externally") \ + M(CallerUsers, "caller_users", \ + "number of blocks reached from a conditional instruction, in the caller") \ + M(CallerConditionallyExecutedBlocks, "caller_conditionally_executed_blocks", \ + "number of blocks reached from a conditional instruction, in the caller") \ + M(CallerBasicBlockCount, "caller_basic_block_count", \ + "number of basic blocks in the caller") \ + M(CalleeConditionallyExecutedBlocks, "callee_conditionally_executed_blocks", \ + "number of blocks reached from a conditional instruction, in the callee") \ + M(CalleeUsers, "callee_users", \ + "number of blocks reached from a conditional instruction, in the callee") + +enum class FeatureIndex : size_t { +#define POPULATE_INDICES(INDEX_NAME, NAME, COMMENT) INDEX_NAME, + INLINE_FEATURE_ITERATOR(POPULATE_INDICES) +#undef POPULATE_INDICES + NumberOfFeatures +}; + +constexpr size_t NumberOfFeatures = + static_cast(FeatureIndex::NumberOfFeatures); + +extern const std::array FeatureNameMap; + +extern const char *const DecisionName; +extern const char *const DefaultDecisionName; +extern const char *const RewardName; + +using InlineFeatures = std::vector; + +} // namespace llvm +#endif // LLVM_ANALYSIS_INLINEMODELFEATUREMAPS_H diff --git a/llvm/include/llvm/Analysis/MLInlineAdvisor.h b/llvm/include/llvm/Analysis/MLInlineAdvisor.h new file mode 100644 --- /dev/null +++ b/llvm/include/llvm/Analysis/MLInlineAdvisor.h @@ -0,0 +1,107 @@ +//===- MLInlineAdvisor.h - ML - based InlineAdvisor factories ---*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef LLVM_ANALYSIS_MLINLINEADVISOR_H +#define LLVM_ANALYSIS_MLINLINEADVISOR_H + +#include "llvm/Analysis/CallGraph.h" +#include "llvm/Analysis/InlineAdvisor.h" +#include "llvm/Analysis/MLModelRunner.h" +#include "llvm/IR/PassManager.h" + +#include +#include + +namespace llvm { +class Module; +class MLInlineAdvice; + +class MLInlineAdvisor : public InlineAdvisor { +public: + MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM, + std::unique_ptr ModelRunner); + + CallGraph *callGraph() const { return CG.get(); } + virtual ~MLInlineAdvisor() = default; + + void onPassEntry() override; + + std::unique_ptr getAdvice(CallBase &CB) override; + + int64_t getIRSize(const Function &F) const { return F.getInstructionCount(); } + void onSuccessfulInlining(const MLInlineAdvice &Advice, + bool CalleeWasDeleted); + + bool isForcedToStop() const { return ForceStop; } + int64_t getLocalCalls(Function &F); + const MLModelRunner &getModelRunner() const { return *ModelRunner.get(); } + +protected: + virtual std::unique_ptr + getMandatoryAdvice(CallBase &CB, OptimizationRemarkEmitter &ORE); + + virtual std::unique_ptr + getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE); + + Module &M; + std::unique_ptr ModelRunner; + +private: + int64_t getModuleIRSize() const; + + std::unique_ptr CG; + + int64_t NodeCount = 0; + int64_t EdgeCount = 0; + std::map FunctionLevels; + const int32_t InitialIRSize = 0; + int32_t CurrentIRSize = 0; + + bool ForceStop = false; +}; + +/// InlineAdvice that tracks changes post inlining. For that reason, it only +/// overrides the "successful inlining" extension points. +class MLInlineAdvice : public InlineAdvice { +public: + MLInlineAdvice(MLInlineAdvisor *Advisor, CallBase &CB, + OptimizationRemarkEmitter &ORE, bool Recommendation) + : InlineAdvice(Advisor, CB, ORE, Recommendation), + CallerIRSize(Advisor->isForcedToStop() ? 0 + : Advisor->getIRSize(*Caller)), + CalleeIRSize(Advisor->isForcedToStop() ? 0 + : Advisor->getIRSize(*Callee)), + CallerAndCalleeEdges(Advisor->isForcedToStop() + ? 0 + : (Advisor->getLocalCalls(*Caller) + + Advisor->getLocalCalls(*Callee))) {} + virtual ~MLInlineAdvice() = default; + + void recordInliningImpl() override; + void recordInliningWithCalleeDeletedImpl() override; + void recordUnsuccessfulInliningImpl(const InlineResult &Result) override; + void recordUnattemptedInliningImpl() override; + + Function *getCaller() const { return Caller; } + Function *getCallee() const { return Callee; } + + const int64_t CallerIRSize; + const int64_t CalleeIRSize; + const int64_t CallerAndCalleeEdges; + +private: + void reportContextForRemark(DiagnosticInfoOptimizationBase &OR); + + MLInlineAdvisor *getAdvisor() const { + return static_cast(Advisor); + }; +}; + +} // namespace llvm + +#endif // LLVM_ANALYSIS_MLINLINEADVISOR_H \ No newline at end of file diff --git a/llvm/include/llvm/Analysis/MLModelRunner.h b/llvm/include/llvm/Analysis/MLModelRunner.h new file mode 100644 --- /dev/null +++ b/llvm/include/llvm/Analysis/MLModelRunner.h @@ -0,0 +1,39 @@ +//===- MLModelRunner.h ---- ML model runner interface -----------*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// + +#ifndef LLVM_ANALYSIS_MLMODELRUNNER_H +#define LLVM_ANALYSIS_MLMODELRUNNER_H + +#include "llvm/Analysis/InlineModelFeatureMaps.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/PassManager.h" + +namespace llvm { + +/// MLModelRunner interface: abstraction of a mechanism for evaluating a +/// tensorflow "saved model". +class MLModelRunner { +public: + // Disallows copy and assign. + MLModelRunner(const MLModelRunner &) = delete; + MLModelRunner &operator=(const MLModelRunner &) = delete; + virtual ~MLModelRunner() = default; + + virtual bool run() = 0; + virtual void setFeature(FeatureIndex Index, int64_t Value) = 0; + virtual int64_t getFeature(int Index) const = 0; + +protected: + MLModelRunner(LLVMContext &Ctx) : Ctx(Ctx) {} + + LLVMContext &Ctx; +}; +} // namespace llvm + +#endif // LLVM_ANALYSIS_MLMODELRUNNER_H diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt --- a/llvm/lib/Analysis/CMakeLists.txt +++ b/llvm/lib/Analysis/CMakeLists.txt @@ -1,3 +1,18 @@ +set(CommonMLSources MLInlineAdvisor.cpp) +set(ReleaseModeMLSources ReleaseModeModelRunner.cpp) + +if (DEFINED LLVM_HAVE_TF_AOT) + include(TensorFlowCompile) + tfcompile(models/inliner serve action InlinerSizeModel llvm::InlinerSizeModel) + list(APPEND ReleaseModeMLSources + $ + ${GENERATED_OBJS} + ) + set(MLPolicySources ${CommonMLSources} ${ReleaseModeMLSources}) +else() + set(LLVM_OPTIONAL_SOURCES ${CommonMLSources} ${ReleaseModeMLSources}) +endif() + add_llvm_component_library(LLVMAnalysis AliasAnalysis.cpp AliasAnalysisEvaluator.cpp @@ -102,6 +117,7 @@ ValueTracking.cpp VectorUtils.cpp VFABIDemangling.cpp + ${MLPolicySources} ADDITIONAL_HEADER_DIRS ${LLVM_MAIN_INCLUDE_DIR}/llvm/Analysis diff --git a/llvm/lib/Analysis/InlineAdvisor.cpp b/llvm/lib/Analysis/InlineAdvisor.cpp --- a/llvm/lib/Analysis/InlineAdvisor.cpp +++ b/llvm/lib/Analysis/InlineAdvisor.cpp @@ -155,7 +155,9 @@ // To be added subsequently under conditional compilation. break; case InliningAdvisorMode::Release: - // To be added subsequently under conditional compilation. +#ifdef LLVM_HAVE_TF_AOT + Advisor = llvm::getReleaseModeAdvisor(M, MAM); +#endif break; } return !!Advisor; diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp new file mode 100644 --- /dev/null +++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp @@ -0,0 +1,301 @@ +//===- MLInlineAdvisor.cpp - machine learned InlineAdvisor ----------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// This file implements the interface between the inliner and a learned model. +// It delegates model evaluation to either the AOT compiled model (the +// 'release' mode) or a runtime-loaded model (the 'development' case). +// +//===----------------------------------------------------------------------===// +#include +#include +#include + +#include "llvm/ADT/SCCIterator.h" +#include "llvm/Analysis/CallGraph.h" +#include "llvm/Analysis/InlineCost.h" +#include "llvm/Analysis/InlineFeaturesAnalysis.h" +#include "llvm/Analysis/MLInlineAdvisor.h" +#include "llvm/Analysis/MLModelRunner.h" +#include "llvm/Analysis/OptimizationRemarkEmitter.h" +#include "llvm/Analysis/TargetLibraryInfo.h" +#include "llvm/Analysis/TargetTransformInfo.h" +#include "llvm/IR/InstIterator.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/PassManager.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/Path.h" + +using namespace llvm; + +#define DEBUG_TYPE "inline-ml" + +static cl::opt SizeIncreaseThreshold( + "ml-advisor-size-increase-threshold", cl::Hidden, + cl::desc("Maximum factor by which expected native size may increase before " + "blocking any further inlining."), + cl::init(2.0)); + +const std::array llvm::FeatureNameMap{ +#define POPULATE_NAMES(INDEX_NAME, NAME, COMMENT) NAME, + INLINE_FEATURE_ITERATOR(POPULATE_NAMES) +#undef POPULATE_NAMES +}; + +const char *const llvm::DecisionName = "inlining_decision"; +const char *const llvm::DefaultDecisionName = "inlining_default"; +const char *const llvm::RewardName = "delta_size"; + +CallBase *getInlinableCS(Instruction &I) { + if (auto *CS = dyn_cast(&I)) + if (Function *Callee = CS->getCalledFunction()) { + if (!Callee->isDeclaration()) { + return CS; + } + } + return nullptr; +} + +MLInlineAdvisor::MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM, + std::unique_ptr Runner) + : InlineAdvisor( + MAM.getResult(M).getManager()), + M(M), ModelRunner(std::move(Runner)), CG(new CallGraph(M)), + InitialIRSize(getModuleIRSize()), CurrentIRSize(InitialIRSize) { + assert(ModelRunner); + + // Extract the 'call site height' feature - the position of a call site + // relative to the farthest statically reachable SCC node. We don't mutate + // this value while inlining happens. Empirically, this feature proved + // critical in behavioral cloning - i.e. training a model to mimic the manual + // heuristic's decisions - and, thus, equally important for training for + // improvement. + for (auto I = scc_begin(CG.get()); !I.isAtEnd(); ++I) { + const std::vector &CGNodes = *I; + unsigned Level = 0; + for (auto *CGNode : CGNodes) { + Function *F = CGNode->getFunction(); + if (!F || F->isDeclaration()) + continue; + for (auto &I : instructions(F)) { + if (auto *CS = getInlinableCS(I)) { + auto *Called = CS->getCalledFunction(); + auto Pos = FunctionLevels.find(Called); + // In bottom up traversal, an inlinable callee is either in the + // same SCC, or to a function in a visited SCC. So not finding its + // level means we haven't visited it yet, meaning it's in this SCC. + if (Pos == FunctionLevels.end()) + continue; + Level = std::max(Level, Pos->second + 1); + } + } + } + for (auto *CGNode : CGNodes) { + Function *F = CGNode->getFunction(); + if (F && !F->isDeclaration()) + FunctionLevels[F] = Level; + } + } +} + +void MLInlineAdvisor::onPassEntry() { + // Function passes executed between InlinerPass runs may have changed the + // module-wide features. + NodeCount = 0; + EdgeCount = 0; + for (auto &F : M) + if (!F.isDeclaration()) { + ++NodeCount; + EdgeCount += getLocalCalls(F); + } +} + +int64_t MLInlineAdvisor::getLocalCalls(Function &F) { + return FAM.getResult(F).DirectCallsToDefinedFunctions; +} + +// Update the internal state of the advisor, and force invalidate feature +// analysis. Currently, we maintain minimal (and very simple) global state - the +// number of functions and the number of static calls. We also keep track of the +// total IR size in this module, to stop misbehaving policies at a certain bloat +// factor (SizeIncreaseThreshold) +void MLInlineAdvisor::onSuccessfulInlining(const MLInlineAdvice &Advice, + bool CalleeWasDeleted) { + assert(!ForceStop); + Function *Caller = Advice.getCaller(); + Function *Callee = Advice.getCallee(); + + // The caller features aren't valid anymore. + FAM.invalidate(*Caller); + int64_t IRSizeAfter = + getIRSize(*Caller) + (CalleeWasDeleted ? 0 : Advice.CalleeIRSize); + CurrentIRSize += IRSizeAfter - (Advice.CallerIRSize + Advice.CalleeIRSize); + if (CurrentIRSize > SizeIncreaseThreshold * InitialIRSize) + ForceStop = true; + + // We can delta-update module-wide features. We know the inlining only changed + // the caller, and maybe the callee (by deleting the latter). + // Nodes are simple to update. + // For edges, we 'forget' the edges that the caller and callee used to have + // before inlining, and add back what they currently have together. + int64_t NewCallerAndCalleeEdges = + FAM.getResult(*Caller) + .DirectCallsToDefinedFunctions; + + if (CalleeWasDeleted) + --NodeCount; + else + NewCallerAndCalleeEdges += FAM.getResult(*Callee) + .DirectCallsToDefinedFunctions; + EdgeCount += (NewCallerAndCalleeEdges - Advice.CallerAndCalleeEdges); + assert(CurrentIRSize >= 0 && EdgeCount >= 0 && NodeCount >= 0); +} + +int64_t MLInlineAdvisor::getModuleIRSize() const { + int64_t Ret = 0; + for (auto &F : CG->getModule()) + if (!F.isDeclaration()) + Ret += getIRSize(F); + return Ret; +} + +std::unique_ptr MLInlineAdvisor::getAdvice(CallBase &CB) { + auto &Caller = *CB.getCaller(); + auto &Callee = *CB.getCalledFunction(); + + auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & { + return FAM.getResult(F); + }; + auto GetTLI = [&](Function &F) -> const TargetLibraryInfo & { + return FAM.getResult(F); + }; + + auto &TIR = FAM.getResult(Callee); + auto &ORE = FAM.getResult(Caller); + + auto TrivialDecision = + llvm::getAttributeBasedInliningDecision(CB, &Callee, TIR, GetTLI); + + // If this is a "never inline" case, there won't be any changes to internal + // state we need to track, so we can just return the base InlineAdvice, which + // will do nothing interesting. + // Same thing if this is a recursive case. + if ((TrivialDecision.hasValue() && !TrivialDecision->isSuccess()) || + &Caller == &Callee) + return std::make_unique(this, CB, ORE, false); + + bool Mandatory = TrivialDecision.hasValue() && TrivialDecision->isSuccess(); + + // If we need to stop, we won't want to track anymore any state changes, so + // we just return the base InlineAdvice, which acts as a noop. + if (ForceStop) { + ORE.emit([&] { + return OptimizationRemarkMissed(DEBUG_TYPE, "ForceStop", &CB) + << "Won't attempt inlining because module size grew too much."; + }); + return std::make_unique(this, CB, ORE, Mandatory); + } + + int CostEstimate = 0; + if (!Mandatory) { + auto IsCallSiteInlinable = + llvm::getInliningCostEstimate(CB, TIR, GetAssumptionCache); + if (!IsCallSiteInlinable) { + // We can't inline this for correctness reasons, so return the base + // InlineAdvice, as we don't care about tracking any state changes (which + // won't happen). + return std::make_unique(this, CB, ORE, false); + } + CostEstimate = *IsCallSiteInlinable; + } + + if (Mandatory) + return getMandatoryAdvice(CB, ORE); + + auto NrCtantParams = 0; + for (auto I = CB.arg_begin(), E = CB.arg_end(); I != E; ++I) { + NrCtantParams += (isa(*I)); + } + + auto &CallerBefore = FAM.getResult(Caller); + auto &CalleeBefore = FAM.getResult(Callee); + + ModelRunner->setFeature(FeatureIndex::CalleeBasicBlockCount, + CalleeBefore.BasicBlockCount); + ModelRunner->setFeature(FeatureIndex::CallSiteHeight, + FunctionLevels[&Caller]); + ModelRunner->setFeature(FeatureIndex::NodeCount, NodeCount); + ModelRunner->setFeature(FeatureIndex::NrCtantParams, NrCtantParams); + ModelRunner->setFeature(FeatureIndex::CostEstimate, CostEstimate); + ModelRunner->setFeature(FeatureIndex::EdgeCount, EdgeCount); + ModelRunner->setFeature(FeatureIndex::CallerUsers, CallerBefore.Uses); + ModelRunner->setFeature(FeatureIndex::CallerConditionallyExecutedBlocks, + CallerBefore.BlocksReachedFromConditionalInstruction); + ModelRunner->setFeature(FeatureIndex::CallerBasicBlockCount, + CallerBefore.BasicBlockCount); + ModelRunner->setFeature(FeatureIndex::CalleeConditionallyExecutedBlocks, + CalleeBefore.BlocksReachedFromConditionalInstruction); + ModelRunner->setFeature(FeatureIndex::CalleeUsers, CalleeBefore.Uses); + return getAdviceFromModel(CB, ORE); +} + +std::unique_ptr +MLInlineAdvisor::getAdviceFromModel(CallBase &CB, + OptimizationRemarkEmitter &ORE) { + return std::make_unique(this, CB, ORE, ModelRunner->run()); +} + +std::unique_ptr +MLInlineAdvisor::getMandatoryAdvice(CallBase &CB, + OptimizationRemarkEmitter &ORE) { + return std::make_unique(this, CB, ORE, true); +} + +void MLInlineAdvice::reportContextForRemark( + DiagnosticInfoOptimizationBase &OR) { + using namespace ore; + OR << NV("Callee", Callee->getName()); + for (size_t I = 0; I < NumberOfFeatures; ++I) + OR << NV(FeatureNameMap[I], getAdvisor()->getModelRunner().getFeature(I)); + OR << NV("ShouldInline", isInliningRecommended()); +} + +void MLInlineAdvice::recordInliningImpl() { + ORE.emit([&]() { + OptimizationRemark R(DEBUG_TYPE, "InliningSuccess", DLoc, Block); + reportContextForRemark(R); + return R; + }); + getAdvisor()->onSuccessfulInlining(*this, /*CalleeWasDeleted*/ false); +} + +void MLInlineAdvice::recordInliningWithCalleeDeletedImpl() { + ORE.emit([&]() { + OptimizationRemark R(DEBUG_TYPE, "InliningSuccessWithCalleeDeleted", DLoc, + Block); + reportContextForRemark(R); + return R; + }); + getAdvisor()->onSuccessfulInlining(*this, /*CalleeWasDeleted*/ true); +} + +void MLInlineAdvice::recordUnsuccessfulInliningImpl( + const InlineResult &Result) { + ORE.emit([&]() { + OptimizationRemarkMissed R(DEBUG_TYPE, "InliningAttemptedAndUnsuccessful", + DLoc, Block); + reportContextForRemark(R); + return R; + }); +} +void MLInlineAdvice::recordUnattemptedInliningImpl() { + ORE.emit([&]() { + OptimizationRemarkMissed R(DEBUG_TYPE, "IniningNotAttempted", DLoc, Block); + reportContextForRemark(R); + return R; + }); +} \ No newline at end of file diff --git a/llvm/lib/Analysis/ReleaseModeModelRunner.cpp b/llvm/lib/Analysis/ReleaseModeModelRunner.cpp new file mode 100644 --- /dev/null +++ b/llvm/lib/Analysis/ReleaseModeModelRunner.cpp @@ -0,0 +1,87 @@ +//===- ReleaseModeModelRunner.cpp - Fast, precompiled model runner -------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// This file implements a model runner wrapping an AOT compiled ML model. +// Only inference is supported. +// +//===----------------------------------------------------------------------===// + +#include "llvm/Analysis/InlineModelFeatureMaps.h" +#include "llvm/Analysis/MLInlineAdvisor.h" + +// codegen-ed file +#include "InlinerSizeModel.h" // NOLINT + +#include +#include + +using namespace llvm; +namespace { + +static const char *const FeedPrefix = "feed_"; +static const char *const FetchPrefix = "fetch_"; + +/// MLModelRunner - production mode implementation. It uses a AOT-compiled +/// SavedModel for efficient execution. +class ReleaseModeModelRunner final : public MLModelRunner { +public: + ReleaseModeModelRunner(LLVMContext &Ctx); + virtual ~ReleaseModeModelRunner() = default; + + bool run() override; + + void setFeature(FeatureIndex Index, int64_t Value) override; + int64_t getFeature(int Index) const override; + +private: + std::vector FeatureIndices; + int32_t ResultIndex = -1; + std::unique_ptr CompiledModel; +}; +} // namespace + +ReleaseModeModelRunner::ReleaseModeModelRunner(LLVMContext &Ctx) + : MLModelRunner(Ctx), + CompiledModel(std::make_unique()) { + assert(CompiledModel && "The CompiledModel should be valid"); + + FeatureIndices.reserve(NumberOfFeatures); + + for (size_t I = 0; I < NumberOfFeatures; ++I) { + const int Index = + CompiledModel->LookupArgIndex(FeedPrefix + FeatureNameMap[I]); + assert(Index >= 0 && "Cannot find Feature in inlining model"); + FeatureIndices[I] = Index; + } + + ResultIndex = + CompiledModel->LookupResultIndex(std::string(FetchPrefix) + DecisionName); + assert(ResultIndex >= 0 && "Cannot find DecisionName in inlining model"); +} + +int64_t ReleaseModeModelRunner::getFeature(int Index) const { + return *static_cast( + CompiledModel->arg_data(FeatureIndices[Index])); +} + +void ReleaseModeModelRunner::setFeature(FeatureIndex Index, int64_t Value) { + *static_cast(CompiledModel->arg_data( + FeatureIndices[static_cast(Index)])) = Value; +} + +bool ReleaseModeModelRunner::run() { + CompiledModel->Run(); + return static_cast( + *static_cast(CompiledModel->result_data(ResultIndex))); +} + +std::unique_ptr +llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) { + auto AOTRunner = std::make_unique(M.getContext()); + return std::make_unique(M, MAM, std::move(AOTRunner)); +} diff --git a/llvm/lib/Analysis/models/inliner/saved_model.pb b/llvm/lib/Analysis/models/inliner/saved_model.pb new file mode 100644 index 0000000000000000000000000000000000000000..0000000000000000000000000000000000000000 GIT binary patch literal 0 Hc$@:5: ; preds = %1 + store i32 2, i32* %2, align 4 + br label %12 + +;