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 bool WebAssemblyRegColoring::runOnMachineFunction(MachineFunction &MF) {
76   LLVM_DEBUG({
77     dbgs() << "********** Register Coloring **********\n"
78            << "********** Function: " << MF.getName() << '\n';
79   });
80 
81   // If there are calls to setjmp or sigsetjmp, don't perform coloring. Virtual
82   // registers could be modified before the longjmp is executed, resulting in
83   // the wrong value being used afterwards. (See <rdar://problem/8007500>.)
84   // TODO: Does WebAssembly need to care about setjmp for register coloring?
85   if (MF.exposesReturnsTwice())
86     return false;
87 
88   MachineRegisterInfo *MRI = &MF.getRegInfo();
89   LiveIntervals *Liveness = &getAnalysis<LiveIntervals>();
90   const MachineBlockFrequencyInfo *MBFI =
91       &getAnalysis<MachineBlockFrequencyInfo>();
92   WebAssemblyFunctionInfo &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();
93 
94   // Gather all register intervals into a list and sort them.
95   unsigned NumVRegs = MRI->getNumVirtRegs();
96   SmallVector<LiveInterval *, 0> SortedIntervals;
97   SortedIntervals.reserve(NumVRegs);
98 
99   LLVM_DEBUG(dbgs() << "Interesting register intervals:\n");
100   for (unsigned I = 0; I < NumVRegs; ++I) {
101     unsigned VReg = Register::index2VirtReg(I);
102     if (MFI.isVRegStackified(VReg))
103       continue;
104     // Skip unused registers, which can use $drop.
105     if (MRI->use_empty(VReg))
106       continue;
107 
108     LiveInterval *LI = &Liveness->getInterval(VReg);
109     assert(LI->weight() == 0.0f);
110     LI->setWeight(computeWeight(MRI, MBFI, VReg));
111     LLVM_DEBUG(LI->dump());
112     SortedIntervals.push_back(LI);
113   }
114   LLVM_DEBUG(dbgs() << '\n');
115 
116   // Sort them to put arguments first (since we don't want to rename live-in
117   // registers), by weight next, and then by position.
118   // TODO: Investigate more intelligent sorting heuristics. For starters, we
119   // should try to coalesce adjacent live intervals before non-adjacent ones.
120   llvm::sort(SortedIntervals, [MRI](LiveInterval *LHS, LiveInterval *RHS) {
121     if (MRI->isLiveIn(LHS->reg()) != MRI->isLiveIn(RHS->reg()))
122       return MRI->isLiveIn(LHS->reg());
123     if (LHS->weight() != RHS->weight())
124       return LHS->weight() > RHS->weight();
125     if (LHS->empty() || RHS->empty())
126       return !LHS->empty() && RHS->empty();
127     return *LHS < *RHS;
128   });
129 
130   LLVM_DEBUG(dbgs() << "Coloring register intervals:\n");
131   SmallVector<unsigned, 16> SlotMapping(SortedIntervals.size(), -1u);
132   SmallVector<SmallVector<LiveInterval *, 4>, 16> Assignments(
133       SortedIntervals.size());
134   BitVector UsedColors(SortedIntervals.size());
135   bool Changed = false;
136   for (size_t I = 0, E = SortedIntervals.size(); I < E; ++I) {
137     LiveInterval *LI = SortedIntervals[I];
138     unsigned Old = LI->reg();
139     size_t Color = I;
140     const TargetRegisterClass *RC = MRI->getRegClass(Old);
141 
142     // Check if it's possible to reuse any of the used colors.
143     if (!MRI->isLiveIn(Old))
144       for (unsigned C : UsedColors.set_bits()) {
145         if (MRI->getRegClass(SortedIntervals[C]->reg()) != RC)
146           continue;
147         for (LiveInterval *OtherLI : Assignments[C])
148           if (!OtherLI->empty() && OtherLI->overlaps(*LI))
149             goto continue_outer;
150         Color = C;
151         break;
152       continue_outer:;
153       }
154 
155     unsigned New = SortedIntervals[Color]->reg();
156     SlotMapping[I] = New;
157     Changed |= Old != New;
158     UsedColors.set(Color);
159     Assignments[Color].push_back(LI);
160     // If we reassigned the stack pointer, update the debug frame base info.
161     if (Old != New && MFI.isFrameBaseVirtual() && MFI.getFrameBaseVreg() == Old)
162       MFI.setFrameBaseVreg(New);
163     LLVM_DEBUG(dbgs() << "Assigning vreg" << Register::virtReg2Index(LI->reg())
164                       << " to vreg" << Register::virtReg2Index(New) << "\n");
165   }
166   if (!Changed)
167     return false;
168 
169   // Rewrite register operands.
170   for (size_t I = 0, E = SortedIntervals.size(); I < E; ++I) {
171     unsigned Old = SortedIntervals[I]->reg();
172     unsigned New = SlotMapping[I];
173     if (Old != New)
174       MRI->replaceRegWith(Old, New);
175   }
176   return true;
177 }
178