[llvm] [AMDGPU] Fix RewriteMFMAFormSchedStage (PR #194887)

[Matt Arsenault via llvm-commits]

================
@@ -2674,37 +2681,72 @@ bool RewriteMFMAFormStage::rewrite(
             unsigned UpdateRegion = LMI->second;
             DAG.Regions[UpdateRegion].second = VGPRCopy;
             LastMIToRegion.erase(RD);
+            // If RD was also the first MI of the next region, update that
+            // region's start boundary to keep the two sides in sync.
+            DenseMap<MachineInstr *, unsigned>::iterator FMI =
+                FirstMIToRegion.find(RD);
+            if (FMI != FirstMIToRegion.end()) {
+              DAG.Regions[FMI->second].first = VGPRCopy;
+              FirstMIToRegion.erase(RD);
+            }
           }
         }
       }
     }
 
     DenseSet<MachineOperand *> &DstRegSet = ReplaceMap[DstReg];
+
+    // Collect same-block uses; defer cross-block uses to ReachingUseTracker.
+    SmallVector<MachineOperand *, 4> SameBlockUses;
     for (MachineOperand *RU : DstReachingUseCopies) {
       MachineBasicBlock *RUBlock = RU->getParent()->getParent();
-      // Just keep track of the reaching use of this register by block. After we
-      // have scanned all the MFMAs we can find optimal insert pts.
       if (RUBlock != MI->getParent()) {
         ReachingUseTracker[RUBlock->getNumber()][DstReg].insert(RU);
         continue;
       }
+      SameBlockUses.push_back(RU);
+    }
 
-      // Special case, the use is in the same block as the MFMA. Insert the copy
-      // just before the use.
+    // One COPY for all same-block uses, placed before the earliest use to
----------------
arsenm wrote:

This feels kind of brute force. Can't you track the earliest position in the above loop building SameBlockUses?

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