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