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/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