Grab bag of minor cleanups. Export some statistics about the number of
[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       MachineBasicBlock::iterator opI = opBlock.end();
294       MachineInstr *opMI = *--opI;
295
296       // must backtrack over ALL the branches in the previous block, until no
297       // more
298       while (MII.isBranch(opMI->getOpcode()) && opI != opBlock.begin())
299         opMI = *--opI;
300
301       // move back to the first branch instruction so new instructions
302       // are inserted right in front of it and not in front of a non-branch
303       if (!MII.isBranch(opMI->getOpcode()))
304         ++opI;
305       
306       // Retrieve the constant value from this op, move it to target
307       // register of the phi
308       if (opVal.isImmediate()) {
309         opI = RegInfo->moveImm2Reg(opBlock, opI, physReg,
310                                    (unsigned) opVal.getImmedValue(),
311                                    dataSize);
312         saveVirtRegToStack(opBlock, opI, virtualReg, physReg);
313       } else {
314         // Allocate a physical register and add a move in the BB
315         unsigned opVirtualReg = (unsigned) opVal.getAllocatedRegNum();
316         unsigned opPhysReg;
317         opI = moveUseToReg(opBlock, opI, opVirtualReg, physReg);
318
319         // Save that register value to the stack of the TARGET REG
320         saveVirtRegToStack(opBlock, opI, virtualReg, physReg);
321       }
322       
323       // make regs available to other instructions
324       clearAllRegs();
325     }
326     
327     // really delete the instruction
328     delete MI;
329   }
330 }
331
332
333 void RegAllocSimple::AllocateBasicBlock(MachineBasicBlock &MBB) {
334   // Handle PHI instructions specially: add moves to each pred block
335   EliminatePHINodes(MBB);
336   
337   //loop over each basic block
338   for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
339     MachineInstr *MI = *I;
340     
341     // a preliminary pass that will invalidate any registers that
342     // are used by the instruction (including implicit uses)
343     invalidatePhysRegs(MI);
344     
345     // Loop over uses, move from memory into registers
346     for (int i = MI->getNumOperands() - 1; i >= 0; --i) {
347       MachineOperand &op = MI->getOperand(i);
348       
349       if (op.isVirtualRegister()) {
350         unsigned virtualReg = (unsigned) op.getAllocatedRegNum();
351         DEBUG(std::cerr << "op: " << op << "\n");
352         DEBUG(std::cerr << "\t inst[" << i << "]: ";
353               MI->print(std::cerr, TM));
354         
355         // make sure the same virtual register maps to the same physical
356         // register in any given instruction
357         unsigned physReg;
358         if (VirtReg2PhysRegMap.find(virtualReg) != VirtReg2PhysRegMap.end()) {
359           physReg = VirtReg2PhysRegMap[virtualReg];
360         } else {
361           if (op.opIsDef()) {
362             if (TM.getInstrInfo().isTwoAddrInstr(MI->getOpcode()) && i == 0) {
363               // must be same register number as the first operand
364               // This maps a = b + c into b += c, and saves b into a's spot
365               physReg = MI->getOperand(1).getAllocatedRegNum();
366             } else {
367               physReg = getFreeReg(virtualReg);
368             }
369             MachineBasicBlock::iterator J = I;
370             J = saveVirtRegToStack(MBB, ++J, virtualReg, physReg);
371             I = --J;
372           } else {
373             I = moveUseToReg(MBB, I, virtualReg, physReg);
374           }
375           VirtReg2PhysRegMap[virtualReg] = physReg;
376         }
377         MI->SetMachineOperandReg(i, physReg);
378         DEBUG(std::cerr << "virt: " << virtualReg << 
379               ", phys: " << op.getAllocatedRegNum() << "\n");
380       }
381     }
382     
383     clearAllRegs();
384     VirtReg2PhysRegMap.clear();
385   }
386 }
387
388 /// runOnMachineFunction - Register allocate the whole function
389 ///
390 bool RegAllocSimple::runOnMachineFunction(MachineFunction &Fn) {
391   DEBUG(std::cerr << "Machine Function " << "\n");
392   MF = &Fn;
393
394   for (MachineFunction::iterator MBB = Fn.begin(), MBBe = Fn.end();
395        MBB != MBBe; ++MBB)
396     AllocateBasicBlock(*MBB);
397
398   // add prologue we should preserve callee-save registers...
399   RegInfo->emitPrologue(Fn, NumBytesAllocated);
400
401   const MachineInstrInfo &MII = TM.getInstrInfo();
402
403   // add epilogue to restore the callee-save registers
404   // loop over the basic block
405   for (MachineFunction::iterator MBB = Fn.begin(), MBBe = Fn.end();
406        MBB != MBBe; ++MBB) {
407     // check if last instruction is a RET
408     if (MII.isReturn(MBB->back()->getOpcode())) {
409       // this block has a return instruction, add epilogue
410       RegInfo->emitEpilogue(*MBB, NumBytesAllocated);
411     }
412   }
413
414   cleanupAfterFunction();
415   return false;  // We never modify the LLVM itself.
416 }
417
418 Pass *createSimpleX86RegisterAllocator(TargetMachine &TM) {
419   return new RegAllocSimple(TM);
420 }