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

[Petr Kurapov via llvm-commits]

================
@@ -2803,35 +2817,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");
----------------
kurapov-peter wrote:

I was playing with some artificially constructed IR to understand the checks here and have the following finding. I think this generally does a good job of sidestepping issues with `RedefMap` having one carrier per register, but there are still cases for when this assertion would shoot. Here's one example (I trimmed it, so the IR looks a bit dumb, but at least it's short):
```
# RUN: llc -mtriple=amdgpu9.0a-amd-amdhsa -run-pass=machine-scheduler \
# RUN:     -amdgpu-disable-rewrite-mfma-form-sched-stage=false \
# RUN:     -verify-machineinstrs -o /dev/null %s

--- |
  define void @redefmap_clobber() #0 { ret void }
  attributes #0 = { "amdgpu-flat-work-group-size"="1,256" }
...
---
name:            redefmap_clobber
tracksRegLiveness: true
body: |
  bb.0:
    successors: %bb.1(0x80000000)
    liveins: $vgpr0, $sgpr4_sgpr5, $sgpr6

    %p0:vreg_1024 = IMPLICIT_DEF
    %p1:vreg_1024 = IMPLICIT_DEF
    %p2:vreg_1024 = IMPLICIT_DEF
    %p3:vreg_1024 = IMPLICIT_DEF
    %p4:vreg_1024 = IMPLICIT_DEF
    %p5:vreg_1024 = IMPLICIT_DEF
    %p6:vreg_1024 = IMPLICIT_DEF
    %p7:vreg_1024 = IMPLICIT_DEF
    %p8:vreg_1024 = IMPLICIT_DEF
    %p9:vreg_1024 = IMPLICIT_DEF
    %p10:vreg_1024 = IMPLICIT_DEF
    %p11:vreg_1024 = IMPLICIT_DEF
    %p12:vreg_1024 = IMPLICIT_DEF
    %s0:vreg_64_align2 = IMPLICIT_DEF
    %s1:vreg_64_align2 = IMPLICIT_DEF
    %v0:vgpr_32 = IMPLICIT_DEF
    %pad128:vreg_128_align2 = IMPLICIT_DEF
    %acc:vreg_128_align2 = IMPLICIT_DEF
    SCHED_BARRIER 0
    S_BRANCH %bb.1

  bb.1:
    successors: %bb.2(0x80000000)
    %b:vreg_128_align2 = nofpexcept V_MFMA_F32_16X16X16F16_vgprcd_e64 %s0, %s1, %acc, 0, 0, 0, implicit $mode, implicit $exec
    S_BRANCH %bb.2

  bb.2:
    successors: %bb.1(0x40000000), %bb.3(0x40000000)
    %acc:vreg_128_align2 = nofpexcept V_MFMA_F32_16X16X16F16_vgprcd_e64 %s0, %s1, %pad128, 0, 0, 0, implicit $mode, implicit $exec
    SCHED_BARRIER 0
    %acc.sub0:vreg_128_align2 = nofpexcept V_ADD_F32_e32 %v0, %v0, implicit $mode, implicit $exec
    S_CBRANCH_SCC1 %bb.1, implicit undef $scc
    S_BRANCH %bb.3

  bb.3:
    %sink:vreg_128_align2 = COPY %acc
    %sinkb:vreg_128_align2 = COPY %b
    KILL %p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7, %p8, %p9, %p10, %p11, %p12
    S_ENDPGM 0, implicit %sink, implicit %sinkb
...
```
What happens here is that the redef map accumulates two values for %acc, then adds a bridge copy for the first mfma which is orphaned. This is reachable just because `findReachingUses` is blind to the subreg def.

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