332 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8///9/// \file10/// This file implements a virtual register coloring pass.11///12/// WebAssembly doesn't have a fixed number of registers, but it is still13/// 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"27using namespace llvm;28 29#define DEBUG_TYPE "wasm-reg-coloring"30 31namespace {32class WebAssemblyRegColoring final : public MachineFunctionPass {33public:34 static char ID; // Pass identification, replacement for typeid35 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<LiveIntervalsWrapperPass>();44 AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();45 AU.addPreserved<MachineBlockFrequencyInfoWrapperPass>();46 AU.addPreservedID(MachineDominatorsID);47 MachineFunctionPass::getAnalysisUsage(AU);48 }49 50 bool runOnMachineFunction(MachineFunction &MF) override;51 52private:53};54} // end anonymous namespace55 56char WebAssemblyRegColoring::ID = 0;57INITIALIZE_PASS(WebAssemblyRegColoring, DEBUG_TYPE,58 "Minimize number of registers used", false, false)59 60FunctionPass *llvm::createWebAssemblyRegColoring() {61 return new WebAssemblyRegColoring();62}63 64// Compute the total spill weight for VReg.65static 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 end77// of a BB, because DBG_VALUE's don't have slot index themselves.78// Adapted from RegisterCoalescer::buildVRegToDbgValueMap.79static DenseMap<Register, std::vector<std::pair<SlotIndex, MachineInstr *>>>80buildVRegToDbgValueMap(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 the87 // 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 the101 // 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 register129// changes.130static void undefInvalidDbgValues(131 const LiveIntervals *Liveness,132 ArrayRef<SmallVector<LiveInterval *, 4>> Assignments,133 DenseMap<Register, std::vector<std::pair<SlotIndex, MachineInstr *>>>134 &DbgVRegToValues) {135#ifndef NDEBUG136 DenseSet<Register> SeenRegs;137#endif138 for (const auto &CoalescedIntervals : Assignments) {139 if (CoalescedIntervals.empty())140 continue;141 for (LiveInterval *LI : CoalescedIntervals) {142 Register Reg = LI->reg();143#ifndef NDEBUG144 // Ensure we don't process the same register twice145 assert(SeenRegs.insert(Reg).second);146#endif147 auto RegMapIt = DbgVRegToValues.find(Reg);148 if (RegMapIt == DbgVRegToValues.end())149 continue;150 SlotIndex LastSlot;151 bool LastUndefResult = false;152 for (auto [Slot, DbgValue] : RegMapIt->second) {153 // All consecutive DBG_VALUEs have the same slot because the slot154 // indices they have is the one for the first non-debug instruction155 // after it, because DBG_VALUEs don't have slot index themselves. Before156 // doing live range queries, quickly check if the current DBG_VALUE has157 // the same slot index as the previous one, in which case we should do158 // the same. Note that RegMapIt->second, the vector of {SlotIndex,159 // DBG_VALUE}, is sorted by SlotIndex, which is necessary for this160 // check.161 if (Slot == LastSlot) {162 if (LastUndefResult) {163 LLVM_DEBUG(dbgs() << "Undefed: " << *DbgValue << "\n");164 DbgValue->setDebugValueUndef();165 }166 continue;167 }168 LastSlot = Slot;169 LastUndefResult = false;170 for (LiveInterval *OtherLI : CoalescedIntervals) {171 if (LI == OtherLI)172 continue;173 174 // This DBG_VALUE has 'Reg' (the current LiveInterval's register) as175 // its operand. If this DBG_VALUE's slot index is within other176 // registers' live ranges, this DBG_VALUE should be undefed. For177 // example, suppose %0 and %1 are to be coalesced into %0.178 // ; %0's live range starts179 // %0 = value_0180 // DBG_VALUE %0, !"a", ... (a)181 // DBG_VALUE %1, !"b", ... (b)182 // use %0183 // ; %0's live range ends184 // ...185 // ; %1's live range starts186 // %1 = value_1187 // DBG_VALUE %0, !"c", ... (c)188 // DBG_VALUE %1, !"d", ... (d)189 // use %1190 // ; %1's live range ends191 //192 // In this code, (b) and (c) should be set to undef. After the two193 // registers are coalesced, (b) will incorrectly say the variable194 // "b"'s value is 'value_0', and (c) will also incorrectly say the195 // variable "c"'s value is value_1. Note it doesn't actually matter196 // which register they are coalesced into (%0 or %1); (b) and (c)197 // should be set to undef as well if they are coalesced into %1.198 //199 // This happens DBG_VALUEs are not included when computing live200 // ranges.201 //202 // Note that it is not possible for this DBG_VALUE to be203 // simultaneously within 'Reg''s live range and one of other coalesced204 // registers' live ranges because if their live ranges overlapped they205 // would have not been selected as a coalescing candidate in the first206 // place.207 auto *SegmentIt = OtherLI->find(Slot);208 if (SegmentIt != OtherLI->end() && SegmentIt->contains(Slot)) {209 LLVM_DEBUG(dbgs() << "Undefed: " << *DbgValue << "\n");210 DbgValue->setDebugValueUndef();211 LastUndefResult = true;212 break;213 }214 }215 }216 }217 }218}219 220bool WebAssemblyRegColoring::runOnMachineFunction(MachineFunction &MF) {221 LLVM_DEBUG({222 dbgs() << "********** Register Coloring **********\n"223 << "********** Function: " << MF.getName() << '\n';224 });225 226 // If there are calls to setjmp or sigsetjmp, don't perform coloring. Virtual227 // registers could be modified before the longjmp is executed, resulting in228 // the wrong value being used afterwards.229 // TODO: Does WebAssembly need to care about setjmp for register coloring?230 if (MF.exposesReturnsTwice())231 return false;232 233 MachineRegisterInfo *MRI = &MF.getRegInfo();234 LiveIntervals *Liveness = &getAnalysis<LiveIntervalsWrapperPass>().getLIS();235 const MachineBlockFrequencyInfo *MBFI =236 &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();237 WebAssemblyFunctionInfo &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();238 239 // We don't preserve SSA form.240 MRI->leaveSSA();241 242 // Gather all register intervals into a list and sort them.243 unsigned NumVRegs = MRI->getNumVirtRegs();244 SmallVector<LiveInterval *, 0> SortedIntervals;245 SortedIntervals.reserve(NumVRegs);246 247 // Record DBG_VALUEs and their SlotIndexes.248 auto DbgVRegToValues = buildVRegToDbgValueMap(MF, Liveness);249 250 LLVM_DEBUG(dbgs() << "Interesting register intervals:\n");251 for (unsigned I = 0; I < NumVRegs; ++I) {252 Register VReg = Register::index2VirtReg(I);253 if (MFI.isVRegStackified(VReg))254 continue;255 // Skip unused registers, which can use $drop.256 if (MRI->use_empty(VReg))257 continue;258 259 LiveInterval *LI = &Liveness->getInterval(VReg);260 assert(LI->weight() == 0.0f);261 LI->setWeight(computeWeight(MRI, MBFI, VReg));262 LLVM_DEBUG(LI->dump());263 SortedIntervals.push_back(LI);264 }265 LLVM_DEBUG(dbgs() << '\n');266 267 // Sort them to put arguments first (since we don't want to rename live-in268 // registers), by weight next, and then by position.269 // TODO: Investigate more intelligent sorting heuristics. For starters, we270 // should try to coalesce adjacent live intervals before non-adjacent ones.271 llvm::sort(SortedIntervals, [MRI](LiveInterval *LHS, LiveInterval *RHS) {272 if (MRI->isLiveIn(LHS->reg()) != MRI->isLiveIn(RHS->reg()))273 return MRI->isLiveIn(LHS->reg());274 if (LHS->weight() != RHS->weight())275 return LHS->weight() > RHS->weight();276 if (LHS->empty() || RHS->empty())277 return !LHS->empty() && RHS->empty();278 return *LHS < *RHS;279 });280 281 LLVM_DEBUG(dbgs() << "Coloring register intervals:\n");282 SmallVector<unsigned, 16> SlotMapping(SortedIntervals.size(), -1u);283 SmallVector<SmallVector<LiveInterval *, 4>, 16> Assignments(284 SortedIntervals.size());285 BitVector UsedColors(SortedIntervals.size());286 bool Changed = false;287 for (size_t I = 0, E = SortedIntervals.size(); I < E; ++I) {288 LiveInterval *LI = SortedIntervals[I];289 Register Old = LI->reg();290 size_t Color = I;291 const TargetRegisterClass *RC = MRI->getRegClass(Old);292 293 // Check if it's possible to reuse any of the used colors.294 if (!MRI->isLiveIn(Old))295 for (unsigned C : UsedColors.set_bits()) {296 if (MRI->getRegClass(SortedIntervals[C]->reg()) != RC)297 continue;298 for (LiveInterval *OtherLI : Assignments[C])299 if (!OtherLI->empty() && OtherLI->overlaps(*LI))300 goto continue_outer;301 Color = C;302 break;303 continue_outer:;304 }305 306 Register New = SortedIntervals[Color]->reg();307 SlotMapping[I] = New;308 Changed |= Old != New;309 UsedColors.set(Color);310 Assignments[Color].push_back(LI);311 // If we reassigned the stack pointer, update the debug frame base info.312 if (Old != New && MFI.isFrameBaseVirtual() && MFI.getFrameBaseVreg() == Old)313 MFI.setFrameBaseVreg(New);314 LLVM_DEBUG(dbgs() << "Assigning vreg " << printReg(LI->reg()) << " to vreg "315 << printReg(New) << "\n");316 }317 if (!Changed)318 return false;319 320 // Set DBG_VALUEs that will be invalid after coalescing to undef.321 undefInvalidDbgValues(Liveness, Assignments, DbgVRegToValues);322 323 // Rewrite register operands.324 for (size_t I = 0, E = SortedIntervals.size(); I < E; ++I) {325 Register Old = SortedIntervals[I]->reg();326 unsigned New = SlotMapping[I];327 if (Old != New)328 MRI->replaceRegWith(Old, New);329 }330 return true;331}332