[AMDGPU] Start refactoring GCNSchedStrategy

[kerbowa]

Summary

Tries to make the different scheduling stages a bit more self contained and
modifiable. Intended to be NFC. Preface to other changes.

[nhaehnle]

This is a big diff that's kind of difficult to review, but I don't see anything obviously wrong with it and I like the overall direction.

llvm/lib/Target/AMDGPU/GCNSchedStrategy.h

43

Maybe remove the extra whitespace?

[rampitec]

Nice refactoring! A lot of code, but indeed looks like NFC to me.

[kerbowa]

Minor fixes.

[rampitec]

Still LGTM.

llvm/lib/Target/AMDGPU/GCNSchedStrategy.h

 class SIRegisterInfo;
 class GCNSubtarget;
 
 /// This is a minimal scheduler strategy.  The main difference between this
 /// and the GenericScheduler is that GCNSchedStrategy uses different
 /// heuristics to determine excess/critical pressure sets.  Its goal is to
 /// maximize kernel occupancy (i.e. maximum number of waves per simd).
 class GCNMaxOccupancySchedStrategy final : public GenericScheduler {
-  friend class GCNScheduleDAGMILive;
-
   SUnit *pickNodeBidirectional(bool &IsTopNode);
 
   void pickNodeFromQueue(SchedBoundary &Zone, const CandPolicy &ZonePolicy,
                          const RegPressureTracker &RPTracker,
                          SchedCandidate &Cand);
 
   void initCandidate(SchedCandidate &Cand, SUnit *SU,
                      bool AtTop, const RegPressureTracker &RPTracker,
                      const SIRegisterInfo *SRI,
                      unsigned SGPRPressure, unsigned VGPRPressure);
 
   std::vector<unsigned> Pressure;
+
   std::vector<unsigned> MaxPressure;
 
   unsigned SGPRExcessLimit;
+
   unsigned VGPRExcessLimit;
-  unsigned SGPRCriticalLimit;
-  unsigned VGPRCriticalLimit;
 
   unsigned TargetOccupancy;
 
+  MachineFunction *MF;
+
+public:
   // schedule() have seen a clustered memory operation. Set it to false
   // before a region scheduling to know if the region had such clusters.
   bool HasClusteredNodes;
 
   // schedule() have seen an excess register pressure and had to track
   // register pressure for actual scheduling heuristics.
   bool HasExcessPressure;
 
-  MachineFunction *MF;
+  unsigned SGPRCriticalLimit;
+
+  unsigned VGPRCriticalLimit;
 
-public:
   GCNMaxOccupancySchedStrategy(const MachineSchedContext *C);
 
   SUnit *pickNode(bool &IsTopNode) override;
 
   void initialize(ScheduleDAGMI *DAG) override;
 
+  unsigned getTargetOccupancy() { return TargetOccupancy; }
+
   void setTargetOccupancy(unsigned Occ) { TargetOccupancy = Occ; }
 };
 
-class GCNScheduleDAGMILive final : public ScheduleDAGMILive {
-
-  enum : unsigned {
-    Collect,
-    InitialSchedule,
-    UnclusteredReschedule,
-    ClusteredLowOccupancyReschedule,
-    PreRARematerialize,
-    LastStage = PreRARematerialize
-  };
+enum class GCNSchedStageID : unsigned {
+  InitialSchedule = 0,
+  UnclusteredReschedule = 1,
+  ClusteredLowOccupancyReschedule = 2,
+  PreRARematerialize = 3,
 
+#ifndef NDEBUG
+raw_ostream &operator<<(raw_ostream &OS, const GCNSchedStageID &StageID);
+#endif
+
+inline GCNSchedStageID &operator++(GCNSchedStageID &Stage, int) {
+  assert(Stage != GCNSchedStageID::PreRARematerialize);
+  Stage = static_cast<GCNSchedStageID>(static_cast<unsigned>(Stage) + 1);
+  return Stage;
+}
+
+inline GCNSchedStageID nextStage(const GCNSchedStageID Stage) {
+  return static_cast<GCNSchedStageID>(static_cast<unsigned>(Stage) + 1);
+}
+
+inline bool operator>(GCNSchedStageID &LHS, GCNSchedStageID &RHS) {
+  return static_cast<unsigned>(LHS) > static_cast<unsigned>(RHS);
+}
+
+class GCNScheduleDAGMILive final : public ScheduleDAGMILive {
+  friend class GCNSchedStage;
+  friend class InitialScheduleStage;
+  friend class UnclusteredRescheduleStage;
+  friend class ClusteredLowOccStage;
+  friend class PreRARematStage;
+
   const GCNSubtarget &ST;
 
   SIMachineFunctionInfo &MFI;
 
   // Occupancy target at the beginning of function scheduling cycle.
   unsigned StartingOccupancy;
 
   // Minimal real occupancy recorder for the function.
   unsigned MinOccupancy;
 
-  // Scheduling stage number.
-  unsigned Stage;
-
-  // Current region index.
-  size_t RegionIdx;
-
   // Vector of regions recorder for later rescheduling
   SmallVector<std::pair<MachineBasicBlock::iterator,
                         MachineBasicBlock::iterator>, 32> Regions;
 
   // Records if a region is not yet scheduled, or schedule has been reverted,
   // or we generally desire to reschedule it.
   BitVector RescheduleRegions;
 
   // Record regions which use clustered loads/stores.
   BitVector RegionsWithClusters;
 
   // Record regions with high register pressure.
   BitVector RegionsWithHighRP;
 
   // Regions that has the same occupancy as the latest MinOccupancy
   BitVector RegionsWithMinOcc;
 
   // Region live-in cache.
   SmallVector<GCNRPTracker::LiveRegSet, 32> LiveIns;
 
   // Region pressure cache.
   SmallVector<GCNRegPressure, 32> Pressure;
 
+  // Temporary basic block live-in cache.
+  DenseMap<const MachineBasicBlock *, GCNRPTracker::LiveRegSet> MBBLiveIns;
+
+  DenseMap<MachineInstr *, GCNRPTracker::LiveRegSet> BBLiveInMap;
+
+  DenseMap<MachineInstr *, GCNRPTracker::LiveRegSet> getBBLiveInMap() const;
+
+  // Return current region pressure.
+  GCNRegPressure getRealRegPressure(unsigned RegionIdx) const;
+
+  // Compute and cache live-ins and pressure for all regions in block.
+  void computeBlockPressure(unsigned RegionIdx, const MachineBasicBlock *MBB);
+
+  // Update region boundaries when removing MI or inserting NewMI before MI.
+  void updateRegionBoundaries(
+      SmallVectorImpl<std::pair<MachineBasicBlock::iterator,
+                                MachineBasicBlock::iterator>> &RegionBoundaries,
+      MachineBasicBlock::iterator MI, MachineInstr *NewMI,
+      bool Removing = false);
+
+  void runSchedStages();
+
+public:
+  GCNScheduleDAGMILive(MachineSchedContext *C,
+                       std::unique_ptr<MachineSchedStrategy> S);
+
+  void schedule() override;
+
+  void finalizeSchedule() override;
+};
+
+// GCNSchedStrategy applies multiple scheduling stages to a function.
+class GCNSchedStage {
+protected:
+  GCNScheduleDAGMILive &DAG;
+
+  GCNMaxOccupancySchedStrategy &S;
+
+  MachineFunction &MF;
+
+  SIMachineFunctionInfo &MFI;
+
+  const GCNSubtarget &ST;
+
+  const GCNSchedStageID StageID;
+
+  // The current block being scheduled.
+  MachineBasicBlock *CurrentMBB = nullptr;
+
+  // Current region index.
+  unsigned RegionIdx = 0;
+
+  // Record the original order of instructions before scheduling.
+  std::vector<MachineInstr *> Unsched;
+
+  // RP before scheduling the current region.
+  GCNRegPressure PressureBefore;
+
+  // RP after scheduling the current region.
+  GCNRegPressure PressureAfter;
+
+  GCNSchedStage(GCNSchedStageID StageID, GCNScheduleDAGMILive &DAG);
+
+public:
+  // Initialize state for a scheduling stage. Returns false if the current stage
+  // should be skipped.
+  virtual bool initGCNSchedStage();
+
+  // Finalize state after finishing a scheduling pass on the function.
+  virtual void finalizeGCNSchedStage();
+
+  // Setup for scheduling a region. Returns false if the current region should
+  // be skipped.
+  virtual bool initGCNRegion();
+
+  // Track whether a new region is also a new MBB.
+  void setupNewBlock();
+
+  // Finalize state after scheudling a region.
+  virtual void finalizeGCNRegion();
+
+  // Check result of scheduling.
+  void checkScheduling();
+
+  // Returns true if scheduling should be reverted.
+  virtual bool shouldRevertScheduling(unsigned WavesAfter);
+
+  // Returns true if the new schedule may result in more spilling.
+  bool mayCauseSpilling(unsigned WavesAfter);
+
+  // Attempt to revert scheduling for this region.
+  void revertScheduling();
+
+  void advanceRegion() { RegionIdx++; }
+
+  virtual ~GCNSchedStage() = default;
+};
+
+class InitialScheduleStage : public GCNSchedStage {
+public:
+  void finalizeGCNRegion() override;
+
+  bool shouldRevertScheduling(unsigned WavesAfter) override;
+
+  InitialScheduleStage(GCNSchedStageID StageID, GCNScheduleDAGMILive &DAG)
+      : GCNSchedStage(StageID, DAG) {}
+};
+
+class UnclusteredRescheduleStage : public GCNSchedStage {
+private:
+  std::vector<std::unique_ptr<ScheduleDAGMutation>> SavedMutations;
+
+public:
+  bool initGCNSchedStage() override;
+
+  void finalizeGCNSchedStage() override;
+
+  bool initGCNRegion() override;
+
+  bool shouldRevertScheduling(unsigned WavesAfter) override;
+
+  UnclusteredRescheduleStage(GCNSchedStageID StageID, GCNScheduleDAGMILive &DAG)
+      : GCNSchedStage(StageID, DAG) {}
+};
+
+// Retry function scheduling if we found resulting occupancy and it is
+// lower than used for other scheduling passes. This will give more freedom
+// to schedule low register pressure blocks.
+class ClusteredLowOccStage : public GCNSchedStage {
+public:
+  bool initGCNSchedStage() override;
+
+  bool initGCNRegion() override;
+
+  bool shouldRevertScheduling(unsigned WavesAfter) override;
+
+  ClusteredLowOccStage(GCNSchedStageID StageID, GCNScheduleDAGMILive &DAG)
+      : GCNSchedStage(StageID, DAG) {}
+};
+
+class PreRARematStage : public GCNSchedStage {
+private:
   // Each region at MinOccupancy will have their own list of trivially
   // rematerializable instructions we can remat to reduce RP. The list maps an
   // instruction to the position we should remat before, usually the MI using
   // the rematerializable instruction.
   MapVector<unsigned, MapVector<MachineInstr *, MachineInstr *>>
       RematerializableInsts;
 
   // Map a trivially remateriazable def to a list of regions at MinOccupancy
   // that has the defined reg as a live-in.
   DenseMap<MachineInstr *, SmallVector<unsigned, 4>> RematDefToLiveInRegions;
 
-  // Temporary basic block live-in cache.
-  DenseMap<const MachineBasicBlock*, GCNRPTracker::LiveRegSet> MBBLiveIns;
-
-  DenseMap<MachineInstr *, GCNRPTracker::LiveRegSet> BBLiveInMap;
-  DenseMap<MachineInstr *, GCNRPTracker::LiveRegSet> getBBLiveInMap() const;
-
   // Collect all trivially rematerializable VGPR instructions with a single def
   // and single use outside the defining block into RematerializableInsts.
   void collectRematerializableInstructions();
 
   bool isTriviallyReMaterializable(const MachineInstr &MI);
 
   // TODO: Should also attempt to reduce RP of SGPRs and AGPRs
   // Attempt to reduce RP of VGPR by sinking trivially rematerializable
   // instructions. Returns true if we were able to sink instruction(s).
   bool sinkTriviallyRematInsts(const GCNSubtarget &ST,
                                const TargetInstrInfo *TII);
 
-  // Return current region pressure.
-  GCNRegPressure getRealRegPressure() const;
-
-  // Compute and cache live-ins and pressure for all regions in block.
-  void computeBlockPressure(const MachineBasicBlock *MBB);
-
-  // Update region boundaries when removing MI or inserting NewMI before MI.
-  void updateRegionBoundaries(
-      SmallVectorImpl<std::pair<MachineBasicBlock::iterator,
-                                MachineBasicBlock::iterator>> &RegionBoundaries,
-      MachineBasicBlock::iterator MI, MachineInstr *NewMI,
-      bool Removing = false);
-
 public:
-  GCNScheduleDAGMILive(MachineSchedContext *C,
-                       std::unique_ptr<MachineSchedStrategy> S);
+  bool initGCNSchedStage() override;
 
-  void schedule() override;
+  bool initGCNRegion() override;
 
-  void finalizeSchedule() override;
+  bool shouldRevertScheduling(unsigned WavesAfter) override;
+
+  PreRARematStage(GCNSchedStageID StageID, GCNScheduleDAGMILive &DAG)
+      : GCNSchedStage(StageID, DAG) {}
 };
 
 } // End namespace llvm
 
 #endif // LLVM_LIB_TARGET_AMDGPU_GCNSCHEDSTRATEGY_H

llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp

 //===-- GCNSchedStrategy.cpp - GCN Scheduler Strategy ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 /// \file
 /// This contains a MachineSchedStrategy implementation for maximizing wave
 /// occupancy on GCN hardware.
+///
+/// This pass will apply multiple scheduling stages to the same function.
+/// Regions are first recorded in GCNScheduleDAGMILive::schedule. The actual
+/// entry point for the scheduling of those regions is
+/// GCNScheduleDAGMILive::runSchedStages.
+
+/// Generally, the reason for having multiple scheduling stages is to account
+/// for the kernel-wide effect of register usage on occupancy.  Usually, only a
+/// few scheduling regions will have register pressure high enough to limit
+/// occupancy for the kernel, so constraints can be relaxed to improve ILP in
+/// other regions.
+///
 //===----------------------------------------------------------------------===//
 
 #include "GCNSchedStrategy.h"
 #include "SIMachineFunctionInfo.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
 
 #define DEBUG_TYPE "machine-scheduler"
 
 using namespace llvm;
 
 GCNMaxOccupancySchedStrategy::GCNMaxOccupancySchedStrategy(
-    const MachineSchedContext *C) :
-    GenericScheduler(C), TargetOccupancy(0), HasClusteredNodes(false),
-    HasExcessPressure(false), MF(nullptr) { }
+    const MachineSchedContext *C)
+    : GenericScheduler(C), TargetOccupancy(0), MF(nullptr),
+      HasClusteredNodes(false), HasExcessPressure(false) {}
 
 void GCNMaxOccupancySchedStrategy::initialize(ScheduleDAGMI *DAG) {
   GenericScheduler::initialize(DAG);
 
   MF = &DAG->MF;
 
   const GCNSubtarget &ST = MF->getSubtarget<GCNSubtarget>();
 
if (!HasClusteredNodes && SU->getInstr()->mayLoadOrStore()) {
     }
   }
 
   LLVM_DEBUG(dbgs() << "Scheduling SU(" << SU->NodeNum << ") "
                     << *SU->getInstr());
   return SU;
 }
 
-GCNScheduleDAGMILive::GCNScheduleDAGMILive(MachineSchedContext *C,
-                        std::unique_ptr<MachineSchedStrategy> S) :
-  ScheduleDAGMILive(C, std::move(S)),
-  ST(MF.getSubtarget<GCNSubtarget>()),
+GCNScheduleDAGMILive::GCNScheduleDAGMILive(
+    MachineSchedContext *C, std::unique_ptr<MachineSchedStrategy> S)
+    : ScheduleDAGMILive(C, std::move(S)), ST(MF.getSubtarget<GCNSubtarget>()),
       MFI(*MF.getInfo<SIMachineFunctionInfo>()),
-  StartingOccupancy(MFI.getOccupancy()),
-  MinOccupancy(StartingOccupancy), Stage(Collect), RegionIdx(0) {
+      StartingOccupancy(MFI.getOccupancy()), MinOccupancy(StartingOccupancy) {
 
   LLVM_DEBUG(dbgs() << "Starting occupancy is " << StartingOccupancy << ".\n");
 }
 
 void GCNScheduleDAGMILive::schedule() {
-  if (Stage == Collect) {
-    // Just record regions at the first pass.
-    Regions.push_back(std::make_pair(RegionBegin, RegionEnd));
-    return;
-  }
-
-  std::vector<MachineInstr*> Unsched;
-  Unsched.reserve(NumRegionInstrs);
-  for (auto &I : *this) {
-    Unsched.push_back(&I);
-  }
-
-  GCNRegPressure PressureBefore;
-  if (LIS) {
-    PressureBefore = Pressure[RegionIdx];
-
-    LLVM_DEBUG(dbgs() << "Pressure before scheduling:\nRegion live-ins:";
-               GCNRPTracker::printLiveRegs(dbgs(), LiveIns[RegionIdx], MRI);
-               dbgs() << "Region live-in pressure:  ";
-               llvm::getRegPressure(MRI, LiveIns[RegionIdx]).print(dbgs());
-               dbgs() << "Region register pressure: ";
-               PressureBefore.print(dbgs()));
-  }
-
-  GCNMaxOccupancySchedStrategy &S = (GCNMaxOccupancySchedStrategy&)*SchedImpl;
-  // Set HasClusteredNodes to true for late stages where we have already
-  // collected it. That way pickNode() will not scan SDep's when not needed.
-  S.HasClusteredNodes = Stage > InitialSchedule;
-  S.HasExcessPressure = false;
-  ScheduleDAGMILive::schedule();
-  Regions[RegionIdx] = std::make_pair(RegionBegin, RegionEnd);
-  RescheduleRegions[RegionIdx] = false;
-  if (Stage == InitialSchedule && S.HasClusteredNodes)
-    RegionsWithClusters[RegionIdx] = true;
-  if (S.HasExcessPressure)
-    RegionsWithHighRP[RegionIdx] = true;
-
-  if (!LIS)
-    return;
-
-  // Check the results of scheduling.
-  auto PressureAfter = getRealRegPressure();
-
-  LLVM_DEBUG(dbgs() << "Pressure after scheduling: ";
-             PressureAfter.print(dbgs()));
-
-  if (PressureAfter.getSGPRNum() <= S.SGPRCriticalLimit &&
-      PressureAfter.getVGPRNum(ST.hasGFX90AInsts()) <= S.VGPRCriticalLimit) {
-    Pressure[RegionIdx] = PressureAfter;
-    RegionsWithMinOcc[RegionIdx] =
-        PressureAfter.getOccupancy(ST) == MinOccupancy;
-
-    LLVM_DEBUG(dbgs() << "Pressure in desired limits, done.\n");
-    return;
-  }
-
-  unsigned WavesAfter =
-      std::min(S.TargetOccupancy, PressureAfter.getOccupancy(ST));
-  unsigned WavesBefore =
-      std::min(S.TargetOccupancy, PressureBefore.getOccupancy(ST));
-  LLVM_DEBUG(dbgs() << "Occupancy before scheduling: " << WavesBefore
-                    << ", after " << WavesAfter << ".\n");
-
-  // We may not be able to keep the current target occupancy because of the just
-  // scheduled region. We might still be able to revert scheduling if the
-  // occupancy before was higher, or if the current schedule has register
-  // pressure higher than the excess limits which could lead to more spilling.
-  unsigned NewOccupancy = std::max(WavesAfter, WavesBefore);
-
-  // Allow memory bound functions to drop to 4 waves if not limited by an
-  // attribute.
-  if (WavesAfter < WavesBefore && WavesAfter < MinOccupancy &&
-      WavesAfter >= MFI.getMinAllowedOccupancy()) {
-    LLVM_DEBUG(dbgs() << "Function is memory bound, allow occupancy drop up to "
-                      << MFI.getMinAllowedOccupancy() << " waves\n");
-    NewOccupancy = WavesAfter;
-  }
-
-  if (NewOccupancy < MinOccupancy) {
-    MinOccupancy = NewOccupancy;
-    MFI.limitOccupancy(MinOccupancy);
-    RegionsWithMinOcc.reset();
-    LLVM_DEBUG(dbgs() << "Occupancy lowered for the function to "
-                      << MinOccupancy << ".\n");
-  }
-
-  unsigned MaxVGPRs = ST.getMaxNumVGPRs(MF);
-  unsigned MaxSGPRs = ST.getMaxNumSGPRs(MF);
-  if (PressureAfter.getVGPRNum(false) > MaxVGPRs ||
-      PressureAfter.getAGPRNum() > MaxVGPRs ||
-      PressureAfter.getSGPRNum() > MaxSGPRs) {
-    RescheduleRegions[RegionIdx] = true;
-    RegionsWithHighRP[RegionIdx] = true;
-  }
-
-  // If this condition is true, then either the occupancy before and after
-  // scheduling is the same, or we are allowing the occupancy to drop because
-  // the function is memory bound. Even if we are OK with the current occupancy,
-  // we still need to verify that we will not introduce any extra chance of
-  // spilling.
-  if (WavesAfter >= MinOccupancy) {
-    if (Stage == UnclusteredReschedule &&
-        !PressureAfter.less(ST, PressureBefore)) {
-      LLVM_DEBUG(dbgs() << "Unclustered reschedule did not help.\n");
-    } else if (WavesAfter > MFI.getMinWavesPerEU() ||
-        PressureAfter.less(ST, PressureBefore) ||
-        !RescheduleRegions[RegionIdx]) {
-      Pressure[RegionIdx] = PressureAfter;
-      RegionsWithMinOcc[RegionIdx] =
-          PressureAfter.getOccupancy(ST) == MinOccupancy;
-      if (!RegionsWithClusters[RegionIdx] &&
-          (Stage + 1) == UnclusteredReschedule)
-        RescheduleRegions[RegionIdx] = false;
-      return;
-    } else {
-      LLVM_DEBUG(dbgs() << "New pressure will result in more spilling.\n");
-    }
-  }
-
-  RegionsWithMinOcc[RegionIdx] =
-      PressureBefore.getOccupancy(ST) == MinOccupancy;
-  LLVM_DEBUG(dbgs() << "Attempting to revert scheduling.\n");
-  RescheduleRegions[RegionIdx] = RegionsWithClusters[RegionIdx] ||
-                                 (Stage + 1) != UnclusteredReschedule;
-  RegionEnd = RegionBegin;
-  int SkippedDebugInstr = 0;
-  for (MachineInstr *MI : Unsched) {
-    if (MI->isDebugInstr()) {
-      ++SkippedDebugInstr;
-      continue;
-    }
-
-    if (MI->getIterator() != RegionEnd) {
-      BB->remove(MI);
-      BB->insert(RegionEnd, MI);
-      if (!MI->isDebugInstr())
-        LIS->handleMove(*MI, true);
-    }
-    // Reset read-undef flags and update them later.
-    for (auto &Op : MI->operands())
-      if (Op.isReg() && Op.isDef())
-        Op.setIsUndef(false);
-    RegisterOperands RegOpers;
-    RegOpers.collect(*MI, *TRI, MRI, ShouldTrackLaneMasks, false);
-    if (!MI->isDebugInstr()) {
-      if (ShouldTrackLaneMasks) {
-        // Adjust liveness and add missing dead+read-undef flags.
-        SlotIndex SlotIdx = LIS->getInstructionIndex(*MI).getRegSlot();
-        RegOpers.adjustLaneLiveness(*LIS, MRI, SlotIdx, MI);
-      } else {
-        // Adjust for missing dead-def flags.
-        RegOpers.detectDeadDefs(*MI, *LIS);
-      }
-    }
-    RegionEnd = MI->getIterator();
-    ++RegionEnd;
-    LLVM_DEBUG(dbgs() << "Scheduling " << *MI);
-  }
-
-  // After reverting schedule, debug instrs will now be at the end of the block
-  // and RegionEnd will point to the first debug instr. Increment RegionEnd
-  // pass debug instrs to the actual end of the scheduling region.
-  while (SkippedDebugInstr-- > 0)
-    ++RegionEnd;
-
-  // If Unsched.front() instruction is a debug instruction, this will actually
-  // shrink the region since we moved all debug instructions to the end of the
-  // block. Find the first instruction that is not a debug instruction.
-  RegionBegin = Unsched.front()->getIterator();
-  if (RegionBegin->isDebugInstr()) {
-    for (MachineInstr *MI : Unsched) {
-      if (MI->isDebugInstr())
-        continue;
-      RegionBegin = MI->getIterator();
-      break;
-    }
-  }
+  // Collect all scheduling regions. The actual scheduling is performed in
+  // GCNScheduleDAGMILive::finalizeSchedule.
 
-  // Then move the debug instructions back into their correct place and set
-  // RegionBegin and RegionEnd if needed.
-  placeDebugValues();
-
-  Regions[RegionIdx] = std::make_pair(RegionBegin, RegionEnd);
-}
-
-GCNRegPressure GCNScheduleDAGMILive::getRealRegPressure() const {
+GCNRegPressure
+GCNScheduleDAGMILive::getRealRegPressure(unsigned RegionIdx) const {
   GCNDownwardRPTracker RPTracker(*LIS);
   RPTracker.advance(begin(), end(), &LiveIns[RegionIdx]);
   return RPTracker.moveMaxPressure();
 }
 
-void GCNScheduleDAGMILive::computeBlockPressure(const MachineBasicBlock *MBB) {
+void GCNScheduleDAGMILive::computeBlockPressure(unsigned RegionIdx,
+                                                const MachineBasicBlock *MBB) {
   GCNDownwardRPTracker RPTracker(*LIS);
 
   // If the block has the only successor then live-ins of that successor are
   // live-outs of the current block. We can reuse calculated live set if the
   // successor will be sent to scheduling past current block.
   const MachineBasicBlock *OnlySucc = nullptr;
   if (MBB->succ_size() == 1 && !(*MBB->succ_begin())->empty()) {
     SlotIndexes *Ind = LIS->getSlotIndexes();
if (LiveInIt != MBBLiveIns.end()) {
     I = Rgn.first;
     auto LRS = BBLiveInMap.lookup(NonDbgMI);
 #ifdef EXPENSIVE_CHECKS
     assert(isEqual(getLiveRegsBefore(*NonDbgMI, *LIS), LRS));
 #endif
     RPTracker.reset(*I, &LRS);
   }
 
-  for ( ; ; ) {
+  for (;;) {
     I = RPTracker.getNext();
 
     if (Regions[CurRegion].first == I || NonDbgMI == I) {
       LiveIns[CurRegion] = RPTracker.getLiveRegs();
       RPTracker.clearMaxPressure();
     }
 
     if (Regions[CurRegion].second == I) {
do {
     do {
       ++I;
     } while (I != E && I->first->getParent() == BB);
   } while (I != E);
   return getLiveRegMap(BBStarters, false /*After*/, *LIS);
 }
 
 void GCNScheduleDAGMILive::finalizeSchedule() {
-  LLVM_DEBUG(dbgs() << "All regions recorded, starting actual scheduling.\n");
-
+  // Start actual scheduling here. This function is called by the base
+  // MachineScheduler after all regions have been recorded by
+  // GCNScheduleDAGMILive::schedule().
   LiveIns.resize(Regions.size());
   Pressure.resize(Regions.size());
   RescheduleRegions.resize(Regions.size());
   RegionsWithClusters.resize(Regions.size());
   RegionsWithHighRP.resize(Regions.size());
   RegionsWithMinOcc.resize(Regions.size());
   RescheduleRegions.set();
   RegionsWithClusters.reset();
   RegionsWithHighRP.reset();
   RegionsWithMinOcc.reset();
 
+  runSchedStages();
+}
+
+void GCNScheduleDAGMILive::runSchedStages() {
+  LLVM_DEBUG(dbgs() << "All regions recorded, starting actual scheduling.\n");
+  InitialScheduleStage S0(GCNSchedStageID::InitialSchedule, *this);
+  UnclusteredRescheduleStage S1(GCNSchedStageID::UnclusteredReschedule, *this);
+  ClusteredLowOccStage S2(GCNSchedStageID::ClusteredLowOccupancyReschedule,
+                          *this);
+  PreRARematStage S3(GCNSchedStageID::PreRARematerialize, *this);
+  GCNSchedStage *SchedStages[] = {&S0, &S1, &S2, &S3};
+
   if (!Regions.empty())
     BBLiveInMap = getBBLiveInMap();
 
-  std::vector<std::unique_ptr<ScheduleDAGMutation>> SavedMutations;
+  for (auto *Stage : SchedStages) {
+    if (!Stage->initGCNSchedStage())
+      continue;
 
-  do {
-    Stage++;
-    RegionIdx = 0;
-    MachineBasicBlock *MBB = nullptr;
+    for (auto Region : Regions) {
+      RegionBegin = Region.first;
+      RegionEnd = Region.second;
+      // Setup for scheduling the region and check whether it should be skipped.
+      if (!Stage->initGCNRegion()) {
+        Stage->advanceRegion();
+        exitRegion();
+        continue;
+      }
 
-    if (Stage > InitialSchedule) {
-      if (!LIS)
-        break;
+      ScheduleDAGMILive::schedule();
+      Stage->finalizeGCNRegion();
+    }
+
+    Stage->finalizeGCNSchedStage();
+  }
+}
+
+#ifndef NDEBUG
+raw_ostream &llvm::operator<<(raw_ostream &OS, const GCNSchedStageID &StageID) {
+  switch (StageID) {
+  case GCNSchedStageID::InitialSchedule:
+    OS << "Initial Schedule";
+    break;
+  case GCNSchedStageID::UnclusteredReschedule:
+    OS << "Unclustered Reschedule";
+    break;
+  case GCNSchedStageID::ClusteredLowOccupancyReschedule:
+    OS << "Clustered Low Occupancy Reschedule";
+    break;
+  case GCNSchedStageID::PreRARematerialize:
+    OS << "Pre-RA Rematerialize";
+  }
+  return OS;
+}
+#endif
 
-      // Retry function scheduling if we found resulting occupancy and it is
-      // lower than used for first pass scheduling. This will give more freedom
-      // to schedule low register pressure blocks.
-      // Code is partially copied from MachineSchedulerBase::scheduleRegions().
+GCNSchedStage::GCNSchedStage(GCNSchedStageID StageID, GCNScheduleDAGMILive &DAG)
+    : DAG(DAG), S(static_cast<GCNMaxOccupancySchedStrategy &>(*DAG.SchedImpl)),
+      MF(DAG.MF), MFI(DAG.MFI), ST(DAG.ST), StageID(StageID) {}
 
-      if (Stage == UnclusteredReschedule) {
-        if (RescheduleRegions.none())
-          continue;
-        LLVM_DEBUG(dbgs() <<
-          "Retrying function scheduling without clustering.\n");
-      }
+bool GCNSchedStage::initGCNSchedStage() {
+  if (!DAG.LIS)
+    return false;
+
+  LLVM_DEBUG(dbgs() << "Starting scheduling stage: " << StageID << "\n");
+  return true;
 
-      if (Stage == ClusteredLowOccupancyReschedule) {
-        if (StartingOccupancy <= MinOccupancy)
-          break;
+bool UnclusteredRescheduleStage::initGCNSchedStage() {
+  if (!GCNSchedStage::initGCNSchedStage())
+    return false;
+
+  if (DAG.RescheduleRegions.none())
+    return false;
+
+  SavedMutations.swap(DAG.Mutations);
+
+  LLVM_DEBUG(dbgs() << "Retrying function scheduling without clustering.\n");
+  return true;
+}
+
+bool ClusteredLowOccStage::initGCNSchedStage() {
+  if (!GCNSchedStage::initGCNSchedStage())
+    return false;
+
+  // Don't bother trying to improve ILP in lower RP regions if occupancy has not
+  // been dropped. All regions will have already been scheduled with the ideal
+  // occupancy targets.
+  if (DAG.StartingOccupancy <= DAG.MinOccupancy)
+    return false;
 
-        LLVM_DEBUG(
-            dbgs()
-            << "Retrying function scheduling with lowest recorded occupancy "
-            << MinOccupancy << ".\n");
-      }
+      dbgs() << "Retrying function scheduling with lowest recorded occupancy "
+             << DAG.MinOccupancy << ".\n");
+  return true;
 
-      if (Stage == PreRARematerialize) {
-        if (RegionsWithMinOcc.none() || Regions.size() == 1)
-          break;
+bool PreRARematStage::initGCNSchedStage() {
+  if (!GCNSchedStage::initGCNSchedStage())
+    return false;
+
+  if (DAG.RegionsWithMinOcc.none() || DAG.Regions.size() == 1)
+    return false;
 
-        const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
-        const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
-        // Check maximum occupancy
-        if (ST.computeOccupancy(MF.getFunction(), MFI.getLDSSize()) ==
-            MinOccupancy)
-          break;
+      DAG.MinOccupancy)
+    return false;
 
-        // FIXME: This pass will invalidate cached MBBLiveIns for regions
-        // inbetween the defs and region we sinked the def to. Cached pressure
-        // for regions where a def is sinked from will also be invalidated. Will
-        // need to be fixed if there is another pass after this pass.
-        static_assert(LastStage == PreRARematerialize,
-                      "Passes after PreRARematerialize are not supported");
 
-        collectRematerializableInstructions();
-        if (RematerializableInsts.empty() || !sinkTriviallyRematInsts(ST, TII))
-          break;
+    return false;
 
-        LLVM_DEBUG(
-            dbgs() << "Retrying function scheduling with improved occupancy of "
-                   << MinOccupancy << " from rematerializing\n");
-      }
-    }
+             << DAG.MinOccupancy << " from rematerializing\n");
+  return true;
+
+void GCNSchedStage::finalizeGCNSchedStage() {
+  DAG.finishBlock();
+  LLVM_DEBUG(dbgs() << "Ending scheduling stage: " << StageID << "\n");
 
-    if (Stage == UnclusteredReschedule)
-      SavedMutations.swap(Mutations);
+void UnclusteredRescheduleStage::finalizeGCNSchedStage() {
+  SavedMutations.swap(DAG.Mutations);
 
-    for (auto Region : Regions) {
-      if (((Stage == UnclusteredReschedule || Stage == PreRARematerialize) &&
-           !RescheduleRegions[RegionIdx]) ||
-          (Stage == ClusteredLowOccupancyReschedule &&
-           !RegionsWithClusters[RegionIdx] && !RegionsWithHighRP[RegionIdx])) {
+  GCNSchedStage::finalizeGCNSchedStage();
+}
 
-        ++RegionIdx;
-        continue;
-      }
+bool GCNSchedStage::initGCNRegion() {
+  // Check whether this new region is also a new block.
+  if (DAG.RegionBegin->getParent() != CurrentMBB)
+    setupNewBlock();
+
+  unsigned NumRegionInstrs = std::distance(DAG.begin(), DAG.end());
+  DAG.enterRegion(CurrentMBB, DAG.begin(), DAG.end(), NumRegionInstrs);
+
+  // Skip empty scheduling regions (0 or 1 schedulable instructions).
+  if (DAG.begin() == DAG.end() || DAG.begin() == std::prev(DAG.end()))
+    return false;
+
+  LLVM_DEBUG(dbgs() << "********** MI Scheduling **********\n");
+  LLVM_DEBUG(dbgs() << MF.getName() << ":" << printMBBReference(*CurrentMBB)
+                    << " " << CurrentMBB->getName()
+                    << "\n  From: " << *DAG.begin() << "    To: ";
+             if (DAG.RegionEnd != CurrentMBB->end()) dbgs() << *DAG.RegionEnd;
+             else dbgs() << "End";
+             dbgs() << " RegionInstrs: " << NumRegionInstrs << '\n');
+
+  // Save original instruction order before scheduling for possible revert.
+  Unsched.clear();
+  Unsched.reserve(DAG.NumRegionInstrs);
+  for (auto &I : DAG)
+    Unsched.push_back(&I);
+
+  PressureBefore = DAG.Pressure[RegionIdx];
+
+  LLVM_DEBUG(
+      dbgs() << "Pressure before scheduling:\nRegion live-ins:";
+      GCNRPTracker::printLiveRegs(dbgs(), DAG.LiveIns[RegionIdx], DAG.MRI);
+      dbgs() << "Region live-in pressure:  ";
+      llvm::getRegPressure(DAG.MRI, DAG.LiveIns[RegionIdx]).print(dbgs());
+      dbgs() << "Region register pressure: "; PressureBefore.print(dbgs()));
+
+  // Set HasClusteredNodes to true for late stages where we have already
+  // collected it. That way pickNode() will not scan SDep's when not needed.
+  S.HasClusteredNodes = StageID > GCNSchedStageID::InitialSchedule;
+  S.HasExcessPressure = false;
+
+  return true;
 
-      RegionBegin = Region.first;
-      RegionEnd = Region.second;
+bool UnclusteredRescheduleStage::initGCNRegion() {
+  if (!DAG.RescheduleRegions[RegionIdx])
+    return false;
 
-      if (RegionBegin->getParent() != MBB) {
-        if (MBB) finishBlock();
-        MBB = RegionBegin->getParent();
-        startBlock(MBB);
-        if (Stage == InitialSchedule)
-          computeBlockPressure(MBB);
-      }
+  return GCNSchedStage::initGCNRegion();
 
-      unsigned NumRegionInstrs = std::distance(begin(), end());
-      enterRegion(MBB, begin(), end(), NumRegionInstrs);
+bool ClusteredLowOccStage::initGCNRegion() {
+  // We may need to reschedule this region if it doesn't have clusters so it
+  // wasn't rescheduled in the last stage, or if we found it was testing
+  // critical register pressure limits in the unclustered reschedule stage. The
+  // later is because we may not have been able to raise the min occupancy in
+  // the previous stage so the region may be overly constrained even if it was
+  // already rescheduled.
+  if (!DAG.RegionsWithClusters[RegionIdx] && !DAG.RegionsWithHighRP[RegionIdx])
+    return false;
 
-      // Skip empty scheduling regions (0 or 1 schedulable instructions).
-      if (begin() == end() || begin() == std::prev(end())) {
-        exitRegion();
-        ++RegionIdx;
-        continue;
-      }
+  return GCNSchedStage::initGCNRegion();
+}
+
+bool PreRARematStage::initGCNRegion() {
+  if (!DAG.RescheduleRegions[RegionIdx])
+    return false;
+
+  return GCNSchedStage::initGCNRegion();
+}
+
+void GCNSchedStage::setupNewBlock() {
+  if (CurrentMBB)
+    DAG.finishBlock();
+
+  CurrentMBB = DAG.RegionBegin->getParent();
+  DAG.startBlock(CurrentMBB);
+  // Get real RP for the region if it hasn't be calculated before. After the
+  // initial schedule stage real RP will be collected after scheduling.
+  if (StageID == GCNSchedStageID::InitialSchedule)
+    DAG.computeBlockPressure(RegionIdx, CurrentMBB);
+}
+
+void GCNSchedStage::finalizeGCNRegion() {
+  DAG.Regions[RegionIdx] = std::make_pair(DAG.RegionBegin, DAG.RegionEnd);
+  DAG.RescheduleRegions[RegionIdx] = false;
+  if (S.HasExcessPressure)
+    DAG.RegionsWithHighRP[RegionIdx] = true;
+
+  // Revert scheduling if we have dropped occupancy or there is some other
+  // reason that the original schedule is better.
+  checkScheduling();
+
+  DAG.exitRegion();
+  RegionIdx++;
+}
+
+void InitialScheduleStage::finalizeGCNRegion() {
+  // Record which regions have clustered nodes for the next unclustered
+  // reschedule stage.
+  assert(nextStage(StageID) == GCNSchedStageID::UnclusteredReschedule);
+  if (S.HasClusteredNodes)
+    DAG.RegionsWithClusters[RegionIdx] = true;
+
+  GCNSchedStage::finalizeGCNRegion();
+}
+
+void GCNSchedStage::checkScheduling() {
+  // Check the results of scheduling.
+  PressureAfter = DAG.getRealRegPressure(RegionIdx);
+  LLVM_DEBUG(dbgs() << "Pressure after scheduling: ";
+             PressureAfter.print(dbgs()));
+
+  if (PressureAfter.getSGPRNum() <= S.SGPRCriticalLimit &&
+      PressureAfter.getVGPRNum(ST.hasGFX90AInsts()) <= S.VGPRCriticalLimit) {
+    DAG.Pressure[RegionIdx] = PressureAfter;
+    DAG.RegionsWithMinOcc[RegionIdx] =
+        PressureAfter.getOccupancy(ST) == DAG.MinOccupancy;
+
+    // Early out if we have achieve the occupancy target.
+    LLVM_DEBUG(dbgs() << "Pressure in desired limits, done.\n");
+    return;
+  }
+
+  unsigned WavesAfter =
+      std::min(S.getTargetOccupancy(), PressureAfter.getOccupancy(ST));
+  unsigned WavesBefore =
+      std::min(S.getTargetOccupancy(), PressureBefore.getOccupancy(ST));
+  LLVM_DEBUG(dbgs() << "Occupancy before scheduling: " << WavesBefore
+                    << ", after " << WavesAfter << ".\n");
+
+  // We may not be able to keep the current target occupancy because of the just
+  // scheduled region. We might still be able to revert scheduling if the
+  // occupancy before was higher, or if the current schedule has register
+  // pressure higher than the excess limits which could lead to more spilling.
+  unsigned NewOccupancy = std::max(WavesAfter, WavesBefore);
+
+  // Allow memory bound functions to drop to 4 waves if not limited by an
+  // attribute.
+  if (WavesAfter < WavesBefore && WavesAfter < DAG.MinOccupancy &&
+      WavesAfter >= MFI.getMinAllowedOccupancy()) {
+    LLVM_DEBUG(dbgs() << "Function is memory bound, allow occupancy drop up to "
+                      << MFI.getMinAllowedOccupancy() << " waves\n");
+    NewOccupancy = WavesAfter;
+  }
+
+  if (NewOccupancy < DAG.MinOccupancy) {
+    DAG.MinOccupancy = NewOccupancy;
+    MFI.limitOccupancy(DAG.MinOccupancy);
+    DAG.RegionsWithMinOcc.reset();
+    LLVM_DEBUG(dbgs() << "Occupancy lowered for the function to "
+                      << DAG.MinOccupancy << ".\n");
+  }
+
+  unsigned MaxVGPRs = ST.getMaxNumVGPRs(MF);
+  unsigned MaxSGPRs = ST.getMaxNumSGPRs(MF);
+  if (PressureAfter.getVGPRNum(false) > MaxVGPRs ||
+      PressureAfter.getAGPRNum() > MaxVGPRs ||
+      PressureAfter.getSGPRNum() > MaxSGPRs) {
+    DAG.RescheduleRegions[RegionIdx] = true;
+    DAG.RegionsWithHighRP[RegionIdx] = true;
+  }
+
+  // Revert if this region's schedule would cause a drop in occupancy or
+  // spilling.
+  if (shouldRevertScheduling(WavesAfter)) {
+    revertScheduling();
+  } else {
+    DAG.Pressure[RegionIdx] = PressureAfter;
+    DAG.RegionsWithMinOcc[RegionIdx] =
+        PressureAfter.getOccupancy(ST) == DAG.MinOccupancy;
+  }
+}
+
+bool GCNSchedStage::shouldRevertScheduling(unsigned WavesAfter) {
+  if (WavesAfter < DAG.MinOccupancy)
+    return true;
+
+  return false;
+}
+
+bool InitialScheduleStage::shouldRevertScheduling(unsigned WavesAfter) {
+  if (GCNSchedStage::shouldRevertScheduling(WavesAfter))
+    return true;
+
+  if (mayCauseSpilling(WavesAfter))
+    return true;
+
+  assert(nextStage(StageID) == GCNSchedStageID::UnclusteredReschedule);
+  // Don't reschedule the region in the next stage if it doesn't have clusters.
+  if (!DAG.RegionsWithClusters[RegionIdx])
+    DAG.RescheduleRegions[RegionIdx] = false;
+
+  return false;
+}
+
+bool UnclusteredRescheduleStage::shouldRevertScheduling(unsigned WavesAfter) {
+  if (GCNSchedStage::shouldRevertScheduling(WavesAfter))
+    return true;
+
+  // If RP is not reduced in the unclustred reschedule stage, revert to the old
+  // schedule.
+  if (!PressureAfter.less(ST, PressureBefore)) {
+    LLVM_DEBUG(dbgs() << "Unclustered reschedule did not help.\n");
+    return true;
+  }
+
+  return false;
+}
+
+bool ClusteredLowOccStage::shouldRevertScheduling(unsigned WavesAfter) {
+  if (GCNSchedStage::shouldRevertScheduling(WavesAfter))
+    return true;
+
+  if (mayCauseSpilling(WavesAfter))
+    return true;
+
+  return false;
+}
+
+bool PreRARematStage::shouldRevertScheduling(unsigned WavesAfter) {
+  if (GCNSchedStage::shouldRevertScheduling(WavesAfter))
+    return true;
+
+  if (mayCauseSpilling(WavesAfter))
+    return true;
+
+  return false;
+}
+
+bool GCNSchedStage::mayCauseSpilling(unsigned WavesAfter) {
+  if (WavesAfter <= MFI.getMinWavesPerEU() &&
+      !PressureAfter.less(ST, PressureBefore) &&
+      DAG.RescheduleRegions[RegionIdx]) {
+    LLVM_DEBUG(dbgs() << "New pressure will result in more spilling.\n");
+    return true;
+  }
+
+  return false;
+}
+
+void GCNSchedStage::revertScheduling() {
+  DAG.RegionsWithMinOcc[RegionIdx] =
+      PressureBefore.getOccupancy(ST) == DAG.MinOccupancy;
+  LLVM_DEBUG(dbgs() << "Attempting to revert scheduling.\n");
+  DAG.RescheduleRegions[RegionIdx] =
+      DAG.RegionsWithClusters[RegionIdx] ||
+      (nextStage(StageID)) != GCNSchedStageID::UnclusteredReschedule;
+  DAG.RegionEnd = DAG.RegionBegin;
+  int SkippedDebugInstr = 0;
+  for (MachineInstr *MI : Unsched) {
+    if (MI->isDebugInstr()) {
+      ++SkippedDebugInstr;
 
-      LLVM_DEBUG(dbgs() << "********** MI Scheduling **********\n");
-      LLVM_DEBUG(dbgs() << MF.getName() << ":" << printMBBReference(*MBB) << " "
-                        << MBB->getName() << "\n  From: " << *begin()
-                        << "    To: ";
-                 if (RegionEnd != MBB->end()) dbgs() << *RegionEnd;
-                 else dbgs() << "End";
-                 dbgs() << " RegionInstrs: " << NumRegionInstrs << '\n');
+    if (MI->getIterator() != DAG.RegionEnd) {
+      DAG.BB->remove(MI);
+      DAG.BB->insert(DAG.RegionEnd, MI);
+      if (!MI->isDebugInstr())
+        DAG.LIS->handleMove(*MI, true);
+    }
+
+    // Reset read-undef flags and update them later.
+    for (auto &Op : MI->operands())
+      if (Op.isReg() && Op.isDef())
+        Op.setIsUndef(false);
+    RegisterOperands RegOpers;
+    RegOpers.collect(*MI, *DAG.TRI, DAG.MRI, DAG.ShouldTrackLaneMasks, false);
+    if (!MI->isDebugInstr()) {
+      if (DAG.ShouldTrackLaneMasks) {
+        // Adjust liveness and add missing dead+read-undef flags.
+        SlotIndex SlotIdx = DAG.LIS->getInstructionIndex(*MI).getRegSlot();
+        RegOpers.adjustLaneLiveness(*DAG.LIS, DAG.MRI, SlotIdx, MI);
+      } else {
+        // Adjust for missing dead-def flags.
+        RegOpers.detectDeadDefs(*MI, *DAG.LIS);
+      }
+    }
+    DAG.RegionEnd = MI->getIterator();
+    ++DAG.RegionEnd;
+    LLVM_DEBUG(dbgs() << "Scheduling " << *MI);
+  }
 
-      schedule();
+  // After reverting schedule, debug instrs will now be at the end of the block
+  // and RegionEnd will point to the first debug instr. Increment RegionEnd
+  // pass debug instrs to the actual end of the scheduling region.
+  while (SkippedDebugInstr-- > 0)
+    ++DAG.RegionEnd;
 
-      exitRegion();
-      ++RegionIdx;
-    }
-    finishBlock();
+  // If Unsched.front() instruction is a debug instruction, this will actually
+  // shrink the region since we moved all debug instructions to the end of the
+  // block. Find the first instruction that is not a debug instruction.
+  DAG.RegionBegin = Unsched.front()->getIterator();
+  if (DAG.RegionBegin->isDebugInstr()) {
+    for (MachineInstr *MI : Unsched) {
+      if (MI->isDebugInstr())
+        continue;
+      DAG.RegionBegin = MI->getIterator();
+      break;
+    }
 
-    if (Stage == UnclusteredReschedule)
-      SavedMutations.swap(Mutations);
-  } while (Stage != LastStage);
+  // Then move the debug instructions back into their correct place and set
+  // RegionBegin and RegionEnd if needed.
+  DAG.placeDebugValues();
+
+  DAG.Regions[RegionIdx] = std::make_pair(DAG.RegionBegin, DAG.RegionEnd);
 }
 
-void GCNScheduleDAGMILive::collectRematerializableInstructions() {
-  const SIRegisterInfo *SRI = static_cast<const SIRegisterInfo *>(TRI);
-  for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {
+void PreRARematStage::collectRematerializableInstructions() {
+  const SIRegisterInfo *SRI = static_cast<const SIRegisterInfo *>(DAG.TRI);
+  for (unsigned I = 0, E = DAG.MRI.getNumVirtRegs(); I != E; ++I) {
     Register Reg = Register::index2VirtReg(I);
-    if (!LIS->hasInterval(Reg))
+    if (!DAG.LIS->hasInterval(Reg))
       continue;
 
     // TODO: Handle AGPR and SGPR rematerialization
-    if (!SRI->isVGPRClass(MRI.getRegClass(Reg)) || !MRI.hasOneDef(Reg) ||
-        !MRI.hasOneNonDBGUse(Reg))
+    if (!SRI->isVGPRClass(DAG.MRI.getRegClass(Reg)) ||
+        !DAG.MRI.hasOneDef(Reg) || !DAG.MRI.hasOneNonDBGUse(Reg))
       continue;
 
-    MachineOperand *Op = MRI.getOneDef(Reg);
+    MachineOperand *Op = DAG.MRI.getOneDef(Reg);
     MachineInstr *Def = Op->getParent();
     if (Op->getSubReg() != 0 || !isTriviallyReMaterializable(*Def))
       continue;
 
-    MachineInstr *UseI = &*MRI.use_instr_nodbg_begin(Reg);
+    MachineInstr *UseI = &*DAG.MRI.use_instr_nodbg_begin(Reg);
     if (Def->getParent() == UseI->getParent())
       continue;
 
     // We are only collecting defs that are defined in another block and are
     // live-through or used inside regions at MinOccupancy. This means that the
     // register must be in the live-in set for the region.
     bool AddedToRematList = false;
-    for (unsigned I = 0, E = Regions.size(); I != E; ++I) {
-      auto It = LiveIns[I].find(Reg);
-      if (It != LiveIns[I].end() && !It->second.none()) {
-        if (RegionsWithMinOcc[I]) {
+    for (unsigned I = 0, E = DAG.Regions.size(); I != E; ++I) {
+      auto It = DAG.LiveIns[I].find(Reg);
+      if (It != DAG.LiveIns[I].end() && !It->second.none()) {
+        if (DAG.RegionsWithMinOcc[I]) {
           RematerializableInsts[I][Def] = UseI;
           AddedToRematList = true;
         }
 
         // Collect regions with rematerializable reg as live-in to avoid
         // searching later when updating RP.
         RematDefToLiveInRegions[Def].push_back(I);
       }
     }
     if (!AddedToRematList)
       RematDefToLiveInRegions.erase(Def);
   }
 }
 
-bool GCNScheduleDAGMILive::sinkTriviallyRematInsts(const GCNSubtarget &ST,
-                                                   const TargetInstrInfo *TII) {
+bool PreRARematStage::sinkTriviallyRematInsts(const GCNSubtarget &ST,
   // Temporary copies of cached variables we will be modifying and replacing if
   // sinking succeeds.
   SmallVector<
       std::pair<MachineBasicBlock::iterator, MachineBasicBlock::iterator>, 32>
       NewRegions;
   DenseMap<unsigned, GCNRPTracker::LiveRegSet> NewLiveIns;
   DenseMap<unsigned, GCNRegPressure> NewPressure;
   BitVector NewRescheduleRegions;
+  LiveIntervals *LIS = DAG.LIS;
 
-  NewRegions.resize(Regions.size());
-  NewRescheduleRegions.resize(Regions.size());
+  NewRegions.resize(DAG.Regions.size());
+  NewRescheduleRegions.resize(DAG.Regions.size());
 
   // Collect only regions that has a rematerializable def as a live-in.
   SmallSet<unsigned, 16> ImpactedRegions;
   for (const auto &It : RematDefToLiveInRegions)
     ImpactedRegions.insert(It.second.begin(), It.second.end());
 
   // Make copies of register pressure and live-ins cache that will be updated
   // as we rematerialize.
   for (auto Idx : ImpactedRegions) {
-    NewPressure[Idx] = Pressure[Idx];
-    NewLiveIns[Idx] = LiveIns[Idx];
+    NewPressure[Idx] = DAG.Pressure[Idx];
+    NewLiveIns[Idx] = DAG.LiveIns[Idx];
   }
-  NewRegions = Regions;
+  NewRegions = DAG.Regions;
   NewRescheduleRegions.reset();
 
   DenseMap<MachineInstr *, MachineInstr *> InsertedMIToOldDef;
   bool Improved = false;
   for (auto I : ImpactedRegions) {
-    if (!RegionsWithMinOcc[I])
+    if (!DAG.RegionsWithMinOcc[I])
       continue;
 
     Improved = false;
     int VGPRUsage = NewPressure[I].getVGPRNum(ST.hasGFX90AInsts());
     int SGPRUsage = NewPressure[I].getSGPRNum();
 
     // TODO: Handle occupancy drop due to AGPR and SGPR.
     // Check if cause of occupancy drop is due to VGPR usage and not SGPR.
-    if (ST.getOccupancyWithNumSGPRs(SGPRUsage) == MinOccupancy)
+    if (ST.getOccupancyWithNumSGPRs(SGPRUsage) == DAG.MinOccupancy)
       break;
 
     // The occupancy of this region could have been improved by a previous
     // iteration's sinking of defs.
-    if (NewPressure[I].getOccupancy(ST) > MinOccupancy) {
+    if (NewPressure[I].getOccupancy(ST) > DAG.MinOccupancy) {
       NewRescheduleRegions[I] = true;
       Improved = true;
       continue;
     }
 
     // First check if we have enough trivially rematerializable instructions to
     // improve occupancy. Optimistically assume all instructions we are able to
     // sink decreased RP.
     int TotalSinkableRegs = 0;
     for (const auto &It : RematerializableInsts[I]) {
       MachineInstr *Def = It.first;
       Register DefReg = Def->getOperand(0).getReg();
       TotalSinkableRegs +=
           SIRegisterInfo::getNumCoveredRegs(NewLiveIns[I][DefReg]);
     }
     int VGPRsAfterSink = VGPRUsage - TotalSinkableRegs;
     unsigned OptimisticOccupancy = ST.getOccupancyWithNumVGPRs(VGPRsAfterSink);
     // If in the most optimistic scenario, we cannot improve occupancy, then do
     // not attempt to sink any instructions.
-    if (OptimisticOccupancy <= MinOccupancy)
+    if (OptimisticOccupancy <= DAG.MinOccupancy)
       break;
 
     unsigned ImproveOccupancy = 0;
     SmallVector<MachineInstr *, 4> SinkedDefs;
     for (auto &It : RematerializableInsts[I]) {
       MachineInstr *Def = It.first;
       MachineBasicBlock::iterator InsertPos =
           MachineBasicBlock::iterator(It.second);
       Register Reg = Def->getOperand(0).getReg();
       // Rematerialize MI to its use block. Since we are only rematerializing
       // instructions that do not have any virtual reg uses, we do not need to
       // call LiveRangeEdit::allUsesAvailableAt() and
       // LiveRangeEdit::canRematerializeAt().
       TII->reMaterialize(*InsertPos->getParent(), InsertPos, Reg,
-                         Def->getOperand(0).getSubReg(), *Def, *TRI);
+                         Def->getOperand(0).getSubReg(), *Def, *DAG.TRI);
       MachineInstr *NewMI = &*(--InsertPos);
       LIS->InsertMachineInstrInMaps(*NewMI);
       LIS->removeInterval(Reg);
       LIS->createAndComputeVirtRegInterval(Reg);
       InsertedMIToOldDef[NewMI] = Def;
 
       // Update region boundaries in scheduling region we sinked from since we
       // may sink an instruction that was at the beginning or end of its region
-      updateRegionBoundaries(NewRegions, Def, /*NewMI =*/nullptr,
+      DAG.updateRegionBoundaries(NewRegions, Def, /*NewMI =*/nullptr,
                                  /*Removing =*/true);
 
       // Update region boundaries in region we sinked to.
-      updateRegionBoundaries(NewRegions, InsertPos, NewMI);
+      DAG.updateRegionBoundaries(NewRegions, InsertPos, NewMI);
 
       LaneBitmask PrevMask = NewLiveIns[I][Reg];
       // FIXME: Also update cached pressure for where the def was sinked from.
       // Update RP for all regions that has this reg as a live-in and remove
       // the reg from all regions as a live-in.
       for (auto Idx : RematDefToLiveInRegions[Def]) {
         NewLiveIns[Idx].erase(Reg);
-        if (InsertPos->getParent() != Regions[Idx].first->getParent()) {
+        if (InsertPos->getParent() != DAG.Regions[Idx].first->getParent()) {
           // Def is live-through and not used in this block.
-          NewPressure[Idx].inc(Reg, PrevMask, LaneBitmask::getNone(), MRI);
+          NewPressure[Idx].inc(Reg, PrevMask, LaneBitmask::getNone(), DAG.MRI);
         } else {
           // Def is used and rematerialized into this block.
           GCNDownwardRPTracker RPT(*LIS);
           auto *NonDbgMI = &*skipDebugInstructionsForward(
               NewRegions[Idx].first, NewRegions[Idx].second);
           RPT.reset(*NonDbgMI, &NewLiveIns[Idx]);
           RPT.advance(NewRegions[Idx].second);
           NewPressure[Idx] = RPT.moveMaxPressure();
         }
       }
 
       SinkedDefs.push_back(Def);
       ImproveOccupancy = NewPressure[I].getOccupancy(ST);
-      if (ImproveOccupancy > MinOccupancy)
+      if (ImproveOccupancy > DAG.MinOccupancy)
         break;
     }
 
     // Remove defs we just sinked from all regions' list of sinkable defs
     for (auto &Def : SinkedDefs)
       for (auto TrackedIdx : RematDefToLiveInRegions[Def])
         RematerializableInsts[TrackedIdx].erase(Def);
 
-    if (ImproveOccupancy <= MinOccupancy)
+    if (ImproveOccupancy <= DAG.MinOccupancy)
       break;
 
     NewRescheduleRegions[I] = true;
     Improved = true;
   }
 
   if (!Improved) {
     // Occupancy was not improved for all regions that were at MinOccupancy.
bool PreRARematStage::sinkTriviallyRematInsts(const GCNSubtarget &ST,
   }
 
   // Occupancy was improved for all regions.
   for (auto &Entry : InsertedMIToOldDef) {
     MachineInstr *MI = Entry.first;
     MachineInstr *OldMI = Entry.second;
 
     // Remove OldMI from BBLiveInMap since we are sinking it from its MBB.
-    BBLiveInMap.erase(OldMI);
+    DAG.BBLiveInMap.erase(OldMI);
 
     // Remove OldMI and update LIS
     Register Reg = MI->getOperand(0).getReg();
     LIS->RemoveMachineInstrFromMaps(*OldMI);
     OldMI->eraseFromParent();
     LIS->removeInterval(Reg);
     LIS->createAndComputeVirtRegInterval(Reg);
   }
 
   // Update live-ins, register pressure, and regions caches.
   for (auto Idx : ImpactedRegions) {
-    LiveIns[Idx] = NewLiveIns[Idx];
-    Pressure[Idx] = NewPressure[Idx];
-    MBBLiveIns.erase(Regions[Idx].first->getParent());
+    DAG.LiveIns[Idx] = NewLiveIns[Idx];
+    DAG.Pressure[Idx] = NewPressure[Idx];
+    DAG.MBBLiveIns.erase(DAG.Regions[Idx].first->getParent());
   }
-  Regions = NewRegions;
-  RescheduleRegions = NewRescheduleRegions;
+  DAG.Regions = NewRegions;
+  DAG.RescheduleRegions = NewRescheduleRegions;
 
   SIMachineFunctionInfo &MFI = *MF.getInfo<SIMachineFunctionInfo>();
-  MFI.increaseOccupancy(MF, ++MinOccupancy);
+  MFI.increaseOccupancy(MF, ++DAG.MinOccupancy);
 
   return true;
 }
 
 // Copied from MachineLICM
-bool GCNScheduleDAGMILive::isTriviallyReMaterializable(const MachineInstr &MI) {
-  if (!TII->isTriviallyReMaterializable(MI))
+bool PreRARematStage::isTriviallyReMaterializable(const MachineInstr &MI) {
+  if (!DAG.TII->isTriviallyReMaterializable(MI))
     return false;
 
   for (const MachineOperand &MO : MI.operands())
     if (MO.isReg() && MO.isUse() && MO.getReg().isVirtual())
       return false;
 
   return true;
 }