Some simpliciations to the spill/reload interface
[oota-llvm.git] / lib / CodeGen / RegAllocSimple.cpp
1 //===-- RegAllocSimple.cpp - A simple generic register allocator --- ------===//
2 //
3 // This file implements a simple register allocator. *Very* simple.
4 //
5 //===----------------------------------------------------------------------===//
6
7 #include "llvm/CodeGen/MachineFunction.h"
8 #include "llvm/CodeGen/MachineInstr.h"
9 #include "llvm/Target/MachineInstrInfo.h"
10 #include "llvm/Target/TargetMachine.h"
11 #include "Support/Statistic.h"
12 #include <iostream>
13 #include <set>
14
15 #if 0
16 /// PhysRegClassMap - Construct a mapping of physical register numbers to their
17 /// register classes.
18 ///
19 /// NOTE: This class will eventually be pulled out to somewhere shared.
20 ///
21 class PhysRegClassMap {
22   std::map<unsigned, const TargetRegisterClass*> PhysReg2RegClassMap;
23 public:
24   PhysRegClassMap(const MRegisterInfo *RI) {
25     for (MRegisterInfo::const_iterator I = RI->regclass_begin(),
26            E = RI->regclass_end(); I != E; ++I)
27       for (unsigned i=0; i < (*I)->getNumRegs(); ++i)
28         PhysReg2RegClassMap[(*I)->getRegister(i)] = *I;
29   }
30
31   const TargetRegisterClass *operator[](unsigned Reg) {
32     assert(PhysReg2RegClassMap[Reg] && "Register is not a known physreg!");
33     return PhysReg2RegClassMap[Reg];
34   }
35
36   const TargetRegisterClass *get(unsigned Reg) { return operator[](Reg); }
37 };
38 #endif
39
40
41 namespace {
42   Statistic<> NumSpilled ("ra-simple", "Number of registers spilled");
43   Statistic<> NumReloaded("ra-simple", "Number of registers reloaded");
44
45   class RegAllocSimple : public FunctionPass {
46     TargetMachine &TM;
47     MachineFunction *MF;
48     const MRegisterInfo *RegInfo;
49     unsigned NumBytesAllocated;
50     
51     // Maps SSA Regs => offsets on the stack where these values are stored
52     std::map<unsigned, unsigned> VirtReg2OffsetMap;
53
54     // RegsUsed - Keep track of what registers are currently in use.
55     std::set<unsigned> RegsUsed;
56
57     // RegClassIdx - Maps RegClass => which index we can take a register
58     // from. Since this is a simple register allocator, when we need a register
59     // of a certain class, we just take the next available one.
60     std::map<const TargetRegisterClass*, unsigned> RegClassIdx;
61
62   public:
63
64     RegAllocSimple(TargetMachine &tm)
65       : TM(tm), RegInfo(tm.getRegisterInfo()) {
66       RegsUsed.insert(RegInfo->getFramePointer());
67       RegsUsed.insert(RegInfo->getStackPointer());
68
69       cleanupAfterFunction();
70     }
71
72     bool runOnFunction(Function &Fn) {
73       return runOnMachineFunction(MachineFunction::get(&Fn));
74     }
75
76     virtual const char *getPassName() const {
77       return "Simple Register Allocator";
78     }
79
80   private:
81     /// runOnMachineFunction - Register allocate the whole function
82     bool runOnMachineFunction(MachineFunction &Fn);
83
84     /// AllocateBasicBlock - Register allocate the specified basic block.
85     void AllocateBasicBlock(MachineBasicBlock &MBB);
86
87     /// EliminatePHINodes - Eliminate phi nodes by inserting copy instructions
88     /// in predecessor basic blocks.
89     void EliminatePHINodes(MachineBasicBlock &MBB);
90
91
92     /// getStackSpaceFor - This returns the offset of the specified virtual
93     /// register on the stack, allocating space if neccesary.
94     unsigned getStackSpaceFor(unsigned VirtReg, 
95                               const TargetRegisterClass *regClass);
96
97     /// Given a virtual register, return a compatible physical register that is
98     /// currently unused.
99     ///
100     /// Side effect: marks that register as being used until manually cleared
101     ///
102     unsigned getFreeReg(unsigned virtualReg);
103
104     /// Returns all `borrowed' registers back to the free pool
105     void clearAllRegs() {
106       RegClassIdx.clear();
107     }
108
109     /// Invalidates any references, real or implicit, to physical registers
110     ///
111     void invalidatePhysRegs(const MachineInstr *MI) {
112       unsigned Opcode = MI->getOpcode();
113       const MachineInstrDescriptor &Desc = TM.getInstrInfo().get(Opcode);
114       const unsigned *regs = Desc.ImplicitUses;
115       while (*regs)
116         RegsUsed.insert(*regs++);
117
118       regs = Desc.ImplicitDefs;
119       while (*regs)
120         RegsUsed.insert(*regs++);
121     }
122
123     void cleanupAfterFunction() {
124       VirtReg2OffsetMap.clear();
125       NumBytesAllocated = 4;   // FIXME: This is X86 specific
126     }
127
128     /// Moves value from memory into that register
129     unsigned reloadVirtReg(MachineBasicBlock &MBB,
130                            MachineBasicBlock::iterator &I, unsigned VirtReg);
131
132     /// Saves reg value on the stack (maps virtual register to stack value)
133     void spillVirtReg(MachineBasicBlock &MBB, MachineBasicBlock::iterator &I,
134                       unsigned VirtReg, unsigned PhysReg);
135   };
136
137 }
138
139 /// getStackSpaceFor - This allocates space for the specified virtual
140 /// register to be held on the stack.
141 unsigned RegAllocSimple::getStackSpaceFor(unsigned VirtReg,
142                                           const TargetRegisterClass *regClass) {
143   // Find the location VirtReg would belong...
144   std::map<unsigned, unsigned>::iterator I =
145     VirtReg2OffsetMap.lower_bound(VirtReg);
146
147   if (I != VirtReg2OffsetMap.end() && I->first == VirtReg)
148     return I->second;          // Already has space allocated?
149
150   unsigned RegSize = regClass->getDataSize();
151
152   // Align NumBytesAllocated.  We should be using TargetData alignment stuff
153   // to determine this, but we don't know the LLVM type associated with the
154   // virtual register.  Instead, just align to a multiple of the size for now.
155   NumBytesAllocated += RegSize-1;
156   NumBytesAllocated = NumBytesAllocated/RegSize*RegSize;
157   
158   // Assign the slot...
159   VirtReg2OffsetMap.insert(I, std::make_pair(VirtReg, NumBytesAllocated));
160   
161   // Reserve the space!
162   NumBytesAllocated += RegSize;
163   return NumBytesAllocated-RegSize;
164 }
165
166 unsigned RegAllocSimple::getFreeReg(unsigned virtualReg) {
167   const TargetRegisterClass* regClass = MF->getRegClass(virtualReg);
168   
169   unsigned regIdx = RegClassIdx[regClass]++;
170   assert(regIdx < regClass->getNumRegs() && "Not enough registers!");
171   unsigned physReg = regClass->getRegister(regIdx);
172
173   if (RegsUsed.find(physReg) == RegsUsed.end())
174     return physReg;
175   else
176     return getFreeReg(virtualReg);
177 }
178
179 unsigned RegAllocSimple::reloadVirtReg(MachineBasicBlock &MBB,
180                                        MachineBasicBlock::iterator &I,
181                                        unsigned VirtReg) {
182   const TargetRegisterClass* regClass = MF->getRegClass(VirtReg);
183   unsigned stackOffset = getStackSpaceFor(VirtReg, regClass);
184   unsigned PhysReg = getFreeReg(VirtReg);
185
186   // Add move instruction(s)
187   ++NumReloaded;
188   I = RegInfo->loadRegOffset2Reg(MBB, I, PhysReg, RegInfo->getFramePointer(),
189                                  -stackOffset, regClass->getDataSize());
190   return PhysReg;
191 }
192
193 void RegAllocSimple::spillVirtReg(MachineBasicBlock &MBB,
194                                   MachineBasicBlock::iterator &I,
195                                   unsigned VirtReg, unsigned PhysReg)
196 {
197   const TargetRegisterClass* regClass = MF->getRegClass(VirtReg);
198   unsigned stackOffset = getStackSpaceFor(VirtReg, regClass);
199
200   // Add move instruction(s)
201   ++NumSpilled;
202   I = RegInfo->storeReg2RegOffset(MBB, I, PhysReg, RegInfo->getFramePointer(),
203                                   -stackOffset, regClass->getDataSize());
204 }
205
206
207 /// EliminatePHINodes - Eliminate phi nodes by inserting copy instructions in
208 /// predecessor basic blocks.
209 ///
210 void RegAllocSimple::EliminatePHINodes(MachineBasicBlock &MBB) {
211   const MachineInstrInfo &MII = TM.getInstrInfo();
212
213   while (MBB.front()->getOpcode() == MachineInstrInfo::PHI) {
214     MachineInstr *MI = MBB.front();
215     // Unlink the PHI node from the basic block... but don't delete the PHI yet
216     MBB.erase(MBB.begin());
217     
218     DEBUG(std::cerr << "num ops: " << MI->getNumOperands() << "\n");
219     assert(MI->getOperand(0).isVirtualRegister() &&
220            "PHI node doesn't write virt reg?");
221
222     unsigned virtualReg = MI->getOperand(0).getAllocatedRegNum();
223     
224     for (int i = MI->getNumOperands() - 1; i >= 2; i-=2) {
225       MachineOperand &opVal = MI->getOperand(i-1);
226       
227       // Get the MachineBasicBlock equivalent of the BasicBlock that is the
228       // source path the phi
229       MachineBasicBlock &opBlock = *MI->getOperand(i).getMachineBasicBlock();
230
231       // Check to make sure we haven't already emitted the copy for this block.
232       // This can happen because PHI nodes may have multiple entries for the
233       // same basic block.  It doesn't matter which entry we use though, because
234       // all incoming values are guaranteed to be the same for a particular bb.
235       //
236       // Note that this is N^2 in the number of phi node entries, but since the
237       // # of entries is tiny, this is not a problem.
238       //
239       bool HaveNotEmitted = true;
240       for (int op = MI->getNumOperands() - 1; op != i; op -= 2)
241         if (&opBlock == MI->getOperand(op).getMachineBasicBlock()) {
242           HaveNotEmitted = false;
243           break;
244         }
245
246       if (HaveNotEmitted) {
247         MachineBasicBlock::iterator opI = opBlock.end();
248         MachineInstr *opMI = *--opI;
249         
250         // must backtrack over ALL the branches in the previous block
251         while (MII.isBranch(opMI->getOpcode()) && opI != opBlock.begin())
252           opMI = *--opI;
253         
254         // move back to the first branch instruction so new instructions
255         // are inserted right in front of it and not in front of a non-branch
256         if (!MII.isBranch(opMI->getOpcode()))
257           ++opI;
258
259         unsigned dataSize = MF->getRegClass(virtualReg)->getDataSize();
260
261         // Retrieve the constant value from this op, move it to target
262         // register of the phi
263         if (opVal.isImmediate()) {
264           opI = RegInfo->moveImm2Reg(opBlock, opI, virtualReg,
265                                      (unsigned) opVal.getImmedValue(),
266                                      dataSize);
267         } else {
268           opI = RegInfo->moveReg2Reg(opBlock, opI, virtualReg,
269                                      opVal.getAllocatedRegNum(), dataSize);
270         }
271       }
272     }
273     
274     // really delete the PHI instruction now!
275     delete MI;
276   }
277 }
278
279
280 void RegAllocSimple::AllocateBasicBlock(MachineBasicBlock &MBB) {
281   // loop over each instruction
282   for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
283     // Made to combat the incorrect allocation of r2 = add r1, r1
284     std::map<unsigned, unsigned> Virt2PhysRegMap;
285
286     MachineInstr *MI = *I;
287     
288     // a preliminary pass that will invalidate any registers that
289     // are used by the instruction (including implicit uses)
290     invalidatePhysRegs(MI);
291     
292     // Loop over uses, move from memory into registers
293     for (int i = MI->getNumOperands() - 1; i >= 0; --i) {
294       MachineOperand &op = MI->getOperand(i);
295       
296       if (op.isVirtualRegister()) {
297         unsigned virtualReg = (unsigned) op.getAllocatedRegNum();
298         DEBUG(std::cerr << "op: " << op << "\n");
299         DEBUG(std::cerr << "\t inst[" << i << "]: ";
300               MI->print(std::cerr, TM));
301         
302         // make sure the same virtual register maps to the same physical
303         // register in any given instruction
304         unsigned physReg = Virt2PhysRegMap[virtualReg];
305         if (physReg == 0) {
306           if (op.opIsDef()) {
307             if (TM.getInstrInfo().isTwoAddrInstr(MI->getOpcode()) && i == 0) {
308               // must be same register number as the first operand
309               // This maps a = b + c into b += c, and saves b into a's spot
310               assert(MI->getOperand(1).isRegister()  &&
311                      MI->getOperand(1).getAllocatedRegNum() &&
312                      MI->getOperand(1).opIsUse() &&
313                      "Two address instruction invalid!");
314
315               physReg = MI->getOperand(1).getAllocatedRegNum();
316             } else {
317               physReg = getFreeReg(virtualReg);
318             }
319             ++I;
320             spillVirtReg(MBB, I, virtualReg, physReg);
321             --I;
322           } else {
323             physReg = reloadVirtReg(MBB, I, virtualReg);
324             Virt2PhysRegMap[virtualReg] = physReg;
325           }
326         }
327         MI->SetMachineOperandReg(i, physReg);
328         DEBUG(std::cerr << "virt: " << virtualReg << 
329               ", phys: " << op.getAllocatedRegNum() << "\n");
330       }
331     }
332     clearAllRegs();
333   }
334 }
335
336 /// runOnMachineFunction - Register allocate the whole function
337 ///
338 bool RegAllocSimple::runOnMachineFunction(MachineFunction &Fn) {
339   DEBUG(std::cerr << "Machine Function " << "\n");
340   MF = &Fn;
341
342   // First pass: eliminate PHI instructions by inserting copies into predecessor
343   // blocks.
344   for (MachineFunction::iterator MBB = Fn.begin(), MBBe = Fn.end();
345        MBB != MBBe; ++MBB)
346     EliminatePHINodes(*MBB);
347
348   // Loop over all of the basic blocks, eliminating virtual register references
349   for (MachineFunction::iterator MBB = Fn.begin(), MBBe = Fn.end();
350        MBB != MBBe; ++MBB)
351     AllocateBasicBlock(*MBB);
352
353   // Add prologue to the function...
354   RegInfo->emitPrologue(Fn, NumBytesAllocated);
355
356   const MachineInstrInfo &MII = TM.getInstrInfo();
357
358   // Add epilogue to restore the callee-save registers in each exiting block
359   for (MachineFunction::iterator MBB = Fn.begin(), MBBe = Fn.end();
360        MBB != MBBe; ++MBB) {
361     // If last instruction is a return instruction, add an epilogue
362     if (MII.isReturn(MBB->back()->getOpcode()))
363       RegInfo->emitEpilogue(*MBB, NumBytesAllocated);
364   }
365
366   cleanupAfterFunction();
367   return true;
368 }
369
370 Pass *createSimpleX86RegisterAllocator(TargetMachine &TM) {
371   return new RegAllocSimple(TM);
372 }