[AMDGPU] Enable PreRARematerialize scheduling pass with multiple high RP regions

[vangthao]

Summary

Enable the PreRARematerialize pass when there are multiple high RP scheduling
regions present. Require the occupancy in all high RP regions be improved
before finalizing sinking. If any high RP region did not improve in occupancy
then un-do all sinking and restore the state to before the pass.

[rampitec]

llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp

657–658

Fix the formatting as tidy suggests.

731

Unrelated to this patch. Can you separately change this to hasOneNonDBGUse()?

739

This also needs a nodbg use (use_nodbg_instructions).

753

It has one use, so you can break from the loop here.

783

This is expensive to make a full copy of everything, even though you only need a small number of regions.

794–805

It is better to have an inverted logic, initialize it to false and only set to true if actually improved anything.

863

Has this *Reg* as a live-in.

881

Don't you need to remove Reg from NewLiveIns at this point for all regions?

910

It is probably also possible to call updateRegionBoundaries to restore the state.

911

return false;

940–941

Actually ++MinOccupancy suggests 1 step improvement. It can be higher, but you do not really know as you would need to check all other blocks too.

942

return true;

[vangthao]

Fix formatting and comments.
Reduce required copies of cached Pressure and Liveins.
Update Pressure and Liveins for all regions where the reg is live-in when sinking.

[vangthao]

llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp

739

If there is a debug use before the first nondbg use, this will sink the MI after the debug use.

753

We also want to collect the def if it is live-through. If we only collect uses, we may skip a trivially rematerializable def that is used in a low RP region but is increasing RP for a high RP region due to being live-through.

881

Updated to also remove Reg from all regions' LiveIns set and re-calculate RP for regions where the Reg was live-in including regions not at MinOccupancy. This should keep cached Pressure and LiveIns up to date after sinking.

910

I am not sure if we can call updateRegionBoundaries. If we sink two defs and both regions are now empty, it will not be possible to tell which def belonged to which region because both will point to the end of block iterator.

[rampitec]

llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp

739

The point is: codegen with and without debug info shall not differ. With this code it will.

789

There can be a lot of regions, I'd better switch to a map indexed by Idx rather than live huge unused chunks of vectors.

1000

Return the newline.

[vangthao]

Rebase and address review comments.

[vangthao]

llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp

881

Actually we also need to update the region where the def was sunk from so cached Pressure is not up to date. Added a FIXME.

[rampitec]

LGTM

[vangthao]

Rebase

llvm/lib/Target/AMDGPU/GCNSchedStrategy.h

 //===-- GCNSchedStrategy.h - GCN Scheduler Strategy -*- C++ -*-------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TARGET_AMDGPU_GCNSCHEDSTRATEGY_H
 #define LLVM_LIB_TARGET_AMDGPU_GCNSCHEDSTRATEGY_H
 
 #include "GCNRegPressure.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/CodeGen/MachineScheduler.h"
 
 namespace llvm {
 
 class SIMachineFunctionInfo;
 class SIRegisterInfo;
 class GCNSubtarget;
 
class GCNScheduleDAGMILive final : public ScheduleDAGMILive {
   BitVector RegionsWithMinOcc;
 
   // Region live-in cache.
   SmallVector<GCNRPTracker::LiveRegSet, 32> LiveIns;
 
   // Region pressure cache.
   SmallVector<GCNRegPressure, 32> Pressure;
 
-  // List of trivially rematerializable instructions we can remat to reduce RP.
-  // First MI is the MI to remat and second MI is the position we should remat
-  // before, usually the MI using the rematerializable instruction.
-  SmallVector<std::pair<MachineInstr *, MachineInstr *>> RematerializableInsts;
+  // 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(unsigned HighRPIdx);
+  void collectRematerializableInstructions();
 
   bool isTriviallyReMaterializable(const MachineInstr &MI, AAResults *AA);
 
   // 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, unsigned HighRPIdx);
+                               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(MachineBasicBlock::iterator MI,
-                              MachineInstr *NewMI, bool Removing = false);
+  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);
 
   void schedule() override;
 
   void finalizeSchedule() override;
 };
 
 } // End namespace llvm
 
 #endif // LLVM_LIB_TARGET_AMDGPU_GCNSCHEDSTRATEGY_H

llvm/lib/Target/AMDGPU/GCNSchedStrategy.cpp

if (Stage > InitialSchedule) {
 
         LLVM_DEBUG(
             dbgs()
             << "Retrying function scheduling with lowest recorded occupancy "
             << MinOccupancy << ".\n");
       }
 
       if (Stage == PreRARematerialize) {
-        if (RegionsWithMinOcc.count() != 1 || Regions.size() == 1)
+        if (RegionsWithMinOcc.none() || Regions.size() == 1)
           break;
 
         const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
         const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
         // Check maximum occupancy
         if (ST.computeOccupancy(MF.getFunction(), MFI.getLDSSize()) ==
             MinOccupancy)
           break;
 
-        // FIXME: This pass will invalidate cached LiveIns, MBBLiveIns and
-        // Pressure for regions inbetween the defs and region we sinked the def
-        // to. Will need to be fixed if there is another pass after this pass.
+        // 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");
 
-        unsigned HighRPIdx = RegionsWithMinOcc.find_first();
-        collectRematerializableInstructions(HighRPIdx);
-        if (RematerializableInsts.empty() ||
-            !sinkTriviallyRematInsts(ST, TII, HighRPIdx))
+        collectRematerializableInstructions();
+        if (RematerializableInsts.empty() || !sinkTriviallyRematInsts(ST, TII))
           break;
 
         LLVM_DEBUG(
             dbgs() << "Retrying function scheduling with improved occupancy of "
                    << MinOccupancy << " from rematerializing\n");
       }
     }
 
do {
     }
     finishBlock();
 
     if (Stage == UnclusteredReschedule)
       SavedMutations.swap(Mutations);
   } while (Stage != LastStage);
 }
 
-void GCNScheduleDAGMILive::collectRematerializableInstructions(
-    unsigned HighRPIdx) {
+void GCNScheduleDAGMILive::collectRematerializableInstructions() {
   const SIRegisterInfo *SRI = static_cast<const SIRegisterInfo *>(TRI);
-  const GCNRPTracker::LiveRegSet &HighRPLiveIns = LiveIns[HighRPIdx];
   for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {
     Register Reg = Register::index2VirtReg(I);
     if (!LIS->hasInterval(Reg))
       continue;
 
     // TODO: Handle AGPR and SGPR rematerialization
     if (!SRI->isVGPRClass(MRI.getRegClass(Reg)) || !MRI.hasOneDef(Reg) ||
         !MRI.hasOneNonDBGUse(Reg))
       continue;
 
-    // We are only collecting defs that are live-through or defined in another
-    // block and used inside this region. This means that the register must be
-    // in the live-in set for this region, else skip this def.
-    if (HighRPLiveIns.find(Reg) == HighRPLiveIns.end())
-      continue;
-
     MachineOperand *Op = MRI.getOneDef(Reg);
     MachineInstr *Def = Op->getParent();
     if (Op->getSubReg() != 0 || !isTriviallyReMaterializable(*Def, AA))
       continue;
 
     MachineInstr *UseI = &*MRI.use_instr_nodbg_begin(Reg);
     if (Def->getParent() == UseI->getParent())
       continue;
 
-    RematerializableInsts.push_back(std::make_pair(Def, UseI));
+    // 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]) {
+          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,
-                                                   unsigned HighRPIdx) {
-  RescheduleRegions.reset();
-  GCNRPTracker::LiveRegSet NewLiveIns;
-  // We may not need to rematerialize all instructions. Keep a list of
-  // instructions we are rematerializing at the end.
-  SmallVector<std::pair<MachineInstr *, MachineInstr *>, 4>
-      TrivialRematDefsToSink;
-
-  GCNRegPressure RegionPressure = Pressure[HighRPIdx];
-  int VGPRUsage = RegionPressure.getVGPRNum(ST.hasGFX90AInsts());
-  int SGPRUsage = RegionPressure.getSGPRNum();
+                                                   const TargetInstrInfo *TII) {
+  // 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;
+
+  NewRegions.resize(Regions.size());
+  NewRescheduleRegions.resize(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];
+  }
+  NewRegions = Regions;
+  NewRescheduleRegions.reset();
+
+  DenseMap<MachineInstr *, MachineInstr *> InsertedMIToOldDef;
+  bool Improved = false;
+  for (auto I : ImpactedRegions) {
+    if (!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.
-  if (ST.getOccupancyWithNumVGPRs(VGPRUsage) > MinOccupancy ||
-      ST.getOccupancyWithNumSGPRs(SGPRUsage) == MinOccupancy)
-    return false;
+    // Check if cause of occupancy drop is due to VGPR usage and not SGPR.
+    if (ST.getOccupancyWithNumSGPRs(SGPRUsage) == 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) {
+      NewRescheduleRegions[I] = true;
+      Improved = true;
+      continue;
+    }
 
-  NewLiveIns.copyFrom(LiveIns[HighRPIdx]);
     // 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 (auto &It : RematerializableInsts) {
-    Register DefReg = It.first->getOperand(0).getReg();
-    TotalSinkableRegs += SIRegisterInfo::getNumCoveredRegs(NewLiveIns[DefReg]);
+    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)
-    return false;
+      break;
 
-  // Keep a list of newly rematerialized instructions so that we can easily
-  // undo if occupancy is not improved.
-  DenseMap<MachineInstr *, MachineInstr *> InsertedMIToOldDef;
-  GCNDownwardRPTracker RPT(*LIS);
-  auto *NonDbgMI = &*skipDebugInstructionsForward(Regions[HighRPIdx].first,
-                                                  Regions[HighRPIdx].second);
     unsigned ImproveOccupancy = 0;
-  for (auto &It : RematerializableInsts) {
+    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().
-    NewLiveIns[Reg] = LaneBitmask::getNone();
       TII->reMaterialize(*InsertPos->getParent(), InsertPos, Reg,
                          Def->getOperand(0).getSubReg(), *Def, *TRI);
       MachineInstr *NewMI = &*(--InsertPos);
       LIS->InsertMachineInstrInMaps(*NewMI);
       LIS->removeInterval(Reg);
       LIS->createAndComputeVirtRegInterval(Reg);
       InsertedMIToOldDef[NewMI] = Def;
 
-    // FIXME: Need better way to update RP without re-iterating over region
-    RPT.reset(*NonDbgMI, &NewLiveIns);
-    RPT.advance(Regions[HighRPIdx].second);
-    GCNRegPressure RPAfterSinking = RPT.moveMaxPressure();
-    ImproveOccupancy = RPAfterSinking.getOccupancy(ST);
+      // 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,
+                             /*Removing =*/true);
+
+      // Update region boundaries in region we sinked to.
+      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()) {
+          // Def is live-through and not used in this block.
+          NewPressure[Idx].inc(Reg, PrevMask, LaneBitmask::getNone(), 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)
         break;
     }
 
-  if (ImproveOccupancy <= MinOccupancy) {
-    // Occupancy is not improved. Undo sinking for the region
+    // 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)
+      break;
+
+    NewRescheduleRegions[I] = true;
+    Improved = true;
+  }
+
+  if (!Improved) {
+    // Occupancy was not improved for all regions that were at MinOccupancy.
+    // Undo sinking and remove newly rematerialized instructions.
     for (auto &Entry : InsertedMIToOldDef) {
       MachineInstr *MI = Entry.first;
       MachineInstr *OldMI = Entry.second;
       Register Reg = MI->getOperand(0).getReg();
       LIS->RemoveMachineInstrFromMaps(*MI);
       MI->eraseFromParent();
       OldMI->clearRegisterDeads(Reg);
       LIS->removeInterval(Reg);
       LIS->createAndComputeVirtRegInterval(Reg);
     }
     return false;
   }
 
-  // Occupancy is improved.
+  // Occupancy was improved for all regions.
   for (auto &Entry : InsertedMIToOldDef) {
     MachineInstr *MI = Entry.first;
     MachineInstr *OldMI = Entry.second;
-    // 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(OldMI, /*NewMI =*/nullptr, /*Removing =*/true);
 
     // Remove OldMI from BBLiveInMap since we are sinking it from its MBB.
     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 region boundaries in region we sinked to.
-    MachineBasicBlock::iterator InsertPos =
-        std::next(MachineBasicBlock::iterator(MI));
-    updateRegionBoundaries(InsertPos, MI);
   }
 
-  // Update cached live-ins and register pressure after rematerializing
-  LiveIns[HighRPIdx].copyFrom(NewLiveIns);
-  MBBLiveIns.erase(Regions[HighRPIdx].first->getParent());
-
-  GCNDownwardRPTracker RPTracker(*LIS);
-  RPTracker.advance(Regions[HighRPIdx].first, Regions[HighRPIdx].second,
-                    &LiveIns[HighRPIdx]);
-  Pressure[HighRPIdx] = RPTracker.moveMaxPressure();
+  // 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());
+  }
+  Regions = NewRegions;
+  RescheduleRegions = NewRescheduleRegions;
 
   SIMachineFunctionInfo &MFI = *MF.getInfo<SIMachineFunctionInfo>();
   MFI.increaseOccupancy(MF, ++MinOccupancy);
-  RescheduleRegions[HighRPIdx] = true;
 
   return true;
 }
 
 // Copied from MachineLICM
 bool GCNScheduleDAGMILive::isTriviallyReMaterializable(const MachineInstr &MI,
                                                        AAResults *AA) {
   if (!TII->isTriviallyReMaterializable(MI, AA))
     return false;
 
   for (const MachineOperand &MO : MI.operands())
     if (MO.isReg() && MO.isUse() && MO.getReg().isVirtual())
       return false;
 
   return true;
 }
 
 // When removing, we will have to check both beginning and ending of the region.
 // When inserting, we will only have to check if we are inserting NewMI in front
 // of a scheduling region and do not need to check the ending since we will only
 // ever be inserting before an already existing MI.
 void GCNScheduleDAGMILive::updateRegionBoundaries(
+    SmallVectorImpl<std::pair<MachineBasicBlock::iterator,
+                              MachineBasicBlock::iterator>> &RegionBoundaries,
     MachineBasicBlock::iterator MI, MachineInstr *NewMI, bool Removing) {
-  unsigned I = 0, E = Regions.size();
+  unsigned I = 0, E = RegionBoundaries.size();
   // Search for first region of the block where MI is located
-  while (I != E && MI->getParent() != Regions[I].first->getParent())
+  while (I != E && MI->getParent() != RegionBoundaries[I].first->getParent())
     ++I;
 
   for (; I != E; ++I) {
-    if (MI->getParent() != Regions[I].first->getParent())
+    if (MI->getParent() != RegionBoundaries[I].first->getParent())
       return;
 
-    if (Removing && MI == Regions[I].first && MI == Regions[I].second) {
+    if (Removing && MI == RegionBoundaries[I].first &&
+        MI == RegionBoundaries[I].second) {
       // MI is in a region with size 1, after removing, the region will be
       // size 0, set RegionBegin and RegionEnd to pass end of block iterator.
-      Regions[I] =
+      RegionBoundaries[I] =
           std::make_pair(MI->getParent()->end(), MI->getParent()->end());
       return;
     }
-    if (MI == Regions[I].first) {
+    if (MI == RegionBoundaries[I].first) {
       if (Removing)
-        Regions[I] = std::make_pair(std::next(MI), Regions[I].second);
+        RegionBoundaries[I] =
+            std::make_pair(std::next(MI), RegionBoundaries[I].second);
       else
         // Inserted NewMI in front of region, set new RegionBegin to NewMI
-        Regions[I] = std::make_pair(MachineBasicBlock::iterator(NewMI),
-                                    Regions[I].second);
+        RegionBoundaries[I] = std::make_pair(MachineBasicBlock::iterator(NewMI),
+                                             RegionBoundaries[I].second);
       return;
     }
-    if (Removing && MI == Regions[I].second) {
-      Regions[I] = std::make_pair(Regions[I].first, std::prev(MI));
+    if (Removing && MI == RegionBoundaries[I].second) {
+      RegionBoundaries[I] =
+          std::make_pair(RegionBoundaries[I].first, std::prev(MI));
       return;
     }
   }
 }

llvm/test/CodeGen/AMDGPU/machine-scheduler-sink-trivial-remats.mir

; predcessors: %bb.1
     S_NOP 0, implicit %10, implicit %11
     S_NOP 0, implicit %12, implicit %13
     S_NOP 0, implicit %14, implicit %15
     S_NOP 0, implicit %16, implicit %17
     S_NOP 0, implicit %18, implicit %19
     S_NOP 0, implicit %20, implicit %21
     S_ENDPGM 0
 ...
-# TODO: Handle multiple blocks
 ---
 name:            test_occ_9_two_block_high_rp
 tracksRegLiveness: true
 machineFunctionInfo:
   isEntryFunction: true
 body:             |
   ; GFX908-LABEL: name: test_occ_9_two_block_high_rp
   ; GFX908: bb.0:
body:             |
   ; GFX908-NEXT:   %15:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 15, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %16:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 16, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %17:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 17, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %18:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 18, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %19:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 19, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %20:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 20, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %21:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 21, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %22:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 22, implicit $exec, implicit $mode, implicit-def $m0
-  ; GFX908-NEXT:   %23:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 23, implicit $exec, implicit $mode
   ; GFX908-NEXT: {{  $}}
   ; GFX908-NEXT: bb.1:
   ; GFX908-NEXT:   successors: %bb.2(0x80000000)
   ; GFX908-NEXT: {{  $}}
   ; GFX908-NEXT:   %24:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 24, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   S_NOP 0, implicit %24
+  ; GFX908-NEXT:   %23:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 23, implicit $exec, implicit $mode
   ; GFX908-NEXT:   S_NOP 0, implicit %23
   ; GFX908-NEXT: {{  $}}
   ; GFX908-NEXT: bb.2:
   ; GFX908-NEXT:   successors: %bb.3(0x80000000)
   ; GFX908-NEXT: {{  $}}
-  ; GFX908-NEXT:   %25:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 25, implicit $exec, implicit $mode
   ; GFX908-NEXT:   S_NOP 0
   ; GFX908-NEXT: {{  $}}
   ; GFX908-NEXT: bb.3:
   ; GFX908-NEXT:   successors: %bb.4(0x80000000)
   ; GFX908-NEXT: {{  $}}
   ; GFX908-NEXT:   %26:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 26, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   S_NOP 0, implicit %26
+  ; GFX908-NEXT:   %25:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 25, implicit $exec, implicit $mode
   ; GFX908-NEXT:   S_NOP 0, implicit %25
   ; GFX908-NEXT: {{  $}}
   ; GFX908-NEXT: bb.4:
   ; GFX908-NEXT:   S_NOP 0, implicit %0, implicit %1
   ; GFX908-NEXT:   S_NOP 0, implicit %2, implicit %3
   ; GFX908-NEXT:   S_NOP 0, implicit %4, implicit %5
   ; GFX908-NEXT:   S_NOP 0, implicit %6, implicit %7
   ; GFX908-NEXT:   S_NOP 0, implicit %8, implicit %9
; predcessors: %bb.3
     S_NOP 0, implicit %12, implicit %13
     S_NOP 0, implicit %14, implicit %15
     S_NOP 0, implicit %16, implicit %17
     S_NOP 0, implicit %18, implicit %19
     S_NOP 0, implicit %20, implicit %21
     S_NOP 0, implicit %22
     S_ENDPGM 0
 ...
-# TODO: Handle multiple blocks
 ---
 name:            test_occ_9_two_block_high_rp_minimum_sinking
 tracksRegLiveness: true
 machineFunctionInfo:
   isEntryFunction: true
 body:             |
   ; GFX908-LABEL: name: test_occ_9_two_block_high_rp_minimum_sinking
   ; GFX908: bb.0:
body:             |
   ; GFX908-NEXT:   %14:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 14, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %15:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 15, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %16:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 16, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %17:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 17, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %18:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 18, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %19:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 19, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %20:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 20, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   %21:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 21, implicit $exec, implicit $mode, implicit-def $m0
-  ; GFX908-NEXT:   %22:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 22, implicit $exec, implicit $mode
   ; GFX908-NEXT:   %23:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 23, implicit $exec, implicit $mode
   ; GFX908-NEXT: {{  $}}
   ; GFX908-NEXT: bb.1:
   ; GFX908-NEXT:   successors: %bb.2(0x80000000)
   ; GFX908-NEXT: {{  $}}
   ; GFX908-NEXT:   %24:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 24, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   S_NOP 0, implicit %24
+  ; GFX908-NEXT:   %22:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 22, implicit $exec, implicit $mode
   ; GFX908-NEXT:   S_NOP 0, implicit %23, implicit %22
   ; GFX908-NEXT: {{  $}}
   ; GFX908-NEXT: bb.2:
   ; GFX908-NEXT:   successors: %bb.3(0x80000000)
   ; GFX908-NEXT: {{  $}}
-  ; GFX908-NEXT:   %25:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 25, implicit $exec, implicit $mode
   ; GFX908-NEXT:   %26:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 26, implicit $exec, implicit $mode
   ; GFX908-NEXT:   S_NOP 0
   ; GFX908-NEXT: {{  $}}
   ; GFX908-NEXT: bb.3:
   ; GFX908-NEXT:   successors: %bb.4(0x80000000)
   ; GFX908-NEXT: {{  $}}
   ; GFX908-NEXT:   %27:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 27, implicit $exec, implicit $mode, implicit-def $m0
   ; GFX908-NEXT:   S_NOP 0, implicit %27
+  ; GFX908-NEXT:   %25:vgpr_32 = nofpexcept V_CVT_I32_F64_e32 25, implicit $exec, implicit $mode
   ; GFX908-NEXT:   S_NOP 0, implicit %25, implicit %26
   ; GFX908-NEXT: {{  $}}
   ; GFX908-NEXT: bb.4:
   ; GFX908-NEXT:   S_NOP 0, implicit %0, implicit %1
   ; GFX908-NEXT:   S_NOP 0, implicit %2, implicit %3
   ; GFX908-NEXT:   S_NOP 0, implicit %4, implicit %5
   ; GFX908-NEXT:   S_NOP 0, implicit %6, implicit %7
   ; GFX908-NEXT:   S_NOP 0, implicit %8, implicit %9