Start allocating stack space at [ebp-4] to not overwrite the return address.
[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/Function.h"
8 #include "llvm/iTerminators.h"
9 #include "llvm/Type.h"
10 #include "llvm/Constants.h"
11 #include "llvm/Pass.h"
12 #include "llvm/CodeGen/MachineInstr.h"
13 #include "llvm/CodeGen/MachineFunction.h"
14 #include "llvm/CodeGen/MachineInstrBuilder.h"
15 #include "llvm/Target/MachineInstrInfo.h"
16 #include "llvm/Target/MRegisterInfo.h"
17 #include "llvm/Target/MachineRegInfo.h"
18 #include "llvm/Target/TargetMachine.h"
19 #include "llvm/Support/InstVisitor.h"
20 #include "Support/Statistic.h"
21 #include <map>
22
23 namespace {
24   struct RegAllocSimple : public FunctionPass {
25     TargetMachine &TM;
26     MachineBasicBlock *CurrMBB;
27     MachineFunction *MF;
28     unsigned maxOffset;
29     const MRegisterInfo *RegInfo;
30     unsigned NumBytesAllocated, ByteAlignment;
31     
32     // Maps SSA Regs => offsets on the stack where these values are stored
33     // FIXME: change name to VirtReg2OffsetMap
34     std::map<unsigned, unsigned> RegMap;
35
36     // Maps SSA Regs => physical regs
37     std::map<unsigned, unsigned> SSA2PhysRegMap;
38
39     // Maps physical register to their register classes
40     std::map<unsigned, const TargetRegisterClass*> PhysReg2RegClassMap;
41
42     // Made to combat the incorrect allocation of r2 = add r1, r1
43     std::map<unsigned, unsigned> VirtReg2PhysRegMap;
44     
45     // Maps RegClass => which index we can take a register from. Since this is a
46     // simple register allocator, when we need a register of a certain class, we
47     // just take the next available one.
48     std::map<unsigned, unsigned> RegsUsed;
49     std::map<const TargetRegisterClass*, unsigned> RegClassIdx;
50
51     RegAllocSimple(TargetMachine &tm) : TM(tm), CurrMBB(0), maxOffset(0), 
52                                         RegInfo(tm.getRegisterInfo()),
53                                         ByteAlignment(4)
54     {
55       // build reverse mapping for physReg -> register class
56       RegInfo->buildReg2RegClassMap(PhysReg2RegClassMap);
57
58       RegsUsed[RegInfo->getFramePointer()] = 1;
59       RegsUsed[RegInfo->getStackPointer()] = 1;
60
61       cleanupAfterFunction();
62     }
63
64     bool isAvailableReg(unsigned Reg) {
65       // assert(Reg < MRegisterInfo::FirstVirtualReg && "...");
66       return RegsUsed.find(Reg) == RegsUsed.end();
67     }
68
69     ///
70     unsigned allocateStackSpaceFor(unsigned VirtReg, 
71                                    const TargetRegisterClass *regClass);
72
73     /// Given size (in bytes), returns a register that is currently unused
74     /// Side effect: marks that register as being used until manually cleared
75     unsigned getFreeReg(unsigned virtualReg);
76
77     /// Returns all `borrowed' registers back to the free pool
78     void clearAllRegs() {
79         RegClassIdx.clear();
80     }
81
82     void cleanupAfterFunction() {
83       RegMap.clear();
84       SSA2PhysRegMap.clear();
85       NumBytesAllocated = 4;
86     }
87
88     /// Moves value from memory into that register
89     MachineBasicBlock::iterator
90     moveUseToReg (MachineBasicBlock::iterator I, unsigned VirtReg,
91                   unsigned &PhysReg);
92
93     /// Saves reg value on the stack (maps virtual register to stack value)
94     MachineBasicBlock::iterator
95     saveVirtRegToStack (MachineBasicBlock::iterator I, unsigned VirtReg,
96                         unsigned PhysReg);
97
98     MachineBasicBlock::iterator
99     savePhysRegToStack (MachineBasicBlock::iterator I, unsigned PhysReg);
100
101     /// runOnFunction - Top level implementation of instruction selection for
102     /// the entire function.
103     ///
104     bool runOnMachineFunction(MachineFunction &Fn);
105
106     bool runOnFunction(Function &Fn) {
107       return runOnMachineFunction(MachineFunction::get(&Fn));
108     }
109   };
110
111 }
112
113 unsigned RegAllocSimple::allocateStackSpaceFor(unsigned VirtReg,
114                                             const TargetRegisterClass *regClass)
115 {
116   if (RegMap.find(VirtReg) == RegMap.end()) {
117 #if 0
118     unsigned size = regClass->getDataSize();
119     unsigned over = NumBytesAllocated - (NumBytesAllocated % ByteAlignment);
120     if (size >= ByteAlignment - over) {
121       // need to pad by (ByteAlignment - over)
122       NumBytesAllocated += ByteAlignment - over;
123     }
124     RegMap[VirtReg] = NumBytesAllocated;
125     NumBytesAllocated += size;
126 #endif
127     // FIXME: forcing each arg to take 4 bytes on the stack
128     RegMap[VirtReg] = NumBytesAllocated;
129     NumBytesAllocated += 4;
130   }
131   return RegMap[VirtReg];
132 }
133
134 unsigned RegAllocSimple::getFreeReg(unsigned virtualReg) {
135   const TargetRegisterClass* regClass = MF->getRegClass(virtualReg);
136   unsigned physReg;
137   assert(regClass);
138   if (RegClassIdx.find(regClass) != RegClassIdx.end()) {
139     unsigned regIdx = RegClassIdx[regClass]++;
140     assert(regIdx < regClass->getNumRegs() && "Not enough registers!");
141     physReg = regClass->getRegister(regIdx);
142   } else {
143     physReg = regClass->getRegister(0);
144     // assert(physReg < regClass->getNumRegs() && "No registers in class!");
145     RegClassIdx[regClass] = 1;
146   }
147
148   if (isAvailableReg(physReg))
149     return physReg;
150   else {
151     return getFreeReg(virtualReg);
152   }
153 }
154
155 MachineBasicBlock::iterator
156 RegAllocSimple::moveUseToReg (MachineBasicBlock::iterator I,
157                               unsigned VirtReg, unsigned &PhysReg)
158 {
159   const TargetRegisterClass* regClass = MF->getRegClass(VirtReg);
160   assert(regClass);
161
162   unsigned stackOffset = allocateStackSpaceFor(VirtReg, regClass);
163   PhysReg = getFreeReg(VirtReg);
164
165   // FIXME: increment the frame pointer
166
167   // Add move instruction(s)
168   return RegInfo->loadRegOffset2Reg(CurrMBB, I, PhysReg,
169                                     RegInfo->getFramePointer(),
170                                     -stackOffset, regClass->getDataSize());
171 }
172
173 MachineBasicBlock::iterator
174 RegAllocSimple::saveVirtRegToStack (MachineBasicBlock::iterator I,
175                                     unsigned VirtReg, unsigned PhysReg)
176 {
177   const TargetRegisterClass* regClass = MF->getRegClass(VirtReg);
178   assert(regClass);
179
180   unsigned stackOffset = allocateStackSpaceFor(VirtReg, regClass);
181
182   // Add move instruction(s)
183   return RegInfo->storeReg2RegOffset(CurrMBB, I, PhysReg,
184                                      RegInfo->getFramePointer(),
185                                      -stackOffset, regClass->getDataSize());
186 }
187
188 MachineBasicBlock::iterator
189 RegAllocSimple::savePhysRegToStack (MachineBasicBlock::iterator I,
190                                     unsigned PhysReg)
191 {
192   const TargetRegisterClass* regClass = MF->getRegClass(PhysReg);
193   assert(regClass);
194
195   unsigned offset = allocateStackSpaceFor(PhysReg, regClass);
196
197   // Add move instruction(s)
198   return RegInfo->storeReg2RegOffset(CurrMBB, I, PhysReg,
199                                      RegInfo->getFramePointer(),
200                                      offset, regClass->getDataSize());
201 }
202
203 bool RegAllocSimple::runOnMachineFunction(MachineFunction &Fn) {
204   cleanupAfterFunction();
205
206   unsigned virtualReg, physReg;
207   DEBUG(std::cerr << "Machine Function " << "\n");
208   MF = &Fn;
209
210   for (MachineFunction::iterator MBB = Fn.begin(), MBBe = Fn.end();
211        MBB != MBBe; ++MBB)
212   {
213     CurrMBB = &(*MBB);
214
215     //loop over each basic block
216     for (MachineBasicBlock::iterator I = MBB->begin(); I != MBB->end(); ++I)
217     {
218       MachineInstr *MI = *I;
219
220       DEBUG(std::cerr << "instr: ";
221             MI->print(std::cerr, TM));
222
223       // FIXME: add a preliminary pass that will invalidate any registers that
224       // are used by the instruction (including implicit uses)
225
226
227       // Loop over each instruction:
228       // uses, move from memory into registers
229       for (int i = MI->getNumOperands() - 1; i >= 0; --i) {
230         MachineOperand &op = MI->getOperand(i);
231
232         if (op.getType() == MachineOperand::MO_SignExtendedImmed ||
233             op.getType() == MachineOperand::MO_UnextendedImmed)
234         {
235           DEBUG(std::cerr << "const\n");
236         } else if (op.isVirtualRegister()) {
237           virtualReg = (unsigned) op.getAllocatedRegNum();
238           DEBUG(std::cerr << "op: " << op << "\n");
239           DEBUG(std::cerr << "\t inst[" << i << "]: ";
240                 MI->print(std::cerr, TM));
241
242           // make sure the same virtual register maps to the same physical
243           // register in any given instruction
244           if (VirtReg2PhysRegMap.find(virtualReg) != VirtReg2PhysRegMap.end()) {
245             physReg = VirtReg2PhysRegMap[virtualReg];
246           } else {
247             if (op.opIsDef()) {
248               if (TM.getInstrInfo().isTwoAddrInstr(MI->getOpcode()) && i == 0) {
249                 // must be same register number as the first operand
250                 // This maps a = b + c into b += c, and saves b into a's spot
251                 physReg = (unsigned) MI->getOperand(1).getAllocatedRegNum();
252               } else {
253                 physReg = getFreeReg(virtualReg);
254               }
255               MachineBasicBlock::iterator J = I;
256               J = saveVirtRegToStack(++J, virtualReg, physReg);
257               I = --J;
258             } else {
259               I = moveUseToReg(I, virtualReg, physReg);
260             }
261             VirtReg2PhysRegMap[virtualReg] = physReg;
262           }
263           MI->SetMachineOperandReg(i, physReg);
264           DEBUG(std::cerr << "virt: " << virtualReg << 
265                 ", phys: " << op.getAllocatedRegNum() << "\n");
266         }
267       }
268
269       clearAllRegs();
270       VirtReg2PhysRegMap.clear();
271     }
272
273   }
274
275   // add prologue we should preserve callee-save registers...
276   MachineFunction::iterator Fi = Fn.begin();
277   MachineBasicBlock *MBB = Fi;
278   MachineBasicBlock::iterator MBBi = MBB->begin();
279   RegInfo->emitPrologue(MBB, MBBi, NumBytesAllocated);
280
281   // add epilogue to restore the callee-save registers
282   // loop over the basic block
283   for (MachineFunction::iterator MBB = Fn.begin(), MBBe = Fn.end();
284        MBB != MBBe; ++MBB)
285   {
286     // check if last instruction is a RET
287     MachineBasicBlock::iterator I = (*MBB).end();
288     MachineInstr *MI = *(--I);
289     const MachineInstrInfo &MII = TM.getInstrInfo();
290     if (MII.isReturn(MI->getOpcode())) {
291       // this block has a return instruction, add epilogue
292       RegInfo->emitEpilogue(MBB, I, NumBytesAllocated);
293     }
294   }
295
296   return false;  // We never modify the LLVM itself.
297 }
298
299 Pass *createSimpleX86RegisterAllocator(TargetMachine &TM) {
300   return new RegAllocSimple(TM);
301 }