Remove extraneous #includes, perform FIXME
[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/MachineFunction.h"
13 #include "llvm/CodeGen/MachineInstrBuilder.h"
14 #include "llvm/Target/MachineInstrInfo.h"
15 #include "llvm/Target/MRegisterInfo.h"
16 #include "llvm/Target/TargetMachine.h"
17 #include "llvm/Support/InstVisitor.h"
18 #include "Support/Statistic.h"
19
20 namespace {
21   struct RegAllocSimple : public FunctionPass {
22     TargetMachine &TM;
23     MachineBasicBlock *CurrMBB;
24     MachineFunction *MF;
25     unsigned maxOffset;
26     const MRegisterInfo *RegInfo;
27     unsigned NumBytesAllocated, ByteAlignment;
28     
29     // Maps SSA Regs => offsets on the stack where these values are stored
30     std::map<unsigned, unsigned> VirtReg2OffsetMap;
31
32     // Maps SSA Regs => physical regs
33     std::map<unsigned, unsigned> SSA2PhysRegMap;
34
35     // Maps physical register to their register classes
36     std::map<unsigned, const TargetRegisterClass*> PhysReg2RegClassMap;
37
38     // Made to combat the incorrect allocation of r2 = add r1, r1
39     std::map<unsigned, unsigned> VirtReg2PhysRegMap;
40     
41     // Maps RegClass => which index we can take a register from. Since this is a
42     // simple register allocator, when we need a register of a certain class, we
43     // just take the next available one.
44     std::map<unsigned, unsigned> RegsUsed;
45     std::map<const TargetRegisterClass*, unsigned> RegClassIdx;
46
47     RegAllocSimple(TargetMachine &tm) : TM(tm), CurrMBB(0), maxOffset(0), 
48                                         RegInfo(tm.getRegisterInfo()),
49                                         ByteAlignment(4)
50     {
51       // build reverse mapping for physReg -> register class
52       RegInfo->buildReg2RegClassMap(PhysReg2RegClassMap);
53
54       RegsUsed[RegInfo->getFramePointer()] = 1;
55       RegsUsed[RegInfo->getStackPointer()] = 1;
56
57       cleanupAfterFunction();
58     }
59
60     bool isAvailableReg(unsigned Reg) {
61       // assert(Reg < MRegisterInfo::FirstVirtualReg && "...");
62       return RegsUsed.find(Reg) == RegsUsed.end();
63     }
64
65     ///
66     unsigned allocateStackSpaceFor(unsigned VirtReg, 
67                                    const TargetRegisterClass *regClass);
68
69     /// Given size (in bytes), returns a register that is currently unused
70     /// Side effect: marks that register as being used until manually cleared
71     unsigned getFreeReg(unsigned virtualReg);
72
73     /// Returns all `borrowed' registers back to the free pool
74     void clearAllRegs() {
75         RegClassIdx.clear();
76     }
77
78     /// Invalidates any references, real or implicit, to physical registers
79     ///
80     void invalidatePhysRegs(const MachineInstr *MI) {
81       unsigned Opcode = MI->getOpcode();
82       const MachineInstrInfo &MII = TM.getInstrInfo();
83       const MachineInstrDescriptor &Desc = MII.get(Opcode);
84       const unsigned *regs = Desc.ImplicitUses;
85       while (*regs)
86         RegsUsed[*regs++] = 1;
87
88       regs = Desc.ImplicitDefs;
89       while (*regs)
90         RegsUsed[*regs++] = 1;
91     }
92
93     void cleanupAfterFunction() {
94       VirtReg2OffsetMap.clear();
95       SSA2PhysRegMap.clear();
96       NumBytesAllocated = ByteAlignment;
97     }
98
99     /// Moves value from memory into that register
100     MachineBasicBlock::iterator
101     moveUseToReg (MachineBasicBlock *MBB,
102                   MachineBasicBlock::iterator I, unsigned VirtReg,
103                   unsigned &PhysReg);
104
105     /// Saves reg value on the stack (maps virtual register to stack value)
106     MachineBasicBlock::iterator
107     saveVirtRegToStack (MachineBasicBlock *MBB,
108                         MachineBasicBlock::iterator I, unsigned VirtReg,
109                         unsigned PhysReg);
110
111     MachineBasicBlock::iterator
112     savePhysRegToStack (MachineBasicBlock *MBB,
113                         MachineBasicBlock::iterator I, unsigned PhysReg);
114
115     /// runOnFunction - Top level implementation of instruction selection for
116     /// the entire function.
117     ///
118     bool runOnMachineFunction(MachineFunction &Fn);
119
120     bool runOnFunction(Function &Fn) {
121       return runOnMachineFunction(MachineFunction::get(&Fn));
122     }
123   };
124
125 }
126
127 unsigned RegAllocSimple::allocateStackSpaceFor(unsigned VirtReg,
128                                             const TargetRegisterClass *regClass)
129 {
130   if (VirtReg2OffsetMap.find(VirtReg) == VirtReg2OffsetMap.end()) {
131 #if 0
132     unsigned size = regClass->getDataSize();
133     unsigned over = NumBytesAllocated - (NumBytesAllocated % ByteAlignment);
134     if (size >= ByteAlignment - over) {
135       // need to pad by (ByteAlignment - over)
136       NumBytesAllocated += ByteAlignment - over;
137     }
138     VirtReg2OffsetMap[VirtReg] = NumBytesAllocated;
139     NumBytesAllocated += size;
140 #endif
141     // FIXME: forcing each arg to take 4 bytes on the stack
142     VirtReg2OffsetMap[VirtReg] = NumBytesAllocated;
143     NumBytesAllocated += ByteAlignment;
144   }
145   return VirtReg2OffsetMap[VirtReg];
146 }
147
148 unsigned RegAllocSimple::getFreeReg(unsigned virtualReg) {
149   const TargetRegisterClass* regClass = MF->getRegClass(virtualReg);
150   unsigned physReg;
151   assert(regClass);
152   if (RegClassIdx.find(regClass) != RegClassIdx.end()) {
153     unsigned regIdx = RegClassIdx[regClass]++;
154     assert(regIdx < regClass->getNumRegs() && "Not enough registers!");
155     physReg = regClass->getRegister(regIdx);
156   } else {
157     physReg = regClass->getRegister(0);
158     // assert(physReg < regClass->getNumRegs() && "No registers in class!");
159     RegClassIdx[regClass] = 1;
160   }
161
162   if (isAvailableReg(physReg))
163     return physReg;
164   else {
165     return getFreeReg(virtualReg);
166   }
167 }
168
169 MachineBasicBlock::iterator
170 RegAllocSimple::moveUseToReg (MachineBasicBlock *MBB,
171                               MachineBasicBlock::iterator I,
172                               unsigned VirtReg, unsigned &PhysReg)
173 {
174   const TargetRegisterClass* regClass = MF->getRegClass(VirtReg);
175   assert(regClass);
176
177   unsigned stackOffset = allocateStackSpaceFor(VirtReg, regClass);
178   PhysReg = getFreeReg(VirtReg);
179
180   // Add move instruction(s)
181   return RegInfo->loadRegOffset2Reg(MBB, I, PhysReg,
182                                     RegInfo->getFramePointer(),
183                                     -stackOffset, regClass->getDataSize());
184 }
185
186 MachineBasicBlock::iterator
187 RegAllocSimple::saveVirtRegToStack (MachineBasicBlock *MBB,
188                                     MachineBasicBlock::iterator I,
189                                     unsigned VirtReg, unsigned PhysReg)
190 {
191   const TargetRegisterClass* regClass = MF->getRegClass(VirtReg);
192   assert(regClass);
193
194   unsigned stackOffset = allocateStackSpaceFor(VirtReg, regClass);
195
196   // Add move instruction(s)
197   return RegInfo->storeReg2RegOffset(MBB, I, PhysReg,
198                                      RegInfo->getFramePointer(),
199                                      -stackOffset, regClass->getDataSize());
200 }
201
202 MachineBasicBlock::iterator
203 RegAllocSimple::savePhysRegToStack (MachineBasicBlock *MBB,
204                                     MachineBasicBlock::iterator I,
205                                     unsigned PhysReg)
206 {
207   const TargetRegisterClass* regClass = MF->getRegClass(PhysReg);
208   assert(regClass);
209
210   unsigned offset = allocateStackSpaceFor(PhysReg, regClass);
211
212   // Add move instruction(s)
213   return RegInfo->storeReg2RegOffset(MBB, I, PhysReg,
214                                      RegInfo->getFramePointer(),
215                                      offset, regClass->getDataSize());
216 }
217
218 bool RegAllocSimple::runOnMachineFunction(MachineFunction &Fn) {
219   cleanupAfterFunction();
220
221   unsigned virtualReg, physReg;
222   DEBUG(std::cerr << "Machine Function " << "\n");
223   MF = &Fn;
224
225   for (MachineFunction::iterator MBB = Fn.begin(), MBBe = Fn.end();
226        MBB != MBBe; ++MBB)
227   {
228     CurrMBB = &(*MBB);
229
230     // Handle PHI instructions specially: add moves to each pred block
231     while (MBB->front()->getOpcode() == 0) {
232       MachineInstr *MI = MBB->front();
233       // get rid of the phi
234       MBB->erase(MBB->begin());
235     
236       // a preliminary pass that will invalidate any registers that
237       // are used by the instruction (including implicit uses)
238       invalidatePhysRegs(MI);
239
240       DEBUG(std::cerr << "num invalid regs: " << RegsUsed.size() << "\n");
241
242       DEBUG(std::cerr << "num ops: " << MI->getNumOperands() << "\n");
243       MachineOperand &targetReg = MI->getOperand(0);
244
245       // If it's a virtual register, allocate a physical one
246       // otherwise, just use whatever register is there now
247       // note: it MUST be a register -- we're assigning to it
248       virtualReg = (unsigned) targetReg.getAllocatedRegNum();
249       if (targetReg.isVirtualRegister()) {
250         physReg = getFreeReg(virtualReg);
251       } else {
252         physReg = virtualReg;
253       }
254
255       // Find the register class of the target register: should be the
256       // same as the values we're trying to store there
257       const TargetRegisterClass* regClass = PhysReg2RegClassMap[physReg];
258       assert(regClass && "Target register class not found!");
259       unsigned dataSize = regClass->getDataSize();
260
261       for (int i = MI->getNumOperands() - 1; i >= 2; i-=2) {
262         MachineOperand &opVal = MI->getOperand(i-1);
263
264         // Get the MachineBasicBlock equivalent of the BasicBlock that is the
265         // source path the phi
266         MachineBasicBlock *opBlock = MI->getOperand(i).getMachineBasicBlock();
267         MachineBasicBlock::iterator opI = opBlock->end();
268         MachineInstr *opMI = *(--opI);
269         const MachineInstrInfo &MII = TM.getInstrInfo();
270         // must backtrack over ALL the branches in the previous block, until no more
271         while ((MII.isBranch(opMI->getOpcode()) || MII.isReturn(opMI->getOpcode()))
272                && opI != opBlock->begin())
273         {
274           opMI = *(--opI);
275         }
276         // move back to the first branch instruction so new instructions
277         // are inserted right in front of it and not in front of a non-branch
278         ++opI; 
279
280
281         // Retrieve the constant value from this op, move it to target
282         // register of the phi
283         if (opVal.getType() == MachineOperand::MO_SignExtendedImmed ||
284             opVal.getType() == MachineOperand::MO_UnextendedImmed)
285         {
286           opI = RegInfo->moveImm2Reg(opBlock, opI, physReg,
287                                      (unsigned) opVal.getImmedValue(),
288                                      dataSize);
289           saveVirtRegToStack(opBlock, opI, virtualReg, physReg);
290         } else {
291           // Allocate a physical register and add a move in the BB
292           unsigned opVirtualReg = (unsigned) opVal.getAllocatedRegNum();
293           unsigned opPhysReg; // = getFreeReg(opVirtualReg);
294           opI = moveUseToReg(opBlock, opI, opVirtualReg, physReg);
295           //opI = RegInfo->moveReg2Reg(opBlock, opI, physReg, opPhysReg,
296           //                           dataSize);
297           // Save that register value to the stack of the TARGET REG
298           saveVirtRegToStack(opBlock, opI, virtualReg, physReg);
299         }
300
301         // make regs available to other instructions
302         clearAllRegs();
303       }
304       
305       // really delete the instruction
306       delete MI;
307     }
308
309     //loop over each basic block
310     for (MachineBasicBlock::iterator I = MBB->begin(); I != MBB->end(); ++I)
311     {
312       MachineInstr *MI = *I;
313
314       // a preliminary pass that will invalidate any registers that
315       // are used by the instruction (including implicit uses)
316       invalidatePhysRegs(MI);
317
318       // Loop over uses, move from memory into registers
319       for (int i = MI->getNumOperands() - 1; i >= 0; --i) {
320         MachineOperand &op = MI->getOperand(i);
321
322         if (op.getType() == MachineOperand::MO_SignExtendedImmed ||
323             op.getType() == MachineOperand::MO_UnextendedImmed)
324         {
325           DEBUG(std::cerr << "const\n");
326         } else if (op.isVirtualRegister()) {
327           virtualReg = (unsigned) op.getAllocatedRegNum();
328           DEBUG(std::cerr << "op: " << op << "\n");
329           DEBUG(std::cerr << "\t inst[" << i << "]: ";
330                 MI->print(std::cerr, TM));
331
332           // make sure the same virtual register maps to the same physical
333           // register in any given instruction
334           if (VirtReg2PhysRegMap.find(virtualReg) != VirtReg2PhysRegMap.end()) {
335             physReg = VirtReg2PhysRegMap[virtualReg];
336           } else {
337             if (op.opIsDef()) {
338               if (TM.getInstrInfo().isTwoAddrInstr(MI->getOpcode()) && i == 0) {
339                 // must be same register number as the first operand
340                 // This maps a = b + c into b += c, and saves b into a's spot
341                 physReg = (unsigned) MI->getOperand(1).getAllocatedRegNum();
342               } else {
343                 physReg = getFreeReg(virtualReg);
344               }
345               MachineBasicBlock::iterator J = I;
346               J = saveVirtRegToStack(CurrMBB, ++J, virtualReg, physReg);
347               I = --J;
348             } else {
349               I = moveUseToReg(CurrMBB, I, virtualReg, physReg);
350             }
351             VirtReg2PhysRegMap[virtualReg] = physReg;
352           }
353           MI->SetMachineOperandReg(i, physReg);
354           DEBUG(std::cerr << "virt: " << virtualReg << 
355                 ", phys: " << op.getAllocatedRegNum() << "\n");
356         }
357       }
358
359       clearAllRegs();
360       VirtReg2PhysRegMap.clear();
361     }
362
363   }
364
365   // add prologue we should preserve callee-save registers...
366   MachineFunction::iterator Fi = Fn.begin();
367   MachineBasicBlock *MBB = Fi;
368   MachineBasicBlock::iterator MBBi = MBB->begin();
369   RegInfo->emitPrologue(MBB, MBBi, NumBytesAllocated);
370
371   // add epilogue to restore the callee-save registers
372   // loop over the basic block
373   for (MachineFunction::iterator MBB = Fn.begin(), MBBe = Fn.end();
374        MBB != MBBe; ++MBB)
375   {
376     // check if last instruction is a RET
377     MachineBasicBlock::iterator I = (*MBB).end();
378     MachineInstr *MI = *(--I);
379     const MachineInstrInfo &MII = TM.getInstrInfo();
380     if (MII.isReturn(MI->getOpcode())) {
381       // this block has a return instruction, add epilogue
382       RegInfo->emitEpilogue(MBB, I, NumBytesAllocated);
383     }
384   }
385
386   return false;  // We never modify the LLVM itself.
387 }
388
389 Pass *createSimpleX86RegisterAllocator(TargetMachine &TM) {
390   return new RegAllocSimple(TM);
391 }