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 /// \brief 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. Note that this produces output
35 // which is not valid WebAssembly, though it may be more convenient for writing
36 // LLVM unit tests with.
37 static cl::opt<bool> DisableWebAssemblyExplicitLocals(
38     "disable-wasm-explicit-locals", cl::ReallyHidden,
39     cl::desc("WebAssembly: Disable emission of get_local/set_local."),
40     cl::init(false));
41 
42 namespace {
43 class WebAssemblyExplicitLocals final : public MachineFunctionPass {
44   StringRef getPassName() const override {
45     return "WebAssembly Explicit Locals";
46   }
47 
48   void getAnalysisUsage(AnalysisUsage &AU) const override {
49     AU.setPreservesCFG();
50     AU.addPreserved<MachineBlockFrequencyInfo>();
51     MachineFunctionPass::getAnalysisUsage(AU);
52   }
53 
54   bool runOnMachineFunction(MachineFunction &MF) override;
55 
56 public:
57   static char ID; // Pass identification, replacement for typeid
58   WebAssemblyExplicitLocals() : MachineFunctionPass(ID) {}
59 };
60 } // end anonymous namespace
61 
62 char WebAssemblyExplicitLocals::ID = 0;
63 FunctionPass *llvm::createWebAssemblyExplicitLocals() {
64   return new WebAssemblyExplicitLocals();
65 }
66 
67 /// Return a local id number for the given register, assigning it a new one
68 /// if it doesn't yet have one.
69 static unsigned getLocalId(DenseMap<unsigned, unsigned> &Reg2Local,
70                            unsigned &CurLocal, unsigned Reg) {
71   auto P = Reg2Local.insert(std::make_pair(Reg, CurLocal));
72   if (P.second)
73     ++CurLocal;
74   return P.first->second;
75 }
76 
77 /// Get the appropriate drop opcode for the given register class.
78 static unsigned getDropOpcode(const TargetRegisterClass *RC) {
79   if (RC == &WebAssembly::I32RegClass)
80     return WebAssembly::DROP_I32;
81   if (RC == &WebAssembly::I64RegClass)
82     return WebAssembly::DROP_I64;
83   if (RC == &WebAssembly::F32RegClass)
84     return WebAssembly::DROP_F32;
85   if (RC == &WebAssembly::F64RegClass)
86     return WebAssembly::DROP_F64;
87   if (RC == &WebAssembly::V128RegClass)
88     return WebAssembly::DROP_V128;
89   if (RC == &WebAssembly::EXCEPT_REFRegClass)
90     return WebAssembly::DROP_EXCEPT_REF;
91   llvm_unreachable("Unexpected register class");
92 }
93 
94 /// Get the appropriate get_local opcode for the given register class.
95 static unsigned getGetLocalOpcode(const TargetRegisterClass *RC) {
96   if (RC == &WebAssembly::I32RegClass)
97     return WebAssembly::GET_LOCAL_I32;
98   if (RC == &WebAssembly::I64RegClass)
99     return WebAssembly::GET_LOCAL_I64;
100   if (RC == &WebAssembly::F32RegClass)
101     return WebAssembly::GET_LOCAL_F32;
102   if (RC == &WebAssembly::F64RegClass)
103     return WebAssembly::GET_LOCAL_F64;
104   if (RC == &WebAssembly::V128RegClass)
105     return WebAssembly::GET_LOCAL_V128;
106   if (RC == &WebAssembly::EXCEPT_REFRegClass)
107     return WebAssembly::GET_LOCAL_EXCEPT_REF;
108   llvm_unreachable("Unexpected register class");
109 }
110 
111 /// Get the appropriate set_local opcode for the given register class.
112 static unsigned getSetLocalOpcode(const TargetRegisterClass *RC) {
113   if (RC == &WebAssembly::I32RegClass)
114     return WebAssembly::SET_LOCAL_I32;
115   if (RC == &WebAssembly::I64RegClass)
116     return WebAssembly::SET_LOCAL_I64;
117   if (RC == &WebAssembly::F32RegClass)
118     return WebAssembly::SET_LOCAL_F32;
119   if (RC == &WebAssembly::F64RegClass)
120     return WebAssembly::SET_LOCAL_F64;
121   if (RC == &WebAssembly::V128RegClass)
122     return WebAssembly::SET_LOCAL_V128;
123   if (RC == &WebAssembly::EXCEPT_REFRegClass)
124     return WebAssembly::SET_LOCAL_EXCEPT_REF;
125   llvm_unreachable("Unexpected register class");
126 }
127 
128 /// Get the appropriate tee_local opcode for the given register class.
129 static unsigned getTeeLocalOpcode(const TargetRegisterClass *RC) {
130   if (RC == &WebAssembly::I32RegClass)
131     return WebAssembly::TEE_LOCAL_I32;
132   if (RC == &WebAssembly::I64RegClass)
133     return WebAssembly::TEE_LOCAL_I64;
134   if (RC == &WebAssembly::F32RegClass)
135     return WebAssembly::TEE_LOCAL_F32;
136   if (RC == &WebAssembly::F64RegClass)
137     return WebAssembly::TEE_LOCAL_F64;
138   if (RC == &WebAssembly::V128RegClass)
139     return WebAssembly::TEE_LOCAL_V128;
140   if (RC == &WebAssembly::EXCEPT_REFRegClass)
141     return WebAssembly::TEE_LOCAL_EXCEPT_REF;
142   llvm_unreachable("Unexpected register class");
143 }
144 
145 /// Get the type associated with the given register class.
146 static MVT typeForRegClass(const TargetRegisterClass *RC) {
147   if (RC == &WebAssembly::I32RegClass)
148     return MVT::i32;
149   if (RC == &WebAssembly::I64RegClass)
150     return MVT::i64;
151   if (RC == &WebAssembly::F32RegClass)
152     return MVT::f32;
153   if (RC == &WebAssembly::F64RegClass)
154     return MVT::f64;
155   if (RC == &WebAssembly::EXCEPT_REFRegClass)
156     return MVT::ExceptRef;
157   llvm_unreachable("unrecognized register class");
158 }
159 
160 /// Given a MachineOperand of a stackified vreg, return the instruction at the
161 /// start of the expression tree.
162 static MachineInstr *FindStartOfTree(MachineOperand &MO,
163                                      MachineRegisterInfo &MRI,
164                                      WebAssemblyFunctionInfo &MFI) {
165   unsigned Reg = MO.getReg();
166   assert(MFI.isVRegStackified(Reg));
167   MachineInstr *Def = MRI.getVRegDef(Reg);
168 
169   // Find the first stackified use and proceed from there.
170   for (MachineOperand &DefMO : Def->explicit_uses()) {
171     if (!DefMO.isReg())
172       continue;
173     return FindStartOfTree(DefMO, MRI, MFI);
174   }
175 
176   // If there were no stackified uses, we've reached the start.
177   return Def;
178 }
179 
180 bool WebAssemblyExplicitLocals::runOnMachineFunction(MachineFunction &MF) {
181   DEBUG(dbgs() << "********** Make Locals Explicit **********\n"
182                   "********** Function: "
183                << MF.getName() << '\n');
184 
185   // Disable this pass if directed to do so.
186   if (DisableWebAssemblyExplicitLocals)
187     return false;
188 
189   // Disable this pass if we aren't doing direct wasm object emission.
190   if (MF.getSubtarget<WebAssemblySubtarget>()
191         .getTargetTriple().isOSBinFormatELF())
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] = MI.getOperand(1).getImm();
212     MI.eraseFromParent();
213     Changed = true;
214   }
215 
216   // Start assigning local numbers after the last parameter.
217   unsigned CurLocal = 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.isDebugValue() || 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             BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc))
285                 .addReg(NewReg);
286           } else {
287             unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
288             unsigned Opc = getSetLocalOpcode(RC);
289             BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc))
290                 .addImm(LocalId)
291                 .addReg(NewReg);
292           }
293           MI.getOperand(0).setReg(NewReg);
294           MFI.stackifyVReg(NewReg);
295           Changed = true;
296         }
297       }
298 
299       // Insert get_locals for any uses that aren't stackified yet.
300       MachineInstr *InsertPt = &MI;
301       for (MachineOperand &MO : reverse(MI.explicit_uses())) {
302         if (!MO.isReg())
303           continue;
304 
305         unsigned OldReg = MO.getReg();
306 
307         // Inline asm may have a def in the middle of the operands. Our contract
308         // with inline asm register operands is to provide local indices as
309         // immediates.
310         if (MO.isDef()) {
311           assert(MI.getOpcode() == TargetOpcode::INLINEASM);
312           unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
313           MRI.removeRegOperandFromUseList(&MO);
314           MO = MachineOperand::CreateImm(LocalId);
315           continue;
316         }
317 
318         // If we see a stackified register, prepare to insert subsequent
319         // get_locals before the start of its tree.
320         if (MFI.isVRegStackified(OldReg)) {
321           InsertPt = FindStartOfTree(MO, MRI, MFI);
322           continue;
323         }
324 
325         // Our contract with inline asm register operands is to provide local
326         // indices as immediates.
327         if (MI.getOpcode() == TargetOpcode::INLINEASM) {
328           unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
329           MRI.removeRegOperandFromUseList(&MO);
330           MO = MachineOperand::CreateImm(LocalId);
331           continue;
332         }
333 
334         // Insert a get_local.
335         unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
336         const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
337         unsigned NewReg = MRI.createVirtualRegister(RC);
338         unsigned Opc = getGetLocalOpcode(RC);
339         InsertPt =
340             BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc), NewReg)
341                 .addImm(LocalId);
342         MO.setReg(NewReg);
343         MFI.stackifyVReg(NewReg);
344         Changed = true;
345       }
346 
347       // Coalesce and eliminate COPY instructions.
348       if (WebAssembly::isCopy(MI)) {
349         MRI.replaceRegWith(MI.getOperand(1).getReg(),
350                            MI.getOperand(0).getReg());
351         MI.eraseFromParent();
352         Changed = true;
353       }
354     }
355   }
356 
357   // Define the locals.
358   // TODO: Sort the locals for better compression.
359   MFI.setNumLocals(CurLocal - MFI.getParams().size());
360   for (size_t i = 0, e = MRI.getNumVirtRegs(); i < e; ++i) {
361     unsigned Reg = TargetRegisterInfo::index2VirtReg(i);
362     auto I = Reg2Local.find(Reg);
363     if (I == Reg2Local.end() || I->second < MFI.getParams().size())
364       continue;
365 
366     MFI.setLocal(I->second - MFI.getParams().size(),
367                  typeForRegClass(MRI.getRegClass(Reg)));
368     Changed = true;
369   }
370 
371 #ifndef NDEBUG
372   // Assert that all registers have been stackified at this point.
373   for (const MachineBasicBlock &MBB : MF) {
374     for (const MachineInstr &MI : MBB) {
375       if (MI.isDebugValue() || MI.isLabel())
376         continue;
377       for (const MachineOperand &MO : MI.explicit_operands()) {
378         assert(
379             (!MO.isReg() || MRI.use_empty(MO.getReg()) ||
380              MFI.isVRegStackified(MO.getReg())) &&
381             "WebAssemblyExplicitLocals failed to stackify a register operand");
382       }
383     }
384   }
385 #endif
386 
387   return Changed;
388 }
389