1//===- MLRegAllocPriorityAdvisor.cpp - ML priority advisor-----------------===//
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// Implementation of the ML priority advisor and reward injection pass
10//
11//===----------------------------------------------------------------------===//
12
13#include "AllocationOrder.h"
14#include "RegAllocGreedy.h"
15#include "llvm/Analysis/AliasAnalysis.h"
16#include "llvm/Analysis/InteractiveModelRunner.h"
17#include "llvm/Analysis/MLModelRunner.h"
18#include "llvm/Analysis/ReleaseModeModelRunner.h"
19#include "llvm/Analysis/TensorSpec.h"
20#include "llvm/CodeGen/CalcSpillWeights.h"
21#include "llvm/CodeGen/LiveRegMatrix.h"
22#include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
23#include "llvm/CodeGen/MachineFunction.h"
24#include "llvm/CodeGen/MachineLoopInfo.h"
25#include "llvm/CodeGen/MachineRegisterInfo.h"
26#include "llvm/CodeGen/Passes.h"
27#include "llvm/CodeGen/RegAllocPriorityAdvisor.h"
28#include "llvm/CodeGen/RegisterClassInfo.h"
29#include "llvm/CodeGen/SlotIndexes.h"
30#include "llvm/CodeGen/VirtRegMap.h"
31#include "llvm/InitializePasses.h"
32#include "llvm/Pass.h"
33#include "llvm/PassRegistry.h"
34#include "llvm/Support/CommandLine.h"
35
36#include <cmath>
37#include <limits>
38
39#if defined(LLVM_HAVE_TFLITE)
40#include "llvm/Analysis/ModelUnderTrainingRunner.h"
41#include "llvm/Analysis/NoInferenceModelRunner.h"
42#include "llvm/Analysis/Utils/TrainingLogger.h"
43#include "llvm/IR/Module.h"
44#endif
45
46using namespace llvm;
47
48static cl::opt<std::string> InteractiveChannelBaseName(
49 "regalloc-priority-interactive-channel-base", cl::Hidden,
50 cl::desc(
51 "Base file path for the interactive mode. The incoming filename should "
52 "have the name <regalloc-priority-interactive-channel-base>.in, while "
53 "the outgoing name should be "
54 "<regalloc-priority-interactive-channel-base>.out"));
55
56using CompiledModelType = NoopSavedModelImpl;
57
58// Options that only make sense in development mode
59#ifdef LLVM_HAVE_TFLITE
60#include "RegAllocScore.h"
61#include "llvm/Analysis/Utils/TFUtils.h"
62
63static cl::opt<std::string> TrainingLog(
64 "regalloc-priority-training-log", cl::Hidden,
65 cl::desc("Training log for the register allocator priority model"));
66
67static cl::opt<std::string> ModelUnderTraining(
68 "regalloc-priority-model", cl::Hidden,
69 cl::desc("The model being trained for register allocation priority"));
70
71#endif // #ifdef LLVM_HAVE_TFLITE
72
73namespace llvm {
74
75static const std::vector<int64_t> PerLiveRangeShape{1};
76
77#define RA_PRIORITY_FEATURES_LIST(M) \
78 M(int64_t, li_size, PerLiveRangeShape, "size") \
79 M(int64_t, stage, PerLiveRangeShape, "stage") \
80 M(float, weight, PerLiveRangeShape, "weight")
81
82#define DecisionName "priority"
83static const TensorSpec DecisionSpec =
84 TensorSpec::createSpec<float>(DecisionName, Shape: {1});
85
86
87// Named features index.
88enum FeatureIDs {
89#define _FEATURE_IDX(_, name, __, ___) name,
90 RA_PRIORITY_FEATURES_LIST(_FEATURE_IDX)
91#undef _FEATURE_IDX
92 FeatureCount
93};
94
95class MLPriorityAdvisor : public RegAllocPriorityAdvisor {
96public:
97 MLPriorityAdvisor(const MachineFunction &MF, const RAGreedy &RA,
98 SlotIndexes *const Indexes, MLModelRunner *Runner);
99
100protected:
101 const RegAllocPriorityAdvisor &getDefaultAdvisor() const {
102 return static_cast<const RegAllocPriorityAdvisor &>(DefaultAdvisor);
103 }
104
105 // The assumption is that if the Runner could not be constructed, we emit-ed
106 // error, and we shouldn't be asking for it here.
107 const MLModelRunner &getRunner() const { return *Runner; }
108 float getPriorityImpl(const LiveInterval &LI) const;
109 unsigned getPriority(const LiveInterval &LI) const override;
110
111private:
112 const DefaultPriorityAdvisor DefaultAdvisor;
113 MLModelRunner *const Runner;
114};
115
116#define _DECL_FEATURES(type, name, shape, _) \
117 TensorSpec::createSpec<type>(#name, shape),
118
119static const std::vector<TensorSpec> InputFeatures{
120 {RA_PRIORITY_FEATURES_LIST(_DECL_FEATURES)},
121};
122#undef _DECL_FEATURES
123
124// ===================================
125// Release (AOT) - specifics
126// ===================================
127class ReleaseModePriorityAdvisorProvider final
128 : public RegAllocPriorityAdvisorProvider {
129public:
130 ReleaseModePriorityAdvisorProvider()
131 : RegAllocPriorityAdvisorProvider(AdvisorMode::Release) {}
132 std::unique_ptr<RegAllocPriorityAdvisor>
133 getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
134 SlotIndexes &SI) override {
135 if (!Runner) {
136 if (InteractiveChannelBaseName.empty())
137 Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
138 args&: MF.getFunction().getContext(), args: InputFeatures, DecisionName);
139 else
140 Runner = std::make_unique<InteractiveModelRunner>(
141 args&: MF.getFunction().getContext(), args: InputFeatures, args: DecisionSpec,
142 args: InteractiveChannelBaseName + ".out",
143 args: InteractiveChannelBaseName + ".in");
144 }
145 return std::make_unique<MLPriorityAdvisor>(args: MF, args: RA, args: &SI, args: Runner.get());
146 }
147
148private:
149 std::unique_ptr<MLModelRunner> Runner;
150};
151
152class ReleaseModePriorityAdvisorAnalysisLegacy final
153 : public RegAllocPriorityAdvisorAnalysisLegacy {
154public:
155 ReleaseModePriorityAdvisorAnalysisLegacy()
156 : RegAllocPriorityAdvisorAnalysisLegacy(AdvisorMode::Release) {}
157 // support for isa<> and dyn_cast.
158 static bool classof(const RegAllocPriorityAdvisorAnalysisLegacy *R) {
159 return R->getAdvisorMode() == AdvisorMode::Release;
160 }
161
162private:
163 void getAnalysisUsage(AnalysisUsage &AU) const override {
164 AU.setPreservesAll();
165 AU.addRequired<SlotIndexesWrapperPass>();
166 RegAllocPriorityAdvisorAnalysisLegacy::getAnalysisUsage(AU);
167 }
168
169 bool doInitialization(Module &M) override {
170 Provider = std::make_unique<ReleaseModePriorityAdvisorProvider>();
171 return false;
172 }
173};
174
175// ===================================
176// Development mode-specifics
177// ===================================
178//
179// Features we log
180#ifdef LLVM_HAVE_TFLITE
181static const TensorSpec Reward = TensorSpec::createSpec<float>("reward", {1});
182
183#define _DECL_TRAIN_FEATURES(type, name, shape, _) \
184 TensorSpec::createSpec<type>(std::string("action_") + #name, shape),
185
186static const std::vector<TensorSpec> TrainingInputFeatures{
187 {RA_PRIORITY_FEATURES_LIST(_DECL_TRAIN_FEATURES)
188 TensorSpec::createSpec<float>("action_discount", {1}),
189 TensorSpec::createSpec<int32_t>("action_step_type", {1}),
190 TensorSpec::createSpec<float>("action_reward", {1})}};
191#undef _DECL_TRAIN_FEATURES
192
193class DevelopmentModePriorityAdvisor : public MLPriorityAdvisor {
194public:
195 DevelopmentModePriorityAdvisor(const MachineFunction &MF, const RAGreedy &RA,
196 SlotIndexes *const Indexes,
197 MLModelRunner *Runner, Logger *Log)
198 : MLPriorityAdvisor(MF, RA, Indexes, Runner), Log(Log) {}
199
200private:
201 unsigned getPriority(const LiveInterval &LI) const override;
202 Logger *const Log;
203};
204
205class DevelopmentModePriorityAdvisorProvider final
206 : public RegAllocPriorityAdvisorProvider {
207
208public:
209 // Save all the logs (when requested).
210 DevelopmentModePriorityAdvisorProvider(LLVMContext &Ctx)
211 : RegAllocPriorityAdvisorProvider(AdvisorMode::Development) {
212 if (ModelUnderTraining.empty() && TrainingLog.empty()) {
213 Ctx.emitError("Regalloc development mode should be requested with at "
214 "least logging enabled and/or a training model");
215 return;
216 }
217 if (ModelUnderTraining.empty())
218 Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
219 else
220 Runner = ModelUnderTrainingRunner::createAndEnsureValid(
221 Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
222 if (!Runner) {
223 Ctx.emitError("Regalloc: could not set up the model runner");
224 return;
225 }
226 if (TrainingLog.empty())
227 return;
228 std::error_code EC;
229 auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
230 if (EC) {
231 Ctx.emitError(EC.message() + ":" + TrainingLog);
232 return;
233 }
234 std::vector<TensorSpec> LFS = InputFeatures;
235 if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
236 append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
237 // We always log the output; in particular, if we're not evaluating, we
238 // don't have an output spec json file. That's why we handle the
239 // 'normal' output separately.
240 LFS.push_back(DecisionSpec);
241
242 Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
243 /*IncludeReward*/ true);
244 }
245
246 void logRewardIfNeeded(const MachineFunction &MF,
247 llvm::function_ref<float()> GetReward) override {
248 if (!Log || !Log->hasAnyObservationForContext(MF.getName()))
249 return;
250 // The function pass manager would run all the function passes for a
251 // function, so we assume the last context belongs to this function. If
252 // this invariant ever changes, we can implement at that time switching
253 // contexts. At this point, it'd be an error
254 if (Log->currentContext() != MF.getName()) {
255 MF.getFunction().getContext().emitError(
256 "The training log context shouldn't have had changed.");
257 }
258 if (Log->hasObservationInProgress())
259 Log->logReward<float>(GetReward());
260 }
261
262 std::unique_ptr<RegAllocPriorityAdvisor>
263 getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
264 SlotIndexes &SI) override {
265 if (!Runner)
266 return nullptr;
267 if (Log) {
268 Log->switchContext(MF.getName());
269 }
270 return std::make_unique<DevelopmentModePriorityAdvisor>(
271 MF, RA, &SI, Runner.get(), Log.get());
272 }
273
274 std::unique_ptr<MLModelRunner> Runner;
275 std::unique_ptr<Logger> Log;
276};
277
278class DevelopmentModePriorityAdvisorAnalysisLegacy final
279 : public RegAllocPriorityAdvisorAnalysisLegacy {
280public:
281 DevelopmentModePriorityAdvisorAnalysisLegacy()
282 : RegAllocPriorityAdvisorAnalysisLegacy(AdvisorMode::Development) {}
283
284 // support for isa<> and dyn_cast.
285 static bool classof(const RegAllocPriorityAdvisorAnalysisLegacy *R) {
286 return R->getAdvisorMode() == AdvisorMode::Development;
287 }
288
289 void logRewardIfNeeded(const MachineFunction &MF,
290 llvm::function_ref<float()> GetReward) override {
291 Provider->logRewardIfNeeded(MF, GetReward);
292 }
293
294private:
295 void getAnalysisUsage(AnalysisUsage &AU) const override {
296 AU.setPreservesAll();
297 AU.addRequired<SlotIndexesWrapperPass>();
298 RegAllocPriorityAdvisorAnalysisLegacy::getAnalysisUsage(AU);
299 }
300
301 // Save all the logs (when requested).
302 bool doInitialization(Module &M) override {
303 Provider = std::make_unique<DevelopmentModePriorityAdvisorProvider>(
304 M.getContext());
305 return false;
306 ;
307 }
308};
309#endif //#ifdef LLVM_HAVE_TFLITE
310
311} // namespace llvm
312
313RegAllocPriorityAdvisorAnalysisLegacy *
314llvm::createReleaseModePriorityAdvisorAnalysis() {
315 return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
316 !InteractiveChannelBaseName.empty()
317 ? new ReleaseModePriorityAdvisorAnalysisLegacy()
318 : nullptr;
319}
320
321MLPriorityAdvisor::MLPriorityAdvisor(const MachineFunction &MF,
322 const RAGreedy &RA,
323 SlotIndexes *const Indexes,
324 MLModelRunner *Runner)
325 : RegAllocPriorityAdvisor(MF, RA, Indexes), DefaultAdvisor(MF, RA, Indexes),
326 Runner(std::move(Runner)) {
327 assert(this->Runner);
328 Runner->switchContext(Name: MF.getName());
329}
330
331// Converting a NaN or an out-of-range float advice to unsigned is undefined.
332// Saturate instead. A NaN is a model error, so also assert on it.
333static unsigned convertAdviceToPriority(double Advice) {
334 assert(!std::isnan(Advice) && "model produced a NaN priority");
335 if (!(Advice > 0.0))
336 return 0;
337 if (Advice >= static_cast<double>(std::numeric_limits<unsigned>::max()))
338 return std::numeric_limits<unsigned>::max();
339 return static_cast<unsigned>(Advice);
340}
341
342float MLPriorityAdvisor::getPriorityImpl(const LiveInterval &LI) const {
343 const unsigned Size = LI.getSize();
344 LiveRangeStage Stage = RA.getExtraInfo().getStage(VirtReg: LI);
345
346 *Runner->getTensor<int64_t>(FeatureID: 0) = static_cast<int64_t>(Size);
347 *Runner->getTensor<int64_t>(FeatureID: 1) = static_cast<int64_t>(Stage);
348 *Runner->getTensor<float>(FeatureID: 2) = static_cast<float>(LI.weight());
349
350 return Runner->evaluate<float>();
351}
352
353unsigned MLPriorityAdvisor::getPriority(const LiveInterval &LI) const {
354 return convertAdviceToPriority(Advice: getPriorityImpl(LI));
355}
356
357#ifdef LLVM_HAVE_TFLITE
358RegAllocPriorityAdvisorAnalysisLegacy *
359llvm::createDevelopmentModePriorityAdvisorAnalysis() {
360 return new DevelopmentModePriorityAdvisorAnalysisLegacy();
361}
362
363unsigned
364DevelopmentModePriorityAdvisor::getPriority(const LiveInterval &LI) const {
365 unsigned Prio = 0;
366
367 if (isa<ModelUnderTrainingRunner>(getRunner())) {
368 Prio = convertAdviceToPriority(MLPriorityAdvisor::getPriorityImpl(LI));
369 } else {
370 Prio = getDefaultAdvisor().getPriority(LI);
371 }
372
373 if (TrainingLog.empty())
374 return Prio;
375
376 // TODO(mtrofin): when we support optional rewards, this can go away. In the
377 // meantime, we log the "pretend" reward (0) for the previous observation
378 // before starting a new one.
379 if (Log->hasObservationInProgress())
380 Log->logReward<float>(0.0);
381
382 Log->startObservation();
383 size_t CurrentFeature = 0;
384 for (; CurrentFeature < InputFeatures.size(); ++CurrentFeature) {
385 Log->logTensorValue(CurrentFeature,
386 reinterpret_cast<const char *>(
387 getRunner().getTensorUntyped(CurrentFeature)));
388 }
389
390 if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner())) {
391 for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size();
392 ++I, ++CurrentFeature)
393 Log->logTensorValue(
394 CurrentFeature,
395 reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I)));
396 }
397
398 float Ret = static_cast<float>(Prio);
399 Log->logTensorValue(CurrentFeature, reinterpret_cast<const char *>(&Ret));
400 Log->endObservation();
401
402 return Prio;
403}
404
405RegAllocPriorityAdvisorProvider *
406llvm::createDevelopmentModePriorityAdvisorProvider(LLVMContext &Ctx) {
407 return new DevelopmentModePriorityAdvisorProvider(Ctx);
408}
409
410#endif // #ifdef LLVM_HAVE_TFLITE
411
412RegAllocPriorityAdvisorProvider *
413llvm::createReleaseModePriorityAdvisorProvider() {
414 return new ReleaseModePriorityAdvisorProvider();
415}
416