path: root/contrib/libs/llvm12/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/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 <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. 
  int64_t DefaultDecision = 0; 
 
  /// What we advised. When training off the default policy, this is the same as 
  /// DefaultDecision. 
  int64_t AdvisedDecision = 0; 
 
  /// 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: 
  StringRef LogFileName; 
  const ModelUnderTrainingRunner *const MUTR; 
  std::unique_ptr<Logger> L; 
  std::vector<bool> Effects; 
  /// There's at least one output. We'll set this to a different value if MUTR 
  /// is avaliable. 
  size_t OutputCount = 1; 
  /// Set these 2 clearly OOB, to make sure we set them later. 
  size_t DefaultDecisionPos = std::numeric_limits<size_t>::max(); 
  size_t DecisionPos = std::numeric_limits<size_t>::max(); 
}; 
 
/// 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> 
  getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE) override; 
 
  Optional<size_t> getNativeSizeEstimate(const Function &F) const; 
 
private: 
  bool isLogging() const { return !!Logger; } 
  std::unique_ptr<MLInlineAdvice> getMandatoryAdviceImpl(CallBase &CB) override; 
 
  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 int64_t DefaultDecision; 
  const int64_t 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<LoggedFeatureSpec> &outputLoggedFeatureSpecs() const { 
    return OutputSpecs; 
  } 
 
  const Optional<TFModelEvaluator::EvaluationResult> & 
  lastEvaluationResult() const { 
    return LastEvaluationResult; 
  } 
 
private: 
  std::unique_ptr<TFModelEvaluator> Evaluator; 
  std::vector<LoggedFeatureSpec> OutputSpecs; 
  Optional<TFModelEvaluator::EvaluationResult> LastEvaluationResult; 
 
  // 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) { 
  // The first output is the inlining decision. 
  if (MUTR) 
    OutputCount = MUTR->outputLoggedFeatureSpecs().size(); 
  std::vector<LoggedFeatureSpec> FT; 
 
  for (size_t I = 0; I < NumberOfFeatures; ++I) 
    FT.push_back( 
        {TensorSpec::createSpec<int64_t>(FeatureNameMap.at(I), {1}), None}); 
  if (MUTR && MUTR->outputLoggedFeatureSpecs().size() > 1) 
    append_range(FT, drop_begin(MUTR->outputLoggedFeatureSpecs())); 
 
  DefaultDecisionPos = FT.size(); 
  FT.push_back( 
      {TensorSpec::createSpec<int64_t>(DefaultDecisionName, {1}), None}); 
 
  DecisionPos = FT.size(); 
  FT.push_back({TensorSpec::createSpec<int64_t>(DecisionName, {1}), None}); 
 
  L = std::make_unique<Logger>( 
      FT, TensorSpec::createSpec<int64_t>(RewardName, {1}), 
      InlineSizeEstimatorAnalysis::isEvaluatorRequested()); 
} 
 
/// Log one inlining event. 
void TrainingLogger::logInlineEvent(const InlineEvent &Event, 
                                    const MLModelRunner &ModelRunner) { 
  size_t CurrentFeature = 0; 
  for (; CurrentFeature < NumberOfFeatures; ++CurrentFeature) { 
    int64_t F = ModelRunner.getFeature(CurrentFeature); 
    L->logTensorValue(CurrentFeature, &F); 
  } 
 
  for (size_t I = 1; I < OutputCount; ++I) { 
    const auto &Result = *MUTR->lastEvaluationResult(); 
    auto &Spec = MUTR->outputLoggedFeatureSpecs()[I].Spec; 
    const char *RawData = 
        reinterpret_cast<const char *>(Result.getUntypedTensorValue(I)); 
    L->logTensorValue(CurrentFeature, RawData, 
                      Spec.getElementCount() * Spec.getElementByteSize()); 
    ++CurrentFeature; 
  } 
 
  assert(CurrentFeature == DefaultDecisionPos); 
  L->logTensorValue(DefaultDecisionPos, &Event.DefaultDecision); 
  L->logTensorValue(DecisionPos, &Event.AdvisedDecision); 
  if (InlineSizeEstimatorAnalysis::isEvaluatorRequested()) 
    L->logReward(Event.Reward); 
 
  // For debugging / later use 
  Effects.push_back(Event.Effect); 
} 
 
void TrainingLogger::print() { 
  std::error_code EC; 
  raw_fd_ostream OutFile(LogFileName, EC); 
  L->print(OutFile); 
} 
 
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::getMandatoryAdviceImpl(CallBase &CB) { 
  return std::make_unique<LoggingMLInlineAdvice>( 
      /*Advisor=*/this, 
      /*CB=*/CB, /*ORE=*/getCallerORE(CB), /*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})); 
  append_range(InputSpecs, TrainingOnlyFeatures); 
  if (auto MaybeOutSpecs = 
          loadOutputSpecs(Ctx, DecisionName, ModelPath, TFOutputSpecOverride)) 
    OutputSpecs = std::move(*MaybeOutSpecs); 
  else 
    return; 
 
  Evaluator = std::make_unique<TFModelEvaluator>( 
      ModelPath, InputSpecs, [&](size_t I) { return OutputSpecs[I].Spec; }, 
      OutputSpecs.size()); 
  if (!Evaluator || !Evaluator->isValid()) { 
    Ctx.emitError("Failed to create inliner saved model evaluator"); 
    Evaluator.reset(); 
    return; 
  } 
} 
 
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)