[NFC][MLGO] ML Regalloc Priority Advisor

[Flpha0830]

Change unsigned to float
The Priority Advisor need the unsigned to become float because ML likes that more in this case

llvm/include/llvm/InitializePasses.h

 void initializePruneEHPass(PassRegistry&);
 void initializeRABasicPass(PassRegistry&);
 void initializePseudoProbeInserterPass(PassRegistry &);
 void initializeRAGreedyPass(PassRegistry&);
 void initializeReachingDefAnalysisPass(PassRegistry&);
 void initializeReassociateLegacyPassPass(PassRegistry&);
 void initializeRedundantDbgInstEliminationPass(PassRegistry&);
 void initializeRegAllocEvictionAdvisorAnalysisPass(PassRegistry &);
+void initializeRegAllocPriorityAdvisorAnalysisPass(PassRegistry &);
 void initializeRegAllocFastPass(PassRegistry&);
 void initializeRegAllocScoringPass(PassRegistry &);
 void initializeRegBankSelectPass(PassRegistry&);
 void initializeRegToMemLegacyPass(PassRegistry&);
 void initializeRegUsageInfoCollectorPass(PassRegistry&);
 void initializeRegUsageInfoPropagationPass(PassRegistry&);
 void initializeRegionInfoPassPass(PassRegistry&);
 void initializeRegionOnlyPrinterPass(PassRegistry&);

llvm/lib/Analysis/models/gen-regalloc-priority-test-model.py

+"""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)

llvm/lib/CodeGen/CMakeLists.txt

add_llvm_component_library(LLVMCodeGen
   MachineSSAUpdater.cpp
   MachineStripDebug.cpp
   MachineTraceMetrics.cpp
   MachineVerifier.cpp
   MIRFSDiscriminator.cpp
   MIRSampleProfile.cpp
   MIRYamlMapping.cpp
   MLRegallocEvictAdvisor.cpp
+  MLRegallocPriorityAdvisor.cpp
   ModuloSchedule.cpp
   MultiHazardRecognizer.cpp
   PatchableFunction.cpp
   MBFIWrapper.cpp
   MIRPrinter.cpp
   MIRPrintingPass.cpp
   MacroFusion.cpp
   NonRelocatableStringpool.cpp
add_llvm_component_library(LLVMCodeGen
   PseudoSourceValue.cpp
   RDFGraph.cpp
   RDFLiveness.cpp
   RDFRegisters.cpp
   ReachingDefAnalysis.cpp
   RegAllocBase.cpp
   RegAllocBasic.cpp
   RegAllocEvictionAdvisor.cpp
+  RegAllocPriorityAdvisor.cpp
   RegAllocFast.cpp
   RegAllocGreedy.cpp
   RegAllocPBQP.cpp
   RegAllocScore.cpp
   RegisterClassInfo.cpp
   RegisterCoalescer.cpp
   RegisterPressure.cpp
   RegisterScavenging.cpp

llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp

+#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<std::string> TrainingLog(
+    "regalloc-prio-training-log", cl::Hidden,
+    cl::desc("Training log for the register allocator priority model"));
+
+static cl::opt<std::string> ModelUnderTraining(
+    "regalloc-prio-model", cl::Hidden,
+    cl::desc("The model being trained for register allocation priority"));
+
+namespace llvm {
+
+static const std::vector<int64_t> 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<const RegAllocPriorityAdvisor &>(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<type>(#name, shape),
+
+static const std::vector<TensorSpec> 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<RegAllocPriorityAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+    if (!Runner)
+      Runner = std::make_unique<ReleaseModeModelRunner<NoopSavedModelImpl>>(
+          MF.getFunction().getContext(), InputFeatures, DecisionName);
+    return std::make_unique<MLPriorityAdvisor>(MF, RA, Runner.get());
+  }
+  std::unique_ptr<ReleaseModeModelRunner<NoopSavedModelImpl>> Runner;
+};
+
+static const TensorSpec Output =
+    TensorSpec::createSpec<float>(DecisionName, {1});
+static const TensorSpec Reward = TensorSpec::createSpec<float>("reward", {1});
+
+#define _DECL_TRAIN_FEATURES(type, name, shape, _)                             \
+  TensorSpec::createSpec<type>(std::string("action_") + #name, shape),
+
+static const std::vector<TensorSpec> TrainingInputFeatures{
+    {RA_PRIORITY_FEATURES_LIST(_DECL_TRAIN_FEATURES)
+         TensorSpec::createSpec<float>("action_discount", {1}),
+     TensorSpec::createSpec<int32_t>("action_step_type", {1}),
+     TensorSpec::createSpec<float>("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<raw_fd_ostream>(TrainingLog, EC);
+    if (EC) {
+      M.getContext().emitError(EC.message() + ":" + TrainingLog);
+      return false;
+    }
+    Logger::flushLogs(*OS, LogMap);
+    return false;
+  }
+
+  std::unique_ptr<RegAllocPriorityAdvisor>
+  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<NoInferenceModelRunner>(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<LoggedFeatureSpec> LFS;
+      for (const auto &FS : InputFeatures)
+        LFS.push_back({FS, None});
+      if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(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<Logger>(LFS, Reward, /*IncludeReward*/ true)));
+      assert(I.second);
+      Log = I.first->second.get();
+    }
+
+    return std::make_unique<DevelopmentModePriorityAdvisor>(MF, RA,
+                                                            Runner.get(), Log);
+  }
+
+  std::unique_ptr<MLModelRunner> Runner;
+  StringMap<std::unique_ptr<Logger>> 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<int64_t>(0) = static_cast<int64_t>(Size);
+  *Runner->getTensor<int64_t>(1) = static_cast<int64_t>(Stage);
+  *Runner->getTensor<float>(2) = static_cast<float>(Weight);
+
+  float Ret = Runner->evaluate<float>();
+  return Ret;
+}
+
+float DevelopmentModePriorityAdvisor::tryFindPriority(unsigned Prio,
+                                                      unsigned Size,
+                                                      LiveRangeStage Stage,
+                                                      float Weight) const {
+  float Ret = 0;
+
+  if (isa<ModelUnderTrainingRunner>(getRunner())) {
+    Ret = MLPriorityAdvisor::tryFindPriority(Prio, Size, Stage, Weight);
+  } else {
+    Ret = static_cast<float>(Prio);
+  }
+
+  if (TrainingLog.empty())
+    return Ret;
+
+  size_t CurrentFeature = 0;
+  for (; CurrentFeature < InputFeatures.size(); ++CurrentFeature) {
+    Log->logSpecifiedTensorValue(
+        CurrentFeature, reinterpret_cast<const char *>(
+                            getRunner().getTensorUntyped(CurrentFeature)));
+  }
+
+  if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner())) {
+    for (size_t I = 1; I < MUTR->outputLoggedFeatureSpecs().size();
+         ++I, ++CurrentFeature)
+      Log->logSpecifiedTensorValue(
+          CurrentFeature,
+          reinterpret_cast<const char *>(
+              MUTR->lastEvaluationResult()->getUntypedTensorValue(I)));
+  }
+  Log->logFloatValue(CurrentFeature, &Ret);
+
+  return Ret;
+}

llvm/lib/CodeGen/RegAllocGreedy.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_REGALLOCGREEDY_H_
 #define LLVM_CODEGEN_REGALLOCGREEDY_H_
 
 #include "InterferenceCache.h"
 #include "RegAllocBase.h"
 #include "RegAllocEvictionAdvisor.h"
+#include "RegAllocPriorityAdvisor.h"
 #include "SpillPlacement.h"
 #include "SplitKit.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/IndexedMap.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #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"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
 #include "llvm/CodeGen/Spiller.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
public:
   VirtRegMap *getVirtRegMap() const { return VRM; }
   const RegisterClassInfo &getRegClassInfo() const { return RegClassInfo; }
   const ExtraRegInfo &getExtraInfo() const { return *ExtraInfo; }
   size_t getQueueSize() const { return Queue.size(); }
   // end (interface to eviction advisers)
 
 private:
   // Convenient shortcuts.
-  using PQueue = std::priority_queue<std::pair<unsigned, unsigned>>;
+  using PQueue = std::priority_queue<std::pair<float, unsigned>>;
   using SmallLISet = SmallPtrSet<const LiveInterval *, 4>;
 
   // We need to track all tentative recolorings so we can roll back any
   // successful and unsuccessful recoloring attempts.
   using RecoloringStack =
       SmallVector<std::pair<const LiveInterval *, MCRegister>, 8>;
 
   // context
private:
 
   // state
   std::unique_ptr<Spiller> SpillerInstance;
   PQueue Queue;
   std::unique_ptr<VirtRegAuxInfo> VRAI;
   Optional<ExtraRegInfo> ExtraInfo;
   std::unique_ptr<RegAllocEvictionAdvisor> EvictAdvisor;
 
+  std::unique_ptr<RegAllocPriorityAdvisor> 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.
   enum CutOffStage {
     // No cutoffs encountered
     CO_None = 0,
 
     // lcr-max-depth cutoff encountered

llvm/lib/CodeGen/RegAllocGreedy.cpp

 //===----------------------------------------------------------------------===//
 
 #include "RegAllocGreedy.h"
 #include "AllocationOrder.h"
 #include "InterferenceCache.h"
 #include "LiveDebugVariables.h"
 #include "RegAllocBase.h"
 #include "RegAllocEvictionAdvisor.h"
+#include "RegAllocScore.h"
 #include "SpillPlacement.h"
 #include "SplitKit.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/IndexedMap.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #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"
 #include "llvm/CodeGen/LiveInterval.h"
 #include "llvm/CodeGen/LiveIntervalUnion.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/LiveRangeEdit.h"
 #include "llvm/CodeGen/LiveRegMatrix.h"

 INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
 INITIALIZE_PASS_DEPENDENCY(VirtRegMap)
 INITIALIZE_PASS_DEPENDENCY(LiveRegMatrix)
 INITIALIZE_PASS_DEPENDENCY(EdgeBundles)
 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)
 
 #ifndef NDEBUG
 const char *const RAGreedy::StageName[] = {
     "RS_New",
     "RS_Assign",
     "RS_Split",
void RAGreedy::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<VirtRegMap>();
   AU.addPreserved<VirtRegMap>();
   AU.addRequired<LiveRegMatrix>();
   AU.addPreserved<LiveRegMatrix>();
   AU.addRequired<EdgeBundles>();
   AU.addRequired<SpillPlacement>();
   AU.addRequired<MachineOptimizationRemarkEmitterPass>();
   AU.addRequired<RegAllocEvictionAdvisorAnalysis>();
+  AU.addRequired<RegAllocPriorityAdvisorAnalysis>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
 //===----------------------------------------------------------------------===//
 //                     LiveRangeEdit delegate methods
 //===----------------------------------------------------------------------===//
 
 bool RAGreedy::LRE_CanEraseVirtReg(Register VirtReg) {
if (Stage == RS_Split) {
 
     // Mark a higher bit to prioritize global and local above RS_Split.
     Prio |= (1u << 31);
 
     // Boost ranges that have a physical register hint.
     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); }
 
 const LiveInterval *RAGreedy::dequeue(PQueue &CurQueue) {
   if (CurQueue.empty())
     return nullptr;
   LiveInterval *LI = &LIS->getInterval(~CurQueue.top().second);
LLVM_DEBUG(dbgs() << "********** GREEDY REGISTER ALLOCATION **********\n"
                     << "********** Function: " << mf.getName() << '\n');
 
   MF = &mf;
   TII = MF->getSubtarget().getInstrInfo();
 
   if (VerifyEnabled)
     MF->verify(this, "Before greedy register allocator");
 
+  PriorityAdvisor =
+      getAnalysis<RegAllocPriorityAdvisorAnalysis>().getAdvisor(*MF, *this);
+
   RegAllocBase::init(getAnalysis<VirtRegMap>(),
                      getAnalysis<LiveIntervals>(),
                      getAnalysis<LiveRegMatrix>());
   Indexes = &getAnalysis<SlotIndexes>();
   MBFI = &getAnalysis<MachineBlockFrequencyInfo>();
   DomTree = &getAnalysis<MachineDominatorTree>();
   ORE = &getAnalysis<MachineOptimizationRemarkEmitterPass>().getORE();
   Loops = &getAnalysis<MachineLoopInfo>();
bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
   allocatePhysRegs();
   tryHintsRecoloring();
 
   if (VerifyEnabled)
     MF->verify(this, "Before post optimization");
   postOptimization();
   reportStats();
 
+  PriorityAdvisor->logReward(static_cast<float>(
+      calculateRegAllocScore(mf, getAnalysis<MachineBlockFrequencyInfo>(),
+                             getAnalysis<AAResultsWrapperPass>().getAAResults())
+          .getScore()));
+
   releaseMemory();
   return true;
 }

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

+#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<RegAllocPriorityAdvisor>
+  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>();
+
+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

llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp

+#include "RegAllocGreedy.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+
+using namespace llvm;
+
+static cl::opt<RegAllocPriorityAdvisorAnalysis::AdvisorMode> 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<RegAllocPriorityAdvisor>
+  getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+    return std::make_unique<DefaultPriorityAdvisor>(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<RegAllocPriorityAdvisorAnalysis>() {
+  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<float>(Prio);
+}
 No newline at end of file

llvm/test/CodeGen/MLRegalloc/Inputs/reference-prio-log-noml.txt

+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  }  }  }  } } "
+ }
+}

llvm/test/CodeGen/MLRegalloc/dev-mode-prio-logging.ll

+; 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\"
+