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

[Lucas Ramirez 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,
+                           const SIRegisterInfo *SRI) {
+  const MachineInstr *UseMI = Use->getParent();
+  // COPY is a generic opcode with no operand constraint; it can read an AGPR.
+  if (UseMI->isCopy())
+    return true;
+  // The use's static register-class constraint must admit AGPRs.
+  const TargetRegisterClass *UseRC =
+      UseMI->getRegClassConstraint(Use->getOperandNo(), TII, SRI);
+  return UseRC && SRI->hasAGPRs(UseRC);
+}
+
+bool RewriteMFMAFormStage::isRecolorSafe(
----------------
lucas-rami wrote:

nit: for the sake of API clarity, can you make one version of this function each for each operand type (`isSrc2RecolorSafe`/`isDstRecolorSafe`), both of which can eventually call this "operand-generic" function.

The `src2` version also never uses `DstReachingUses` and AFAIU there will always be at least one reaching use for `dst` so you don't really need `IsDst`: just check `DstReachingUses.empty()` in its place in the generic implementation.

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