WebAssemblyRegColoring.cpp revision 360784
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(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(VReg))
70    Weight += LiveIntervals::getSpillWeight(MO.isDef(), MO.isUse(), MBFI,
71                                            *MO.getParent());
72  return Weight;
73}
74
75bool 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->weight = 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    LLVM_DEBUG(dbgs() << "Assigning vreg" << Register::virtReg2Index(LI->reg)
161                      << " to vreg" << Register::virtReg2Index(New) << "\n");
162  }
163  if (!Changed)
164    return false;
165
166  // Rewrite register operands.
167  for (size_t I = 0, E = SortedIntervals.size(); I < E; ++I) {
168    unsigned Old = SortedIntervals[I]->reg;
169    unsigned New = SlotMapping[I];
170    if (Old != New)
171      MRI->replaceRegWith(Old, New);
172  }
173  return true;
174}
175