From 47cacb42869e7629c3183365fb660b070472f426 Mon Sep 17 00:00:00 2001
From: Sagar Maheshwari <sagar.maheshwari@amd.com>
Date: Fri, 22 Nov 2024 13:03:22 +0530
Subject: [PATCH] [DCE] Ensure reserved simplifiable registers are live across
 call boundaries.

---
 .../CodeGen/DeadMachineInstructionElim.cpp    | 30 ++++++++++++++++---
 .../AIE/aie2/live-reserved-regs-call.ll       |  8 ++---
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/CodeGen/DeadMachineInstructionElim.cpp b/llvm/lib/CodeGen/DeadMachineInstructionElim.cpp
index 64a0d2a3e5ef..4ae740e8f3f0 100644
--- a/llvm/lib/CodeGen/DeadMachineInstructionElim.cpp
+++ b/llvm/lib/CodeGen/DeadMachineInstructionElim.cpp
@@ -154,9 +154,24 @@ bool DeadMachineInstructionElimImpl::runImpl(MachineFunction &MF,
   return AnyChanges;
 }
 
+SmallVector<MCPhysReg, 16>
+getSimplifiableReservedRegs(const MachineRegisterInfo *MRI) {
+  BitVector ReservedRegs = MRI->getReservedRegs();
+  SmallVector<MCPhysReg, 16> SimplifiableReservedRegs;
+  for (MCPhysReg PhysReg : ReservedRegs.set_bits()) {
+    if (MRI->canSimplifyPhysReg(PhysReg)) {
+      SimplifiableReservedRegs.push_back(PhysReg);
+    }
+  }
+  return SimplifiableReservedRegs;
+}
+
 bool DeadMachineInstructionElimImpl::eliminateDeadMI(MachineFunction &MF) {
   bool AnyChanges = false;
 
+  SmallVector<MCPhysReg, 16> SimplifiableReservedRegs =
+      getSimplifiableReservedRegs(MRI);
+
   // Loop over all instructions in all blocks, from bottom to top, so that it's
   // more likely that chains of dependent but ultimately dead instructions will
   // be cleaned up.
@@ -165,10 +180,8 @@ bool DeadMachineInstructionElimImpl::eliminateDeadMI(MachineFunction &MF) {
 
     // Reserved registers are considered always live, so consider them as
     // live-outs for MBB. Inside MBB, dead assignments can still be detected.
-    for (MCPhysReg PhysReg : MRI->getReservedRegs().set_bits()) {
-      if (MRI->canSimplifyPhysReg(PhysReg)) {
-        LivePhysRegs.addReg(PhysReg);
-      }
+    for (MCPhysReg PhysReg : SimplifiableReservedRegs) {
+      LivePhysRegs.addReg(PhysReg);
     }
 
     // Now scan the instructions and delete dead ones, tracking physreg
@@ -187,6 +200,15 @@ bool DeadMachineInstructionElimImpl::eliminateDeadMI(MachineFunction &MF) {
       }
 
       LivePhysRegs.stepBackward(MI);
+
+      // If the instruction is a call, ensure reserved registers are made live.
+      // This also ensures that reserved registers are preserved across call
+      // boundaries.
+      if (MI.isCall()) {
+        for (MCPhysReg PhysReg : SimplifiableReservedRegs) {
+          LivePhysRegs.addReg(PhysReg);
+        }
+      }
     }
   }
 
diff --git a/llvm/test/CodeGen/AIE/aie2/live-reserved-regs-call.ll b/llvm/test/CodeGen/AIE/aie2/live-reserved-regs-call.ll
index b7245818da49..5e37e33d74da 100644
--- a/llvm/test/CodeGen/AIE/aie2/live-reserved-regs-call.ll
+++ b/llvm/test/CodeGen/AIE/aie2/live-reserved-regs-call.ll
@@ -14,12 +14,12 @@ define void @caller1() {
 ; CHECK-LABEL: caller1:
 ; CHECK:         .p2align 4
 ; CHECK-NEXT:  // %bb.0: // %entry
-; CHECK-NEXT:    nopb ; nopa ; nops ; jl #callee1; nopv
-; CHECK-NEXT:    nopx // Delay Slot 5
+; CHECK-NEXT:    nopa ; nopb ; jl #callee1; nops
+; CHECK-NEXT:    nop // Delay Slot 5
 ; CHECK-NEXT:    paddb [sp], #32 // Delay Slot 4
 ; CHECK-NEXT:    st lr, [sp, #-32] // 4-byte Folded Spill Delay Slot 3
-; CHECK-NEXT:    nop // Delay Slot 2
-; CHECK-NEXT:    nop // Delay Slot 1
+; CHECK-NEXT:    mov crSat, #1 // Delay Slot 2
+; CHECK-NEXT:    mov crRnd, #12 // Delay Slot 1
 ; CHECK-NEXT:    lda lr, [sp, #-32] // 4-byte Folded Reload
 ; CHECK-NEXT:    nop
 ; CHECK-NEXT:    nop