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/Analysis/models/gen-regalloc-priority-test-model.py b/llvm/lib/Analysis/models/gen-regalloc-priority-test-model.py new file mode 100644 --- /dev/null +++ b/llvm/lib/Analysis/models/gen-regalloc-priority-test-model.py @@ -0,0 +1,95 @@ +"""Generate a mock model for LLVM tests for Register Allocation. +The generated model is not a neural net - it is just a tf.function with the +correct input and output parameters. +""" +## By construction, the mock model will always output the first liverange that can be evicted. + +import os +import sys +import tensorflow as tf +POLICY_DECISION_LABEL = 'priority' +POLICY_OUTPUT_SPEC = """ +[ + { + "logging_name": "priority", + "tensor_spec": { + "name": "StatefulPartitionedCall", + "port": 0, + "type": "float", + "shape": [ + 1 + ] + } + } +] +""" +PER_LIVEINTERVAL_INT64_FEATURE_LIST = [ + 'li_size', 'stage' +] +PER_LIVEINTERVAL_FLOAT32_FEATURE_LIST = ['weight' +] +PER_LIVEINTERVAL_FEATURE_LIST = PER_LIVEINTERVAL_FLOAT32_FEATURE_LIST + \ + PER_LIVEINTERVAL_INT64_FEATURE_LIST +CONTEXT_FEATURE_LIST = ('discount', 'reward', 'step_type') + + +def get_input_signature(): + """Returns (time_step_spec, action_spec) for LLVM register allocation.""" + inputs = dict( + (key, tf.TensorSpec(dtype=tf.int64, shape=(), name=key)) + for key in PER_LIVEINTERVAL_INT64_FEATURE_LIST) + inputs.update( + dict((key, + tf.TensorSpec(dtype=tf.float32, shape=(), name=key)) + for key in PER_LIVEINTERVAL_FLOAT32_FEATURE_LIST)) + inputs.update( + dict((key, tf.TensorSpec(dtype=tf.float32, shape=(), name=key)) + for key in ['discount', 'reward'])) + inputs.update( + dict((key, tf.TensorSpec(dtype=tf.int32, shape=(), name=key)) + for key in ['step_type'])) + return inputs + + +def get_output_spec_path(path): + return os.path.join(path, 'output_spec.json') + + +def build_mock_model(path): + """Build and save the mock model with the given signature.""" + module = tf.Module() + # We have to set this useless variable in order for the TF C API to correctly + # intake it + module.var = tf.Variable(0, dtype=tf.float32) + + def action(*inputs): + s1 = tf.reduce_sum([ + tf.cast(inputs[0][key], tf.float32) for key in PER_LIVEINTERVAL_FEATURE_LIST + ], + axis=0) + s2 = tf.reduce_sum( + [tf.cast(inputs[0][key], tf.float32) for key in CONTEXT_FEATURE_LIST]) + # Add a large number so s won't be 0. + s = s1 + s2 + result = s + module.var + return {POLICY_DECISION_LABEL: result} + module.action = tf.function()(action) + action = { + 'action': module.action.get_concrete_function(get_input_signature()) + } + + tf.saved_model.save(module, path, signatures=action) + output_spec_path = get_output_spec_path(path) + with open(output_spec_path, 'w') as f: + print(f'Writing output spec to {output_spec_path}.') + f.write(POLICY_OUTPUT_SPEC) + + +def main(argv): + assert len(argv) == 2 + model_path = argv[1] + build_mock_model(model_path) + + +if __name__ == '__main__': + main(sys.argv) 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/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 diff --git a/llvm/test/CodeGen/MLRegalloc/Inputs/reference-prio-log-noml.txt b/llvm/test/CodeGen/MLRegalloc/Inputs/reference-prio-log-noml.txt new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/MLRegalloc/Inputs/reference-prio-log-noml.txt @@ -0,0 +1,361 @@ +fields { + key: "SyFgets" + value { + string_value: "feature_lists { + feature_list { + key: \"li_size\" value { + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } } } + feature_list { + key: \"priority\" value { + feature { float_list { value: 2.68435814e+09 } } + feature { float_list { value: 2.68435814e+09 } } + feature { float_list { value: 2.68435866e+09 } } + feature { float_list { value: 2.68435789e+09 } } + feature { float_list { value: 3.75810074e+09 } } + feature { float_list { value: 3.7580969e+09 } } + feature { float_list { value: 2.14748518e+09 } } + feature { float_list { value: 2.14748493e+09 } } + feature { float_list { value: 2.68435814e+09 } } + feature { float_list { value: 2.14748493e+09 } } + feature { float_list { value: 2.14748493e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.68435763e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.14748467e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.68435763e+09 } } + feature { float_list { value: 2.68435763e+09 } } + feature { float_list { value: 2.1474839e+09 } } + feature { float_list { value: 2.1474839e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.1474839e+09 } } + feature { float_list { value: 3.22122547e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.14748365e+09 } } + feature { float_list { value: 2.14748493e+09 } } + feature { float_list { value: 2.14748493e+09 } } + feature { float_list { value: 2.68435814e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.68435584e+09 } } + feature { float_list { value: 2.68435482e+09 } } + feature { float_list { value: 2.68435482e+09 } } + feature { float_list { value: 2.68435763e+09 } } + feature { float_list { value: 2.68435584e+09 } } + feature { float_list { value: 2.14748365e+09 } } + feature { float_list { value: 3.75810074e+09 } } + feature { float_list { value: 2.68435814e+09 } } + feature { float_list { value: 3584 } } + feature { float_list { value: 2.68435814e+09 } } + feature { float_list { value: 3550 } } + feature { float_list { value: 2958 } } + feature { float_list { value: 2974 } } + feature { float_list { value: 2.68435814e+09 } } + feature { float_list { value: 3534 } } + feature { float_list { value: 2.68435866e+09 } } + feature { float_list { value: 3998 } } + feature { float_list { value: 3.75809946e+09 } } + feature { float_list { value: 3.75810074e+09 } } + feature { float_list { value: 4336 } } + feature { float_list { value: 2.68435814e+09 } } + feature { float_list { value: 3.7580969e+09 } } + feature { float_list { value: 2.68435482e+09 } } + feature { float_list { value: 2.68435482e+09 } } + feature { float_list { value: 2.68435814e+09 } } + feature { float_list { value: 2.68435482e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.68435814e+09 } } + feature { float_list { value: 2.14748493e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.68435814e+09 } } + feature { float_list { value: 2.14748493e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.68435456e+09 } } + feature { float_list { value: 2.68435456e+09 } } } } + feature_list { + key: \"reward\" value { + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 36.9006081 } } } } + feature_list { + key: \"stage\" value { + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } + feature { int64_list { value: 0 } } } } + feature_list { + key: \"weight\" value { + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } + feature { float_list { value: 0 } } } } } " + } +} diff --git a/llvm/test/CodeGen/MLRegalloc/dev-mode-prio-logging.ll b/llvm/test/CodeGen/MLRegalloc/dev-mode-prio-logging.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/MLRegalloc/dev-mode-prio-logging.ll @@ -0,0 +1,30 @@ +; REQUIRES: have_tf_api +; REQUIRES: x86_64-linux +; +; Check that we log correctly, both with a learned policy, and the default policy +; +; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-prio-enable-advisor=development \ +; RUN: -regalloc-prio-training-log=%t1 -tfutils-text-log < %S/Inputs/input.ll +; RUN: sed -i 's/ \+/ /g' %t1 +; RUN: sed -i 's/\\n key:/\n key:/g' %t1 +; RUN: sed -i 's/\\n feature/\n feature/g' %t1 +; RUN: sed -i 's/\\n/ /g' %t1 +; RUN: FileCheck --input-file %t1 %s --check-prefixes=CHECK,NOML +; RUN: diff %t1 %S/Inputs/reference-prio-log-noml.txt + +; RUN: rm -rf %t && mkdir %t +; RUN: %python %S/../../../lib/Analysis/models/gen-regalloc-priority-test-model.py %t +; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-prio-enable-advisor=development \ +; RUN: -regalloc-prio-training-log=%t2 -tfutils-text-log -regalloc-prio-model=%t < %S/Inputs/input.ll +; RUN: sed -i 's/ \+/ /g' %t2 +; RUN: sed -i 's/\\n key:/\n key:/g' %t2 +; RUN: sed -i 's/\\n feature/\n feature/g' %t2 +; RUN: sed -i 's/\\n/ /g' %t2 +; RUN: FileCheck --input-file %t2 %s --check-prefixes=CHECK,ML + +; CHECK-NOT: nan +; CHECK-LABEL: key: \"priority\" +; NOML-NEXT: feature { float_list { value: 2.68435814e+09 } } +; ML-NEXT: feature { float_list { value: 3551 } } +; CHECK-LABEL: key: \"reward\" +