//===- 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/Config/config.h"
#if defined(LLVM_HAVE_TF_API)

#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 "llvm/Support/Path.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> TFModelUnderTrainingPath(
    "ml-inliner-model-under-training", cl::Hidden,
    cl::desc(R"(Path to SavedModel from the previous training iteration.
The directory is also expected to contain a JSON specification of the 
outputs expected to be logged, where the first entry must be the 
inlining decision. The file containing the specification should be 
called output_spec.json. The expected JSON value is an array of 
dictionaries. Each dictionary should have 2 keys: 

- "tensor_spec, followed by the TensorSpec description of the
output; and 
- "logging_name", a string indicating the name to use when
logging the output values. 

Example:
[
  {
    "logging_name" : "some_name", 
    "tensor_spec" : { 
      "name" : "model_name", 
      "port" : 0,
      "shape" : [2, 3],
      "type" : "float"
      }
  }
]

The first value must always correspond to the decision.)"));

static cl::opt<std::string> TFOutputSpecOverride(
    "ml-inliner-output-spec-override", cl::Hidden,
    cl::desc("Override the path to the output spec json file. See "
             "-ml-inliner-model-under-training documentation for the "
             "specification of that file."));

static cl::opt<std::string> TFFeedPrefix("ml-inliner-trained-model-feed-prefix",
                                         cl::Hidden, cl::init("action_"),
                                         cl::desc("Prefix for feature names."));

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 ModelUnderTrainingRunner;
class TrainingLogger final {
public:
  TrainingLogger(StringRef LogFileName, const ModelUnderTrainingRunner *MUTR);

  /// Log one inlining event.
  void logInlineEvent(const InlineEvent &Event,
                      const MLModelRunner &ModelRunner);

  /// Print the stored tensors.
  void print();

private:
  /// Write the values of one tensor as a list.
  template <typename T>
  void writeTensorValues(raw_fd_ostream &OutFile, const char *TensorData,
                         size_t ElemCount) const {
    OutFile << "[";
    const T *TypedData = reinterpret_cast<const T *>(TensorData);
    for (size_t I = 0; I < ElemCount; ++I) {
      if (I > 0)
        OutFile << ", ";
      OutFile << TypedData[I];
    }
    OutFile << "]";
  }

  /// Write a list of tensors as a sequence of TensorFlow FeatureList protobufs.
  /// The tensors are assumed to be stored contiguously, in row-major format,
  /// in the TensorData buffer. Each tensor has the shape given by Spec. The
  /// feature name in the output is either the provided LoggingName, if
  /// specified, otherwise it's the name of the tensor (as given by Spec).
  template <typename T>
  void
  writeTensorsAsFeatureLists(raw_fd_ostream &OutFile, const TensorSpec &Spec,
                             const T *TensorData, size_t TensorCount,
                             Optional<StringRef> LoggingName = None) const {
    writeRawTensorsAsFeatureLists(OutFile, Spec,
                                  reinterpret_cast<const char *>(TensorData),
                                  TensorCount, LoggingName);
  }

  /// Untyped implementation of the API above.
  void
  writeRawTensorsAsFeatureLists(raw_fd_ostream &OutFile, const TensorSpec &Spec,
                                const char *TensorData, size_t TensorCount,
                                Optional<StringRef> LoggingName = None) const {
    const char *FieldName = "<invalid>";
    std::function<void(const char *)> ValueWriter;
    // The 'Feature' protobuf only has 3 possible fields: float_list,
    // int64_list, or bytes_list, so we capture int32 values as int64. We don't
    // support any other types.
    if (Spec.isElementType<int64_t>()) {
      FieldName = "int64_list";
      ValueWriter = [&](const char *Data) {
        writeTensorValues<int64_t>(OutFile, Data, Spec.getElementCount());
      };
    } else if (Spec.isElementType<int32_t>()) {
      FieldName = "int64_list";
      ValueWriter = [&](const char *Data) {
        writeTensorValues<int32_t>(OutFile, Data, Spec.getElementCount());
      };

    } else if (Spec.isElementType<float>()) {
      FieldName = "float_list";
      ValueWriter = [&](const char *Data) {
        writeTensorValues<float>(OutFile, Data, Spec.getElementCount());
      };

    } else
      llvm_unreachable("Unsupported tensor type.");

    OutFile << "  feature_list: {\n";
    OutFile << "    key: "
            << "\"" << (LoggingName ? *LoggingName : Spec.name()) << "\" ";
    OutFile << "value: {\n";
    size_t TensorByteSize = Spec.getElementCount() * Spec.getElementByteSize();
    for (const char *P = TensorData,
                    *E = TensorData + TensorByteSize * TensorCount;
         P < E; P += TensorByteSize) {
      OutFile << "      feature: { " << FieldName << ": { value: ";
      ValueWriter(P);
      OutFile << " } }\n";
    }
    OutFile << "    }\n";
    OutFile << "  }\n";
  }

  StringRef LogFileName;
  const ModelUnderTrainingRunner *const MUTR;
  std::vector<InlineFeatures> Features;
  std::vector<int64_t> DefaultDecisions;
  // We store all outputs as data blobs, but we always expect to have one, the
  // first one, representing the decision. While we could track that separately,
  // for uniformity, we store it, generically, here.
  std::vector<std::vector<char>> Outputs;
  std::vector<bool> Effects;
  std::vector<int64_t> Rewards;
};

/// An extension of the MLInlineAdvisor for the 'development' mode, targeting
/// the offline training scenario. Note that training happens outside of the
/// compiler, this facility is concerned with producing training data ("logs").
/// This InlineAdvisor can operate in the following modes:
///
/// 1) collect logs for the default policy. This is useful for bootstrapping
/// training, which will be considerably faster by starting from a reasonable
/// policy.
///
/// 2) collect logs for the ML policy, using a model from a previous
/// training. Potentially, that model uses internally some small random
/// perturbation of its weights, to induce exploration (setting this up is the
/// responsibility of the training algorithm). The logs would then be used to
/// retrain and improve on this model.
///
/// 3) use the provided model, with no logging. This is useful for end to end
/// validation - the model, in this case, is a release candidate and shouldn't
/// have random perturbations. It is a convenience feature: rather than needing
/// to take the release candidate model and compile it in 'release' mode,
/// validate it, then potentially discard it, it's easier to just pass the model
/// to the compiler, albeit compilation would be slower, as a one-off. Once the
/// model behaves satisfactorily, it can be compiled AOT, for efficiency, in
/// release mode. The expectation is that a well-trained model provides a good
/// policy over a sufficiently diverse codebase, over many changes (i.e.
/// training happens seldom).
class DevelopmentModeMLInlineAdvisor : public MLInlineAdvisor {
public:
  DevelopmentModeMLInlineAdvisor(
      Module &M, ModuleAnalysisManager &MAM,
      std::unique_ptr<MLModelRunner> ModelRunner,
      std::function<bool(CallBase &)> GetDefaultAdvice, bool IsDoingInference,
      std::unique_ptr<TrainingLogger> Logger);

  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;

  Optional<size_t> getNativeSizeEstimate(const Function &F) const;

private:
  bool isLogging() const { return !!Logger; }

  std::function<bool(CallBase &)> GetDefaultAdvice;
  const bool IsDoingInference;
  std::unique_ptr<TrainingLogger> Logger;

  const Optional<int32_t> InitialNativeSize;
  Optional<int32_t> CurrentNativeSize;
};

/// 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,
                        Optional<size_t> CallerSizeEstimateBefore,
                        Optional<size_t> CalleeSizeEstimateBefore,
                        bool DefaultDecision, bool Mandatory = false)
      : MLInlineAdvice(Advisor, CB, ORE, Recommendation), Logger(Logger),
        CallerSizeEstimateBefore(CallerSizeEstimateBefore),
        CalleeSizeEstimateBefore(CalleeSizeEstimateBefore),
        DefaultDecision(DefaultDecision), Mandatory(Mandatory) {}

  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 (InlineSizeEstimatorAnalysis::isEvaluatorRequested() &&
        !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 (InlineSizeEstimatorAnalysis::isEvaluatorRequested() &&
        !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) {
    if (Mandatory)
      return;
    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 Optional<size_t> CallerSizeEstimateBefore;
  const Optional<size_t> CalleeSizeEstimateBefore;
  const bool DefaultDecision;
  const bool Mandatory;
};

/// 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;
};

/// ModelUnderTrainingRunner - training mode implementation. It uses TF C APIs
/// to dynamically load and evaluate a TF SavedModel
/// (https://www.tensorflow.org/guide/saved_model). Runtime performance is
/// sacrificed for ease of use while training.
class ModelUnderTrainingRunner final : public MLModelRunner {
public:
  ModelUnderTrainingRunner(LLVMContext &Ctx, const std::string &ModelPath);

  bool run() override;

  // Disallows copy and assign.
  ModelUnderTrainingRunner(const ModelUnderTrainingRunner &) = delete;
  ModelUnderTrainingRunner &
  operator=(const ModelUnderTrainingRunner &) = delete;

  void setFeature(FeatureIndex Index, int64_t Value) override;
  int64_t getFeature(int Index) const override;
  bool isValid() const { return !!Evaluator; }

  const std::vector<std::string> outputNames() const { return OutputNames; }

  const std::vector<TensorSpec> outputSpecs() const { return OutputSpecs; }

  const Optional<TFModelEvaluator::EvaluationResult> &
  lastEvaluationResult() const {
    return LastEvaluationResult;
  }

private:
  std::unique_ptr<TFModelEvaluator> Evaluator;
  std::vector<std::string> OutputNames;
  std::vector<TensorSpec> OutputSpecs;
  Optional<TFModelEvaluator::EvaluationResult> LastEvaluationResult;

  bool loadOutputSpecs(LLVMContext &Ctx, StringRef FileName);

  // The training framework needs some additional features.
  const std::vector<TensorSpec> TrainingOnlyFeatures{
      TensorSpec::createSpec<int64_t>(TFFeedPrefix + "inlining_default", {1}),
      TensorSpec::createSpec<float>(TFFeedPrefix + "discount", {1}),
      TensorSpec::createSpec<float>(TFFeedPrefix + "reward", {1}),
      TensorSpec::createSpec<int32_t>(TFFeedPrefix + "step_type", {1})};
};
} // namespace

TrainingLogger::TrainingLogger(StringRef LogFileName,
                               const ModelUnderTrainingRunner *MUTR)
    : LogFileName(LogFileName), MUTR(MUTR) {
  for (size_t I = 0; I < NumberOfFeatures; ++I)
    Features.push_back(InlineFeatures());

  // The first output is the inlining decision.
  auto OutputCount = MUTR ? MUTR->outputSpecs().size() : 1;
  Outputs.assign(OutputCount, std::vector<char>());
}

/// Log one inlining event.
void TrainingLogger::logInlineEvent(const InlineEvent &Event,
                                    const MLModelRunner &ModelRunner) {
  for (size_t I = 0; I < NumberOfFeatures; ++I)
    Features[I].push_back(ModelRunner.getFeature(I));

  Effects.push_back(Event.Effect);
  Rewards.push_back(Event.Reward);
  DefaultDecisions.push_back(Event.DefaultDecision);
  int64_t Advice = static_cast<int64_t>(Event.AdvisedDecision);
  const char *AdviceData = reinterpret_cast<const char *>(&Advice);
  Outputs[0].insert(Outputs[0].end(), AdviceData, AdviceData + sizeof(int64_t));
  for (size_t I = 1; I < Outputs.size(); ++I) {
    const auto &Result = *MUTR->lastEvaluationResult();
    auto &Spec = MUTR->outputSpecs()[I];
    const char *RawData =
        reinterpret_cast<const char *>(Result.getUntypedTensorValue(I));
    Outputs[I].insert(Outputs[I].end(), RawData,
                      RawData +
                          Spec.getElementCount() * Spec.getElementByteSize());
  }
}

void TrainingLogger::print() {
  std::error_code EC;
  raw_fd_ostream OutFile(LogFileName, EC);
  size_t NumberOfRecords = Rewards.size();
  if (NumberOfRecords == 0)
    return;

  OutFile << "feature_lists: {\n";
  for (size_t I = 0; I < Features.size(); ++I)
    writeTensorsAsFeatureLists(
        OutFile, TensorSpec::createSpec<int64_t>(FeatureNameMap.at(I), {1}),
        Features[I].data(), NumberOfRecords);

  writeTensorsAsFeatureLists(
      OutFile, TensorSpec::createSpec<int64_t>(DefaultDecisionName, {1}),
      DefaultDecisions.data(), NumberOfRecords);

  writeRawTensorsAsFeatureLists(
      OutFile, TensorSpec::createSpec<int64_t>(DecisionName, {1}),
      Outputs[0].data(), NumberOfRecords);

  if (InlineSizeEstimatorAnalysis::isEvaluatorRequested())
    writeTensorsAsFeatureLists(OutFile,
                               TensorSpec::createSpec<int64_t>(RewardName, {1}),
                               Rewards.data(), NumberOfRecords);

  for (size_t I = 1; I < Outputs.size(); ++I)
    writeRawTensorsAsFeatureLists(OutFile, MUTR->outputSpecs()[I],
                                  Outputs[I].data(), NumberOfRecords,
                                  StringRef(MUTR->outputNames()[I]));

  OutFile << "}\n";
}

DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor(
    Module &M, ModuleAnalysisManager &MAM,
    std::unique_ptr<MLModelRunner> ModelRunner,
    std::function<bool(CallBase &)> GetDefaultAdvice, bool IsDoingInference,
    std::unique_ptr<TrainingLogger> Logger)
    : MLInlineAdvisor(M, MAM, std::move(ModelRunner)),
      GetDefaultAdvice(GetDefaultAdvice), IsDoingInference(IsDoingInference),
      Logger(std::move(Logger)),
      InitialNativeSize(isLogging() ? getTotalSizeEstimate() : 0),
      CurrentNativeSize(InitialNativeSize) {
  // We cannot have the case of neither inference nor logging.
  assert(IsDoingInference || isLogging());
}

DevelopmentModeMLInlineAdvisor::~DevelopmentModeMLInlineAdvisor() {
  if (isLogging())
    Logger->print();
}

Optional<size_t>
DevelopmentModeMLInlineAdvisor::getNativeSizeEstimate(const Function &F) const {
  if (!InlineSizeEstimatorAnalysis::isEvaluatorRequested())
    return None;
  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, /*Mandatory*/ 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() {
  if (!InlineSizeEstimatorAnalysis::isEvaluatorRequested())
    return 0;
  size_t Ret = 0;
  for (auto &F : M) {
    if (F.isDeclaration())
      continue;
    if (isFunctionDeleted(&F))
      continue;
    Ret += *getNativeSizeEstimate(F);
  }
  return Ret;
}

ModelUnderTrainingRunner::ModelUnderTrainingRunner(LLVMContext &Ctx,
                                                   const std::string &ModelPath)
    : MLModelRunner(Ctx) {
  std::vector<TensorSpec> InputSpecs;
  for (size_t I = 0; I < NumberOfFeatures; ++I)
    InputSpecs.push_back(
        TensorSpec::createSpec<int64_t>(TFFeedPrefix + FeatureNameMap[I], {1}));
  InputSpecs.insert(InputSpecs.end(), TrainingOnlyFeatures.begin(),
                    TrainingOnlyFeatures.end());
  SmallVector<char, 128> OutputSpecsPath;
  StringRef OutputSpecPath = TFOutputSpecOverride;
  if (OutputSpecPath.empty()) {
    llvm::sys::path::append(OutputSpecsPath, ModelPath, "output_spec.json");
    OutputSpecPath = {OutputSpecsPath.data(), OutputSpecsPath.size()};
  }
  if (!loadOutputSpecs(Ctx, OutputSpecPath))
    return;

  Evaluator =
      std::make_unique<TFModelEvaluator>(ModelPath, InputSpecs, OutputSpecs);
  if (!Evaluator || !Evaluator->isValid()) {
    Ctx.emitError("Failed to create inliner saved model evaluator");
    Evaluator.reset();
    return;
  }
}

bool ModelUnderTrainingRunner::loadOutputSpecs(LLVMContext &Ctx,
                                               StringRef FileName) {
  auto BufferOrError = MemoryBuffer::getFileOrSTDIN(FileName);
  if (!BufferOrError) {
    Ctx.emitError("Error opening output specs file: " + FileName + " : " +
                  BufferOrError.getError().message());
    return false;
  }
  auto ParsedJSONValues = json::parse(BufferOrError.get()->getBuffer());
  if (!ParsedJSONValues) {
    Ctx.emitError("Could not parse specs file: " + FileName);
    return false;
  }
  auto ValuesArray = ParsedJSONValues->getAsArray();
  if (!ValuesArray) {
    Ctx.emitError("Expected an array of {tensor_spec:<TensorSpec>, "
                  "logging_name:<name>} dictionaries");
    return false;
  }

  for (const auto &Value : *ValuesArray)
    if (const auto *Obj = Value.getAsObject())
      if (const auto *SpecPart = Obj->get("tensor_spec"))
        if (auto TensorSpec = getTensorSpecFromJSON(Ctx, *SpecPart))
          if (auto LoggingName = Obj->getString("logging_name")) {
            if (!TensorSpec->isElementType<int64_t>() &&
                !TensorSpec->isElementType<int32_t>() &&
                !TensorSpec->isElementType<float>()) {
              Ctx.emitError(
                  "Only int64, int32, and float tensors are supported. "
                  "Found unsupported type for tensor named " +
                  TensorSpec->name());
              return false;
            }
            OutputNames.push_back(LoggingName->str());
            OutputSpecs.push_back(*TensorSpec);
          }

  if (ValuesArray->size() != OutputNames.size()) {
    Ctx.emitError(
        "Unable to parse output spec. It should be a json file containing an "
        "array of dictionaries. Each dictionary must have a 'tensor_spec' key, "
        "with a json object describing a TensorSpec; and a 'logging_name' key, "
        "which is a string to use as name when logging this tensor in the "
        "training log.");
    return false;
  }
  assert(OutputNames.size() == OutputSpecs.size());
  if (OutputNames.empty() || OutputNames[0] != DecisionName) {
    Ctx.emitError("The first output spec must describe the decision tensor, "
                  "and must have the logging_name " +
                  StringRef(DecisionName));
    return false;
  }
  return true;
}

bool ModelUnderTrainingRunner::run() {
  LastEvaluationResult = Evaluator->evaluate();
  if (!LastEvaluationResult.hasValue()) {
    Ctx.emitError("Error evaluating model.");
    return false;
  }
  int64_t Decision = *LastEvaluationResult->getTensorValue<int64_t>(0);
  return static_cast<bool>(Decision);
}

int64_t ModelUnderTrainingRunner::getFeature(int Index) const {
  return *Evaluator->getInput<int64_t>(Index);
}

void ModelUnderTrainingRunner::setFeature(FeatureIndex Index, int64_t Value) {
  size_t NumericIndex = static_cast<size_t>(Index);
  *(Evaluator->getInput<int64_t>(NumericIndex)) = Value;
}

std::unique_ptr<InlineAdvisor> llvm::getDevelopmentModeAdvisor(
    Module &M, ModuleAnalysisManager &MAM,
    std::function<bool(CallBase &)> GetDefaultAdvice) {
  auto &Ctx = M.getContext();
  std::unique_ptr<MLModelRunner> Runner;
  ModelUnderTrainingRunner *MUTRPtr = nullptr;
  bool IsDoingInference = false;
  if (TFModelUnderTrainingPath.empty())
    Runner.reset(new NoInferenceModelRunner(Ctx));
  else {
    auto MUTR = std::make_unique<ModelUnderTrainingRunner>(
        Ctx, TFModelUnderTrainingPath);
    if (!MUTR || !MUTR->isValid()) {
      Ctx.emitError("Could not load the policy model from the provided path");
      return nullptr;
    }
    IsDoingInference = true;
    MUTRPtr = MUTR.get();
    Runner = std::move(MUTR);
  }
  std::unique_ptr<TrainingLogger> Logger;
  if (!TrainingLog.empty())
    Logger = std::make_unique<TrainingLogger>(TrainingLog, MUTRPtr);

  return std::make_unique<DevelopmentModeMLInlineAdvisor>(
      M, MAM, std::move(Runner), GetDefaultAdvice, IsDoingInference,
      std::move(Logger));
}
#endif // defined(LLVM_HAVE_TF_API)