[llvm] [AMDGPU] Constrain AV->VReg if we do not exceed RP thresholds (PR #150086)

[Lucas Ramirez via llvm-commits]

================
@@ -1036,6 +1046,78 @@ bool GCNSchedStage::initGCNSchedStage() {
   return true;
 }
 
+bool AVGPRRewriteScheduleStage::reconstrainRegClass(
+    Register Reg, const TargetRegisterClass *NewRC) const {
+  const SIInstrInfo *TII = MF.getSubtarget<GCNSubtarget>().getInstrInfo();
+  const TargetRegisterClass *OldRC = DAG.MRI.getRegClass(Reg);
+  const TargetRegisterInfo *TRI = DAG.MRI.getTargetRegisterInfo();
+  const TargetRegisterClass *ConstrainRC = NewRC;
+  const SIRegisterInfo *SRI = MF.getSubtarget<GCNSubtarget>().getRegisterInfo();
+
+  // Stop early if there is nothing to do.
+  if (!NewRC || NewRC == OldRC)
+    return false;
+
+  // Accumulate constraints from all uses.
+  for (MachineOperand &MO : DAG.MRI.reg_nodbg_operands(Reg)) {
+    // Apply the effect of the given operand to NewRC.
+    MachineInstr *MI = MO.getParent();
+    unsigned OpNo = &MO - &MI->getOperand(0);
+    ConstrainRC = MI->getRegClassConstraintEffect(OpNo, ConstrainRC, TII, TRI);
+    if (!ConstrainRC)
+      return false;
+    if (MI->isCopy()) {
+      MachineOperand &OtherOp = MI->getOperand(1 - OpNo);
+      if (!OtherOp.isReg())
+        continue;
+
+      if (!SRI->isVGPR(DAG.MRI, OtherOp.getReg()))
+        return false;
+    }
+  }
+  DAG.MRI.setRegClass(Reg, ConstrainRC);
+  return true;
+}
+
+bool AVGPRRewriteScheduleStage::initGCNSchedStage() {
+  const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
+
+  // The main benefit of AVReg usage is that the register can be assigned to
+  // either VGPR or AGPR. However, for the unified RF case, we should only be
+  // using AGPR if strictly necessary. That is, if the required number of VGPRs
+  // exceeds the addressable limit.  Otherwise, we should be stricly using VGPRs
+  // to minimize cross RC copies. Thus, if we are under this limit, we should
+  // constrain AVReg- > VReg.
----------------
lucas-rami wrote:

```suggestion
  // constrain AVReg -> VReg.
```

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