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/MLRegallocPriorityAdvisor.cpp b/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp new file mode 100644 --- /dev/null +++ b/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp @@ -0,0 +1,290 @@ +#include "AllocationOrder.h" +#include "RegAllocGreedy.h" +#include "llvm/Analysis/AliasAnalysis.h" +#include "llvm/Analysis/MLModelRunner.h" +#include "llvm/Analysis/ReleaseModeModelRunner.h" +#include "llvm/Analysis/TensorSpec.h" +#include "llvm/CodeGen/CalcSpillWeights.h" +#include "llvm/CodeGen/LiveRegMatrix.h" +#include "llvm/CodeGen/MachineBlockFrequencyInfo.h" +#include "llvm/CodeGen/MachineFunction.h" +#include "llvm/CodeGen/MachineLoopInfo.h" +#include "llvm/CodeGen/MachineRegisterInfo.h" +#include "llvm/CodeGen/Passes.h" +#include "llvm/CodeGen/RegisterClassInfo.h" +#include "llvm/CodeGen/VirtRegMap.h" +#include "llvm/InitializePasses.h" +#include "llvm/Pass.h" +#include "llvm/PassRegistry.h" +#include "llvm/Support/CommandLine.h" + +#include "llvm/Analysis/ModelUnderTrainingRunner.h" +#include "llvm/Analysis/NoInferenceModelRunner.h" + +#include "RegAllocScore.h" +#include "llvm/Analysis/Utils/TFUtils.h" + +using namespace llvm; + +static cl::opt TrainingLog( + "regalloc-prio-training-log", cl::Hidden, + cl::desc("Training log for the register allocator priority model")); + +static cl::opt ModelUnderTraining( + "regalloc-prio-model", cl::Hidden, + cl::desc("The model being trained for register allocation priority")); + +namespace llvm { + +static const std::vector PerLiveRangeShape{1}; + +#define RA_PRIORITY_FEATURES_LIST(M) \ + M(int64_t, li_size, PerLiveRangeShape, "size") \ + M(int64_t, stage, PerLiveRangeShape, "stage") \ + M(float, weight, PerLiveRangeShape, "weight") + +#define DecisionName "priority" + +// Named features index. +enum FeatureIDs { +#define _FEATURE_IDX(_, name, __, ___) name, + RA_PRIORITY_FEATURES_LIST(_FEATURE_IDX) +#undef _FEATURE_IDX + FeatureCount +}; + +class MLPriorityAdvisor : public RegAllocPriorityAdvisor { +public: + MLPriorityAdvisor(const MachineFunction &MF, const RAGreedy &RA, + MLModelRunner *Runner); + void logReward(float reward) override {} + +protected: + const RegAllocPriorityAdvisor &getDefaultAdvisor() const { + return static_cast(DefaultAdvisor); + } + + // The assumption is that if the Runner could not be constructed, we emit-ed + // error, and we shouldn't be asking for it here. + const MLModelRunner &getRunner() const { return *Runner; } + + float tryFindPriority(unsigned Prio, unsigned Size, LiveRangeStage Stage, + float Weight) const override; + +private: + const DefaultPriorityAdvisor DefaultAdvisor; + MLModelRunner *const Runner; +}; + +#define _DECL_FEATURES(type, name, shape, _) \ + TensorSpec::createSpec(#name, shape), + +static const std::vector InputFeatures{ + {RA_PRIORITY_FEATURES_LIST(_DECL_FEATURES)}, +}; +#undef _DECL_FEATURES + +// =================================== +// Release (AOT) - specifics +// =================================== +class ReleaseModePriorityAdvisorAnalysis final + : public RegAllocPriorityAdvisorAnalysis { +public: + ReleaseModePriorityAdvisorAnalysis() + : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Release) {} + // support for isa<> and dyn_cast. + static bool classof(const RegAllocPriorityAdvisorAnalysis *R) { + return R->getAdvisorMode() == AdvisorMode::Release; + } + +private: + void getAnalysisUsage(AnalysisUsage &AU) const override { + RegAllocPriorityAdvisorAnalysis::getAnalysisUsage(AU); + } + + std::unique_ptr + getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override { + if (!Runner) + Runner = std::make_unique>( + MF.getFunction().getContext(), InputFeatures, DecisionName); + return std::make_unique(MF, RA, Runner.get()); + } + std::unique_ptr> Runner; +}; + +static const TensorSpec Output = + TensorSpec::createSpec(DecisionName, {1}); +static const TensorSpec Reward = TensorSpec::createSpec("reward", {1}); + +#define _DECL_TRAIN_FEATURES(type, name, shape, _) \ + TensorSpec::createSpec(std::string("action_") + #name, shape), + +static const std::vector TrainingInputFeatures{ + {RA_PRIORITY_FEATURES_LIST(_DECL_TRAIN_FEATURES) + TensorSpec::createSpec("action_discount", {1}), + TensorSpec::createSpec("action_step_type", {1}), + TensorSpec::createSpec("action_reward", {1})}}; +#undef _DECL_TRAIN_FEATURES + +class DevelopmentModePriorityAdvisor : public MLPriorityAdvisor { +public: + DevelopmentModePriorityAdvisor(const MachineFunction &MF, const RAGreedy &RA, + MLModelRunner *Runner, Logger *Log) + : MLPriorityAdvisor(MF, RA, Runner), Log(Log) {} + + void logReward(float reward) override { Log->logFloatFinalReward(reward); } + +private: + float tryFindPriority(unsigned Prio, unsigned Size, LiveRangeStage Stage, + float Weight) const override; + Logger *const Log; +}; + +class DevelopmentModePriorityAdvisorAnalysis final + : public RegAllocPriorityAdvisorAnalysis { +public: + DevelopmentModePriorityAdvisorAnalysis() + : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Development) {} + // support for isa<> and dyn_cast. + static bool classof(const RegAllocPriorityAdvisorAnalysis *R) { + return R->getAdvisorMode() == AdvisorMode::Development; + } + + /// get the logger for the given function, or nullptr if we didn't collect + /// one. This is used to inject the score by the RegAllocScoring pass. + Logger *getLogger(const MachineFunction &MF) const { + auto I = LogMap.find(MF.getName()); + if (I == LogMap.end()) + return nullptr; + return I->second.get(); + } + +private: + void getAnalysisUsage(AnalysisUsage &AU) const override { + RegAllocPriorityAdvisorAnalysis::getAnalysisUsage(AU); + } + + // Save all the logs (when requested). + bool doFinalization(Module &M) override { + if (TrainingLog.empty()) + return false; + std::error_code EC; + auto OS = std::make_unique(TrainingLog, EC); + if (EC) { + M.getContext().emitError(EC.message() + ":" + TrainingLog); + return false; + } + Logger::flushLogs(*OS, LogMap); + return false; + } + + std::unique_ptr + getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override { + + LLVMContext &Ctx = MF.getFunction().getContext(); + if (ModelUnderTraining.empty() && TrainingLog.empty()) { + Ctx.emitError("Regalloc development mode should be requested with at " + "least logging enabled and/or a training model"); + return nullptr; + } + if (!Runner) { + if (ModelUnderTraining.empty()) + Runner = std::make_unique(Ctx, InputFeatures); + else + Runner = ModelUnderTrainingRunner::createAndEnsureValid( + Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures); + if (!Runner) { + Ctx.emitError("Regalloc: could not set up the model runner"); + return nullptr; + } + } + + Logger *Log = nullptr; + if (!TrainingLog.empty()) { + std::vector LFS; + for (const auto &FS : InputFeatures) + LFS.push_back({FS, None}); + if (auto *MUTR = dyn_cast(Runner.get())) + if (MUTR->outputLoggedFeatureSpecs().size() > 1) + append_range(LFS, drop_begin(MUTR->outputLoggedFeatureSpecs())); + // We always log the output; in particular, if we're not evaluating, we + // don't have an output spec json file. That's why we handle the + // 'normal' output separately. + LFS.push_back({Output, None}); + auto I = LogMap.insert(std::make_pair( + MF.getFunction().getName(), + std::make_unique(LFS, Reward, /*IncludeReward*/ true))); + assert(I.second); + Log = I.first->second.get(); + } + + return std::make_unique(MF, RA, + Runner.get(), Log); + } + + std::unique_ptr Runner; + StringMap> LogMap; +}; + +} // namespace llvm + +RegAllocPriorityAdvisorAnalysis *llvm::createDevelopmentModePriorityAdvisor() { + return new DevelopmentModePriorityAdvisorAnalysis(); +} + +RegAllocPriorityAdvisorAnalysis *llvm::createReleaseModePriorityAdvisor() { + return new ReleaseModePriorityAdvisorAnalysis(); +} + +MLPriorityAdvisor::MLPriorityAdvisor(const MachineFunction &MF, + const RAGreedy &RA, MLModelRunner *Runner) + : RegAllocPriorityAdvisor(MF, RA), DefaultAdvisor(MF, RA), + Runner(std::move(Runner)) { + assert(this->Runner); +} + +float MLPriorityAdvisor::tryFindPriority(unsigned Prio, unsigned Size, + LiveRangeStage Stage, + float Weight) const { + *Runner->getTensor(0) = static_cast(Size); + *Runner->getTensor(1) = static_cast(Stage); + *Runner->getTensor(2) = static_cast(Weight); + + float Ret = Runner->evaluate(); + return Ret; +} + +float DevelopmentModePriorityAdvisor::tryFindPriority(unsigned Prio, + unsigned Size, + LiveRangeStage Stage, + float Weight) const { + float Ret = 0; + + if (isa(getRunner())) { + Ret = MLPriorityAdvisor::tryFindPriority(Prio, Size, Stage, Weight); + } else { + Ret = static_cast(Prio); + } + + if (TrainingLog.empty()) + return Ret; + + size_t CurrentFeature = 0; + for (; CurrentFeature < InputFeatures.size(); ++CurrentFeature) { + Log->logSpecifiedTensorValue( + CurrentFeature, reinterpret_cast( + getRunner().getTensorUntyped(CurrentFeature))); + } + + if (auto *MUTR = dyn_cast(&getRunner())) { + for (size_t I = 1; I < MUTR->outputLoggedFeatureSpecs().size(); + ++I, ++CurrentFeature) + Log->logSpecifiedTensorValue( + CurrentFeature, + reinterpret_cast( + MUTR->lastEvaluationResult()->getUntypedTensorValue(I))); + } + Log->logFloatValue(CurrentFeature, &Ret); + + return Ret; +} 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