[llvm] [MCP] Fix bug in spill copy elimination folding chain across a call. (PR #206858)

[Valery Dmitriev via llvm-commits]

https://github.com/valerydmit updated https://github.com/llvm/llvm-project/pull/206858

>From a51a4cee7d4e0fcfa5f79e74de73af6cbcf7a8a6 Mon Sep 17 00:00:00 2001
From: Valery Dmitriev <valeryd at nvidia.com>
Date: Tue, 30 Jun 2026 15:37:43 -0700
Subject: [PATCH 1/2] [MCP] Fix bug in spill copy elimination folding chain
 across a call

The bug led to having a caller preserved register not being preserved across a call.
---
 llvm/lib/CodeGen/MachineCopyPropagation.cpp   |  2 +-
 .../machine-cp-spill-chain-across-call.mir    | 22 +++++++++++--------
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/CodeGen/MachineCopyPropagation.cpp b/llvm/lib/CodeGen/MachineCopyPropagation.cpp
index ad33cb71b7c56..6f1540f2bea8f 100644
--- a/llvm/lib/CodeGen/MachineCopyPropagation.cpp
+++ b/llvm/lib/CodeGen/MachineCopyPropagation.cpp
@@ -1499,7 +1499,7 @@ void MachineCopyPropagation::eliminateSpillageCopies(MachineBasicBlock &MBB) {
     LLVM_DEBUG(dbgs() << "MCP: Searching paired spill for reload: ");
     LLVM_DEBUG(MI.dump());
     MachineInstr *MaybeSpill =
-        Tracker.findLastSeenDefInCopy(MI, Src, *TRI, *TII, UseCopyInstr);
+        Tracker.findAvailCopy(MI, Src, *TRI, *TII, UseCopyInstr);
     bool MaybeSpillIsChained = ChainLeader.count(MaybeSpill);
     if (!MaybeSpillIsChained && MaybeSpill &&
         IsSpillReloadPair(*MaybeSpill, MI)) {
diff --git a/llvm/test/CodeGen/AArch64/machine-cp-spill-chain-across-call.mir b/llvm/test/CodeGen/AArch64/machine-cp-spill-chain-across-call.mir
index fe5fe3b0274e5..5810a55d9c359 100644
--- a/llvm/test/CodeGen/AArch64/machine-cp-spill-chain-across-call.mir
+++ b/llvm/test/CodeGen/AArch64/machine-cp-spill-chain-across-call.mir
@@ -1,15 +1,14 @@
 # RUN: llc -mtriple=aarch64 -run-pass=machine-cp -enable-spill-copy-elim -o - %s | FileCheck %s
 
-# TODO -- remove this restriction with the follow up patch as it actually fixes it.
-# UNSUPPORTED: expensive_checks
-
-# Test to show a bug in eliminateSpillageCopies.
+# eliminateSpillageCopies must not fold a spill-reload chain that straddles a
+# call if source is a caller preserved register.
 # Here the spill copy "$x23 = COPY $x16" sits *before* the call and its
 # source $x16 (caller-saved) may be clobbered by the call, while its
 # destination $x23 is callee-saved. Because $x23 stays available in the
-# CopyTracker across the call, the chain builder paired this spill with the
-# post-call reload "$x16 = COPY $x23" and folded the whole chain, which
-# led to $x16 not being preserved across the call.
+# If not checked with regmask, the chain builder could otherwise pair this
+# spill with the post-call reload "$x16 = COPY $x23" # and fold the whole chain
+# leaving $x16 not preserved across the call.
+# Check that the chain is left untouched.
 
 --- |
   declare void @callee()
@@ -33,9 +32,14 @@ body:             |
     ; The value of $x16 must still be spilled into the callee-saved $x23 before
     ; the call (this copy is erased if the chain is wrongly folded).
     ; CHECK:      $x26 = COPY {{.*}}$x23
+    ; CHECK:      $x23 = COPY {{.*}}$x16
     ; CHECK:      BL @callee
-    ; CHECK:      $x23 = COPY {{.*}}$x20
-    ; CHECK:      $x20 = COPY {{.*}}$x23
+    ; The inner spill must stay "$x16 = COPY $x20" and must not be rewritten to
+    ; target the outermost temporary as "$x23 = COPY $x20".
+    ; CHECK:      $x16 = COPY {{.*}}$x20
+    ; CHECK:      $x20 = COPY {{.*}}$x16
+    ; And $x16 must be restored from $x23 after the inner reload.
+    ; CHECK:      $x16 = COPY {{.*}}$x23
 
     renamable $x26 = COPY killed renamable $x23
     renamable $x23 = COPY killed renamable $x16

>From b5fa68789b3f718ecb6cd188030d106fd4edaf7a Mon Sep 17 00:00:00 2001
From: Valery Dmitriev <valeryd at nvidia.com>
Date: Wed, 1 Jul 2026 09:17:31 -0700
Subject: [PATCH 2/2] adjust PPC test

---
 .../CodeGen/PowerPC/mcp-elim-eviction-chain.mir   | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/llvm/test/CodeGen/PowerPC/mcp-elim-eviction-chain.mir b/llvm/test/CodeGen/PowerPC/mcp-elim-eviction-chain.mir
index 456f26ca322a0..81a9582b1995f 100644
--- a/llvm/test/CodeGen/PowerPC/mcp-elim-eviction-chain.mir
+++ b/llvm/test/CodeGen/PowerPC/mcp-elim-eviction-chain.mir
@@ -204,9 +204,10 @@ body: |
     ; CHECK: liveins: $x17, $x16, $x15, $x14, $x3
     ; CHECK-NEXT: {{  $}}
     ; CHECK-NEXT: renamable $x18 = COPY renamable $x17
-    ; CHECK-NEXT: $x17 = COPY renamable $x3
+    ; CHECK-NEXT: $x17 = COPY renamable $x14
+    ; CHECK-NEXT: renamable $x14 = COPY renamable $x3
     ; CHECK-NEXT: BL8_NOP @foo, csr_ppc64, implicit-def dead $lr8, implicit $rm, implicit-def $x3, implicit $x3
-    ; CHECK-NEXT: renamable $x3 = COPY $x17
+    ; CHECK-NEXT: renamable $x3 = COPY renamable $x14
     ; CHECK-NEXT: BLR8 implicit $lr8, implicit undef $rm, implicit $x3
     renamable $x18 = COPY renamable $x17
     renamable $x17 = COPY renamable $x16
@@ -236,11 +237,13 @@ body: |
     ; CHECK-NEXT: {{  $}}
     ; CHECK-NEXT: renamable $x21 = COPY renamable $x20
     ; CHECK-NEXT: renamable $x18 = COPY renamable $x17
-    ; CHECK-NEXT: $x17 = COPY renamable $x3
-    ; CHECK-NEXT: $x20 = COPY renamable $x4
+    ; CHECK-NEXT: renamable $x20 = COPY renamable $x19
+    ; CHECK-NEXT: $x17 = COPY renamable $x14
+    ; CHECK-NEXT: renamable $x14 = COPY renamable $x3
+    ; CHECK-NEXT: renamable $x19 = COPY renamable $x4
     ; CHECK-NEXT: BL8_NOP @foo, csr_ppc64, implicit-def dead $lr8, implicit $rm, implicit-def $x3, implicit $x3, implicit-def $x4, implicit $x4
-    ; CHECK-NEXT: renamable $x3 = COPY $x17
-    ; CHECK-NEXT: renamable $x4 = COPY $x20
+    ; CHECK-NEXT: renamable $x3 = COPY renamable $x14
+    ; CHECK-NEXT: renamable $x4 = COPY renamable $x19
     ; CHECK-NEXT: BLR8 implicit $lr8, implicit undef $rm, implicit $x3, implicit $x4
     renamable $x21 = COPY renamable $x20
     renamable $x18 = COPY renamable $x17