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