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 "Utils/WebAssemblyUtilities.h"
20#include "WebAssembly.h"
21#include "WebAssemblyDebugValueManager.h"
22#include "WebAssemblyMachineFunctionInfo.h"
23#include "WebAssemblySubtarget.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::FUNCREFRegClass)
100    return WebAssembly::DROP_FUNCREF;
101  if (RC == &WebAssembly::EXTERNREFRegClass)
102    return WebAssembly::DROP_EXTERNREF;
103  llvm_unreachable("Unexpected register class");
104}
105
106/// Get the appropriate local.get opcode for the given register class.
107static unsigned getLocalGetOpcode(const TargetRegisterClass *RC) {
108  if (RC == &WebAssembly::I32RegClass)
109    return WebAssembly::LOCAL_GET_I32;
110  if (RC == &WebAssembly::I64RegClass)
111    return WebAssembly::LOCAL_GET_I64;
112  if (RC == &WebAssembly::F32RegClass)
113    return WebAssembly::LOCAL_GET_F32;
114  if (RC == &WebAssembly::F64RegClass)
115    return WebAssembly::LOCAL_GET_F64;
116  if (RC == &WebAssembly::V128RegClass)
117    return WebAssembly::LOCAL_GET_V128;
118  if (RC == &WebAssembly::FUNCREFRegClass)
119    return WebAssembly::LOCAL_GET_FUNCREF;
120  if (RC == &WebAssembly::EXTERNREFRegClass)
121    return WebAssembly::LOCAL_GET_EXTERNREF;
122  llvm_unreachable("Unexpected register class");
123}
124
125/// Get the appropriate local.set opcode for the given register class.
126static unsigned getLocalSetOpcode(const TargetRegisterClass *RC) {
127  if (RC == &WebAssembly::I32RegClass)
128    return WebAssembly::LOCAL_SET_I32;
129  if (RC == &WebAssembly::I64RegClass)
130    return WebAssembly::LOCAL_SET_I64;
131  if (RC == &WebAssembly::F32RegClass)
132    return WebAssembly::LOCAL_SET_F32;
133  if (RC == &WebAssembly::F64RegClass)
134    return WebAssembly::LOCAL_SET_F64;
135  if (RC == &WebAssembly::V128RegClass)
136    return WebAssembly::LOCAL_SET_V128;
137  if (RC == &WebAssembly::FUNCREFRegClass)
138    return WebAssembly::LOCAL_SET_FUNCREF;
139  if (RC == &WebAssembly::EXTERNREFRegClass)
140    return WebAssembly::LOCAL_SET_EXTERNREF;
141  llvm_unreachable("Unexpected register class");
142}
143
144/// Get the appropriate local.tee opcode for the given register class.
145static unsigned getLocalTeeOpcode(const TargetRegisterClass *RC) {
146  if (RC == &WebAssembly::I32RegClass)
147    return WebAssembly::LOCAL_TEE_I32;
148  if (RC == &WebAssembly::I64RegClass)
149    return WebAssembly::LOCAL_TEE_I64;
150  if (RC == &WebAssembly::F32RegClass)
151    return WebAssembly::LOCAL_TEE_F32;
152  if (RC == &WebAssembly::F64RegClass)
153    return WebAssembly::LOCAL_TEE_F64;
154  if (RC == &WebAssembly::V128RegClass)
155    return WebAssembly::LOCAL_TEE_V128;
156  if (RC == &WebAssembly::FUNCREFRegClass)
157    return WebAssembly::LOCAL_TEE_FUNCREF;
158  if (RC == &WebAssembly::EXTERNREFRegClass)
159    return WebAssembly::LOCAL_TEE_EXTERNREF;
160  llvm_unreachable("Unexpected register class");
161}
162
163/// Get the type associated with the given register class.
164static MVT typeForRegClass(const TargetRegisterClass *RC) {
165  if (RC == &WebAssembly::I32RegClass)
166    return MVT::i32;
167  if (RC == &WebAssembly::I64RegClass)
168    return MVT::i64;
169  if (RC == &WebAssembly::F32RegClass)
170    return MVT::f32;
171  if (RC == &WebAssembly::F64RegClass)
172    return MVT::f64;
173  if (RC == &WebAssembly::V128RegClass)
174    return MVT::v16i8;
175  if (RC == &WebAssembly::FUNCREFRegClass)
176    return MVT::funcref;
177  if (RC == &WebAssembly::EXTERNREFRegClass)
178    return MVT::externref;
179  llvm_unreachable("unrecognized register class");
180}
181
182/// Given a MachineOperand of a stackified vreg, return the instruction at the
183/// start of the expression tree.
184static MachineInstr *findStartOfTree(MachineOperand &MO,
185                                     MachineRegisterInfo &MRI,
186                                     const WebAssemblyFunctionInfo &MFI) {
187  Register Reg = MO.getReg();
188  assert(MFI.isVRegStackified(Reg));
189  MachineInstr *Def = MRI.getVRegDef(Reg);
190
191  // If this instruction has any non-stackified defs, it is the start
192  for (auto DefReg : Def->defs()) {
193    if (!MFI.isVRegStackified(DefReg.getReg())) {
194      return Def;
195    }
196  }
197
198  // Find the first stackified use and proceed from there.
199  for (MachineOperand &DefMO : Def->explicit_uses()) {
200    if (!DefMO.isReg())
201      continue;
202    return findStartOfTree(DefMO, MRI, MFI);
203  }
204
205  // If there were no stackified uses, we've reached the start.
206  return Def;
207}
208
209bool WebAssemblyExplicitLocals::runOnMachineFunction(MachineFunction &MF) {
210  LLVM_DEBUG(dbgs() << "********** Make Locals Explicit **********\n"
211                       "********** Function: "
212                    << MF.getName() << '\n');
213
214  bool Changed = false;
215  MachineRegisterInfo &MRI = MF.getRegInfo();
216  WebAssemblyFunctionInfo &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();
217  const auto *TII = MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
218
219  // Map non-stackified virtual registers to their local ids.
220  DenseMap<unsigned, unsigned> Reg2Local;
221
222  // Handle ARGUMENTS first to ensure that they get the designated numbers.
223  for (MachineBasicBlock::iterator I = MF.begin()->begin(),
224                                   E = MF.begin()->end();
225       I != E;) {
226    MachineInstr &MI = *I++;
227    if (!WebAssembly::isArgument(MI.getOpcode()))
228      break;
229    Register Reg = MI.getOperand(0).getReg();
230    assert(!MFI.isVRegStackified(Reg));
231    auto Local = static_cast<unsigned>(MI.getOperand(1).getImm());
232    Reg2Local[Reg] = Local;
233    checkFrameBase(MFI, Local, Reg);
234
235    // Update debug value to point to the local before removing.
236    WebAssemblyDebugValueManager(&MI).replaceWithLocal(Local);
237
238    MI.eraseFromParent();
239    Changed = true;
240  }
241
242  // Start assigning local numbers after the last parameter and after any
243  // already-assigned locals.
244  unsigned CurLocal = static_cast<unsigned>(MFI.getParams().size());
245  CurLocal += static_cast<unsigned>(MFI.getLocals().size());
246
247  // Precompute the set of registers that are unused, so that we can insert
248  // drops to their defs.
249  BitVector UseEmpty(MRI.getNumVirtRegs());
250  for (unsigned I = 0, E = MRI.getNumVirtRegs(); I < E; ++I)
251    UseEmpty[I] = MRI.use_empty(Register::index2VirtReg(I));
252
253  // Visit each instruction in the function.
254  for (MachineBasicBlock &MBB : MF) {
255    for (MachineInstr &MI : llvm::make_early_inc_range(MBB)) {
256      assert(!WebAssembly::isArgument(MI.getOpcode()));
257
258      if (MI.isDebugInstr() || MI.isLabel())
259        continue;
260
261      if (MI.getOpcode() == WebAssembly::IMPLICIT_DEF) {
262        MI.eraseFromParent();
263        Changed = true;
264        continue;
265      }
266
267      // Replace tee instructions with local.tee. The difference is that tee
268      // instructions have two defs, while local.tee instructions have one def
269      // and an index of a local to write to.
270      if (WebAssembly::isTee(MI.getOpcode())) {
271        assert(MFI.isVRegStackified(MI.getOperand(0).getReg()));
272        assert(!MFI.isVRegStackified(MI.getOperand(1).getReg()));
273        Register OldReg = MI.getOperand(2).getReg();
274        const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
275
276        // Stackify the input if it isn't stackified yet.
277        if (!MFI.isVRegStackified(OldReg)) {
278          unsigned LocalId = getLocalId(Reg2Local, MFI, CurLocal, OldReg);
279          Register NewReg = MRI.createVirtualRegister(RC);
280          unsigned Opc = getLocalGetOpcode(RC);
281          BuildMI(MBB, &MI, MI.getDebugLoc(), TII->get(Opc), NewReg)
282              .addImm(LocalId);
283          MI.getOperand(2).setReg(NewReg);
284          MFI.stackifyVReg(MRI, NewReg);
285        }
286
287        // Replace the TEE with a LOCAL_TEE.
288        unsigned LocalId =
289            getLocalId(Reg2Local, MFI, CurLocal, MI.getOperand(1).getReg());
290        unsigned Opc = getLocalTeeOpcode(RC);
291        BuildMI(MBB, &MI, MI.getDebugLoc(), TII->get(Opc),
292                MI.getOperand(0).getReg())
293            .addImm(LocalId)
294            .addReg(MI.getOperand(2).getReg());
295
296        WebAssemblyDebugValueManager(&MI).replaceWithLocal(LocalId);
297
298        MI.eraseFromParent();
299        Changed = true;
300        continue;
301      }
302
303      // Insert local.sets for any defs that aren't stackified yet.
304      for (auto &Def : MI.defs()) {
305        Register OldReg = Def.getReg();
306        if (!MFI.isVRegStackified(OldReg)) {
307          const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
308          Register NewReg = MRI.createVirtualRegister(RC);
309          auto InsertPt = std::next(MI.getIterator());
310          if (UseEmpty[Register::virtReg2Index(OldReg)]) {
311            unsigned Opc = getDropOpcode(RC);
312            MachineInstr *Drop =
313                BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc))
314                    .addReg(NewReg);
315            // After the drop instruction, this reg operand will not be used
316            Drop->getOperand(0).setIsKill();
317            if (MFI.isFrameBaseVirtual() && OldReg == MFI.getFrameBaseVreg())
318              MFI.clearFrameBaseVreg();
319          } else {
320            unsigned LocalId = getLocalId(Reg2Local, MFI, CurLocal, OldReg);
321            unsigned Opc = getLocalSetOpcode(RC);
322
323            WebAssemblyDebugValueManager(&MI).replaceWithLocal(LocalId);
324
325            BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc))
326                .addImm(LocalId)
327                .addReg(NewReg);
328          }
329          // This register operand of the original instruction is now being used
330          // by the inserted drop or local.set instruction, so make it not dead
331          // yet.
332          Def.setReg(NewReg);
333          Def.setIsDead(false);
334          MFI.stackifyVReg(MRI, NewReg);
335          Changed = true;
336        }
337      }
338
339      // Insert local.gets for any uses that aren't stackified yet.
340      MachineInstr *InsertPt = &MI;
341      for (MachineOperand &MO : reverse(MI.explicit_uses())) {
342        if (!MO.isReg())
343          continue;
344
345        Register OldReg = MO.getReg();
346
347        // Inline asm may have a def in the middle of the operands. Our contract
348        // with inline asm register operands is to provide local indices as
349        // immediates.
350        if (MO.isDef()) {
351          assert(MI.isInlineAsm());
352          unsigned LocalId = getLocalId(Reg2Local, MFI, CurLocal, OldReg);
353          // If this register operand is tied to another operand, we can't
354          // change it to an immediate. Untie it first.
355          MI.untieRegOperand(MI.getOperandNo(&MO));
356          MO.ChangeToImmediate(LocalId);
357          continue;
358        }
359
360        // If we see a stackified register, prepare to insert subsequent
361        // local.gets before the start of its tree.
362        if (MFI.isVRegStackified(OldReg)) {
363          InsertPt = findStartOfTree(MO, MRI, MFI);
364          continue;
365        }
366
367        // Our contract with inline asm register operands is to provide local
368        // indices as immediates.
369        if (MI.isInlineAsm()) {
370          unsigned LocalId = getLocalId(Reg2Local, MFI, CurLocal, OldReg);
371          // Untie it first if this reg operand is tied to another operand.
372          MI.untieRegOperand(MI.getOperandNo(&MO));
373          MO.ChangeToImmediate(LocalId);
374          continue;
375        }
376
377        // Insert a local.get.
378        unsigned LocalId = getLocalId(Reg2Local, MFI, CurLocal, OldReg);
379        const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
380        Register NewReg = MRI.createVirtualRegister(RC);
381        unsigned Opc = getLocalGetOpcode(RC);
382        // Use a InsertPt as our DebugLoc, since MI may be discontinuous from
383        // the where this local is being inserted, causing non-linear stepping
384        // in the debugger or function entry points where variables aren't live
385        // yet. Alternative is previous instruction, but that is strictly worse
386        // since it can point at the previous statement.
387        // See crbug.com/1251909, crbug.com/1249745
388        InsertPt = BuildMI(MBB, InsertPt, InsertPt->getDebugLoc(),
389                           TII->get(Opc), NewReg).addImm(LocalId);
390        MO.setReg(NewReg);
391        MFI.stackifyVReg(MRI, NewReg);
392        Changed = true;
393      }
394
395      // Coalesce and eliminate COPY instructions.
396      if (WebAssembly::isCopy(MI.getOpcode())) {
397        MRI.replaceRegWith(MI.getOperand(1).getReg(),
398                           MI.getOperand(0).getReg());
399        MI.eraseFromParent();
400        Changed = true;
401      }
402    }
403  }
404
405  // Define the locals.
406  // TODO: Sort the locals for better compression.
407  MFI.setNumLocals(CurLocal - MFI.getParams().size());
408  for (unsigned I = 0, E = MRI.getNumVirtRegs(); I < E; ++I) {
409    Register Reg = Register::index2VirtReg(I);
410    auto RL = Reg2Local.find(Reg);
411    if (RL == Reg2Local.end() || RL->second < MFI.getParams().size())
412      continue;
413
414    MFI.setLocal(RL->second - MFI.getParams().size(),
415                 typeForRegClass(MRI.getRegClass(Reg)));
416    Changed = true;
417  }
418
419#ifndef NDEBUG
420  // Assert that all registers have been stackified at this point.
421  for (const MachineBasicBlock &MBB : MF) {
422    for (const MachineInstr &MI : MBB) {
423      if (MI.isDebugInstr() || MI.isLabel())
424        continue;
425      for (const MachineOperand &MO : MI.explicit_operands()) {
426        assert(
427            (!MO.isReg() || MRI.use_empty(MO.getReg()) ||
428             MFI.isVRegStackified(MO.getReg())) &&
429            "WebAssemblyExplicitLocals failed to stackify a register operand");
430      }
431    }
432  }
433#endif
434
435  return Changed;
436}
437