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