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