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