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

[Matt Arsenault via llvm-commits]

================
@@ -2316,40 +2316,69 @@ void GCNSchedStage::modifyRegionSchedule(unsigned RegionIdx,
   DAG.Regions[RegionIdx].first = MIOrder.front();
 }
 
-/// Returns true if reaching def \p RD will be in AGPR form after the rewrite
-/// and so needs no bridge copy: a candidate MFMA in \p RewriteSet, an
-/// AV_MOV_*_IMM_PSEUDO, or a copy from a candidate src2 reg in \p CandSrc2Regs.
-/// A non-candidate MFMA stays in VGPR form and still needs a bridge.
-static bool isReachingDefAGPRForm(
-    MachineInstr *RD, const SmallPtrSetImpl<MachineInstr *> &RewriteSet,
-    const DenseSet<Register> &CandSrc2Regs, const SIInstrInfo &TII) {
-  if (TII.isMAI(*RD))
-    return RewriteSet.contains(RD);
-  if (RD->getOpcode() == AMDGPU::AV_MOV_B32_IMM_PSEUDO ||
-      RD->getOpcode() == AMDGPU::AV_MOV_B64_IMM_PSEUDO)
+static bool
+isRewriteCandidateMAI(const MachineInstr *MI, const SIInstrInfo *TII,
+                      const SmallPtrSetImpl<MachineInstr *> &RewriteCandsSet) {
+  return TII->isMAI(*MI) && RewriteCandsSet.contains(MI);
+}
+
+static bool canWriteAGPR(const MachineInstr *MI, const SIInstrInfo *TII,
+                         const SIRegisterInfo *SRI) {
+  // A writer can write its result into an AGPR lane iff its def operand's
+  // register-class constraint admits AGPRs (AV/agnostic classes do, plain VGPR
+  // classes don't). COPY is a generic opcode with no operand constraint, so
+  // special-case it.
+  if (MI->isCopy())
     return true;
-  if (RD->isCopy() && CandSrc2Regs.contains(RD->getOperand(1).getReg()))
-    return true;
-  return false;
+  const MCInstrDesc &Desc = MI->getDesc();
+  if (Desc.getNumDefs() == 0)
+    return false;
+  const TargetRegisterClass *DefRC = TII->getRegClass(Desc, 0);
+  return DefRC && SRI->hasAGPRs(DefRC);
 }
 
-bool RewriteMFMAFormStage::hasUseRequiringVGPR(
-    ArrayRef<SlotIndex> Src2ReachingDefs,
-    const SmallPtrSetImpl<MachineInstr *> &RewriteSet) {
-  for (SlotIndex RDIdx : Src2ReachingDefs) {
-    const MachineInstr *RD = DAG.LIS->getInstructionFromIndex(RDIdx);
-    SmallVector<MachineOperand *, 8> ReachingUses;
-    findReachingUses(RD, DAG.LIS, ReachingUses);
-    for (const MachineOperand *UseMO : ReachingUses) {
-      const MachineInstr *UseMI = UseMO->getParent();
-      if (UseMI->isCopy())
-        continue;
-      if (TII->isMAI(*UseMI) && RewriteSet.contains(UseMI))
+static bool useAcceptsAGPR(const MachineOperand *Use, const SIInstrInfo *TII,
----------------
arsenm wrote:

This is unnecessarily special casing. You should check if the operand class is compatible with the class you want, not that it's specifically AGPR 

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