1//===- MLRegAllocEvictAdvisor.cpp - ML eviction 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 eviction advisor and reward injection pass
10//
11//===----------------------------------------------------------------------===//
12
13#include "AllocationOrder.h"
14#include "RegAllocGreedy.h"
15#include "llvm/Analysis/InteractiveModelRunner.h"
16#include "llvm/Analysis/MLModelRunner.h"
17#include "llvm/Analysis/TensorSpec.h"
18#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
19#if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL) || defined(LLVM_HAVE_TFLITE)
20#include "llvm/Analysis/ModelUnderTrainingRunner.h"
21#include "llvm/Analysis/NoInferenceModelRunner.h"
22#include "llvm/Analysis/Utils/TrainingLogger.h"
23#endif
24#include "MLRegAllocEvictAdvisor.h"
25#include "llvm/Analysis/ReleaseModeModelRunner.h"
26#include "llvm/CodeGen/CalcSpillWeights.h"
27#include "llvm/CodeGen/LiveRegMatrix.h"
28#include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
29#include "llvm/CodeGen/MachineFunction.h"
30#include "llvm/CodeGen/MachineLoopInfo.h"
31#include "llvm/CodeGen/MachineRegisterInfo.h"
32#include "llvm/CodeGen/Passes.h"
33#include "llvm/CodeGen/RegisterClassInfo.h"
34#include "llvm/CodeGen/VirtRegMap.h"
35#include "llvm/IR/Module.h"
36#include "llvm/InitializePasses.h"
37#include "llvm/Pass.h"
38#include "llvm/PassRegistry.h"
39#include "llvm/Support/CommandLine.h"
40#include "llvm/Support/ErrorHandling.h"
41
42#include <array>
43#include <bitset>
44#include <memory>
45
46using namespace llvm;
47
48#define DEBUG_TYPE "ml-regalloc"
49
50// Generated header in release (AOT) mode
51#if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL)
52#include "RegAllocEvictModel.h"
53using CompiledModelType = RegAllocEvictModel;
54#else
55using CompiledModelType = NoopSavedModelImpl;
56#endif
57
58static cl::opt<std::string> InteractiveChannelBaseName(
59 "regalloc-evict-interactive-channel-base", cl::Hidden,
60 cl::desc(
61 "Base file path for the interactive mode. The incoming filename should "
62 "have the name <regalloc-evict-interactive-channel-base>.in, while the "
63 "outgoing name should be "
64 "<regalloc-evict-interactive-channel-base>.out"));
65
66static cl::opt<unsigned> MaxEvictionCount(
67 "mlregalloc-max-eviction-count", cl::Hidden,
68 cl::desc("The maximum number of times a live range can be "
69 "evicted before preventing it from being evicted"),
70 cl::init(Val: 100));
71
72// Options that only make sense in development mode
73#ifdef LLVM_HAVE_TFLITE
74#include "RegAllocScore.h"
75#include "llvm/Analysis/Utils/TFUtils.h"
76
77static cl::opt<std::string> TrainingLog(
78 "regalloc-training-log", cl::Hidden,
79 cl::desc("Training log for the register allocator eviction model"));
80
81static cl::opt<std::string> ModelUnderTraining(
82 "regalloc-model", cl::Hidden,
83 cl::desc("The model being trained for register allocation eviction"));
84
85#endif // #ifdef LLVM_HAVE_TFLITE
86
87/// The score injection pass.
88/// This pass calculates the score for a function and inserts it in the log, but
89/// this happens only in development mode. It's a no-op otherwise.
90namespace llvm {
91extern cl::opt<unsigned> EvictInterferenceCutoff;
92} // namespace llvm
93
94namespace {
95class RegAllocScoring : public MachineFunctionPass {
96public:
97 static char ID;
98
99 RegAllocScoring() : MachineFunctionPass(ID) {}
100
101 ~RegAllocScoring() override = default;
102
103 StringRef getPassName() const override {
104 return "Register Allocation Pass Scoring";
105 }
106
107 /// RegAllocReward analysis usage.
108 void getAnalysisUsage(AnalysisUsage &AU) const override {
109 AU.setPreservesAll();
110 AU.addRequired<RegAllocEvictionAdvisorAnalysisLegacy>();
111 AU.addRequired<RegAllocPriorityAdvisorAnalysisLegacy>();
112 AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
113 MachineFunctionPass::getAnalysisUsage(AU);
114 }
115
116 /// Performs this pass
117 bool runOnMachineFunction(MachineFunction &) override;
118};
119} // namespace
120
121char RegAllocScoring::ID = 0;
122FunctionPass *llvm::createRegAllocScoringPass() {
123 return new RegAllocScoring();
124}
125
126INITIALIZE_PASS(RegAllocScoring, "regallocscoringpass",
127 "Register Allocation Scoring Pass", false, false)
128
129// ===================================
130// Common ML Advisor declarations
131// ===================================
132namespace {
133// Most features are as described above, so we'll reuse this vector in defining
134// them.
135static const std::vector<int64_t> PerLiveRangeShape{1, NumberOfInterferences};
136
137// --------------
138// Features table
139// --------------
140// For each interfering live range (incl. the candidate) we collect a number of
141// features. However, because the features are of different types (and because
142// of ML best practices), we organize the tensors per feature, not per
143// candidate. Each such tensor has a scalar value corresponding to the
144// interferring live range at that position, in the order in AllocationOrder.
145// The last position corresponds to the virt reg seeking allocation.
146// Exception to all that is the progression feature, which is just a scalar (see
147// its documentation for details).
148// Note on naming: the "_by_max" are normalized using the largest value of that
149// tensor, as observed in the current decision making stage (i.e. for the
150// current call to the advisor's tryFindEvictionCandidate)
151//
152// The feature list format: type, name, shape, documentation.
153// Note: we can really just use int64 and float, hence the modeling of some
154// bools as int64 values.
155#define RA_EVICT_FEATURES_LIST(M) \
156 M(int64_t, mask, PerLiveRangeShape, \
157 "boolean values, 0 for unavailable candidates (i.e. if a position is 0, " \
158 "it " \
159 "can't be evicted)") \
160 M(int64_t, is_free, PerLiveRangeShape, \
161 "boolean values, 1 if this phys reg is actually free (no interferences)") \
162 M(float, nr_urgent, PerLiveRangeShape, \
163 "number of 'urgent' intervals, normalized. Urgent are those that are OK " \
164 "to break cascades") \
165 M(float, nr_broken_hints, PerLiveRangeShape, \
166 "if this position were evicted, how many broken hints would there be") \
167 M(int64_t, is_hint, PerLiveRangeShape, \
168 "is this a preferred phys reg for the candidate") \
169 M(int64_t, is_local, PerLiveRangeShape, \
170 "is this live range local to a basic block") \
171 M(float, nr_rematerializable, PerLiveRangeShape, \
172 "nr rematerializable ranges") \
173 M(float, nr_defs_and_uses, PerLiveRangeShape, \
174 "bb freq - weighed nr defs and uses") \
175 M(float, weighed_reads_by_max, PerLiveRangeShape, \
176 "bb freq - weighed nr of reads, normalized") \
177 M(float, weighed_writes_by_max, PerLiveRangeShape, \
178 "bb feq - weighed nr of writes, normalized") \
179 M(float, weighed_read_writes_by_max, PerLiveRangeShape, \
180 "bb freq - weighed nr of uses that are both read and writes, normalized") \
181 M(float, weighed_indvars_by_max, PerLiveRangeShape, \
182 "bb freq - weighed nr of uses that are indvars, normalized") \
183 M(float, hint_weights_by_max, PerLiveRangeShape, \
184 "bb freq - weighed nr of uses that are hints, normalized") \
185 M(float, start_bb_freq_by_max, PerLiveRangeShape, \
186 "the freq in the start block, normalized") \
187 M(float, end_bb_freq_by_max, PerLiveRangeShape, \
188 "freq of end block, normalized") \
189 M(float, hottest_bb_freq_by_max, PerLiveRangeShape, \
190 "hottest BB freq, normalized") \
191 M(float, liverange_size, PerLiveRangeShape, \
192 "size (instr index diff) of the LR") \
193 M(float, use_def_density, PerLiveRangeShape, \
194 "the max weight, as computed by the manual heuristic") \
195 M(int64_t, max_stage, PerLiveRangeShape, \
196 "largest stage of an interval in this LR") \
197 M(int64_t, min_stage, PerLiveRangeShape, \
198 "lowest stage of an interval in this LR") \
199 M(float, progress, {1}, "ratio of current queue size to initial size")
200
201// The model learns to pick one of the mask == 1 interferences. This is the
202// name of the output tensor. The contract with the model is that the output
203// will be guaranteed to be to a mask == 1 position. Using a macro here to
204// avoid 'not used' warnings (and keep cond compilation to a minimum)
205#define DecisionName "index_to_evict"
206static const TensorSpec DecisionSpec =
207 TensorSpec::createSpec<int64_t>(DecisionName, Shape: {1});
208
209// Named features index.
210enum FeatureIDs {
211#define _FEATURE_IDX_SIMPLE(_, name, __, ___) name
212#define _FEATURE_IDX(A, B, C, D) _FEATURE_IDX_SIMPLE(A, B, C, D),
213 RA_EVICT_FEATURES_LIST(_FEATURE_IDX) FeatureCount,
214#undef _FEATURE_IDX
215#undef _FEATURE_IDX_SIMPLE
216};
217
218// The ML advisor will typically have a sparse input to the evaluator, because
219// various phys regs won't be available. It's easier (maintenance-wise) to
220// bulk-reset the state of the evaluator each time we are about to use it
221// again.
222template <typename T> size_t getTotalSize(const std::vector<int64_t> &Shape) {
223 size_t Ret = sizeof(T);
224 for (const auto V : Shape)
225 Ret *= V;
226 return Ret;
227}
228
229void resetInputs(MLModelRunner &Runner) {
230#define _RESET(TYPE, NAME, SHAPE, __) \
231 std::memset(Runner.getTensorUntyped(FeatureIDs::NAME), 0, \
232 getTotalSize<TYPE>(SHAPE));
233 RA_EVICT_FEATURES_LIST(_RESET)
234#undef _RESET
235}
236
237// Per-live interval components that get aggregated into the feature values
238// that will be passed to the evaluator.
239struct LIFeatureComponents {
240 double R = 0;
241 double W = 0;
242 double RW = 0;
243 double IndVarUpdates = 0;
244 double HintWeights = 0.0;
245 int64_t NumDefsAndUses = 0;
246 float HottestBlockFreq = 0.0;
247 bool IsRemat = false;
248};
249
250using CandidateRegList =
251 std::array<std::pair<MCRegister, bool>, NumberOfInterferences>;
252using FeaturesListNormalizer =
253 llvm::SmallVector<float, FeatureIDs::FeatureCount>;
254
255/// The ML evictor (commonalities between release and development mode)
256class MLEvictAdvisor : public RegAllocEvictionAdvisor {
257public:
258 MLEvictAdvisor(const MachineFunction &MF, const RAGreedy &RA,
259 MLModelRunner *Runner, const MachineBlockFrequencyInfo &MBFI,
260 const MachineLoopInfo &Loops);
261
262protected:
263 const RegAllocEvictionAdvisor &getDefaultAdvisor() const {
264 return static_cast<const RegAllocEvictionAdvisor &>(DefaultAdvisor);
265 }
266
267 // The assumption is that if the Runner could not be constructed, we emit-ed
268 // error, and we shouldn't be asking for it here.
269 const MLModelRunner &getRunner() const { return *Runner; }
270
271 /// This just calls Evaluate on the Runner, but in the development mode
272 /// case, if we're just capturing the log of the default advisor, it needs
273 /// to call the latter instead, so we need to pass all the necessary
274 /// parameters for it. In the development case, it will also log.
275 virtual int64_t
276 tryFindEvictionCandidatePosition(const LiveInterval &VirtReg,
277 const AllocationOrder &Order,
278 unsigned OrderLimit, uint8_t CostPerUseLimit,
279 const SmallVirtRegSet &FixedRegisters) const;
280
281 /// Load the features of the given VirtReg (allocated or not) at column Pos,
282 /// but if that can't be evicted, return false instead.
283 bool
284 loadInterferenceFeatures(const LiveInterval &VirtReg, MCRegister PhysReg,
285 bool IsHint, const SmallVirtRegSet &FixedRegisters,
286 llvm::SmallVectorImpl<float> &Largest, size_t Pos,
287 SmallVectorImpl<LRStartEndInfo> &LRPosInfo) const;
288
289private:
290 static float getInitialQueueSize(const MachineFunction &MF);
291
292 MCRegister tryFindEvictionCandidate(
293 const LiveInterval &VirtReg, const AllocationOrder &Order,
294 uint8_t CostPerUseLimit,
295 const SmallVirtRegSet &FixedRegisters) const override;
296
297 void extractFeatures(const SmallVectorImpl<const LiveInterval *> &Intervals,
298 llvm::SmallVectorImpl<float> &Largest, size_t Pos,
299 int64_t IsHint, int64_t LocalIntfsCount, float NumUrgent,
300 SmallVectorImpl<LRStartEndInfo> &LRPosInfo) const;
301
302 // Point-in-time: we didn't learn this, so we always delegate to the
303 // default.
304 bool canEvictHintInterference(
305 const LiveInterval &VirtReg, MCRegister PhysReg,
306 const SmallVirtRegSet &FixedRegisters) const override {
307 return getDefaultAdvisor().canEvictHintInterference(VirtReg, PhysReg,
308 FixedRegisters);
309 }
310
311 const LIFeatureComponents &
312 getLIFeatureComponents(const LiveInterval &LI) const;
313
314 // Hold on to a default advisor for:
315 // 1) the implementation of canEvictHintInterference, because we didn't
316 // learn that nuance yet; 2) for bootstrapping (logging) in the development
317 // mode case.
318 const DefaultEvictionAdvisor DefaultAdvisor;
319 MLModelRunner *const Runner;
320 const MachineBlockFrequencyInfo &MBFI;
321 const MachineLoopInfo &Loops;
322
323 // Indices of those features we don't want to normalize.
324 // This could be static and shared, but its initialization is non-trivial.
325 std::bitset<FeatureIDs::FeatureCount> DoNotNormalize;
326 const float InitialQSize;
327
328 using RegID = unsigned;
329 mutable DenseMap<RegID, LIFeatureComponents> CachedFeatures;
330
331 mutable DenseMap<unsigned, unsigned> VirtRegEvictionCounts;
332
333 void onEviction(Register RegBeingEvicted) const {
334 // If we cannot find the virtual register in the map, we just assume it has
335 // not been evicted before and thus has a value of zero (which is what the
336 // subscript operator returns by default).
337 ++VirtRegEvictionCounts[RegBeingEvicted.id()];
338 }
339
340 unsigned getEvictionCount(Register Reg) const {
341 auto EvictionCountIt = VirtRegEvictionCounts.find(Val: Reg.id());
342 if (EvictionCountIt != VirtRegEvictionCounts.end())
343 return EvictionCountIt->second;
344 return 0;
345 }
346};
347
348#define _DECL_FEATURES(type, name, shape, _) \
349 TensorSpec::createSpec<type>(#name, shape),
350
351// ===================================
352// Release (AOT) - specifics
353// ===================================
354/// Common provider for legacy and new pass managers.
355class ReleaseModeEvictionAdvisorProvider final
356 : public RegAllocEvictionAdvisorProvider {
357public:
358 ReleaseModeEvictionAdvisorProvider(LLVMContext &Ctx)
359 : RegAllocEvictionAdvisorProvider(AdvisorMode::Release, Ctx) {
360 InputFeatures = {RA_EVICT_FEATURES_LIST(_DECL_FEATURES)};
361 }
362 // support for isa<> and dyn_cast.
363 static bool classof(const RegAllocEvictionAdvisorProvider *R) {
364 return R->getAdvisorMode() == AdvisorMode::Release;
365 }
366
367 std::unique_ptr<RegAllocEvictionAdvisor>
368 getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
369 MachineBlockFrequencyInfo *MBFI, MachineLoopInfo *Loops) override {
370 if (!Runner) {
371 if (InteractiveChannelBaseName.empty())
372 Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
373 args&: MF.getFunction().getContext(), args&: InputFeatures, DecisionName);
374 else
375 Runner = std::make_unique<InteractiveModelRunner>(
376 args&: MF.getFunction().getContext(), args&: InputFeatures, args: DecisionSpec,
377 args: InteractiveChannelBaseName + ".out",
378 args: InteractiveChannelBaseName + ".in");
379 }
380 assert(MBFI && Loops &&
381 "Invalid provider state: must have analysis available");
382 return std::make_unique<MLEvictAdvisor>(args: MF, args: RA, args: Runner.get(), args&: *MBFI,
383 args&: *Loops);
384 }
385
386private:
387 std::vector<TensorSpec> InputFeatures;
388 std::unique_ptr<MLModelRunner> Runner;
389};
390
391class ReleaseModeEvictionAdvisorAnalysisLegacy final
392 : public RegAllocEvictionAdvisorAnalysisLegacy {
393public:
394 ReleaseModeEvictionAdvisorAnalysisLegacy()
395 : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Release) {}
396
397 void logRewardIfNeeded(const MachineFunction &MF,
398 llvm::function_ref<float()> GetReward) override {
399 // No-op in release mode
400 }
401
402 bool doInitialization(Module &M) override {
403 Provider =
404 std::make_unique<ReleaseModeEvictionAdvisorProvider>(args&: M.getContext());
405 return false;
406 }
407
408 static bool classof(const RegAllocEvictionAdvisorAnalysisLegacy *R) {
409 return R->getAdvisorMode() == AdvisorMode::Release;
410 }
411
412 void getAnalysisUsage(AnalysisUsage &AU) const override {
413 AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
414 AU.addRequired<MachineLoopInfoWrapperPass>();
415 RegAllocEvictionAdvisorAnalysisLegacy::getAnalysisUsage(AU);
416 }
417};
418
419// ===================================
420// Development mode-specifics
421// ===================================
422//
423// Features we log
424#ifdef LLVM_HAVE_TFLITE
425static const TensorSpec Reward = TensorSpec::createSpec<float>("reward", {1});
426
427// Features we bind on the model. The tensor names have a prefix, and we also
428// need to include some tensors that are expected to be present by the
429// training algo.
430// TODO: can we just get rid of these?
431#define _DECL_TRAIN_FEATURES(type, name, shape, _) \
432 TensorSpec::createSpec<type>(std::string("action_") + #name, shape),
433
434class DevelopmentModeEvictAdvisor : public MLEvictAdvisor {
435public:
436 DevelopmentModeEvictAdvisor(const MachineFunction &MF, const RAGreedy &RA,
437 MLModelRunner *Runner,
438 const MachineBlockFrequencyInfo &MBFI,
439 const MachineLoopInfo &Loops, Logger *Log)
440 : MLEvictAdvisor(MF, RA, Runner, MBFI, Loops), Log(Log) {}
441
442private:
443 int64_t tryFindEvictionCandidatePosition(
444 const LiveInterval &VirtReg, const AllocationOrder &Order,
445 unsigned OrderLimit, uint8_t CostPerUseLimit,
446 const SmallVirtRegSet &FixedRegisters) const override;
447
448 Logger *const Log;
449};
450
451class DevelopmentModeEvictionAdvisorProvider final
452 : public RegAllocEvictionAdvisorProvider {
453public:
454 DevelopmentModeEvictionAdvisorProvider(LLVMContext &Ctx)
455 : RegAllocEvictionAdvisorProvider(AdvisorMode::Development, Ctx) {
456 InputFeatures = {RA_EVICT_FEATURES_LIST(_DECL_FEATURES)};
457 TrainingInputFeatures = {
458 RA_EVICT_FEATURES_LIST(_DECL_TRAIN_FEATURES)
459 TensorSpec::createSpec<float>("action_discount", {1}),
460 TensorSpec::createSpec<int32_t>("action_step_type", {1}),
461 TensorSpec::createSpec<float>("action_reward", {1})};
462 if (ModelUnderTraining.empty() && TrainingLog.empty()) {
463 Ctx.emitError("Regalloc development mode should be requested with at "
464 "least logging enabled and/or a training model");
465 return;
466 }
467 if (ModelUnderTraining.empty())
468 Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
469 else
470 Runner = ModelUnderTrainingRunner::createAndEnsureValid(
471 Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
472 if (!Runner) {
473 Ctx.emitError("Regalloc: could not set up the model runner");
474 return;
475 }
476 if (TrainingLog.empty())
477 return;
478 std::error_code EC;
479 auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
480 if (EC) {
481 Ctx.emitError(EC.message() + ":" + TrainingLog);
482 return;
483 }
484 std::vector<TensorSpec> LFS = InputFeatures;
485 if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
486 append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
487 // We always log the output; in particular, if we're not evaluating, we
488 // don't have an output spec json file. That's why we handle the
489 // 'normal' output separately.
490 LFS.push_back(DecisionSpec);
491
492 Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
493 /*IncludeReward*/ true);
494 return;
495 }
496
497 // support for isa<> and dyn_cast.
498 static bool classof(const RegAllocEvictionAdvisorProvider *R) {
499 return R->getAdvisorMode() == AdvisorMode::Development;
500 }
501
502 void logRewardIfNeeded(const MachineFunction &MF,
503 llvm::function_ref<float()> GetReward) override {
504 if (!Log || !Log->hasAnyObservationForContext(MF.getName()))
505 return;
506 // The function pass manager would run all the function passes for a
507 // function, so we assume the last context belongs to this function. If
508 // this invariant ever changes, we can implement at that time switching
509 // contexts. At this point, it'd be an error
510 if (Log->currentContext() != MF.getName()) {
511 MF.getFunction().getContext().emitError(
512 "The training log context shouldn't have had changed.");
513 }
514 if (Log->hasObservationInProgress())
515 Log->logReward<float>(GetReward());
516 }
517
518 std::unique_ptr<RegAllocEvictionAdvisor>
519 getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
520 MachineBlockFrequencyInfo *MBFI, MachineLoopInfo *Loops) override {
521 if (!Runner)
522 return nullptr;
523 if (Log)
524 Log->switchContext(MF.getName());
525 assert(MBFI && Loops &&
526 "Invalid provider state: must have analysis available");
527 return std::make_unique<DevelopmentModeEvictAdvisor>(
528 MF, RA, Runner.get(), *MBFI, *Loops, Log.get());
529 }
530
531private:
532 std::vector<TensorSpec> InputFeatures;
533 std::vector<TensorSpec> TrainingInputFeatures;
534
535 std::unique_ptr<MLModelRunner> Runner;
536 std::unique_ptr<Logger> Log;
537};
538
539class DevelopmentModeEvictionAdvisorAnalysisLegacy final
540 : public RegAllocEvictionAdvisorAnalysisLegacy {
541public:
542 DevelopmentModeEvictionAdvisorAnalysisLegacy()
543 : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Development) {}
544
545 bool doInitialization(Module &M) override {
546 Provider = std::make_unique<DevelopmentModeEvictionAdvisorProvider>(
547 M.getContext());
548 return false;
549 }
550
551 void logRewardIfNeeded(const MachineFunction &MF,
552 llvm::function_ref<float()> GetReward) override {
553 Provider->logRewardIfNeeded(MF, GetReward);
554 }
555
556 // support for isa<> and dyn_cast.
557 static bool classof(const RegAllocEvictionAdvisorAnalysisLegacy *R) {
558 return R->getAdvisorMode() == AdvisorMode::Development;
559 }
560
561 void getAnalysisUsage(AnalysisUsage &AU) const override {
562 AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
563 AU.addRequired<MachineLoopInfoWrapperPass>();
564 RegAllocEvictionAdvisorAnalysisLegacy::getAnalysisUsage(AU);
565 }
566};
567
568#endif // #ifdef LLVM_HAVE_TFLITE
569} // namespace
570
571float MLEvictAdvisor::getInitialQueueSize(const MachineFunction &MF) {
572 auto &MRI = MF.getRegInfo();
573 unsigned NumUsedRegs = 0;
574 for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {
575 Register Reg = Register::index2VirtReg(Index: I);
576 if (!MRI.reg_nodbg_empty(RegNo: Reg))
577 ++NumUsedRegs;
578 }
579 return static_cast<float>(NumUsedRegs);
580}
581
582MLEvictAdvisor::MLEvictAdvisor(const MachineFunction &MF, const RAGreedy &RA,
583 MLModelRunner *Runner,
584 const MachineBlockFrequencyInfo &MBFI,
585 const MachineLoopInfo &Loops)
586 : RegAllocEvictionAdvisor(MF, RA), DefaultAdvisor(MF, RA),
587 Runner(std::move(Runner)), MBFI(MBFI), Loops(Loops),
588 InitialQSize(MLEvictAdvisor::getInitialQueueSize(MF)) {
589 assert(this->Runner);
590 Runner->switchContext(Name: MF.getName());
591 DoNotNormalize.set(position: FeatureIDs::mask);
592 DoNotNormalize.set(position: FeatureIDs::is_free);
593 DoNotNormalize.set(position: FeatureIDs::is_hint);
594 DoNotNormalize.set(position: FeatureIDs::is_local);
595 DoNotNormalize.set(position: FeatureIDs::min_stage);
596 DoNotNormalize.set(position: FeatureIDs::max_stage);
597 DoNotNormalize.set(position: FeatureIDs::progress);
598}
599
600int64_t MLEvictAdvisor::tryFindEvictionCandidatePosition(
601 const LiveInterval &, const AllocationOrder &, unsigned, uint8_t,
602 const SmallVirtRegSet &) const {
603 int64_t Ret = Runner->evaluate<int64_t>();
604 assert(Ret >= 0);
605 assert(Ret <= CandidateVirtRegPos);
606 return Ret;
607}
608
609bool MLEvictAdvisor::loadInterferenceFeatures(
610 const LiveInterval &VirtReg, MCRegister PhysReg, bool IsHint,
611 const SmallVirtRegSet &FixedRegisters,
612 llvm::SmallVectorImpl<float> &Largest, size_t Pos,
613 llvm::SmallVectorImpl<LRStartEndInfo> &LRPosInfo) const {
614 // It is only possible to evict virtual register interference.
615 if (Matrix->checkInterference(VirtReg, PhysReg) > LiveRegMatrix::IK_VirtReg) {
616 // leave unavailable
617 return false;
618 }
619
620 const bool IsLocal = LIS->intervalIsInOneMBB(LI: VirtReg);
621 int64_t LocalIntfs = 0;
622 float NumUrgent = 0.0f;
623
624 // The cascade tracking is the same as in the default advisor
625 unsigned Cascade = RA.getExtraInfo().getCascadeOrCurrentNext(Reg: VirtReg.reg());
626
627 SmallVector<const LiveInterval *, MaxInterferences> InterferingIntervals;
628 for (MCRegUnit Unit : TRI->regunits(Reg: PhysReg)) {
629 LiveIntervalUnion::Query &Q = Matrix->query(LR: VirtReg, RegUnit: Unit);
630 // Different from the default heuristic, we don't make any assumptions
631 // about what having more than 10 results in the query may mean.
632 const auto &IFIntervals = Q.interferingVRegs(MaxInterferingRegs: EvictInterferenceCutoff);
633 if (IFIntervals.empty() && InterferingIntervals.empty())
634 continue;
635 if (IFIntervals.size() >= EvictInterferenceCutoff)
636 return false;
637 InterferingIntervals.append(in_start: IFIntervals.begin(), in_end: IFIntervals.end());
638 for (const LiveInterval *Intf : reverse(C: IFIntervals)) {
639 assert(Intf->reg().isVirtual() &&
640 "Only expecting virtual register interference from query");
641 // This is the same set of legality checks as in the default case: don't
642 // try to evict fixed regs or 'done' ones. Also don't break cascades,
643 // except in the urgent case, with the same nuances used in the default
644 // heuristic.
645 // We could try sharing this between the advisors, but it may end up
646 // more complex than it is right now.
647 if (FixedRegisters.count(V: Intf->reg()))
648 return false;
649 if (RA.getExtraInfo().getStage(VirtReg: *Intf) == RS_Done)
650 return false;
651 bool Urgent =
652 !VirtReg.isSpillable() &&
653 (Intf->isSpillable() ||
654 RegClassInfo.getNumAllocatableRegs(RC: MRI->getRegClass(Reg: VirtReg.reg())) <
655 RegClassInfo.getNumAllocatableRegs(
656 RC: MRI->getRegClass(Reg: Intf->reg())));
657
658 unsigned IntfCascade = RA.getExtraInfo().getCascade(Reg: Intf->reg());
659 // There is a potential that the model could be adversarial and
660 // continually evict live ranges over and over again, leading to a
661 // large amount of compile time being spent in regalloc. If we hit the
662 // threshold, prevent the range from being evicted. We still let the
663 // range through if it is urgent as we are required to produce an
664 // eviction if the candidate is not spillable.
665 if (getEvictionCount(Reg: Intf->reg()) > MaxEvictionCount && !Urgent)
666 return false;
667
668 // Only evict older cascades or live ranges without a cascade.
669 if (Cascade <= IntfCascade) {
670 if (!Urgent)
671 return false;
672 ++NumUrgent;
673 }
674
675 LocalIntfs += (IsLocal && LIS->intervalIsInOneMBB(LI: *Intf) &&
676 (!EnableLocalReassign || !canReassign(VirtReg: *Intf, FromReg: PhysReg)));
677 }
678 }
679 // OK, so if we made it this far, this LR is an eviction candidate, load its
680 // features.
681 extractFeatures(Intervals: InterferingIntervals, Largest, Pos, IsHint, LocalIntfsCount: LocalIntfs,
682 NumUrgent, LRPosInfo);
683 return true;
684}
685
686MCRegister MLEvictAdvisor::tryFindEvictionCandidate(
687 const LiveInterval &VirtReg, const AllocationOrder &Order,
688 uint8_t CostPerUseLimit, const SmallVirtRegSet &FixedRegisters) const {
689 auto MaybeOrderLimit = getOrderLimit(VirtReg, Order, CostPerUseLimit);
690 if (!MaybeOrderLimit)
691 return MCRegister::NoRegister;
692 unsigned OrderLimit = *MaybeOrderLimit;
693
694 // The heuristic sets initial costs such as, if CostPerUseLimit is
695 // max<uint8_t>, then any of the costs of the legally-evictable intervals
696 // would be lower. When that happens, one of those will be selected.
697 // Therefore, we allow the candidate be selected, unless the candidate is
698 // unspillable, in which case it would be incorrect to not find a register
699 // for it.
700 const bool MustFindEviction =
701 (!VirtReg.isSpillable() && CostPerUseLimit == static_cast<uint8_t>(~0u));
702 // Number of available candidates - if 0, no need to continue.
703 size_t Available = 0;
704 // Make sure we don't have leftover partial state from an attempt where we
705 // had no available candidates and bailed out early.
706 resetInputs(Runner&: *Runner);
707
708 // Track the index->register mapping because AllocationOrder doesn't do that
709 // and we'd have to scan it.
710 // Also track their mask, to write asserts/debug.
711 CandidateRegList Regs;
712 Regs.fill(u: {0, false});
713
714 // Track the largest value of features seen during this eviction session. We
715 // only normalize (some of) the float features, but it's just simpler to
716 // dimension 'Largest' to all the features, especially since we have the
717 // 'DoNotNormalize' list.
718 FeaturesListNormalizer Largest(FeatureIDs::FeatureCount, 0.0);
719
720 // Same overal idea as in the default eviction policy - we visit the values
721 // of AllocationOrder one at a time. If it's not legally available, we mask
722 // off the corresponding feature column (==do nothing because we already
723 // reset all the features to 0) Use Pos to capture the column we load
724 // features at - in AllocationOrder order.
725 size_t Pos = 0;
726 SmallVector<LRStartEndInfo, NumberOfInterferences> LRPosInfo;
727 for (auto I = Order.begin(), E = Order.getOrderLimitEnd(OrderLimit); I != E;
728 ++I, ++Pos) {
729 MCRegister PhysReg = *I;
730 assert(!Regs[Pos].second);
731 assert(PhysReg);
732 if (!canAllocatePhysReg(CostPerUseLimit, PhysReg)) {
733 continue;
734 }
735 if (loadInterferenceFeatures(VirtReg, PhysReg, IsHint: I.isHint(), FixedRegisters,
736 Largest, Pos, LRPosInfo)) {
737 ++Available;
738 Regs[Pos] = std::make_pair(x&: PhysReg, y: true);
739 }
740 }
741 if (Available == 0) {
742 // Nothing to decide, nothing to learn.
743 assert(!MustFindEviction);
744 return MCRegister::NoRegister;
745 }
746 const size_t ValidPosLimit = Pos;
747 // If we must find eviction, the candidate should be masked out of the
748 // decision making process.
749 Regs[CandidateVirtRegPos].second = !MustFindEviction;
750 if (!MustFindEviction)
751 extractFeatures(Intervals: SmallVector<const LiveInterval *, 1>(1, &VirtReg), Largest,
752 Pos: CandidateVirtRegPos, /*IsHint*/ 0,
753 /*LocalIntfsCount*/ 0,
754 /*NumUrgent*/ 0.0, LRPosInfo);
755 assert(InitialQSize > 0.0 && "We couldn't have gotten here if we had "
756 "nothing to allocate initially.");
757 // Normalize the features.
758 for (auto &V : Largest)
759 V = V ? V : 1.0;
760 for (size_t FeatureIndex = 0; FeatureIndex < FeatureIDs::FeatureCount;
761 ++FeatureIndex) {
762 if (DoNotNormalize.test(position: FeatureIndex))
763 continue;
764 for (size_t Pos = 0; Pos < NumberOfInterferences; ++Pos) {
765 Runner->getTensor<float>(FeatureID: FeatureIndex)[Pos] /= Largest[FeatureIndex];
766 }
767 }
768 *Runner->getTensor<float>(FeatureID: FeatureIDs::progress) =
769 static_cast<float>(RA.getQueueSize()) / InitialQSize;
770
771 // Get a decision.
772 size_t CandidatePos = tryFindEvictionCandidatePosition(
773 VirtReg, Order, OrderLimit, CostPerUseLimit, FixedRegisters);
774 // The contract with the ML side is that CandidatePos is mask == 1 (i.e.
775 // Regs[CandidatePos].second)
776 assert(Regs[CandidatePos].second);
777 if (CandidatePos == CandidateVirtRegPos) {
778 onEviction(RegBeingEvicted: VirtReg.reg());
779 assert(!MustFindEviction);
780 return MCRegister::NoRegister;
781 }
782 assert(CandidatePos < ValidPosLimit);
783 (void)ValidPosLimit;
784
785 // Update information about how many times the virtual registers being
786 // evicted have been evicted so that we can prevent the model from evicting
787 // the same ranges continually and eating compile time.
788 for (MCRegUnit Unit : TRI->regunits(Reg: Regs[CandidatePos].first)) {
789 LiveIntervalUnion::Query &Q = Matrix->query(LR: VirtReg, RegUnit: Unit);
790 const auto &IFIntervals = Q.interferingVRegs(MaxInterferingRegs: EvictInterferenceCutoff);
791 for (const LiveInterval *Intf : reverse(C: IFIntervals)) {
792 onEviction(RegBeingEvicted: Intf->reg());
793 }
794 }
795
796 return Regs[CandidatePos].first;
797}
798
799const LIFeatureComponents &
800MLEvictAdvisor::getLIFeatureComponents(const LiveInterval &LI) const {
801 RegID ID = LI.reg().id();
802 LIFeatureComponents Empty;
803 auto I = CachedFeatures.insert(KV: std::make_pair(x&: ID, y&: Empty));
804 LIFeatureComponents &Ret = I.first->getSecond();
805 if (!I.second)
806 return Ret;
807
808 SmallPtrSet<MachineInstr *, 8> Visited;
809 const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo();
810
811 for (MachineRegisterInfo::reg_instr_nodbg_iterator
812 I = MRI->reg_instr_nodbg_begin(RegNo: LI.reg()),
813 E = MRI->reg_instr_nodbg_end();
814 I != E;) {
815 MachineInstr *MI = &*(I++);
816
817 ++Ret.NumDefsAndUses;
818 if (!Visited.insert(Ptr: MI).second)
819 continue;
820
821 if (MI->isIdentityCopy() || MI->isImplicitDef())
822 continue;
823
824 bool Reads, Writes;
825 std::tie(args&: Reads, args&: Writes) = MI->readsWritesVirtualRegister(Reg: LI.reg());
826
827 float Freq = MBFI.getBlockFreqRelativeToEntryBlock(MBB: MI->getParent());
828 Ret.HottestBlockFreq = std::max(a: Freq, b: Ret.HottestBlockFreq);
829
830 Ret.R += (Reads && !Writes) * Freq;
831 Ret.W += (!Reads && Writes) * Freq;
832 Ret.RW += (Reads && Writes) * Freq;
833
834 auto *MBB = MI->getParent();
835 auto *Loop = Loops.getLoopFor(BB: MBB);
836 bool IsExiting = Loop ? Loop->isLoopExiting(BB: MBB) : false;
837
838 if (Writes && IsExiting && LIS->isLiveOutOfMBB(LR: LI, mbb: MBB))
839 Ret.IndVarUpdates += Freq;
840
841 if (MI->isCopy() && VirtRegAuxInfo::copyHint(MI, Reg: LI.reg(), TRI, MRI: *MRI))
842 Ret.HintWeights += Freq;
843 }
844 Ret.IsRemat = VirtRegAuxInfo::isRematerializable(
845 LI, LIS: *LIS, VRM: *VRM, MRI: *MRI, TII: *MF.getSubtarget().getInstrInfo());
846 return Ret;
847}
848
849// Overall, this currently mimics what we do for weight calculation, but instead
850// of accummulating the various features, we keep them separate.
851void MLEvictAdvisor::extractFeatures(
852 const SmallVectorImpl<const LiveInterval *> &Intervals,
853 llvm::SmallVectorImpl<float> &Largest, size_t Pos, int64_t IsHint,
854 int64_t LocalIntfsCount, float NumUrgent,
855 SmallVectorImpl<LRStartEndInfo> &LRPosInfo) const {
856 int64_t NumDefsAndUses = 0;
857 int64_t NumBrokenHints = 0;
858 double R = 0.0;
859 double W = 0.0;
860 double RW = 0.0;
861 double IndVarUpdates = 0.0;
862 double HintWeights = 0.0;
863 float StartBBFreq = 0.0;
864 float EndBBFreq = 0.0;
865 float HottestBlockFreq = 0.0;
866 int32_t NumRematerializable = 0;
867 float TotalWeight = 0.0;
868
869 SlotIndex EndSI = LIS->getSlotIndexes()->getZeroIndex();
870 SlotIndex StartSI = LIS->getSlotIndexes()->getLastIndex();
871 int64_t MaxStage = 0;
872 int64_t MinStage =
873 Intervals.empty() ? 0 : std::numeric_limits<int64_t>::max();
874
875 for (const auto *L : Intervals) {
876 const LiveInterval &LI = *L;
877 MaxStage = std::max<int64_t>(
878 a: MaxStage, b: static_cast<int64_t>(RA.getExtraInfo().getStage(VirtReg: LI)));
879 MinStage = std::min<int64_t>(
880 a: MinStage, b: static_cast<int64_t>(RA.getExtraInfo().getStage(VirtReg: LI)));
881
882 TotalWeight = std::max(a: TotalWeight, b: LI.weight());
883
884 if (LI.beginIndex() < StartSI)
885 StartSI = LI.beginIndex();
886
887 if (LI.endIndex() > EndSI)
888 EndSI = LI.endIndex();
889 const LIFeatureComponents &LIFC = getLIFeatureComponents(LI);
890 NumBrokenHints += VRM->hasPreferredPhys(VirtReg: LI.reg());
891
892 NumDefsAndUses += LIFC.NumDefsAndUses;
893 HottestBlockFreq = std::max(a: HottestBlockFreq, b: LIFC.HottestBlockFreq);
894 R += LIFC.R;
895 W += LIFC.W;
896 RW += LIFC.RW;
897
898 IndVarUpdates += LIFC.IndVarUpdates;
899
900 HintWeights += LIFC.HintWeights;
901 NumRematerializable += LIFC.IsRemat;
902 }
903 size_t Size = 0;
904 if (!Intervals.empty()) {
905 StartBBFreq =
906 MBFI.getBlockFreqRelativeToEntryBlock(MBB: LIS->getMBBFromIndex(index: StartSI));
907 if (EndSI >= LIS->getSlotIndexes()->getLastIndex())
908 EndSI = LIS->getSlotIndexes()->getLastIndex().getPrevIndex();
909 EndBBFreq =
910 MBFI.getBlockFreqRelativeToEntryBlock(MBB: LIS->getMBBFromIndex(index: EndSI));
911 Size = StartSI.distance(other: EndSI);
912 }
913 // Set the features at the column 'Pos'.
914#define SET(ID, TYPE, VAL) \
915 do { \
916 Runner->getTensor<TYPE>(FeatureIDs::ID)[Pos] = static_cast<TYPE>(VAL); \
917 if (!DoNotNormalize.test(FeatureIDs::ID)) \
918 Largest[FeatureIDs::ID] = \
919 std::max(Largest[FeatureIDs::ID], static_cast<float>(VAL)); \
920 } while (false)
921 SET(mask, int64_t, 1);
922 SET(is_free, int64_t, Intervals.empty());
923 SET(nr_urgent, float, NumUrgent);
924 SET(nr_broken_hints, float, NumBrokenHints);
925 SET(is_hint, int64_t, IsHint);
926 SET(is_local, int64_t, LocalIntfsCount);
927 SET(nr_rematerializable, float, NumRematerializable);
928 SET(nr_defs_and_uses, float, NumDefsAndUses);
929 SET(weighed_reads_by_max, float, R);
930 SET(weighed_writes_by_max, float, W);
931 SET(weighed_read_writes_by_max, float, RW);
932 SET(weighed_indvars_by_max, float, IndVarUpdates);
933 SET(hint_weights_by_max, float, HintWeights);
934 SET(start_bb_freq_by_max, float, StartBBFreq);
935 SET(end_bb_freq_by_max, float, EndBBFreq);
936 SET(hottest_bb_freq_by_max, float, HottestBlockFreq);
937 SET(liverange_size, float, Size);
938 SET(use_def_density, float, TotalWeight);
939 SET(max_stage, int64_t, MaxStage);
940 SET(min_stage, int64_t, MinStage);
941#undef SET
942}
943
944// Development mode-specific implementations
945#ifdef LLVM_HAVE_TFLITE
946
947RegAllocEvictionAdvisorAnalysisLegacy *
948llvm::createDevelopmentModeAdvisorAnalysisLegacy() {
949 return new DevelopmentModeEvictionAdvisorAnalysisLegacy();
950}
951
952int64_t DevelopmentModeEvictAdvisor::tryFindEvictionCandidatePosition(
953 const LiveInterval &VirtReg, const AllocationOrder &Order,
954 unsigned OrderLimit, uint8_t CostPerUseLimit,
955 const SmallVirtRegSet &FixedRegisters) const {
956 int64_t Ret = 0;
957 if (isa<ModelUnderTrainingRunner>(getRunner())) {
958 Ret = MLEvictAdvisor::tryFindEvictionCandidatePosition(
959 VirtReg, Order, OrderLimit, CostPerUseLimit, FixedRegisters);
960 } else {
961 MCRegister PhysReg = getDefaultAdvisor().tryFindEvictionCandidate(
962 VirtReg, Order, CostPerUseLimit, FixedRegisters);
963 // Find the index of the selected PhysReg. We need it for logging,
964 // otherwise this is wasted cycles (but so would starting development mode
965 // without a model nor logging)
966 if (!PhysReg)
967 Ret = CandidateVirtRegPos;
968 else
969 for (auto I = Order.begin(), E = Order.getOrderLimitEnd(OrderLimit);
970 I != E; ++I, ++Ret)
971 if (*I == PhysReg)
972 break;
973 }
974 if (TrainingLog.empty())
975 return Ret;
976 // TODO(mtrofin): when we support optional rewards, this can go away. In the
977 // meantime, we log the "pretend" reward (0) for the previous observation
978 // before starting a new one.
979 if (Log->hasObservationInProgress())
980 Log->logReward<float>(0.0);
981
982 Log->startObservation();
983 size_t CurrentFeature = 0;
984 size_t FeatureCount = FeatureIDs::FeatureCount;
985 for (; CurrentFeature < FeatureCount; ++CurrentFeature) {
986 Log->logTensorValue(CurrentFeature,
987 reinterpret_cast<const char *>(
988 getRunner().getTensorUntyped(CurrentFeature)));
989 }
990 if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner()))
991 for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size();
992 ++I, ++CurrentFeature)
993 Log->logTensorValue(
994 CurrentFeature,
995 reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I)));
996 // The output is right after the features and the extra outputs
997 Log->logTensorValue(CurrentFeature, reinterpret_cast<const char *>(&Ret));
998 Log->endObservation();
999 return Ret;
1000}
1001
1002bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
1003 std::optional<float> CachedReward;
1004 auto GetReward = [&]() {
1005 if (!CachedReward)
1006 CachedReward = static_cast<float>(
1007 calculateRegAllocScore(
1008 MF, getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI())
1009 .getScore());
1010 return *CachedReward;
1011 };
1012
1013 getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().logRewardIfNeeded(
1014 MF, GetReward);
1015 getAnalysis<RegAllocPriorityAdvisorAnalysisLegacy>().logRewardIfNeeded(
1016 MF, GetReward);
1017 return false;
1018}
1019#endif // #ifdef LLVM_HAVE_TFLITE
1020
1021RegAllocEvictionAdvisorProvider *
1022llvm::createReleaseModeAdvisorProvider(LLVMContext &Ctx) {
1023 return new ReleaseModeEvictionAdvisorProvider(Ctx);
1024}
1025
1026RegAllocEvictionAdvisorProvider *
1027llvm::createDevelopmentModeAdvisorProvider(LLVMContext &Ctx) {
1028#if defined(LLVM_HAVE_TFLITE)
1029 return new DevelopmentModeEvictionAdvisorProvider(Ctx);
1030#endif
1031 return nullptr;
1032}
1033
1034RegAllocEvictionAdvisorAnalysisLegacy *
1035llvm::createReleaseModeAdvisorAnalysisLegacy() {
1036 return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
1037 !InteractiveChannelBaseName.empty()
1038 ? new ReleaseModeEvictionAdvisorAnalysisLegacy()
1039 : nullptr;
1040}
1041
1042// In all cases except development mode, we don't need scoring.
1043#if !defined(LLVM_HAVE_TFLITE)
1044bool RegAllocScoring::runOnMachineFunction(MachineFunction &) { return false; }
1045#endif
1046