diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h --- a/llvm/include/llvm/InitializePasses.h +++ b/llvm/include/llvm/InitializePasses.h @@ -369,6 +369,7 @@ void initializeReassociateLegacyPassPass(PassRegistry&); void initializeRedundantDbgInstEliminationPass(PassRegistry&); void initializeRegAllocEvictionAdvisorAnalysisPass(PassRegistry &); +void initializeRegAllocPriorityAdvisorAnalysisPass(PassRegistry &); void initializeRegAllocFastPass(PassRegistry&); void initializeRegAllocScoringPass(PassRegistry &); void initializeRegBankSelectPass(PassRegistry&); diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt --- a/llvm/lib/CodeGen/CMakeLists.txt +++ b/llvm/lib/CodeGen/CMakeLists.txt @@ -144,6 +144,7 @@ MIRSampleProfile.cpp MIRYamlMapping.cpp MLRegallocEvictAdvisor.cpp + MLRegallocPriorityAdvisor.cpp ModuloSchedule.cpp MultiHazardRecognizer.cpp PatchableFunction.cpp @@ -171,6 +172,7 @@ RegAllocBase.cpp RegAllocBasic.cpp RegAllocEvictionAdvisor.cpp + RegAllocPriorityAdvisor.cpp RegAllocFast.cpp RegAllocGreedy.cpp RegAllocPBQP.cpp diff --git a/llvm/lib/CodeGen/RegAllocGreedy.h b/llvm/lib/CodeGen/RegAllocGreedy.h --- a/llvm/lib/CodeGen/RegAllocGreedy.h +++ b/llvm/lib/CodeGen/RegAllocGreedy.h @@ -15,6 +15,7 @@ #include "InterferenceCache.h" #include "RegAllocBase.h" #include "RegAllocEvictionAdvisor.h" +#include "RegAllocPriorityAdvisor.h" #include "SpillPlacement.h" #include "SplitKit.h" #include "llvm/ADT/ArrayRef.h" @@ -26,6 +27,9 @@ #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" #include "llvm/Analysis/AliasAnalysis.h" +#include "llvm/Analysis/NoInferenceModelRunner.h" +#include "llvm/Analysis/TensorSpec.h" +#include "llvm/Analysis/Utils/TFUtils.h" #include "llvm/CodeGen/CalcSpillWeights.h" #include "llvm/CodeGen/LiveInterval.h" #include "llvm/CodeGen/LiveRangeEdit.h" @@ -151,7 +155,7 @@ private: // Convenient shortcuts. - using PQueue = std::priority_queue>; + using PQueue = std::priority_queue>; using SmallLISet = SmallPtrSet; // We need to track all tentative recolorings so we can roll back any @@ -183,6 +187,8 @@ Optional ExtraInfo; std::unique_ptr EvictAdvisor; + std::unique_ptr PriorityAdvisor; + // Enum CutOffStage to keep a track whether the register allocation failed // because of the cutoffs encountered in last chance recoloring. // Note: This is used as bitmask. New value should be next power of 2. diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp --- a/llvm/lib/CodeGen/RegAllocGreedy.cpp +++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp @@ -17,6 +17,7 @@ #include "LiveDebugVariables.h" #include "RegAllocBase.h" #include "RegAllocEvictionAdvisor.h" +#include "RegAllocScore.h" #include "SpillPlacement.h" #include "SplitKit.h" #include "llvm/ADT/ArrayRef.h" @@ -29,6 +30,7 @@ #include "llvm/ADT/Statistic.h" #include "llvm/ADT/StringRef.h" #include "llvm/Analysis/AliasAnalysis.h" +#include "llvm/Analysis/ModelUnderTrainingRunner.h" #include "llvm/Analysis/OptimizationRemarkEmitter.h" #include "llvm/CodeGen/CalcSpillWeights.h" #include "llvm/CodeGen/EdgeBundles.h" @@ -157,6 +159,7 @@ INITIALIZE_PASS_DEPENDENCY(SpillPlacement) INITIALIZE_PASS_DEPENDENCY(MachineOptimizationRemarkEmitterPass) INITIALIZE_PASS_DEPENDENCY(RegAllocEvictionAdvisorAnalysis) +INITIALIZE_PASS_DEPENDENCY(RegAllocPriorityAdvisorAnalysis) INITIALIZE_PASS_END(RAGreedy, "greedy", "Greedy Register Allocator", false, false) @@ -224,6 +227,7 @@ AU.addRequired(); AU.addRequired(); AU.addRequired(); + AU.addRequired(); MachineFunctionPass::getAnalysisUsage(AU); } @@ -347,9 +351,12 @@ if (VRM->hasKnownPreference(Reg)) Prio |= (1u << 30); } + + float Ret = PriorityAdvisor->tryFindPriority(Prio, Size, Stage, LI->weight()); + // The virtual register number is a tie breaker for same-sized ranges. // Give lower vreg numbers higher priority to assign them first. - CurQueue.push(std::make_pair(Prio, ~Reg)); + CurQueue.push(std::make_pair(Ret, ~Reg)); } const LiveInterval *RAGreedy::dequeue() { return dequeue(Queue); } @@ -2688,6 +2695,9 @@ if (VerifyEnabled) MF->verify(this, "Before greedy register allocator"); + PriorityAdvisor = + getAnalysis().getAdvisor(*MF, *this); + RegAllocBase::init(getAnalysis(), getAnalysis(), getAnalysis()); @@ -2736,6 +2746,11 @@ postOptimization(); reportStats(); + PriorityAdvisor->logReward(static_cast( + calculateRegAllocScore(mf, getAnalysis(), + getAnalysis().getAAResults()) + .getScore())); + releaseMemory(); return true; } diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h new file mode 100644 --- /dev/null +++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h @@ -0,0 +1,67 @@ +#include "llvm/Pass.h" + +namespace llvm { + +class MachineFunction; +class RAGreedy; + +class RegAllocPriorityAdvisor { +public: + RegAllocPriorityAdvisor(const RegAllocPriorityAdvisor &) = delete; + RegAllocPriorityAdvisor(RegAllocPriorityAdvisor &&) = delete; + virtual ~RegAllocPriorityAdvisor() = default; + + virtual float tryFindPriority(unsigned Prio, unsigned Size, + LiveRangeStage Stage, float Weight) const = 0; + virtual void logReward(float reward) = 0; + +protected: + RegAllocPriorityAdvisor(const MachineFunction &MF, const RAGreedy &RA); + + const MachineFunction &MF; + const RAGreedy &RA; +}; + +class RegAllocPriorityAdvisorAnalysis : public ImmutablePass { +public: + enum class AdvisorMode : int { Default, Release, Development }; + + RegAllocPriorityAdvisorAnalysis(AdvisorMode Mode) + : ImmutablePass(ID), Mode(Mode){}; + static char ID; + + /// Get an advisor for the given context (i.e. machine function, etc) + virtual std::unique_ptr + getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0; + AdvisorMode getAdvisorMode() const { return Mode; } + +protected: + // This analysis preserves everything, and subclasses may have additional + // requirements. + void getAnalysisUsage(AnalysisUsage &AU) const override { + AU.setPreservesAll(); + } + +private: + StringRef getPassName() const override; + const AdvisorMode Mode; +}; + +template <> Pass *callDefaultCtor(); + +RegAllocPriorityAdvisorAnalysis *createReleaseModePriorityAdvisor(); + +RegAllocPriorityAdvisorAnalysis *createDevelopmentModePriorityAdvisor(); + +class DefaultPriorityAdvisor : public RegAllocPriorityAdvisor { +public: + DefaultPriorityAdvisor(const MachineFunction &MF, const RAGreedy &RA) + : RegAllocPriorityAdvisor(MF, RA) {} + + void logReward(float reward) override{}; + +private: + float tryFindPriority(unsigned Prio, unsigned Size, LiveRangeStage Stage, + float Weight) const override; +}; +} // namespace llvm \ No newline at end of file diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp new file mode 100644 --- /dev/null +++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp @@ -0,0 +1,94 @@ +#include "RegAllocGreedy.h" +#include "llvm/CodeGen/MachineFunction.h" +#include "llvm/InitializePasses.h" +#include "llvm/Pass.h" + +using namespace llvm; + +static cl::opt Mode( + "regalloc-prio-enable-advisor", cl::Hidden, + cl::init(RegAllocPriorityAdvisorAnalysis::AdvisorMode::Default), + cl::desc("Enable regalloc advisor mode"), + cl::values( + clEnumValN(RegAllocPriorityAdvisorAnalysis::AdvisorMode::Default, + "default", "Default"), + clEnumValN(RegAllocPriorityAdvisorAnalysis::AdvisorMode::Release, + "release", "precompiled"), + clEnumValN(RegAllocPriorityAdvisorAnalysis::AdvisorMode::Development, + "development", "for training"))); + +char RegAllocPriorityAdvisorAnalysis::ID = 0; +INITIALIZE_PASS(RegAllocPriorityAdvisorAnalysis, "regalloc-prio", + "RegAllocPriorityAdvisorAnalysis", false, true) + +namespace { +class DefaultPriorityAdvisorAnalysis final + : public RegAllocPriorityAdvisorAnalysis { +public: + DefaultPriorityAdvisorAnalysis(bool NotAsRequested) + : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Default), + NotAsRequested(NotAsRequested) {} + + // support for isa<> and dyn_cast. + static bool classof(const RegAllocPriorityAdvisorAnalysis *R) { + return R->getAdvisorMode() == AdvisorMode::Default; + } + +private: + std::unique_ptr + getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override { + return std::make_unique(MF, RA); + } + bool doInitialization(Module &M) override { + if (NotAsRequested) + M.getContext().emitError("Requested regalloc priority advisor analysis " + "could be created. Using default"); + return RegAllocPriorityAdvisorAnalysis::doInitialization(M); + } + const bool NotAsRequested; +}; +} // namespace + +template <> Pass *llvm::callDefaultCtor() { + Pass *Ret = nullptr; + switch (Mode) { + case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Default: + Ret = new DefaultPriorityAdvisorAnalysis(/*NotAsRequested*/ false); + break; + case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Development: +#if defined(LLVM_HAVE_TF_API) + Ret = createDevelopmentModePriorityAdvisor(); +#endif + break; + case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Release: +#if defined(LLVM_HAVE_TF_AOT) + Ret = createReleaseModePriorityAdvisor(); +#endif + break; + } + if (Ret) + return Ret; + return new DefaultPriorityAdvisorAnalysis(/*NotAsRequested*/ true); +} + +StringRef RegAllocPriorityAdvisorAnalysis::getPassName() const { + switch (getAdvisorMode()) { + case AdvisorMode::Default: + return "Default Regalloc Priority Advisor"; + case AdvisorMode::Release: + return "Release mode Regalloc Priority Advisor"; + case AdvisorMode::Development: + return "Development mode Regalloc Priority Advisor"; + } + llvm_unreachable("Unknown advisor kind"); +} + +RegAllocPriorityAdvisor::RegAllocPriorityAdvisor(const MachineFunction &MF, + const RAGreedy &RA) + : MF(MF), RA(RA) {} + +float DefaultPriorityAdvisor::tryFindPriority(unsigned Prio, unsigned Size, + LiveRangeStage Stage, + float Weight) const { + return static_cast(Prio); +} \ No newline at end of file