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