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