1 //===-- WebAssemblyRegColoring.cpp - Register coloring --------------------===// 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 /// \file 10 /// This file implements a virtual register coloring pass. 11 /// 12 /// WebAssembly doesn't have a fixed number of registers, but it is still 13 /// desirable to minimize the total number of registers used in each function. 14 /// 15 /// This code is modeled after lib/CodeGen/StackSlotColoring.cpp. 16 /// 17 //===----------------------------------------------------------------------===// 18 19 #include "WebAssembly.h" 20 #include "WebAssemblyMachineFunctionInfo.h" 21 #include "llvm/CodeGen/LiveIntervals.h" 22 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h" 23 #include "llvm/CodeGen/MachineRegisterInfo.h" 24 #include "llvm/CodeGen/Passes.h" 25 #include "llvm/Support/Debug.h" 26 #include "llvm/Support/raw_ostream.h" 27 using namespace llvm; 28 29 #define DEBUG_TYPE "wasm-reg-coloring" 30 31 namespace { 32 class WebAssemblyRegColoring final : public MachineFunctionPass { 33 public: 34 static char ID; // Pass identification, replacement for typeid 35 WebAssemblyRegColoring() : MachineFunctionPass(ID) {} 36 37 StringRef getPassName() const override { 38 return "WebAssembly Register Coloring"; 39 } 40 41 void getAnalysisUsage(AnalysisUsage &AU) const override { 42 AU.setPreservesCFG(); 43 AU.addRequired<LiveIntervals>(); 44 AU.addRequired<MachineBlockFrequencyInfo>(); 45 AU.addPreserved<MachineBlockFrequencyInfo>(); 46 AU.addPreservedID(MachineDominatorsID); 47 MachineFunctionPass::getAnalysisUsage(AU); 48 } 49 50 bool runOnMachineFunction(MachineFunction &MF) override; 51 52 private: 53 }; 54 } // end anonymous namespace 55 56 char WebAssemblyRegColoring::ID = 0; 57 INITIALIZE_PASS(WebAssemblyRegColoring, DEBUG_TYPE, 58 "Minimize number of registers used", false, false) 59 60 FunctionPass *llvm::createWebAssemblyRegColoring() { 61 return new WebAssemblyRegColoring(); 62 } 63 64 // Compute the total spill weight for VReg. 65 static float computeWeight(const MachineRegisterInfo *MRI, 66 const MachineBlockFrequencyInfo *MBFI, 67 unsigned VReg) { 68 float Weight = 0.0f; 69 for (MachineOperand &MO : MRI->reg_nodbg_operands(VReg)) 70 Weight += LiveIntervals::getSpillWeight(MO.isDef(), MO.isUse(), MBFI, 71 *MO.getParent()); 72 return Weight; 73 } 74 75 // Create a map of "Register -> vector of <SlotIndex, DBG_VALUE>". 76 // The SlotIndex is the slot index of the next non-debug instruction or the end 77 // of a BB, because DBG_VALUE's don't have slot index themselves. 78 // Adapted from RegisterCoalescer::buildVRegToDbgValueMap. 79 static DenseMap<Register, std::vector<std::pair<SlotIndex, MachineInstr *>>> 80 buildVRegToDbgValueMap(MachineFunction &MF, const LiveIntervals *Liveness) { 81 DenseMap<Register, std::vector<std::pair<SlotIndex, MachineInstr *>>> 82 DbgVRegToValues; 83 const SlotIndexes *Slots = Liveness->getSlotIndexes(); 84 SmallVector<MachineInstr *, 8> ToInsert; 85 86 // After collecting a block of DBG_VALUEs into ToInsert, enter them into the 87 // map. 88 auto CloseNewDVRange = [&DbgVRegToValues, &ToInsert](SlotIndex Slot) { 89 for (auto *X : ToInsert) { 90 for (const auto &Op : X->debug_operands()) { 91 if (Op.isReg() && Op.getReg().isVirtual()) 92 DbgVRegToValues[Op.getReg()].push_back({Slot, X}); 93 } 94 } 95 96 ToInsert.clear(); 97 }; 98 99 // Iterate over all instructions, collecting them into the ToInsert vector. 100 // Once a non-debug instruction is found, record the slot index of the 101 // collected DBG_VALUEs. 102 for (auto &MBB : MF) { 103 SlotIndex CurrentSlot = Slots->getMBBStartIdx(&MBB); 104 105 for (auto &MI : MBB) { 106 if (MI.isDebugValue()) { 107 if (any_of(MI.debug_operands(), [](const MachineOperand &MO) { 108 return MO.isReg() && MO.getReg().isVirtual(); 109 })) 110 ToInsert.push_back(&MI); 111 } else if (!MI.isDebugOrPseudoInstr()) { 112 CurrentSlot = Slots->getInstructionIndex(MI); 113 CloseNewDVRange(CurrentSlot); 114 } 115 } 116 117 // Close range of DBG_VALUEs at the end of blocks. 118 CloseNewDVRange(Slots->getMBBEndIdx(&MBB)); 119 } 120 121 // Sort all DBG_VALUEs we've seen by slot number. 122 for (auto &Pair : DbgVRegToValues) 123 llvm::sort(Pair.second); 124 return DbgVRegToValues; 125 } 126 127 // After register coalescing, some DBG_VALUEs will be invalid. Set them undef. 128 // This function has to run before the actual coalescing, i.e., the register 129 // changes. 130 static void undefInvalidDbgValues( 131 const LiveIntervals *Liveness, 132 const ArrayRef<SmallVector<LiveInterval *, 4>> &Assignments, 133 DenseMap<Register, std::vector<std::pair<SlotIndex, MachineInstr *>>> 134 &DbgVRegToValues) { 135 #ifndef NDEBUG 136 DenseSet<Register> SeenRegs; 137 #endif 138 for (size_t I = 0, E = Assignments.size(); I < E; ++I) { 139 const auto &CoalescedIntervals = Assignments[I]; 140 if (CoalescedIntervals.empty()) 141 continue; 142 for (LiveInterval *LI : CoalescedIntervals) { 143 Register Reg = LI->reg(); 144 #ifndef NDEBUG 145 // Ensure we don't process the same register twice 146 assert(SeenRegs.insert(Reg).second); 147 #endif 148 auto RegMapIt = DbgVRegToValues.find(Reg); 149 if (RegMapIt == DbgVRegToValues.end()) 150 continue; 151 SlotIndex LastSlot; 152 bool LastUndefResult = false; 153 for (auto [Slot, DbgValue] : RegMapIt->second) { 154 // All consecutive DBG_VALUEs have the same slot because the slot 155 // indices they have is the one for the first non-debug instruction 156 // after it, because DBG_VALUEs don't have slot index themselves. Before 157 // doing live range queries, quickly check if the current DBG_VALUE has 158 // the same slot index as the previous one, in which case we should do 159 // the same. Note that RegMapIt->second, the vector of {SlotIndex, 160 // DBG_VALUE}, is sorted by SlotIndex, which is necessary for this 161 // check. 162 if (Slot == LastSlot) { 163 if (LastUndefResult) { 164 LLVM_DEBUG(dbgs() << "Undefed: " << *DbgValue << "\n"); 165 DbgValue->setDebugValueUndef(); 166 } 167 continue; 168 } 169 LastSlot = Slot; 170 LastUndefResult = false; 171 for (LiveInterval *OtherLI : CoalescedIntervals) { 172 if (LI == OtherLI) 173 continue; 174 175 // This DBG_VALUE has 'Reg' (the current LiveInterval's register) as 176 // its operand. If this DBG_VALUE's slot index is within other 177 // registers' live ranges, this DBG_VALUE should be undefed. For 178 // example, suppose %0 and %1 are to be coalesced into %0. 179 // ; %0's live range starts 180 // %0 = value_0 181 // DBG_VALUE %0, !"a", ... (a) 182 // DBG_VALUE %1, !"b", ... (b) 183 // use %0 184 // ; %0's live range ends 185 // ... 186 // ; %1's live range starts 187 // %1 = value_1 188 // DBG_VALUE %0, !"c", ... (c) 189 // DBG_VALUE %1, !"d", ... (d) 190 // use %1 191 // ; %1's live range ends 192 // 193 // In this code, (b) and (c) should be set to undef. After the two 194 // registers are coalesced, (b) will incorrectly say the variable 195 // "b"'s value is 'value_0', and (c) will also incorrectly say the 196 // variable "c"'s value is value_1. Note it doesn't actually matter 197 // which register they are coalesced into (%0 or %1); (b) and (c) 198 // should be set to undef as well if they are coalesced into %1. 199 // 200 // This happens DBG_VALUEs are not included when computing live 201 // ranges. 202 // 203 // Note that it is not possible for this DBG_VALUE to be 204 // simultaneously within 'Reg''s live range and one of other coalesced 205 // registers' live ranges because if their live ranges overlapped they 206 // would have not been selected as a coalescing candidate in the first 207 // place. 208 auto *SegmentIt = OtherLI->find(Slot); 209 if (SegmentIt != OtherLI->end() && SegmentIt->contains(Slot)) { 210 LLVM_DEBUG(dbgs() << "Undefed: " << *DbgValue << "\n"); 211 DbgValue->setDebugValueUndef(); 212 LastUndefResult = true; 213 break; 214 } 215 } 216 } 217 } 218 } 219 } 220 221 bool WebAssemblyRegColoring::runOnMachineFunction(MachineFunction &MF) { 222 LLVM_DEBUG({ 223 dbgs() << "********** Register Coloring **********\n" 224 << "********** Function: " << MF.getName() << '\n'; 225 }); 226 227 // If there are calls to setjmp or sigsetjmp, don't perform coloring. Virtual 228 // registers could be modified before the longjmp is executed, resulting in 229 // the wrong value being used afterwards. (See <rdar://problem/8007500>.) 230 // TODO: Does WebAssembly need to care about setjmp for register coloring? 231 if (MF.exposesReturnsTwice()) 232 return false; 233 234 MachineRegisterInfo *MRI = &MF.getRegInfo(); 235 LiveIntervals *Liveness = &getAnalysis<LiveIntervals>(); 236 const MachineBlockFrequencyInfo *MBFI = 237 &getAnalysis<MachineBlockFrequencyInfo>(); 238 WebAssemblyFunctionInfo &MFI = *MF.getInfo<WebAssemblyFunctionInfo>(); 239 240 // We don't preserve SSA form. 241 MRI->leaveSSA(); 242 243 // Gather all register intervals into a list and sort them. 244 unsigned NumVRegs = MRI->getNumVirtRegs(); 245 SmallVector<LiveInterval *, 0> SortedIntervals; 246 SortedIntervals.reserve(NumVRegs); 247 248 // Record DBG_VALUEs and their SlotIndexes. 249 auto DbgVRegToValues = buildVRegToDbgValueMap(MF, Liveness); 250 251 LLVM_DEBUG(dbgs() << "Interesting register intervals:\n"); 252 for (unsigned I = 0; I < NumVRegs; ++I) { 253 Register VReg = Register::index2VirtReg(I); 254 if (MFI.isVRegStackified(VReg)) 255 continue; 256 // Skip unused registers, which can use $drop. 257 if (MRI->use_empty(VReg)) 258 continue; 259 260 LiveInterval *LI = &Liveness->getInterval(VReg); 261 assert(LI->weight() == 0.0f); 262 LI->setWeight(computeWeight(MRI, MBFI, VReg)); 263 LLVM_DEBUG(LI->dump()); 264 SortedIntervals.push_back(LI); 265 } 266 LLVM_DEBUG(dbgs() << '\n'); 267 268 // Sort them to put arguments first (since we don't want to rename live-in 269 // registers), by weight next, and then by position. 270 // TODO: Investigate more intelligent sorting heuristics. For starters, we 271 // should try to coalesce adjacent live intervals before non-adjacent ones. 272 llvm::sort(SortedIntervals, [MRI](LiveInterval *LHS, LiveInterval *RHS) { 273 if (MRI->isLiveIn(LHS->reg()) != MRI->isLiveIn(RHS->reg())) 274 return MRI->isLiveIn(LHS->reg()); 275 if (LHS->weight() != RHS->weight()) 276 return LHS->weight() > RHS->weight(); 277 if (LHS->empty() || RHS->empty()) 278 return !LHS->empty() && RHS->empty(); 279 return *LHS < *RHS; 280 }); 281 282 LLVM_DEBUG(dbgs() << "Coloring register intervals:\n"); 283 SmallVector<unsigned, 16> SlotMapping(SortedIntervals.size(), -1u); 284 SmallVector<SmallVector<LiveInterval *, 4>, 16> Assignments( 285 SortedIntervals.size()); 286 BitVector UsedColors(SortedIntervals.size()); 287 bool Changed = false; 288 for (size_t I = 0, E = SortedIntervals.size(); I < E; ++I) { 289 LiveInterval *LI = SortedIntervals[I]; 290 Register Old = LI->reg(); 291 size_t Color = I; 292 const TargetRegisterClass *RC = MRI->getRegClass(Old); 293 294 // Check if it's possible to reuse any of the used colors. 295 if (!MRI->isLiveIn(Old)) 296 for (unsigned C : UsedColors.set_bits()) { 297 if (MRI->getRegClass(SortedIntervals[C]->reg()) != RC) 298 continue; 299 for (LiveInterval *OtherLI : Assignments[C]) 300 if (!OtherLI->empty() && OtherLI->overlaps(*LI)) 301 goto continue_outer; 302 Color = C; 303 break; 304 continue_outer:; 305 } 306 307 Register New = SortedIntervals[Color]->reg(); 308 SlotMapping[I] = New; 309 Changed |= Old != New; 310 UsedColors.set(Color); 311 Assignments[Color].push_back(LI); 312 // If we reassigned the stack pointer, update the debug frame base info. 313 if (Old != New && MFI.isFrameBaseVirtual() && MFI.getFrameBaseVreg() == Old) 314 MFI.setFrameBaseVreg(New); 315 LLVM_DEBUG(dbgs() << "Assigning vreg" << Register::virtReg2Index(LI->reg()) 316 << " to vreg" << Register::virtReg2Index(New) << "\n"); 317 } 318 if (!Changed) 319 return false; 320 321 // Set DBG_VALUEs that will be invalid after coalescing to undef. 322 undefInvalidDbgValues(Liveness, Assignments, DbgVRegToValues); 323 324 // Rewrite register operands. 325 for (size_t I = 0, E = SortedIntervals.size(); I < E; ++I) { 326 Register Old = SortedIntervals[I]->reg(); 327 unsigned New = SlotMapping[I]; 328 if (Old != New) 329 MRI->replaceRegWith(Old, New); 330 } 331 return true; 332 } 333