[llvm] [MCP] Avoid Folding Spill Chains Over Across Source Clobbers (PR #207169)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-regalloc

Author: Jack Styles (Stylie777)

<details>
<summary>Changes</summary>

Previously, MachineCopyPropogation::eliminateSpillageCopies was able to fold COPY instructions that were part of a spill/reload chain over a different COPY instruction that would clobber the original source. This was missed, and as a result would lead to the original value of the source being lost.

To ensure the original value is retained, the uses of the COPY's outside of the chain are tracked, and then when folding is attempted, any instructions that are unexpected will ensure the Spill Chain is retained for that source, maintaining the original value.

Fixes: #<!-- -->206839

---
Full diff: https://github.com/llvm/llvm-project/pull/207169.diff


2 Files Affected:

- (modified) llvm/lib/CodeGen/MachineCopyPropagation.cpp (+27) 
- (added) llvm/test/CodeGen/AArch64/aarch64-no-fold-over-clobbered-source.mir (+33) 


``````````diff
diff --git a/llvm/lib/CodeGen/MachineCopyPropagation.cpp b/llvm/lib/CodeGen/MachineCopyPropagation.cpp
index ad33cb71b7c56..a9b3cf68cc2c3 100644
--- a/llvm/lib/CodeGen/MachineCopyPropagation.cpp
+++ b/llvm/lib/CodeGen/MachineCopyPropagation.cpp
@@ -1337,6 +1337,8 @@ void MachineCopyPropagation::eliminateSpillageCopies(MachineBasicBlock &MBB) {
   // If a COPY's Source has use or def until next COPY defines the Source,
   // we put the COPY in this set to keep property#2.
   DenseSet<const MachineInstr *> CopySourceInvalid;
+  // Track uses of a register as a source for a previous COPY
+  DenseMap<const MachineInstr *, SmallVector<const MachineInstr *>> CopySourceDefCopies;
 
   auto TryFoldSpillageCopies =
       [&, this](const SmallVectorImpl<MachineInstr *> &SC,
@@ -1376,6 +1378,26 @@ void MachineCopyPropagation::eliminateSpillageCopies(MachineBasicBlock &MBB) {
           }
         };
 
+        auto HasUnexpectedCopySourceDef = [&](const MachineInstr *Copy, const MachineInstr *ExpectedDef) {
+          auto CopySouceDefPair = CopySourceDefCopies.find(Copy);
+          if (CopySouceDefPair == CopySourceDefCopies.end())
+            return false;
+
+          for (const MachineInstr *DefCopy : CopySouceDefPair->second)
+            if (DefCopy != ExpectedDef)
+              return true;
+          return false;
+        };
+        // Do not fold across an unexpected COPY def of a chains copy's source.
+        // The paired reload may be needed to restore the original source value.
+        // More info: https://github.com/llvm/llvm-project/issues/206839
+        for (size_t I = 1; I < SC.size(); ++I) {
+          if (HasUnexpectedCopySourceDef(SC[I], SC[I - 1]))
+            return;
+          if (I + 1 > SC.size() && HasUnexpectedCopySourceDef(RC[I], RC[I + 1]))
+            return;
+        }
+
         DestSourcePair InnerMostSpillCopy =
             *isCopyInstr(*SC[0], *TII, UseCopyInstr);
         DestSourcePair OuterMostSpillCopy =
@@ -1495,6 +1517,11 @@ void MachineCopyPropagation::eliminateSpillageCopies(MachineBasicBlock &MBB) {
     }
 
     auto [Dst, Src] = getDstSrcMCRegs(*CopyOperands);
+    // Track potential unexpected COPY def's that use a source in the chain. These
+    // cannot be folded across to preserve the original source value.
+    if (MachineInstr *LastUseCopy = Tracker.findLastSeenUseInCopy(Dst, *TRI)) {
+      CopySourceDefCopies[LastUseCopy].push_back(&MI);
+    }
     // Check if we can find a pair spill-reload copy.
     LLVM_DEBUG(dbgs() << "MCP: Searching paired spill for reload: ");
     LLVM_DEBUG(MI.dump());
diff --git a/llvm/test/CodeGen/AArch64/aarch64-no-fold-over-clobbered-source.mir b/llvm/test/CodeGen/AArch64/aarch64-no-fold-over-clobbered-source.mir
new file mode 100644
index 0000000000000..4908f6f159d0a
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/aarch64-no-fold-over-clobbered-source.mir
@@ -0,0 +1,33 @@
+# RUN: llc -mtriple=aarch64 -run-pass=machine-cp -enable-spill-copy-elim -verify-machineinstrs -o - %s | FileCheck %s
+# When running MachineCopyPropagation, we need to make sure COPY's do not get removed
+# when calls outside of the chain clobber a source included in the chain.
+# See: https://github.com/llvm/llvm-project/issues/206839
+---
+name: spill_chain_copy_clobbers_source
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $w10, $w13, $w23, $w25, $x7
+
+    ; CHECK: renamable $w10 = COPY killed renamable $w13
+    ; CHECK: renamable $w13 = COPY killed renamable $w10
+
+    renamable $w11 = COPY killed renamable $w10  ; S1
+    renamable $w10 = COPY killed renamable $w13  ; S2
+
+    renamable $w13 = COPY killed renamable $w23  ; clobbers S2 source
+    renamable $w15 = COPY killed renamable $w13
+    HINT 0, implicit killed $w15
+
+    renamable $w13 = COPY killed renamable $w25  ; S3
+
+    HINT 0, implicit killed $x7
+
+    renamable $w25 = COPY killed renamable $w13  ; R1
+    renamable $w13 = COPY killed renamable $w10  ; R2
+    renamable $w10 = COPY killed renamable $w11  ; R3
+
+    HINT 0, implicit killed $w10
+    HINT 0, implicit killed $w13
+
+...

``````````

</details>


https://github.com/llvm/llvm-project/pull/207169