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 358