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

[via llvm-commits]

================
@@ -2314,40 +2314,63 @@ 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 || !RegAGPRClass)
+    return false;
+  int DefOpIdx =
+      MI->findRegisterDefOperandIdx(Reg, /*TRI=*/nullptr, false, false);
+  return DefOpIdx >= 0 && MI->getRegClassConstraintEffect(
+                              DefOpIdx, RegAGPRClass, TII, SRI) != nullptr;
+}
+
+static bool useAcceptsAGPR(const MachineOperand *Use,
+                           const TargetRegisterClass *RegAGPRClass,
+                           const SIInstrInfo *TII, const SIRegisterInfo *SRI) {
+  const MachineInstr *UseMI = Use->getParent();
+  if (!RegAGPRClass)
+    return false;
+  return UseMI->getRegClassConstraintEffect(Use->getOperandNo(), RegAGPRClass,
+                                            TII, SRI) != nullptr;
+}
+
+bool RewriteMFMAFormStage::isRecolorSafe(
+    Register Reg, ArrayRef<MachineOperand *> DstReachingUses,
+    const SmallPtrSetImpl<MachineInstr *> &RewriteCandsSet, bool IsDst) {
+  const TargetRegisterClass *RegAGPRClass =
+      SRI->getEquivalentAGPRClass(DAG.MRI.getRegClass(Reg));
+  for (MachineInstr &DefMI : DAG.MRI.def_instructions(Reg)) {
+    // A candidate MFMA def is rewritten to AGPR form (it produces the AGPR
+    // result directly), so it does not constrain the recolor.
+    if (isRewriteCandidateMAI(&DefMI, TII, RewriteCandsSet))
+      continue;
+    if (!canWriteAGPR(&DefMI, Reg, RegAGPRClass, TII, SRI))
+      return false;
+    SmallVector<MachineOperand *, 8> DefReachingUses;
+    findReachingUses(&DefMI, DAG.LIS, DefReachingUses);
----------------
xgxanq wrote:

> Really should be reaching users?

Yes. Rather than classifying by def type, we can split the ReachingUses into two sets — rewrite-related and rewrite-unrelated — on the use side.

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