351 lines · 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-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements a model runner using TFLite, allowing the10// 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/MLInlineAdvisor.h"20#include "llvm/Analysis/ModelUnderTrainingRunner.h"21#include "llvm/Analysis/NoInferenceModelRunner.h"22#include "llvm/Analysis/Utils/TFUtils.h"23#include "llvm/Analysis/Utils/TrainingLogger.h"24#include "llvm/IR/LLVMContext.h"25#include "llvm/IR/Module.h"26#include "llvm/Support/CommandLine.h"27#include "llvm/Support/ManagedStatic.h"28 29#include <optional>30#include <vector>31 32using namespace llvm;33 34static cl::opt<std::string> TrainingLog(35 "training-log", cl::Hidden,36 cl::desc("Path where the development - mode inlining log is saved."));37 38static cl::opt<std::string> TFModelUnderTrainingPath(39 "ml-inliner-model-under-training", cl::Hidden,40 cl::desc(R"(Path to SavedModel from the previous training iteration.41The directory is also expected to contain a JSON specification of the 42outputs expected to be logged, where the first entry must be the 43inlining decision. The file containing the specification should be 44called output_spec.json. The expected JSON value is an array of 45dictionaries. Each dictionary should have 2 keys: 46 47- "tensor_spec, followed by the TensorSpec description of the48output; and 49- "logging_name", a string indicating the name to use when50logging the output values. 51 52Example:53[54 {55 "logging_name" : "some_name", 56 "tensor_spec" : { 57 "name" : "model_name", 58 "port" : 0,59 "shape" : [2, 3],60 "type" : "float"61 }62 }63]64 65The first value must always correspond to the decision.)"));66 67static cl::opt<std::string> TFOutputSpecOverride(68 "ml-inliner-output-spec-override", cl::Hidden,69 cl::desc("Override the path to the output spec json file. See "70 "-ml-inliner-model-under-training documentation for the "71 "specification of that file."));72 73static cl::opt<std::string> TFFeedPrefix("ml-inliner-trained-model-feed-prefix",74 cl::Hidden, cl::init("action_"),75 cl::desc("Prefix for feature names."));76 77namespace {78/// An InlineEvent, used by TrainingLogger.79struct InlineEvent {80 /// What the default policy's decision would have been.81 int64_t DefaultDecision = 0;82 83 /// What we advised. When training off the default policy, this is the same as84 /// DefaultDecision.85 int64_t AdvisedDecision = 0;86 87 /// What actually happened. This would be 'false' in the case of an inline88 /// error, even if AdvisedDecision were true, otherwise it agrees with89 /// AdvisedDecision.90 bool Effect = false;91};92 93/// Collect data we may use for training a model.94class TrainingLogger final {95public:96 TrainingLogger(StringRef LogFileName, const ModelUnderTrainingRunner *MUTR,97 const std::vector<TensorSpec> &FeatureMap);98 99 /// Log one inlining event.100 void logInlineEvent(const InlineEvent &Event,101 const MLModelRunner &ModelRunner);102 103private:104 StringRef LogFileName;105 const ModelUnderTrainingRunner *const MUTR;106 const std::vector<TensorSpec> &FeatureMap;107 108 std::unique_ptr<Logger> L;109 BitVector Effects;110 /// Set these 2 clearly OOB, to make sure we set them later.111 size_t DefaultDecisionPos = std::numeric_limits<size_t>::max();112 size_t DecisionPos = std::numeric_limits<size_t>::max();113};114 115/// An extension of the MLInlineAdvisor for the 'development' mode, targeting116/// the offline training scenario. Note that training happens outside of the117/// compiler, this facility is concerned with producing training data ("logs").118/// This InlineAdvisor can operate in the following modes:119///120/// 1) collect logs for the default policy. This is useful for bootstrapping121/// training, which will be considerably faster by starting from a reasonable122/// policy.123///124/// 2) collect logs for the ML policy, using a model from a previous125/// training. Potentially, that model uses internally some small random126/// perturbation of its weights, to induce exploration (setting this up is the127/// responsibility of the training algorithm). The logs would then be used to128/// retrain and improve on this model.129///130/// 3) use the provided model, with no logging. This is useful for end to end131/// validation - the model, in this case, is a release candidate and shouldn't132/// have random perturbations. It is a convenience feature: rather than needing133/// to take the release candidate model and compile it in 'release' mode,134/// validate it, then potentially discard it, it's easier to just pass the model135/// to the compiler, albeit compilation would be slower, as a one-off. Once the136/// model behaves satisfactorily, it can be compiled AOT, for efficiency, in137/// release mode. The expectation is that a well-trained model provides a good138/// policy over a sufficiently diverse codebase, over many changes (i.e.139/// training happens seldom).140class DevelopmentModeMLInlineAdvisor : public MLInlineAdvisor {141public:142 DevelopmentModeMLInlineAdvisor(143 Module &M, ModuleAnalysisManager &MAM,144 std::function<145 std::unique_ptr<MLModelRunner>(const std::vector<TensorSpec> &)>146 GetModelRunner,147 std::function<bool(CallBase &)> GetDefaultAdvice);148 149 std::unique_ptr<MLInlineAdvice>150 getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE) override;151 152private:153 bool isLogging() const { return !!Logger; }154 std::unique_ptr<MLInlineAdvice> getMandatoryAdviceImpl(CallBase &CB) override;155 156 const bool IsDoingInference;157 std::unique_ptr<TrainingLogger> Logger;158};159 160/// A variant of MLInlineAdvice that tracks all non-trivial inlining161/// decisions, for training/logging.162class LoggingMLInlineAdvice : public MLInlineAdvice {163public:164 LoggingMLInlineAdvice(DevelopmentModeMLInlineAdvisor *Advisor, CallBase &CB,165 OptimizationRemarkEmitter &ORE, bool Recommendation,166 TrainingLogger &Logger, bool DefaultDecision,167 bool Mandatory = false)168 : MLInlineAdvice(Advisor, CB, ORE, Recommendation), Logger(Logger),169 DefaultDecision(DefaultDecision), Mandatory(Mandatory) {}170 171 virtual ~LoggingMLInlineAdvice() = default;172 173private:174 DevelopmentModeMLInlineAdvisor *getAdvisor() const {175 return static_cast<DevelopmentModeMLInlineAdvisor *>(Advisor);176 }177 void recordInliningImpl() override {178 MLInlineAdvice::recordInliningImpl();179 log(/*Success=*/true);180 }181 182 void recordInliningWithCalleeDeletedImpl() override {183 MLInlineAdvice::recordInliningWithCalleeDeletedImpl();184 log(/*Success=*/true);185 }186 187 void recordUnsuccessfulInliningImpl(const InlineResult &Result) override {188 MLInlineAdvice::recordUnsuccessfulInliningImpl(Result);189 log(/*Success=*/false);190 }191 192 void recordUnattemptedInliningImpl() override {193 MLInlineAdvice::recordUnattemptedInliningImpl();194 log(/*Success=*/false);195 }196 197 void log(bool Success) {198 if (Mandatory)199 return;200 InlineEvent Event;201 Event.AdvisedDecision = isInliningRecommended();202 Event.DefaultDecision = DefaultDecision;203 Event.Effect = Success;204 Logger.logInlineEvent(Event, getAdvisor()->getModelRunner());205 }206 207 TrainingLogger &Logger;208 const int64_t DefaultDecision;209 const int64_t Mandatory;210};211 212static const std::vector<TensorSpec> TrainingOnlyFeatures{213 TensorSpec::createSpec<float>(TFFeedPrefix + "discount", {1}),214 TensorSpec::createSpec<float>(TFFeedPrefix + "reward", {1}),215 TensorSpec::createSpec<int32_t>(TFFeedPrefix + "step_type", {1})};216 217// add TFFeedPrefix to the names and also add the "TrainingOnlyFeatures" which218// the model runner needs to see present. We don't set them ourselves or219// interact with them.220static const std::vector<TensorSpec>221convertInputFeatures(const std::vector<TensorSpec> &OriginalFeatures) {222 std::vector<TensorSpec> InputSpecs;223 for (const auto &Feature : OriginalFeatures)224 InputSpecs.push_back(TensorSpec(TFFeedPrefix + Feature.name(), Feature));225 append_range(InputSpecs, TrainingOnlyFeatures);226 return InputSpecs;227}228 229} // namespace230 231TrainingLogger::TrainingLogger(StringRef LogFileName,232 const ModelUnderTrainingRunner *MUTR,233 const std::vector<TensorSpec> &FeatureMap)234 : LogFileName(LogFileName), MUTR(MUTR), FeatureMap(FeatureMap) {235 // The first output is the inlining decision.236 std::vector<TensorSpec> FT(FeatureMap.begin(), FeatureMap.end());237 238 if (MUTR)239 append_range(FT, MUTR->extraOutputsForLoggingSpecs());240 241 DefaultDecisionPos = FT.size();242 FT.push_back(DefaultDecisionSpec);243 244 DecisionPos = FT.size();245 FT.push_back(InlineDecisionSpec);246 std::error_code EC;247 auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);248 if (EC)249 dbgs() << (EC.message() + ":" + TrainingLog);250 251 L = std::make_unique<Logger>(std::move(OS), FT,252 TensorSpec::createSpec<int64_t>(RewardName, {1}),253 false);254 L->switchContext("");255}256 257/// Log one inlining event.258void TrainingLogger::logInlineEvent(const InlineEvent &Event,259 const MLModelRunner &ModelRunner) {260 L->startObservation();261 size_t CurrentFeature = 0;262 for (; CurrentFeature < FeatureMap.size(); ++CurrentFeature)263 L->logTensorValue(CurrentFeature,264 reinterpret_cast<const char *>(265 ModelRunner.getTensorUntyped(CurrentFeature)));266 267 if (MUTR)268 for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size(); ++I) {269 const char *RawData =270 reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I));271 L->logTensorValue(CurrentFeature, RawData);272 ++CurrentFeature;273 }274 275 assert(CurrentFeature == DefaultDecisionPos);276 L->logTensorValue(DefaultDecisionPos,277 reinterpret_cast<const char *>(&Event.DefaultDecision));278 L->logTensorValue(DecisionPos,279 reinterpret_cast<const char *>(&Event.AdvisedDecision));280 L->endObservation();281 282 // For debugging / later use283 Effects.push_back(Event.Effect);284}285 286DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor(287 Module &M, ModuleAnalysisManager &MAM,288 std::function<289 std::unique_ptr<MLModelRunner>(const std::vector<TensorSpec> &)>290 GetModelRunner,291 std::function<bool(CallBase &)> GetDefaultAdvice)292 : MLInlineAdvisor(M, MAM, GetModelRunner, GetDefaultAdvice),293 IsDoingInference(isa<ModelUnderTrainingRunner>(getModelRunner())) {294 // We cannot have the case of neither inference nor logging.295 if (!TrainingLog.empty())296 Logger = std::make_unique<TrainingLogger>(297 TrainingLog, dyn_cast<ModelUnderTrainingRunner>(ModelRunner.get()),298 getFeatureMap());299 assert(IsDoingInference || isLogging());300}301 302std::unique_ptr<MLInlineAdvice>303DevelopmentModeMLInlineAdvisor::getMandatoryAdviceImpl(CallBase &CB) {304 return std::make_unique<LoggingMLInlineAdvice>(305 /*Advisor=*/this,306 /*CB=*/CB, /*ORE=*/getCallerORE(CB), /*Recommendation=*/true,307 /*Logger=*/*Logger,308 /*DefaultDecision=*/true, /*Mandatory*/ true);309}310 311std::unique_ptr<MLInlineAdvice>312DevelopmentModeMLInlineAdvisor::getAdviceFromModel(313 CallBase &CB, OptimizationRemarkEmitter &ORE) {314 if (IsDoingInference && !isLogging())315 return MLInlineAdvisor::getAdviceFromModel(CB, ORE);316 317 bool DefaultAdvice = GetDefaultAdvice(CB);318 auto Recommendation =319 IsDoingInference ? static_cast<bool>(ModelRunner->evaluate<int64_t>())320 : DefaultAdvice;321 return std::make_unique<LoggingMLInlineAdvice>(322 /*Advisor=*/this,323 /*CB=*/CB, /*ORE=*/ORE, /*Recommendation=*/Recommendation,324 /*Logger=*/*Logger,325 /*DefaultDecision=*/DefaultAdvice);326}327 328std::unique_ptr<InlineAdvisor> llvm::getDevelopmentModeAdvisor(329 Module &M, ModuleAnalysisManager &MAM,330 std::function<bool(CallBase &)> GetDefaultAdvice) {331 auto &Ctx = M.getContext();332 auto RunnerFactory = [&](const std::vector<TensorSpec> &InputFeatures)333 -> std::unique_ptr<MLModelRunner> {334 std::unique_ptr<MLModelRunner> Runner;335 const std::vector<TensorSpec> ConvertedFeatures =336 convertInputFeatures(InputFeatures);337 if (TFModelUnderTrainingPath.empty())338 Runner.reset(new NoInferenceModelRunner(Ctx, ConvertedFeatures));339 else340 Runner = ModelUnderTrainingRunner::createAndEnsureValid(341 Ctx, TFModelUnderTrainingPath, DecisionName, ConvertedFeatures,342 TFOutputSpecOverride);343 if (!Runner)344 return nullptr;345 return Runner;346 };347 return std::make_unique<DevelopmentModeMLInlineAdvisor>(M, MAM, RunnerFactory,348 GetDefaultAdvice);349}350#endif // defined(LLVM_HAVE_TFLITE)351