1 //===- MLRegAllocEvictAdvisor.cpp - ML eviction 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 eviction advisor and reward injection pass 10 // 11 //===----------------------------------------------------------------------===// 12 13 #include "llvm/Config/config.h" 14 15 #if defined(LLVM_HAVE_TF_AOT) || defined(LLVM_HAVE_TF_API) 16 17 #include "RegAllocEvictionAdvisor.h" 18 #include "llvm/Analysis/MLModelRunner.h" 19 #include "llvm/Analysis/ModelUnderTrainingRunner.h" 20 #include "llvm/Analysis/NoInferenceModelRunner.h" 21 #include "llvm/Analysis/Utils/TFUtils.h" 22 #include "llvm/CodeGen/CalcSpillWeights.h" 23 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h" 24 #include "llvm/CodeGen/MachineFunction.h" 25 #include "llvm/CodeGen/MachineLoopInfo.h" 26 #include "llvm/CodeGen/RegisterClassInfo.h" 27 #include "llvm/CodeGen/VirtRegMap.h" 28 #include "llvm/InitializePasses.h" 29 #include "llvm/Pass.h" 30 #include "llvm/PassRegistry.h" 31 #include "llvm/Support/CommandLine.h" 32 #include "llvm/Support/ErrorHandling.h" 33 #include "llvm/Target/TargetMachine.h" 34 35 #include <memory> 36 37 using namespace llvm; 38 39 #define DEBUG_TYPE "ml-regalloc" 40 41 namespace { 42 // This is the maximum number of interfererring ranges. That's the number of 43 // distinct AllocationOrder values, which comes from MCRegisterClass::RegsSize. 44 // For X86, that's 32. 45 // TODO: find a way to get this, statically, in a programmatic way. 46 static const int64_t MaxInterferences = 32; 47 48 // Logically, we can think of the feature set given to the evaluator as a 2D 49 // matrix. The rows are the features (see next). The columns correspond to the 50 // interferences. We treat the candidate virt reg as an 'interference', too, as 51 // its feature set is the same as that of the interferring ranges. So we'll have 52 // MaxInterferences + 1 columns and by convention, we will use the last column 53 // for the virt reg seeking allocation. 54 static const int64_t CandidateVirtRegPos = MaxInterferences; 55 static const int64_t NumberOfInterferences = CandidateVirtRegPos + 1; 56 57 // Most features are as described above, so we'll reuse this vector in defining 58 // them. 59 static const std::vector<int64_t> PerLiveRangeShape{1, NumberOfInterferences}; 60 61 // -------------- 62 // Features table 63 // -------------- 64 // For each interfering live range (incl. the candidate) we collect a number of 65 // features. However, because the features are of different types (and because 66 // of ML best practices), we organize the tensors per feature, not per 67 // candidate. Each such tensor has a scalar value corresponding to the 68 // interferring live range at that position, in the order in AllocationOrder. 69 // The last position corresponds to the virt reg seeking allocation. 70 // Exception to all that is the progression feature, which is just a scalar (see 71 // its documentation for details). 72 // Note on naming: the "_by_max" are normalized using the largest value of that 73 // tensor, as observed in the current decision making stage (i.e. for the 74 // current call to the advisor's tryFindEvictionCandidate) 75 // 76 // The feature list format: type, name, shape, documentation. 77 // Note: we can really just use int64 and float, hence the modeling of some 78 // bools as int64 values. 79 #define RA_EVICT_FEATURES_LIST(M) \ 80 M(int64_t, mask, PerLiveRangeShape, \ 81 "boolean values, 0 for unavailable candidates (i.e. if a position is 0, " \ 82 "it " \ 83 "can't be evicted)") \ 84 M(int64_t, is_free, PerLiveRangeShape, \ 85 "boolean values, 1 if this phys reg is actually free (no interferences)") \ 86 M(float, nr_urgent, PerLiveRangeShape, \ 87 "number of 'urgent' intervals, normalized. Urgent are those that are OK " \ 88 "to break cascades") \ 89 M(float, nr_broken_hints, PerLiveRangeShape, \ 90 "if this position were evicted, how many broken hints would there be") \ 91 M(int64_t, is_hint, PerLiveRangeShape, \ 92 "is this a preferred phys reg for the candidate") \ 93 M(int64_t, is_local, PerLiveRangeShape, \ 94 "is this live range local to a basic block") \ 95 M(float, nr_rematerializable, PerLiveRangeShape, \ 96 "nr rematerializable ranges") \ 97 M(float, nr_defs_and_uses, PerLiveRangeShape, \ 98 "bb freq - weighed nr defs and uses") \ 99 M(float, weighed_reads_by_max, PerLiveRangeShape, \ 100 "bb freq - weighed nr of reads, normalized") \ 101 M(float, weighed_writes_by_max, PerLiveRangeShape, \ 102 "bb feq - weighed nr of writes, normalized") \ 103 M(float, weighed_read_writes_by_max, PerLiveRangeShape, \ 104 "bb freq - weighed nr of uses that are both read and writes, normalized") \ 105 M(float, weighed_indvars_by_max, PerLiveRangeShape, \ 106 "bb freq - weighed nr of uses that are indvars, normalized") \ 107 M(float, hint_weights_by_max, PerLiveRangeShape, \ 108 "bb freq - weighed nr of uses that are hints, normalized") \ 109 M(float, start_bb_freq_by_max, PerLiveRangeShape, \ 110 "the freq in the start block, normalized") \ 111 M(float, end_bb_freq_by_max, PerLiveRangeShape, \ 112 "freq of end block, normalized") \ 113 M(float, hottest_bb_freq_by_max, PerLiveRangeShape, \ 114 "hottest BB freq, normalized") \ 115 M(float, liverange_size, PerLiveRangeShape, \ 116 "size (instr index diff) of the LR") \ 117 M(float, use_def_density, PerLiveRangeShape, \ 118 "the max weight, as computed by the manual heuristic") \ 119 M(int64_t, max_stage, PerLiveRangeShape, \ 120 "largest stage of an interval in this LR") \ 121 M(int64_t, min_stage, PerLiveRangeShape, \ 122 "lowest stage of an interval in this LR") \ 123 M(float, progress, {1}, "ratio of current queue size to initial size") 124 125 // The model learns to pick one of the mask == 1 interferences. This is the name 126 // of the output tensor. 127 // The contract with the model is that the output will be guaranteed to be to a 128 // mask == 1 position. 129 const char *const DecisionName = "index_to_evict"; 130 131 // Named features index. 132 enum FeatureIDs { 133 #define _FEATURE_IDX(_, name, __, ___) name, 134 RA_EVICT_FEATURES_LIST(_FEATURE_IDX) 135 #undef _FEATURE_IDX 136 FeatureCount 137 }; 138 139 // The ML advisor will typically have a sparse input to the evaluator, because 140 // various phys regs won't be available. It's easier (maintenance-wise) to 141 // bulk-reset the state of the evaluator each time we are about to use it again. 142 template <typename T> size_t getTotalSize(const std::vector<int64_t> &Shape) { 143 size_t Ret = sizeof(T); 144 for (const auto V : Shape) 145 Ret *= V; 146 return Ret; 147 } 148 149 void resetInputs(MLModelRunner &Runner) { 150 #define _RESET(TYPE, NAME, SHAPE, __) \ 151 std::memset(Runner.getTensorUntyped(FeatureIDs::NAME), 0, \ 152 getTotalSize<TYPE>(SHAPE)); 153 RA_EVICT_FEATURES_LIST(_RESET) 154 #undef _RESET 155 } 156 157 // Development mode-specifics 158 #ifdef LLVM_HAVE_TF_API 159 #define _DECL_FEATURES(type, name, shape, _) \ 160 TensorSpec::createSpec<type>(#name, shape), 161 162 static const std::vector<TensorSpec> InputFeatures{ 163 {RA_EVICT_FEATURES_LIST(_DECL_FEATURES)}}; 164 #undef _DECL_FEATURES 165 static const TensorSpec Output = 166 TensorSpec::createSpec<int64_t>(DecisionName, {1}); 167 const char *const RewardName = "reward"; 168 static const TensorSpec Reward = 169 TensorSpec::createSpec<int64_t>(RewardName, {1}); 170 171 #endif //#ifdef LLVM_HAVE_TF_API 172 } // namespace 173 #endif // defined(LLVM_HAVE_TF_AOT) || defined(LLVM_HAVE_TF_API) 174