[SystemZ] Improve shouldCoalesce() for i128. (#74942)

The SystemZ implementation of shouldCoalesce() is merely a workaround
for the fact that regalloc can run out of registers when extending 128-bit
intervals with subreg (GPR64/GPR32) COPYs. 

This patch adds more freedom to the coalescer as it now only checks that
the subreg interval is local to MBB and does not have too many physreg
clobbers.

Author: JonPsson1

llvm/lib/Target/SystemZ/SystemZRegisterInfo.cpp

@@ -377,57 +377,39 @@ SystemZRegisterInfo::eliminateFrameIndex(MachineBasicBlock::iterator MI,
 }
 
 bool SystemZRegisterInfo::shouldCoalesce(MachineInstr *MI,
-                                  const TargetRegisterClass *SrcRC,
-                                  unsigned SubReg,
-                                  const TargetRegisterClass *DstRC,
-                                  unsigned DstSubReg,
-                                  const TargetRegisterClass *NewRC,
-                                  LiveIntervals &LIS) const {
+                                         const TargetRegisterClass *SrcRC,
+                                         unsigned SubReg,
+                                         const TargetRegisterClass *DstRC,
+                                         unsigned DstSubReg,
+                                         const TargetRegisterClass *NewRC,
+                                         LiveIntervals &LIS) const {
   assert (MI->isCopy() && "Only expecting COPY instructions");
 
   // Coalesce anything which is not a COPY involving a subreg to/from GR128.
   if (!(NewRC->hasSuperClassEq(&SystemZ::GR128BitRegClass) &&
         (getRegSizeInBits(*SrcRC) <= 64 || getRegSizeInBits(*DstRC) <= 64)))
     return true;
 
-  // Allow coalescing of a GR128 subreg COPY only if the live ranges are small
-  // and local to one MBB with not too much interferring registers. Otherwise
+  // Allow coalescing of a GR128 subreg COPY only if the subreg liverange is
+  // local to one MBB with not too many interferring physreg clobbers. Otherwise
   // regalloc may run out of registers.
+  unsigned SubregOpIdx = getRegSizeInBits(*SrcRC) == 128 ? 0 : 1;
+  LiveInterval &LI = LIS.getInterval(MI->getOperand(SubregOpIdx).getReg());
 
-  unsigned WideOpNo = (getRegSizeInBits(*SrcRC) == 128 ? 1 : 0);
-  Register GR128Reg = MI->getOperand(WideOpNo).getReg();
-  Register GRNarReg = MI->getOperand((WideOpNo == 1) ? 0 : 1).getReg();
-  LiveInterval &IntGR128 = LIS.getInterval(GR128Reg);
-  LiveInterval &IntGRNar = LIS.getInterval(GRNarReg);
-
-  // Check that the two virtual registers are local to MBB.
+  // Check that the subreg is local to MBB.
   MachineBasicBlock *MBB = MI->getParent();
-  MachineInstr *FirstMI_GR128 =
-    LIS.getInstructionFromIndex(IntGR128.beginIndex());
-  MachineInstr *FirstMI_GRNar =
-    LIS.getInstructionFromIndex(IntGRNar.beginIndex());
-  MachineInstr *LastMI_GR128 = LIS.getInstructionFromIndex(IntGR128.endIndex());
-  MachineInstr *LastMI_GRNar = LIS.getInstructionFromIndex(IntGRNar.endIndex());
-  if ((!FirstMI_GR128 || FirstMI_GR128->getParent() != MBB) ||
-      (!FirstMI_GRNar || FirstMI_GRNar->getParent() != MBB) ||
-      (!LastMI_GR128 || LastMI_GR128->getParent() != MBB) ||
-      (!LastMI_GRNar || LastMI_GRNar->getParent() != MBB))
+  MachineInstr *FirstMI = LIS.getInstructionFromIndex(LI.beginIndex());
+  MachineInstr *LastMI = LIS.getInstructionFromIndex(LI.endIndex());
+  if (!FirstMI || FirstMI->getParent() != MBB ||
+      !LastMI || LastMI->getParent() != MBB)
     return false;
 
-  MachineBasicBlock::iterator MII = nullptr, MEE = nullptr;
-  if (WideOpNo == 1) {
-    MII = FirstMI_GR128;
-    MEE = LastMI_GRNar;
-  } else {
-    MII = FirstMI_GRNar;
-    MEE = LastMI_GR128;
-  }
-
   // Check if coalescing seems safe by finding the set of clobbered physreg
   // pairs in the region.
   BitVector PhysClobbered(getNumRegs());
-  MEE++;
-  for (; MII != MEE; ++MII) {
+  for (MachineBasicBlock::iterator MII = FirstMI,
+                                   MEE = std::next(LastMI->getIterator());
+       MII != MEE; ++MII)
     for (const MachineOperand &MO : MII->operands())
       if (MO.isReg() && MO.getReg().isPhysical()) {
         for (MCPhysReg SI : superregs_inclusive(MO.getReg()))
@@ -436,7 +418,6 @@ bool SystemZRegisterInfo::shouldCoalesce(MachineInstr *MI,
             break;
           }
       }
-  }
 
   // Demand an arbitrary margin of free regs.
   unsigned const DemandedFreeGR128 = 3;

llvm/test/CodeGen/SystemZ/atomicrmw-ops-i128.ll

@@ -14,13 +14,11 @@ define i128 @atomicrmw_xchg(ptr %src, i128 %b) {
 ; CHECK-NEXT:    stmg %r12, %r15, 96(%r15)
 ; CHECK-NEXT:    .cfi_offset %r12, -64
 ; CHECK-NEXT:    .cfi_offset %r13, -56
-; CHECK-NEXT:    .cfi_offset %r14, -48
 ; CHECK-NEXT:    .cfi_offset %r15, -40
-; CHECK-NEXT:    lg %r14, 8(%r4)
+; CHECK-NEXT:    lg %r1, 8(%r4)
 ; CHECK-NEXT:    lg %r0, 0(%r4)
 ; CHECK-NEXT:    lg %r4, 8(%r3)
 ; CHECK-NEXT:    lg %r5, 0(%r3)
-; CHECK-NEXT:    lgr %r1, %r14
 ; CHECK-NEXT:  .LBB0_1: # %atomicrmw.start
 ; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:    lgr %r12, %r5

llvm/test/CodeGen/SystemZ/atomicrmw-xchg-07.ll

@@ -4,11 +4,10 @@
 
 define void @f1(ptr align 16 %ret, ptr align 16 %src, ptr align 16 %b) {
 ; CHECK-LABEL: f1:
-; CHECK:       lg      %r14, 8(%r4)
+; CHECK:       lg      %r1, 8(%r4)
 ; CHECK-NEXT:  lg      %r0, 0(%r4)
 ; CHECK-NEXT:  lg      %r4, 8(%r3)
 ; CHECK-NEXT:  lg      %r5, 0(%r3)
-; CHECK-NEXT:  lgr     %r1, %r14
 ; CHECK-NEXT:.LBB0_1:                          # %atomicrmw.start
 ; CHECK-NEXT:                                  # =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:  lgr     %r12, %r5