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/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 the
48output; and
49- "logging_name", a string indicating the name to use when
50logging 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 as
84 /// DefaultDecision.
85 int64_t AdvisedDecision = 0;
86
87 /// What actually happened. This would be 'false' in the case of an inline
88 /// error, even if AdvisedDecision were true, otherwise it agrees with
89 /// 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, targeting
116/// the offline training scenario. Note that training happens outside of the
117/// 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 bootstrapping
121/// training, which will be considerably faster by starting from a reasonable
122/// policy.
123///
124/// 2) collect logs for the ML policy, using a model from a previous
125/// training. Potentially, that model uses internally some small random
126/// perturbation of its weights, to induce exploration (setting this up is the
127/// responsibility of the training algorithm). The logs would then be used to
128/// retrain and improve on this model.
129///
130/// 3) use the provided model, with no logging. This is useful for end to end
131/// validation - the model, in this case, is a release candidate and shouldn't
132/// have random perturbations. It is a convenience feature: rather than needing
133/// 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 model
135/// to the compiler, albeit compilation would be slower, as a one-off. Once the
136/// model behaves satisfactorily, it can be compiled AOT, for efficiency, in
137/// release mode. The expectation is that a well-trained model provides a good
138/// 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 inlining
161/// 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" which
218// the model runner needs to see present. We don't set them ourselves or
219// 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} // namespace
230
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 use
283 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 else
340 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