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