RegAllocPBQP.cpp

void RegAllocPBQP::addStackInterval(const LiveInterval *spilled,
                                    MachineRegisterInfo* mri) {
  int stackSlot = vrm->getStackSlot(spilled->reg);

  if (stackSlot == VirtRegMap::NO_STACK_SLOT)
    return;

  const TargetRegisterClass *RC = mri->getRegClass(spilled->reg);
  LiveInterval &stackInterval = lss->getOrCreateInterval(stackSlot, RC);

  VNInfo *vni;
  if (stackInterval.getNumValNums() != 0)
    vni = stackInterval.getValNumInfo(0);
  else
    vni = stackInterval.getNextValue(
      SlotIndex(), 0, lss->getVNInfoAllocator());

  LiveInterval &rhsInterval = lis->getInterval(spilled->reg);
  stackInterval.MergeRangesInAsValue(rhsInterval, vni);
}

bool RegAllocPBQP::mapPBQPToRegAllocOld(const PBQP::Solution &solution) {

  // Set to true if we have any spills
  bool anotherRoundNeeded = false;

  // Clear the existing allocation.
  vrm->clearAllVirt();

  // Iterate over the nodes mapping the PBQP solution to a register assignment.
  for (unsigned node = 0; node < node2LI.size(); ++node) {
    unsigned virtReg = node2LI[node]->reg,
             allocSelection = solution.getSelection(problemNodes[node]);


    // If the PBQP solution is non-zero it's a physical register...
    if (allocSelection != 0) {
      // Get the physical reg, subtracting 1 to account for the spill option.
      unsigned physReg = allowedSets[node][allocSelection - 1];

      DEBUG(dbgs() << "VREG " << virtReg << " -> "
            << tri->getName(physReg) << " (Option: " << allocSelection << ")\n");

      assert(physReg != 0);

      // Add to the virt reg map and update the used phys regs.
      vrm->assignVirt2Phys(virtReg, physReg);
    }
    // ...Otherwise it's a spill.
    else {

      // Make sure we ignore this virtual reg on the next round
      // of allocation
      vregsToAlloc.erase(virtReg);

      // Insert spill ranges for this live range
      const LiveInterval *spillInterval = node2LI[node];
      double oldSpillWeight = spillInterval->weight;
      SmallVector<LiveInterval*, 8> spillIs;
      rmf->rememberUseDefs(spillInterval);
      std::vector<LiveInterval*> newSpills =
        lis->addIntervalsForSpills(*spillInterval, spillIs, loopInfo, *vrm);
      addStackInterval(spillInterval, mri);
      rmf->rememberSpills(spillInterval, newSpills);

      (void) oldSpillWeight;
      DEBUG(dbgs() << "VREG " << virtReg << " -> SPILLED (Option: 0, Cost: "
                   << oldSpillWeight << ", New vregs: ");

      // Copy any newly inserted live intervals into the list of regs to
      // allocate.
      for (std::vector<LiveInterval*>::const_iterator
           itr = newSpills.begin(), end = newSpills.end();
           itr != end; ++itr) {

        assert(!(*itr)->empty() && "Empty spill range.");

        DEBUG(dbgs() << (*itr)->reg << " ");

        vregsToAlloc.insert((*itr)->reg);
      }

      DEBUG(dbgs() << ")\n");

      // We need another round if spill intervals were added.
      anotherRoundNeeded |= !newSpills.empty();
    }
  }

  return !anotherRoundNeeded;
}

bool RegAllocPBQP::mapPBQPToRegAlloc(const PBQPRAProblem &problem,
                                     const PBQP::Solution &solution) {
  // Set to true if we have any spills
  bool anotherRoundNeeded = false;

  // Clear the existing allocation.
  vrm->clearAllVirt();

  const PBQP::Graph &g = problem.getGraph();
  // Iterate over the nodes mapping the PBQP solution to a register
  // assignment.
  for (PBQP::Graph::ConstNodeItr node = g.nodesBegin(),
                                 nodeEnd = g.nodesEnd();
       node != nodeEnd; ++node) {
    unsigned vreg = problem.getVRegForNode(node);
    unsigned alloc = solution.getSelection(node);

    if (problem.isPRegOption(vreg, alloc)) {
      unsigned preg = problem.getPRegForOption(vreg, alloc);    
      DEBUG(dbgs() << "VREG " << vreg << " -> " << tri->getName(preg) << "\n");
      assert(preg != 0 && "Invalid preg selected.");
      vrm->assignVirt2Phys(vreg, preg);      
    } else if (problem.isSpillOption(vreg, alloc)) {
      vregsToAlloc.erase(vreg);
      const LiveInterval* spillInterval = &lis->getInterval(vreg);
      double oldWeight = spillInterval->weight;
      SmallVector<LiveInterval*, 8> spillIs;
      rmf->rememberUseDefs(spillInterval);
      std::vector<LiveInterval*> newSpills =
        lis->addIntervalsForSpills(*spillInterval, spillIs, loopInfo, *vrm);
      addStackInterval(spillInterval, mri);
      rmf->rememberSpills(spillInterval, newSpills);

      (void) oldWeight;
      DEBUG(dbgs() << "VREG " << vreg << " -> SPILLED (Cost: "
                   << oldWeight << ", New vregs: ");

      // Copy any newly inserted live intervals into the list of regs to
      // allocate.
      for (std::vector<LiveInterval*>::const_iterator
           itr = newSpills.begin(), end = newSpills.end();
           itr != end; ++itr) {
        assert(!(*itr)->empty() && "Empty spill range.");
        DEBUG(dbgs() << (*itr)->reg << " ");
        vregsToAlloc.insert((*itr)->reg);
      }

      DEBUG(dbgs() << ")\n");

      // We need another round if spill intervals were added.
      anotherRoundNeeded |= !newSpills.empty();
    } else {
      assert(false && "Unknown allocation option.");
    }
  }

  return !anotherRoundNeeded;
}


void RegAllocPBQP::finalizeAlloc() const {
  typedef LiveIntervals::iterator LIIterator;
  typedef LiveInterval::Ranges::const_iterator LRIterator;

  // First allocate registers for the empty intervals.
  for (RegSet::const_iterator
         itr = emptyIntervalVRegs.begin(), end = emptyIntervalVRegs.end();
         itr != end; ++itr) {
    LiveInterval *li = &lis->getInterval(*itr);

    unsigned physReg = vrm->getRegAllocPref(li->reg);

    if (physReg == 0) {
      const TargetRegisterClass *liRC = mri->getRegClass(li->reg);
      physReg = *liRC->allocation_order_begin(*mf);
    }

    vrm->assignVirt2Phys(li->reg, physReg);
  }

  // Finally iterate over the basic blocks to compute and set the live-in sets.
  SmallVector<MachineBasicBlock*, 8> liveInMBBs;
  MachineBasicBlock *entryMBB = &*mf->begin();

  for (LIIterator liItr = lis->begin(), liEnd = lis->end();
       liItr != liEnd; ++liItr) {

    const LiveInterval *li = liItr->second;
    unsigned reg = 0;

    // Get the physical register for this interval
    if (TargetRegisterInfo::isPhysicalRegister(li->reg)) {
      reg = li->reg;
    }
    else if (vrm->isAssignedReg(li->reg)) {
      reg = vrm->getPhys(li->reg);
    }
    else {
      // Ranges which are assigned a stack slot only are ignored.
      continue;
    }

    if (reg == 0) {
      // Filter out zero regs - they're for intervals that were spilled.
      continue;
    }

    // Iterate over the ranges of the current interval...
    for (LRIterator lrItr = li->begin(), lrEnd = li->end();
         lrItr != lrEnd; ++lrItr) {

      // Find the set of basic blocks which this range is live into...
      if (lis->findLiveInMBBs(lrItr->start, lrItr->end,  liveInMBBs)) {
        // And add the physreg for this interval to their live-in sets.
        for (unsigned i = 0; i < liveInMBBs.size(); ++i) {
          if (liveInMBBs[i] != entryMBB) {
            if (!liveInMBBs[i]->isLiveIn(reg)) {
              liveInMBBs[i]->addLiveIn(reg);
            }
          }
        }
        liveInMBBs.clear();
      }
    }
  }

}

bool RegAllocPBQP::runOnMachineFunction(MachineFunction &MF) {

  mf = &MF;
  tm = &mf->getTarget();
  tri = tm->getRegisterInfo();
  tii = tm->getInstrInfo();
  mri = &mf->getRegInfo(); 

  lis = &getAnalysis<LiveIntervals>();
  lss = &getAnalysis<LiveStacks>();
  loopInfo = &getAnalysis<MachineLoopInfo>();
  rmf = &getAnalysis<RenderMachineFunction>();

  vrm = &getAnalysis<VirtRegMap>();


  DEBUG(dbgs() << "PBQP Register Allocating for " << mf->getFunction()->getName() << "\n");

  // Allocator main loop:
  //
  // * Map current regalloc problem to a PBQP problem
  // * Solve the PBQP problem
  // * Map the solution back to a register allocation
  // * Spill if necessary
  //
  // This process is continued till no more spills are generated.

  // Find the vreg intervals in need of allocation.
  findVRegIntervalsToAlloc();

  // If there are non-empty intervals allocate them using pbqp.
  if (!vregsToAlloc.empty()) {

    bool pbqpAllocComplete = false;
    unsigned round = 0;

    if (!pbqpBuilder) {
      while (!pbqpAllocComplete) {
        DEBUG(dbgs() << "  PBQP Regalloc round " << round << ":\n");

        PBQP::Graph problem = constructPBQPProblemOld();
        PBQP::Solution solution =
          PBQP::HeuristicSolver<PBQP::Heuristics::Briggs>::solve(problem);

        pbqpAllocComplete = mapPBQPToRegAllocOld(solution);

        ++round;
      }
    } else {
      while (!pbqpAllocComplete) {
        DEBUG(dbgs() << "  PBQP Regalloc round " << round << ":\n");

        std::auto_ptr<PBQPRAProblem> problem =
          builder->build(mf, lis, loopInfo, vregsToAlloc);
        PBQP::Solution solution =
          PBQP::HeuristicSolver<PBQP::Heuristics::Briggs>::solve(
            problem->getGraph());

        pbqpAllocComplete = mapPBQPToRegAlloc(*problem, solution);

        ++round;
      }
    }
  }

  // Finalise allocation, allocate empty ranges.
  finalizeAlloc();

  rmf->renderMachineFunction("After PBQP register allocation.", vrm);

  vregsToAlloc.clear();
  emptyIntervalVRegs.clear();
  li2Node.clear();
  node2LI.clear();
  allowedSets.clear();
  problemNodes.clear();

  DEBUG(dbgs() << "Post alloc VirtRegMap:\n" << *vrm << "\n");

  // Run rewriter
  std::auto_ptr<VirtRegRewriter> rewriter(createVirtRegRewriter());

  rewriter->runOnMachineFunction(*mf, *vrm, lis);

  return true;
}

FunctionPass* llvm::createPBQPRegisterAllocator(
                                           std::auto_ptr<PBQPBuilder> builder) {
  return new RegAllocPBQP(builder);
}

FunctionPass* llvm::createDefaultPBQPRegisterAllocator() {
  if (pbqpCoalescing) {
    return createPBQPRegisterAllocator(
             std::auto_ptr<PBQPBuilder>(new PBQPBuilderWithCoalescing()));
  } // else
  return createPBQPRegisterAllocator(
           std::auto_ptr<PBQPBuilder>(new PBQPBuilder()));
}

#undef DEBUG_TYPE