[llvm] [llvm][AArch64] Preserve regmask when expanding the BLR_BTI pseudo instruction (PR #73927)

[David Spickett via llvm-commits]

https://github.com/DavidSpickett updated https://github.com/llvm/llvm-project/pull/73927

>From 53be1483508b69b63c4e104f61d595a1aed146ae Mon Sep 17 00:00:00 2001
From: David Spickett <david.spickett at linaro.org>
Date: Thu, 30 Nov 2023 11:17:19 +0000
Subject: [PATCH] [llvm][AArch64] Preserve regmask when expanding the BLR_BTI
 pseudo instruction

Fixes #73787

Not doing so lead to us making use of a register after the call, which
has been clobbered by the call.

Added an MIR test that runs only the pseudo expasion pass.
---
 .../AArch64/AArch64ExpandPseudoInsts.cpp      |  1 +
 .../AArch64/blr-bti-preserves-regmask.mir     | 63 +++++++++++++++++++
 llvm/test/CodeGen/AArch64/kcfi-bti.ll         |  2 +-
 3 files changed, 65 insertions(+), 1 deletion(-)
 create 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 ac26f4d4fbe66ae..3748f671f12371f 100644
--- a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
@@ -837,6 +837,7 @@ bool AArch64ExpandPseudo::expandCALL_BTI(MachineBasicBlock &MBB,
       BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(Opc)).getInstr();
   Call->addOperand(CallTarget);
   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-regmask.mir b/llvm/test/CodeGen/AArch64/blr-bti-preserves-regmask.mir
new file mode 100644
index 000000000000000..38a4ca70eaf9786
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/blr-bti-preserves-regmask.mir
@@ -0,0 +1,63 @@
+# 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:    }
+
+# Generated from C, then simplified:
+# void _setjmp();
+# void a(int b) {
+#   _setjmp();
+#   for (; b;)
+#     ;
+# }
+
+--- |
+  define void @a() {
+    ret void
+  }
+
+  declare void @_setjmp(...)
+...
+---
+name: a
+stack:
+  - { id: 0, type: spill-slot, offset: -16, size: 8, alignment: 16, callee-saved-register: '$x19' }
+  - { id: 1, type: spill-slot, offset: -24, size: 8, alignment: 8, callee-saved-register: '$lr' }
+  - { id: 2, type: spill-slot, offset: -32, size: 8, alignment: 8, callee-saved-register: '$fp' }
+body: |
+  bb.0:
+    successors: %bb.2, %bb.1
+    liveins: $w0, $lr, $x19
+
+    frame-setup PAUTH_PROLOGUE implicit-def $lr, implicit $lr, implicit $sp
+    early-clobber $sp = frame-setup STPXpre killed $fp, killed $lr, $sp, -4 :: (store (s64) into %stack.2), (store (s64) into %stack.1)
+    frame-setup CFI_INSTRUCTION def_cfa_offset 32
+    frame-setup STRXui killed $x19, $sp, 2 :: (store (s64) into %stack.0)
+    $fp = frame-setup ADDXri $sp, 0, 0
+    frame-setup CFI_INSTRUCTION def_cfa $w29, 32
+    frame-setup CFI_INSTRUCTION offset $w19, -16
+    frame-setup CFI_INSTRUCTION offset $w30, -24
+    frame-setup CFI_INSTRUCTION offset $w29, -32
+    $w19 = ORRWrr $wzr, $w0
+    BLR_BTI @_setjmp, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp
+    CBZW killed renamable $w19, %bb.2
+
+  bb.1:
+    B %bb.1
+
+  bb.2:
+    frame-destroy CFI_INSTRUCTION def_cfa $wsp, 32
+    $x19 = frame-destroy LDRXui $sp, 2 :: (load (s64) from %stack.0)
+    early-clobber $sp, $fp, $lr = frame-destroy LDPXpost $sp, 4 :: (load (s64) from %stack.2), (load (s64) from %stack.1)
+    frame-destroy CFI_INSTRUCTION def_cfa_offset 0
+    frame-destroy PAUTH_EPILOGUE implicit-def $lr, implicit $lr, implicit $sp
+    frame-destroy CFI_INSTRUCTION restore $w19
+    frame-destroy CFI_INSTRUCTION restore $w30
+    frame-destroy CFI_INSTRUCTION restore $w29
+    RET_ReallyLR
+...
diff --git a/llvm/test/CodeGen/AArch64/kcfi-bti.ll b/llvm/test/CodeGen/AArch64/kcfi-bti.ll
index 0e8dbad1f7c7595..12cde4371e15b1a 100644
--- a/llvm/test/CodeGen/AArch64/kcfi-bti.ll
+++ b/llvm/test/CodeGen/AArch64/kcfi-bti.ll
@@ -49,7 +49,7 @@ define void @f2(ptr noundef %x) !kcfi_type !2 {
 
 ; KCFI:       BUNDLE{{.*}} {
 ; KCFI-NEXT:    KCFI_CHECK $x0, 12345678, implicit-def $x9, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
-; KCFI-NEXT:    BLR killed $x0, implicit-def $lr, implicit $sp
+; KCFI-NEXT:    BLR killed $x0, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit-def dead $lr, implicit $sp, implicit-def $sp
 ; KCFI-NEXT:    HINT 36
 ; KCFI-NEXT:  }