[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))
-          continue;
+          NeedsAGPRToVGPRCopy = false;
 
         // For any user of the result of the MFMA which is not an MFMA, we
-        // insert a copy. For a given register, we will only insert one copy
-        // per user block.
-        CopyForUse[UserMI->getParent()].insert(RUOp->getReg());
-
-        if (TII->isMAI(*UserMI))
+        // record a copy location. For a given register, we only record one
+        // copy per user block.
+        if (NeedsAGPRToVGPRCopy)
+          CopyForUse[UserMI->getParent()].insert(RUOp->getReg());
+
+        // If the dst can be wholly recolored to AGPR, its reaching defs are
+        // reclassified along with it, so no per-def bridge copy is needed.
+        if (DstRecolorSafe)
           continue;
-
         SmallVector<SlotIndex, 8> DstUsesReachingDefs;
         findReachingDefs(*RUOp, DAG.LIS, DstUsesReachingDefs);
 
         for (SlotIndex RDIndex : DstUsesReachingDefs) {
           MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIndex);
-          if (TII->isMAI(*RD))
+          if (TII->isMAI(*RD) && RewriteSet.contains(RD))
----------------
lucas-rami wrote:

Same here

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