af9b6716888826a9e1d0b989860a553ce6e23aa1
[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     // Maps RegClass => which index we can take a register from. Since this is a
43     // simple register allocator, when we need a register of a certain class, we
44     // just take the next available one.
45     std::map<unsigned, unsigned> RegsUsed;
46     std::map<const TargetRegisterClass*, unsigned> RegClassIdx;
47
48     RegAllocSimple(TargetMachine &tm) : TM(tm), CurrMBB(0), maxOffset(0), 
49                                         RegInfo(tm.getRegisterInfo()),
50                                         NumBytesAllocated(0), ByteAlignment(4)
51     {
52       // build reverse mapping for physReg -> register class
53       RegInfo->buildReg2RegClassMap(PhysReg2RegClassMap);
54
55       RegsUsed[RegInfo->getFramePointer()] = 1;
56       RegsUsed[RegInfo->getStackPointer()] = 1;
57     }
58
59     bool isAvailableReg(unsigned Reg) {
60       // assert(Reg < MRegisterInfo::FirstVirtualReg && "...");
61       return RegsUsed.find(Reg) == RegsUsed.end();
62     }
63
64     ///
65     unsigned allocateStackSpaceFor(unsigned VirtReg, 
66                                    const TargetRegisterClass *regClass);
67
68     /// Given size (in bytes), returns a register that is currently unused
69     /// Side effect: marks that register as being used until manually cleared
70     unsigned getFreeReg(unsigned virtualReg);
71
72     /// Returns all `borrowed' registers back to the free pool
73     void clearAllRegs() {
74         RegClassIdx.clear();
75     }
76
77     void cleanupAfterFunction() {
78       RegMap.clear();
79       SSA2PhysRegMap.clear();
80       NumBytesAllocated = 0;
81     }
82
83     /// Moves value from memory into that register
84     MachineBasicBlock::iterator
85     moveUseToReg (MachineBasicBlock::iterator I, unsigned VirtReg,
86                   unsigned &PhysReg);
87
88     /// Saves reg value on the stack (maps virtual register to stack value)
89     MachineBasicBlock::iterator
90     saveVirtRegToStack (MachineBasicBlock::iterator I, unsigned VirtReg,
91                         unsigned PhysReg);
92
93     MachineBasicBlock::iterator
94     savePhysRegToStack (MachineBasicBlock::iterator I, unsigned PhysReg);
95
96     /// runOnFunction - Top level implementation of instruction selection for
97     /// the entire function.
98     ///
99     bool runOnMachineFunction(MachineFunction &Fn);
100
101     bool runOnFunction(Function &Fn) {
102       return runOnMachineFunction(MachineFunction::get(&Fn));
103     }
104   };
105
106 }
107
108 unsigned RegAllocSimple::allocateStackSpaceFor(unsigned VirtReg,
109                                             const TargetRegisterClass *regClass)
110 {
111   if (RegMap.find(VirtReg) == RegMap.end()) {
112     unsigned size = regClass->getDataSize();
113     unsigned over = NumBytesAllocated - (NumBytesAllocated % ByteAlignment);
114     if (size >= ByteAlignment - over) {
115       // need to pad by (ByteAlignment - over)
116       NumBytesAllocated += ByteAlignment - over;
117     }
118     RegMap[VirtReg] = NumBytesAllocated;
119     NumBytesAllocated += size;
120   }
121   return RegMap[VirtReg];
122 }
123
124 unsigned RegAllocSimple::getFreeReg(unsigned virtualReg) {
125   const TargetRegisterClass* regClass = MF->getRegClass(virtualReg);
126   unsigned physReg;
127   assert(regClass);
128   if (RegClassIdx.find(regClass) != RegClassIdx.end()) {
129     unsigned regIdx = RegClassIdx[regClass]++;
130     assert(regIdx < regClass->getNumRegs() && "Not enough registers!");
131     physReg = regClass->getRegister(regIdx);
132   } else {
133     physReg = regClass->getRegister(0);
134     // assert(physReg < regClass->getNumRegs() && "No registers in class!");
135     RegClassIdx[regClass] = 1;
136   }
137
138   if (isAvailableReg(physReg))
139     return physReg;
140   else {
141     return getFreeReg(virtualReg);
142   }
143 }
144
145 MachineBasicBlock::iterator
146 RegAllocSimple::moveUseToReg (MachineBasicBlock::iterator I,
147                               unsigned VirtReg, unsigned &PhysReg)
148 {
149   const TargetRegisterClass* regClass = MF->getRegClass(VirtReg);
150   assert(regClass);
151
152   unsigned stackOffset = allocateStackSpaceFor(VirtReg, regClass);
153   PhysReg = getFreeReg(VirtReg);
154
155   // FIXME: increment the frame pointer
156
157   // Add move instruction(s)
158   return RegInfo->loadRegOffset2Reg(CurrMBB, I, PhysReg,
159                                     RegInfo->getFramePointer(),
160                                     -stackOffset, regClass->getDataSize());
161 }
162
163 MachineBasicBlock::iterator
164 RegAllocSimple::saveVirtRegToStack (MachineBasicBlock::iterator I,
165                                     unsigned VirtReg, unsigned PhysReg)
166 {
167   const TargetRegisterClass* regClass = MF->getRegClass(VirtReg);
168   assert(regClass);
169
170   unsigned stackOffset = allocateStackSpaceFor(VirtReg, regClass);
171
172   // Add move instruction(s)
173   return RegInfo->storeReg2RegOffset(CurrMBB, I, PhysReg,
174                                      RegInfo->getFramePointer(),
175                                      -stackOffset, regClass->getDataSize());
176 }
177
178 MachineBasicBlock::iterator
179 RegAllocSimple::savePhysRegToStack (MachineBasicBlock::iterator I,
180                                     unsigned PhysReg)
181 {
182   const TargetRegisterClass* regClass = MF->getRegClass(PhysReg);
183   assert(regClass);
184
185   unsigned offset = allocateStackSpaceFor(PhysReg, regClass);
186
187   // Add move instruction(s)
188   return RegInfo->storeReg2RegOffset(CurrMBB, I, PhysReg,
189                                      RegInfo->getFramePointer(),
190                                      offset, regClass->getDataSize());
191 }
192
193 bool RegAllocSimple::runOnMachineFunction(MachineFunction &Fn) {
194   cleanupAfterFunction();
195
196   unsigned virtualReg, physReg;
197   DEBUG(std::cerr << "Machine Function " << "\n");
198   MF = &Fn;
199
200   for (MachineFunction::iterator MBB = Fn.begin(), MBBe = Fn.end();
201        MBB != MBBe; ++MBB)
202   {
203     CurrMBB = &(*MBB);
204
205     //loop over each basic block
206     for (MachineBasicBlock::iterator I = MBB->begin(); I != MBB->end(); ++I)
207     {
208       MachineInstr *MI = *I;
209
210       DEBUG(std::cerr << "instr: ";
211             MI->print(std::cerr, TM));
212
213       // FIXME: add a preliminary pass that will invalidate any registers that
214       // are used by the instruction (including implicit uses)
215
216
217       // Loop over each instruction:
218       // uses, move from memory into registers
219       for (int i = MI->getNumOperands() - 1; i >= 0; --i) {
220         MachineOperand &op = MI->getOperand(i);
221
222         if (op.getType() == MachineOperand::MO_SignExtendedImmed ||
223             op.getType() == MachineOperand::MO_UnextendedImmed)
224         {
225           DEBUG(std::cerr << "const\n");
226         } else if (op.isVirtualRegister()) {
227           virtualReg = (unsigned) op.getAllocatedRegNum();
228           // save register to stack if it's a def
229           DEBUG(std::cerr << "op: " << op << "\n");
230           DEBUG(std::cerr << "\t inst[" << i << "]: ";
231                 MI->print(std::cerr, TM));
232           if (op.opIsDef()) {
233             physReg = getFreeReg(virtualReg);
234             MachineBasicBlock::iterator J = I;
235             J = saveVirtRegToStack(++J, virtualReg, physReg);
236             I = --J;
237           } else {
238             I = moveUseToReg(I, virtualReg, physReg);
239           }
240           MI->SetMachineOperandReg(i, physReg);
241           DEBUG(std::cerr << "virt: " << virtualReg << 
242                 ", phys: " << op.getAllocatedRegNum() << "\n");
243         }
244       }
245
246       clearAllRegs();
247     }
248
249   }
250
251   // add prologue we should preserve callee-save registers...
252   MachineFunction::iterator Fi = Fn.begin();
253   MachineBasicBlock *MBB = Fi;
254   MachineBasicBlock::iterator MBBi = MBB->begin();
255   RegInfo->emitPrologue(MBB, MBBi, NumBytesAllocated);
256
257   // add epilogue to restore the callee-save registers
258   // loop over the basic block
259   for (MachineFunction::iterator MBB = Fn.begin(), MBBe = Fn.end();
260        MBB != MBBe; ++MBB)
261   {
262     // check if last instruction is a RET
263     MachineBasicBlock::iterator I = (*MBB).end();
264     MachineInstr *MI = *(--I);
265     const MachineInstrInfo &MII = TM.getInstrInfo();
266     if (MII.isReturn(MI->getOpcode())) {
267       // this block has a return instruction, add epilogue
268       RegInfo->emitEpilogue(MBB, I, NumBytesAllocated);
269     }
270   }
271
272   return false;  // We never modify the LLVM itself.
273 }
274
275 Pass *createSimpleX86RegisterAllocator(TargetMachine &TM) {
276   return new RegAllocSimple(TM);
277 }