[RegisterCoalescer] Improve register allocation for return values by limiting rematerialization

Author: rez5427

clang/test/CodeGen/msp430-abi-complex.c

@@ -202,7 +202,7 @@ float _Complex complex_float_res(void) {
 // CHECK-DAG: clr r12
 // CHECK-DAG: mov #16256, r13
   __imag__ res = -1;
-// CHECK-DAG: clr r14
+// CHECK-DAG: mov r12, r14
 // CHECK-DAG: mov #-16512, r15
   return res;
 // CHECK:     ret

llvm/lib/CodeGen/RegisterCoalescer.cpp

@@ -1326,6 +1326,69 @@ bool RegisterCoalescer::reMaterializeDef(const CoalescerPair &CP,
   if (!TII->isAsCheapAsAMove(*DefMI))
     return false;
 
+  // Skip rematerialization for physical registers used as return values within
+  // the same basic block to enable better coalescing.
+  if (DstReg.isPhysical()) {
+    MachineBasicBlock *MBB = CopyMI->getParent();
+    if (DefMI->getParent() == MBB) {
+      // Check if there's already an identical instruction before CopyMI
+      // If so, allow rematerialization to avoid redundant instructions
+      bool FoundCopy = false;
+      for (MachineInstr &MI : *MBB) {
+        if (&MI == CopyMI) {
+          FoundCopy = true;
+          continue;
+        }
+
+        // Before CopyMI: check for duplicate instructions
+        if (!FoundCopy && &MI != DefMI &&
+            MI.isIdenticalTo(*DefMI, MachineInstr::IgnoreDefs)) {
+          break; // Found duplicate, allow rematerialization
+        } else if (FoundCopy) {
+          // After CopyMI: check if used as return register
+          // If the register is redefined, it's not a return register
+          if (MI.modifiesRegister(DstReg, TRI))
+            break;
+          // If there's a return instruction that uses this register, skip remat
+          if (MI.isReturn() && MI.readsRegister(DstReg, TRI)) {
+            // Exception: if DefMI is moving a constant and SrcReg has no other
+            // uses (besides copies), rematerialization is beneficial to
+            // eliminate the def
+            if (DefMI->isMoveImmediate()) {
+              // Quick check: if there's only one use and it's this copy,
+              // definitely remat
+              if (MRI->hasOneNonDBGUse(SrcReg)) {
+                LLVM_DEBUG(dbgs()
+                           << "\tAllow remat: single use constant move\n");
+                break;
+              }
+
+              // Check all uses to see if they're all copies
+              bool OnlyUsedByCopies = true;
+              unsigned UseCount = 0;
+              for (const MachineOperand &MO : MRI->use_operands(SrcReg)) {
+                const MachineInstr *UseMI = MO.getParent();
+                if (!UseMI->isCopy() && !UseMI->isSubregToReg()) {
+                  OnlyUsedByCopies = false;
+                  break;
+                }
+                UseCount++;
+              }
+
+              if (OnlyUsedByCopies && UseCount > 0) {
+                break;
+              }
+            }
+
+            LLVM_DEBUG(dbgs() << "\tSkip remat for return register: "
+                              << printReg(DstReg, TRI) << '\n');
+            return false;
+          }
+        }
+      }
+    }
+  }
+
   if (!TII->isReMaterializable(*DefMI))
     return false;
 

llvm/test/CodeGen/AArch64/aarch64_win64cc_vararg.ll

@@ -52,9 +52,8 @@ define win64cc ptr @f9(i64 %a0, i64 %a1, i64 %a2, i64 %a3, i64 %a4, i64 %a5, i64
 ; CHECK-LABEL: f9:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    str x18, [sp, #-16]! // 8-byte Folded Spill
-; CHECK-NEXT:    add x8, sp, #24
 ; CHECK-NEXT:    add x0, sp, #24
-; CHECK-NEXT:    str x8, [sp, #8]
+; CHECK-NEXT:    str x0, [sp, #8]
 ; CHECK-NEXT:    ldr x18, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
 ;
@@ -78,9 +77,8 @@ define win64cc ptr @f8(i64 %a0, i64 %a1, i64 %a2, i64 %a3, i64 %a4, i64 %a5, i64
 ; CHECK-LABEL: f8:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    str x18, [sp, #-16]! // 8-byte Folded Spill
-; CHECK-NEXT:    add x8, sp, #16
 ; CHECK-NEXT:    add x0, sp, #16
-; CHECK-NEXT:    str x8, [sp, #8]
+; CHECK-NEXT:    str x0, [sp, #8]
 ; CHECK-NEXT:    ldr x18, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
 ;
@@ -104,10 +102,9 @@ define win64cc ptr @f7(i64 %a0, i64 %a1, i64 %a2, i64 %a3, i64 %a4, i64 %a5, i64
 ; CHECK-LABEL: f7:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    str x18, [sp, #-32]! // 8-byte Folded Spill
-; CHECK-NEXT:    add x8, sp, #24
 ; CHECK-NEXT:    add x0, sp, #24
 ; CHECK-NEXT:    str x7, [sp, #24]
-; CHECK-NEXT:    str x8, [sp, #8]
+; CHECK-NEXT:    str x0, [sp, #8]
 ; CHECK-NEXT:    ldr x18, [sp], #32 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
 ;

llvm/test/CodeGen/AArch64/arm64-neon-copy.ll

@@ -2156,6 +2156,7 @@ define <4 x i16> @concat_vector_v4i16_const() {
 ; CHECK-LABEL: concat_vector_v4i16_const:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
  %r = shufflevector <1 x i16> zeroinitializer, <1 x i16> undef, <4 x i32> zeroinitializer
  ret <4 x i16> %r
@@ -2183,6 +2184,7 @@ define <8 x i8> @concat_vector_v8i8_const() {
 ; CHECK-LABEL: concat_vector_v8i8_const:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
  %r = shufflevector <1 x i8> zeroinitializer, <1 x i8> undef, <8 x i32> zeroinitializer
  ret <8 x i8> %r

llvm/test/CodeGen/AArch64/arm64-vector-ext.ll

@@ -20,6 +20,7 @@ define void @func30(%T0_30 %v0, ptr %p1) {
 define <1 x i32> @autogen_SD7918() {
 ; CHECK-LABEL: autogen_SD7918
 ; CHECK: movi.2d v0, #0000000000000000
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT: ret
   %I29 = insertelement <1 x i1> zeroinitializer, i1 false, i32 0
   %ZE = zext <1 x i1> %I29 to <1 x i32>

llvm/test/CodeGen/AArch64/arm64-vshuffle.ll

@@ -4,6 +4,7 @@ define <8 x i1> @test1() {
 ; CHECK-LABEL: test1:
 ; CHECK:       ; %bb.0: ; %entry
 ; CHECK-NEXT:    movi.16b v0, #0
+; CHECK-NEXT:    ; kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
 entry:
   %Shuff = shufflevector <8 x i1> <i1 0, i1 1, i1 2, i1 3, i1 4, i1 5, i1 6,

llvm/test/CodeGen/AArch64/bitcast.ll

@@ -8,6 +8,7 @@ define <4 x i16> @foo1(<2 x i32> %a) {
 ; CHECK-SD-LABEL: foo1:
 ; CHECK-SD:       // %bb.0:
 ; CHECK-SD-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-SD-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-SD-NEXT:    ret
 ;
 ; CHECK-GI-LABEL: foo1:
@@ -28,6 +29,7 @@ define <4 x i16> @foo2(<2 x i32> %a) {
 ; CHECK-SD-LABEL: foo2:
 ; CHECK-SD:       // %bb.0:
 ; CHECK-SD-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-SD-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-SD-NEXT:    ret
 ;
 ; CHECK-GI-LABEL: foo2:

llvm/test/CodeGen/AArch64/combine-mul.ll

@@ -18,6 +18,7 @@ define <4 x i1> @PR48683_vec(<4 x i32> %x) {
 ; CHECK-LABEL: PR48683_vec:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
   %a = mul <4 x i32> %x, %x
   %b = and <4 x i32> %a, <i32 2, i32 2, i32 2, i32 2>
@@ -29,6 +30,7 @@ define <4 x i1> @PR48683_vec_undef(<4 x i32> %x) {
 ; CHECK-LABEL: PR48683_vec_undef:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
   %a = mul <4 x i32> %x, %x
   %b = and <4 x i32> %a, <i32 2, i32 2, i32 2, i32 undef>

llvm/test/CodeGen/AArch64/ext-narrow-index.ll

@@ -251,6 +251,7 @@ define <8 x i8> @i8_zero_off22(<16 x i8> %arg1) {
 ; CHECK-SD-LABEL: i8_zero_off22:
 ; CHECK-SD:       // %bb.0: // %entry
 ; CHECK-SD-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-SD-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-SD-NEXT:    ret
 ;
 ; CHECK-GISEL-LABEL: i8_zero_off22:
@@ -302,6 +303,7 @@ define <4 x i16> @i16_zero_off8(<8 x i16> %arg1) {
 ; CHECK-LABEL: i16_zero_off8:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
 entry:
   %shuffle = shufflevector <8 x i16> %arg1, <8 x i16> zeroinitializer, <4 x i32> <i32 8, i32 9, i32 10, i32 11>
@@ -346,6 +348,7 @@ define <2 x i32> @i32_zero_off4(<4 x i32> %arg1) {
 ; CHECK-LABEL: i32_zero_off4:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
 entry:
   %shuffle = shufflevector <4 x i32> %arg1, <4 x i32> zeroinitializer, <2 x i32> <i32 4, i32 5>

llvm/test/CodeGen/AArch64/fast-isel-const-float.ll

@@ -9,6 +9,7 @@ define float @select_fp_const() {
 ; GISEL-LABEL: select_fp_const:
 ; GISEL:       // %bb.0: // %entry
 ; GISEL-NEXT:    movi v0.2s, #79, lsl #24
+; GISEL-NEXT:    // kill: def $s0 killed $s0 killed $d0
 ; GISEL-NEXT:    ret
 ;
 ; FISEL-LABEL: select_fp_const: