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
42using namespace llvm;
43
44static 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
52using 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
59static cl::opt<std::string> TrainingLog(
60 "regalloc-priority-training-log", cl::Hidden,
61 cl::desc("Training log for the register allocator priority model"));
62
63static 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
69namespace llvm {
70
71static 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"
79static const TensorSpec DecisionSpec =
80 TensorSpec::createSpec<float>(DecisionName, Shape: {1});
81
82
83// Named features index.
84enum FeatureIDs {
85#define _FEATURE_IDX(_, name, __, ___) name,
86 RA_PRIORITY_FEATURES_LIST(_FEATURE_IDX)
87#undef _FEATURE_IDX
88 FeatureCount
89};
90
91class MLPriorityAdvisor : public RegAllocPriorityAdvisor {
92public:
93 MLPriorityAdvisor(const MachineFunction &MF, const RAGreedy &RA,
94 SlotIndexes *const Indexes, MLModelRunner *Runner);
95
96protected:
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
107private:
108 const DefaultPriorityAdvisor DefaultAdvisor;
109 MLModelRunner *const Runner;
110};
111
112#define _DECL_FEATURES(type, name, shape, _) \
113 TensorSpec::createSpec<type>(#name, shape),
114
115static const std::vector<TensorSpec> InputFeatures{
116 {RA_PRIORITY_FEATURES_LIST(_DECL_FEATURES)},
117};
118#undef _DECL_FEATURES
119
120// ===================================
121// Release (AOT) - specifics
122// ===================================
123class ReleaseModePriorityAdvisorAnalysis final
124 : public RegAllocPriorityAdvisorAnalysis {
125public:
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
133private:
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
164static 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
169static 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
176class DevelopmentModePriorityAdvisor : public MLPriorityAdvisor {
177public:
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
183private:
184 unsigned getPriority(const LiveInterval &LI) const override;
185 Logger *const Log;
186};
187
188class DevelopmentModePriorityAdvisorAnalysis final
189 : public RegAllocPriorityAdvisorAnalysis {
190public:
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
214private:
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
278RegAllocPriorityAdvisorAnalysis *llvm::createReleaseModePriorityAdvisor() {
279 return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
280 !InteractiveChannelBaseName.empty()
281 ? new ReleaseModePriorityAdvisorAnalysis()
282 : nullptr;
283}
284
285MLPriorityAdvisor::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
295float 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
306unsigned MLPriorityAdvisor::getPriority(const LiveInterval &LI) const {
307 return static_cast<unsigned>(getPriorityImpl(LI));
308}
309
310#ifdef LLVM_HAVE_TFLITE
311RegAllocPriorityAdvisorAnalysis *llvm::createDevelopmentModePriorityAdvisor() {
312 return new DevelopmentModePriorityAdvisorAnalysis();
313}
314
315unsigned
316DevelopmentModePriorityAdvisor::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

source code of llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp