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

[Lucas Ramirez via llvm-commits]

================
@@ -2427,51 +2446,53 @@ bool RewriteMFMAFormStage::initHeuristics(
         SmallVector<SlotIndex, 8> Src2ReachingDefs;
         findReachingDefs(*Src2, DAG.LIS, Src2ReachingDefs);
 
-        // If src2 has a use that must remain VGPR, it cannot be reclassified to
-        // AGPR.
-        bool Src2NeedsVGPR = hasUseRequiringVGPR(Src2ReachingDefs, RewriteSet);
-        Src2NeedsVGPRCache[&MI] = Src2NeedsVGPR;
-
-        for (SlotIndex RDIdx : Src2ReachingDefs) {
-          MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIdx);
-          if (!Src2NeedsVGPR &&
-              isReachingDefAGPRForm(RD, RewriteSet, CandSrc2Regs, *TII))
-            continue;
-          CopyForDef.insert(RD);
+        bool Src2RecolorSafe =
+            isRecolorSafe(Src2->getReg(), {}, RewriteSet, /*IsDst=*/false);
+        if (!Src2RecolorSafe) {
+          for (SlotIndex RDIdx : Src2ReachingDefs) {
+            MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIdx);
+            CopyForDef.insert(RD);
+          }
         }
       }
 
       MachineOperand &Dst = MI.getOperand(0);
       SmallVector<MachineOperand *, 8> DstReachingUses;
 
       findReachingUses(&MI, DAG.LIS, DstReachingUses);
+      bool DstRecolorSafe =
+          isRecolorSafe(Dst.getReg(), DstReachingUses, RewriteSet,
+                        /*IsDst=*/true);
 
       for (MachineOperand *RUOp : DstReachingUses) {
         MachineInstr *UserMI = RUOp->getParent();
+        bool NeedsAGPRToVGPRCopy = true;
         // Group members read the AGPR result directly.
         if (TII->isMAI(*UserMI) && RewriteSet.contains(UserMI))
----------------
lucas-rami wrote:

Use `isRewriteCandidateMAI`

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