[llvm] [AMDGPU] Fix AGPR_32 reg assign for mfma scale ops (PR #168964)

[Matt Arsenault via llvm-commits]

================
@@ -152,9 +152,26 @@ bool AMDGPURewriteAGPRCopyMFMAImpl::recomputeRegClassExceptRewritable(
 
       // We can swap the classes of dst + src2 as a pair to AGPR, so ignore the
       // effects of rewrite candidates. It just so happens that we can use
-      // either AGPR or VGPR in src0/src1, so don't bother checking the
-      // constraint effects of the individual operands.
+      // either AGPR or VGPR in src0/src1. We still need to check constraint
+      // effects for scale variant, which does not allow AGPR.
       if (isRewriteCandidate(*MI)) {
+
+        int AGPROp = AMDGPU::getMFMASrcCVDstAGPROp(MI->getOpcode());
+        MachineInstrBuilder TmpMIB =
+            BuildMI(*MI->getParent(), MI->getIterator(), MI->getDebugLoc(),
+                    TII.get(AGPROp));
+        for (const MachineOperand &TmpMO : MI->operands())
+          TmpMIB.add(TmpMO);
+        MachineInstr *TmpMI = TmpMIB.getInstr();
+        unsigned OpNo = &MO - &MI->getOperand(0);
+        const TargetRegisterClass *EquivalentAGPRRegClass =
+            TRI.getEquivalentAGPRClass(MRI.getRegClass(Reg));
+        const TargetRegisterClass *Allowed = TmpMI->getRegClassConstraintEffect(
----------------
arsenm wrote:

You want TargetInstrInfo::getRegClass to get the static constraint of the known operand (alternatively, you could check that the use is one of the known src0/src1 operands and not the _scale name)

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