1 //===-- WebAssemblyExplicitLocals.cpp - Make Locals Explicit --------------===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 ///
10 /// \file
11 /// This file converts any remaining registers into WebAssembly locals.
12 ///
13 /// After register stackification and register coloring, convert non-stackified
14 /// registers into locals, inserting explicit get_local and set_local
15 /// instructions.
16 ///
17 //===----------------------------------------------------------------------===//
18 
19 #include "MCTargetDesc/WebAssemblyMCTargetDesc.h"
20 #include "WebAssembly.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::EXCEPT_REFRegClass)
95     return WebAssembly::DROP_EXCEPT_REF;
96   llvm_unreachable("Unexpected register class");
97 }
98 
99 /// Get the appropriate get_local opcode for the given register class.
100 static unsigned getGetLocalOpcode(const TargetRegisterClass *RC) {
101   if (RC == &WebAssembly::I32RegClass)
102     return WebAssembly::GET_LOCAL_I32;
103   if (RC == &WebAssembly::I64RegClass)
104     return WebAssembly::GET_LOCAL_I64;
105   if (RC == &WebAssembly::F32RegClass)
106     return WebAssembly::GET_LOCAL_F32;
107   if (RC == &WebAssembly::F64RegClass)
108     return WebAssembly::GET_LOCAL_F64;
109   if (RC == &WebAssembly::V128RegClass)
110     return WebAssembly::GET_LOCAL_V128;
111   if (RC == &WebAssembly::EXCEPT_REFRegClass)
112     return WebAssembly::GET_LOCAL_EXCEPT_REF;
113   llvm_unreachable("Unexpected register class");
114 }
115 
116 /// Get the appropriate set_local opcode for the given register class.
117 static unsigned getSetLocalOpcode(const TargetRegisterClass *RC) {
118   if (RC == &WebAssembly::I32RegClass)
119     return WebAssembly::SET_LOCAL_I32;
120   if (RC == &WebAssembly::I64RegClass)
121     return WebAssembly::SET_LOCAL_I64;
122   if (RC == &WebAssembly::F32RegClass)
123     return WebAssembly::SET_LOCAL_F32;
124   if (RC == &WebAssembly::F64RegClass)
125     return WebAssembly::SET_LOCAL_F64;
126   if (RC == &WebAssembly::V128RegClass)
127     return WebAssembly::SET_LOCAL_V128;
128   if (RC == &WebAssembly::EXCEPT_REFRegClass)
129     return WebAssembly::SET_LOCAL_EXCEPT_REF;
130   llvm_unreachable("Unexpected register class");
131 }
132 
133 /// Get the appropriate tee_local opcode for the given register class.
134 static unsigned getTeeLocalOpcode(const TargetRegisterClass *RC) {
135   if (RC == &WebAssembly::I32RegClass)
136     return WebAssembly::TEE_LOCAL_I32;
137   if (RC == &WebAssembly::I64RegClass)
138     return WebAssembly::TEE_LOCAL_I64;
139   if (RC == &WebAssembly::F32RegClass)
140     return WebAssembly::TEE_LOCAL_F32;
141   if (RC == &WebAssembly::F64RegClass)
142     return WebAssembly::TEE_LOCAL_F64;
143   if (RC == &WebAssembly::V128RegClass)
144     return WebAssembly::TEE_LOCAL_V128;
145   if (RC == &WebAssembly::EXCEPT_REFRegClass)
146     return WebAssembly::TEE_LOCAL_EXCEPT_REF;
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::EXCEPT_REFRegClass)
161     return MVT::ExceptRef;
162   llvm_unreachable("unrecognized register class");
163 }
164 
165 /// Given a MachineOperand of a stackified vreg, return the instruction at the
166 /// start of the expression tree.
167 static MachineInstr *findStartOfTree(MachineOperand &MO,
168                                      MachineRegisterInfo &MRI,
169                                      WebAssemblyFunctionInfo &MFI) {
170   unsigned Reg = MO.getReg();
171   assert(MFI.isVRegStackified(Reg));
172   MachineInstr *Def = MRI.getVRegDef(Reg);
173 
174   // Find the first stackified use and proceed from there.
175   for (MachineOperand &DefMO : Def->explicit_uses()) {
176     if (!DefMO.isReg())
177       continue;
178     return findStartOfTree(DefMO, MRI, MFI);
179   }
180 
181   // If there were no stackified uses, we've reached the start.
182   return Def;
183 }
184 
185 bool WebAssemblyExplicitLocals::runOnMachineFunction(MachineFunction &MF) {
186   LLVM_DEBUG(dbgs() << "********** Make Locals Explicit **********\n"
187                        "********** Function: "
188                     << MF.getName() << '\n');
189 
190   // Disable this pass if directed to do so.
191   if (WasmDisableExplicitLocals)
192     return false;
193 
194   bool Changed = false;
195   MachineRegisterInfo &MRI = MF.getRegInfo();
196   WebAssemblyFunctionInfo &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();
197   const auto *TII = MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
198 
199   // Map non-stackified virtual registers to their local ids.
200   DenseMap<unsigned, unsigned> Reg2Local;
201 
202   // Handle ARGUMENTS first to ensure that they get the designated numbers.
203   for (MachineBasicBlock::iterator I = MF.begin()->begin(),
204                                    E = MF.begin()->end();
205        I != E;) {
206     MachineInstr &MI = *I++;
207     if (!WebAssembly::isArgument(MI))
208       break;
209     unsigned Reg = MI.getOperand(0).getReg();
210     assert(!MFI.isVRegStackified(Reg));
211     Reg2Local[Reg] = static_cast<unsigned>(MI.getOperand(1).getImm());
212     MI.eraseFromParent();
213     Changed = true;
214   }
215 
216   // Start assigning local numbers after the last parameter.
217   unsigned CurLocal = static_cast<unsigned>(MFI.getParams().size());
218 
219   // Precompute the set of registers that are unused, so that we can insert
220   // drops to their defs.
221   BitVector UseEmpty(MRI.getNumVirtRegs());
222   for (unsigned I = 0, E = MRI.getNumVirtRegs(); I < E; ++I)
223     UseEmpty[I] = MRI.use_empty(TargetRegisterInfo::index2VirtReg(I));
224 
225   // Visit each instruction in the function.
226   for (MachineBasicBlock &MBB : MF) {
227     for (MachineBasicBlock::iterator I = MBB.begin(), E = MBB.end(); I != E;) {
228       MachineInstr &MI = *I++;
229       assert(!WebAssembly::isArgument(MI));
230 
231       if (MI.isDebugInstr() || MI.isLabel())
232         continue;
233 
234       // Replace tee instructions with tee_local. The difference is that tee
235       // instructins have two defs, while tee_local instructions have one def
236       // and an index of a local to write to.
237       if (WebAssembly::isTee(MI)) {
238         assert(MFI.isVRegStackified(MI.getOperand(0).getReg()));
239         assert(!MFI.isVRegStackified(MI.getOperand(1).getReg()));
240         unsigned OldReg = MI.getOperand(2).getReg();
241         const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
242 
243         // Stackify the input if it isn't stackified yet.
244         if (!MFI.isVRegStackified(OldReg)) {
245           unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
246           unsigned NewReg = MRI.createVirtualRegister(RC);
247           unsigned Opc = getGetLocalOpcode(RC);
248           BuildMI(MBB, &MI, MI.getDebugLoc(), TII->get(Opc), NewReg)
249               .addImm(LocalId);
250           MI.getOperand(2).setReg(NewReg);
251           MFI.stackifyVReg(NewReg);
252         }
253 
254         // Replace the TEE with a TEE_LOCAL.
255         unsigned LocalId =
256             getLocalId(Reg2Local, CurLocal, MI.getOperand(1).getReg());
257         unsigned Opc = getTeeLocalOpcode(RC);
258         BuildMI(MBB, &MI, MI.getDebugLoc(), TII->get(Opc),
259                 MI.getOperand(0).getReg())
260             .addImm(LocalId)
261             .addReg(MI.getOperand(2).getReg());
262 
263         MI.eraseFromParent();
264         Changed = true;
265         continue;
266       }
267 
268       // Insert set_locals for any defs that aren't stackified yet. Currently
269       // we handle at most one def.
270       assert(MI.getDesc().getNumDefs() <= 1);
271       if (MI.getDesc().getNumDefs() == 1) {
272         unsigned OldReg = MI.getOperand(0).getReg();
273         if (!MFI.isVRegStackified(OldReg)) {
274           const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
275           unsigned NewReg = MRI.createVirtualRegister(RC);
276           auto InsertPt = std::next(MachineBasicBlock::iterator(&MI));
277           if (MI.getOpcode() == WebAssembly::IMPLICIT_DEF) {
278             MI.eraseFromParent();
279             Changed = true;
280             continue;
281           }
282           if (UseEmpty[TargetRegisterInfo::virtReg2Index(OldReg)]) {
283             unsigned Opc = getDropOpcode(RC);
284             MachineInstr *Drop =
285                 BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc))
286                     .addReg(NewReg);
287             // After the drop instruction, this reg operand will not be used
288             Drop->getOperand(0).setIsKill();
289           } else {
290             unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
291             unsigned Opc = getSetLocalOpcode(RC);
292             BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc))
293                 .addImm(LocalId)
294                 .addReg(NewReg);
295           }
296           MI.getOperand(0).setReg(NewReg);
297           // This register operand is now being used by the inserted drop
298           // instruction, so make it undead.
299           MI.getOperand(0).setIsDead(false);
300           MFI.stackifyVReg(NewReg);
301           Changed = true;
302         }
303       }
304 
305       // Insert get_locals for any uses that aren't stackified yet.
306       MachineInstr *InsertPt = &MI;
307       for (MachineOperand &MO : reverse(MI.explicit_uses())) {
308         if (!MO.isReg())
309           continue;
310 
311         unsigned OldReg = MO.getReg();
312 
313         // Inline asm may have a def in the middle of the operands. Our contract
314         // with inline asm register operands is to provide local indices as
315         // immediates.
316         if (MO.isDef()) {
317           assert(MI.getOpcode() == TargetOpcode::INLINEASM);
318           unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
319           MRI.removeRegOperandFromUseList(&MO);
320           MO = MachineOperand::CreateImm(LocalId);
321           continue;
322         }
323 
324         // If we see a stackified register, prepare to insert subsequent
325         // get_locals before the start of its tree.
326         if (MFI.isVRegStackified(OldReg)) {
327           InsertPt = findStartOfTree(MO, MRI, MFI);
328           continue;
329         }
330 
331         // Our contract with inline asm register operands is to provide local
332         // indices as immediates.
333         if (MI.getOpcode() == TargetOpcode::INLINEASM) {
334           unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
335           MRI.removeRegOperandFromUseList(&MO);
336           MO = MachineOperand::CreateImm(LocalId);
337           continue;
338         }
339 
340         // Insert a get_local.
341         unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
342         const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
343         unsigned NewReg = MRI.createVirtualRegister(RC);
344         unsigned Opc = getGetLocalOpcode(RC);
345         InsertPt =
346             BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc), NewReg)
347                 .addImm(LocalId);
348         MO.setReg(NewReg);
349         MFI.stackifyVReg(NewReg);
350         Changed = true;
351       }
352 
353       // Coalesce and eliminate COPY instructions.
354       if (WebAssembly::isCopy(MI)) {
355         MRI.replaceRegWith(MI.getOperand(1).getReg(),
356                            MI.getOperand(0).getReg());
357         MI.eraseFromParent();
358         Changed = true;
359       }
360     }
361   }
362 
363   // Define the locals.
364   // TODO: Sort the locals for better compression.
365   MFI.setNumLocals(CurLocal - MFI.getParams().size());
366   for (unsigned I = 0, E = MRI.getNumVirtRegs(); I < E; ++I) {
367     unsigned Reg = TargetRegisterInfo::index2VirtReg(I);
368     auto RL = Reg2Local.find(Reg);
369     if (RL == Reg2Local.end() || RL->second < MFI.getParams().size())
370       continue;
371 
372     MFI.setLocal(RL->second - MFI.getParams().size(),
373                  typeForRegClass(MRI.getRegClass(Reg)));
374     Changed = true;
375   }
376 
377 #ifndef NDEBUG
378   // Assert that all registers have been stackified at this point.
379   for (const MachineBasicBlock &MBB : MF) {
380     for (const MachineInstr &MI : MBB) {
381       if (MI.isDebugInstr() || MI.isLabel())
382         continue;
383       for (const MachineOperand &MO : MI.explicit_operands()) {
384         assert(
385             (!MO.isReg() || MRI.use_empty(MO.getReg()) ||
386              MFI.isVRegStackified(MO.getReg())) &&
387             "WebAssemblyExplicitLocals failed to stackify a register operand");
388       }
389     }
390   }
391 #endif
392 
393   return Changed;
394 }
395