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<LiveIntervals>();
44 AU.addRequired<MachineBlockFrequencyInfo>();
45 AU.addPreserved<MachineBlockFrequencyInfo>();
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 (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(Val: 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(Pos: Slot);
209 if (SegmentIt != OtherLI->end() && SegmentIt->contains(I: Slot)) {
210 LLVM_DEBUG(dbgs() << "Undefed: " << *DbgValue << "\n");
211 DbgValue->setDebugValueUndef();
212 LastUndefResult = true;
213 break;
214 }
215 }
216 }
217 }
218 }
219}
220
221bool 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.
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(N: 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(Index: I);
254 if (MFI.isVRegStackified(VReg))
255 continue;
256 // Skip unused registers, which can use $drop.
257 if (MRI->use_empty(RegNo: VReg))
258 continue;
259
260 LiveInterval *LI = &Liveness->getInterval(Reg: VReg);
261 assert(LI->weight() == 0.0f);
262 LI->setWeight(computeWeight(MRI, MBFI, VReg));
263 LLVM_DEBUG(LI->dump());
264 SortedIntervals.push_back(Elt: 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(C&: SortedIntervals, Comp: [MRI](LiveInterval *LHS, LiveInterval *RHS) {
273 if (MRI->isLiveIn(Reg: LHS->reg()) != MRI->isLiveIn(Reg: RHS->reg()))
274 return MRI->isLiveIn(Reg: 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(Reg: Old);
293
294 // Check if it's possible to reuse any of the used colors.
295 if (!MRI->isLiveIn(Reg: Old))
296 for (unsigned C : UsedColors.set_bits()) {
297 if (MRI->getRegClass(Reg: SortedIntervals[C]->reg()) != RC)
298 continue;
299 for (LiveInterval *OtherLI : Assignments[C])
300 if (!OtherLI->empty() && OtherLI->overlaps(other: *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(Elt: 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(FromReg: Old, ToReg: New);
330 }
331 return true;
332}
333

source code of llvm/lib/Target/WebAssembly/WebAssemblyRegColoring.cpp