lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp

   1 //===-- WebAssemblyExplicitLocals.cpp - Make Locals Explicit --------------===//
   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 converts any remaining registers into WebAssembly locals.
  11 ///
  12 /// After register stackification and register coloring, convert non-stackified
  13 /// registers into locals, inserting explicit local.get and local.set
  14 /// instructions.
  15 ///
  16 //===----------------------------------------------------------------------===//
  17
  18 #include "MCTargetDesc/WebAssemblyMCTargetDesc.h"
  19 #include "WebAssembly.h"
  20 #include "WebAssemblyMachineFunctionInfo.h"
  21 #include "WebAssemblySubtarget.h"
  22 #include "WebAssemblyUtilities.h"
  23 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
  24 #include "llvm/CodeGen/MachineInstrBuilder.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-explicit-locals"
  32
  33 // A command-line option to disable this pass, and keep implicit locals
  34 // for the purpose of testing with lit/llc ONLY.
  35 // This produces output which is not valid WebAssembly, and is not supported
  36 // by assemblers/disassemblers and other MC based tools.
  37 static cl::opt<bool> WasmDisableExplicitLocals(
  38     "wasm-disable-explicit-locals", cl::Hidden,
  39     cl::desc("WebAssembly: output implicit locals in"
  40              " instruction output for test purposes only."),
  41     cl::init(false));
  42
  43 namespace {
  44 class WebAssemblyExplicitLocals final : public MachineFunctionPass {
  45   StringRef getPassName() const override {
  46     return "WebAssembly Explicit Locals";
  47   }
  48
  49   void getAnalysisUsage(AnalysisUsage &AU) const override {
  50     AU.setPreservesCFG();
  51     AU.addPreserved<MachineBlockFrequencyInfo>();
  52     MachineFunctionPass::getAnalysisUsage(AU);
  53   }
  54
  55   bool runOnMachineFunction(MachineFunction &MF) override;
  56
  57 public:
  58   static char ID; // Pass identification, replacement for typeid
  59   WebAssemblyExplicitLocals() : MachineFunctionPass(ID) {}
  60 };
  61 } // end anonymous namespace
  62
  63 char WebAssemblyExplicitLocals::ID = 0;
  64 INITIALIZE_PASS(WebAssemblyExplicitLocals, DEBUG_TYPE,
  65                 "Convert registers to WebAssembly locals", false, false)
  66
  67 FunctionPass *llvm::createWebAssemblyExplicitLocals() {
  68   return new WebAssemblyExplicitLocals();
  69 }
  70
  71 /// Return a local id number for the given register, assigning it a new one
  72 /// if it doesn't yet have one.
  73 static unsigned getLocalId(DenseMap<unsigned, unsigned> &Reg2Local,
  74                            unsigned &CurLocal, unsigned Reg) {
  75   auto P = Reg2Local.insert(std::make_pair(Reg, CurLocal));
  76   if (P.second)
  77     ++CurLocal;
  78   return P.first->second;
  79 }
  80
  81 /// Get the appropriate drop opcode for the given register class.
  82 static unsigned getDropOpcode(const TargetRegisterClass *RC) {
  83   if (RC == &WebAssembly::I32RegClass)
  84     return WebAssembly::DROP_I32;
  85   if (RC == &WebAssembly::I64RegClass)
  86     return WebAssembly::DROP_I64;
  87   if (RC == &WebAssembly::F32RegClass)
  88     return WebAssembly::DROP_F32;
  89   if (RC == &WebAssembly::F64RegClass)
  90     return WebAssembly::DROP_F64;
  91   if (RC == &WebAssembly::V128RegClass)
  92     return WebAssembly::DROP_V128;
  93   if (RC == &WebAssembly::EXNREFRegClass)
  94     return WebAssembly::DROP_EXNREF;
  95   llvm_unreachable("Unexpected register class");
  96 }
  97
  98 /// Get the appropriate local.get opcode for the given register class.
  99 static unsigned getLocalGetOpcode(const TargetRegisterClass *RC) {
 100   if (RC == &WebAssembly::I32RegClass)
 101     return WebAssembly::LOCAL_GET_I32;
 102   if (RC == &WebAssembly::I64RegClass)
 103     return WebAssembly::LOCAL_GET_I64;
 104   if (RC == &WebAssembly::F32RegClass)
 105     return WebAssembly::LOCAL_GET_F32;
 106   if (RC == &WebAssembly::F64RegClass)
 107     return WebAssembly::LOCAL_GET_F64;
 108   if (RC == &WebAssembly::V128RegClass)
 109     return WebAssembly::LOCAL_GET_V128;
 110   if (RC == &WebAssembly::EXNREFRegClass)
 111     return WebAssembly::LOCAL_GET_EXNREF;
 112   llvm_unreachable("Unexpected register class");
 113 }
 114
 115 /// Get the appropriate local.set opcode for the given register class.
 116 static unsigned getLocalSetOpcode(const TargetRegisterClass *RC) {
 117   if (RC == &WebAssembly::I32RegClass)
 118     return WebAssembly::LOCAL_SET_I32;
 119   if (RC == &WebAssembly::I64RegClass)
 120     return WebAssembly::LOCAL_SET_I64;
 121   if (RC == &WebAssembly::F32RegClass)
 122     return WebAssembly::LOCAL_SET_F32;
 123   if (RC == &WebAssembly::F64RegClass)
 124     return WebAssembly::LOCAL_SET_F64;
 125   if (RC == &WebAssembly::V128RegClass)
 126     return WebAssembly::LOCAL_SET_V128;
 127   if (RC == &WebAssembly::EXNREFRegClass)
 128     return WebAssembly::LOCAL_SET_EXNREF;
 129   llvm_unreachable("Unexpected register class");
 130 }
 131
 132 /// Get the appropriate local.tee opcode for the given register class.
 133 static unsigned getLocalTeeOpcode(const TargetRegisterClass *RC) {
 134   if (RC == &WebAssembly::I32RegClass)
 135     return WebAssembly::LOCAL_TEE_I32;
 136   if (RC == &WebAssembly::I64RegClass)
 137     return WebAssembly::LOCAL_TEE_I64;
 138   if (RC == &WebAssembly::F32RegClass)
 139     return WebAssembly::LOCAL_TEE_F32;
 140   if (RC == &WebAssembly::F64RegClass)
 141     return WebAssembly::LOCAL_TEE_F64;
 142   if (RC == &WebAssembly::V128RegClass)
 143     return WebAssembly::LOCAL_TEE_V128;
 144   if (RC == &WebAssembly::EXNREFRegClass)
 145     return WebAssembly::LOCAL_TEE_EXNREF;
 146   llvm_unreachable("Unexpected register class");
 147 }
 148
 149 /// Get the type associated with the given register class.
 150 static MVT typeForRegClass(const TargetRegisterClass *RC) {
 151   if (RC == &WebAssembly::I32RegClass)
 152     return MVT::i32;
 153   if (RC == &WebAssembly::I64RegClass)
 154     return MVT::i64;
 155   if (RC == &WebAssembly::F32RegClass)
 156     return MVT::f32;
 157   if (RC == &WebAssembly::F64RegClass)
 158     return MVT::f64;
 159   if (RC == &WebAssembly::V128RegClass)
 160     return MVT::v16i8;
 161   if (RC == &WebAssembly::EXNREFRegClass)
 162     return MVT::exnref;
 163   llvm_unreachable("unrecognized register class");
 164 }
 165
 166 /// Given a MachineOperand of a stackified vreg, return the instruction at the
 167 /// start of the expression tree.
 168 static MachineInstr *findStartOfTree(MachineOperand &MO,
 169                                      MachineRegisterInfo &MRI,
 170                                      WebAssemblyFunctionInfo &MFI) {
 171   Register Reg = MO.getReg();
 172   assert(MFI.isVRegStackified(Reg));
 173   MachineInstr *Def = MRI.getVRegDef(Reg);
 174
 175   // Find the first stackified use and proceed from there.
 176   for (MachineOperand &DefMO : Def->explicit_uses()) {
 177     if (!DefMO.isReg())
 178       continue;
 179     return findStartOfTree(DefMO, MRI, MFI);
 180   }
 181
 182   // If there were no stackified uses, we've reached the start.
 183   return Def;
 184 }
 185
 186 bool WebAssemblyExplicitLocals::runOnMachineFunction(MachineFunction &MF) {
 187   LLVM_DEBUG(dbgs() << "********** Make Locals Explicit **********\n"
 188                        "********** Function: "
 189                     << MF.getName() << '\n');
 190
 191   // Disable this pass if directed to do so.
 192   if (WasmDisableExplicitLocals)
 193     return false;
 194
 195   bool Changed = false;
 196   MachineRegisterInfo &MRI = MF.getRegInfo();
 197   WebAssemblyFunctionInfo &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();
 198   const auto *TII = MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
 199
 200   // Map non-stackified virtual registers to their local ids.
 201   DenseMap<unsigned, unsigned> Reg2Local;
 202
 203   // Handle ARGUMENTS first to ensure that they get the designated numbers.
 204   for (MachineBasicBlock::iterator I = MF.begin()->begin(),
 205                                    E = MF.begin()->end();
 206        I != E;) {
 207     MachineInstr &MI = *I++;
 208     if (!WebAssembly::isArgument(MI.getOpcode()))
 209       break;
 210     Register Reg = MI.getOperand(0).getReg();
 211     assert(!MFI.isVRegStackified(Reg));
 212     Reg2Local[Reg] = static_cast<unsigned>(MI.getOperand(1).getImm());
 213     MI.eraseFromParent();
 214     Changed = true;
 215   }
 216
 217   // Start assigning local numbers after the last parameter.
 218   unsigned CurLocal = static_cast<unsigned>(MFI.getParams().size());
 219
 220   // Precompute the set of registers that are unused, so that we can insert
 221   // drops to their defs.
 222   BitVector UseEmpty(MRI.getNumVirtRegs());
 223   for (unsigned I = 0, E = MRI.getNumVirtRegs(); I < E; ++I)
 224     UseEmpty[I] = MRI.use_empty(Register::index2VirtReg(I));
 225
 226   // Visit each instruction in the function.
 227   for (MachineBasicBlock &MBB : MF) {
 228     for (MachineBasicBlock::iterator I = MBB.begin(), E = MBB.end(); I != E;) {
 229       MachineInstr &MI = *I++;
 230       assert(!WebAssembly::isArgument(MI.getOpcode()));
 231
 232       if (MI.isDebugInstr() || MI.isLabel())
 233         continue;
 234
 235       // Replace tee instructions with local.tee. The difference is that tee
 236       // instructions have two defs, while local.tee instructions have one def
 237       // and an index of a local to write to.
 238       if (WebAssembly::isTee(MI.getOpcode())) {
 239         assert(MFI.isVRegStackified(MI.getOperand(0).getReg()));
 240         assert(!MFI.isVRegStackified(MI.getOperand(1).getReg()));
 241         Register OldReg = MI.getOperand(2).getReg();
 242         const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
 243
 244         // Stackify the input if it isn't stackified yet.
 245         if (!MFI.isVRegStackified(OldReg)) {
 246           unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
 247           Register NewReg = MRI.createVirtualRegister(RC);
 248           unsigned Opc = getLocalGetOpcode(RC);
 249           BuildMI(MBB, &MI, MI.getDebugLoc(), TII->get(Opc), NewReg)
 250               .addImm(LocalId);
 251           MI.getOperand(2).setReg(NewReg);
 252           MFI.stackifyVReg(NewReg);
 253         }
 254
 255         // Replace the TEE with a LOCAL_TEE.
 256         unsigned LocalId =
 257             getLocalId(Reg2Local, CurLocal, MI.getOperand(1).getReg());
 258         unsigned Opc = getLocalTeeOpcode(RC);
 259         BuildMI(MBB, &MI, MI.getDebugLoc(), TII->get(Opc),
 260                 MI.getOperand(0).getReg())
 261             .addImm(LocalId)
 262             .addReg(MI.getOperand(2).getReg());
 263
 264         MI.eraseFromParent();
 265         Changed = true;
 266         continue;
 267       }
 268
 269       // Insert local.sets for any defs that aren't stackified yet. Currently
 270       // we handle at most one def.
 271       assert(MI.getDesc().getNumDefs() <= 1);
 272       if (MI.getDesc().getNumDefs() == 1) {
 273         Register OldReg = MI.getOperand(0).getReg();
 274         if (!MFI.isVRegStackified(OldReg)) {
 275           const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
 276           Register NewReg = MRI.createVirtualRegister(RC);
 277           auto InsertPt = std::next(MI.getIterator());
 278           if (MI.getOpcode() == WebAssembly::IMPLICIT_DEF) {
 279             MI.eraseFromParent();
 280             Changed = true;
 281             continue;
 282           }
 283           if (UseEmpty[Register::virtReg2Index(OldReg)]) {
 284             unsigned Opc = getDropOpcode(RC);
 285             MachineInstr *Drop =
 286                 BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc))
 287                     .addReg(NewReg);
 288             // After the drop instruction, this reg operand will not be used
 289             Drop->getOperand(0).setIsKill();
 290           } else {
 291             unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
 292             unsigned Opc = getLocalSetOpcode(RC);
 293             BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc))
 294                 .addImm(LocalId)
 295                 .addReg(NewReg);
 296           }
 297           MI.getOperand(0).setReg(NewReg);
 298           // This register operand of the original instruction is now being used
 299           // by the inserted drop or local.set instruction, so make it not dead
 300           // yet.
 301           MI.getOperand(0).setIsDead(false);
 302           MFI.stackifyVReg(NewReg);
 303           Changed = true;
 304         }
 305       }
 306
 307       // Insert local.gets for any uses that aren't stackified yet.
 308       MachineInstr *InsertPt = &MI;
 309       for (MachineOperand &MO : reverse(MI.explicit_uses())) {
 310         if (!MO.isReg())
 311           continue;
 312
 313         Register OldReg = MO.getReg();
 314
 315         // Inline asm may have a def in the middle of the operands. Our contract
 316         // with inline asm register operands is to provide local indices as
 317         // immediates.
 318         if (MO.isDef()) {
 319           assert(MI.isInlineAsm());
 320           unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
 321           // If this register operand is tied to another operand, we can't
 322           // change it to an immediate. Untie it first.
 323           MI.untieRegOperand(MI.getOperandNo(&MO));
 324           MO.ChangeToImmediate(LocalId);
 325           continue;
 326         }
 327
 328         // If we see a stackified register, prepare to insert subsequent
 329         // local.gets before the start of its tree.
 330         if (MFI.isVRegStackified(OldReg)) {
 331           InsertPt = findStartOfTree(MO, MRI, MFI);
 332           continue;
 333         }
 334
 335         // Our contract with inline asm register operands is to provide local
 336         // indices as immediates.
 337         if (MI.isInlineAsm()) {
 338           unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
 339           // Untie it first if this reg operand is tied to another operand.
 340           MI.untieRegOperand(MI.getOperandNo(&MO));
 341           MO.ChangeToImmediate(LocalId);
 342           continue;
 343         }
 344
 345         // Insert a local.get.
 346         unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
 347         const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
 348         Register NewReg = MRI.createVirtualRegister(RC);
 349         unsigned Opc = getLocalGetOpcode(RC);
 350         InsertPt =
 351             BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc), NewReg)
 352                 .addImm(LocalId);
 353         MO.setReg(NewReg);
 354         MFI.stackifyVReg(NewReg);
 355         Changed = true;
 356       }
 357
 358       // Coalesce and eliminate COPY instructions.
 359       if (WebAssembly::isCopy(MI.getOpcode())) {
 360         MRI.replaceRegWith(MI.getOperand(1).getReg(),
 361                            MI.getOperand(0).getReg());
 362         MI.eraseFromParent();
 363         Changed = true;
 364       }
 365     }
 366   }
 367
 368   // Define the locals.
 369   // TODO: Sort the locals for better compression.
 370   MFI.setNumLocals(CurLocal - MFI.getParams().size());
 371   for (unsigned I = 0, E = MRI.getNumVirtRegs(); I < E; ++I) {
 372     unsigned Reg = Register::index2VirtReg(I);
 373     auto RL = Reg2Local.find(Reg);
 374     if (RL == Reg2Local.end() || RL->second < MFI.getParams().size())
 375       continue;
 376
 377     MFI.setLocal(RL->second - MFI.getParams().size(),
 378                  typeForRegClass(MRI.getRegClass(Reg)));
 379     Changed = true;
 380   }
 381
 382 #ifndef NDEBUG
 383   // Assert that all registers have been stackified at this point.
 384   for (const MachineBasicBlock &MBB : MF) {
 385     for (const MachineInstr &MI : MBB) {
 386       if (MI.isDebugInstr() || MI.isLabel())
 387         continue;
 388       for (const MachineOperand &MO : MI.explicit_operands()) {
 389         assert(
 390             (!MO.isReg() || MRI.use_empty(MO.getReg()) ||
 391              MFI.isVRegStackified(MO.getReg())) &&
 392             "WebAssemblyExplicitLocals failed to stackify a register operand");
 393       }
 394     }
 395   }
 396 #endif
 397
 398   return Changed;
 399 }