19 #if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL) || defined(LLVM_HAVE_TF_API)
44 #define DEBUG_TYPE "ml-regalloc"
47 #if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL)
48 #include "RegallocEvictModel.h"
55 #ifdef LLVM_HAVE_TF_API
61 cl::desc(
"Training log for the register allocator eviction model"));
65 cl::desc(
"The model being trained for register allocation eviction"));
67 #endif // #ifdef LLVM_HAVE_TF_API
86 return "Register Allocation Pass Scoring";
108 "Register Allocation Scoring Pass",
false,
false)
118 static const int64_t MaxInterferences = 32;
126 static const int64_t CandidateVirtRegPos = MaxInterferences;
127 static const int64_t NumberOfInterferences = CandidateVirtRegPos + 1;
131 static const std::vector<int64_t> PerLiveRangeShape{1, NumberOfInterferences};
151 #define RA_EVICT_FEATURES_LIST(M) \
152 M(int64_t, mask, PerLiveRangeShape, \
153 "boolean values, 0 for unavailable candidates (i.e. if a position is 0, " \
155 "can't be evicted)") \
156 M(int64_t, is_free, PerLiveRangeShape, \
157 "boolean values, 1 if this phys reg is actually free (no interferences)") \
158 M(float, nr_urgent, PerLiveRangeShape, \
159 "number of 'urgent' intervals, normalized. Urgent are those that are OK " \
160 "to break cascades") \
161 M(float, nr_broken_hints, PerLiveRangeShape, \
162 "if this position were evicted, how many broken hints would there be") \
163 M(int64_t, is_hint, PerLiveRangeShape, \
164 "is this a preferred phys reg for the candidate") \
165 M(int64_t, is_local, PerLiveRangeShape, \
166 "is this live range local to a basic block") \
167 M(float, nr_rematerializable, PerLiveRangeShape, \
168 "nr rematerializable ranges") \
169 M(float, nr_defs_and_uses, PerLiveRangeShape, \
170 "bb freq - weighed nr defs and uses") \
171 M(float, weighed_reads_by_max, PerLiveRangeShape, \
172 "bb freq - weighed nr of reads, normalized") \
173 M(float, weighed_writes_by_max, PerLiveRangeShape, \
174 "bb feq - weighed nr of writes, normalized") \
175 M(float, weighed_read_writes_by_max, PerLiveRangeShape, \
176 "bb freq - weighed nr of uses that are both read and writes, normalized") \
177 M(float, weighed_indvars_by_max, PerLiveRangeShape, \
178 "bb freq - weighed nr of uses that are indvars, normalized") \
179 M(float, hint_weights_by_max, PerLiveRangeShape, \
180 "bb freq - weighed nr of uses that are hints, normalized") \
181 M(float, start_bb_freq_by_max, PerLiveRangeShape, \
182 "the freq in the start block, normalized") \
183 M(float, end_bb_freq_by_max, PerLiveRangeShape, \
184 "freq of end block, normalized") \
185 M(float, hottest_bb_freq_by_max, PerLiveRangeShape, \
186 "hottest BB freq, normalized") \
187 M(float, liverange_size, PerLiveRangeShape, \
188 "size (instr index diff) of the LR") \
189 M(float, use_def_density, PerLiveRangeShape, \
190 "the max weight, as computed by the manual heuristic") \
191 M(int64_t, max_stage, PerLiveRangeShape, \
192 "largest stage of an interval in this LR") \
193 M(int64_t, min_stage, PerLiveRangeShape, \
194 "lowest stage of an interval in this LR") \
195 M(float, progress, {1}, "ratio of current queue size to initial size")
203 #define DecisionName "index_to_evict"
207 #define _FEATURE_IDX(_, name, __, ___) name,
216 template <
typename T>
size_t getTotalSize(
const std::vector<int64_t> &Shape) {
217 size_t Ret =
sizeof(
T);
218 for (
const auto V : Shape)
224 #define _RESET(TYPE, NAME, SHAPE, __) \
225 std::memset(Runner.getTensorUntyped(FeatureIDs::NAME), 0, \
226 getTotalSize<TYPE>(SHAPE));
233 struct LIFeatureComponents {
237 double IndVarUpdates = 0;
238 double HintWeights = 0.0;
239 int64_t NrDefsAndUses = 0;
240 float HottestBlockFreq = 0.0;
241 bool IsRemat =
false;
244 using CandidateRegList =
245 std::array<std::pair<MCRegister, bool>, NumberOfInterferences>;
246 using FeaturesListNormalizer = std::array<float, FeatureIDs::FeatureCount>;
269 tryFindEvictionCandidatePosition(
const LiveInterval &VirtReg,
271 unsigned OrderLimit, uint8_t CostPerUseLimit,
279 std::array<float, FeatureIDs::FeatureCount> &Largest,
287 uint8_t CostPerUseLimit,
291 std::array<float, FeatureIDs::FeatureCount> &Largest,
292 size_t Pos, int64_t IsHint, int64_t LocalIntfsCount,
293 float NrUrgent)
const;
296 bool canEvictHintInterference(
299 return getDefaultAdvisor().canEvictHintInterference(VirtReg, PhysReg,
303 const LIFeatureComponents &
317 std::bitset<FeatureIDs::FeatureCount> DoNotNormalize;
318 const float InitialQSize;
320 using RegID = unsigned;
324 #define _DECL_FEATURES(type, name, shape, _) \
325 TensorSpec::createSpec<type>(#name, shape),
327 static const std::vector<TensorSpec> InputFeatures{
330 #undef _DECL_FEATURES
334 class ReleaseModeEvictionAdvisorAnalysis final
337 ReleaseModeEvictionAdvisorAnalysis()
341 return R->getAdvisorMode() == AdvisorMode::Release;
351 std::unique_ptr<RegAllocEvictionAdvisor>
354 Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
356 return std::make_unique<MLEvictAdvisor>(
357 MF,
RA, Runner.get(), getAnalysis<MachineBlockFrequencyInfo>(),
358 getAnalysis<MachineLoopInfo>());
360 std::unique_ptr<ReleaseModeModelRunner<CompiledModelType>> Runner;
368 #ifdef LLVM_HAVE_TF_API
371 static const TensorSpec Reward = TensorSpec::createSpec<float>(
"reward", {1});
377 #define _DECL_TRAIN_FEATURES(type, name, shape, _) \
378 TensorSpec::createSpec<type>(std::string("action_") + #name, shape),
380 static const std::vector<TensorSpec> TrainingInputFeatures{
382 TensorSpec::createSpec<float>(
"action_discount", {1}),
383 TensorSpec::createSpec<int32_t>(
"action_step_type", {1}),
384 TensorSpec::createSpec<float>(
"action_reward", {1})}};
385 #undef _DECL_TRAIN_FEATURES
387 class DevelopmentModeEvictAdvisor :
public MLEvictAdvisor {
393 : MLEvictAdvisor(MF,
RA, Runner, MBFI,
Loops), Log(Log) {}
396 int64_t tryFindEvictionCandidatePosition(
398 unsigned OrderLimit, uint8_t CostPerUseLimit,
404 class DevelopmentModeEvictionAdvisorAnalysis final
407 DevelopmentModeEvictionAdvisorAnalysis()
411 return R->getAdvisorMode() == AdvisorMode::Development;
418 if (
I == LogMap.end())
420 return I->second.get();
431 bool doFinalization(
Module &
M)
override {
432 if (TrainingLog.empty())
435 auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
437 M.getContext().emitError(EC.message() +
":" + TrainingLog);
440 Logger::flushLogs(*OS, LogMap);
444 std::unique_ptr<RegAllocEvictionAdvisor>
447 if (ModelUnderTraining.empty() && TrainingLog.empty()) {
448 Ctx.
emitError(
"Regalloc development mode should be requested with at "
449 "least logging enabled and/or a training model");
453 if (ModelUnderTraining.empty())
454 Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
456 Runner = ModelUnderTrainingRunner::createAndEnsureValid(
457 Ctx, ModelUnderTraining,
DecisionName, TrainingInputFeatures);
459 Ctx.
emitError(
"Regalloc: could not set up the model runner");
464 Logger *Log =
nullptr;
465 if (!TrainingLog.empty()) {
466 std::vector<LoggedFeatureSpec> LFS;
467 for (
const auto &
FS : InputFeatures)
468 LFS.push_back({
FS,
None});
469 if (
auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
470 if (MUTR->outputLoggedFeatureSpecs().size() > 1)
475 LFS.push_back({Output,
None});
476 auto I = LogMap.insert(std::make_pair(
478 std::make_unique<Logger>(LFS, Reward,
true)));
480 Log =
I.first->second.get();
482 return std::make_unique<DevelopmentModeEvictAdvisor>(
483 MF,
RA, Runner.get(), getAnalysis<MachineBlockFrequencyInfo>(),
484 getAnalysis<MachineLoopInfo>(), Log);
487 std::unique_ptr<MLModelRunner> Runner;
490 #endif //#ifdef LLVM_HAVE_TF_API
511 InitialQSize(MLEvictAdvisor::getInitialQueueSize(MF)) {
513 DoNotNormalize.set(FeatureIDs::mask);
514 DoNotNormalize.set(FeatureIDs::is_free);
515 DoNotNormalize.set(FeatureIDs::is_hint);
517 DoNotNormalize.set(FeatureIDs::min_stage);
518 DoNotNormalize.set(FeatureIDs::max_stage);
519 DoNotNormalize.set(FeatureIDs::progress);
522 int64_t MLEvictAdvisor::tryFindEvictionCandidatePosition(
531 bool MLEvictAdvisor::loadInterferenceFeatures(
533 const SmallVirtRegSet &FixedRegisters, FeaturesListNormalizer &Largest,
536 if (
Matrix->checkInterference(VirtReg, PhysReg) > LiveRegMatrix::IK_VirtReg) {
541 const bool IsLocal = LIS->intervalIsInOneMBB(VirtReg);
542 int64_t LocalIntfs = 0;
543 float NrUrgent = 0.0f;
546 unsigned Cascade =
RA.getExtraInfo().getCascadeOrCurrentNext(VirtReg.
reg());
554 if (IFIntervals.empty() && InterferingIntervals.empty())
558 InterferingIntervals.
append(IFIntervals.begin(), IFIntervals.end());
560 assert(Register::isVirtualRegister(Intf->reg()) &&
561 "Only expecting virtual register interference from query");
568 if (FixedRegisters.
count(Intf->reg()))
570 if (
RA.getExtraInfo().getStage(*Intf) ==
RS_Done)
574 (Intf->isSpillable() ||
576 RegClassInfo.getNumAllocatableRegs(
579 unsigned IntfCascade =
RA.getExtraInfo().getCascade(Intf->reg());
580 if (Cascade <= IntfCascade) {
586 LocalIntfs += (IsLocal && LIS->intervalIsInOneMBB(*Intf) &&
587 (!EnableLocalReassign || !canReassign(*Intf, PhysReg)));
592 extractFeatures(InterferingIntervals, Largest, Pos, IsHint, LocalIntfs,
597 MCRegister MLEvictAdvisor::tryFindEvictionCandidate(
599 uint8_t CostPerUseLimit,
const SmallVirtRegSet &FixedRegisters)
const {
600 auto MaybeOrderLimit = getOrderLimit(VirtReg, Order, CostPerUseLimit);
601 if (!MaybeOrderLimit)
602 return MCRegister::NoRegister;
603 unsigned OrderLimit = *MaybeOrderLimit;
611 const bool MustFindEviction =
612 (!VirtReg.
isSpillable() && CostPerUseLimit ==
static_cast<uint8_t
>(~0u));
617 resetInputs(*Runner);
622 CandidateRegList Regs;
623 Regs.fill({0,
false});
629 FeaturesListNormalizer Largest;
642 assert(!Regs[Pos].second);
644 if (!canAllocatePhysReg(CostPerUseLimit, PhysReg)) {
647 if (loadInterferenceFeatures(VirtReg, PhysReg,
I.isHint(), FixedRegisters,
650 Regs[Pos] = std::make_pair(PhysReg,
true);
653 if (Available == 0) {
655 assert(!MustFindEviction);
656 return MCRegister::NoRegister;
658 const size_t ValidPosLimit = Pos;
661 Regs[CandidateVirtRegPos].second = !MustFindEviction;
662 if (!MustFindEviction)
664 CandidateVirtRegPos, 0, 0,
666 assert(InitialQSize > 0.0 &&
"We couldn't have gotten here if we had "
667 "nothing to allocate initially.");
669 for (
auto &V : Largest)
675 for (
size_t Pos = 0; Pos < NumberOfInterferences; ++Pos) {
679 *Runner->
getTensor<
float>(FeatureIDs::progress) =
680 static_cast<float>(
RA.getQueueSize()) / InitialQSize;
683 size_t CandidatePos = tryFindEvictionCandidatePosition(
684 VirtReg, Order, OrderLimit, CostPerUseLimit, FixedRegisters);
687 assert(Regs[CandidatePos].second);
688 if (CandidatePos == CandidateVirtRegPos) {
689 assert(!MustFindEviction);
690 return MCRegister::NoRegister;
692 assert(CandidatePos < ValidPosLimit);
694 return Regs[CandidatePos].first;
697 const LIFeatureComponents &
698 MLEvictAdvisor::getLIFeatureComponents(
const LiveInterval &LI)
const {
700 LIFeatureComponents
Empty;
701 auto I = CachedFeatures.insert(std::make_pair(
ID,
Empty));
702 LIFeatureComponents &
Ret =
I.first->getSecond();
719 if (
MI->isIdentityCopy() ||
MI->isImplicitDef())
723 std::tie(Reads, Writes) =
MI->readsWritesVirtualRegister(LI.
reg());
728 Ret.R += (Reads && !Writes) * Freq;
729 Ret.W += (!Reads && Writes) * Freq;
730 Ret.RW += (Reads && Writes) * Freq;
732 auto *
MBB =
MI->getParent();
736 if (Writes && IsExiting && LIS->isLiveOutOfMBB(LI,
MBB))
737 Ret.IndVarUpdates += Freq;
740 Ret.HintWeights += Freq;
742 Ret.IsRemat = VirtRegAuxInfo::isRematerializable(
749 void MLEvictAdvisor::extractFeatures(
751 std::array<float, FeatureIDs::FeatureCount> &Largest,
size_t Pos,
752 int64_t IsHint, int64_t LocalIntfsCount,
float NrUrgent)
const {
753 int64_t NrDefsAndUses = 0;
754 int64_t NrBrokenHints = 0;
758 double IndVarUpdates = 0.0;
759 double HintWeights = 0.0;
760 float StartBBFreq = 0.0;
761 float EndBBFreq = 0.0;
762 float HottestBlockFreq = 0.0;
763 int32_t NrRematerializable = 0;
764 float TotalWeight = 0.0;
766 SlotIndex EndSI = LIS->getSlotIndexes()->getZeroIndex();
767 SlotIndex StartSI = LIS->getSlotIndexes()->getLastIndex();
768 int64_t MaxStage = 0;
772 for (
const auto *L : Intervals) {
774 MaxStage = std::max<int64_t>(
775 MaxStage,
static_cast<int64_t
>(
RA.getExtraInfo().getStage(LI)));
776 MinStage = std::min<int64_t>(
777 MinStage,
static_cast<int64_t
>(
RA.getExtraInfo().getStage(LI)));
786 const LIFeatureComponents &LIFC = getLIFeatureComponents(LI);
787 NrBrokenHints += VRM->hasPreferredPhys(LI.
reg());
789 NrDefsAndUses += LIFC.NrDefsAndUses;
790 HottestBlockFreq =
std::max(HottestBlockFreq, LIFC.HottestBlockFreq);
795 IndVarUpdates += LIFC.IndVarUpdates;
797 HintWeights += LIFC.HintWeights;
798 NrRematerializable += LIFC.IsRemat;
801 if (!Intervals.empty()) {
804 if (EndSI >= LIS->getSlotIndexes()->getLastIndex())
805 EndSI = LIS->getSlotIndexes()->getLastIndex().
getPrevIndex();
811 #define SET(ID, TYPE, VAL) \
813 Runner->getTensor<TYPE>(FeatureIDs::ID)[Pos] = static_cast<TYPE>(VAL); \
814 if (!DoNotNormalize.test(FeatureIDs::ID)) \
815 Largest[FeatureIDs::ID] = \
816 std::max(Largest[FeatureIDs::ID], static_cast<float>(VAL)); \
818 SET(mask, int64_t, 1);
819 SET(is_free, int64_t, Intervals.empty());
820 SET(nr_urgent,
float, NrUrgent);
821 SET(nr_broken_hints,
float, NrBrokenHints);
822 SET(is_hint, int64_t, IsHint);
824 SET(nr_rematerializable,
float, NrRematerializable);
825 SET(nr_defs_and_uses,
float, NrDefsAndUses);
826 SET(weighed_reads_by_max,
float,
R);
827 SET(weighed_writes_by_max,
float,
W);
828 SET(weighed_read_writes_by_max,
float, RW);
829 SET(weighed_indvars_by_max,
float, IndVarUpdates);
830 SET(hint_weights_by_max,
float, HintWeights);
831 SET(start_bb_freq_by_max,
float, StartBBFreq);
832 SET(end_bb_freq_by_max,
float, EndBBFreq);
833 SET(hottest_bb_freq_by_max,
float, HottestBlockFreq);
834 SET(liverange_size,
float,
Size);
835 SET(use_def_density,
float, TotalWeight);
836 SET(max_stage, int64_t, MaxStage);
837 SET(min_stage, int64_t, MinStage);
842 #ifdef LLVM_HAVE_TF_API
844 return new DevelopmentModeEvictionAdvisorAnalysis();
847 int64_t DevelopmentModeEvictAdvisor::tryFindEvictionCandidatePosition(
849 unsigned OrderLimit, uint8_t CostPerUseLimit,
852 if (isa<ModelUnderTrainingRunner>(getRunner())) {
853 Ret = MLEvictAdvisor::tryFindEvictionCandidatePosition(
854 VirtReg, Order, OrderLimit, CostPerUseLimit, FixedRegisters);
856 MCRegister PhysReg = getDefaultAdvisor().tryFindEvictionCandidate(
857 VirtReg, Order, CostPerUseLimit, FixedRegisters);
862 Ret = CandidateVirtRegPos;
869 if (TrainingLog.empty())
871 size_t CurrentFeature = 0;
872 for (; CurrentFeature < FeatureIDs::FeatureCount; ++CurrentFeature) {
873 Log->logSpecifiedTensorValue(
874 CurrentFeature,
reinterpret_cast<const char *
>(
875 getRunner().getTensorUntyped(CurrentFeature)));
877 if (
auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner()))
878 for (
size_t I = 1;
I < MUTR->outputLoggedFeatureSpecs().
size();
879 ++
I, ++CurrentFeature)
880 Log->logSpecifiedTensorValue(
882 reinterpret_cast<const char *
>(
883 MUTR->lastEvaluationResult()->getUntypedTensorValue(
I)));
885 Log->logInt64Value(CurrentFeature, &
Ret);
890 if (
auto *DevModeAnalysis = dyn_cast<DevelopmentModeEvictionAdvisorAnalysis>(
891 &getAnalysis<RegAllocEvictionAdvisorAnalysis>()))
892 if (
auto *Log = DevModeAnalysis->getLogger(MF))
893 Log->logFloatFinalReward(
static_cast<float>(
895 MF, getAnalysis<MachineBlockFrequencyInfo>(),
896 getAnalysis<AAResultsWrapperPass>().getAAResults())
901 #endif // #ifdef LLVM_HAVE_TF_API
904 return new ReleaseModeEvictionAdvisorAnalysis();
908 #if !defined(LLVM_HAVE_TF_API)