contrib/llvm/lib/Target/WebAssembly/WebAssemblyRegColoring.cpp

   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 /// 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/LiveIntervals.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 namespace {
  33 class WebAssemblyRegColoring final : public MachineFunctionPass {
  34 public:
  35   static char ID; // Pass identification, replacement for typeid
  36   WebAssemblyRegColoring() : MachineFunctionPass(ID) {}
  37
  38   StringRef getPassName() const override {
  39     return "WebAssembly Register Coloring";
  40   }
  41
  42   void getAnalysisUsage(AnalysisUsage &AU) const override {
  43     AU.setPreservesCFG();
  44     AU.addRequired<LiveIntervals>();
  45     AU.addRequired<MachineBlockFrequencyInfo>();
  46     AU.addPreserved<MachineBlockFrequencyInfo>();
  47     AU.addPreservedID(MachineDominatorsID);
  48     MachineFunctionPass::getAnalysisUsage(AU);
  49   }
  50
  51   bool runOnMachineFunction(MachineFunction &MF) override;
  52
  53 private:
  54 };
  55 } // end anonymous namespace
  56
  57 char WebAssemblyRegColoring::ID = 0;
  58 INITIALIZE_PASS(WebAssemblyRegColoring, DEBUG_TYPE,
  59                 "Minimize number of registers used", false, false)
  60
  61 FunctionPass *llvm::createWebAssemblyRegColoring() {
  62   return new WebAssemblyRegColoring();
  63 }
  64
  65 // Compute the total spill weight for VReg.
  66 static float computeWeight(const MachineRegisterInfo *MRI,
  67                            const MachineBlockFrequencyInfo *MBFI,
  68                            unsigned VReg) {
  69   float weight = 0.0f;
  70   for (MachineOperand &MO : MRI->reg_nodbg_operands(VReg))
  71     weight += LiveIntervals::getSpillWeight(MO.isDef(), MO.isUse(), MBFI,
  72                                             *MO.getParent());
  73   return weight;
  74 }
  75
  76 bool WebAssemblyRegColoring::runOnMachineFunction(MachineFunction &MF) {
  77   LLVM_DEBUG({
  78     dbgs() << "********** Register Coloring **********\n"
  79            << "********** Function: " << MF.getName() << '\n';
  80   });
  81
  82   // If there are calls to setjmp or sigsetjmp, don't perform coloring. Virtual
  83   // registers could be modified before the longjmp is executed, resulting in
  84   // the wrong value being used afterwards. (See <rdar://problem/8007500>.)
  85   // TODO: Does WebAssembly need to care about setjmp for register coloring?
  86   if (MF.exposesReturnsTwice())
  87     return false;
  88
  89   MachineRegisterInfo *MRI = &MF.getRegInfo();
  90   LiveIntervals *Liveness = &getAnalysis<LiveIntervals>();
  91   const MachineBlockFrequencyInfo *MBFI =
  92       &getAnalysis<MachineBlockFrequencyInfo>();
  93   WebAssemblyFunctionInfo &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();
  94
  95   // Gather all register intervals into a list and sort them.
  96   unsigned NumVRegs = MRI->getNumVirtRegs();
  97   SmallVector<LiveInterval *, 0> SortedIntervals;
  98   SortedIntervals.reserve(NumVRegs);
  99
 100   LLVM_DEBUG(dbgs() << "Interesting register intervals:\n");
 101   for (unsigned i = 0; i < NumVRegs; ++i) {
 102     unsigned VReg = TargetRegisterInfo::index2VirtReg(i);
 103     if (MFI.isVRegStackified(VReg))
 104       continue;
 105     // Skip unused registers, which can use $drop.
 106     if (MRI->use_empty(VReg))
 107       continue;
 108
 109     LiveInterval *LI = &Liveness->getInterval(VReg);
 110     assert(LI->weight == 0.0f);
 111     LI->weight = computeWeight(MRI, MBFI, VReg);
 112     LLVM_DEBUG(LI->dump());
 113     SortedIntervals.push_back(LI);
 114   }
 115   LLVM_DEBUG(dbgs() << '\n');
 116
 117   // Sort them to put arguments first (since we don't want to rename live-in
 118   // registers), by weight next, and then by position.
 119   // TODO: Investigate more intelligent sorting heuristics. For starters, we
 120   // should try to coalesce adjacent live intervals before non-adjacent ones.
 121   llvm::sort(SortedIntervals.begin(), SortedIntervals.end(),
 122              [MRI](LiveInterval *LHS, LiveInterval *RHS) {
 123                if (MRI->isLiveIn(LHS->reg) != MRI->isLiveIn(RHS->reg))
 124                  return MRI->isLiveIn(LHS->reg);
 125                if (LHS->weight != RHS->weight)
 126                  return LHS->weight > RHS->weight;
 127                if (LHS->empty() || RHS->empty())
 128                  return !LHS->empty() && RHS->empty();
 129                return *LHS < *RHS;
 130              });
 131
 132   LLVM_DEBUG(dbgs() << "Coloring register intervals:\n");
 133   SmallVector<unsigned, 16> SlotMapping(SortedIntervals.size(), -1u);
 134   SmallVector<SmallVector<LiveInterval *, 4>, 16> Assignments(
 135       SortedIntervals.size());
 136   BitVector UsedColors(SortedIntervals.size());
 137   bool Changed = false;
 138   for (size_t i = 0, e = SortedIntervals.size(); i < e; ++i) {
 139     LiveInterval *LI = SortedIntervals[i];
 140     unsigned Old = LI->reg;
 141     size_t Color = i;
 142     const TargetRegisterClass *RC = MRI->getRegClass(Old);
 143
 144     // Check if it's possible to reuse any of the used colors.
 145     if (!MRI->isLiveIn(Old))
 146       for (unsigned C : UsedColors.set_bits()) {
 147         if (MRI->getRegClass(SortedIntervals[C]->reg) != RC)
 148           continue;
 149         for (LiveInterval *OtherLI : Assignments[C])
 150           if (!OtherLI->empty() && OtherLI->overlaps(*LI))
 151             goto continue_outer;
 152         Color = C;
 153         break;
 154       continue_outer:;
 155       }
 156
 157     unsigned New = SortedIntervals[Color]->reg;
 158     SlotMapping[i] = New;
 159     Changed |= Old != New;
 160     UsedColors.set(Color);
 161     Assignments[Color].push_back(LI);
 162     LLVM_DEBUG(
 163         dbgs() << "Assigning vreg" << TargetRegisterInfo::virtReg2Index(LI->reg)
 164                << " to vreg" << TargetRegisterInfo::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 }