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 namespace {
35 class WebAssemblyExplicitLocals final : public MachineFunctionPass {
36   StringRef getPassName() const override {
37     return "WebAssembly Explicit Locals";
38   }
39 
40   void getAnalysisUsage(AnalysisUsage &AU) const override {
41     AU.setPreservesCFG();
42     AU.addPreserved<MachineBlockFrequencyInfo>();
43     MachineFunctionPass::getAnalysisUsage(AU);
44   }
45 
46   bool runOnMachineFunction(MachineFunction &MF) override;
47 
48 public:
49   static char ID; // Pass identification, replacement for typeid
50   WebAssemblyExplicitLocals() : MachineFunctionPass(ID) {}
51 };
52 } // end anonymous namespace
53 
54 char WebAssemblyExplicitLocals::ID = 0;
55 FunctionPass *llvm::createWebAssemblyExplicitLocals() {
56   return new WebAssemblyExplicitLocals();
57 }
58 
59 /// Return a local id number for the given register, assigning it a new one
60 /// if it doesn't yet have one.
61 static unsigned getLocalId(DenseMap<unsigned, unsigned> &Reg2Local,
62                            unsigned &CurLocal, unsigned Reg) {
63   return Reg2Local.insert(std::make_pair(Reg, CurLocal++)).first->second;
64 }
65 
66 /// Get the appropriate get_local opcode for the given register class.
67 static unsigned getGetLocalOpcode(const TargetRegisterClass *RC) {
68   if (RC == &WebAssembly::I32RegClass)
69     return WebAssembly::GET_LOCAL_I32;
70   if (RC == &WebAssembly::I64RegClass)
71     return WebAssembly::GET_LOCAL_I64;
72   if (RC == &WebAssembly::F32RegClass)
73     return WebAssembly::GET_LOCAL_F32;
74   if (RC == &WebAssembly::F64RegClass)
75     return WebAssembly::GET_LOCAL_F64;
76   if (RC == &WebAssembly::V128RegClass)
77     return WebAssembly::GET_LOCAL_V128;
78   llvm_unreachable("Unexpected register class");
79 }
80 
81 /// Get the appropriate set_local opcode for the given register class.
82 static unsigned getSetLocalOpcode(const TargetRegisterClass *RC) {
83   if (RC == &WebAssembly::I32RegClass)
84     return WebAssembly::SET_LOCAL_I32;
85   if (RC == &WebAssembly::I64RegClass)
86     return WebAssembly::SET_LOCAL_I64;
87   if (RC == &WebAssembly::F32RegClass)
88     return WebAssembly::SET_LOCAL_F32;
89   if (RC == &WebAssembly::F64RegClass)
90     return WebAssembly::SET_LOCAL_F64;
91   if (RC == &WebAssembly::V128RegClass)
92     return WebAssembly::SET_LOCAL_V128;
93   llvm_unreachable("Unexpected register class");
94 }
95 
96 /// Get the appropriate tee_local opcode for the given register class.
97 static unsigned getTeeLocalOpcode(const TargetRegisterClass *RC) {
98   if (RC == &WebAssembly::I32RegClass)
99     return WebAssembly::TEE_LOCAL_I32;
100   if (RC == &WebAssembly::I64RegClass)
101     return WebAssembly::TEE_LOCAL_I64;
102   if (RC == &WebAssembly::F32RegClass)
103     return WebAssembly::TEE_LOCAL_F32;
104   if (RC == &WebAssembly::F64RegClass)
105     return WebAssembly::TEE_LOCAL_F64;
106   if (RC == &WebAssembly::V128RegClass)
107     return WebAssembly::TEE_LOCAL_V128;
108   llvm_unreachable("Unexpected register class");
109 }
110 
111 /// Get the type associated with the given register class.
112 static MVT typeForRegClass(const TargetRegisterClass *RC) {
113   if (RC == &WebAssembly::I32RegClass)
114     return MVT::i32;
115   if (RC == &WebAssembly::I64RegClass)
116     return MVT::i64;
117   if (RC == &WebAssembly::F32RegClass)
118     return MVT::f32;
119   if (RC == &WebAssembly::F64RegClass)
120     return MVT::f64;
121   llvm_unreachable("unrecognized register class");
122 }
123 
124 /// Given a MachineOperand of a stackified vreg, return the instruction at the
125 /// start of the expression tree.
126 static MachineInstr *FindStartOfTree(MachineOperand &MO,
127                                      MachineRegisterInfo &MRI,
128                                      WebAssemblyFunctionInfo &MFI) {
129   unsigned Reg = MO.getReg();
130   assert(MFI.isVRegStackified(Reg));
131   MachineInstr *Def = MRI.getVRegDef(Reg);
132 
133   // Find the first stackified use and proceed from there.
134   for (MachineOperand &DefMO : Def->explicit_uses()) {
135     if (!DefMO.isReg())
136       continue;
137     return FindStartOfTree(DefMO, MRI, MFI);
138   }
139 
140   // If there were no stackified uses, we've reached the start.
141   return Def;
142 }
143 
144 bool WebAssemblyExplicitLocals::runOnMachineFunction(MachineFunction &MF) {
145   DEBUG(dbgs() << "********** Make Locals Explicit **********\n"
146                   "********** Function: "
147                << MF.getName() << '\n');
148 
149   // Disable this pass if we aren't doing direct wasm object emission.
150   if (MF.getSubtarget<WebAssemblySubtarget>()
151         .getTargetTriple().isOSBinFormatELF())
152     return false;
153 
154   bool Changed = false;
155   MachineRegisterInfo &MRI = MF.getRegInfo();
156   WebAssemblyFunctionInfo &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();
157   const auto *TII = MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
158 
159   // Map non-stackified virtual registers to their local ids.
160   DenseMap<unsigned, unsigned> Reg2Local;
161 
162   // Handle ARGUMENTS first to ensure that they get the designated numbers.
163   for (MachineBasicBlock::iterator I = MF.begin()->begin(),
164                                    E = MF.begin()->end();
165        I != E;) {
166     MachineInstr &MI = *I++;
167     if (!WebAssembly::isArgument(MI))
168       break;
169     unsigned Reg = MI.getOperand(0).getReg();
170     assert(!MFI.isVRegStackified(Reg));
171     Reg2Local[Reg] = MI.getOperand(1).getImm();
172     MI.eraseFromParent();
173     Changed = true;
174   }
175 
176   // Start assigning local numbers after the last parameter.
177   unsigned CurLocal = MFI.getParams().size();
178 
179   // Visit each instruction in the function.
180   for (MachineBasicBlock &MBB : MF) {
181     for (MachineBasicBlock::iterator I = MBB.begin(), E = MBB.end(); I != E;) {
182       MachineInstr &MI = *I++;
183       assert(!WebAssembly::isArgument(MI));
184 
185       if (MI.isDebugValue() || MI.isLabel())
186         continue;
187 
188       // Replace tee instructions with tee_local. The difference is that tee
189       // instructins have two defs, while tee_local instructions have one def
190       // and an index of a local to write to.
191       if (WebAssembly::isTee(MI)) {
192         assert(MFI.isVRegStackified(MI.getOperand(0).getReg()));
193         assert(!MFI.isVRegStackified(MI.getOperand(1).getReg()));
194         unsigned OldReg = MI.getOperand(2).getReg();
195         const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
196 
197         // Stackify the input if it isn't stackified yet.
198         if (!MFI.isVRegStackified(OldReg)) {
199           unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
200           unsigned NewReg = MRI.createVirtualRegister(RC);
201           unsigned Opc = getGetLocalOpcode(RC);
202           BuildMI(MBB, &MI, MI.getDebugLoc(), TII->get(Opc), NewReg)
203               .addImm(LocalId);
204           MI.getOperand(2).setReg(NewReg);
205           MFI.stackifyVReg(NewReg);
206         }
207 
208         // Replace the TEE with a TEE_LOCAL.
209         unsigned LocalId =
210             getLocalId(Reg2Local, CurLocal, MI.getOperand(1).getReg());
211         unsigned Opc = getTeeLocalOpcode(RC);
212         BuildMI(MBB, &MI, MI.getDebugLoc(), TII->get(Opc),
213                 MI.getOperand(0).getReg())
214             .addImm(LocalId)
215             .addReg(MI.getOperand(2).getReg());
216 
217         MI.eraseFromParent();
218         Changed = true;
219         continue;
220       }
221 
222       // Insert set_locals for any defs that aren't stackified yet. Currently
223       // we handle at most one def.
224       assert(MI.getDesc().getNumDefs() <= 1);
225       if (MI.getDesc().getNumDefs() == 1) {
226         unsigned OldReg = MI.getOperand(0).getReg();
227         if (!MFI.isVRegStackified(OldReg) && !MRI.use_empty(OldReg)) {
228           unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
229           const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
230           unsigned NewReg = MRI.createVirtualRegister(RC);
231           auto InsertPt = std::next(MachineBasicBlock::iterator(&MI));
232           unsigned Opc = getSetLocalOpcode(RC);
233           BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc))
234               .addImm(LocalId)
235               .addReg(NewReg);
236           MI.getOperand(0).setReg(NewReg);
237           MFI.stackifyVReg(NewReg);
238           Changed = true;
239         }
240       }
241 
242       // Insert get_locals for any uses that aren't stackified yet.
243       MachineInstr *InsertPt = &MI;
244       for (MachineOperand &MO : reverse(MI.explicit_uses())) {
245         if (!MO.isReg())
246           continue;
247 
248         unsigned OldReg = MO.getReg();
249 
250         // If we see a stackified register, prepare to insert subsequent
251         // get_locals before the start of its tree.
252         if (MFI.isVRegStackified(OldReg)) {
253           InsertPt = FindStartOfTree(MO, MRI, MFI);
254           continue;
255         }
256 
257         // Insert a get_local.
258         unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
259         const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
260         unsigned NewReg = MRI.createVirtualRegister(RC);
261         unsigned Opc = getGetLocalOpcode(RC);
262         InsertPt =
263             BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc), NewReg)
264                 .addImm(LocalId);
265         MO.setReg(NewReg);
266         MFI.stackifyVReg(NewReg);
267         Changed = true;
268       }
269 
270       // Coalesce and eliminate COPY instructions.
271       if (WebAssembly::isCopy(MI)) {
272         MRI.replaceRegWith(MI.getOperand(1).getReg(),
273                            MI.getOperand(0).getReg());
274         MI.eraseFromParent();
275         Changed = true;
276       }
277     }
278   }
279 
280   // Define the locals.
281   for (size_t i = 0, e = MRI.getNumVirtRegs(); i < e; ++i) {
282     unsigned Reg = TargetRegisterInfo::index2VirtReg(i);
283     auto I = Reg2Local.find(Reg);
284     if (I == Reg2Local.end() || I->second < MFI.getParams().size())
285       continue;
286 
287     MFI.addLocal(typeForRegClass(MRI.getRegClass(Reg)));
288     Changed = true;
289   }
290 
291 #ifndef NDEBUG
292   // Assert that all registers have been stackified at this point.
293   for (const MachineBasicBlock &MBB : MF) {
294     for (const MachineInstr &MI : MBB) {
295       if (MI.isDebugValue() || MI.isLabel())
296         continue;
297       for (const MachineOperand &MO : MI.explicit_operands()) {
298         assert(
299             (!MO.isReg() || MRI.use_empty(MO.getReg()) ||
300              MFI.isVRegStackified(MO.getReg())) &&
301             "WebAssemblyExplicitLocals failed to stackify a register operand");
302       }
303     }
304   }
305 #endif
306 
307   return Changed;
308 }
309