1 files changed, 71 insertions, 10 deletions

diff --git a/lib/CodeGen/MachineScheduler.cpp b/lib/CodeGen/MachineScheduler.cpp
index 4314e48f9f..7e871a737b 100644
--- a/lib/CodeGen/MachineScheduler.cpp
+++ b/lib/CodeGen/MachineScheduler.cpp
@@ -322,15 +322,21 @@ class ScheduleDAGMI : public ScheduleDAGInstrs {
   RegisterClassInfo *RegClassInfo;
   MachineSchedStrategy *SchedImpl;
 
+  MachineBasicBlock::iterator LiveRegionEnd;
+
   // Register pressure in this region computed by buildSchedGraph.
   IntervalPressure RegPressure;
   RegPressureTracker RPTracker;
 
   /// The top of the unscheduled zone.
   MachineBasicBlock::iterator CurrentTop;
+  IntervalPressure TopPressure;
+  RegPressureTracker TopRPTracker;
 
   /// The bottom of the unscheduled zone.
   MachineBasicBlock::iterator CurrentBottom;
+  IntervalPressure BotPressure;
+  RegPressureTracker BotRPTracker;
 
   /// The number of instructions scheduled so far. Used to cut off the
   /// scheduler at the point determined by misched-cutoff.
@@ -339,8 +345,8 @@ public:
   ScheduleDAGMI(MachineSchedContext *C, MachineSchedStrategy *S):
     ScheduleDAGInstrs(*C->MF, *C->MLI, *C->MDT, /*IsPostRA=*/false, C->LIS),
     AA(C->AA), RegClassInfo(C->RegClassInfo), SchedImpl(S),
-    RPTracker(RegPressure), CurrentTop(), CurrentBottom(),
-    NumInstrsScheduled(0) {}
+    RPTracker(RegPressure), CurrentTop(), TopRPTracker(TopPressure),
+    CurrentBottom(), BotRPTracker(BotPressure), NumInstrsScheduled(0) {}
 
   ~ScheduleDAGMI() {
     delete SchedImpl;
@@ -349,6 +355,15 @@ public:
   MachineBasicBlock::iterator top() const { return CurrentTop; }
   MachineBasicBlock::iterator bottom() const { return CurrentBottom; }
 
+  /// Get current register pressure for the top scheduled instructions.
+  const IntervalPressure &getTopPressure() const { return TopPressure; }
+
+  /// Get current register pressure for the bottom scheduled instructions.
+  const IntervalPressure &getBotPressure() const { return BotPressure; }
+
+  /// Get register pressure for the entire scheduling region before scheduling.
+  const IntervalPressure &getRegPressure() const { return RegPressure; }
+
   /// Implement the ScheduleDAGInstrs interface for handling the next scheduling
   /// region. This covers all instructions in a block, while schedule() may only
   /// cover a subset.
@@ -362,6 +377,8 @@ public:
   void schedule();
 
 protected:
+  void initRegPressure();
+
   void moveInstruction(MachineInstr *MI, MachineBasicBlock::iterator InsertPos);
   bool checkSchedLimit();
 
@@ -436,6 +453,9 @@ void ScheduleDAGMI::moveInstruction(MachineInstr *MI,
   // Fix RegionBegin if another instruction moves above the first instruction.
   if (RegionBegin == InsertPos)
     RegionBegin = MI;
+  // Fix TopRPTracker if we move something above CurrentTop.
+  if (CurrentTop == InsertPos)
+    TopRPTracker.setPos(MI);
 }
 
 bool ScheduleDAGMI::checkSchedLimit() {
@@ -459,23 +479,55 @@ void ScheduleDAGMI::enterRegion(MachineBasicBlock *bb,
                                 unsigned endcount)
 {
   ScheduleDAGInstrs::enterRegion(bb, begin, end, endcount);
-  // Setup the register pressure tracker to begin tracking at the end of this
-  // region.
-  RPTracker.init(&MF, RegClassInfo, LIS, BB, end);
+
+  // For convenience remember the end of the liveness region.
+  LiveRegionEnd =
+    (RegionEnd == bb->end()) ? RegionEnd : llvm::next(RegionEnd);
+}
+
+// Setup the register pressure trackers for the top scheduled top and bottom
+// scheduled regions.
+void ScheduleDAGMI::initRegPressure() {
+  TopRPTracker.init(&MF, RegClassInfo, LIS, BB, RegionBegin);
+  BotRPTracker.init(&MF, RegClassInfo, LIS, BB, LiveRegionEnd);
+
+  // Close the RPTracker to finalize live ins.
+  RPTracker.closeRegion();
+
+  // Initialize the live ins and live outs.
+  TopRPTracker.addLiveRegs(RPTracker.getPressure().LiveInRegs);
+  BotRPTracker.addLiveRegs(RPTracker.getPressure().LiveOutRegs);
+
+  // Close one end of the tracker so we can call
+  // getMaxUpward/DownwardPressureDelta before advancing across any
+  // instructions. This converts currently live regs into live ins/outs.
+  TopRPTracker.closeTop();
+  BotRPTracker.closeBottom();
+
+  // Account for liveness generated by the region boundary.
+  if (LiveRegionEnd != RegionEnd)
+    BotRPTracker.recede();
+
+  assert(BotRPTracker.getPos() == RegionEnd && "Can't find the region bottom");
 }
 
 /// schedule - Called back from MachineScheduler::runOnMachineFunction
 /// after setting up the current scheduling region. [RegionBegin, RegionEnd)
 /// only includes instructions that have DAG nodes, not scheduling boundaries.
 void ScheduleDAGMI::schedule() {
-  while(RPTracker.getPos() != RegionEnd) {
-    bool Moved = RPTracker.recede();
-    assert(Moved && "Regpressure tracker cannot find RegionEnd"); (void)Moved;
-  }
+  // Initialize the register pressure tracker used by buildSchedGraph.
+  RPTracker.init(&MF, RegClassInfo, LIS, BB, LiveRegionEnd);
+
+  // Account for liveness generate by the region boundary.
+  if (LiveRegionEnd != RegionEnd)
+    RPTracker.recede();
 
-  // Build the DAG.
+  // Build the DAG, and compute current register pressure.
   buildSchedGraph(AA, &RPTracker);
 
+  // Initialize top/bottom trackers after computing region pressure.
+  initRegPressure();
+
   DEBUG(dbgs() << "********** MI Scheduling **********\n");
   DEBUG(for (unsigned su = 0, e = SUnits.size(); su != e; ++su)
           SUnits[su].dumpAll(this));
@@ -517,6 +569,11 @@ void ScheduleDAGMI::schedule() {
         CurrentTop = nextIfDebug(++CurrentTop, CurrentBottom);
       else
         moveInstruction(MI, CurrentTop);
+
+      // Update top scheduled pressure.
+      TopRPTracker.advance();
+      assert(TopRPTracker.getPos() == CurrentTop && "out of sync");
+
       // Release dependent instructions for scheduling.
       releaseSuccessors(SU);
     }
@@ -532,6 +589,10 @@ void ScheduleDAGMI::schedule() {
         moveInstruction(MI, CurrentBottom);
         CurrentBottom = MI;
       }
+      // Update bottom scheduled pressure.
+      BotRPTracker.recede();
+      assert(BotRPTracker.getPos() == CurrentBottom && "out of sync");
+
       // Release dependent instructions for scheduling.
       releasePredecessors(SU);
     }