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 @@ -142,6 +142,7 @@ MIRSampleProfile.cpp MIRYamlMapping.cpp MLRegallocEvictAdvisor.cpp + MLRegallocPriorityAdvisor.cpp ModuloSchedule.cpp MultiHazardRecognizer.cpp PatchableFunction.cpp diff --git a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp --- a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp @@ -90,6 +90,7 @@ void getAnalysisUsage(AnalysisUsage &AU) const override { AU.setPreservesAll(); AU.addRequired(); + AU.addRequired(); AU.addRequired(); MachineFunctionPass::getAnalysisUsage(AU); } @@ -420,6 +421,12 @@ return I->second.get(); } + virtual void logRewardIfNeeded(const MachineFunction &MF, + float Reward) override { + if (auto *Log = this->getLogger(MF)) + Log->logFloatFinalReward(Reward); + } + private: std::vector InputFeatures; std::vector TrainingInputFeatures; @@ -890,12 +897,12 @@ } bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) { - if (auto *DevModeAnalysis = dyn_cast( - &getAnalysis())) - if (auto *Log = DevModeAnalysis->getLogger(MF)) - Log->logFloatFinalReward(static_cast( - calculateRegAllocScore(MF, getAnalysis()) - .getScore())); + float Reward = static_cast( + calculateRegAllocScore(MF, getAnalysis()) + .getScore()); + + getAnalysis().logRewardIfNeeded(MF, Reward); + getAnalysis().logRewardIfNeeded(MF, Reward); return false; } 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,330 @@ +//===- MLRegAllocPriorityAdvisor.cpp - ML priority advisor-----------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// Implementation of the ML priority advisor and reward injection pass +// +//===----------------------------------------------------------------------===// + +#include "AllocationOrder.h" +#include "RegAllocGreedy.h" +#include "RegAllocPriorityAdvisor.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/SlotIndexes.h" +#include "llvm/CodeGen/VirtRegMap.h" +#include "llvm/InitializePasses.h" +#include "llvm/Pass.h" +#include "llvm/PassRegistry.h" +#include "llvm/Support/CommandLine.h" + +#if defined(LLVM_HAVE_TF_API) +#include "llvm/Analysis/ModelUnderTrainingRunner.h" +#include "llvm/Analysis/NoInferenceModelRunner.h" +#include "llvm/Analysis/Utils/TrainingLogger.h" +#endif + +using namespace llvm; + +// Options that only make sense in development mode +#ifdef LLVM_HAVE_TF_API +#include "RegAllocScore.h" +#include "llvm/Analysis/Utils/TFUtils.h" + +static cl::opt TrainingLog( + "regalloc-priority-training-log", cl::Hidden, + cl::desc("Training log for the register allocator priority model")); + +static cl::opt ModelUnderTraining( + "regalloc-priority-model", cl::Hidden, + cl::desc("The model being trained for register allocation priority")); + +#endif // #ifdef LLVM_HAVE_TF_API + +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, + SlotIndexes *const Indexes, MLModelRunner *Runner); + +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; } + + unsigned getPriority(const LiveInterval &LI) 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 { + AU.addRequired(); + 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, &getAnalysis(), Runner.get()); + } + std::unique_ptr> Runner; +}; + +// =================================== +// Development mode-specifics +// =================================== +// +// Features we log +#ifdef LLVM_HAVE_TF_API + +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, + SlotIndexes *const Indexes, + MLModelRunner *Runner, Logger *Log) + : MLPriorityAdvisor(MF, RA, Indexes, Runner), Log(Log) {} + +private: + unsigned getPriority(const LiveInterval &LI) 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. + virtual Logger *getLogger(const MachineFunction &MF) const { + auto I = LogMap.find(MF.getName()); + if (I == LogMap.end()) + return nullptr; + return I->second.get(); + } + + virtual void logRewardIfNeeded(const MachineFunction &MF, + float Reward) override { + if (auto *Log = this->getLogger(MF)) + Log->logFloatFinalReward(Reward); + } + +private: + void getAnalysisUsage(AnalysisUsage &AU) const override { + AU.addRequired(); + 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, &getAnalysis(), Runner.get(), Log); + } + + std::unique_ptr Runner; + StringMap> LogMap; +}; +#endif //#ifdef LLVM_HAVE_TF_API + +} // namespace llvm + +RegAllocPriorityAdvisorAnalysis *llvm::createReleaseModePriorityAdvisor() { + return new ReleaseModePriorityAdvisorAnalysis(); +} + +MLPriorityAdvisor::MLPriorityAdvisor(const MachineFunction &MF, + const RAGreedy &RA, + SlotIndexes *const Indexes, + MLModelRunner *Runner) + : RegAllocPriorityAdvisor(MF, RA, Indexes), DefaultAdvisor(MF, RA, Indexes), + Runner(std::move(Runner)) { + assert(this->Runner); +} + +unsigned MLPriorityAdvisor::getPriority(const LiveInterval &LI) const { + const unsigned Size = LI.getSize(); + LiveRangeStage Stage = RA.getExtraInfo().getStage(LI); + + *Runner->getTensor(0) = static_cast(Size); + *Runner->getTensor(1) = static_cast(Stage); + *Runner->getTensor(2) = static_cast(LI.weight()); + + unsigned Prio = static_cast(Runner->evaluate()); + return Prio; +} + +#ifdef LLVM_HAVE_TF_API +RegAllocPriorityAdvisorAnalysis *llvm::createDevelopmentModePriorityAdvisor() { + return new DevelopmentModePriorityAdvisorAnalysis(); +} + +unsigned +DevelopmentModePriorityAdvisor::getPriority(const LiveInterval &LI) const { + unsigned Prio = 0; + + if (isa(getRunner())) { + Prio = MLPriorityAdvisor::getPriority(LI); + } else { + Prio = getDefaultAdvisor().getPriority(LI); + } + + if (TrainingLog.empty()) + return Prio; + + 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))); + } + + float Ret = static_cast(Prio); + Log->logFloatValue(CurrentFeature, &Ret); + + return Prio; +} + +#endif // #ifdef LLVM_HAVE_TF_API diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h --- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h +++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h @@ -177,6 +177,7 @@ virtual std::unique_ptr getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0; AdvisorMode getAdvisorMode() const { return Mode; } + virtual void logRewardIfNeeded(const MachineFunction &MF, float Reward){}; protected: // This analysis preserves everything, and subclasses may have additional diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h --- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h +++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h @@ -68,6 +68,7 @@ virtual std::unique_ptr getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0; AdvisorMode getAdvisorMode() const { return Mode; } + virtual void logRewardIfNeeded(const MachineFunction &MF, float Reward){}; protected: // This analysis preserves everything, and subclasses may have additional diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp --- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp +++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp @@ -36,6 +36,10 @@ INITIALIZE_PASS(RegAllocPriorityAdvisorAnalysis, "regalloc-priority", "Regalloc priority policy", false, true) +#ifdef LLVM_HAVE_TF_AOT_REGALLOCPRIORITYMODEL +#define LLVM_HAVE_TF_AOT +#endif + namespace { class DefaultPriorityAdvisorAnalysis final : public RegAllocPriorityAdvisorAnalysis { @@ -76,10 +80,14 @@ Ret = new DefaultPriorityAdvisorAnalysis(/*NotAsRequested*/ false); break; case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Development: - // TODO: add implementation +#if defined(LLVM_HAVE_TF_API) + Ret = createDevelopmentModePriorityAdvisor(); +#endif break; case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Release: - // TODO: add implementation +#if defined(LLVM_HAVE_TF_AOT) + Ret = createReleaseModePriorityAdvisor(); +#endif break; } if (Ret)