[llvm] bdceefe - [llvm] Release-mode ML InlineAdvisor
Mircea Trofin via llvm-commits
llvm-commits at lists.llvm.org
Wed Jun 24 08:19:59 PDT 2020
Author: Mircea Trofin
Date: 2020-06-24T08:18:42-07:00
New Revision: bdceefe95ba6a3057947705ae7a48bfcbaed2f64
URL: https://github.com/llvm/llvm-project/commit/bdceefe95ba6a3057947705ae7a48bfcbaed2f64
DIFF: https://github.com/llvm/llvm-project/commit/bdceefe95ba6a3057947705ae7a48bfcbaed2f64.diff
LOG: [llvm] Release-mode ML InlineAdvisor
Summary:
This implementation uses a pre-trained model which is statically
compiled into a native function.
RFC: http://lists.llvm.org/pipermail/llvm-dev/2020-April/140763.html
Reviewers: davidxl, jdoerfert, dblaikie
Subscribers: mgorny, eraman, hiraditya, arphaman, llvm-commits
Tags: #llvm
Differential Revision: https://reviews.llvm.org/D81515
Added:
llvm/cmake/modules/TensorFlowCompile.cmake
llvm/include/llvm/Analysis/InlineModelFeatureMaps.h
llvm/include/llvm/Analysis/MLInlineAdvisor.h
llvm/include/llvm/Analysis/MLModelRunner.h
llvm/lib/Analysis/MLInlineAdvisor.cpp
llvm/lib/Analysis/ReleaseModeModelRunner.cpp
llvm/lib/Analysis/models/inliner/saved_model.pb
llvm/lib/Analysis/models/inliner/variables/variables.data-00000-of-00002
llvm/lib/Analysis/models/inliner/variables/variables.data-00001-of-00002
llvm/lib/Analysis/models/inliner/variables/variables.index
llvm/test/Transforms/Inline/ML/Inputs/test-module.ll
llvm/test/Transforms/Inline/ML/bounds-checks.ll
llvm/test/Transforms/Inline/ML/ml-test-release-mode.ll
Modified:
llvm/CMakeLists.txt
llvm/include/llvm/Analysis/InlineAdvisor.h
llvm/lib/Analysis/CMakeLists.txt
llvm/lib/Analysis/InlineAdvisor.cpp
llvm/test/Bindings/Go/lit.local.cfg
llvm/test/Transforms/Inline/inlining-advisor-default.ll
llvm/test/lit.cfg.py
llvm/test/lit.site.cfg.py.in
Removed:
################################################################################
diff --git a/llvm/CMakeLists.txt b/llvm/CMakeLists.txt
index 023733300090..3c53cb99512d 100644
--- a/llvm/CMakeLists.txt
+++ b/llvm/CMakeLists.txt
@@ -962,6 +962,25 @@ if( MINGW AND NOT "${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang" )
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
index 000000000000..a8ba56e67d1f
--- /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
diff erent 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
index 66a848e865e1..3480d93385a8 100644
--- a/llvm/include/llvm/Analysis/InlineAdvisor.h
+++ b/llvm/include/llvm/Analysis/InlineAdvisor.h
@@ -203,6 +203,11 @@ class InlineAdvisorAnalysis : public AnalysisInfoMixin<InlineAdvisorAnalysis> {
Result run(Module &M, ModuleAnalysisManager &MAM) { return Result(M, MAM); }
};
+#ifdef LLVM_HAVE_TF_AOT
+std::unique_ptr<InlineAdvisor>
+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
index 000000000000..8da442cc4a53
--- /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 <array>
+#include <string>
+#include <vector>
+
+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<size_t>(FeatureIndex::NumberOfFeatures);
+
+extern const std::array<std::string, NumberOfFeatures> FeatureNameMap;
+
+extern const char *const DecisionName;
+extern const char *const DefaultDecisionName;
+extern const char *const RewardName;
+
+using InlineFeatures = std::vector<int64_t>;
+
+} // 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
index 000000000000..cbe3b1f1f4e6
--- /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 <memory>
+#include <unordered_map>
+
+namespace llvm {
+class Module;
+class MLInlineAdvice;
+
+class MLInlineAdvisor : public InlineAdvisor {
+public:
+ MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM,
+ std::unique_ptr<MLModelRunner> ModelRunner);
+
+ CallGraph *callGraph() const { return CG.get(); }
+ virtual ~MLInlineAdvisor() = default;
+
+ void onPassEntry() override;
+
+ std::unique_ptr<InlineAdvice> 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<MLInlineAdvice>
+ getMandatoryAdvice(CallBase &CB, OptimizationRemarkEmitter &ORE);
+
+ virtual std::unique_ptr<MLInlineAdvice>
+ getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE);
+
+ Module &M;
+ std::unique_ptr<MLModelRunner> ModelRunner;
+
+private:
+ int64_t getModuleIRSize() const;
+
+ std::unique_ptr<CallGraph> CG;
+
+ int64_t NodeCount = 0;
+ int64_t EdgeCount = 0;
+ std::map<const Function *, unsigned> 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<MLInlineAdvisor *>(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
index 000000000000..7cfa6efedf10
--- /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
index dc0f8584ae23..eaf9670f00dd 100644
--- 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
+ $<TARGET_OBJECTS:tf_xla_runtime_objects>
+ ${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 @@ add_llvm_component_library(LLVMAnalysis
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
index 04e562eea8e5..9a3e5fa0df72 100644
--- a/llvm/lib/Analysis/InlineAdvisor.cpp
+++ b/llvm/lib/Analysis/InlineAdvisor.cpp
@@ -155,7 +155,9 @@ bool InlineAdvisorAnalysis::Result::tryCreate(InlineParams Params,
// 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
index 000000000000..45873f260f23
--- /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 <limits>
+#include <unordered_map>
+#include <unordered_set>
+
+#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<float> 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<std::string, NumberOfFeatures> 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<CallBase>(&I))
+ if (Function *Callee = CS->getCalledFunction()) {
+ if (!Callee->isDeclaration()) {
+ return CS;
+ }
+ }
+ return nullptr;
+}
+
+MLInlineAdvisor::MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM,
+ std::unique_ptr<MLModelRunner> Runner)
+ : InlineAdvisor(
+ MAM.getResult<FunctionAnalysisManagerModuleProxy>(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<CallGraphNode *> &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<InlineFeaturesAnalysis>(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<InlineFeaturesAnalysis>(*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<InlineFeaturesAnalysis>(*Caller)
+ .DirectCallsToDefinedFunctions;
+
+ if (CalleeWasDeleted)
+ --NodeCount;
+ else
+ NewCallerAndCalleeEdges += FAM.getResult<InlineFeaturesAnalysis>(*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<InlineAdvice> MLInlineAdvisor::getAdvice(CallBase &CB) {
+ auto &Caller = *CB.getCaller();
+ auto &Callee = *CB.getCalledFunction();
+
+ auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & {
+ return FAM.getResult<AssumptionAnalysis>(F);
+ };
+ auto GetTLI = [&](Function &F) -> const TargetLibraryInfo & {
+ return FAM.getResult<TargetLibraryAnalysis>(F);
+ };
+
+ auto &TIR = FAM.getResult<TargetIRAnalysis>(Callee);
+ auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(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<InlineAdvice>(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<InlineAdvice>(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<InlineAdvice>(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<Constant>(*I));
+ }
+
+ auto &CallerBefore = FAM.getResult<InlineFeaturesAnalysis>(Caller);
+ auto &CalleeBefore = FAM.getResult<InlineFeaturesAnalysis>(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<MLInlineAdvice>
+MLInlineAdvisor::getAdviceFromModel(CallBase &CB,
+ OptimizationRemarkEmitter &ORE) {
+ return std::make_unique<MLInlineAdvice>(this, CB, ORE, ModelRunner->run());
+}
+
+std::unique_ptr<MLInlineAdvice>
+MLInlineAdvisor::getMandatoryAdvice(CallBase &CB,
+ OptimizationRemarkEmitter &ORE) {
+ return std::make_unique<MLInlineAdvice>(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
index 000000000000..4c0ffbc17ff7
--- /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 <memory>
+#include <vector>
+
+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<int32_t> FeatureIndices;
+ int32_t ResultIndex = -1;
+ std::unique_ptr<llvm::InlinerSizeModel> CompiledModel;
+};
+} // namespace
+
+ReleaseModeModelRunner::ReleaseModeModelRunner(LLVMContext &Ctx)
+ : MLModelRunner(Ctx),
+ CompiledModel(std::make_unique<llvm::InlinerSizeModel>()) {
+ 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<int64_t *>(
+ CompiledModel->arg_data(FeatureIndices[Index]));
+}
+
+void ReleaseModeModelRunner::setFeature(FeatureIndex Index, int64_t Value) {
+ *static_cast<int64_t *>(CompiledModel->arg_data(
+ FeatureIndices[static_cast<size_t>(Index)])) = Value;
+}
+
+bool ReleaseModeModelRunner::run() {
+ CompiledModel->Run();
+ return static_cast<bool>(
+ *static_cast<int64_t *>(CompiledModel->result_data(ResultIndex)));
+}
+
+std::unique_ptr<InlineAdvisor>
+llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) {
+ auto AOTRunner = std::make_unique<ReleaseModeModelRunner>(M.getContext());
+ return std::make_unique<MLInlineAdvisor>(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 000000000000..5488989454f7
Binary files /dev/null and b/llvm/lib/Analysis/models/inliner/saved_model.pb
diff er
diff --git a/llvm/lib/Analysis/models/inliner/variables/variables.data-00000-of-00002 b/llvm/lib/Analysis/models/inliner/variables/variables.data-00000-of-00002
new file mode 100644
index 000000000000..58ebd0fc9871
Binary files /dev/null and b/llvm/lib/Analysis/models/inliner/variables/variables.data-00000-of-00002
diff er
diff --git a/llvm/lib/Analysis/models/inliner/variables/variables.data-00001-of-00002 b/llvm/lib/Analysis/models/inliner/variables/variables.data-00001-of-00002
new file mode 100644
index 000000000000..1f1f1b151a71
Binary files /dev/null and b/llvm/lib/Analysis/models/inliner/variables/variables.data-00001-of-00002
diff er
diff --git a/llvm/lib/Analysis/models/inliner/variables/variables.index b/llvm/lib/Analysis/models/inliner/variables/variables.index
new file mode 100644
index 000000000000..318d5a2443c2
Binary files /dev/null and b/llvm/lib/Analysis/models/inliner/variables/variables.index
diff er
diff --git a/llvm/test/Bindings/Go/lit.local.cfg b/llvm/test/Bindings/Go/lit.local.cfg
index 3021fc64a750..e32737079e4a 100644
--- a/llvm/test/Bindings/Go/lit.local.cfg
+++ b/llvm/test/Bindings/Go/lit.local.cfg
@@ -9,6 +9,9 @@ if not 'go' in config.root.llvm_bindings:
if not config.root.include_go_tests:
config.unsupported = True
+if config.have_tf_aot:
+ config.unsupported = True
+
def find_executable(executable, path=None):
if path is None:
path = os.environ['PATH']
diff --git a/llvm/test/Transforms/Inline/ML/Inputs/test-module.ll b/llvm/test/Transforms/Inline/ML/Inputs/test-module.ll
new file mode 100644
index 000000000000..b8279e5db6a0
--- /dev/null
+++ b/llvm/test/Transforms/Inline/ML/Inputs/test-module.ll
@@ -0,0 +1,64 @@
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-grtev4-linux-gnu"
+
+declare void @external_fct(i32)
+
+define dso_local i32 @top() {
+ %a = call i32 @multiplier(i32 5)
+ %b = call i32 @adder(i32 10)
+ %ret = add nsw i32 %a, %b
+ call void @external_fct(i32 %ret)
+ ret i32 %ret
+}
+
+define internal dso_local i32 @adder(i32) {
+ %2 = alloca i32, align 4
+ store i32 %0, i32* %2, align 4
+ %3 = load i32, i32* %2, align 4
+ %4 = call i32 @multiplier(i32 %3)
+ %5 = load i32, i32* %2, align 4
+ %6 = call i32 @switcher(i32 1)
+ %7 = add nsw i32 %4, %6
+ ret i32 %7
+}
+
+define internal i32 @multiplier(i32) {
+ %2 = alloca i32, align 4
+ store i32 %0, i32* %2, align 4
+ %3 = load i32, i32* %2, align 4
+ %4 = load i32, i32* %2, align 4
+ %5 = mul nsw i32 %3, %4
+ ret i32 %5
+}
+
+define i32 @switcher(i32) {
+ %2 = alloca i32, align 4
+ %3 = alloca i32, align 4
+ store i32 %0, i32* %3, align 4
+ %4 = load i32, i32* %3, align 4
+ switch i32 %4, label %11 [
+ i32 1, label %5
+ i32 2, label %6
+ ]
+
+; <label>:5: ; preds = %1
+ store i32 2, i32* %2, align 4
+ br label %12
+
+; <label>:6: ; preds = %1
+ %7 = load i32, i32* %3, align 4
+ %8 = load i32, i32* %3, align 4
+ %9 = call i32 @multiplier(i32 %8)
+ %10 = add nsw i32 %7, %9
+ store i32 %10, i32* %2, align 4
+ br label %12
+
+; <label>:11: ; preds = %1
+ %adder.result = call i32 @adder(i32 2)
+ store i32 %adder.result, i32* %2, align 4
+ br label %12
+
+; <label>:12: ; preds = %11, %6, %5
+ %13 = load i32, i32* %2, align 4
+ ret i32 %13
+}
\ No newline at end of file
diff --git a/llvm/test/Transforms/Inline/ML/bounds-checks.ll b/llvm/test/Transforms/Inline/ML/bounds-checks.ll
new file mode 100644
index 000000000000..b98e23acb427
--- /dev/null
+++ b/llvm/test/Transforms/Inline/ML/bounds-checks.ll
@@ -0,0 +1,41 @@
+; Test behavior when inlining policy grows size out of control.
+; In all cases, the end result is the same: mandatory inlinings must happen.
+; However, when we discover we 'trip' over the artificially-low size increase
+; factor, we don't inline anymore.
+; REQUIRES: have_tf_aot
+; RUN: opt -passes=scc-oz-module-inliner -enable-ml-inliner=release -ml-advisor-size-increase-threshold=10.0 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=NOBOUNDS
+; RUN: opt -passes=scc-oz-module-inliner -enable-ml-inliner=release -ml-advisor-size-increase-threshold=1.0 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=BOUNDS
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-grtev4-linux-gnu"
+
+declare i64 @f1()
+
+define i64 @f2() #0 {
+ %r = call i64 @f1()
+ %r2 = add i64 13, %r
+ ret i64 %r2
+}
+
+define i64 @some_function() {
+ %r = call i64 @f1()
+ %r2 = add i64 13, %r
+ ret i64 %r2
+}
+
+define i64 @top() {
+ %r = call i64 @f2()
+ %r2 = call i64 @some_function()
+ %r3 = add i64 %r, %r2
+ ret i64 %r3
+}
+
+attributes #0 = { alwaysinline }
+
+; CHECK-LABEL: @top
+; f2 must always be inlined, so we won't find a call to it in @top()
+; CHECK-NOT: call i64 @f2
+; @some-function isn't mandatory, and when we set the increase threshold too low,
+; it won't be inlined.
+; NOBOUNDS-NOT: @some_function
+; BOUNDS: call i64 @some_function
\ No newline at end of file
diff --git a/llvm/test/Transforms/Inline/ML/ml-test-release-mode.ll b/llvm/test/Transforms/Inline/ML/ml-test-release-mode.ll
new file mode 100644
index 000000000000..1ac0efd3cd0d
--- /dev/null
+++ b/llvm/test/Transforms/Inline/ML/ml-test-release-mode.ll
@@ -0,0 +1,14 @@
+; The default inliner doesn't elide @adder, it believes it's too costly to inline
+; adder into switcher. The ML inliner carries out that inlining, resulting in
+; a smaller result (part of it is that adder gets elided).
+;
+; This test uses Inputs/test-module.ll, as it will share it with a similar test
+; for the 'development' mode.
+;
+; REQUIRES: have_tf_aot
+; RUN: opt -passes=scc-oz-module-inliner -enable-ml-inliner=release -S < %S/Inputs/test-module.ll 2>&1 | FileCheck %s --check-prefix=CHECK
+; RUN: opt -passes=scc-oz-module-inliner -enable-ml-inliner=default -S < %S/Inputs/test-module.ll 2>&1 | FileCheck %s --check-prefix=DEFAULT
+
+; CHECK-NOT: @adder
+; DEFAULT-LABEL: @adder
+; DEFAULT-NEXT: %2 = mul
\ No newline at end of file
diff --git a/llvm/test/Transforms/Inline/inlining-advisor-default.ll b/llvm/test/Transforms/Inline/inlining-advisor-default.ll
index 808ddd97dc6b..7ee3e176bbea 100644
--- a/llvm/test/Transforms/Inline/inlining-advisor-default.ll
+++ b/llvm/test/Transforms/Inline/inlining-advisor-default.ll
@@ -1,6 +1,6 @@
; Check that, in the absence of dependencies, we emit an error message when
; trying to use ML-driven inlining.
-;
+; REQUIRES: !have_tf_aot
; RUN: not opt -passes=scc-oz-module-inliner -enable-ml-inliner=development -S < %s 2>&1 | FileCheck %s
; RUN: not opt -passes=scc-oz-module-inliner -enable-ml-inliner=release -S < %s 2>&1 | FileCheck %s
diff --git a/llvm/test/lit.cfg.py b/llvm/test/lit.cfg.py
index 49d345566b65..a3a97dd3b5c8 100644
--- a/llvm/test/lit.cfg.py
+++ b/llvm/test/lit.cfg.py
@@ -219,6 +219,9 @@ def get_asan_rtlib():
if not config.build_shared_libs and not config.link_llvm_dylib:
config.available_features.add('static-libs')
+if config.have_tf_aot:
+ config.available_features.add("have_tf_aot")
+
def have_cxx_shared_library():
readobj_exe = lit.util.which('llvm-readobj', config.llvm_tools_dir)
if not readobj_exe:
diff --git a/llvm/test/lit.site.cfg.py.in b/llvm/test/lit.site.cfg.py.in
index 3797830862f7..67c1302a7fb1 100644
--- a/llvm/test/lit.site.cfg.py.in
+++ b/llvm/test/lit.site.cfg.py.in
@@ -48,6 +48,7 @@ config.have_opt_viewer_modules = @LLVM_HAVE_OPT_VIEWER_MODULES@
config.libcxx_used = @LLVM_LIBCXX_USED@
config.has_plugins = @LLVM_ENABLE_PLUGINS@
config.linked_bye_extension = @LLVM_BYE_LINK_INTO_TOOLS@
+config.have_tf_aot = ("@LLVM_HAVE_TF_AOT@" == "ON")
# Support substitution of the tools_dir with user parameters. This is
# used when we can't determine the tool dir at configuration time.
More information about the llvm-commits
mailing list