[llvm] [AMDGPU][Scheduler] Refactor ArchVGPR rematerialization during scheduling (PR #125885)

[Jeffrey Byrnes via llvm-commits]

================
@@ -1989,45 +2040,115 @@ bool PreRARematStage::isTriviallyReMaterializable(const MachineInstr &MI) {
   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.
+/// Identifies the parent MBB of a \p Region. For an empty region, this runs in
+/// linear time in the number of MBBs in \p MF; otherwise it runs in constant
+/// time.
+static MachineBasicBlock *getRegionMBB(MachineFunction &MF,
+                                       RegionBoundaries &Region) {
+  if (Region.first != Region.second)
+    return Region.first->getParent();
+  // The boundaries of an empty region may not point to a valid MI from which we
+  // can get the parent MBB, so we have to iterate over all the MF's MBBs.
+  for (MachineBasicBlock &MBB : MF)
+    if (MBB.end() == Region.second)
+      return &MBB;
+  llvm_unreachable("region's parent MBB cannot be identified");
+}
+
+void PreRARematStage::finalizeGCNSchedStage() {
+  // We consider that reducing spilling is always beneficial so we never
+  // rollback rematerializations in such cases. It's also possible that
+  // rescheduling lowers occupancy over the one achived just through remats, in
+  // which case we do not want to rollback either.
+  if (!IncreaseOccupancy || AchievedOcc == TargetOcc)
+    return;
+
+  REMAT_DEBUG(dbgs() << "Rollbacking all rematerializations\n");
+  const auto *TII =
+      static_cast<const SIInstrInfo *>(MF.getSubtarget().getInstrInfo());
+
+  // Rollback the rematerializations.
+  for (const auto &[_, Remat] : Rematerializations) {
+    MachineInstr &RematMI = *Remat.RematMI;
+    MachineBasicBlock::iterator InsertPos(DAG.Regions[Remat.DefRegion].second);
+    MachineBasicBlock *MBB = getRegionMBB(MF, DAG.Regions[Remat.DefRegion]);
+    Register Reg = RematMI.getOperand(0).getReg();
+    unsigned SubReg = RematMI.getOperand(0).getSubReg();
+
+    // Re-rematerialize MI at the end of its original region. Note that it may
+    // not be rematerialized exactly in the same position as originally within
+    // the region, but it should not matter much. Since we are only
+    // rematerializing instructions that do not have any virtual reg uses, we
----------------
jrbyrnes wrote:

we remat instructions with virtual reg uses so the comment is incorrect. However, all the uses will be live out to the region so the logic is still valid.

https://github.com/llvm/llvm-project/pull/125885