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

[Matt Arsenault via llvm-commits]

================
@@ -2803,35 +2816,92 @@ bool RewriteMFMAFormStage::rewrite(
 
     findReachingUses(MI, DAG.LIS, DstReachingUses);
 
+    // An already-redefined dst reuses its mapped reg, so treat it as unsafe to
+    // recolor and bridge its reaching defs instead.
+    bool DstAlreadyRedef = RedefMap.contains(DstReg);
+    bool DstRecolorSafe =
+        !DstAlreadyRedef &&
+        isRecolorSafe(DstReg, DstReachingUses, RewriteCandsSet, /*IsDst=*/true);
+    const TargetRegisterClass *DstAGPRClass =
+        SRI->getEquivalentAGPRClass(DAG.MRI.getRegClass(DstReg));
     for (MachineOperand *RUOp : DstReachingUses) {
       MachineInstr *UserMI = RUOp->getParent();
-      // Group members read the AGPR result directly.
-      if (TII->isMAI(*UserMI) && RewriteCandsSet.contains(UserMI))
-        continue;
-
-      // If there is a non mai reaching use, then we need a copy.
-      if (find(DstReachingUseCopies, RUOp) == DstReachingUseCopies.end())
+      // Decide whether this reaching use can read the dst's AGPR form directly
+      // or needs an AGPR->VGPR bridge copy.
+      //   - A group-member MFMA always reads the AGPR result directly.
+      //   - Any other user can skip the bridge only when the dst is recolored
+      //     to AGPR (DstRecolorSafe) and its operand accepts an AGPR. When the
+      //     dst is unsafe, its original reg stays VGPR, so every non-MFMA user
+      //     must go through a bridge copy.
+      bool CanReadAGPR =
+          TII->isMAI(*UserMI)
+              ? RewriteCandsSet.contains(UserMI)
+              : DstRecolorSafe &&
+                    userAcceptsAGPR(UserMI, DstReg, DstAGPRClass, TII, SRI);
+      if (!CanReadAGPR &&
+          find(DstReachingUseCopies, RUOp) == DstReachingUseCopies.end())
         DstReachingUseCopies.push_back(RUOp);
-
-      // Non-rewritten MAI: its defs aren't being reclassified.
-      if (TII->isMAI(*UserMI))
+      // If the dst is wholly recolored to AGPR, its reaching defs are
+      // reclassified along with it, so none of them need a bridge copy.
+      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 (isRewriteCandidateMAI(RD, TII, RewriteCandsSet))
           continue;
-
-        // If there is a non mai reaching def of this reaching use, then we will
-        // need a copy.
+        // A non-candidate reaching def must be bridged to VGPR; record it once
+        // (dedup against DstUseDefsReplace).
         if (find(DstUseDefsReplace, RD) == DstUseDefsReplace.end())
           DstUseDefsReplace.push_back(RD);
       }
     }
+    // The dst has no reaching uses and cannot be recolored: create a fresh
+    // reg to carry the AGPR-form value and record the mapping, leaving the
+    // original dst reg in VGPR form.
+    if (DstReachingUses.empty() && !DstRecolorSafe) {
+      // Exclusion must already have dropped any dst that was bridged as an
+      // earlier MFMA's src2, so it cannot be pre-mapped when we reach here.
+      assert(
+          !RedefMap.contains(DstReg) &&
+          "empty-use dst unexpectedly already mapped -- exclusion missed it");
+      const TargetRegisterClass *DstRC = DAG.MRI.getRegClass(DstReg);
+      const TargetRegisterClass *VGPRRC = SRI->getEquivalentVGPRClass(DstRC);
+      MappedReg = DAG.MRI.createVirtualRegister(VGPRRC);
+      RedefMap[DstReg] = MappedReg;
----------------
arsenm wrote:

can replace the assert above with insert + llvm_unreachable 

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