arm64: Handle SME's ZA in setjmp/longjmp.

AAPCS64 requires that setjmp and longjmp turn ZA off:
https://github.com/ARM-software/abi-aa/blob/main/aapcs64/aapcs64.rst#setjmp-and-longjmp

__arm_za_disable is provided as a builtin by compiler-rt and check if
SME is available or not. It is described in the AAPCS64:
https://github.com/ARM-software/abi-aa/blob/main/aapcs64/aapcs64.rst#81sme-support-routines

New tests are added to make sure this works as expected.

Test: bionic-unit-tests --gtest_filter='setjmp.*'
(cherry picked from https://googleplex-android-review.googlesource.com/q/commit:6f106f79b772c007319abec67e47912bb169c08f)
Merged-In: I4e63a26b4719fa6aebf6e87ee90dbe3d7f76769b
Change-Id: I4e63a26b4719fa6aebf6e87ee90dbe3d7f76769b
diff --git a/libc/arch-arm64/bionic/setjmp.S b/libc/arch-arm64/bionic/setjmp.S
index c408998..e94e5f4 100644
--- a/libc/arch-arm64/bionic/setjmp.S
+++ b/libc/arch-arm64/bionic/setjmp.S
@@ -114,6 +114,9 @@
   .cfi_rel_offset x0, 0
   .cfi_rel_offset x30, 8
 
+  // Commit SME's ZA lazy save. Note that the call preserves x1.
+  bl __arm_za_disable
+
   // Get the cookie and store it along with the signal flag.
   mov x0, x1
   bl __bionic_setjmp_cookie_get
@@ -183,6 +186,17 @@
 
 // void siglongjmp(sigjmp_buf env, int value);
 ENTRY_WEAK_FOR_NATIVE_BRIDGE(siglongjmp)
+  // First of all, disable SME's ZA, so that it does not interfere
+  // with anything else. Note that __arm_za_disable is guaranteed to
+  // preserve x0 and x1.
+  str x30, [sp, #-16]!
+  .cfi_adjust_cfa_offset 16
+  .cfi_rel_offset x30, 0
+  bl __arm_za_disable
+  ldr x30, [sp], #16
+  .cfi_adjust_cfa_offset -16
+  .cfi_restore x30
+
   // Check the checksum before doing anything.
   m_calculate_checksum x12, x0, x2
   ldr x2, [x0, #(_JB_CHECKSUM * 8)]
diff --git a/tests/setjmp_test.cpp b/tests/setjmp_test.cpp
index 9469285..836aadc 100644
--- a/tests/setjmp_test.cpp
+++ b/tests/setjmp_test.cpp
@@ -18,6 +18,7 @@
 
 #include <setjmp.h>
 #include <stdlib.h>
+#include <sys/auxv.h>
 #include <sys/syscall.h>
 #include <unistd.h>
 
@@ -364,3 +365,42 @@
   GTEST_SKIP() << "tests uses functions not in glibc";
 #endif
 }
+
+#if defined(__aarch64__)
+TEST(setjmp, sigsetjmp_sme) {
+  if (!(getauxval(AT_HWCAP2) & HWCAP2_SME)) {
+    GTEST_SKIP() << "SME is not enabled on device.";
+  }
+
+  uint64_t svcr, za_state;
+  sigjmp_buf jb;
+  __asm__ __volatile__(".arch_extension sme; smstart za");
+  sigsetjmp(jb, 0);
+  __asm__ __volatile__(".arch_extension sme; mrs %0, SVCR" : "=r"(svcr));
+  __asm__ __volatile__(".arch_extension sme; smstop za");  // Turn ZA off anyway.
+  za_state = svcr & 0x2UL;
+  ASSERT_EQ(0UL, za_state);
+}
+
+TEST(setjmp, siglongjmp_sme) {
+  if (!(getauxval(AT_HWCAP2) & HWCAP2_SME)) {
+    GTEST_SKIP() << "SME is not enabled on device.";
+  }
+
+  uint64_t svcr, za_state;
+  int value;
+  sigjmp_buf jb;
+  if ((value = sigsetjmp(jb, 0)) == 0) {
+    __asm__ __volatile__(".arch_extension sme; smstart za");
+    siglongjmp(jb, 789);
+    __asm__ __volatile__(".arch_extension sme; smstop za");
+    FAIL();  // Unreachable.
+  } else {
+    __asm__ __volatile__(".arch_extension sme; mrs %0, SVCR" : "=r"(svcr));
+    __asm__ __volatile__(".arch_extension sme; smstop za");  // Turn ZA off anyway.
+    za_state = svcr & 0x2UL;
+    ASSERT_EQ(789, value);
+    ASSERT_EQ(0UL, za_state);
+  }
+}
+#endif