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