[llvm] [llvm][AArch64] Copy all operands when expanding BLR_BTI bundle (PR #78267)

[David Spickett via llvm-commits]

https://github.com/DavidSpickett created https://github.com/llvm/llvm-project/pull/78267

Fixes #77915

Previously I based the operand copying on expandCALL_RVMARKER but did not understand it properly at the time. This lead to me dropping the arguments of the function being branched to.

This fixes that by copying all operands from the BLR_BTI to the BL/BLR without skipping anything.

I've updated the existing test by adding function arguments.

>From 8541e26d46e1109f94a8f0ab163a23b8aa1a414a Mon Sep 17 00:00:00 2001
From: David Spickett <david.spickett at linaro.org>
Date: Tue, 16 Jan 2024 11:10:24 +0000
Subject: [PATCH] [llvm][AArch64] Copy all operands when expanding BLR_BTI
 bundle

Fixes #77915

Previously I based the operand copying on expandCALL_RVMARKER
but did not understand it properly at the time. This lead to us
dropping the arguments to the function being branched to.

This fixes that by copying all operands from the BLR_BTI to the
BL/BLR without skipping anything.

I've updated the existing test by adding function arguments.
---
 .../AArch64/AArch64ExpandPseudoInsts.cpp      |  6 +++--
 .../AArch64/blr-bti-preserves-operands.mir    | 24 +++++++++++++++++++
 .../AArch64/blr-bti-preserves-regmask.mir     | 23 ------------------
 3 files changed, 28 insertions(+), 25 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/blr-bti-preserves-operands.mir
 delete mode 100644 llvm/test/CodeGen/AArch64/blr-bti-preserves-regmask.mir

diff --git a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
index bb7f4d907ffd7f..c45cb8eba89723 100644
--- a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
@@ -842,9 +842,11 @@ bool AArch64ExpandPseudo::expandCALL_BTI(MachineBasicBlock &MBB,
   unsigned Opc = CallTarget.isGlobal() ? AArch64::BL : AArch64::BLR;
   MachineInstr *Call =
       BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(Opc)).getInstr();
-  Call->addOperand(CallTarget);
+
+  for (const MachineOperand &MO : MI.operands())
+    Call->addOperand(MO);
+
   Call->setCFIType(*MBB.getParent(), MI.getCFIType());
-  Call->copyImplicitOps(*MBB.getParent(), MI);
 
   MachineInstr *BTI =
       BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(AArch64::HINT))
diff --git a/llvm/test/CodeGen/AArch64/blr-bti-preserves-operands.mir b/llvm/test/CodeGen/AArch64/blr-bti-preserves-operands.mir
new file mode 100644
index 00000000000000..3b2f10348bd773
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/blr-bti-preserves-operands.mir
@@ -0,0 +1,24 @@
+# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass=aarch64-expand-pseudo -o - %s | FileCheck %s
+
+# When expanding a BLR_BTI, we should copy all the operands to the branch in the
+# bundle. Otherwise we could end up using a register after the BL which was
+# clobbered by the function that was called, or overwriting an argument to that
+# function before we make the call.
+# CHECK:    BUNDLE implicit-def $lr, implicit-def $w30, implicit-def $sp, implicit-def $wsp, implicit $x0, implicit $w1, implicit $sp {
+# CHECK:      BL @_setjmp, $x0, $w1, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit-def dead $lr, implicit $sp, implicit-def $sp
+# CHECK:      HINT 36
+# CHECK:    }
+
+--- |
+  define void @a() {
+    ret void
+  }
+
+  declare void @_setjmp(...)
+...
+---
+name: a
+body: |
+  bb.0:
+    BLR_BTI @_setjmp, $x0, $w1, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp
+...
diff --git a/llvm/test/CodeGen/AArch64/blr-bti-preserves-regmask.mir b/llvm/test/CodeGen/AArch64/blr-bti-preserves-regmask.mir
deleted file mode 100644
index 91652c6e20c8f8..00000000000000
--- a/llvm/test/CodeGen/AArch64/blr-bti-preserves-regmask.mir
+++ /dev/null
@@ -1,23 +0,0 @@
-# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass=aarch64-expand-pseudo -o - %s | FileCheck %s
-
-# When expanding a BLR_BTI, we should keep the regmask that was attached to it.
-# Otherwise we could end up using a register after the BL which was clobbered by
-# the function that was called.
-# CHECK:    BUNDLE implicit-def $lr, implicit-def $w30, implicit-def $sp, implicit-def $wsp, implicit $sp {
-# CHECK:      BL @_setjmp, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit-def dead $lr, implicit $sp, implicit-def $sp
-# CHECK:      HINT 36
-# CHECK:    }
-
---- |
-  define void @a() {
-    ret void
-  }
-
-  declare void @_setjmp(...)
-...
----
-name: a
-body: |
-  bb.0:
-    BLR_BTI @_setjmp, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp
-...