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, {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             MF.getFunction().getContext(), InputFeatures, DecisionName);
146       else
147         Runner = std::make_unique<InteractiveModelRunner>(
148             MF.getFunction().getContext(), InputFeatures, DecisionSpec,
149             InteractiveChannelBaseName + ".out",
150             InteractiveChannelBaseName + ".in");
151     }
152     return std::make_unique<MLPriorityAdvisor>(
153         MF, RA, &getAnalysis<SlotIndexes>(), 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(MF.getName());
293 }
294 
295 float MLPriorityAdvisor::getPriorityImpl(const LiveInterval &LI) const {
296   const unsigned Size = LI.getSize();
297   LiveRangeStage Stage = RA.getExtraInfo().getStage(LI);
298 
299   *Runner->getTensor<int64_t>(0) = static_cast<int64_t>(Size);
300   *Runner->getTensor<int64_t>(1) = static_cast<int64_t>(Stage);
301   *Runner->getTensor<float>(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