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

[Jonas Paulsson via llvm-commits]

https://github.com/JonPsson1 created https://github.com/llvm/llvm-project/pull/74942

The SystemZ implementation of shouldCoalesce() is merely a workaround for the fact that regalloc can run out of registers when building 128-bit registers out of subregs. Too many long GPR128 live ranges resulting from coalescing must be avoided. The heuristic used is to only allow this when the resulting live range is relatively small inside a single MBB.

This patch adds a bit more freedom to the coalescer by allowing coalescing when COPY:ing a subreg into a 128-bit reg where the subreg was defined close above and only has one use. This is an obvious case that should not cause any problems.

This helps @atomicrmw_xchg() in atomicrmw-ops-i128.ll both on main and on "[SystemZ] Support i128 as legal type in VRs
#74625".



>From 29faa5485b5051c4a41b8d6217dd1ebc4899c41a Mon Sep 17 00:00:00 2001
From: Jonas Paulsson <paulson1 at linux.ibm.com>
Date: Sat, 9 Dec 2023 10:25:10 -0600
Subject: [PATCH] Improve for PAIR128

---
 llvm/lib/Target/SystemZ/SystemZRegisterInfo.cpp | 11 ++++++++++-
 llvm/test/CodeGen/SystemZ/atomicrmw-ops-i128.ll |  4 +---
 llvm/test/CodeGen/SystemZ/atomicrmw-xchg-07.ll  |  3 +--
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Target/SystemZ/SystemZRegisterInfo.cpp b/llvm/lib/Target/SystemZ/SystemZRegisterInfo.cpp
index 4d6b94da3a275..e822dbf586ab7 100644
--- a/llvm/lib/Target/SystemZ/SystemZRegisterInfo.cpp
+++ b/llvm/lib/Target/SystemZ/SystemZRegisterInfo.cpp
@@ -384,6 +384,7 @@ bool SystemZRegisterInfo::shouldCoalesce(MachineInstr *MI,
                                   const TargetRegisterClass *NewRC,
                                   LiveIntervals &LIS) const {
   assert (MI->isCopy() && "Only expecting COPY instructions");
+  const MachineRegisterInfo *MRI = &MI->getMF()->getRegInfo();
 
   // Coalesce anything which is not a COPY involving a subreg to/from GR128.
   if (!(NewRC->hasSuperClassEq(&SystemZ::GR128BitRegClass) &&
@@ -400,7 +401,6 @@ bool SystemZRegisterInfo::shouldCoalesce(MachineInstr *MI,
   LiveInterval &IntGR128 = LIS.getInterval(GR128Reg);
   LiveInterval &IntGRNar = LIS.getInterval(GRNarReg);
 
-  // Check that the two virtual registers are local to MBB.
   MachineBasicBlock *MBB = MI->getParent();
   MachineInstr *FirstMI_GR128 =
     LIS.getInstructionFromIndex(IntGR128.beginIndex());
@@ -408,6 +408,15 @@ bool SystemZRegisterInfo::shouldCoalesce(MachineInstr *MI,
     LIS.getInstructionFromIndex(IntGRNar.beginIndex());
   MachineInstr *LastMI_GR128 = LIS.getInstructionFromIndex(IntGR128.endIndex());
   MachineInstr *LastMI_GRNar = LIS.getInstructionFromIndex(IntGRNar.endIndex());
+
+  // Allow cases (like PAIR128) where there is just one use of the narrow
+  // source reg defined close to MI.
+  if (WideOpNo == 0 && MRI->hasOneUse(GRNarReg) &&
+      FirstMI_GRNar && FirstMI_GRNar->getParent() == MBB &&
+      std::distance(FirstMI_GRNar->getIterator(), MI->getIterator()) < 5)
+    return true;
+
+  // Check that the two virtual registers are local to MBB.
   if ((!FirstMI_GR128 || FirstMI_GR128->getParent() != MBB) ||
       (!FirstMI_GRNar || FirstMI_GRNar->getParent() != MBB) ||
       (!LastMI_GR128 || LastMI_GR128->getParent() != MBB) ||
diff --git a/llvm/test/CodeGen/SystemZ/atomicrmw-ops-i128.ll b/llvm/test/CodeGen/SystemZ/atomicrmw-ops-i128.ll
index 0e8f044680222..41b0964be05e5 100644
--- a/llvm/test/CodeGen/SystemZ/atomicrmw-ops-i128.ll
+++ b/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
diff --git a/llvm/test/CodeGen/SystemZ/atomicrmw-xchg-07.ll b/llvm/test/CodeGen/SystemZ/atomicrmw-xchg-07.ll
index b9e29599af7ee..18fa89e6ca6cb 100644
--- a/llvm/test/CodeGen/SystemZ/atomicrmw-xchg-07.ll
+++ b/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