[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

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)

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

include/llvm/Analysis/InlineAdvisor.h

@@ -203,6 +203,11 @@ public:
   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.
 

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

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

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

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

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;

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;
+  });
+}

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));
+}

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']

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

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

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

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
 

test/lit.cfg.py

@@ -219,6 +219,9 @@ else:
 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:

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.