llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

   1 //===- DevelopmentModeInlineAdvisor.cpp - runtime-loadable model runner  --===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file implements a model runner using TFLite, allowing the
  10 // loading of a model from a command line option.
  11 //
  12 //===----------------------------------------------------------------------===//
  13 #include "llvm/Analysis/TensorSpec.h"
  14 #include "llvm/Config/config.h"
  15 #if defined(LLVM_HAVE_TFLITE)
  16
  17 #include "llvm/ADT/BitVector.h"
  18 #include "llvm/Analysis/CallGraph.h"
  19 #include "llvm/Analysis/InlineSizeEstimatorAnalysis.h"
  20 #include "llvm/Analysis/MLInlineAdvisor.h"
  21 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
  22 #include "llvm/Analysis/NoInferenceModelRunner.h"
  23 #include "llvm/Analysis/Utils/TFUtils.h"
  24 #include "llvm/Analysis/Utils/TrainingLogger.h"
  25 #include "llvm/IR/LLVMContext.h"
  26 #include "llvm/IR/Module.h"
  27 #include "llvm/Support/CommandLine.h"
  28 #include "llvm/Support/ManagedStatic.h"
  29
  30 #include <vector>
  31 #include <optional>
  32
  33 using namespace llvm;
  34
  35 static cl::opt<std::string> TrainingLog(
  36     "training-log", cl::Hidden,
  37     cl::desc("Path where the development - mode inlining log is saved."));
  38
  39 static cl::opt<std::string> TFModelUnderTrainingPath(
  40     "ml-inliner-model-under-training", cl::Hidden,
  41     cl::desc(R"(Path to SavedModel from the previous training iteration.
  42 The directory is also expected to contain a JSON specification of the
  43 outputs expected to be logged, where the first entry must be the
  44 inlining decision. The file containing the specification should be
  45 called output_spec.json. The expected JSON value is an array of
  46 dictionaries. Each dictionary should have 2 keys:
  47
  48 - "tensor_spec, followed by the TensorSpec description of the
  49 output; and
  50 - "logging_name", a string indicating the name to use when
  51 logging the output values.
  52
  53 Example:
  54 [
  55   {
  56     "logging_name" : "some_name",
  57     "tensor_spec" : {
  58       "name" : "model_name",
  59       "port" : 0,
  60       "shape" : [2, 3],
  61       "type" : "float"
  62       }
  63   }
  64 ]
  65
  66 The first value must always correspond to the decision.)"));
  67
  68 static cl::opt<std::string> TFOutputSpecOverride(
  69     "ml-inliner-output-spec-override", cl::Hidden,
  70     cl::desc("Override the path to the output spec json file. See "
  71              "-ml-inliner-model-under-training documentation for the "
  72              "specification of that file."));
  73
  74 static cl::opt<std::string> TFFeedPrefix("ml-inliner-trained-model-feed-prefix",
  75                                          cl::Hidden, cl::init("action_"),
  76                                          cl::desc("Prefix for feature names."));
  77
  78 namespace {
  79 /// An InlineEvent, used by TrainingLogger.
  80 struct InlineEvent {
  81   /// What the default policy's decision would have been.
  82   int64_t DefaultDecision = 0;
  83
  84   /// What we advised. When training off the default policy, this is the same as
  85   /// DefaultDecision.
  86   int64_t AdvisedDecision = 0;
  87
  88   /// What actually happened. This would be 'false' in the case of an inline
  89   /// error, even if AdvisedDecision were true, otherwise it agrees with
  90   /// AdvisedDecision.
  91   bool Effect = false;
  92
  93   /// What the change in size was: size_after - size_before
  94   int64_t Reward = 0;
  95 };
  96
  97 /// Collect data we may use for training a model.
  98 class TrainingLogger final {
  99 public:
 100   TrainingLogger(StringRef LogFileName, const ModelUnderTrainingRunner *MUTR);
 101
 102   /// Log one inlining event.
 103   void logInlineEvent(const InlineEvent &Event,
 104                       const MLModelRunner &ModelRunner);
 105
 106 private:
 107   StringRef LogFileName;
 108   const ModelUnderTrainingRunner *const MUTR;
 109   std::unique_ptr<Logger> L;
 110   BitVector Effects;
 111   /// Set these 2 clearly OOB, to make sure we set them later.
 112   size_t DefaultDecisionPos = std::numeric_limits<size_t>::max();
 113   size_t DecisionPos = std::numeric_limits<size_t>::max();
 114 };
 115
 116 /// An extension of the MLInlineAdvisor for the 'development' mode, targeting
 117 /// the offline training scenario. Note that training happens outside of the
 118 /// compiler, this facility is concerned with producing training data ("logs").
 119 /// This InlineAdvisor can operate in the following modes:
 120 ///
 121 /// 1) collect logs for the default policy. This is useful for bootstrapping
 122 /// training, which will be considerably faster by starting from a reasonable
 123 /// policy.
 124 ///
 125 /// 2) collect logs for the ML policy, using a model from a previous
 126 /// training. Potentially, that model uses internally some small random
 127 /// perturbation of its weights, to induce exploration (setting this up is the
 128 /// responsibility of the training algorithm). The logs would then be used to
 129 /// retrain and improve on this model.
 130 ///
 131 /// 3) use the provided model, with no logging. This is useful for end to end
 132 /// validation - the model, in this case, is a release candidate and shouldn't
 133 /// have random perturbations. It is a convenience feature: rather than needing
 134 /// to take the release candidate model and compile it in 'release' mode,
 135 /// validate it, then potentially discard it, it's easier to just pass the model
 136 /// to the compiler, albeit compilation would be slower, as a one-off. Once the
 137 /// model behaves satisfactorily, it can be compiled AOT, for efficiency, in
 138 /// release mode. The expectation is that a well-trained model provides a good
 139 /// policy over a sufficiently diverse codebase, over many changes (i.e.
 140 /// training happens seldom).
 141 class DevelopmentModeMLInlineAdvisor : public MLInlineAdvisor {
 142 public:
 143   DevelopmentModeMLInlineAdvisor(
 144       Module &M, ModuleAnalysisManager &MAM,
 145       std::unique_ptr<MLModelRunner> ModelRunner,
 146       std::function<bool(CallBase &)> GetDefaultAdvice,
 147       std::unique_ptr<TrainingLogger> Logger);
 148
 149   size_t getTotalSizeEstimate();
 150
 151   void updateNativeSizeEstimate(int64_t Change) {
 152     *CurrentNativeSize += Change;
 153   }
 154   void resetNativeSize(Function *F) {
 155     PreservedAnalyses PA = PreservedAnalyses::all();
 156     PA.abandon<InlineSizeEstimatorAnalysis>();
 157     FAM.invalidate(*F, PA);
 158   }
 159
 160   std::unique_ptr<MLInlineAdvice>
 161   getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE) override;
 162
 163   std::optional<size_t> getNativeSizeEstimate(const Function &F) const;
 164
 165 private:
 166   bool isLogging() const { return !!Logger; }
 167   std::unique_ptr<MLInlineAdvice> getMandatoryAdviceImpl(CallBase &CB) override;
 168
 169   const bool IsDoingInference;
 170   std::unique_ptr<TrainingLogger> Logger;
 171
 172   const std::optional<int32_t> InitialNativeSize;
 173   std::optional<int32_t> CurrentNativeSize;
 174 };
 175
 176 /// A variant of MLInlineAdvice that tracks all non-trivial inlining
 177 /// decisions, for training/logging.
 178 class LoggingMLInlineAdvice : public MLInlineAdvice {
 179 public:
 180   LoggingMLInlineAdvice(DevelopmentModeMLInlineAdvisor *Advisor, CallBase &CB,
 181                         OptimizationRemarkEmitter &ORE, bool Recommendation,
 182                         TrainingLogger &Logger,
 183                         std::optional<size_t> CallerSizeEstimateBefore,
 184                         std::optional<size_t> CalleeSizeEstimateBefore,
 185                         bool DefaultDecision, bool Mandatory = false)
 186       : MLInlineAdvice(Advisor, CB, ORE, Recommendation), Logger(Logger),
 187         CallerSizeEstimateBefore(CallerSizeEstimateBefore),
 188         CalleeSizeEstimateBefore(CalleeSizeEstimateBefore),
 189         DefaultDecision(DefaultDecision), Mandatory(Mandatory) {}
 190
 191   virtual ~LoggingMLInlineAdvice() = default;
 192
 193 private:
 194   DevelopmentModeMLInlineAdvisor *getAdvisor() const {
 195     return static_cast<DevelopmentModeMLInlineAdvisor *>(Advisor);
 196   }
 197   void recordInliningImpl() override {
 198     MLInlineAdvice::recordInliningImpl();
 199     getAdvisor()->resetNativeSize(Caller);
 200     int Reward = std::numeric_limits<int>::max();
 201     if (InlineSizeEstimatorAnalysis::isEvaluatorRequested() &&
 202         !getAdvisor()->isForcedToStop()) {
 203       int NativeSizeAfter = *getAdvisor()->getNativeSizeEstimate(*Caller) +
 204                             *CalleeSizeEstimateBefore;
 205       Reward = NativeSizeAfter -
 206                (*CallerSizeEstimateBefore + *CalleeSizeEstimateBefore);
 207       getAdvisor()->updateNativeSizeEstimate(Reward);
 208     }
 209     log(Reward, /*Success=*/true);
 210   }
 211
 212   void recordInliningWithCalleeDeletedImpl() override {
 213     MLInlineAdvice::recordInliningWithCalleeDeletedImpl();
 214     getAdvisor()->resetNativeSize(Caller);
 215     if (InlineSizeEstimatorAnalysis::isEvaluatorRequested() &&
 216         !getAdvisor()->isForcedToStop()) {
 217       int NativeSizeAfter = *getAdvisor()->getNativeSizeEstimate(*Caller);
 218       int Reward = NativeSizeAfter -
 219                    (*CallerSizeEstimateBefore + *CalleeSizeEstimateBefore);
 220       getAdvisor()->updateNativeSizeEstimate(Reward);
 221       log(Reward, /*Success=*/true);
 222     } else {
 223       log(NoReward, /*Success=*/true);
 224     }
 225   }
 226
 227   void recordUnsuccessfulInliningImpl(const InlineResult &Result) override {
 228     MLInlineAdvice::recordUnsuccessfulInliningImpl(Result);
 229     log(NoReward, /*Success=*/false);
 230   }
 231
 232   void recordUnattemptedInliningImpl() override {
 233     MLInlineAdvice::recordUnattemptedInliningImpl();
 234     log(NoReward, /*Success=*/false);
 235   }
 236
 237   void log(int64_t Reward, bool Success) {
 238     if (Mandatory)
 239       return;
 240     InlineEvent Event;
 241     Event.AdvisedDecision = isInliningRecommended();
 242     Event.DefaultDecision = DefaultDecision;
 243     Event.Effect = Success;
 244     Event.Reward = Reward;
 245     Logger.logInlineEvent(Event, getAdvisor()->getModelRunner());
 246   }
 247
 248   static const int64_t NoReward = 0;
 249   TrainingLogger &Logger;
 250   const std::optional<size_t> CallerSizeEstimateBefore;
 251   const std::optional<size_t> CalleeSizeEstimateBefore;
 252   const int64_t DefaultDecision;
 253   const int64_t Mandatory;
 254 };
 255
 256 static const std::vector<TensorSpec> TrainingOnlyFeatures{
 257     TensorSpec::createSpec<float>(TFFeedPrefix + "discount", {1}),
 258     TensorSpec::createSpec<float>(TFFeedPrefix + "reward", {1}),
 259     TensorSpec::createSpec<int32_t>(TFFeedPrefix + "step_type", {1})};
 260
 261 static const std::vector<TensorSpec> getInputFeatures() {
 262   std::vector<TensorSpec> InputSpecs;
 263   for (size_t I = 0; I < NumberOfFeatures; ++I)
 264     InputSpecs.push_back(TensorSpec::createSpec<int64_t>(
 265         TFFeedPrefix + FeatureMap[I].name(), FeatureMap[I].shape()));
 266   append_range(InputSpecs, TrainingOnlyFeatures);
 267   return InputSpecs;
 268 }
 269
 270 } // namespace
 271
 272 TrainingLogger::TrainingLogger(StringRef LogFileName,
 273                                const ModelUnderTrainingRunner *MUTR)
 274     : LogFileName(LogFileName), MUTR(MUTR) {
 275   // The first output is the inlining decision.
 276   std::vector<TensorSpec> FT(FeatureMap.begin(), FeatureMap.end());
 277
 278   if (MUTR)
 279     append_range(FT, MUTR->extraOutputsForLoggingSpecs());
 280
 281   DefaultDecisionPos = FT.size();
 282   FT.push_back(DefaultDecisionSpec);
 283
 284   DecisionPos = FT.size();
 285   FT.push_back(InlineDecisionSpec);
 286   std::error_code EC;
 287   auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
 288   if (EC)
 289     dbgs() << (EC.message() + ":" + TrainingLog);
 290
 291   L = std::make_unique<Logger>(
 292       std::move(OS), FT, TensorSpec::createSpec<int64_t>(RewardName, {1}),
 293       InlineSizeEstimatorAnalysis::isEvaluatorRequested());
 294   L->switchContext("");
 295 }
 296
 297 /// Log one inlining event.
 298 void TrainingLogger::logInlineEvent(const InlineEvent &Event,
 299                                     const MLModelRunner &ModelRunner) {
 300   L->startObservation();
 301   size_t CurrentFeature = 0;
 302   for (; CurrentFeature < NumberOfFeatures; ++CurrentFeature)
 303     L->logTensorValue(CurrentFeature,
 304                       reinterpret_cast<const char *>(
 305                           ModelRunner.getTensorUntyped(CurrentFeature)));
 306
 307   if (MUTR)
 308     for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size(); ++I) {
 309       const char *RawData =
 310           reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I));
 311       L->logTensorValue(CurrentFeature, RawData);
 312       ++CurrentFeature;
 313     }
 314
 315   assert(CurrentFeature == DefaultDecisionPos);
 316   L->logTensorValue(DefaultDecisionPos,
 317                     reinterpret_cast<const char *>(&Event.DefaultDecision));
 318   L->logTensorValue(DecisionPos,
 319                     reinterpret_cast<const char *>(&Event.AdvisedDecision));
 320   L->endObservation();
 321   if (InlineSizeEstimatorAnalysis::isEvaluatorRequested())
 322     L->logReward(Event.Reward);
 323
 324   // For debugging / later use
 325   Effects.push_back(Event.Effect);
 326 }
 327
 328 DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor(
 329     Module &M, ModuleAnalysisManager &MAM,
 330     std::unique_ptr<MLModelRunner> ModelRunner,
 331     std::function<bool(CallBase &)> GetDefaultAdvice,
 332     std::unique_ptr<TrainingLogger> Logger)
 333     : MLInlineAdvisor(M, MAM, std::move(ModelRunner), GetDefaultAdvice),
 334       IsDoingInference(isa<ModelUnderTrainingRunner>(getModelRunner())),
 335       Logger(std::move(Logger)),
 336       InitialNativeSize(isLogging() ? getTotalSizeEstimate() : 0),
 337       CurrentNativeSize(InitialNativeSize) {
 338   // We cannot have the case of neither inference nor logging.
 339   assert(IsDoingInference || isLogging());
 340 }
 341
 342 std::optional<size_t>
 343 DevelopmentModeMLInlineAdvisor::getNativeSizeEstimate(const Function &F) const {
 344   if (!InlineSizeEstimatorAnalysis::isEvaluatorRequested())
 345     return std::nullopt;
 346   auto &R =
 347       FAM.getResult<InlineSizeEstimatorAnalysis>(const_cast<Function &>(F));
 348   if (!R) {
 349     F.getParent()->getContext().emitError(
 350         "Native size estimator is not present.");
 351     return 0;
 352   }
 353   return *R;
 354 }
 355
 356 std::unique_ptr<MLInlineAdvice>
 357 DevelopmentModeMLInlineAdvisor::getMandatoryAdviceImpl(CallBase &CB) {
 358   return std::make_unique<LoggingMLInlineAdvice>(
 359       /*Advisor=*/this,
 360       /*CB=*/CB, /*ORE=*/getCallerORE(CB), /*Recommendation=*/true,
 361       /*Logger=*/*Logger,
 362       /*CallerSizeEstimateBefore=*/getNativeSizeEstimate(*CB.getCaller()),
 363       /*CalleeSizeEstimateBefore=*/
 364       getNativeSizeEstimate(*CB.getCalledFunction()),
 365       /*DefaultDecision=*/true, /*Mandatory*/ true);
 366 }
 367
 368 std::unique_ptr<MLInlineAdvice>
 369 DevelopmentModeMLInlineAdvisor::getAdviceFromModel(
 370     CallBase &CB, OptimizationRemarkEmitter &ORE) {
 371   if (IsDoingInference && !isLogging())
 372     return MLInlineAdvisor::getAdviceFromModel(CB, ORE);
 373
 374   bool DefaultAdvice = GetDefaultAdvice(CB);
 375   auto Recommendation =
 376       IsDoingInference ? static_cast<bool>(ModelRunner->evaluate<int64_t>())
 377                        : DefaultAdvice;
 378   return std::make_unique<LoggingMLInlineAdvice>(
 379       /*Advisor=*/this,
 380       /*CB=*/CB, /*ORE=*/ORE, /*Recommendation=*/Recommendation,
 381       /*Logger=*/*Logger,
 382       /*CallerSizeEstimateBefore=*/getNativeSizeEstimate(*CB.getCaller()),
 383       /*CalleeSizeEstimateBefore=*/
 384       getNativeSizeEstimate(*CB.getCalledFunction()),
 385       /*DefaultDecision=*/DefaultAdvice);
 386 }
 387
 388 size_t DevelopmentModeMLInlineAdvisor::getTotalSizeEstimate() {
 389   if (!InlineSizeEstimatorAnalysis::isEvaluatorRequested())
 390     return 0;
 391   size_t Ret = 0;
 392   for (auto &F : M) {
 393     if (F.isDeclaration())
 394       continue;
 395     Ret += *getNativeSizeEstimate(F);
 396   }
 397   return Ret;
 398 }
 399
 400 std::unique_ptr<InlineAdvisor> llvm::getDevelopmentModeAdvisor(
 401     Module &M, ModuleAnalysisManager &MAM,
 402     std::function<bool(CallBase &)> GetDefaultAdvice) {
 403   auto &Ctx = M.getContext();
 404   std::unique_ptr<MLModelRunner> Runner;
 405   if (TFModelUnderTrainingPath.empty())
 406     Runner.reset(new NoInferenceModelRunner(Ctx, getInputFeatures()));
 407   else
 408     Runner = ModelUnderTrainingRunner::createAndEnsureValid(
 409         Ctx, TFModelUnderTrainingPath, DecisionName, getInputFeatures(),
 410         TFOutputSpecOverride);
 411   if (!Runner)
 412     return nullptr;
 413   std::unique_ptr<TrainingLogger> Logger;
 414   if (!TrainingLog.empty())
 415     Logger = std::make_unique<TrainingLogger>(
 416         TrainingLog, dyn_cast<ModelUnderTrainingRunner>(Runner.get()));
 417
 418   return std::make_unique<DevelopmentModeMLInlineAdvisor>(
 419       M, MAM, std::move(Runner), GetDefaultAdvice, std::move(Logger));
 420 }
 421 #endif // defined(LLVM_HAVE_TFLITE)