[llvm][Support] Add ARM64e support for split stacks

[Bigcheese]

ARM64e needs to use the authenticated version of blr for this call as the target function pointer is signed. This also adds the pacibsp/retab pair. This is not strictly required, but is a standard part of the prolog/epilog.

[llvmbot]

@llvm/pr-subscribers-llvm-support

Author: Michael Spencer (Bigcheese)

Changes

ARM64e needs to use the authenticated version of blr for this call as the target function pointer is signed. This also adds the pacibsp/retab pair. This is not strictly required, but is a standard part of the prolog/epilog.

---

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

1 Files Affected:

• (modified) llvm/lib/Support/ProgramStack.cpp (+11)

diff --git a/llvm/lib/Support/ProgramStack.cpp b/llvm/lib/Support/ProgramStack.cpp
index 66f74ff660b20..ce396472126cd 100644
--- a/llvm/lib/Support/ProgramStack.cpp
+++ b/llvm/lib/Support/ProgramStack.cpp
@@ -70,6 +70,9 @@ __asm__ (
     ".p2align  2\n\t"
     "_ZN4llvm17runOnNewStackImplEPvPFvS0_ES0_:\n\t"
     ".cfi_startproc\n\t"
+#if __arm64e__
+    "pacibsp\n\t"                            // sign x30 using sp
+#endif
     "mov       x16, sp\n\t"
     "sub       x0, x0, #0x20\n\t"            // subtract space from stack
     "stp       xzr, x16, [x0, #0x00]\n\t"    // save old sp
@@ -81,12 +84,20 @@ __asm__ (
     ".cfi_offset w29, -16\n\t"               // fp
 
     "mov       x0, x2\n\t"                   // Ctx is the only argument
+#if __arm64e__
+    "blraaz    x1\n\t"                       // authenticate x1 then call Fn
+#else
     "blr       x1\n\t"                       // call Fn
+#endif
 
     "ldp       x29, x30, [sp, #0x10]\n\t"    // restore fp, lr
     "ldp       xzr, x16, [sp, #0x00]\n\t"    // load old sp
     "mov       sp, x16\n\t"
+#if __arm64e__
+    "retab\n\t"                              // authenticate x30 then return
+#else
     "ret\n\t"
+#endif
     ".cfi_endproc"
 );
 #endif

[github-actions]

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff HEAD~1 HEAD --extensions cpp -- llvm/lib/Support/ProgramStack.cpp

View the diff from clang-format here.

diff --git a/llvm/lib/Support/ProgramStack.cpp b/llvm/lib/Support/ProgramStack.cpp
index ce3964721..e49a0cf81 100644
--- a/llvm/lib/Support/ProgramStack.cpp
+++ b/llvm/lib/Support/ProgramStack.cpp
@@ -65,41 +65,39 @@ void runOnNewStackImpl(void *Stack, void (*Fn)(void *), void *Ctx) __asm__(
 //
 // When adding new platforms it may be better to move to a .S file with macros
 // for dealing with platform differences.
-__asm__ (
-    ".globl  _ZN4llvm17runOnNewStackImplEPvPFvS0_ES0_\n\t"
-    ".p2align  2\n\t"
-    "_ZN4llvm17runOnNewStackImplEPvPFvS0_ES0_:\n\t"
-    ".cfi_startproc\n\t"
+__asm__(".globl  _ZN4llvm17runOnNewStackImplEPvPFvS0_ES0_\n\t"
+        ".p2align  2\n\t"
+        "_ZN4llvm17runOnNewStackImplEPvPFvS0_ES0_:\n\t"
+        ".cfi_startproc\n\t"
 #if __arm64e__
-    "pacibsp\n\t"                            // sign x30 using sp
+        "pacibsp\n\t" // sign x30 using sp
 #endif
-    "mov       x16, sp\n\t"
-    "sub       x0, x0, #0x20\n\t"            // subtract space from stack
-    "stp       xzr, x16, [x0, #0x00]\n\t"    // save old sp
-    "stp       x29, x30, [x0, #0x10]\n\t"    // save fp, lr
-    "mov       sp, x0\n\t"                   // switch to new stack
-    "add       x29, x0, #0x10\n\t"           // switch to new frame
-    ".cfi_def_cfa w29, 16\n\t"
-    ".cfi_offset w30, -8\n\t"                // lr
-    ".cfi_offset w29, -16\n\t"               // fp
-
-    "mov       x0, x2\n\t"                   // Ctx is the only argument
+        "mov       x16, sp\n\t"
+        "sub       x0, x0, #0x20\n\t"         // subtract space from stack
+        "stp       xzr, x16, [x0, #0x00]\n\t" // save old sp
+        "stp       x29, x30, [x0, #0x10]\n\t" // save fp, lr
+        "mov       sp, x0\n\t"                // switch to new stack
+        "add       x29, x0, #0x10\n\t"        // switch to new frame
+        ".cfi_def_cfa w29, 16\n\t"
+        ".cfi_offset w30, -8\n\t"  // lr
+        ".cfi_offset w29, -16\n\t" // fp
+
+        "mov       x0, x2\n\t" // Ctx is the only argument
 #if __arm64e__
-    "blraaz    x1\n\t"                       // authenticate x1 then call Fn
+        "blraaz    x1\n\t" // authenticate x1 then call Fn
 #else
-    "blr       x1\n\t"                       // call Fn
+        "blr       x1\n\t" // call Fn
 #endif
 
-    "ldp       x29, x30, [sp, #0x10]\n\t"    // restore fp, lr
-    "ldp       xzr, x16, [sp, #0x00]\n\t"    // load old sp
-    "mov       sp, x16\n\t"
+        "ldp       x29, x30, [sp, #0x10]\n\t" // restore fp, lr
+        "ldp       xzr, x16, [sp, #0x00]\n\t" // load old sp
+        "mov       sp, x16\n\t"
 #if __arm64e__
-    "retab\n\t"                              // authenticate x30 then return
+        "retab\n\t" // authenticate x30 then return
 #else
-    "ret\n\t"
+        "ret\n\t"
 #endif
-    ".cfi_endproc"
-);
+        ".cfi_endproc");
 #endif
 } // namespace llvm
 

[ahmedbougacha]

There are many situations where this is incorrect, but this seems obscure, will loudly fail, and can easily be fixed as needed, so I think LGTM for arm64e? We can make it slightly more general with small tweaks on the checks; suggested inline

cc @asl

[ahmedbougacha]

Perhaps:
__has_feature(ptrauth_returns)
here and around the retab

[ahmedbougacha]

And here:
__has_feature(ptrauth_calls)