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