[llvm][AArch64] Copy all operands when expanding BLR_BTI bundle

[DavidSpickett]

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.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: David Spickett (DavidSpickett)

Changes

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

[DavidSpickett]

@fhahn you made changes to RVMARKER recently, perhaps you can review here.

What I'm doing doesn't exactly match RVMARKER, and I'm not sure of the consequences of that. Also, the comment there // Skip register arguments. seems to be outdated, as it does copy the register operands.

[DavidSpickett]

@mstorsjo I tested this using the ffmpeg commands you gave on the issue, but please run it through others if you have them.

[mstorsjo]

@mstorsjo I tested this using the ffmpeg commands you gave on the issue, but please run it through others if you have them.

Thanks! If it passes make fate-checkasm it should cover all the cases I had at hand - I'll try it out in my setup as well.

[mstorsjo]

@mstorsjo I tested this using the ffmpeg commands you gave on the issue, but please run it through others if you have them.

Thanks! If it passes make fate-checkasm it should cover all the cases I had at hand - I'll try it out in my setup as well.

It seems to run fine in my full build/test as well. Looking very good, thanks!

[efriedma-quic]

LGTM

Copying all the operands makes sense; the pseudo-instruction is basically a call instruction, so all the operands should be relevant.

[DavidSpickett]

Going to land this, we can look at RVMARKER later if needed.

[mstorsjo]

Thanks for fixing this issue!

[DavidSpickett]

...and I forgot to try expensive checks. Revert incoming.