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"
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 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<LiveIntervalsWrapperPass>();
44 AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
45 AU.addPreserved<MachineBlockFrequencyInfoWrapperPass>();
46 AU.addPreservedID(ID&: MachineDominatorsID);
47 MachineFunctionPass::getAnalysisUsage(AU);
48 }
49
50 bool runOnMachineFunction(MachineFunction &MF) override;
51
52private:
53};
54} // end anonymous namespace
55
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(Reg: VReg))
70 Weight += LiveIntervals::getSpillWeight(isDef: MO.isDef(), isUse: MO.isUse(), MBFI,
71 MI: *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.
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 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(x: {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: &MBB);
104
105 for (auto &MI : MBB) {
106 if (MI.isDebugValue()) {
107 if (any_of(Range: MI.debug_operands(), P: [](const MachineOperand &MO) {
108 return MO.isReg() && MO.getReg().isVirtual();
109 }))
110 ToInsert.push_back(Elt: &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: &MBB));
119 }
120
121 // Sort all DBG_VALUEs we've seen by slot number.
122 for (auto &Pair : DbgVRegToValues)
123 llvm::sort(C&: 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.
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 NDEBUG
136 DenseSet<Register> SeenRegs;
137#endif
138 for (const auto &CoalescedIntervals : Assignments) {
139 if (CoalescedIntervals.empty())
140 continue;
141 for (LiveInterval *LI : CoalescedIntervals) {
142 Register Reg = LI->reg();
143#ifndef NDEBUG
144 // Ensure we don't process the same register twice
145 assert(SeenRegs.insert(Reg).second);
146#endif
147 auto RegMapIt = DbgVRegToValues.find(Val: 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 slot
154 // indices they have is the one for the first non-debug instruction
155 // after it, because DBG_VALUEs don't have slot index themselves. Before
156 // doing live range queries, quickly check if the current DBG_VALUE has
157 // the same slot index as the previous one, in which case we should do
158 // the same. Note that RegMapIt->second, the vector of {SlotIndex,
159 // DBG_VALUE}, is sorted by SlotIndex, which is necessary for this
160 // 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) as
175 // its operand. If this DBG_VALUE's slot index is within other
176 // registers' live ranges, this DBG_VALUE should be undefed. For
177 // example, suppose %0 and %1 are to be coalesced into %0.
178 // ; %0's live range starts
179 // %0 = value_0
180 // DBG_VALUE %0, !"a", ... (a)
181 // DBG_VALUE %1, !"b", ... (b)
182 // use %0
183 // ; %0's live range ends
184 // ...
185 // ; %1's live range starts
186 // %1 = value_1
187 // DBG_VALUE %0, !"c", ... (c)
188 // DBG_VALUE %1, !"d", ... (d)
189 // use %1
190 // ; %1's live range ends
191 //
192 // In this code, (b) and (c) should be set to undef. After the two
193 // registers are coalesced, (b) will incorrectly say the variable
194 // "b"'s value is 'value_0', and (c) will also incorrectly say the
195 // variable "c"'s value is value_1. Note it doesn't actually matter
196 // 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 live
200 // ranges.
201 //
202 // Note that it is not possible for this DBG_VALUE to be
203 // simultaneously within 'Reg''s live range and one of other coalesced
204 // registers' live ranges because if their live ranges overlapped they
205 // would have not been selected as a coalescing candidate in the first
206 // place.
207 auto *SegmentIt = OtherLI->find(Pos: Slot);
208 if (SegmentIt != OtherLI->end() && SegmentIt->contains(I: 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. Virtual
227 // registers could be modified before the longjmp is executed, resulting in
228 // 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(N: 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(Index: I);
253 if (MFI.isVRegStackified(VReg))
254 continue;
255 // Skip unused registers, which can use $drop.
256 if (MRI->use_empty(RegNo: VReg))
257 continue;
258
259 LiveInterval *LI = &Liveness->getInterval(Reg: VReg);
260 assert(LI->weight() == 0.0f);
261 LI->setWeight(computeWeight(MRI, MBFI, VReg));
262 LLVM_DEBUG(LI->dump());
263 SortedIntervals.push_back(Elt: LI);
264 }
265 LLVM_DEBUG(dbgs() << '\n');
266
267 // Sort them to put arguments first (since we don't want to rename live-in
268 // registers), by weight next, and then by position.
269 // TODO: Investigate more intelligent sorting heuristics. For starters, we
270 // should try to coalesce adjacent live intervals before non-adjacent ones.
271 llvm::sort(C&: SortedIntervals, Comp: [MRI](LiveInterval *LHS, LiveInterval *RHS) {
272 if (MRI->isLiveIn(Reg: LHS->reg()) != MRI->isLiveIn(Reg: RHS->reg()))
273 return MRI->isLiveIn(Reg: 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(Reg: Old);
292
293 // Check if it's possible to reuse any of the used colors.
294 if (!MRI->isLiveIn(Reg: Old))
295 for (unsigned C : UsedColors.set_bits()) {
296 if (MRI->getRegClass(Reg: SortedIntervals[C]->reg()) != RC)
297 continue;
298 for (LiveInterval *OtherLI : Assignments[C])
299 if (!OtherLI->empty() && OtherLI->overlaps(other: *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(Elt: 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(FromReg: Old, ToReg: New);
329 }
330 return true;
331}
332