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

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

@@ -300,6 +300,14 @@ class RegisterCoalescer : private LiveRangeEdit::Delegate {
   bool reMaterializeDef(const CoalescerPair &CP, MachineInstr *CopyMI,
                         bool &IsDefCopy);
 
+  /// Check if rematerialization should be skipped for a physical register
+  /// used as a return value. This helps enable better coalescing by avoiding
+  /// rematerialization when a copy to a return register can be coalesced.
+  /// Returns true if rematerialization should be skipped.
+  bool shouldSkipRematerializationForReturn(Register DstReg, Register SrcReg,
+                                            MachineInstr *DefMI,
+                                            MachineInstr *CopyMI);
+
   /// Return true if a copy involving a physreg should be joined.
   bool canJoinPhys(const CoalescerPair &CP);
 
@@ -1284,6 +1292,78 @@ bool RegisterCoalescer::removePartialRedundancy(const CoalescerPair &CP,
 
 /// Returns true if @p MI defines the full vreg @p Reg, as opposed to just
 /// defining a subregister.
+bool RegisterCoalescer::shouldSkipRematerializationForReturn(
+    Register DstReg, Register SrcReg, MachineInstr *DefMI,
+    MachineInstr *CopyMI) {
+  MachineBasicBlock *MBB = CopyMI->getParent();
+  if (!DstReg.isPhysical() || DefMI->getParent() != MBB || MBB->empty())
+    return false;
+
+  // Check if the last instruction is a return using DstReg
+  const MachineInstr &LastInstr = MBB->back();
+  if (!LastInstr.isReturn() || !LastInstr.readsRegister(DstReg, TRI))
+    return false;
+
+  // Exception: Allow rematerialization for zero-idiom instructions
+  // (e.g., xorps %xmm0, %xmm0) because rematerialization produces
+  // independent zero-latency instructions, which is better than copying
+  const TargetSubtargetInfo &STI = MF->getSubtarget();
+  APInt Mask;
+  if (STI.isZeroIdiom(DefMI, Mask)) {
+    LLVM_DEBUG(dbgs() << "\tAllow remat: zero-idiom instruction\n");
+    return false;
+  }
+
+  // Check for duplicate DefMI before CopyMI
+  bool HasDuplicateDef = false;
+  for (MachineBasicBlock::iterator I = MBB->begin(); &*I != CopyMI; ++I) {
+    if (&*I != DefMI && I->isIdenticalTo(*DefMI, MachineInstr::IgnoreDefs)) {
+      HasDuplicateDef = true;
+      break;
+    }
+  }
+
+  // Check if register is redefined after CopyMI
+  bool RegRedefinedAfterCopy = false;
+  for (auto I = std::next(CopyMI->getIterator()); I != MBB->end(); ++I) {
+    if (I->modifiesRegister(DstReg, TRI)) {
+      RegRedefinedAfterCopy = true;
+      break;
+    }
+  }
+
+  // Only consider skipping remat if there's no duplicate def and register
+  // is not redefined after the copy
+  if (!HasDuplicateDef && !RegRedefinedAfterCopy) {
+    // Exception: Allow remat for constant moves with restricted usage
+    if (DefMI->isMoveImmediate()) {
+      // Single use is fine - allow remat
+      if (MRI->hasOneNonDBGUse(SrcReg))
+        return false;
+
+      // Multiple uses: only allow if all uses are copies
+      bool OnlyUsedByCopies = true;
+      for (const MachineOperand &MO : MRI->use_operands(SrcReg)) {
+        const MachineInstr *UseMI = MO.getParent();
+        if (!UseMI->isCopy() && !UseMI->isSubregToReg()) {
+          OnlyUsedByCopies = false;
+          break;
+        }
+      }
+
+      if (OnlyUsedByCopies)
+        return false;
+    }
+
+    // Skip remat for return register to enable better coalescing
+    LLVM_DEBUG(dbgs() << "\tSkip remat for return register: "
+                      << printReg(DstReg, TRI) << '\n');
+    return true;
+  }
+
+  return false;
+}
+
 static bool definesFullReg(const MachineInstr &MI, Register Reg) {
   assert(!Reg.isPhysical() && "This code cannot handle physreg aliasing");
 
@@ -1326,6 +1406,9 @@ bool RegisterCoalescer::reMaterializeDef(const CoalescerPair &CP,
   if (!TII->isAsCheapAsAMove(*DefMI))
     return false;
 
+  if (shouldSkipRematerializationForReturn(DstReg, SrcReg, DefMI, CopyMI))
+    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

@@ -2173,6 +2173,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
@@ -2200,6 +2201,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:

llvm/test/CodeGen/AArch64/movi64_sve.ll

@@ -12,6 +12,7 @@ define <2 x i64> @movi_1_v2i64() {
 ; SVE-LABEL: movi_1_v2i64:
 ; SVE:       // %bb.0:
 ; SVE-NEXT:    mov z0.d, #1 // =0x1
+; SVE-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; SVE-NEXT:    ret
   ret <2 x i64> splat (i64 1)
 }
@@ -26,6 +27,7 @@ define <2 x i64> @movi_127_v2i64() {
 ; SVE-LABEL: movi_127_v2i64:
 ; SVE:       // %bb.0:
 ; SVE-NEXT:    mov z0.d, #127 // =0x7f
+; SVE-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; SVE-NEXT:    ret
   ret <2 x i64> splat (i64 127)
 }
@@ -40,6 +42,7 @@ define <2 x i64> @movi_m128_v2i64() {
 ; SVE-LABEL: movi_m128_v2i64:
 ; SVE:       // %bb.0:
 ; SVE-NEXT:    mov z0.d, #-128 // =0xffffffffffffff80
+; SVE-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; SVE-NEXT:    ret
   ret <2 x i64> splat (i64 -128)
 }
@@ -54,6 +57,7 @@ define <2 x i64> @movi_256_v2i64() {
 ; SVE-LABEL: movi_256_v2i64:
 ; SVE:       // %bb.0:
 ; SVE-NEXT:    mov z0.d, #256 // =0x100
+; SVE-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; SVE-NEXT:    ret
   ret <2 x i64> splat (i64 256)
 }
@@ -68,6 +72,7 @@ define <2 x i64> @movi_32512_v2i64() {
 ; SVE-LABEL: movi_32512_v2i64:
 ; SVE:       // %bb.0:
 ; SVE-NEXT:    mov z0.d, #32512 // =0x7f00
+; SVE-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; SVE-NEXT:    ret
   ret <2 x i64> splat (i64 32512)
 }
@@ -82,6 +87,7 @@ define <2 x i64> @movi_m32768_v2i64() {
 ; SVE-LABEL: movi_m32768_v2i64:
 ; SVE:       // %bb.0:
 ; SVE-NEXT:    mov z0.d, #-32768 // =0xffffffffffff8000
+; SVE-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; SVE-NEXT:    ret
   ret <2 x i64> splat (i64 -32768)
 }
@@ -98,6 +104,7 @@ define <4 x i32> @movi_v4i32_1() {
 ; SVE-LABEL: movi_v4i32_1:
 ; SVE:       // %bb.0:
 ; SVE-NEXT:    mov z0.d, #127 // =0x7f
+; SVE-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; SVE-NEXT:    ret
   ret <4 x i32> <i32 127, i32 0, i32 127, i32 0>
 }
@@ -112,6 +119,7 @@ define <4 x i32> @movi_v4i32_2() {
 ; SVE-LABEL: movi_v4i32_2:
 ; SVE:       // %bb.0:
 ; SVE-NEXT:    mov z0.d, #32512 // =0x7f00
+; SVE-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; SVE-NEXT:    ret
   ret <4 x i32> <i32 32512, i32 0, i32 32512, i32 0>
 }
@@ -126,6 +134,7 @@ define <8 x i16> @movi_v8i16_1() {
 ; SVE-LABEL: movi_v8i16_1:
 ; SVE:       // %bb.0:
 ; SVE-NEXT:    mov z0.d, #127 // =0x7f
+; SVE-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; SVE-NEXT:    ret
   ret <8 x i16> <i16 127, i16 0, i16 0, i16 0, i16 127, i16 0, i16 0, i16 0>
 }
@@ -140,6 +149,7 @@ define <8 x i16> @movi_v8i16_2() {
 ; SVE-LABEL: movi_v8i16_2:
 ; SVE:       // %bb.0:
 ; SVE-NEXT:    mov z0.d, #32512 // =0x7f00
+; SVE-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; SVE-NEXT:    ret
   ret <8 x i16> <i16 32512, i16 0, i16 0, i16 0, i16 32512, i16 0, i16 0, i16 0>
 }
@@ -154,6 +164,7 @@ define <16 x i8> @movi_v16i8_1() {
 ; SVE-LABEL: movi_v16i8_1:
 ; SVE:       // %bb.0:
 ; SVE-NEXT:    mov z0.d, #127 // =0x7f
+; SVE-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; SVE-NEXT:    ret
   ret <16 x i8> <i8 127, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 127, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0>
 }
@@ -168,6 +179,7 @@ define <16 x i8> @movi_v16i8_2() {
 ; SVE-LABEL: movi_v16i8_2:
 ; SVE:       // %bb.0:
 ; SVE-NEXT:    mov z0.d, #32512 // =0x7f00
+; SVE-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; SVE-NEXT:    ret
   ret <16 x i8> <i8 0, i8 127, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 127, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0>
 }

llvm/test/CodeGen/AArch64/neon-abd.ll

@@ -525,6 +525,7 @@ define <4 x i16> @combine_sabd_4h_zerosign(<4 x i16> %a, <4 x i16> %b) #0 {
 ; CHECK-LABEL: combine_sabd_4h_zerosign:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
   %a.ext = ashr <4 x i16> %a, <i16 7, i16 8, i16 9, i16 10>
   %b.ext = ashr <4 x i16> %b, <i16 11, i16 12, i16 13, i16 14>

llvm/test/CodeGen/AArch64/neon-compare-instructions.ll

@@ -2482,6 +2482,7 @@ define <2 x i32> @fcmal2xfloat(<2 x float> %A, <2 x float> %B) {
 ; CHECK-SD-LABEL: fcmal2xfloat:
 ; CHECK-SD:       // %bb.0:
 ; CHECK-SD-NEXT:    movi v0.2d, #0xffffffffffffffff
+; CHECK-SD-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-SD-NEXT:    ret
 ;
 ; CHECK-GI-LABEL: fcmal2xfloat:
@@ -2535,6 +2536,7 @@ define <2 x i32> @fcmnv2xfloat(<2 x float> %A, <2 x float> %B) {
 ; CHECK-LABEL: fcmnv2xfloat:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
   %tmp3 = fcmp false <2 x float> %A, %B
   %tmp4 = sext <2 x i1> %tmp3 to <2 x i32>

llvm/test/CodeGen/AArch64/neon-mov.ll

@@ -16,6 +16,7 @@ define <8 x i8> @movi8b_0() {
 ; CHECK-LABEL: movi8b_0:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
    ret <8 x i8> zeroinitializer
 }
@@ -48,6 +49,7 @@ define <2 x i32> @movi2s_0() {
 ; CHECK-LABEL: movi2s_0:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
    ret <2 x i32> zeroinitializer
 }
@@ -474,6 +476,7 @@ define <2 x float> @fmov2s_0() {
 ; CHECK-LABEL: fmov2s_0:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
 	ret <2 x float> zeroinitializer
 }

llvm/test/CodeGen/AArch64/remat-const-float-simd.ll

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