[MLGO] ML Regalloc Priority Advisor

[Flpha0830]

Summary

The bulk of the implementation is common between 'release' mode (==AOT-ed
model) and 'development' mode (for training), the main difference is
that in development mode, we may also log features (for training logs),
inject scoring information and then produce the log file.

[mtrofin]

I think we need to separate in a first patch the changes needed to reuse regallocscoring. For those, I think maybe a simple solution right now would be to have the base RegAllocEvictionAdvisorAnalysis and RegAllocPriorityAdvisorAnalysis (the base types) have a virtual void logRewardIfNeeded(float Reward) method. By default it does nothing. Just the Development variant does something. So that way you don't need the additional .h for the scoring, just moving it to a .cpp would be sufficient.

[Flpha0830]

Add logRewardIfNeeded method

[Flpha0830]

Add a model test

[mtrofin]

Overall fine, I think you should still split out the change that introduces the logRewardIfNeeded as a NFC.

For the main change (i.e. this after rebasing with the above NFC), you also need test.

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

386

in the overridden case, drop the virtual.

llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp

122

AU.setPreservesAll(), too.

182

why is this virtual?

197

AU.setPreservesAll()

301

There's a loss here due to the ping-ponging between float and unsigned and then back to float for logging. I imagine we want to stay in float domain for training, though. One way to avoid unnecessary lossiness is to: factor getPriority in the ML case into a getPriorityImpl that returns a float, and then getPriority just casts.

The oddity is that here you'd have to keep track of the "float" priorty as well as the unisgned one, so you'd have a float TrainingPrio and a unsigned Prio; or have Prio be a double and cast it to float when logging and to unsigned when returning, but you still need the getPriorityImpl.

[Flpha0830]

add a getPriorityImpl and update the model test

[mtrofin]

added Yundi for the test model generator, please wait for her feedback before submitting.

On my side, lgtm, left some comments about not trying to enable release mode at this point.

(self-note: there's some boilerplate we can further simplify, I think like ctoring of the dev mode could be more streamlined - I'll take a look into that)

llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp

39

I'd remove this. We don't currently support constructing the release model - that also needs some stuff in CMakeLists.txt. All good, all we need is the development stuff anyway, and all the release stuff is common anyway.

88–90

check LLVM_HAVE_TF_AOT_REGALLOCPRIORITYMODEL here (which nothing sets, which is good right now). We need to do something similar for evictadvisor now, too (not in this patch)

[Flpha0830]

use LLVM_HAVE_TF_AOT_REGALLOCPRIORITYMODEL

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

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

 //===- MLRegAllocEvictAdvisor.cpp - ML eviction 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 eviction advisor and reward injection pass
 //
 //===----------------------------------------------------------------------===//
 
 #include "AllocationOrder.h"
 #include "RegAllocEvictionAdvisor.h"
 #include "RegAllocGreedy.h"
+#include "RegAllocScoring.h"
 #include "llvm/Analysis/MLModelRunner.h"
 #include "llvm/Analysis/TensorSpec.h"
 #if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL) || defined(LLVM_HAVE_TF_API)
 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
 #include "llvm/Analysis/NoInferenceModelRunner.h"
 #include "llvm/Analysis/Utils/TrainingLogger.h"
 #endif
 #include "llvm/Analysis/ReleaseModeModelRunner.h"

 static cl::opt<std::string> ModelUnderTraining(
     "regalloc-model", cl::Hidden,
     cl::desc("The model being trained for register allocation eviction"));
 
 #endif // #ifdef LLVM_HAVE_TF_API
 
 extern cl::opt<unsigned> EvictInterferenceCutoff;
 
-/// The score injection pass.
-/// This pass calculates the score for a function and inserts it in the log, but
-/// this happens only in development mode. It's a no-op otherwise.
-namespace llvm {
-class RegAllocScoring : public MachineFunctionPass {
-public:
-  static char ID;
-
-  RegAllocScoring() : MachineFunctionPass(ID) {
-    initializeRegAllocScoringPass(*PassRegistry::getPassRegistry());
-  }
-
-  ~RegAllocScoring() override = default;
-
-  StringRef getPassName() const override {
-    return "Register Allocation Pass Scoring";
-  }
-
-  /// RegAllocReward analysis usage.
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.setPreservesAll();
-    AU.addRequired<RegAllocEvictionAdvisorAnalysis>();
-    AU.addRequired<MachineBlockFrequencyInfo>();
-    MachineFunctionPass::getAnalysisUsage(AU);
-  }
-
-  /// Performs this pass
-  bool runOnMachineFunction(MachineFunction &) override;
-};
-
-char RegAllocScoring::ID = 0;
-FunctionPass *createRegAllocScoringPass() { return new RegAllocScoring(); }
-
-} // namespace llvm
-
-INITIALIZE_PASS(RegAllocScoring, "regallocscoringpass",
-                "Register Allocation Scoring Pass", false, false)
-
 // ===================================
 // Common ML Advisor declarations
 // ===================================
 namespace {
 // This is the maximum number of interfererring ranges. That's the number of
 // distinct AllocationOrder values, which comes from MCRegisterClass::RegsSize.
 // For X86, that's 32.
 // TODO: find a way to get this, statically, in a programmatic way.
public:
   /// 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:
   std::vector<TensorSpec> InputFeatures;
   std::vector<TensorSpec> TrainingInputFeatures;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<MachineBlockFrequencyInfo>();
     AU.addRequired<MachineLoopInfo>();
     RegAllocEvictionAdvisorAnalysis::getAnalysisUsage(AU);
   }
for (size_t I = 1; I < MUTR->outputLoggedFeatureSpecs().size();
           CurrentFeature,
           reinterpret_cast<const char *>(
               MUTR->lastEvaluationResult()->getUntypedTensorValue(I)));
   // The output is right after the features and the extra outputs
   Log->logInt64Value(CurrentFeature, &Ret);
   return Ret;
 }
 
-bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
+Logger *RegAllocEvictionAdvisorAnalysis::getLogger(
+    const MachineFunction &MF,
+    const RegAllocEvictionAdvisorAnalysis *RegAllocEvictionAdvisorAnalysis) {
   if (auto *DevModeAnalysis = dyn_cast<DevelopmentModeEvictionAdvisorAnalysis>(
-          &getAnalysis<RegAllocEvictionAdvisorAnalysis>()))
+          RegAllocEvictionAdvisorAnalysis))
     if (auto *Log = DevModeAnalysis->getLogger(MF))
-      Log->logFloatFinalReward(static_cast<float>(
-          calculateRegAllocScore(MF, getAnalysis<MachineBlockFrequencyInfo>())
-              .getScore()));
+      return Log;
 
-  return false;
+  return nullptr;
 }
 #endif // #ifdef LLVM_HAVE_TF_API
 
 RegAllocEvictionAdvisorAnalysis *llvm::createReleaseModeAdvisor() {
   return new ReleaseModeEvictionAdvisorAnalysis();
 }
-
-// In all cases except development mode, we don't need scoring.
-#if !defined(LLVM_HAVE_TF_API)
-bool RegAllocScoring::runOnMachineFunction(MachineFunction &) { return false; }
-#endif

llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp

+//===- 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<std::string> TrainingLog(
+    "regalloc-priority-training-log", cl::Hidden,
+    cl::desc("Training log for the register allocator priority model"));
+
+static cl::opt<std::string> 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<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,
+                    SlotIndexes *const Indexes, MLModelRunner *Runner);
+
+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; }
+
+  unsigned getPriority(const LiveInterval &LI) 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 {
+    AU.addRequired<SlotIndexes>();
+    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, &getAnalysis<SlotIndexes>(), Runner.get());
+  }
+  std::unique_ptr<ReleaseModeModelRunner<NoopSavedModelImpl>> Runner;
+};
+
+// ===================================
+// Development mode-specifics
+// ===================================
+//
+// Features we log
+#ifdef LLVM_HAVE_TF_API
+
+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,
+                                 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();
+  }
+
+private:
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<SlotIndexes>();
+    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, &getAnalysis<SlotIndexes>(), Runner.get(), Log);
+  }
+
+  std::unique_ptr<MLModelRunner> Runner;
+  StringMap<std::unique_ptr<Logger>> 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<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>(LI.weight());
+
+  unsigned Prio = static_cast<unsigned>(Runner->evaluate<float>());
+  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<ModelUnderTrainingRunner>(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<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)));
+  }
+
+  float Ret = static_cast<float>(Prio);
+  Log->logFloatValue(CurrentFeature, &Ret);
+
+  return Prio;
+}
+
+Logger *RegAllocPriorityAdvisorAnalysis::getLogger(
+    const MachineFunction &MF,
+    const RegAllocPriorityAdvisorAnalysis *RegAllocPriorityAdvisorAnalysis) {
+  if (auto *DevModeAnalysis = dyn_cast<DevelopmentModePriorityAdvisorAnalysis>(
+          RegAllocPriorityAdvisorAnalysis))
+    if (auto *Log = DevModeAnalysis->getLogger(MF))
+      return Log;
+
+  return nullptr;
+}
+
+#endif // #ifdef LLVM_HAVE_TF_API

llvm/lib/CodeGen/RegAllocEvictionAdvisor.h

 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/CodeGen/Register.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/MC/MCRegister.h"
 #include "llvm/Pass.h"
 
+#ifdef LLVM_HAVE_TF_API
+#include "llvm/Analysis/Utils/TrainingLogger.h"
+#endif // #ifdef LLVM_HAVE_TF_API
 namespace llvm {
 class AllocationOrder;
 class LiveInterval;
 class LiveIntervals;
 class LiveRegMatrix;
 class MachineFunction;
 class MachineRegisterInfo;
 class RegisterClassInfo;
RegAllocEvictionAdvisorAnalysis(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<RegAllocEvictionAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
   AdvisorMode getAdvisorMode() const { return Mode; }
 
+#if defined(LLVM_HAVE_TF_API)
+  static Logger *getLogger(
+      const MachineFunction &MF,
+      const RegAllocEvictionAdvisorAnalysis *RegAllocEvictionAdvisorAnalysis);
+#endif
+
 protected:
   // This analysis preserves everything, and subclasses may have additional
   // requirements.
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
   }
 
 private:

llvm/lib/CodeGen/RegAllocPriorityAdvisor.h

 //===- RegAllocPriorityAdvisor.h - live ranges priority advisor -*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_REGALLOCPRIORITYADVISOR_H
 #define LLVM_CODEGEN_REGALLOCPRIORITYADVISOR_H
 
 #include "RegAllocEvictionAdvisor.h"
 #include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/Pass.h"
 
+#ifdef LLVM_HAVE_TF_API
+#include "llvm/Analysis/Utils/TrainingLogger.h"
+#endif // #ifdef LLVM_HAVE_TF_API
+
 namespace llvm {
 
 class MachineFunction;
 class VirtRegMap;
 class RAGreedy;
 
 /// Interface to the priority advisor, which is responsible for prioritizing
 /// live ranges.
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; }
 
+#if defined(LLVM_HAVE_TF_API)
+  static Logger *getLogger(
+      const MachineFunction &MF,
+      const RegAllocPriorityAdvisorAnalysis *RegAllocPriorityAdvisorAnalysis);
+#endif
+
 protected:
   // This analysis preserves everything, and subclasses may have additional
   // requirements.
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
   }
 
 private:

llvm/lib/CodeGen/RegAllocPriorityAdvisor.cpp

cl::values(
                    "release", "precompiled"),
         clEnumValN(RegAllocPriorityAdvisorAnalysis::AdvisorMode::Development,
                    "development", "for training")));
 
 char RegAllocPriorityAdvisorAnalysis::ID = 0;
 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 {
 public:
   DefaultPriorityAdvisorAnalysis(bool NotAsRequested)
       : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Default),
         NotAsRequested(NotAsRequested) {}
 

 
 template <> Pass *llvm::callDefaultCtor<RegAllocPriorityAdvisorAnalysis>() {
   Pass *Ret = nullptr;
   switch (Mode) {
   case RegAllocPriorityAdvisorAnalysis::AdvisorMode::Default:
     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)
     return Ret;
   return new DefaultPriorityAdvisorAnalysis(/*NotAsRequested*/ true);
 }
 
 StringRef RegAllocPriorityAdvisorAnalysis::getPassName() const {

llvm/lib/CodeGen/RegAllocScoring.h

+//===- RegAllocScoring.h - score injection pass -----------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_REGALLOCSCORING_H
+#define LLVM_CODEGEN_REGALLOCSCORING_H
+
+#include "RegAllocEvictionAdvisor.h"
+#include "RegAllocPriorityAdvisor.h"
+#include "RegAllocScore.h"
+#include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/PassRegistry.h"
+
+using namespace llvm;
+
+/// The score injection pass.
+/// This pass calculates the score for a function and inserts it in the log, but
+/// this happens only in development mode. It's a no-op otherwise.
+namespace llvm {
+
+class RegAllocScoring : public MachineFunctionPass {
+public:
+  static char ID;
+
+  RegAllocScoring() : MachineFunctionPass(ID) {
+    initializeRegAllocScoringPass(*PassRegistry::getPassRegistry());
+  }
+
+  ~RegAllocScoring() override = default;
+
+  StringRef getPassName() const override {
+    return "Register Allocation Pass Scoring";
+  }
+
+  /// RegAllocReward analysis usage.
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+    AU.addRequired<RegAllocEvictionAdvisorAnalysis>();
+    AU.addRequired<RegAllocPriorityAdvisorAnalysis>();
+    AU.addRequired<MachineBlockFrequencyInfo>();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
+  /// Performs this pass
+  bool runOnMachineFunction(MachineFunction &) override;
+};
+
+char RegAllocScoring::ID = 0;
+FunctionPass *createRegAllocScoringPass() { return new RegAllocScoring(); }
+
+} // namespace llvm
+
+INITIALIZE_PASS(RegAllocScoring, "regallocscoringpass",
+                "Register Allocation Scoring Pass", false, false)
+
+#ifdef LLVM_HAVE_TF_API
+
+bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
+  if (auto *Log = RegAllocEvictionAdvisorAnalysis::getLogger(
+          MF, &getAnalysis<RegAllocEvictionAdvisorAnalysis>())) {
+    Log->logFloatFinalReward(static_cast<float>(
+        calculateRegAllocScore(MF, getAnalysis<MachineBlockFrequencyInfo>())
+            .getScore()));
+  }
+
+  if (auto *Log = RegAllocPriorityAdvisorAnalysis::getLogger(
+          MF, &getAnalysis<RegAllocPriorityAdvisorAnalysis>())) {
+    Log->logFloatFinalReward(static_cast<float>(
+        calculateRegAllocScore(MF, getAnalysis<MachineBlockFrequencyInfo>())
+            .getScore()));
+  }
+
+  return false;
+}
+#endif // #ifdef LLVM_HAVE_TF_API
+
+// In all cases except development mode, we don't need scoring.
+#if !defined(LLVM_HAVE_TF_API)
+bool RegAllocScoring::runOnMachineFunction(MachineFunction &) { return false; }
+#endif
+
+#endif // LLVM_CODEGEN_REGALLOCSCORING_H