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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-aarch64

Author: David Spickett (DavidSpickett)

<details>
<summary>Changes</summary>

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.

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


3 Files Affected:

- (modified) llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp (+4-2) 
- (added) llvm/test/CodeGen/AArch64/blr-bti-preserves-operands.mir (+24) 
- (removed) llvm/test/CodeGen/AArch64/blr-bti-preserves-regmask.mir (-23) 


``````````diff
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
-...

``````````

</details>


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