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