[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())
----------------
arsenm wrote:

Maybe should generalize this to the unknown register class, which would cover all the pseudos. Though there is some risk there. In particular I'm concerned about future SSAification where REG_SEQUENCE may appear here 

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