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 "MCTargetDesc/WebAssemblyMCTargetDesc.h" // for WebAssembly::ARGUMENT_*
23 #include "llvm/CodeGen/LiveIntervalAnalysis.h"
24 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
25 #include "llvm/CodeGen/MachineRegisterInfo.h"
26 #include "llvm/CodeGen/Passes.h"
27 #include "llvm/Support/Debug.h"
28 #include "llvm/Support/raw_ostream.h"
29 using namespace llvm;
30 
31 #define DEBUG_TYPE "wasm-reg-coloring"
32 
33 namespace {
34 class WebAssemblyRegColoring final : public MachineFunctionPass {
35 public:
36   static char ID; // Pass identification, replacement for typeid
37   WebAssemblyRegColoring() : MachineFunctionPass(ID) {}
38 
39   const char *getPassName() const override {
40     return "WebAssembly Register Coloring";
41   }
42 
43   void getAnalysisUsage(AnalysisUsage &AU) const override {
44     AU.setPreservesCFG();
45     AU.addRequired<LiveIntervals>();
46     AU.addRequired<MachineBlockFrequencyInfo>();
47     AU.addPreserved<MachineBlockFrequencyInfo>();
48     AU.addPreservedID(MachineDominatorsID);
49     AU.addRequired<SlotIndexes>(); // for ARGUMENT fixups
50     MachineFunctionPass::getAnalysisUsage(AU);
51   }
52 
53   bool runOnMachineFunction(MachineFunction &MF) override;
54 
55 private:
56 };
57 } // end anonymous namespace
58 
59 char WebAssemblyRegColoring::ID = 0;
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   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   // FIXME: If scheduling has moved an ARGUMENT virtual register, move it back,
100   // and recompute liveness. This is a temporary hack.
101   bool MovedArg = false;
102   MachineBasicBlock &EntryMBB = MF.front();
103   MachineBasicBlock::iterator InsertPt = EntryMBB.end();
104   // Look for the first NonArg instruction.
105   for (auto MII = EntryMBB.begin(), MIE = EntryMBB.end(); MII != MIE; ++MII) {
106     MachineInstr *MI = MII;
107     if (MI->getOpcode() != WebAssembly::ARGUMENT_I32 &&
108         MI->getOpcode() != WebAssembly::ARGUMENT_I64 &&
109         MI->getOpcode() != WebAssembly::ARGUMENT_F32 &&
110         MI->getOpcode() != WebAssembly::ARGUMENT_F64) {
111       InsertPt = MII;
112       break;
113     }
114   }
115   // Now move any argument instructions later in the block
116   // to before our first NonArg instruction.
117   for (auto I = InsertPt, E = EntryMBB.end(); I != E; ++I) {
118     MachineInstr *MI = I;
119     if (MI->getOpcode() == WebAssembly::ARGUMENT_I32 ||
120         MI->getOpcode() == WebAssembly::ARGUMENT_I64 ||
121         MI->getOpcode() == WebAssembly::ARGUMENT_F32 ||
122         MI->getOpcode() == WebAssembly::ARGUMENT_F64) {
123       EntryMBB.insert(InsertPt, MI->removeFromParent());
124       MovedArg = true;
125     }
126   }
127   if (MovedArg) {
128     SlotIndexes &Slots = getAnalysis<SlotIndexes>();
129     Liveness->releaseMemory();
130     Slots.releaseMemory();
131     Slots.runOnMachineFunction(MF);
132     Liveness->runOnMachineFunction(MF);
133   }
134 
135   DEBUG(dbgs() << "Interesting register intervals:\n");
136   for (unsigned i = 0; i < NumVRegs; ++i) {
137     unsigned VReg = TargetRegisterInfo::index2VirtReg(i);
138     if (MFI.isVRegStackified(VReg))
139       continue;
140 
141     LiveInterval *LI = &Liveness->getInterval(VReg);
142     assert(LI->weight == 0.0f);
143     LI->weight = computeWeight(MRI, MBFI, VReg);
144     DEBUG(LI->dump());
145     SortedIntervals.push_back(LI);
146   }
147   DEBUG(dbgs() << '\n');
148 
149   // Sort them to put arguments first (since we don't want to rename live-in
150   // registers), by weight next, and then by position.
151   // TODO: Investigate more intelligent sorting heuristics. For starters, we
152   // should try to coalesce adjacent live intervals before non-adjacent ones.
153   std::sort(SortedIntervals.begin(), SortedIntervals.end(),
154             [MRI](LiveInterval *LHS, LiveInterval *RHS) {
155               if (MRI->isLiveIn(LHS->reg) != MRI->isLiveIn(RHS->reg))
156                 return MRI->isLiveIn(LHS->reg);
157               if (LHS->weight != RHS->weight)
158                 return LHS->weight > RHS->weight;
159               if (LHS->empty() || RHS->empty())
160                 return !LHS->empty() && RHS->empty();
161               return *LHS < *RHS;
162             });
163 
164   DEBUG(dbgs() << "Coloring register intervals:\n");
165   SmallVector<unsigned, 16> SlotMapping(SortedIntervals.size(), -1u);
166   SmallVector<SmallVector<LiveInterval *, 4>, 16> Assignments(
167       SortedIntervals.size());
168   BitVector UsedColors(SortedIntervals.size());
169   bool Changed = false;
170   for (size_t i = 0, e = SortedIntervals.size(); i < e; ++i) {
171     LiveInterval *LI = SortedIntervals[i];
172     unsigned Old = LI->reg;
173     size_t Color = i;
174     const TargetRegisterClass *RC = MRI->getRegClass(Old);
175 
176     // Check if it's possible to reuse any of the used colors.
177     if (!MRI->isLiveIn(Old))
178       for (int C(UsedColors.find_first()); C != -1;
179            C = UsedColors.find_next(C)) {
180         if (MRI->getRegClass(SortedIntervals[C]->reg) != RC)
181           continue;
182         for (LiveInterval *OtherLI : Assignments[C])
183           if (!OtherLI->empty() && OtherLI->overlaps(*LI))
184             goto continue_outer;
185         Color = C;
186         break;
187       continue_outer:;
188       }
189 
190     unsigned New = SortedIntervals[Color]->reg;
191     SlotMapping[i] = New;
192     Changed |= Old != New;
193     UsedColors.set(Color);
194     Assignments[Color].push_back(LI);
195     DEBUG(dbgs() << "Assigning vreg"
196                  << TargetRegisterInfo::virtReg2Index(LI->reg) << " to vreg"
197                  << TargetRegisterInfo::virtReg2Index(New) << "\n");
198   }
199   if (!Changed)
200     return false;
201 
202   // Rewrite register operands.
203   for (size_t i = 0, e = SortedIntervals.size(); i < e; ++i) {
204     unsigned Old = SortedIntervals[i]->reg;
205     unsigned New = SlotMapping[i];
206     if (Old != New)
207       MRI->replaceRegWith(Old, New);
208   }
209   return true;
210 }
211