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