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