[libunwind][AArch64] Disable ZA before resuming from unwinding (on Linux)

libunwind/src/Registers.hpp

@@ -20,6 +20,11 @@
 #include "libunwind_ext.h"
 #include "shadow_stack_unwind.h"
 
+#if __has_include(<sys/auxv.h>)
+#include <sys/auxv.h>
+#define HAVE_SYS_AUXV_H
+#endif
+
 namespace libunwind {
 
 // For emulating 128-bit registers
@@ -1828,6 +1833,7 @@ inline const char *Registers_ppc64::getRegisterName(int regNum) {
 /// process.
 class _LIBUNWIND_HIDDEN Registers_arm64;
 extern "C" void __libunwind_Registers_arm64_jumpto(Registers_arm64 *);
+extern "C" bool __libunwind_Registers_arm64_za_disable();
 
 #if defined(_LIBUNWIND_USE_GCS)
 extern "C" void *__libunwind_shstk_get_jump_target() {
@@ -1837,7 +1843,7 @@ extern "C" void *__libunwind_shstk_get_jump_target() {
 
 class _LIBUNWIND_HIDDEN Registers_arm64 {
 public:
-  Registers_arm64();
+  Registers_arm64() = default;
   Registers_arm64(const void *registers);
   Registers_arm64(const Registers_arm64 &);
   Registers_arm64 &operator=(const Registers_arm64 &);
@@ -1855,7 +1861,10 @@ class _LIBUNWIND_HIDDEN Registers_arm64 {
   v128        getVectorRegister(int num) const;
   void        setVectorRegister(int num, v128 value);
   static const char *getRegisterName(int num);
-  void        jumpto() { __libunwind_Registers_arm64_jumpto(this); }
+  void jumpto() {
+    zaDisable();
+    __libunwind_Registers_arm64_jumpto(this);
+  }
   static constexpr int lastDwarfRegNum() {
     return _LIBUNWIND_HIGHEST_DWARF_REGISTER_ARM64;
   }
@@ -1908,25 +1917,43 @@ class _LIBUNWIND_HIDDEN Registers_arm64 {
 private:
   uint64_t lazyGetVG() const;
 
+  void zaDisable() const {
+    if (!_misc_registers.__has_sme)
+      return;
+    if (!__libunwind_Registers_arm64_za_disable())
+      _LIBUNWIND_ABORT("SME ZA disable failed");
+  }
+
+  static bool checkHasSME() {
+#if defined(HAVE_SYS_AUXV_H)
+    constexpr int hwcap2_sme = (1 << 23);
+    unsigned long hwcap2 = getauxval(AT_HWCAP2);
+    return (hwcap2 & hwcap2_sme) != 0;
+#endif
+    // TODO: Support other platforms.
+    return false;
+  }
+
   struct GPRs {
-    uint64_t __x[29]; // x0-x28
-    uint64_t __fp;    // Frame pointer x29
-    uint64_t __lr;    // Link register x30
-    uint64_t __sp;    // Stack pointer x31
-    uint64_t __pc;    // Program counter
-    uint64_t __ra_sign_state; // RA sign state register
+    uint64_t __x[29] = {};        // x0-x28
+    uint64_t __fp = 0;            // Frame pointer x29
+    uint64_t __lr = 0;            // Link register x30
+    uint64_t __sp = 0;            // Stack pointer x31
+    uint64_t __pc = 0;            // Program counter
+    uint64_t __ra_sign_state = 0; // RA sign state register
   };
 
   struct Misc {
-    mutable uint64_t __vg = 0; // Vector Granule
+    mutable uint32_t __vg = 0; // Vector Granule
+    bool __has_sme = checkHasSME();
   };
 
-  GPRs _registers;
+  GPRs _registers = {};
   // Currently only the lower double in 128-bit vectore registers
   // is perserved during unwinding.  We could define new register
   // numbers (> 96) which mean whole vector registers, then this
   // struct would need to change to contain whole vector registers.
-  double _vectorHalfRegisters[32];
+  double _vectorHalfRegisters[32] = {};
 
   // Miscellaneous/virtual registers. These are stored below the GPRs and FPRs
   // as they do not correspond to physical registers, so do not need to be
@@ -1971,10 +1998,6 @@ Registers_arm64::operator=(const Registers_arm64 &other) {
   return *this;
 }
 
-inline Registers_arm64::Registers_arm64() {
-  memset(static_cast<void *>(this), 0, sizeof(*this));
-}
-
 inline bool Registers_arm64::validRegister(int regNum) const {
   if (regNum == UNW_REG_IP)
     return true;

libunwind/src/UnwindRegistersSave.S

@@ -827,6 +827,54 @@ DEFINE_LIBUNWIND_FUNCTION(__unw_getcontext)
   ret
 #endif
 
+//
+// extern "C" bool __libunwind_Registers_arm64_za_disable()
+//
+// On return:
+//  success (true/false) is returned in x0
+//
+  .p2align 2
+DEFINE_LIBUNWIND_FUNCTION(__libunwind_Registers_arm64_za_disable)
+  // If TPIDR2_EL0 is null, the subroutine just disables ZA.
+  .inst 0xd53bd0b0 // mrs x16, TPIDR2_EL0
+  cbz x16, 1f
+
+  // If any of the reserved bytes in the first 16 bytes of the TPIDR2 block are
+  // nonzero, return false (libunwind will then abort).
+  ldrh w14, [x16, #10]
+  cbnz w14, 2f
+  ldr w14, [x16, #12]
+  cbnz w14, 2f
+
+  // If num_za_save_slices is zero, the subroutine just disables ZA.
+  ldrh w14, [x16, #8]
+  cbz x14, 1f
+
+  // If za_save_buffer is NULL, the subroutine just disables ZA.
+  ldr x16, [x16]
+  cbz x16, 1f
+
+  // Store ZA to za_save_buffer.
+  mov x15, xzr
+0:
+  .inst 0xe1206200 // str za[w15,0], [x16]
+  .inst 0x04305830 // addsvl x16, x16, #1
+  add x15, x15, #1
+  cmp x14, x15
+  b.ne 0b
+1:
+  // * Set TPIDR2_EL0 to null.
+  .inst 0xd51bd0bf // msr TPIDR2_EL0, xzr
+  // * Set PSTATE.ZA to 0.
+  .inst 0xd503447f // smstop za
+  // * Return true (success)
+  mov x0, #1
+  ret
+2:
+  // * Return false (failure/invalid TPIDR2 block)
+  mov x0, #0
+  ret
+
 #elif defined(__arm__) && !defined(__APPLE__)
 
 #if !defined(__ARM_ARCH_ISA_ARM)

libunwind/src/libunwind.cpp

@@ -514,4 +514,3 @@ bool logDWARF() {
 }
 
 #endif // NDEBUG
-

libunwind/test/aarch64_za_unwind.pass.cpp

@@ -0,0 +1,117 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// REQUIRES: linux && target={{aarch64-.+}}
+
+#include <libunwind.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/auxv.h>
+
+// Basic test of unwinding with SME lazy saves. This tests libunwind disables ZA
+// (and commits a lazy save of ZA) before resuming from unwinding.
+
+// Note: This test requires SME (and is setup to pass on targets without SME).
+
+static bool checkHasSME() {
+  constexpr int hwcap2_sme = (1 << 23);
+  unsigned long hwcap2 = getauxval(AT_HWCAP2);
+  return (hwcap2 & hwcap2_sme) != 0;
+}
+
+struct TPIDR2Block {
+  void *za_save_buffer;
+  uint64_t num_save_slices;
+};
+
+__attribute__((noinline)) void private_za() {
+  // Note: Lazy save active on entry to function.
+  unw_context_t context;
+  unw_cursor_t cursor;
+
+  unw_getcontext(&context);
+  unw_init_local(&cursor, &context);
+  unw_step(&cursor);
+  unw_resume(&cursor);
+}
+
+bool isZAOn() {
+  register uint64_t svcr asm("x20");
+  asm(".inst 0xd53b4254" : "=r"(svcr));
+  return (svcr & 0b10) != 0;
+}
+
+__attribute__((noinline)) void za_function_with_lazy_save() {
+  register uint64_t tmp asm("x8");
+
+  // SMSTART ZA (should zero ZA)
+  asm(".inst 0xd503457f");
+
+  // RDSVL x8, #1 (read streaming vector length)
+  asm(".inst 0x04bf5828" : "=r"(tmp));
+
+  // Allocate and fill ZA save buffer with 0xAA.
+  size_t buffer_size = tmp * tmp;
+  uint8_t *za_save_buffer = (uint8_t *)alloca(buffer_size);
+  memset(za_save_buffer, 0xAA, buffer_size);
+
+  TPIDR2Block block = {za_save_buffer, tmp};
+  tmp = reinterpret_cast<uint64_t>(&block);
+
+  // MRS TPIDR2_EL0, x8 (setup lazy save of ZA)
+  asm(".inst 0xd51bd0a8" ::"r"(tmp));
+
+  // ZA should be on before unwinding.
+  if (!isZAOn()) {
+    fprintf(stderr, __FILE__ ": fail (ZA not on before call)\n");
+    abort();
+  } else {
+    fprintf(stderr, __FILE__ ": pass (ZA on before call)\n");
+  }
+
+  private_za();
+
+  // ZA should be off after unwinding.
+  if (isZAOn()) {
+    fprintf(stderr, __FILE__ ": fail (ZA on after unwinding)\n");
+    abort();
+  } else {
+    fprintf(stderr, __FILE__ ": pass (ZA off after unwinding)\n");
+  }
+
+  // MRS x8, TPIDR2_EL0 (read TPIDR2_EL0)
+  asm(".inst 0xd53bd0a8" : "=r"(tmp));
+  // ZA should have been saved (TPIDR2_EL0 zero).
+  if (tmp != 0) {
+    fprintf(stderr, __FILE__ ": fail (TPIDR2_EL0 non-null after unwinding)\n");
+    abort();
+  } else {
+    fprintf(stderr, __FILE__ ": pass (TPIDR2_EL0 null after unwinding)\n");
+  }
+
+  // ZA (all zero) should have been saved to the buffer.
+  for (unsigned i = 0; i < buffer_size; ++i) {
+    if (za_save_buffer[i] != 0) {
+      fprintf(stderr,
+              __FILE__ ": fail (za_save_buffer non-zero after unwinding)\n");
+      abort();
+    }
+  }
+  fprintf(stderr, __FILE__ ": pass (za_save_buffer zero'd after unwinding)\n");
+}
+
+int main(int, char **) {
+  if (!checkHasSME()) {
+    fprintf(stderr, __FILE__ ": pass (no SME support)\n");
+    return 0; // Pass (SME is required for this test to run).
+  }
+  za_function_with_lazy_save();
+  return 0;
+}