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

[David Spickett via llvm-commits]

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

This reverts commit 955417ade2648c2b1a4e5f0be697f61570590a88.

The problem with the previous version was that the bundle instruction had arguments like "target arg1 arg2". When it's expanded we produced a BL or BLR which can only accept one argument, the target of the branch.

Now I realise why expandCALL_RVMARKER has to copy them in mutiple steps. The operands for the called function need to be changed to implicit arguments of the branch instruction.

* Copy the branch target.
* Copy all register operands, marking them as implicit.
* Copy any other operands without modifying them.

Prior to any attempt to fix #77915:
BL @_setjmp, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit-def dead $lr, implicit $sp, implicit-def $sp

Which is dropping the use of the arguments for the called function.

My first fix attempt produced:
BL @_setjmp, $x0, $w1, <regmask $fp ...>, implicit-def $lr, implicit $sp, implicit-def dead $lr, implicit $sp, implicit-def $sp

It copied the arguments but as explicit arguments to the BL which only expects 1, failing verification.

With this new change we produce:
BL @_setjmp, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit $x0, implicit $w1, implicit-def dead $lr, implicit $sp, implicit-def $sp

Note specifically the added "implicit $x0, implicit $w1". So BL only has 1 explicit argument, but the arguments to the function are still used.

>From 559662850ff99139454c57377ed4c505b1cb2993 Mon Sep 17 00:00:00 2001
From: David Spickett <david.spickett at linaro.org>
Date: Fri, 19 Jan 2024 12:07:06 +0000
Subject: [PATCH] Reland "[llvm][AArch64] Copy all operands when expanding
 BLR_BTI bundle (#78267)"

This reverts commit 955417ade2648c2b1a4e5f0be697f61570590a88.

The problem with the previous version was that the bundle instruction had arguments
like "target arg1 arg2". When it's expanded we produced a BL or BLR which can only
accept one argument, the target of the branch.

Now I realise why expandCALL_RVMARKER has to copy them in mutiple steps.
The operands for the called function need to be changed to implicit arguments
of the branch instruction.

* Copy the branch target.
* Copy all register operands, marking them as implicit.
* Copy any other operands without modifying them.

Prior to any attempt to fix #77915:
BL @_setjmp, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit-def dead $lr, implicit $sp, implicit-def $sp

Which is dropping the use of the arguments for the called function.

My first fix attempt produced:
BL @_setjmp, $x0, $w1, <regmask $fp ...>, implicit-def $lr, implicit $sp, implicit-def dead $lr, implicit $sp, implicit-def $sp

It copied the arguments but as explicit arguments to the BL which only expects 1, failing verification.

With this new change we produce:
BL @_setjmp, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit $x0, implicit $w1, implicit-def dead $lr, implicit $sp, implicit-def $sp

Note specifically the added "implicit $x0, implicit $w1". So BL only has 1 explicit argument, but the arguments
to the function are still used.
---
 .../AArch64/AArch64ExpandPseudoInsts.cpp      | 18 +++++++++++-
 .../AArch64/blr-bti-preserves-operands.mir    | 28 +++++++++++++++++++
 .../AArch64/blr-bti-preserves-regmask.mir     | 23 ---------------
 3 files changed, 45 insertions(+), 24 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..352c61d48e2fff 100644
--- a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
@@ -843,8 +843,24 @@ bool AArch64ExpandPseudo::expandCALL_BTI(MachineBasicBlock &MBB,
   MachineInstr *Call =
       BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(Opc)).getInstr();
   Call->addOperand(CallTarget);
+
+  // 1 because we already added the branch target above.
+  unsigned RegMaskStartIdx = 1;
+  // The branch is BL <target>, so we cannot attach the arguments of the called
+  // function to it. Those must be added as implicitly used by the branch.
+  while (!MI.getOperand(RegMaskStartIdx).isRegMask()) {
+    auto MOP = MI.getOperand(RegMaskStartIdx);
+    assert(MOP.isReg() && "can only add register operands");
+    Call->addOperand(MachineOperand::CreateReg(
+        MOP.getReg(), /*Def=*/false, /*Implicit=*/true, /*isKill=*/false,
+        /*isDead=*/false, /*isUndef=*/MOP.isUndef()));
+    RegMaskStartIdx++;
+  }
+  for (const MachineOperand &MO :
+       llvm::drop_begin(MI.operands(), RegMaskStartIdx))
+    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..760ae4794e3043
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/blr-bti-preserves-operands.mir
@@ -0,0 +1,28 @@
+# 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 take the branch.
+#
+# The arguments to the call must become implicit arguments, because the branch
+# only expects to get 1 explicit operand which is the branch target.
+
+# CHECK:    BUNDLE implicit-def $lr, implicit-def $w30, implicit-def $sp, implicit-def $wsp, implicit $sp, implicit $x0, implicit $w1 {
+# CHECK:      BL @_setjmp, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit $x0, implicit $w1, 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
-...