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(ID&: 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(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. |
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(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. |
130 | static 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 | |
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. |
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 | |