llvm/lib/CodeGen/MLRegAllocPriorityAdvisor.cpp

   1 //===- MLRegAllocPriorityAdvisor.cpp - ML priority advisor-----------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // Implementation of the ML priority advisor and reward injection pass
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #include "AllocationOrder.h"
  14 #include "RegAllocGreedy.h"
  15 #include "RegAllocPriorityAdvisor.h"
  16 #include "llvm/Analysis/AliasAnalysis.h"
  17 #include "llvm/Analysis/InteractiveModelRunner.h"
  18 #include "llvm/Analysis/MLModelRunner.h"
  19 #include "llvm/Analysis/ReleaseModeModelRunner.h"
  20 #include "llvm/Analysis/TensorSpec.h"
  21 #include "llvm/CodeGen/CalcSpillWeights.h"
  22 #include "llvm/CodeGen/LiveRegMatrix.h"
  23 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
  24 #include "llvm/CodeGen/MachineFunction.h"
  25 #include "llvm/CodeGen/MachineLoopInfo.h"
  26 #include "llvm/CodeGen/MachineRegisterInfo.h"
  27 #include "llvm/CodeGen/Passes.h"
  28 #include "llvm/CodeGen/RegisterClassInfo.h"
  29 #include "llvm/CodeGen/SlotIndexes.h"
  30 #include "llvm/CodeGen/VirtRegMap.h"
  31 #include "llvm/InitializePasses.h"
  32 #include "llvm/Pass.h"
  33 #include "llvm/PassRegistry.h"
  34 #include "llvm/Support/CommandLine.h"
  35
  36 #if defined(LLVM_HAVE_TFLITE)
  37 #include "llvm/Analysis/ModelUnderTrainingRunner.h"
  38 #include "llvm/Analysis/NoInferenceModelRunner.h"
  39 #include "llvm/Analysis/Utils/TrainingLogger.h"
  40 #endif
  41
  42 using namespace llvm;
  43
  44 static cl::opt<std::string> InteractiveChannelBaseName(
  45     "regalloc-priority-interactive-channel-base", cl::Hidden,
  46     cl::desc(
  47         "Base file path for the interactive mode. The incoming filename should "
  48         "have the name <regalloc-priority-interactive-channel-base>.in, while "
  49         "the outgoing name should be "
  50         "<regalloc-priority-interactive-channel-base>.out"));
  51
  52 using CompiledModelType = NoopSavedModelImpl;
  53
  54 // Options that only make sense in development mode
  55 #ifdef LLVM_HAVE_TFLITE
  56 #include "RegAllocScore.h"
  57 #include "llvm/Analysis/Utils/TFUtils.h"
  58
  59 static cl::opt<std::string> TrainingLog(
  60     "regalloc-priority-training-log", cl::Hidden,
  61     cl::desc("Training log for the register allocator priority model"));
  62
  63 static cl::opt<std::string> ModelUnderTraining(
  64     "regalloc-priority-model", cl::Hidden,
  65     cl::desc("The model being trained for register allocation priority"));
  66
  67 #endif // #ifdef LLVM_HAVE_TFLITE
  68
  69 namespace llvm {
  70
  71 static const std::vector<int64_t> PerLiveRangeShape{1};
  72
  73 #define RA_PRIORITY_FEATURES_LIST(M)                                           \
  74   M(int64_t, li_size, PerLiveRangeShape, "size")                               \
  75   M(int64_t, stage, PerLiveRangeShape, "stage")                                \
  76   M(float, weight, PerLiveRangeShape, "weight")
  77
  78 #define DecisionName "priority"
  79 static const TensorSpec DecisionSpec =
  80     TensorSpec::createSpec<float>(DecisionName, {1});
  81
  82
  83 // Named features index.
  84 enum FeatureIDs {
  85 #define _FEATURE_IDX(_, name, __, ___) name,
  86   RA_PRIORITY_FEATURES_LIST(_FEATURE_IDX)
  87 #undef _FEATURE_IDX
  88       FeatureCount
  89 };
  90
  91 class MLPriorityAdvisor : public RegAllocPriorityAdvisor {
  92 public:
  93   MLPriorityAdvisor(const MachineFunction &MF, const RAGreedy &RA,
  94                     SlotIndexes *const Indexes, MLModelRunner *Runner);
  95
  96 protected:
  97   const RegAllocPriorityAdvisor &getDefaultAdvisor() const {
  98     return static_cast<const RegAllocPriorityAdvisor &>(DefaultAdvisor);
  99   }
 100
 101   // The assumption is that if the Runner could not be constructed, we emit-ed
 102   // error, and we shouldn't be asking for it here.
 103   const MLModelRunner &getRunner() const { return *Runner; }
 104   float getPriorityImpl(const LiveInterval &LI) const;
 105   unsigned getPriority(const LiveInterval &LI) const override;
 106
 107 private:
 108   const DefaultPriorityAdvisor DefaultAdvisor;
 109   MLModelRunner *const Runner;
 110 };
 111
 112 #define _DECL_FEATURES(type, name, shape, _)                                   \
 113   TensorSpec::createSpec<type>(#name, shape),
 114
 115 static const std::vector<TensorSpec> InputFeatures{
 116     {RA_PRIORITY_FEATURES_LIST(_DECL_FEATURES)},
 117 };
 118 #undef _DECL_FEATURES
 119
 120 // ===================================
 121 // Release (AOT) - specifics
 122 // ===================================
 123 class ReleaseModePriorityAdvisorAnalysis final
 124     : public RegAllocPriorityAdvisorAnalysis {
 125 public:
 126   ReleaseModePriorityAdvisorAnalysis()
 127       : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Release) {}
 128   // support for isa<> and dyn_cast.
 129   static bool classof(const RegAllocPriorityAdvisorAnalysis *R) {
 130     return R->getAdvisorMode() == AdvisorMode::Release;
 131   }
 132
 133 private:
 134   void getAnalysisUsage(AnalysisUsage &AU) const override {
 135     AU.setPreservesAll();
 136     AU.addRequired<SlotIndexes>();
 137     RegAllocPriorityAdvisorAnalysis::getAnalysisUsage(AU);
 138   }
 139
 140   std::unique_ptr<RegAllocPriorityAdvisor>
 141   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
 142     if (!Runner) {
 143       if (InteractiveChannelBaseName.empty())
 144         Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
 145             MF.getFunction().getContext(), InputFeatures, DecisionName);
 146       else
 147         Runner = std::make_unique<InteractiveModelRunner>(
 148             MF.getFunction().getContext(), InputFeatures, DecisionSpec,
 149             InteractiveChannelBaseName + ".out",
 150             InteractiveChannelBaseName + ".in");
 151     }
 152     return std::make_unique<MLPriorityAdvisor>(
 153         MF, RA, &getAnalysis<SlotIndexes>(), Runner.get());
 154   }
 155   std::unique_ptr<MLModelRunner> Runner;
 156 };
 157
 158 // ===================================
 159 // Development mode-specifics
 160 // ===================================
 161 //
 162 // Features we log
 163 #ifdef LLVM_HAVE_TFLITE
 164 static const TensorSpec Reward = TensorSpec::createSpec<float>("reward", {1});
 165
 166 #define _DECL_TRAIN_FEATURES(type, name, shape, _)                             \
 167   TensorSpec::createSpec<type>(std::string("action_") + #name, shape),
 168
 169 static const std::vector<TensorSpec> TrainingInputFeatures{
 170     {RA_PRIORITY_FEATURES_LIST(_DECL_TRAIN_FEATURES)
 171          TensorSpec::createSpec<float>("action_discount", {1}),
 172      TensorSpec::createSpec<int32_t>("action_step_type", {1}),
 173      TensorSpec::createSpec<float>("action_reward", {1})}};
 174 #undef _DECL_TRAIN_FEATURES
 175
 176 class DevelopmentModePriorityAdvisor : public MLPriorityAdvisor {
 177 public:
 178   DevelopmentModePriorityAdvisor(const MachineFunction &MF, const RAGreedy &RA,
 179                                  SlotIndexes *const Indexes,
 180                                  MLModelRunner *Runner, Logger *Log)
 181       : MLPriorityAdvisor(MF, RA, Indexes, Runner), Log(Log) {}
 182
 183 private:
 184   unsigned getPriority(const LiveInterval &LI) const override;
 185   Logger *const Log;
 186 };
 187
 188 class DevelopmentModePriorityAdvisorAnalysis final
 189     : public RegAllocPriorityAdvisorAnalysis {
 190 public:
 191   DevelopmentModePriorityAdvisorAnalysis()
 192       : RegAllocPriorityAdvisorAnalysis(AdvisorMode::Development) {}
 193   // support for isa<> and dyn_cast.
 194   static bool classof(const RegAllocPriorityAdvisorAnalysis *R) {
 195     return R->getAdvisorMode() == AdvisorMode::Development;
 196   }
 197
 198   void logRewardIfNeeded(const MachineFunction &MF,
 199                          llvm::function_ref<float()> GetReward) override {
 200     if (!Log || !Log->hasAnyObservationForContext(MF.getName()))
 201       return;
 202     // The function pass manager would run all the function passes for a
 203     // function, so we assume the last context belongs to this function. If
 204     // this invariant ever changes, we can implement at that time switching
 205     // contexts. At this point, it'd be an error
 206     if (Log->currentContext() != MF.getName()) {
 207       MF.getFunction().getContext().emitError(
 208           "The training log context shouldn't have had changed.");
 209     }
 210     if (Log->hasObservationInProgress())
 211       Log->logReward<float>(GetReward());
 212   }
 213
 214 private:
 215   void getAnalysisUsage(AnalysisUsage &AU) const override {
 216     AU.setPreservesAll();
 217     AU.addRequired<SlotIndexes>();
 218     RegAllocPriorityAdvisorAnalysis::getAnalysisUsage(AU);
 219   }
 220
 221   // Save all the logs (when requested).
 222   bool doInitialization(Module &M) override {
 223     LLVMContext &Ctx = M.getContext();
 224     if (ModelUnderTraining.empty() && TrainingLog.empty()) {
 225       Ctx.emitError("Regalloc development mode should be requested with at "
 226                     "least logging enabled and/or a training model");
 227       return false;
 228     }
 229     if (ModelUnderTraining.empty())
 230       Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
 231     else
 232       Runner = ModelUnderTrainingRunner::createAndEnsureValid(
 233           Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
 234     if (!Runner) {
 235       Ctx.emitError("Regalloc: could not set up the model runner");
 236       return false;
 237     }
 238     if (TrainingLog.empty())
 239       return false;
 240     std::error_code EC;
 241     auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
 242     if (EC) {
 243       M.getContext().emitError(EC.message() + ":" + TrainingLog);
 244       return false;
 245     }
 246     std::vector<TensorSpec> LFS = InputFeatures;
 247     if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
 248       append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
 249     // We always log the output; in particular, if we're not evaluating, we
 250     // don't have an output spec json file. That's why we handle the
 251     // 'normal' output separately.
 252     LFS.push_back(DecisionSpec);
 253
 254     Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
 255                                    /*IncludeReward*/ true);
 256     return false;
 257   }
 258
 259   std::unique_ptr<RegAllocPriorityAdvisor>
 260   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
 261     if (!Runner)
 262       return nullptr;
 263     if (Log) {
 264       Log->switchContext(MF.getName());
 265     }
 266
 267     return std::make_unique<DevelopmentModePriorityAdvisor>(
 268         MF, RA, &getAnalysis<SlotIndexes>(), Runner.get(), Log.get());
 269   }
 270
 271   std::unique_ptr<MLModelRunner> Runner;
 272   std::unique_ptr<Logger> Log;
 273 };
 274 #endif //#ifdef LLVM_HAVE_TFLITE
 275
 276 } // namespace llvm
 277
 278 RegAllocPriorityAdvisorAnalysis *llvm::createReleaseModePriorityAdvisor() {
 279   return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
 280                  !InteractiveChannelBaseName.empty()
 281              ? new ReleaseModePriorityAdvisorAnalysis()
 282              : nullptr;
 283 }
 284
 285 MLPriorityAdvisor::MLPriorityAdvisor(const MachineFunction &MF,
 286                                      const RAGreedy &RA,
 287                                      SlotIndexes *const Indexes,
 288                                      MLModelRunner *Runner)
 289     : RegAllocPriorityAdvisor(MF, RA, Indexes), DefaultAdvisor(MF, RA, Indexes),
 290       Runner(std::move(Runner)) {
 291   assert(this->Runner);
 292   Runner->switchContext(MF.getName());
 293 }
 294
 295 float MLPriorityAdvisor::getPriorityImpl(const LiveInterval &LI) const {
 296   const unsigned Size = LI.getSize();
 297   LiveRangeStage Stage = RA.getExtraInfo().getStage(LI);
 298
 299   *Runner->getTensor<int64_t>(0) = static_cast<int64_t>(Size);
 300   *Runner->getTensor<int64_t>(1) = static_cast<int64_t>(Stage);
 301   *Runner->getTensor<float>(2) = static_cast<float>(LI.weight());
 302
 303   return Runner->evaluate<float>();
 304 }
 305
 306 unsigned MLPriorityAdvisor::getPriority(const LiveInterval &LI) const {
 307   return static_cast<unsigned>(getPriorityImpl(LI));
 308 }
 309
 310 #ifdef LLVM_HAVE_TFLITE
 311 RegAllocPriorityAdvisorAnalysis *llvm::createDevelopmentModePriorityAdvisor() {
 312   return new DevelopmentModePriorityAdvisorAnalysis();
 313 }
 314
 315 unsigned
 316 DevelopmentModePriorityAdvisor::getPriority(const LiveInterval &LI) const {
 317   double Prio = 0;
 318
 319   if (isa<ModelUnderTrainingRunner>(getRunner())) {
 320     Prio = MLPriorityAdvisor::getPriorityImpl(LI);
 321   } else {
 322     Prio = getDefaultAdvisor().getPriority(LI);
 323   }
 324
 325   if (TrainingLog.empty())
 326     return Prio;
 327
 328   // TODO(mtrofin): when we support optional rewards, this can go away. In the
 329   // meantime, we log the "pretend" reward (0) for the previous observation
 330   // before starting a new one.
 331   if (Log->hasObservationInProgress())
 332     Log->logReward<float>(0.0);
 333
 334   Log->startObservation();
 335   size_t CurrentFeature = 0;
 336   for (; CurrentFeature < InputFeatures.size(); ++CurrentFeature) {
 337     Log->logTensorValue(CurrentFeature,
 338                         reinterpret_cast<const char *>(
 339                             getRunner().getTensorUntyped(CurrentFeature)));
 340   }
 341
 342   if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner())) {
 343     for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size();
 344          ++I, ++CurrentFeature)
 345       Log->logTensorValue(
 346           CurrentFeature,
 347           reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I)));
 348   }
 349
 350   float Ret = static_cast<float>(Prio);
 351   Log->logTensorValue(CurrentFeature, reinterpret_cast<const char *>(&Ret));
 352   Log->endObservation();
 353
 354   return static_cast<unsigned>(Prio);
 355 }
 356
 357 #endif // #ifdef LLVM_HAVE_TFLITE