1 //===-- WebAssemblyRegColoring.cpp - Register coloring --------------------===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 ///
10 /// \file
11 /// \brief This file implements a virtual register coloring pass.
12 ///
13 /// WebAssembly doesn't have a fixed number of registers, but it is still
14 /// desirable to minimize the total number of registers used in each function.
15 ///
16 /// This code is modeled after lib/CodeGen/StackSlotColoring.cpp.
17 ///
18 //===----------------------------------------------------------------------===//
19 
20 #include "WebAssembly.h"
21 #include "WebAssemblyMachineFunctionInfo.h"
22 #include "llvm/CodeGen/LiveIntervalAnalysis.h"
23 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
24 #include "llvm/CodeGen/MachineRegisterInfo.h"
25 #include "llvm/CodeGen/Passes.h"
26 #include "llvm/Support/Debug.h"
27 #include "llvm/Support/raw_ostream.h"
28 using namespace llvm;
29 
30 #define DEBUG_TYPE "wasm-reg-coloring"
31 
32 static cl::opt<bool>
33     DisableRegColoring("disable-wasm-reg-coloring", cl::Hidden, cl::init(false),
34                        cl::desc("Disable WebAssembly register coloring"));
35 
36 namespace {
37 class WebAssemblyRegColoring final : public MachineFunctionPass {
38 public:
39   static char ID; // Pass identification, replacement for typeid
40   WebAssemblyRegColoring() : MachineFunctionPass(ID) {}
41 
42   const char *getPassName() const override {
43     return "WebAssembly Register Coloring";
44   }
45 
46   void getAnalysisUsage(AnalysisUsage &AU) const override {
47     AU.setPreservesCFG();
48     AU.addRequired<LiveIntervals>();
49     AU.addRequired<MachineBlockFrequencyInfo>();
50     AU.addPreserved<MachineBlockFrequencyInfo>();
51     AU.addPreservedID(MachineDominatorsID);
52     MachineFunctionPass::getAnalysisUsage(AU);
53   }
54 
55   bool runOnMachineFunction(MachineFunction &MF) override;
56 
57 private:
58 };
59 } // end anonymous namespace
60 
61 char WebAssemblyRegColoring::ID = 0;
62 FunctionPass *llvm::createWebAssemblyRegColoring() {
63   return new WebAssemblyRegColoring();
64 }
65 
66 // Compute the total spill weight for VReg.
67 static float computeWeight(const MachineRegisterInfo *MRI,
68                            const MachineBlockFrequencyInfo *MBFI,
69                            unsigned VReg) {
70   float weight = 0.0f;
71   for (MachineOperand &MO : MRI->reg_nodbg_operands(VReg))
72     weight += LiveIntervals::getSpillWeight(MO.isDef(), MO.isUse(), MBFI,
73                                             MO.getParent());
74   return weight;
75 }
76 
77 bool WebAssemblyRegColoring::runOnMachineFunction(MachineFunction &MF) {
78   DEBUG({
79     dbgs() << "********** Register Coloring **********\n"
80            << "********** Function: " << MF.getName() << '\n';
81   });
82 
83   if (DisableRegColoring)
84     return false;
85 
86   // If there are calls to setjmp or sigsetjmp, don't perform coloring. Virtual
87   // registers could be modified before the longjmp is executed, resulting in
88   // the wrong value being used afterwards. (See <rdar://problem/8007500>.)
89   // TODO: Does WebAssembly need to care about setjmp for register coloring?
90   if (MF.exposesReturnsTwice())
91     return false;
92 
93   MachineRegisterInfo *MRI = &MF.getRegInfo();
94   LiveIntervals *Liveness = &getAnalysis<LiveIntervals>();
95   const MachineBlockFrequencyInfo *MBFI =
96       &getAnalysis<MachineBlockFrequencyInfo>();
97   WebAssemblyFunctionInfo &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();
98 
99   // Gather all register intervals into a list and sort them.
100   unsigned NumVRegs = MRI->getNumVirtRegs();
101   SmallVector<LiveInterval *, 0> SortedIntervals;
102   SortedIntervals.reserve(NumVRegs);
103 
104   DEBUG(dbgs() << "Interesting register intervals:\n");
105   for (unsigned i = 0; i < NumVRegs; ++i) {
106     unsigned VReg = TargetRegisterInfo::index2VirtReg(i);
107     if (MFI.isVRegStackified(VReg))
108       continue;
109     // Skip unused registers, which can use $discard.
110     if (MRI->use_empty(VReg))
111       continue;
112 
113     LiveInterval *LI = &Liveness->getInterval(VReg);
114     assert(LI->weight == 0.0f);
115     LI->weight = computeWeight(MRI, MBFI, VReg);
116     DEBUG(LI->dump());
117     SortedIntervals.push_back(LI);
118   }
119   DEBUG(dbgs() << '\n');
120 
121   // Sort them to put arguments first (since we don't want to rename live-in
122   // registers), by weight next, and then by position.
123   // TODO: Investigate more intelligent sorting heuristics. For starters, we
124   // should try to coalesce adjacent live intervals before non-adjacent ones.
125   std::sort(SortedIntervals.begin(), SortedIntervals.end(),
126             [MRI](LiveInterval *LHS, LiveInterval *RHS) {
127               if (MRI->isLiveIn(LHS->reg) != MRI->isLiveIn(RHS->reg))
128                 return MRI->isLiveIn(LHS->reg);
129               if (LHS->weight != RHS->weight)
130                 return LHS->weight > RHS->weight;
131               if (LHS->empty() || RHS->empty())
132                 return !LHS->empty() && RHS->empty();
133               return *LHS < *RHS;
134             });
135 
136   DEBUG(dbgs() << "Coloring register intervals:\n");
137   SmallVector<unsigned, 16> SlotMapping(SortedIntervals.size(), -1u);
138   SmallVector<SmallVector<LiveInterval *, 4>, 16> Assignments(
139       SortedIntervals.size());
140   BitVector UsedColors(SortedIntervals.size());
141   bool Changed = false;
142   for (size_t i = 0, e = SortedIntervals.size(); i < e; ++i) {
143     LiveInterval *LI = SortedIntervals[i];
144     unsigned Old = LI->reg;
145     size_t Color = i;
146     const TargetRegisterClass *RC = MRI->getRegClass(Old);
147 
148     // Check if it's possible to reuse any of the used colors.
149     if (!MRI->isLiveIn(Old))
150       for (int C(UsedColors.find_first()); C != -1;
151            C = UsedColors.find_next(C)) {
152         if (MRI->getRegClass(SortedIntervals[C]->reg) != RC)
153           continue;
154         for (LiveInterval *OtherLI : Assignments[C])
155           if (!OtherLI->empty() && OtherLI->overlaps(*LI))
156             goto continue_outer;
157         Color = C;
158         break;
159       continue_outer:;
160       }
161 
162     unsigned New = SortedIntervals[Color]->reg;
163     SlotMapping[i] = New;
164     Changed |= Old != New;
165     UsedColors.set(Color);
166     Assignments[Color].push_back(LI);
167     DEBUG(dbgs() << "Assigning vreg"
168                  << TargetRegisterInfo::virtReg2Index(LI->reg) << " to vreg"
169                  << TargetRegisterInfo::virtReg2Index(New) << "\n");
170   }
171   if (!Changed)
172     return false;
173 
174   // Rewrite register operands.
175   for (size_t i = 0, e = SortedIntervals.size(); i < e; ++i) {
176     unsigned Old = SortedIntervals[i]->reg;
177     unsigned New = SlotMapping[i];
178     if (Old != New)
179       MRI->replaceRegWith(Old, New);
180   }
181   return true;
182 }
183