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 "RegAllocPriorityAdvisor.h"
16 #include "llvm/Analysis/AliasAnalysis.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/RegisterClassInfo.h"
28 #include "llvm/CodeGen/SlotIndexes.h"
29 #include "llvm/CodeGen/VirtRegMap.h"
30 #include "llvm/InitializePasses.h"
31 #include "llvm/Pass.h"
32 #include "llvm/PassRegistry.h"
33 #include "llvm/Support/CommandLine.h"
34
35 #if defined(LLVM_HAVE_TFLITE)
36 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
37 #include "llvm/Analysis/NoInferenceModelRunner.h"
38 #include "llvm/Analysis/Utils/TrainingLogger.h"
39 #endif
40
41 using namespace llvm;
42
43 // Options that only make sense in development mode
44 #ifdef LLVM_HAVE_TFLITE
45 #include "RegAllocScore.h"
46 #include "llvm/Analysis/Utils/TFUtils.h"
47
48 static cl::opt<std::string> TrainingLog(
49 "regalloc-priority-training-log", cl::Hidden,
50 cl::desc("Training log for the register allocator priority model"));
51
52 static cl::opt<std::string> ModelUnderTraining(
53 "regalloc-priority-model", cl::Hidden,
54 cl::desc("The model being trained for register allocation priority"));
55
56 #endif // #ifdef LLVM_HAVE_TFLITE
57
58 namespace llvm {
59
60 static const std::vector<int64_t> PerLiveRangeShape{1};
61
62 #define RA_PRIORITY_FEATURES_LIST(M) \
63 M(int64_t, li_size, PerLiveRangeShape, "size") \
64 M(int64_t, stage, PerLiveRangeShape, "stage") \
65 M(float, weight, PerLiveRangeShape, "weight")
66
67 #define DecisionName "priority"
68
69 // Named features index.
70 enum FeatureIDs {
71 #define _FEATURE_IDX(_, name, __, ___) name,
72 RA_PRIORITY_FEATURES_LIST(_FEATURE_IDX)
73 #undef _FEATURE_IDX
74 FeatureCount
75 };
76
77 class MLPriorityAdvisor : public RegAllocPriorityAdvisor {
78 public:
79 MLPriorityAdvisor(const MachineFunction &MF, const RAGreedy &RA,
80 SlotIndexes *const Indexes, MLModelRunner *Runner);
81
82 protected:
getDefaultAdvisor() const83 const RegAllocPriorityAdvisor &getDefaultAdvisor() const {
84 return static_cast<const RegAllocPriorityAdvisor &>(DefaultAdvisor);
85 }
86
87 // The assumption is that if the Runner could not be constructed, we emit-ed
88 // error, and we shouldn't be asking for it here.
getRunner() const89 const MLModelRunner &getRunner() const { return *Runner; }
90 float getPriorityImpl(const LiveInterval &LI) const;
91 unsigned getPriority(const LiveInterval &LI) const override;
92
93 private:
94 const DefaultPriorityAdvisor DefaultAdvisor;
95 MLModelRunner *const Runner;
96 };
97
98 #define _DECL_FEATURES(type, name, shape, _) \
99 TensorSpec::createSpec<type>(#name, shape),
100
101 static const std::vector<TensorSpec> InputFeatures{
102 {RA_PRIORITY_FEATURES_LIST(_DECL_FEATURES)},
103 };
104 #undef _DECL_FEATURES
105
106 // ===================================
107 // Release (AOT) - specifics
108 // ===================================
109 class ReleaseModePriorityAdvisorAnalysis final
110 : public RegAllocPriorityAdvisorAnalysis {
111 public:
ReleaseModePriorityAdvisorAnalysis()112 ReleaseModePriorityAdvisorAnalysis()
113 : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Release) {}
114 // support for isa<> and dyn_cast.
classof(const RegAllocPriorityAdvisorAnalysis * R)115 static bool classof(const RegAllocPriorityAdvisorAnalysis *R) {
116 return R->getAdvisorMode() == AdvisorMode::Release;
117 }
118
119 private:
getAnalysisUsage(AnalysisUsage & AU) const120 void getAnalysisUsage(AnalysisUsage &AU) const override {
121 AU.setPreservesAll();
122 AU.addRequired<SlotIndexes>();
123 RegAllocPriorityAdvisorAnalysis::getAnalysisUsage(AU);
124 }
125
126 std::unique_ptr<RegAllocPriorityAdvisor>
getAdvisor(const MachineFunction & MF,const RAGreedy & RA)127 getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
128 if (!Runner)
129 Runner = std::make_unique<ReleaseModeModelRunner<NoopSavedModelImpl>>(
130 MF.getFunction().getContext(), InputFeatures, DecisionName);
131 return std::make_unique<MLPriorityAdvisor>(
132 MF, RA, &getAnalysis<SlotIndexes>(), Runner.get());
133 }
134 std::unique_ptr<ReleaseModeModelRunner<NoopSavedModelImpl>> Runner;
135 };
136
137 // ===================================
138 // Development mode-specifics
139 // ===================================
140 //
141 // Features we log
142 #ifdef LLVM_HAVE_TFLITE
143
144 static const TensorSpec Output =
145 TensorSpec::createSpec<float>(DecisionName, {1});
146 static const TensorSpec Reward = TensorSpec::createSpec<float>("reward", {1});
147
148 #define _DECL_TRAIN_FEATURES(type, name, shape, _) \
149 TensorSpec::createSpec<type>(std::string("action_") + #name, shape),
150
151 static const std::vector<TensorSpec> TrainingInputFeatures{
152 {RA_PRIORITY_FEATURES_LIST(_DECL_TRAIN_FEATURES)
153 TensorSpec::createSpec<float>("action_discount", {1}),
154 TensorSpec::createSpec<int32_t>("action_step_type", {1}),
155 TensorSpec::createSpec<float>("action_reward", {1})}};
156 #undef _DECL_TRAIN_FEATURES
157
158 class DevelopmentModePriorityAdvisor : public MLPriorityAdvisor {
159 public:
DevelopmentModePriorityAdvisor(const MachineFunction & MF,const RAGreedy & RA,SlotIndexes * const Indexes,MLModelRunner * Runner,Logger * Log)160 DevelopmentModePriorityAdvisor(const MachineFunction &MF, const RAGreedy &RA,
161 SlotIndexes *const Indexes,
162 MLModelRunner *Runner, Logger *Log)
163 : MLPriorityAdvisor(MF, RA, Indexes, Runner), Log(Log) {}
164
165 private:
166 unsigned getPriority(const LiveInterval &LI) const override;
167 Logger *const Log;
168 };
169
170 class DevelopmentModePriorityAdvisorAnalysis final
171 : public RegAllocPriorityAdvisorAnalysis {
172 public:
DevelopmentModePriorityAdvisorAnalysis()173 DevelopmentModePriorityAdvisorAnalysis()
174 : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Development) {}
175 // support for isa<> and dyn_cast.
classof(const RegAllocPriorityAdvisorAnalysis * R)176 static bool classof(const RegAllocPriorityAdvisorAnalysis *R) {
177 return R->getAdvisorMode() == AdvisorMode::Development;
178 }
179
logRewardIfNeeded(const MachineFunction & MF,llvm::function_ref<float ()> GetReward)180 void logRewardIfNeeded(const MachineFunction &MF,
181 llvm::function_ref<float()> GetReward) override {
182 if (!Log)
183 return;
184 // The function pass manager would run all the function passes for a
185 // function, so we assume the last context belongs to this function. If
186 // this invariant ever changes, we can implement at that time switching
187 // contexts. At this point, it'd be an error
188 if (Log->currentContext() != MF.getName()) {
189 MF.getFunction().getContext().emitError(
190 "The training log context shouldn't have had changed.");
191 }
192 if (Log->hasObservationInProgress())
193 Log->logReward<float>(GetReward());
194 }
195
196 private:
getAnalysisUsage(AnalysisUsage & AU) const197 void getAnalysisUsage(AnalysisUsage &AU) const override {
198 AU.setPreservesAll();
199 AU.addRequired<SlotIndexes>();
200 RegAllocPriorityAdvisorAnalysis::getAnalysisUsage(AU);
201 }
202
203 // Save all the logs (when requested).
doInitialization(Module & M)204 bool doInitialization(Module &M) override {
205 LLVMContext &Ctx = M.getContext();
206 if (ModelUnderTraining.empty() && TrainingLog.empty()) {
207 Ctx.emitError("Regalloc development mode should be requested with at "
208 "least logging enabled and/or a training model");
209 return false;
210 }
211 if (ModelUnderTraining.empty())
212 Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
213 else
214 Runner = ModelUnderTrainingRunner::createAndEnsureValid(
215 Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
216 if (!Runner) {
217 Ctx.emitError("Regalloc: could not set up the model runner");
218 return false;
219 }
220 if (TrainingLog.empty())
221 return false;
222 std::error_code EC;
223 auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
224 if (EC) {
225 M.getContext().emitError(EC.message() + ":" + TrainingLog);
226 return false;
227 }
228 std::vector<TensorSpec> LFS = InputFeatures;
229 if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
230 append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
231 // We always log the output; in particular, if we're not evaluating, we
232 // don't have an output spec json file. That's why we handle the
233 // 'normal' output separately.
234 LFS.push_back(Output);
235
236 Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
237 /*IncludeReward*/ true);
238 return false;
239 }
240
241 std::unique_ptr<RegAllocPriorityAdvisor>
getAdvisor(const MachineFunction & MF,const RAGreedy & RA)242 getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
243 if (!Runner)
244 return nullptr;
245 if (Log) {
246 Log->switchContext(MF.getName());
247 }
248
249 return std::make_unique<DevelopmentModePriorityAdvisor>(
250 MF, RA, &getAnalysis<SlotIndexes>(), Runner.get(), Log.get());
251 }
252
253 std::unique_ptr<MLModelRunner> Runner;
254 std::unique_ptr<Logger> Log;
255 };
256 #endif //#ifdef LLVM_HAVE_TFLITE
257
258 } // namespace llvm
259
createReleaseModePriorityAdvisor()260 RegAllocPriorityAdvisorAnalysis *llvm::createReleaseModePriorityAdvisor() {
261 return new ReleaseModePriorityAdvisorAnalysis();
262 }
263
MLPriorityAdvisor(const MachineFunction & MF,const RAGreedy & RA,SlotIndexes * const Indexes,MLModelRunner * Runner)264 MLPriorityAdvisor::MLPriorityAdvisor(const MachineFunction &MF,
265 const RAGreedy &RA,
266 SlotIndexes *const Indexes,
267 MLModelRunner *Runner)
268 : RegAllocPriorityAdvisor(MF, RA, Indexes), DefaultAdvisor(MF, RA, Indexes),
269 Runner(std::move(Runner)) {
270 assert(this->Runner);
271 }
272
getPriorityImpl(const LiveInterval & LI) const273 float MLPriorityAdvisor::getPriorityImpl(const LiveInterval &LI) const {
274 const unsigned Size = LI.getSize();
275 LiveRangeStage Stage = RA.getExtraInfo().getStage(LI);
276
277 *Runner->getTensor<int64_t>(0) = static_cast<int64_t>(Size);
278 *Runner->getTensor<int64_t>(1) = static_cast<int64_t>(Stage);
279 *Runner->getTensor<float>(2) = static_cast<float>(LI.weight());
280
281 return Runner->evaluate<float>();
282 }
283
getPriority(const LiveInterval & LI) const284 unsigned MLPriorityAdvisor::getPriority(const LiveInterval &LI) const {
285 return static_cast<unsigned>(getPriorityImpl(LI));
286 }
287
288 #ifdef LLVM_HAVE_TFLITE
createDevelopmentModePriorityAdvisor()289 RegAllocPriorityAdvisorAnalysis *llvm::createDevelopmentModePriorityAdvisor() {
290 return new DevelopmentModePriorityAdvisorAnalysis();
291 }
292
293 unsigned
getPriority(const LiveInterval & LI) const294 DevelopmentModePriorityAdvisor::getPriority(const LiveInterval &LI) const {
295 double Prio = 0;
296
297 if (isa<ModelUnderTrainingRunner>(getRunner())) {
298 Prio = MLPriorityAdvisor::getPriorityImpl(LI);
299 } else {
300 Prio = getDefaultAdvisor().getPriority(LI);
301 }
302
303 if (TrainingLog.empty())
304 return Prio;
305
306 // TODO(mtrofin): when we support optional rewards, this can go away. In the
307 // meantime, we log the "pretend" reward (0) for the previous observation
308 // before starting a new one.
309 if (Log->hasObservationInProgress())
310 Log->logReward<float>(0.0);
311
312 Log->startObservation();
313 size_t CurrentFeature = 0;
314 for (; CurrentFeature < InputFeatures.size(); ++CurrentFeature) {
315 Log->logTensorValue(CurrentFeature,
316 reinterpret_cast<const char *>(
317 getRunner().getTensorUntyped(CurrentFeature)));
318 }
319
320 if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner())) {
321 for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size();
322 ++I, ++CurrentFeature)
323 Log->logTensorValue(
324 CurrentFeature,
325 reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I)));
326 }
327
328 float Ret = static_cast<float>(Prio);
329 Log->logTensorValue(CurrentFeature, reinterpret_cast<const char *>(&Ret));
330 Log->endObservation();
331
332 return static_cast<unsigned>(Prio);
333 }
334
335 #endif // #ifdef LLVM_HAVE_TFLITE
336