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

[Matt Arsenault via llvm-commits]

================
@@ -2351,40 +2351,58 @@ 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)
-    return true;
-  if (RD->isCopy() && CandSrc2Regs.contains(RD->getOperand(1).getReg()))
-    return true;
-  return false;
+static bool
+isRewriteCandidateMAI(const MachineInstr *MI, const SIInstrInfo *TII,
+                      const SmallPtrSetImpl<MachineInstr *> &RewriteCandsSet) {
+  return TII->isMAI(*MI) && RewriteCandsSet.contains(MI);
 }
 
-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 canWriteAGPR(const MachineInstr *MI, Register Reg,
+                         const TargetRegisterClass *RegAGPRClass,
+                         const SIInstrInfo *TII, const SIRegisterInfo *SRI) {
+  if (MI->getDesc().getNumDefs() == 0)
+    return false;
+  return MI->getRegClassConstraintEffectForVReg(Reg, RegAGPRClass, TII, SRI) !=
+         nullptr;
+}
+
+static bool userAcceptsAGPR(const MachineInstr *UserMI, Register Reg,
+                            const TargetRegisterClass *RegAGPRClass,
+                            const SIInstrInfo *TII, const SIRegisterInfo *SRI) {
+  return UserMI->getRegClassConstraintEffectForVReg(Reg, RegAGPRClass, TII,
+                                                    SRI) != nullptr;
+}
+
+bool RewriteMFMAFormStage::isRecolorSafe(
----------------
arsenm wrote:

This is pretty algorithmically expensive and is now called more times without any caching.

 O(defs x uses x CFG-walk), is now called for both src2 and dst of every candidate in both initHeuristics and rewrite()

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