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"
30 using namespace llvm;
31 
32 #define DEBUG_TYPE "wasm-explicit-locals"
33 
34 namespace {
35 class 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 
48 public:
49   static char ID; // Pass identification, replacement for typeid
50   WebAssemblyExplicitLocals() : MachineFunctionPass(ID) {}
51 };
52 } // end anonymous namespace
53 
54 char WebAssemblyExplicitLocals::ID = 0;
55 INITIALIZE_PASS(WebAssemblyExplicitLocals, DEBUG_TYPE,
56                 "Convert registers to WebAssembly locals", false, false)
57 
58 FunctionPass *llvm::createWebAssemblyExplicitLocals() {
59   return new WebAssemblyExplicitLocals();
60 }
61 
62 static 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.
76 static 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.
88 static 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.
107 static 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.
126 static 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.
145 static 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.
164 static 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.
184 static 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 
209 bool 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       //
271       // - Before:
272       // TeeReg, Reg = TEE DefReg
273       // INST ..., TeeReg, ...
274       // INST ..., Reg, ...
275       // INST ..., Reg, ...
276       // * DefReg: may or may not be stackified
277       // * Reg: not stackified
278       // * TeeReg: stackified
279       //
280       // - After (when DefReg was already stackified):
281       // TeeReg = LOCAL_TEE LocalId1, DefReg
282       // INST ..., TeeReg, ...
283       // INST ..., Reg, ...
284       // INST ..., Reg, ...
285       // * Reg: mapped to LocalId1
286       // * TeeReg: stackified
287       //
288       // - After (when DefReg was not already stackified):
289       // NewReg = LOCAL_GET LocalId1
290       // TeeReg = LOCAL_TEE LocalId2, NewReg
291       // INST ..., TeeReg, ...
292       // INST ..., Reg, ...
293       // INST ..., Reg, ...
294       // * DefReg: mapped to LocalId1
295       // * Reg: mapped to LocalId2
296       // * TeeReg: stackified
297       if (WebAssembly::isTee(MI.getOpcode())) {
298         assert(MFI.isVRegStackified(MI.getOperand(0).getReg()));
299         assert(!MFI.isVRegStackified(MI.getOperand(1).getReg()));
300         Register DefReg = MI.getOperand(2).getReg();
301         const TargetRegisterClass *RC = MRI.getRegClass(DefReg);
302 
303         // Stackify the input if it isn't stackified yet.
304         if (!MFI.isVRegStackified(DefReg)) {
305           unsigned LocalId = getLocalId(Reg2Local, MFI, CurLocal, DefReg);
306           Register NewReg = MRI.createVirtualRegister(RC);
307           unsigned Opc = getLocalGetOpcode(RC);
308           BuildMI(MBB, &MI, MI.getDebugLoc(), TII->get(Opc), NewReg)
309               .addImm(LocalId);
310           MI.getOperand(2).setReg(NewReg);
311           MFI.stackifyVReg(MRI, NewReg);
312         }
313 
314         // Replace the TEE with a LOCAL_TEE.
315         unsigned LocalId =
316             getLocalId(Reg2Local, MFI, CurLocal, MI.getOperand(1).getReg());
317         unsigned Opc = getLocalTeeOpcode(RC);
318         BuildMI(MBB, &MI, MI.getDebugLoc(), TII->get(Opc),
319                 MI.getOperand(0).getReg())
320             .addImm(LocalId)
321             .addReg(MI.getOperand(2).getReg());
322 
323         WebAssemblyDebugValueManager(&MI).replaceWithLocal(LocalId);
324 
325         MI.eraseFromParent();
326         Changed = true;
327         continue;
328       }
329 
330       // Insert local.sets for any defs that aren't stackified yet.
331       for (auto &Def : MI.defs()) {
332         Register OldReg = Def.getReg();
333         if (!MFI.isVRegStackified(OldReg)) {
334           const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
335           Register NewReg = MRI.createVirtualRegister(RC);
336           auto InsertPt = std::next(MI.getIterator());
337           if (UseEmpty[Register::virtReg2Index(OldReg)]) {
338             unsigned Opc = getDropOpcode(RC);
339             MachineInstr *Drop =
340                 BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc))
341                     .addReg(NewReg);
342             // After the drop instruction, this reg operand will not be used
343             Drop->getOperand(0).setIsKill();
344             if (MFI.isFrameBaseVirtual() && OldReg == MFI.getFrameBaseVreg())
345               MFI.clearFrameBaseVreg();
346           } else {
347             unsigned LocalId = getLocalId(Reg2Local, MFI, CurLocal, OldReg);
348             unsigned Opc = getLocalSetOpcode(RC);
349 
350             WebAssemblyDebugValueManager(&MI).replaceWithLocal(LocalId);
351 
352             BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc))
353                 .addImm(LocalId)
354                 .addReg(NewReg);
355           }
356           // This register operand of the original instruction is now being used
357           // by the inserted drop or local.set instruction, so make it not dead
358           // yet.
359           Def.setReg(NewReg);
360           Def.setIsDead(false);
361           MFI.stackifyVReg(MRI, NewReg);
362           Changed = true;
363         }
364       }
365 
366       // Insert local.gets for any uses that aren't stackified yet.
367       MachineInstr *InsertPt = &MI;
368       for (MachineOperand &MO : reverse(MI.explicit_uses())) {
369         if (!MO.isReg())
370           continue;
371 
372         Register OldReg = MO.getReg();
373 
374         // Inline asm may have a def in the middle of the operands. Our contract
375         // with inline asm register operands is to provide local indices as
376         // immediates.
377         if (MO.isDef()) {
378           assert(MI.isInlineAsm());
379           unsigned LocalId = getLocalId(Reg2Local, MFI, CurLocal, OldReg);
380           // If this register operand is tied to another operand, we can't
381           // change it to an immediate. Untie it first.
382           MI.untieRegOperand(MO.getOperandNo());
383           MO.ChangeToImmediate(LocalId);
384           continue;
385         }
386 
387         // If we see a stackified register, prepare to insert subsequent
388         // local.gets before the start of its tree.
389         if (MFI.isVRegStackified(OldReg)) {
390           InsertPt = findStartOfTree(MO, MRI, MFI);
391           continue;
392         }
393 
394         // Our contract with inline asm register operands is to provide local
395         // indices as immediates.
396         if (MI.isInlineAsm()) {
397           unsigned LocalId = getLocalId(Reg2Local, MFI, CurLocal, OldReg);
398           // Untie it first if this reg operand is tied to another operand.
399           MI.untieRegOperand(MO.getOperandNo());
400           MO.ChangeToImmediate(LocalId);
401           continue;
402         }
403 
404         // Insert a local.get.
405         unsigned LocalId = getLocalId(Reg2Local, MFI, CurLocal, OldReg);
406         const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
407         Register NewReg = MRI.createVirtualRegister(RC);
408         unsigned Opc = getLocalGetOpcode(RC);
409         // Use a InsertPt as our DebugLoc, since MI may be discontinuous from
410         // the where this local is being inserted, causing non-linear stepping
411         // in the debugger or function entry points where variables aren't live
412         // yet. Alternative is previous instruction, but that is strictly worse
413         // since it can point at the previous statement.
414         // See crbug.com/1251909, crbug.com/1249745
415         InsertPt = BuildMI(MBB, InsertPt, InsertPt->getDebugLoc(),
416                            TII->get(Opc), NewReg).addImm(LocalId);
417         MO.setReg(NewReg);
418         MFI.stackifyVReg(MRI, NewReg);
419         Changed = true;
420       }
421 
422       // Coalesce and eliminate COPY instructions.
423       if (WebAssembly::isCopy(MI.getOpcode())) {
424         MRI.replaceRegWith(MI.getOperand(1).getReg(),
425                            MI.getOperand(0).getReg());
426         MI.eraseFromParent();
427         Changed = true;
428       }
429     }
430   }
431 
432   // Define the locals.
433   // TODO: Sort the locals for better compression.
434   MFI.setNumLocals(CurLocal - MFI.getParams().size());
435   for (unsigned I = 0, E = MRI.getNumVirtRegs(); I < E; ++I) {
436     Register Reg = Register::index2VirtReg(I);
437     auto RL = Reg2Local.find(Reg);
438     if (RL == Reg2Local.end() || RL->second < MFI.getParams().size())
439       continue;
440 
441     MFI.setLocal(RL->second - MFI.getParams().size(),
442                  typeForRegClass(MRI.getRegClass(Reg)));
443     Changed = true;
444   }
445 
446 #ifndef NDEBUG
447   // Assert that all registers have been stackified at this point.
448   for (const MachineBasicBlock &MBB : MF) {
449     for (const MachineInstr &MI : MBB) {
450       if (MI.isDebugInstr() || MI.isLabel())
451         continue;
452       for (const MachineOperand &MO : MI.explicit_operands()) {
453         assert(
454             (!MO.isReg() || MRI.use_empty(MO.getReg()) ||
455              MFI.isVRegStackified(MO.getReg())) &&
456             "WebAssemblyExplicitLocals failed to stackify a register operand");
457       }
458     }
459   }
460 #endif
461 
462   return Changed;
463 }
464