[llvm][AArch64] Preserve regmask when expanding the BLR_BTI pseudo instruction

[DavidSpickett]

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

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: David Spickett (DavidSpickett)

Changes

Fixes #73787

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

To fix this I've done as expandCALL_RVMARKER. With a slight simplification because some of that code doesn't seem necessary and didn't effect the test.

Added an MIR test that runs only the pseudo expasion pass.

---

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

3 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp (+7)
• (added) llvm/test/CodeGen/AArch64/blr-bti-preserves-regmask.mir (+63)
• (modified) llvm/test/CodeGen/AArch64/kcfi-bti.ll (+1-1)

diff --git a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
index ac26f4d4fbe66ae..99071ff1b19b45d 100644
--- a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
@@ -838,6 +838,13 @@ bool AArch64ExpandPseudo::expandCALL_BTI(MachineBasicBlock &MBB,
   Call->addOperand(CallTarget);
   Call->setCFIType(*MBB.getParent(), MI.getCFIType());
 
+  unsigned RegMaskStartIdx = 1;
+  while (!MI.getOperand(RegMaskStartIdx).isRegMask())
+    RegMaskStartIdx++;
+  for (const MachineOperand &MO :
+       llvm::drop_begin(MI.operands(), RegMaskStartIdx))
+    Call->addOperand(MO);
+
   MachineInstr *BTI =
       BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(AArch64::HINT))
           // BTI J so that setjmp can to BR to this.
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:  }
 

[DavidSpickett]

With this fix the assembly is:

        mov     w19, w0
        bl      _setjmp
        hint    #36
        cbz     w19, .LBB0_2

Because we now keep csr_aarch64_aapcs around (w19 is callee saved).

[mstorsjo]

Thanks for the quick fix! I'm not familiar with this area of the code so I can't comment on the implementation of the fix, but I can confirm that it does fix my original issue!

[DanielKristofKiss]

LGTM