[llvm] Development-mode InlineAdvisor

[mtrofin]

Summary

This is the InlineAdvisor used in 'development' mode. It enables two
scenarios:

• loading models via a command-line parameter, thus allowing for rapid training iteration, where models can be used for the next exploration phase without requiring recompiling the compiler. This trades off some compilation speed for the added flexibility.
• collecting training logs, in the form of tensorflow.SequenceExample protobufs. We generate these as textual protobufs, which simplifies generation and testing. The protobufs may then be readily consumed by a tensorflow-based training algorithm.

To speed up training, training logs may also be collected from the
'default' training policy. In that case, this InlineAdvisor does not
use a model.

RFC: http://lists.llvm.org/pipermail/llvm-dev/2020-April/140763.html

[mtrofin]

Removed direct access to TF APIs, fully leveraging the TFUtils abstraction instead.

[davidxl]

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

31

TrainedModel sounds misleading as if it has been trained. Perhaps just TrainModel or ModelInTrain? Or TrainingModel?

132

May be documenting like this:

The advisor operates in two modes: 1) log collection, and 2) inference (during training). In the first mode, the default advisor is used while in 2), the MLInlineAdvisor is used.

133

Naming: MLInlineAdvisorForTraining or TrainingModeInlineAdvisor

257

SavedModel --> TrainedModel or ModelBeingTrained

259

Perhaps change DynamicModel name to be TrainModel to make the name consistent with the option.

357

can Inference and Logging both be on?

449

IsDoingInference

[mtrofin]

feedback

[mtrofin]

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

31

It's actually correct, it is trained: this would be a model obtained from a previous iteration of the reinforcement learning training loop. We aren't mutating it here.

132

logging may happen with either advisors, actually. To bootstrap, we use the default, then, to improve the model under training, we use the MLInlineAdvisor.

I expanded the comment to further clarify.

133

We use the term 'development' everywhere else, though - e.g. the command line flag is called 'development'.

257

SavedModel is a tensorflow term. Added a link in the comment.

259

I don't want to suggest it's being trained, though. How about LoadableModelRunner - the key aspect of this ModelRunner is that the model may be loaded from command line;

or, to contrast the release/AOT case ('ReleaseModeModelRunner) - DevelopmentModeModelRunner

wdyt?

357

Yes. I updated the comment earlier on, so now this should be more clear, I think

[davidxl]

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

31

As a reader, I feel the name means 'fully trained' instead of being 'partially trained' in the iteration loop.

259

How about IterativeModelRunner?

[mtrofin]

renames

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

31

renamed

259

ModelUnderTrainingRunner? similar to how I renamed the option.

[davidxl]

lgtm

llvm/include/llvm/Analysis/InlineAdvisor.h

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
 
+#ifdef LLVM_HAVE_TF_API
+std::unique_ptr<InlineAdvisor>
+getDevelopmentModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
+                          std::function<bool(CallBase &)> GetDefaultAdvice);
+#endif
+
 // Default (manual policy) decision making helper APIs. Shared with the legacy
 // pass manager inliner.
 
 /// Return the cost only if the inliner should attempt to inline at the given
 /// CallSite. If we return the cost, we will emit an optimisation remark later
 /// using that cost, so we won't do so from this function. Return None if
 /// inlining should not be attempted.
 Optional<InlineCost>

llvm/lib/Analysis/CMakeLists.txt

 set(CommonMLSources MLInlineAdvisor.cpp)
 set(ReleaseModeMLSources ReleaseModeModelRunner.cpp)
-set(DevelopmentModeMLSources TFUtils.cpp)
+set(DevelopmentModeMLSources
+  DevelopmentModeInlineAdvisor.cpp 
+  TFUtils.cpp
+  )
 
 if (DEFINED LLVM_HAVE_TF_AOT OR DEFINED LLVM_HAVE_TF_API)
   set(MLPolicySources ${CommonMLSources})
   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>

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

+//===- DevelopmentModeInlineAdvisor.cpp - runtime-loadable model runner  --===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a model runner using Tensorflow C APIs, allowing the
+// loading of a model from a command line option.
+//
+//===----------------------------------------------------------------------===//
+#include "llvm/Analysis/CallGraph.h"
+#include "llvm/Analysis/InlineSizeEstimatorAnalysis.h"
+#include "llvm/Analysis/MLInlineAdvisor.h"
+#include "llvm/Analysis/Utils/TFUtils.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/ManagedStatic.h"
+
+#include "tensorflow/c/c_api.h"
+
+#include <vector>
+
+using namespace llvm;
+
+static cl::opt<std::string> TrainingLog(
+    "training-log", cl::Hidden,
+    cl::desc("Path where the development - mode inlining log is saved."));
+
+static cl::opt<std::string> TFTrainedModelPath(
+    "ml-inliner-trained-model", cl::Hidden,
+    cl::desc(
+        "Path to saved model to use as policy during this training session."));
+
+static cl::opt<std::string> TFFeedPrefix("ml-inliner-trained-model-feed-prefix",
+                                         cl::Hidden, cl::init("action_"),
+                                         cl::desc("Prefix for feature names."));
+
+static cl::opt<std::string> TFDecisionName(
+    "ml-inliner-trained-model-decision-name", cl::Hidden,
+    cl::init("StatefulPartitionedCall"),
+    cl::desc("Name of the graph operation representing the decision."));
+
+namespace {
+/// An InlineEvent, used by TrainingLogger.
+struct InlineEvent {
+  /// What the default policy's decision would have been.
+  bool DefaultDecision = false;
+
+  /// What we advised. When training off the default policy, this is the same as
+  /// DefaultDecision.
+  bool AdvisedDecision = false;
+
+  /// What actually happened. This would be 'false' in the case of an inline
+  /// error, even if AdvisedDecision were true, otherwise it agrees with
+  /// AdvisedDecision.
+  bool Effect = false;
+
+  /// What the change in size was: size_after - size_before
+  int64_t Reward = 0;
+};
+
+/// Collect data we may use for training a model, and write it as a textual
+/// Tensorflow SequenceExample
+/// (https://www.tensorflow.org/api_docs/python/tf/train/SequenceExample)
+/// protobuf (https://developers.google.com/protocol-buffers).
+/// Because this is a protobuf, we cannot just stream the events as they come.
+/// Internally, TrainingLogger stores data in column-major format, because that
+/// lines up with how TF SequenceExample represents it.
+class TrainingLogger final {
+public:
+  TrainingLogger() {
+    for (size_t I = 0; I < NumberOfFeatures; ++I) {
+      Features.push_back(InlineFeatures());
+    }
+  }
+
+  /// Log one inlining event.
+  void logInlineEvent(const InlineEvent &Event,
+                      const MLModelRunner &ModelRunner) {
+    for (size_t I = 0; I < NumberOfFeatures; ++I) {
+      Features[I].push_back(ModelRunner.getFeature(I));
+    }
+    Decisions.push_back(Event.AdvisedDecision);
+    Effects.push_back(Event.Effect);
+    Rewards.push_back(Event.Reward);
+    DefaultDecisions.push_back(Event.DefaultDecision);
+  }
+
+  void printTensor(raw_fd_ostream &OutFile) {
+    if (DefaultDecisions.empty())
+      return;
+    OutFile << "feature_lists: {\n";
+
+    for (size_t I = 0; I < Features.size(); I++) {
+      writeTensor(OutFile, FeatureNameMap.at(I), Features[I]);
+    }
+    writeTensor(OutFile, DefaultDecisionName, DefaultDecisions);
+    writeTensor(OutFile, DecisionName, Decisions);
+    writeTensor(OutFile, RewardName, Rewards);
+
+    OutFile << "}\n";
+  }
+
+private:
+  template <typename T>
+  void writeTensor(raw_fd_ostream &OutFile, StringRef TensorName,
+                   const std::vector<T> &Tensor) {
+    OutFile << "  feature_list: {\n";
+    OutFile << "    key: "
+            << "\"" << TensorName << "\" ";
+    OutFile << "value: {\n";
+    for (const auto &Feature : Tensor) {
+      OutFile << "      feature: { int64_list: { value: [" << Feature
+              << "] } }\n";
+    }
+    OutFile << "    }\n";
+    OutFile << "  }\n";
+  }
+
+  std::vector<InlineFeatures> Features;
+  std::vector<bool> DefaultDecisions;
+  std::vector<bool> Decisions;
+  std::vector<bool> Effects;
+  std::vector<int64_t> Rewards;
+  std::vector<bool> Mandatory;
+};
+
+/// An extension of the MLInlineAdvisor which can also produce training logs,
+/// and can default to the default policy when we need to produce training logs
+/// from it.
+class DevelopmentModeMLInlineAdvisor : public MLInlineAdvisor {
+public:
+  DevelopmentModeMLInlineAdvisor(
+      Module &M, ModuleAnalysisManager &MAM,
+      std::unique_ptr<MLModelRunner> ModelRunner,
+      std::function<bool(CallBase &)> GetDefaultAdvice, bool IsDoingInference);
+
+  size_t getTotalSizeEstimate();
+
+  virtual ~DevelopmentModeMLInlineAdvisor();
+  void updateNativeSizeEstimate(int64_t Change) { CurrentNativeSize += Change; }
+  void resetNativeSize(Function *F) {
+    FAM.invalidate<InlineSizeEstimatorAnalysis>(*F);
+  }
+
+  std::unique_ptr<MLInlineAdvice>
+  getMandatoryAdvice(CallBase &CB, OptimizationRemarkEmitter &ORE) override;
+  std::unique_ptr<MLInlineAdvice>
+  getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE) override;
+
+  size_t getNativeSizeEstimate(const Function &F) const;
+
+private:
+  bool isLogging() const { return !TrainingLog.empty(); }
+
+  std::function<bool(CallBase &)> GetDefaultAdvice;
+  TrainingLogger Logger;
+  const bool IsDoingInference;
+
+  const int32_t InitialNativeSize;
+  int32_t CurrentNativeSize = 0;
+};
+
+/// A variant of MLInlineAdvice that tracks all non-trivial inlining
+/// decisions, for training/logging.
+class LoggingMLInlineAdvice : public MLInlineAdvice {
+public:
+  LoggingMLInlineAdvice(DevelopmentModeMLInlineAdvisor *Advisor, CallBase &CB,
+                        OptimizationRemarkEmitter &ORE, bool Recommendation,
+                        TrainingLogger &Logger, size_t CallerSizeEstimateBefore,
+                        size_t CalleeSizeEstimateBefore, bool DefaultDecision)
+      : MLInlineAdvice(Advisor, CB, ORE, Recommendation), Logger(Logger),
+        CallerSizeEstimateBefore(CallerSizeEstimateBefore),
+        CalleeSizeEstimateBefore(CalleeSizeEstimateBefore),
+        DefaultDecision(DefaultDecision) {}
+
+  virtual ~LoggingMLInlineAdvice() = default;
+
+private:
+  DevelopmentModeMLInlineAdvisor *getAdvisor() const {
+    return static_cast<DevelopmentModeMLInlineAdvisor *>(Advisor);
+  }
+  void recordInliningImpl() override {
+    MLInlineAdvice::recordInliningImpl();
+    getAdvisor()->resetNativeSize(Caller);
+    int Reward = std::numeric_limits<int>::max();
+    if (!getAdvisor()->isForcedToStop()) {
+      int NativeSizeAfter = getAdvisor()->getNativeSizeEstimate(*Caller) +
+                            CalleeSizeEstimateBefore;
+      Reward = NativeSizeAfter -
+               (CallerSizeEstimateBefore + CalleeSizeEstimateBefore);
+      getAdvisor()->updateNativeSizeEstimate(Reward);
+    }
+    log(Reward, /*Success=*/true);
+  }
+
+  void recordInliningWithCalleeDeletedImpl() override {
+    MLInlineAdvice::recordInliningWithCalleeDeletedImpl();
+    getAdvisor()->resetNativeSize(Caller);
+    if (!getAdvisor()->isForcedToStop()) {
+      int NativeSizeAfter = getAdvisor()->getNativeSizeEstimate(*Caller);
+      int Reward = NativeSizeAfter -
+                   (CallerSizeEstimateBefore + CalleeSizeEstimateBefore);
+      getAdvisor()->updateNativeSizeEstimate(Reward);
+      log(Reward, /*Success=*/true);
+    }
+  }
+
+  void recordUnsuccessfulInliningImpl(const InlineResult &Result) override {
+    MLInlineAdvice::recordUnsuccessfulInliningImpl(Result);
+    log(NoReward, /*Success=*/false);
+  }
+
+  void recordUnattemptedInliningImpl() override {
+    MLInlineAdvice::recordUnattemptedInliningImpl();
+    log(NoReward, /*Success=*/false);
+  }
+
+  void log(int64_t Reward, bool Success) {
+    InlineEvent Event;
+    Event.AdvisedDecision = isInliningRecommended();
+    Event.DefaultDecision = DefaultDecision;
+    Event.Effect = Success;
+    Event.Reward = Reward;
+    Logger.logInlineEvent(Event, getAdvisor()->getModelRunner());
+  }
+
+  static const int64_t NoReward = 0;
+  TrainingLogger &Logger;
+  const size_t CallerSizeEstimateBefore;
+  const size_t CalleeSizeEstimateBefore;
+  const bool DefaultDecision;
+};
+
+/// A pseudo model runner. We use it to store feature values when collecting
+/// logs for the default policy, but never ask it to 'run'.
+class NoInferenceModelRunner : public MLModelRunner {
+public:
+  NoInferenceModelRunner(LLVMContext &Ctx)
+      : MLModelRunner(Ctx), Features(NumberOfFeatures) {}
+  void setFeature(FeatureIndex Index, int64_t Value) override {
+    Features[static_cast<int>(Index)] = Value;
+  }
+
+  int64_t getFeature(int Index) const override { return Features[Index]; }
+  bool run() override {
+    llvm_unreachable("We shouldn't call run on this model runner.");
+  }
+
+private:
+  InlineFeatures Features;
+};
+
+/// DynamicModelRunner - training mode implementation. It uses TF C APIs to
+/// dynamically load and evaluate a TF SavedModel. Runtime performance is
+/// sacrificed for ease of use while training.
+class DynamicModelRunner final : public MLModelRunner {
+public:
+  DynamicModelRunner(LLVMContext &Ctx, const std::string &ModelPath);
+
+  bool run() override;
+
+  // Disallows copy and assign.
+  DynamicModelRunner(const DynamicModelRunner &) = delete;
+  DynamicModelRunner &operator=(const DynamicModelRunner &) = delete;
+
+  void setFeature(FeatureIndex Index, int64_t Value) override;
+  int64_t getFeature(int Index) const override;
+  bool isValid() const { return !!Evaluator; }
+
+private:
+  std::unique_ptr<TFModelEvaluator> Evaluator;
+
+  // The training framework needs some additional features, that just need to
+  // be set to 0.
+  struct TensorSpec {
+    std::string Name;
+    TF_DataType Type;
+  };
+
+  const std::vector<TensorSpec> TrainingOnlyFeatures{
+      {"inlining_default", TF_INT64},
+      {"discount", TF_FLOAT},
+      {"reward", TF_FLOAT},
+      {"step_type", TF_INT32}};
+};
+} // namespace
+
+DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor(
+    Module &M, ModuleAnalysisManager &MAM,
+    std::unique_ptr<MLModelRunner> ModelRunner,
+    std::function<bool(CallBase &)> GetDefaultAdvice, bool IsDoingInference)
+    : MLInlineAdvisor(M, MAM, std::move(ModelRunner)),
+      GetDefaultAdvice(GetDefaultAdvice), IsDoingInference(IsDoingInference),
+      InitialNativeSize(isLogging() ? getTotalSizeEstimate() : 0),
+      CurrentNativeSize(InitialNativeSize) {
+  // We cannot have the case of neither inference nor logging.
+  assert(IsDoingInference || isLogging());
+}
+
+DevelopmentModeMLInlineAdvisor::~DevelopmentModeMLInlineAdvisor() {
+  if (TrainingLog.empty())
+    return;
+  std::error_code ErrorCode;
+  raw_fd_ostream OutFile(TrainingLog, ErrorCode);
+  Logger.printTensor(OutFile);
+}
+
+size_t
+DevelopmentModeMLInlineAdvisor::getNativeSizeEstimate(const Function &F) const {
+  auto &R =
+      FAM.getResult<InlineSizeEstimatorAnalysis>(const_cast<Function &>(F));
+  if (!R) {
+    F.getParent()->getContext().emitError(
+        "Native size estimator is not present.");
+    return 0;
+  }
+  return *R;
+}
+
+std::unique_ptr<MLInlineAdvice>
+DevelopmentModeMLInlineAdvisor::getMandatoryAdvice(
+    CallBase &CB, OptimizationRemarkEmitter &ORE) {
+  if (!isLogging())
+    return MLInlineAdvisor::getMandatoryAdvice(CB, ORE);
+  return std::make_unique<LoggingMLInlineAdvice>(
+      /*Advisor=*/this,
+      /*CB=*/CB, /*ORE=*/ORE, /*Recommendation=*/true, /*Logger=*/Logger,
+      /*CallerSizeEstimateBefore=*/getNativeSizeEstimate(*CB.getCaller()),
+      /*CalleeSizeEstimateBefore=*/
+      getNativeSizeEstimate(*CB.getCalledFunction()),
+      /*DefaultDecision=*/true);
+}
+
+std::unique_ptr<MLInlineAdvice>
+DevelopmentModeMLInlineAdvisor::getAdviceFromModel(
+    CallBase &CB, OptimizationRemarkEmitter &ORE) {
+  if (IsDoingInference && !isLogging())
+    return MLInlineAdvisor::getAdviceFromModel(CB, ORE);
+
+  bool DefaultAdvice = GetDefaultAdvice(CB);
+  auto Recommendation = IsDoingInference ? ModelRunner->run() : DefaultAdvice;
+  return std::make_unique<LoggingMLInlineAdvice>(
+      /*Advisor=*/this,
+      /*CB=*/CB, /*ORE=*/ORE, /*Recommendation=*/Recommendation,
+      /*Logger=*/Logger,
+      /*CallerSizeEstimateBefore=*/getNativeSizeEstimate(*CB.getCaller()),
+      /*CalleeSizeEstimateBefore=*/
+      getNativeSizeEstimate(*CB.getCalledFunction()),
+      /*DefaultDecision=*/DefaultAdvice);
+}
+
+size_t DevelopmentModeMLInlineAdvisor::getTotalSizeEstimate() {
+  size_t Ret = 0;
+  for (auto &F : M) {
+    if (F.isDeclaration())
+      continue;
+    if (isFunctionDeleted(&F))
+      continue;
+    Ret += getNativeSizeEstimate(F);
+  }
+  return Ret;
+}
+
+DynamicModelRunner::DynamicModelRunner(LLVMContext &Ctx,
+                                       const std::string &ModelPath)
+    : MLModelRunner(Ctx) {
+  std::vector<std::string> InputNames;
+  std::vector<std::string> OutputNames;
+  for (size_t I = 0; I < NumberOfFeatures; ++I)
+    InputNames.push_back(TFFeedPrefix + FeatureNameMap[I]);
+  for (size_t I = 0; I < TrainingOnlyFeatures.size(); ++I)
+    InputNames.push_back(TFFeedPrefix + TrainingOnlyFeatures[I].Name);
+  OutputNames.push_back(TFDecisionName);
+
+  Evaluator =
+      std::make_unique<TFModelEvaluator>(ModelPath, InputNames, OutputNames);
+  if (!Evaluator || !Evaluator->isValid()) {
+    Ctx.emitError("Failed to create inliner saved model evaluator");
+    Evaluator.reset();
+    return;
+  }
+
+  static const std::vector<int64_t> Dim{1};
+
+  size_t InputIndex = 0;
+  for (; InputIndex < NumberOfFeatures; ++InputIndex) {
+    Evaluator->initInput(InputIndex, TF_INT64, Dim);
+  }
+
+  for (; InputIndex < Evaluator->getInput().size(); ++InputIndex) {
+    auto TFType = TrainingOnlyFeatures[InputIndex - NumberOfFeatures].Type;
+    Evaluator->initInput(InputIndex, TFType, Dim);
+  }
+}
+
+bool DynamicModelRunner::run() {
+  auto ER = Evaluator->evaluate();
+  if (!ER.hasValue()) {
+    Ctx.emitError("Error evaluating model.");
+    return false;
+  }
+  int64_t Decision = *ER->getTensorValue<int64_t>(0);
+  return static_cast<bool>(Decision);
+}
+
+int64_t DynamicModelRunner::getFeature(int Index) const {
+  return *static_cast<int64_t *>(TF_TensorData(Evaluator->getInput()[Index]));
+}
+
+void DynamicModelRunner::setFeature(FeatureIndex Index, int64_t Value) {
+  size_t NumericIndex = static_cast<size_t>(Index);
+  *(static_cast<int64_t *>(
+      TF_TensorData(Evaluator->getInput()[NumericIndex]))) = Value;
+}
+
+std::unique_ptr<InlineAdvisor> llvm::getDevelopmentModeAdvisor(
+    Module &M, ModuleAnalysisManager &MAM,
+    std::function<bool(CallBase &)> GetDefaultAdvice) {
+  auto &Ctx = M.getContext();
+  if (TrainingLog.empty() !=
+      !InlineSizeEstimatorAnalysis::isEvaluatorRequested()) {
+    Ctx.emitError("For development mode, if training logs are requested, then "
+                  "a size estimator must be available; either that, or neither "
+                  "are specified.");
+    return nullptr;
+  }
+
+  std::unique_ptr<MLModelRunner> Runner;
+
+  bool DoesInference = false;
+  if (TFTrainedModelPath.empty())
+    Runner.reset(new NoInferenceModelRunner(Ctx));
+  else {
+    Runner = std::make_unique<DynamicModelRunner>(Ctx, TFTrainedModelPath);
+    if (!Runner) {
+      Ctx.emitError("Could not load the policy model from the provided path");
+      return nullptr;
+    }
+    DoesInference = true;
+  }
+  return std::make_unique<DevelopmentModeMLInlineAdvisor>(
+      M, MAM, std::move(Runner), GetDefaultAdvice, DoesInference);
+}
 No newline at end of file

llvm/lib/Analysis/InlineAdvisor.cpp

 bool InlineAdvisorAnalysis::Result::tryCreate(InlineParams Params,
                                               InliningAdvisorMode Mode) {
   auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
   switch (Mode) {
   case InliningAdvisorMode::Default:
     Advisor.reset(new DefaultInlineAdvisor(FAM, Params));
     break;
   case InliningAdvisorMode::Development:
-    // To be added subsequently under conditional compilation.
+#ifdef LLVM_HAVE_TF_API
+    Advisor =
+        llvm::getDevelopmentModeAdvisor(M, MAM, [&FAM, Params](CallBase &CB) {
+          auto OIC = getDefaultInlineAdvice(CB, FAM, Params);
+          return OIC.hasValue();
+        });
+#endif
     break;
   case InliningAdvisorMode::Release:
 #ifdef LLVM_HAVE_TF_AOT
     Advisor = llvm::getReleaseModeAdvisor(M, MAM);
 #endif
     break;
   }
   return !!Advisor;

llvm/test/Bindings/Go/lit.local.cfg

 import os
 import pipes
 import shlex
 import sys
 
 if not 'go' in config.root.llvm_bindings:
     config.unsupported = True
 
 if not config.root.include_go_tests:
     config.unsupported = True
 
-if config.have_tf_aot:
+if config.have_tf_aot or config.have_tf_api:
     config.unsupported = True
 
 def find_executable(executable, path=None):
     if path is None:
         path = os.environ['PATH']
     paths = path.split(os.pathsep)
     base, ext = os.path.splitext(executable)
 

llvm/test/CMakeLists.txt

llvm_canonicalize_cmake_booleans(
   LLVM_LIBXML2_ENABLED
   LLVM_LINK_LLVM_DYLIB
   LLVM_TOOL_LTO_BUILD
   LLVM_USE_INTEL_JITEVENTS
   LLVM_BUILD_EXAMPLES
   LLVM_ENABLE_PLUGINS
   LLVM_BYE_LINK_INTO_TOOLS
   LLVM_HAVE_TF_AOT
+  LLVM_HAVE_TF_API
   )
 
 configure_lit_site_cfg(
   ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.py.in
   ${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg.py
   MAIN_CONFIG
   ${CMAKE_CURRENT_SOURCE_DIR}/lit.cfg.py
   PATHS

llvm/test/Transforms/Inline/ML/bounds-checks-rewards.ll

+; 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 penalize the 'bad' decision.
+; REQUIRES: have_tf_api
+; RUN: opt -passes=scc-oz-module-inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-trained-model=%S/../../../../lib/Analysis/models/inliner -training-log=- -enable-ml-inliner=development -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 -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-trained-model=%S/../../../../lib/Analysis/models/inliner -training-log=- -enable-ml-inliner=development -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: key: "delta_size" value: {
+; NOBOUNDS-NEXT: feature: { int64_list: { value: [10] } }
+; NOBOUNDS-NEXT: feature: { int64_list: { value: [6] } }
+; BOUNDS-NEXT: feature: { int64_list: { value: [2147483647] } }
+; CHECK-NEXT: }
+; 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

llvm/test/Transforms/Inline/ML/development-training-log.ll

+; Test that we can produce a log if we have or do not have a model, in development mode.
+; REQUIRES: have_tf_api
+; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-trained-model=%S/../../../../lib/Analysis/models/inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -S < %s | FileCheck %s 
+; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -S < %s | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-linux-gnu"
+
+declare i32 @f1(i32)
+declare i32 @f2(i32)
+
+define dso_local i32 @branches(i32) {
+  %cond = icmp slt i32 %0, 3
+  br i1 %cond, label %then, label %else
+
+then:
+  %ret.1 = call i32 @f1(i32 %0)
+  br label %last.block
+
+else:
+  %ret.2 = call i32 @f2(i32 %0)
+  br label %last.block
+
+last.block:
+  %ret = phi i32 [%ret.1, %then], [%ret.2, %else]
+  ret i32 %ret
+}
+
+define dso_local i32 @top() {
+  %1 = call i32 @branches(i32 2)
+  ret i32 %1
+}
+
+
+!llvm.module.flags = !{!0}
+!llvm.ident = !{!1}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{!"clang version 7.0.0-6 (tags/RELEASE_700/final)"}
+
+; Check we produce a protobuf that has inlining decisions and rewards.
+; CHECK:     feature_lists: {
+; CHECK:          key: "inlining_decision" value: {
+; CHECK-NEXT:       feature: { int64_list: { value: [1] } }
+; CHECK-NEXT:     }
+; CHECK-NEXT:   }
+; CHECK-NEXT:   feature_list: {
+; CHECK-NEXT:     key: "delta_size" value: {
+; CHECK-NEXT:       feature: { int64_list: { value: [0] } }
+; CHECK-NEXT:     }
+; CHECK-NEXT:   }
 No newline at end of file

llvm/test/Transforms/Inline/ML/ml-test-development-mode.ll

+; 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 shares it with a similar test
+; for the 'release' mode.
+;
+; REQUIRES: have_tf_api
+; RUN: opt -passes=scc-oz-module-inliner -enable-ml-inliner=default -S < %S/Inputs/test-module.ll 2>&1 | FileCheck %S/Inputs/test-module.ll --check-prefix=DEFAULT
+; RUN: opt -passes=scc-oz-module-inliner -enable-ml-inliner=development -ml-inliner-trained-model=%S/../../../../lib/Analysis/models/inliner -S < %S/Inputs/test-module.ll 2>&1 | FileCheck %S/Inputs/test-module.ll --check-prefix=CHECK

llvm/test/Transforms/Inline/inlining-advisor-default.ll

 ; Check that, in the absence of dependencies, we emit an error message when
 ; trying to use ML-driven inlining.
 ; REQUIRES: !have_tf_aot
+; REQUIRES: !have_tf_api
 ; 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
 
 declare i64 @f1()
 
 ; CHECK: Could not setup Inlining Advisor for the requested mode and/or options
 No newline at end of file

llvm/test/lit.cfg.py

 
 # Static libraries are not built if BUILD_SHARED_LIBS is ON.
 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")
 
+if config.have_tf_api:
+    config.available_features.add("have_tf_api")
+
 def have_cxx_shared_library():
     readobj_exe = lit.util.which('llvm-readobj', config.llvm_tools_dir)
     if not readobj_exe:
         print('llvm-readobj not found')
         return False
 
     try:
         readobj_cmd = subprocess.Popen(

llvm/test/lit.site.cfg.py.in

 config.llvm_libxml2_enabled = @LLVM_LIBXML2_ENABLED@
 config.llvm_host_triple = '@LLVM_HOST_TRIPLE@'
 config.host_arch = "@HOST_ARCH@"
 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@
+config.have_tf_api = @LLVM_HAVE_TF_API@
 
 # Support substitution of the tools_dir with user parameters. This is
 # used when we can't determine the tool dir at configuration time.
 try:
     config.llvm_tools_dir = config.llvm_tools_dir % lit_config.params
     config.llvm_shlib_dir = config.llvm_shlib_dir % lit_config.params
 except KeyError:
     e = sys.exc_info()[1]