Use shouldSkipRematerializationForReturn

Author: rez5427

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,80 @@ 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;
+    }
+    if (I->isReturn())
+      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,78 +1408,8 @@ 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 && !MBB->empty()) {
-      // Quick check: is the last instruction a return using DstReg?
-      const MachineInstr &LastInstr = MBB->back();
-      if (LastInstr.isReturn() && LastInstr.readsRegister(DstReg, TRI)) {
-        // This is a return register, perform checks
-
-        // 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");
-        } else {
-          // 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 (MachineBasicBlock::iterator I = std::next(CopyMI->getIterator());
-               I != MBB->end(); ++I) {
-            if (I->modifiesRegister(DstReg, TRI)) {
-              RegRedefinedAfterCopy = true;
-              break;
-            }
-            if (I->isReturn())
-              break;
-          }
-
-          // Skip remat only if: no duplicate def AND reg not redefined
-          if (!HasDuplicateDef && !RegRedefinedAfterCopy) {
-            // Exception: allow remat for constant moves with limited uses
-            if (DefMI->isMoveImmediate()) {
-              if (!MRI->hasOneNonDBGUse(SrcReg)) {
-                // Check 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 || MRI->use_empty(SrcReg)) {
-                  LLVM_DEBUG(dbgs() << "\tSkip remat for return register: "
-                                    << printReg(DstReg, TRI) << '\n');
-                  return false;
-                }
-              }
-            } else {
-              LLVM_DEBUG(dbgs() << "\tSkip remat for return register: "
-                                << printReg(DstReg, TRI) << '\n');
-              return false;
-            }
-          }
-        }
-      }
-    }
-  }
+  if (shouldSkipRematerializationForReturn(DstReg, SrcReg, DefMI, CopyMI))
+    return false;
 
   if (!TII->isReMaterializable(*DefMI))
     return false;

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.load.tr.gfx1250.w32.ll

@@ -365,18 +365,6 @@ define { i32, <3 x i32> } @global_load_tr6_b96_vaddr_no_align2_requirement(ptr a
 }
 
 define { i32, <3 x i32> } @global_load_tr6_b96_saddr_no_align2_requirement(ptr addrspace(1) inreg %addr, ptr addrspace(1) %use) {
-; GFX1250-LABEL: global_load_tr6_b96_saddr_no_align2_requirement:
-; GFX1250:       ; %bb.0:
-; GFX1250-NEXT:    s_wait_loadcnt_dscnt 0x0
-; GFX1250-NEXT:    s_wait_kmcnt 0x0
-; GFX1250-NEXT:    v_mov_b32_e32 v0, 0
-; GFX1250-NEXT:    global_load_tr6_b96 v[2:4], v0, s[0:1] offset:32
-; GFX1250-NEXT:    s_wait_loadcnt 0x0
-; GFX1250-NEXT:    s_wait_xcnt 0x0
-; GFX1250-NEXT:    v_dual_mov_b32 v0, 0 :: v_dual_mov_b32 v1, v2
-; GFX1250-NEXT:    v_dual_mov_b32 v2, v3 :: v_dual_mov_b32 v3, v4
-; GFX1250-NEXT:    s_set_pc_i64 s[30:31]
-;
 ; GFX1250-SDAG-LABEL: global_load_tr6_b96_saddr_no_align2_requirement:
 ; GFX1250-SDAG:       ; %bb.0:
 ; GFX1250-SDAG-NEXT:    s_wait_loadcnt_dscnt 0x0
@@ -395,6 +383,7 @@ define { i32, <3 x i32> } @global_load_tr6_b96_saddr_no_align2_requirement(ptr a
 ; GFX1250-GISEL-NEXT:    v_mov_b32_e32 v0, 0
 ; GFX1250-GISEL-NEXT:    global_load_tr6_b96 v[2:4], v0, s[0:1] offset:32
 ; GFX1250-GISEL-NEXT:    s_wait_loadcnt 0x0
+; GFX1250-GISEL-NEXT:    s_wait_xcnt 0x0
 ; GFX1250-GISEL-NEXT:    v_dual_mov_b32 v0, 0 :: v_dual_mov_b32 v1, v2
 ; GFX1250-GISEL-NEXT:    v_dual_mov_b32 v2, v3 :: v_dual_mov_b32 v3, v4
 ; GFX1250-GISEL-NEXT:    s_set_pc_i64 s[30:31]