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   bool Changed = false;
150   MachineRegisterInfo &MRI = MF.getRegInfo();
151   WebAssemblyFunctionInfo &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();
152   const auto *TII = MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
153 
154   // Map non-stackified virtual registers to their local ids.
155   DenseMap<unsigned, unsigned> Reg2Local;
156 
157   // Handle ARGUMENTS first to ensure that they get the designated numbers.
158   for (MachineBasicBlock::iterator I = MF.begin()->begin(),
159                                    E = MF.begin()->end();
160        I != E;) {
161     MachineInstr &MI = *I++;
162     if (!WebAssembly::isArgument(MI))
163       break;
164     unsigned Reg = MI.getOperand(0).getReg();
165     assert(!MFI.isVRegStackified(Reg));
166     Reg2Local[Reg] = MI.getOperand(1).getImm();
167     MI.eraseFromParent();
168     Changed = true;
169   }
170 
171   // Start assigning local numbers after the last parameter.
172   unsigned CurLocal = MFI.getParams().size();
173 
174   // Visit each instruction in the function.
175   for (MachineBasicBlock &MBB : MF) {
176     for (MachineBasicBlock::iterator I = MBB.begin(), E = MBB.end(); I != E;) {
177       MachineInstr &MI = *I++;
178       assert(!WebAssembly::isArgument(MI));
179 
180       if (MI.isDebugValue() || MI.isLabel())
181         continue;
182 
183       // Replace tee instructions with tee_local. The difference is that tee
184       // instructins have two defs, while tee_local instructions have one def
185       // and an index of a local to write to.
186       if (WebAssembly::isTee(MI)) {
187         assert(MFI.isVRegStackified(MI.getOperand(0).getReg()));
188         assert(!MFI.isVRegStackified(MI.getOperand(1).getReg()));
189         unsigned OldReg = MI.getOperand(2).getReg();
190         const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
191 
192         // Stackify the input if it isn't stackified yet.
193         if (!MFI.isVRegStackified(OldReg)) {
194           unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
195           unsigned NewReg = MRI.createVirtualRegister(RC);
196           unsigned Opc = getGetLocalOpcode(RC);
197           BuildMI(MBB, &MI, MI.getDebugLoc(), TII->get(Opc), NewReg)
198               .addImm(LocalId);
199           MI.getOperand(2).setReg(NewReg);
200           MFI.stackifyVReg(NewReg);
201         }
202 
203         // Replace the TEE with a TEE_LOCAL.
204         unsigned LocalId =
205             getLocalId(Reg2Local, CurLocal, MI.getOperand(1).getReg());
206         unsigned Opc = getTeeLocalOpcode(RC);
207         BuildMI(MBB, &MI, MI.getDebugLoc(), TII->get(Opc),
208                 MI.getOperand(0).getReg())
209             .addImm(LocalId)
210             .addReg(MI.getOperand(2).getReg());
211 
212         MI.eraseFromParent();
213         Changed = true;
214         continue;
215       }
216 
217       // Insert set_locals for any defs that aren't stackified yet. Currently
218       // we handle at most one def.
219       assert(MI.getDesc().getNumDefs() <= 1);
220       if (MI.getDesc().getNumDefs() == 1) {
221         unsigned OldReg = MI.getOperand(0).getReg();
222         if (!MFI.isVRegStackified(OldReg) && !MRI.use_empty(OldReg)) {
223           unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
224           const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
225           unsigned NewReg = MRI.createVirtualRegister(RC);
226           auto InsertPt = std::next(MachineBasicBlock::iterator(&MI));
227           unsigned Opc = getSetLocalOpcode(RC);
228           BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc))
229               .addImm(LocalId)
230               .addReg(NewReg);
231           MI.getOperand(0).setReg(NewReg);
232           MFI.stackifyVReg(NewReg);
233           Changed = true;
234         }
235       }
236 
237       // Insert get_locals for any uses that aren't stackified yet.
238       MachineInstr *InsertPt = &MI;
239       for (MachineOperand &MO : reverse(MI.explicit_uses())) {
240         if (!MO.isReg())
241           continue;
242 
243         unsigned OldReg = MO.getReg();
244 
245         // If we see a stackified register, prepare to insert subsequent
246         // get_locals before the start of its tree.
247         if (MFI.isVRegStackified(OldReg)) {
248           InsertPt = FindStartOfTree(MO, MRI, MFI);
249           continue;
250         }
251 
252         // Insert a get_local.
253         unsigned LocalId = getLocalId(Reg2Local, CurLocal, OldReg);
254         const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
255         unsigned NewReg = MRI.createVirtualRegister(RC);
256         unsigned Opc = getGetLocalOpcode(RC);
257         InsertPt =
258             BuildMI(MBB, InsertPt, MI.getDebugLoc(), TII->get(Opc), NewReg)
259                 .addImm(LocalId);
260         MO.setReg(NewReg);
261         MFI.stackifyVReg(NewReg);
262         Changed = true;
263       }
264 
265       // Coalesce and eliminate COPY instructions.
266       if (WebAssembly::isCopy(MI)) {
267         MRI.replaceRegWith(MI.getOperand(1).getReg(),
268                            MI.getOperand(0).getReg());
269         MI.eraseFromParent();
270         Changed = true;
271       }
272     }
273   }
274 
275   // Define the locals.
276   for (size_t i = 0, e = MRI.getNumVirtRegs(); i < e; ++i) {
277     unsigned Reg = TargetRegisterInfo::index2VirtReg(i);
278     auto I = Reg2Local.find(Reg);
279     if (I == Reg2Local.end() || I->second < MFI.getParams().size())
280       continue;
281 
282     MFI.addLocal(typeForRegClass(MRI.getRegClass(Reg)));
283     Changed = true;
284   }
285 
286 #ifndef NDEBUG
287   // Assert that all registers have been stackified at this point.
288   for (const MachineBasicBlock &MBB : MF) {
289     for (const MachineInstr &MI : MBB) {
290       if (MI.isDebugValue() || MI.isLabel())
291         continue;
292       for (const MachineOperand &MO : MI.explicit_operands()) {
293         assert(
294             (!MO.isReg() || MRI.use_empty(MO.getReg()) ||
295              MFI.isVRegStackified(MO.getReg())) &&
296             "WebAssemblyExplicitLocals failed to stackify a register operand");
297       }
298     }
299   }
300 #endif
301 
302   return Changed;
303 }
304