[llvm] [AMDGPU] Guard RewriteMFMAFormStage recolor against unsafe def/use (PR #217396)

[Romanov Vlad via llvm-commits]

================
@@ -2728,17 +2738,21 @@ bool RewriteMFMAFormStage::rewrite(
       findReachingDefs(*Src2, DAG.LIS, Src2ReachingDefs);
       SmallSetVector<MachineInstr *, 8> Src2DefsReplace;
 
-      // If src2 has a use that must remain VGPR, it cannot be reclassified to
-      // AGPR.
-      bool Src2NeedsVGPR = Src2NeedsVGPRCache.lookup(MI);
-
-      for (SlotIndex RDIndex : Src2ReachingDefs) {
-        MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIndex);
-        if (!Src2NeedsVGPR &&
-            isReachingDefAGPRForm(RD, RewriteCandsSet, RewriteSrc2Regs, *TII))
-          continue;
-
-        Src2DefsReplace.insert(RD);
+      // An already-redefined src2 reuses its mapped reg below, so treat it as
+      // unsafe to recolor and collect the reaching defs that need a bridge.
+      bool Src2AlreadyRedef = RedefMap.contains(Src2Reg);
+      bool Src2RecolorSafe =
+          !Src2AlreadyRedef &&
+          isRecolorSafe(Src2Reg, {}, RewriteCandsSet, /*IsDst=*/false);
+      // src2 cannot be recolored to AGPR: collect its non-candidate reaching
+      // defs to bridge below (candidate MFMAs already produce AGPR directly).
+      if (!Src2RecolorSafe) {
+        for (SlotIndex RDIndex : Src2ReachingDefs) {
+          MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIndex);
+          if (isRewriteCandidateMAI(RD, TII, RewriteCandsSet))
+            continue;
----------------
romanovvlad wrote:

1. Don't we skip `isRewriteCandidateMAI` in the `isRecolorSafe` already? Do we really need this check?
2. We already collected `CopyForDef`. Can't we reuse it here as `Src2DefsReplace`?

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