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 /// \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     // Skip unused registers, which can use $discard.
 141     if (MRI->use_empty(VReg))
 142       continue;
 143
 144     LiveInterval *LI = &Liveness->getInterval(VReg);
 145     assert(LI->weight == 0.0f);
 146     LI->weight = computeWeight(MRI, MBFI, VReg);
 147     DEBUG(LI->dump());
 148     SortedIntervals.push_back(LI);
 149   }
 150   DEBUG(dbgs() << '\n');
 151
 152   // Sort them to put arguments first (since we don't want to rename live-in
 153   // registers), by weight next, and then by position.
 154   // TODO: Investigate more intelligent sorting heuristics. For starters, we
 155   // should try to coalesce adjacent live intervals before non-adjacent ones.
 156   std::sort(SortedIntervals.begin(), SortedIntervals.end(),
 157             [MRI](LiveInterval *LHS, LiveInterval *RHS) {
 158               if (MRI->isLiveIn(LHS->reg) != MRI->isLiveIn(RHS->reg))
 159                 return MRI->isLiveIn(LHS->reg);
 160               if (LHS->weight != RHS->weight)
 161                 return LHS->weight > RHS->weight;
 162               if (LHS->empty() || RHS->empty())
 163                 return !LHS->empty() && RHS->empty();
 164               return *LHS < *RHS;
 165             });
 166
 167   DEBUG(dbgs() << "Coloring register intervals:\n");
 168   SmallVector<unsigned, 16> SlotMapping(SortedIntervals.size(), -1u);
 169   SmallVector<SmallVector<LiveInterval *, 4>, 16> Assignments(
 170       SortedIntervals.size());
 171   BitVector UsedColors(SortedIntervals.size());
 172   bool Changed = false;
 173   for (size_t i = 0, e = SortedIntervals.size(); i < e; ++i) {
 174     LiveInterval *LI = SortedIntervals[i];
 175     unsigned Old = LI->reg;
 176     size_t Color = i;
 177     const TargetRegisterClass *RC = MRI->getRegClass(Old);
 178
 179     // Check if it's possible to reuse any of the used colors.
 180     if (!MRI->isLiveIn(Old))
 181       for (int C(UsedColors.find_first()); C != -1;
 182            C = UsedColors.find_next(C)) {
 183         if (MRI->getRegClass(SortedIntervals[C]->reg) != RC)
 184           continue;
 185         for (LiveInterval *OtherLI : Assignments[C])
 186           if (!OtherLI->empty() && OtherLI->overlaps(*LI))
 187             goto continue_outer;
 188         Color = C;
 189         break;
 190       continue_outer:;
 191       }
 192
 193     unsigned New = SortedIntervals[Color]->reg;
 194     SlotMapping[i] = New;
 195     Changed |= Old != New;
 196     UsedColors.set(Color);
 197     Assignments[Color].push_back(LI);
 198     DEBUG(dbgs() << "Assigning vreg"
 199                  << TargetRegisterInfo::virtReg2Index(LI->reg) << " to vreg"
 200                  << TargetRegisterInfo::virtReg2Index(New) << "\n");
 201   }
 202   if (!Changed)
 203     return false;
 204
 205   // Rewrite register operands.
 206   for (size_t i = 0, e = SortedIntervals.size(); i < e; ++i) {
 207     unsigned Old = SortedIntervals[i]->reg;
 208     unsigned New = SlotMapping[i];
 209     if (Old != New)
 210       MRI->replaceRegWith(Old, New);
 211   }
 212   return true;
 213 }