lib/Target/WebAssembly/WebAssemblyRegColoring.cpp

   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"
  27 using namespace llvm;
  28
  29 #define DEBUG_TYPE "wasm-reg-coloring"
  30
  31 namespace {
  32 class WebAssemblyRegColoring final : public MachineFunctionPass {
  33 public:
  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
  52 private:
  53 };
  54 } // end anonymous namespace
  55
  56 char WebAssemblyRegColoring::ID = 0;
  57 INITIALIZE_PASS(WebAssemblyRegColoring, DEBUG_TYPE,
  58                 "Minimize number of registers used", false, false)
  59
  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   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 }