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