[AMDGPU] Enable SGPR copy folding

[rampitec]

Summary

That used to fail in the lasti testcase function because after
%0:sreg_64.sub0 was folded into %3:sreg_32_xm0_xexec COPY, it
was further folded into S_STORE_DWORD_IMM. Its legal effective
subreg class is SReg_32 while instruction expects more restricted
SReg_32_XM0_EXEC. However, SIInstrInfo::isLegalRegOperand()
passed the legality check and it was caught in the verifier.

Borrowed code from the verifier to check for RC legality.

[arsenm]

LGTM. Typo "lasti" in commit message

llvm/lib/Target/AMDGPU/SIFoldOperands.cpp

while (ImpOpI != ImpOpE) {
       MachineInstr::mop_iterator Tmp = ImpOpI;
       ImpOpI++;
       UseMI->RemoveOperand(UseMI->getOperandNo(Tmp));
     }
     CopiesToReplace.push_back(UseMI);
   } else {
     if (UseMI->isCopy() && OpToFold.isReg() &&
         UseMI->getOperand(0).getReg().isVirtual() &&
-        TRI->isVectorRegister(*MRI, UseMI->getOperand(0).getReg()) &&
         !UseMI->getOperand(1).getSubReg()) {
       LLVM_DEBUG(dbgs() << "Folding " << OpToFold
                         << "\n into " << *UseMI << '\n');
       unsigned Size = TII->getOpSize(*UseMI, 1);
       Register UseReg = OpToFold.getReg();
       UseMI->getOperand(1).setReg(UseReg);
       UseMI->getOperand(1).setSubReg(OpToFold.getSubReg());
       UseMI->getOperand(1).setIsKill(false);

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

bool SIInstrInfo::isLegalRegOperand(const MachineRegisterInfo &MRI,
   if (!MO.isReg())
     return false;
 
   Register Reg = MO.getReg();
   const TargetRegisterClass *RC = Register::isVirtualRegister(Reg)
                                       ? MRI.getRegClass(Reg)
                                       : RI.getPhysRegClass(Reg);
 
-  const SIRegisterInfo *TRI =
-      static_cast<const SIRegisterInfo*>(MRI.getTargetRegisterInfo());
-  RC = TRI->getSubRegClass(RC, MO.getSubReg());
-
-  // In order to be legal, the common sub-class must be equal to the
-  // class of the current operand.  For example:
-  //
-  // v_mov_b32 s0 ; Operand defined as vsrc_b32
-  //              ; RI.getCommonSubClass(s0,vsrc_b32) = sgpr ; LEGAL
-  //
-  // s_sendmsg 0, s0 ; Operand defined as m0reg
-  //                 ; RI.getCommonSubClass(s0,m0reg) = m0reg ; NOT LEGAL
+  const TargetRegisterClass *DRC = RI.getRegClass(OpInfo.RegClass);
+  if (MO.getSubReg()) {
+    const MachineFunction *MF = MO.getParent()->getParent()->getParent();
+    const TargetRegisterClass *SuperRC = RI.getLargestLegalSuperClass(RC, *MF);
+    if (!SuperRC)
+      return false;
 
-  return RI.getCommonSubClass(RC, RI.getRegClass(OpInfo.RegClass)) == RC;
+    DRC = RI.getMatchingSuperRegClass(SuperRC, DRC, MO.getSubReg());
+    if (!DRC)
+      return false;
+  }
+  return RC->hasSuperClassEq(DRC);
 }
 
 bool SIInstrInfo::isLegalVSrcOperand(const MachineRegisterInfo &MRI,
                                      const MCOperandInfo &OpInfo,
                                      const MachineOperand &MO) const {
   if (MO.isReg())
     return isLegalRegOperand(MRI, OpInfo, MO);
 

llvm/test/CodeGen/AMDGPU/fold-sgpr-copy.mir

+# RUN: llc -march=amdgcn -mcpu=gfx900 -run-pass=si-fold-operands,dead-mi-elimination -verify-machineinstrs %s -o - | FileCheck -check-prefix=GCN %s
+
+---
+
+# GCN-LABEL: name: fold_sgpr_to_sgpr_copy_full
+# GCN:      %0:sgpr_32 = IMPLICIT_DEF
+# GCN-NEXT: S_STORE_DWORD_IMM %0, undef $sgpr10_sgpr11, 0, 0, 0
+
+name:            fold_sgpr_to_sgpr_copy_full
+body:             |
+  bb.0:
+
+    %0:sgpr_32 = IMPLICIT_DEF
+    %1:sgpr_32 = COPY %0
+    %2:sgpr_32 = COPY %1
+    S_STORE_DWORD_IMM %2, undef $sgpr10_sgpr11, 0, 0, 0
+...
+
+# GCN-LABEL: name: fold_sgpr_to_sgpr_copy_subreg
+# GCN:      %0:sreg_64 = IMPLICIT_DEF
+# GCN-NEXT: %2:sgpr_32 = COPY %0.sub0
+# GCN-NEXT: S_STORE_DWORD_IMM %2, undef $sgpr10_sgpr11, 0, 0, 0
+
+name:            fold_sgpr_to_sgpr_copy_subreg
+body:             |
+  bb.0:
+
+    %0:sreg_64 = IMPLICIT_DEF
+    %1:sgpr_32 = COPY %0.sub0
+    %2:sgpr_32 = COPY %1
+    S_STORE_DWORD_IMM %2, undef $sgpr10_sgpr11, 0, 0, 0
+...
+
+# GCN-LABEL: name: fold_sgpr_to_sgpr_copy_subreg2
+# GCN:      %0:sreg_64 = IMPLICIT_DEF
+# GCN-NEXT: %3:sreg_32_xm0_xexec = COPY %0.sub0
+# GCN-NEXT: S_STORE_DWORD_IMM %3, undef $sgpr10_sgpr11, 0, 0, 0
+
+name:            fold_sgpr_to_sgpr_copy_subreg2
+body:             |
+  bb.0:
+
+    %0:sreg_64 = IMPLICIT_DEF
+    %1:sgpr_32 = COPY %0.sub0
+    %2:sgpr_32 = COPY %1
+    %3:sreg_32_xm0_xexec = COPY %2
+    S_STORE_DWORD_IMM %3, undef $sgpr10_sgpr11, 0, 0, 0
+...