sme: Enable SME registers and pseudo-registers

The SME (Scalable Matrix Extension) [1] exposes a new matrix register ZA with
variable sizes.  It also exposes a new mode called streaming mode.

Similarly to SVE, the ZA register size is dictated by a vector length, but the
SME vector length is called streaming vetor length. The total size for
ZA in a given moment is svl x svl.

In streaming mode, the SVE registers have their sizes based on svl rather than
the regular vector length (vl).

The feature detection is controlled by the HWCAP2_SME bit, but actual support
should be validated by attempting a ptrace call for one of the new register
sets: NT_ARM_ZA and NT_ARM_SSVE.

Due to its large size, the ZA register is exposed as a vector of bytes, but we
introduce a number of pseudo-registers that gives various different views
into the ZA contents. These can be arranged in a couple categories: tiles and
tile slices.

Tiles are matrices the same size or smaller than ZA.  Tile slices are vectors
which map to ZA's rows/columns in different ways.

A new dynamic target description is provided containing the ZA register, the SVG
register and the SVCR register.  The size of ZA, like the SVE vector registers,
is based on the vector length register SVG (VG for SVE).

This patch enables SME register support for gdb.

[1] https://community.arm.com/arm-community-blogs/b/architectures-and-processors-blog/posts/scalable-matrix-extension-armv9-a-architecture

Co-Authored-By: Ezra Sitorus <ezra.sitorus@arm.com>
Reviewed-by: Thiago Jung Bauermann <thiago.bauermann@linaro.org>

3 files changed, 703 insertions, 34 deletions

diff --git a/gdb/nat/aarch64-scalable-linux-ptrace.c b/gdb/nat/aarch64-scalable-linux-ptrace.c
index 6eea910..d3a50ed 100644
--- a/gdb/nat/aarch64-scalable-linux-ptrace.c
+++ b/gdb/nat/aarch64-scalable-linux-ptrace.c
@@ -1,5 +1,4 @@
-/* Common target dependent routines for AArch64 Scalable Extensions
-   (SVE/SME).
+/* Common native Linux code for the AArch64 scalable extensions: SVE and SME.
 
    Copyright (C) 2018-2023 Free Software Foundation, Inc.
 
@@ -28,6 +27,193 @@
 #include "gdbsupport/common-regcache.h"
 #include "gdbsupport/byte-vector.h"
 #include <endian.h>
+#include "arch/aarch64-scalable-linux.h"
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+bool
+aarch64_has_sve_state (int tid)
+{
+  struct user_sve_header header;
+
+  if (!read_sve_header (tid, header))
+    return false;
+
+  if ((header.flags & SVE_PT_REGS_SVE) == 0)
+    return false;
+
+  if (sizeof (header) == header.size)
+    return false;
+
+  return true;
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+bool
+aarch64_has_ssve_state (int tid)
+{
+  struct user_sve_header header;
+
+  if (!read_ssve_header (tid, header))
+    return false;
+
+  if ((header.flags & SVE_PT_REGS_SVE) == 0)
+    return false;
+
+  if (sizeof (header) == header.size)
+    return false;
+
+  return true;
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+bool
+aarch64_has_za_state (int tid)
+{
+  struct user_za_header header;
+
+  if (!read_za_header (tid, header))
+    return false;
+
+  if (sizeof (header) == header.size)
+    return false;
+
+  return true;
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+bool
+read_sve_header (int tid, struct user_sve_header &header)
+{
+  struct iovec iovec;
+
+  iovec.iov_len = sizeof (header);
+  iovec.iov_base = &header;
+
+  if (ptrace (PTRACE_GETREGSET, tid, NT_ARM_SVE, &iovec) < 0)
+    {
+      /* SVE is not supported.  */
+      return false;
+    }
+  return true;
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+bool
+write_sve_header (int tid, const struct user_sve_header &header)
+{
+  struct iovec iovec;
+
+  iovec.iov_len = sizeof (header);
+  iovec.iov_base = (void *) &header;
+
+  if (ptrace (PTRACE_SETREGSET, tid, NT_ARM_SVE, &iovec) < 0)
+    {
+      /* SVE is not supported.  */
+      return false;
+    }
+  return true;
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+bool
+read_ssve_header (int tid, struct user_sve_header &header)
+{
+  struct iovec iovec;
+
+  iovec.iov_len = sizeof (header);
+  iovec.iov_base = &header;
+
+  if (ptrace (PTRACE_GETREGSET, tid, NT_ARM_SSVE, &iovec) < 0)
+    {
+      /* SSVE is not supported.  */
+      return false;
+    }
+  return true;
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+bool
+write_ssve_header (int tid, const struct user_sve_header &header)
+{
+  struct iovec iovec;
+
+  iovec.iov_len = sizeof (header);
+  iovec.iov_base = (void *) &header;
+
+  if (ptrace (PTRACE_SETREGSET, tid, NT_ARM_SSVE, &iovec) < 0)
+    {
+      /* SSVE is not supported.  */
+      return false;
+    }
+  return true;
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+bool
+read_za_header (int tid, struct user_za_header &header)
+{
+  struct iovec iovec;
+
+  iovec.iov_len = sizeof (header);
+  iovec.iov_base = &header;
+
+  if (ptrace (PTRACE_GETREGSET, tid, NT_ARM_ZA, &iovec) < 0)
+    {
+      /* ZA is not supported.  */
+      return false;
+    }
+  return true;
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+bool
+write_za_header (int tid, const struct user_za_header &header)
+{
+  struct iovec iovec;
+
+  iovec.iov_len = sizeof (header);
+  iovec.iov_base = (void *) &header;
+
+  if (ptrace (PTRACE_SETREGSET, tid, NT_ARM_ZA, &iovec) < 0)
+    {
+      /* ZA is not supported.  */
+      return false;
+    }
+  return true;
+}
+
+/* Given VL, the streaming vector length for SME, return true if it is valid
+   and false otherwise.  */
+
+static bool
+aarch64_sme_vl_valid (size_t vl)
+{
+  return (vl == 16 || vl == 32 || vl == 64 || vl == 128 || vl == 256);
+}
+
+/* Given VL, the vector length for SVE, return true if it is valid and false
+   otherwise.  SVE_state is true when the check is for the SVE register set.
+   Otherwise the check is for the SSVE register set.  */
+
+static bool
+aarch64_sve_vl_valid (const bool sve_state, size_t vl)
+{
+  if (sve_state)
+    return sve_vl_valid (vl);
+
+  /* We have an active SSVE state, where the valid vector length values are
+     more restrictive.  */
+  return aarch64_sme_vl_valid (vl);
+}
 
 /* See nat/aarch64-scalable-linux-ptrace.h.  */
 
@@ -36,23 +222,25 @@ aarch64_sve_get_vq (int tid)
 {
   struct iovec iovec;
   struct user_sve_header header;
-
   iovec.iov_len = sizeof (header);
   iovec.iov_base = &header;
 
-  /* Ptrace gives the vector length in bytes.  Convert it to VQ, the number of
-     128bit chunks in a Z register.  We use VQ because 128bits is the minimum
-     a Z register can increase in size.  */
-
-  if (ptrace (PTRACE_GETREGSET, tid, NT_ARM_SVE, &iovec) < 0)
+  /* Figure out which register set to use for the request.  The vector length
+     for SVE can be different from the vector length for SSVE.  */
+  bool has_sve_state = !aarch64_has_ssve_state (tid);
+  if (ptrace (PTRACE_GETREGSET, tid, has_sve_state? NT_ARM_SVE : NT_ARM_SSVE,
+	      &iovec) < 0)
     {
       /* SVE is not supported.  */
       return 0;
     }
 
+  /* Ptrace gives the vector length in bytes.  Convert it to VQ, the number of
+     128bit chunks in a Z register.  We use VQ because 128 bits is the minimum
+     a Z register can increase in size.  */
   uint64_t vq = sve_vq_from_vl (header.vl);
 
-  if (!sve_vl_valid (header.vl))
+  if (!aarch64_sve_vl_valid (has_sve_state, header.vl))
     {
       warning (_("Invalid SVE state from kernel; SVE disabled."));
       return 0;
@@ -72,15 +260,20 @@ aarch64_sve_set_vq (int tid, uint64_t vq)
   iovec.iov_len = sizeof (header);
   iovec.iov_base = &header;
 
-  if (ptrace (PTRACE_GETREGSET, tid, NT_ARM_SVE, &iovec) < 0)
+  /* Figure out which register set to use for the request.  The vector length
+     for SVE can be different from the vector length for SSVE.  */
+  bool has_sve_state = !aarch64_has_ssve_state (tid);
+  if (ptrace (PTRACE_GETREGSET, tid, has_sve_state? NT_ARM_SVE : NT_ARM_SSVE,
+	      &iovec) < 0)
     {
-      /* SVE is not supported.  */
+      /* SVE/SSVE is not supported.  */
       return false;
     }
 
   header.vl = sve_vl_from_vq (vq);
 
-  if (ptrace (PTRACE_SETREGSET, tid, NT_ARM_SVE, &iovec) < 0)
+  if (ptrace (PTRACE_SETREGSET, tid, has_sve_state? NT_ARM_SVE : NT_ARM_SSVE,
+	      &iovec) < 0)
     {
       /* Vector length change failed.  */
       return false;
@@ -120,13 +313,108 @@ aarch64_sve_set_vq (int tid, struct reg_buffer_common *reg_buf)
 
 /* See nat/aarch64-scalable-linux-ptrace.h.  */
 
+uint64_t
+aarch64_za_get_svq (int tid)
+{
+  struct user_za_header header;
+  if (!read_za_header (tid, header))
+    return 0;
+
+  uint64_t vq = sve_vq_from_vl (header.vl);
+
+  if (!aarch64_sve_vl_valid (false, header.vl))
+    {
+      warning (_("Invalid ZA state from kernel; ZA disabled."));
+      return 0;
+    }
+
+  return vq;
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+bool
+aarch64_za_set_svq (int tid, uint64_t vq)
+{
+  struct iovec iovec;
+
+  /* Read the NT_ARM_ZA header.  */
+  struct user_za_header header;
+  if (!read_za_header (tid, header))
+    {
+      /* ZA is not supported.  */
+      return false;
+    }
+
+  /* If the size is the correct one already, don't update it.  If we do
+     update the streaming vector length, we will invalidate the register
+     state for ZA, and we do not want that.  */
+  if (header.vl == sve_vl_from_vq (vq))
+    return true;
+
+  /* The streaming vector length is about to get updated.  Set the new value
+     in the NT_ARM_ZA header and adjust the size as well.  */
+
+  header.vl = sve_vl_from_vq (vq);
+  header.size = sizeof (struct user_za_header);
+
+  /* Update the NT_ARM_ZA register set with the new streaming vector
+     length.  */
+  iovec.iov_len = sizeof (header);
+  iovec.iov_base = &header;
+
+  if (ptrace (PTRACE_SETREGSET, tid, NT_ARM_ZA, &iovec) < 0)
+    {
+      /* Streaming vector length change failed.  */
+      return false;
+    }
+
+  /* At this point we have successfully adjusted the streaming vector length
+     for the NT_ARM_ZA register set, and it should have no payload
+     (no ZA state).  */
+
+  return true;
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+bool
+aarch64_za_set_svq (int tid, const struct reg_buffer_common *reg_buf,
+		    int svg_regnum)
+{
+  uint64_t reg_svg = 0;
+
+  /* The svg register may not be valid if we've not collected any value yet.
+     This can happen, for example,  if we're restoring the regcache after an
+     inferior function call, and the svg register comes after the Z
+     registers.  */
+  if (reg_buf->get_register_status (svg_regnum) != REG_VALID)
+    {
+      /* If svg is not available yet, fetch it from ptrace.  The svg value from
+	 ptrace is likely the correct one.  */
+      uint64_t svq = aarch64_za_get_svq (tid);
+
+      /* If something went wrong, just bail out.  */
+      if (svq == 0)
+	return false;
+
+      reg_svg = sve_vg_from_vq (svq);
+    }
+  else
+    reg_buf->raw_collect (svg_regnum, &reg_svg);
+
+  return aarch64_za_set_svq (tid, sve_vq_from_vg (reg_svg));
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
 gdb::byte_vector
 aarch64_fetch_sve_regset (int tid)
 {
   uint64_t vq = aarch64_sve_get_vq (tid);
 
   if (vq == 0)
-    perror_with_name (_("Unable to fetch SVE vector length"));
+    perror_with_name (_("Unable to fetch SVE/SSVE vector length"));
 
   /* A ptrace call with NT_ARM_SVE will return a header followed by either a
      dump of all the SVE and FP registers, or an fpsimd structure (identical to
@@ -139,8 +427,10 @@ aarch64_fetch_sve_regset (int tid)
   iovec.iov_base = sve_state.data ();
   iovec.iov_len = sve_state.size ();
 
-  if (ptrace (PTRACE_GETREGSET, tid, NT_ARM_SVE, &iovec) < 0)
-    perror_with_name (_("Unable to fetch SVE registers"));
+  bool has_sve_state = !aarch64_has_ssve_state (tid);
+  if (ptrace (PTRACE_GETREGSET, tid, has_sve_state? NT_ARM_SVE : NT_ARM_SSVE,
+	      &iovec) < 0)
+    perror_with_name (_("Unable to fetch SVE/SSVE registers"));
 
   return sve_state;
 }
@@ -155,8 +445,82 @@ aarch64_store_sve_regset (int tid, const gdb::byte_vector &sve_state)
   iovec.iov_base = (void *) sve_state.data ();
   iovec.iov_len = sve_state.size ();
 
-  if (ptrace (PTRACE_SETREGSET, tid, NT_ARM_SVE, &iovec) < 0)
-    perror_with_name (_("Unable to store SVE registers"));
+  bool has_sve_state = !aarch64_has_ssve_state (tid);
+  if (ptrace (PTRACE_SETREGSET, tid, has_sve_state? NT_ARM_SVE : NT_ARM_SSVE,
+	      &iovec) < 0)
+    perror_with_name (_("Unable to store SVE/SSVE registers"));
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+gdb::byte_vector
+aarch64_fetch_za_regset (int tid)
+{
+  struct user_za_header header;
+  if (!read_za_header (tid, header))
+    error (_("Failed to read NT_ARM_ZA header."));
+
+  if (!aarch64_sme_vl_valid (header.vl))
+    error (_("Found invalid vector length for NT_ARM_ZA."));
+
+  struct iovec iovec;
+  iovec.iov_len = header.size;
+  gdb::byte_vector za_state (header.size);
+  iovec.iov_base = za_state.data ();
+
+  if (ptrace (PTRACE_GETREGSET, tid, NT_ARM_ZA, &iovec) < 0)
+    perror_with_name (_("Failed to fetch NT_ARM_ZA register set."));
+
+  return za_state;
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+void
+aarch64_store_za_regset (int tid, const gdb::byte_vector &za_state)
+{
+  struct iovec iovec;
+  /* We need to cast from (const void *) here.  */
+  iovec.iov_base = (void *) za_state.data ();
+  iovec.iov_len = za_state.size ();
+
+  if (ptrace (PTRACE_SETREGSET, tid, NT_ARM_ZA, &iovec) < 0)
+    perror_with_name (_("Failed to write to the NT_ARM_ZA register set."));
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+void
+aarch64_initialize_za_regset (int tid)
+{
+  /* First fetch the NT_ARM_ZA header so we can fetch the streaming vector
+     length.  */
+  struct user_za_header header;
+  if (!read_za_header (tid, header))
+    error (_("Failed to read NT_ARM_ZA header."));
+
+  /* The vector should be default-initialized to zero, and we should account
+     for the payload as well.  */
+  std::vector<gdb_byte> za_new_state (ZA_PT_SIZE (sve_vq_from_vl (header.vl)));
+
+  /* Adjust the header size since we are adding the initialized ZA
+     payload.  */
+  header.size = ZA_PT_SIZE (sve_vq_from_vl (header.vl));
+
+  /* Overlay the modified header onto the new ZA state.  */
+  const gdb_byte *base = (gdb_byte *) &header;
+  memcpy (za_new_state.data (), base, sizeof (user_za_header));
+
+  /* Set the ptrace request up and update the NT_ARM_ZA register set.  */
+  struct iovec iovec;
+  iovec.iov_len = za_new_state.size ();
+  iovec.iov_base = za_new_state.data ();
+
+  if (ptrace (PTRACE_SETREGSET, tid, NT_ARM_ZA, &iovec) < 0)
+    perror_with_name (_("Failed to initialize the NT_ARM_ZA register set."));
+
+  /* The NT_ARM_ZA register set should now contain a zero-initialized ZA
+     payload.  */
 }
 
 /* If we are running in BE mode, byteswap the contents
@@ -451,3 +815,177 @@ aarch64_sve_regs_copy_from_reg_buf (int tid,
      passed on to ptrace.  */
   aarch64_store_sve_regset (tid, new_state);
 }
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+void
+aarch64_za_regs_copy_to_reg_buf (int tid, struct reg_buffer_common *reg_buf,
+				 int za_regnum, int svg_regnum,
+				 int svcr_regnum)
+{
+  /* Fetch the current ZA state from the thread.  */
+  gdb::byte_vector za_state = aarch64_fetch_za_regset (tid);
+
+  /* Sanity check.  */
+  gdb_assert (!za_state.empty ());
+
+  char *base = (char *) za_state.data ();
+  struct user_za_header *header = (struct user_za_header *) base;
+
+  /* If we have ZA state, read it.  Otherwise, make the contents of ZA
+     in the register cache all zeroes.  This is how we present the ZA
+     state when it is not initialized.  */
+  uint64_t svcr_value = 0;
+  if (aarch64_has_za_state (tid))
+    {
+      /* Sanity check the data in the header.  */
+      if (!sve_vl_valid (header->vl)
+	  || ZA_PT_SIZE (sve_vq_from_vl (header->vl)) != header->size)
+	{
+	  error (_("Found invalid streaming vector length in NT_ARM_ZA"
+		   " register set"));
+	}
+
+      reg_buf->raw_supply (za_regnum, base + ZA_PT_ZA_OFFSET);
+      svcr_value |= SVCR_ZA_BIT;
+    }
+  else
+    {
+      size_t za_bytes = header->vl * header->vl;
+      gdb_byte za_zeroed[za_bytes];
+      memset (za_zeroed, 0, za_bytes);
+      reg_buf->raw_supply (za_regnum, za_zeroed);
+    }
+
+  /* Handle the svg and svcr registers separately.  We need to calculate
+     their values manually, as the Linux Kernel doesn't expose those
+     explicitly.  */
+  svcr_value |= aarch64_has_ssve_state (tid)? SVCR_SM_BIT : 0;
+  uint64_t svg_value = sve_vg_from_vl (header->vl);
+
+  /* Update the contents of svg and svcr registers.  */
+  reg_buf->raw_supply (svg_regnum, &svg_value);
+  reg_buf->raw_supply (svcr_regnum, &svcr_value);
+
+  /* The register buffer should now contain the updated copy of the NT_ARM_ZA
+     state.  */
+}
+
+/* See nat/aarch64-scalable-linux-ptrace.h.  */
+
+void
+aarch64_za_regs_copy_from_reg_buf (int tid,
+				   struct reg_buffer_common *reg_buf,
+				   int za_regnum, int svg_regnum,
+				   int svcr_regnum)
+{
+  /* REG_BUF contains the updated ZA state.  We need to extract that state
+     and write it to the thread TID.  */
+
+
+  /* First check if there is a change to the streaming vector length.  Two
+     outcomes are possible here:
+
+     1 - The streaming vector length in the register cache differs from the
+	 one currently on the thread state.  This means that we will need to
+	 update the NT_ARM_ZA register set to reflect the new streaming vector
+	 length.
+
+     2 - The streaming vector length in the register cache is the same as in
+	 the thread state.  This means we do not need to update the NT_ARM_ZA
+	 register set for a new streaming vector length, and we only need to
+	 deal with changes to za, svg and svcr.
+
+     None of the two possibilities above imply that the ZA state actually
+     exists.  They only determine what needs to be done with any ZA content
+     based on the state of the streaming vector length.  */
+
+  /* First fetch the NT_ARM_ZA header so we can fetch the streaming vector
+     length.  */
+  struct user_za_header header;
+  if (!read_za_header (tid, header))
+    error (_("Failed to read NT_ARM_ZA header."));
+
+  /* Fetch the current streaming vector length.  */
+  uint64_t old_svg = sve_vg_from_vl (header.vl);
+
+  /* Fetch the (potentially) new streaming vector length.  */
+  uint64_t new_svg;
+  reg_buf->raw_collect (svg_regnum, &new_svg);
+
+  /* Did the streaming vector length change?  */
+  bool svg_changed = new_svg != old_svg;
+
+  /* First store the streaming vector length to the thread.  This is done
+     first to ensure the ptrace buffers read from the kernel are the correct
+     size.  If the streaming vector length is the same as the current one, it
+     won't be updated.  */
+  if (!aarch64_za_set_svq (tid, reg_buf, svg_regnum))
+    error (_("Unable to set svg register"));
+
+  bool has_za_state = aarch64_has_za_state (tid);
+
+  size_t za_bytes = sve_vl_from_vg (old_svg) * sve_vl_from_vg (old_svg);
+  gdb_byte za_zeroed[za_bytes];
+  memset (za_zeroed, 0, za_bytes);
+
+  /* If the streaming vector length changed, zero out the contents of ZA in
+     the register cache.  Otherwise, we will need to update the ZA contents
+     in the thread with the ZA contents from the register cache, and they will
+     differ in size.  */
+  if (svg_changed)
+    reg_buf->raw_supply (za_regnum, za_zeroed);
+
+  /* When we update svg, we don't automatically initialize the ZA buffer.  If
+     we have no ZA state and the ZA register contents in the register cache are
+     zero, just return and leave the ZA register cache contents as zero.  */
+  if (!has_za_state
+      && reg_buf->raw_compare (za_regnum, za_zeroed, 0))
+    {
+      /* No ZA state in the thread or in the register cache.  This was likely
+	 just an adjustment of the streaming vector length.  Let this fall
+	 through and update svcr and svg in the register cache.  */
+    }
+  else
+    {
+      /* If there is no ZA state but the register cache contains ZA data, we
+	 need to initialize the ZA data through ptrace.  First we initialize
+	 all the bytes of ZA to zero.  */
+      if (!has_za_state
+	  && !reg_buf->raw_compare (za_regnum, za_zeroed, 0))
+	aarch64_initialize_za_regset (tid);
+
+      /* From this point onwards, it is assumed we have a ZA payload in
+	 the NT_ARM_ZA register set for this thread, and we need to update
+	 such state based on the contents of the register cache.  */
+
+      /* Fetch the current ZA state from the thread.  */
+      gdb::byte_vector za_state = aarch64_fetch_za_regset (tid);
+
+      char *base = (char *) za_state.data ();
+      struct user_za_header *za_header = (struct user_za_header *) base;
+      uint64_t svq = sve_vq_from_vl (za_header->vl);
+
+      /* Sanity check the data in the header.  */
+      if (!sve_vl_valid (za_header->vl)
+	  || ZA_PT_SIZE (svq) != za_header->size)
+	error (_("Invalid vector length or payload size when reading ZA."));
+
+      /* Overwrite the ZA state contained in the thread with the ZA state from
+	 the register cache.  */
+      if (REG_VALID == reg_buf->get_register_status (za_regnum))
+	reg_buf->raw_collect (za_regnum, base + ZA_PT_ZA_OFFSET);
+
+      /* Write back the ZA state to the thread's NT_ARM_ZA register set.  */
+      aarch64_store_za_regset (tid, za_state);
+    }
+
+  /* Update svcr and svg accordingly.  */
+  uint64_t svcr_value = 0;
+  svcr_value |= aarch64_has_ssve_state (tid)? SVCR_SM_BIT : 0;
+  svcr_value |= aarch64_has_za_state (tid)? SVCR_ZA_BIT : 0;
+  reg_buf->raw_supply (svcr_regnum, &svcr_value);
+
+  /* At this point we have written the data contained in the register cache to
+     the thread's NT_ARM_ZA register set.  */
+}
diff --git a/gdb/nat/aarch64-scalable-linux-ptrace.h b/gdb/nat/aarch64-scalable-linux-ptrace.h
index 167782c..d609933 100644
--- a/gdb/nat/aarch64-scalable-linux-ptrace.h
+++ b/gdb/nat/aarch64-scalable-linux-ptrace.h
@@ -1,5 +1,5 @@
-/* Common target dependent definitions for AArch64 Scalable Extensions
-   (SVE/SME).
+/* Common native Linux definitions for the AArch64 scalable
+   extensions: SVE and SME.
 
    Copyright (C) 2018-2023 Free Software Foundation, Inc.
 
@@ -31,19 +31,58 @@
    result when <asm/ptrace.h> is included before <sys/ptrace.h>.  */
 #include <sys/ptrace.h>
 #include <asm/ptrace.h>
-
-#ifndef SVE_SIG_ZREGS_SIZE
+#include <stdarg.h>
+#include "aarch64-scalable-linux-ptrace.h"
 #include "aarch64-scalable-linux-sigcontext.h"
-#endif
 
 /* Indicates whether a SVE ptrace header is followed by SVE registers or a
    fpsimd structure.  */
-
 #define HAS_SVE_STATE(header) ((header).flags & SVE_PT_REGS_SVE)
 
+/* Return true if there is an active SVE state in TID.
+   Return false otherwise.  */
+bool aarch64_has_sve_state (int tid);
+
+/* Return true if there is an active SSVE state in TID.
+   Return false otherwise.  */
+bool aarch64_has_ssve_state (int tid);
+
+/* Return true if there is an active ZA state in TID.
+   Return false otherwise.  */
+bool aarch64_has_za_state (int tid);
+
+/* Given TID, read the SVE header into HEADER.
+
+   Return true if successful, false otherwise.  */
+bool read_sve_header (int tid, struct user_sve_header &header);
+
+/* Given TID, store the SVE HEADER.
+
+   Return true if successful, false otherwise.  */
+bool write_sve_header (int tid, const struct user_sve_header &header);
+
+/* Given TID, read the SSVE header into HEADER.
+
+   Return true if successful, false otherwise.  */
+bool read_ssve_header (int tid, struct user_sve_header &header);
+
+/* Given TID, store the SSVE HEADER.
+
+   Return true if successful, false otherwise.  */
+bool write_ssve_header (int tid, const struct user_sve_header &header);
+
+/* Given TID, read the ZA header into HEADER.
+
+   Return true if successful, false otherwise.  */
+bool read_za_header (int tid, struct user_za_header &header);
+
+/* Given TID, store the ZA HEADER.
+
+   Return true if successful, false otherwise.  */
+bool write_za_header (int tid, const struct user_za_header &header);
+
 /* Read VQ for the given tid using ptrace.  If SVE is not supported then zero
    is returned (on a system that supports SVE, then VQ cannot be zero).  */
-
 uint64_t aarch64_sve_get_vq (int tid);
 
 /* Set VQ in the kernel for the given tid, using either the value VQ or
@@ -52,27 +91,64 @@ uint64_t aarch64_sve_get_vq (int tid);
 bool aarch64_sve_set_vq (int tid, uint64_t vq);
 bool aarch64_sve_set_vq (int tid, struct reg_buffer_common *reg_buf);
 
-/* Read the current SVE register set from thread TID and return its data
-   through a byte vector.  */
+/* Read the streaming mode vq (svq) for the given TID.  If the ZA state is not
+   supported or active, return 0.  */
+uint64_t aarch64_za_get_svq (int tid);
+
+/* Set the vector quotient (vq) in the kernel for the given TID using the
+   value VQ.
 
+   Return true if successful, false otherwise.  */
+bool aarch64_za_set_svq (int tid, uint64_t vq);
+bool aarch64_za_set_svq (int tid, const struct reg_buffer_common *reg_buf,
+			 int svg_regnum);
+
+/* Given TID, return the SVE/SSVE data as a vector of bytes.  */
 extern gdb::byte_vector aarch64_fetch_sve_regset (int tid);
 
-/* Write the SVE contents from SVE_STATE to thread TID.  */
+/* Write the SVE/SSVE contents from SVE_STATE to TID.  */
+extern void aarch64_store_sve_regset (int tid,
+				      const gdb::byte_vector &sve_state);
 
-extern void
-aarch64_store_sve_regset (int tid, const gdb::byte_vector &sve_state);
+/* Given TID, return the ZA data as a vector of bytes.  */
+extern gdb::byte_vector aarch64_fetch_za_regset (int tid);
 
-/* Given a thread id TID and a register buffer REG_BUF, update the register
-   buffer with the SVE state from thread TID.  */
+/* Write ZA_STATE for TID.  */
+extern void aarch64_store_za_regset (int tid, const gdb::byte_vector &za_state);
 
+/* Given TID, initialize the ZA register set so the header contains the right
+   size.  The bytes of the ZA register are initialized to zero.  */
+extern void aarch64_initialize_za_regset (int tid);
+
+/* Given a register buffer REG_BUF, update it with SVE/SSVE register data
+   from SVE_STATE.  */
 extern void
 aarch64_sve_regs_copy_to_reg_buf (int tid, struct reg_buffer_common *reg_buf);
 
-/* Given a thread id TID and a register buffer REG_BUF containing SVE
+/* Given a thread id TID and a register buffer REG_BUF containing SVE/SSVE
    register data, write the SVE data to thread TID.  */
-
 extern void
 aarch64_sve_regs_copy_from_reg_buf (int tid,
 				    struct reg_buffer_common *reg_buf);
 
+/* Given a thread id TID and a register buffer REG_BUF, update the register
+   buffer with the ZA state from thread TID.
+
+   ZA_REGNUM, SVG_REGNUM and SVCR_REGNUM are the register numbers for ZA,
+   SVG and SVCR registers.  */
+extern void aarch64_za_regs_copy_to_reg_buf (int tid,
+					     struct reg_buffer_common *reg_buf,
+					     int za_regnum, int svg_regnum,
+					     int svcr_regnum);
+
+/* Given a thread id TID and a register buffer REG_BUF containing ZA register
+   data, write the ZA data to thread TID.
+
+   ZA_REGNUM, SVG_REGNUM and SVCR_REGNUM are the register numbers for ZA,
+   SVG and SVCR registers.  */
+extern void
+aarch64_za_regs_copy_from_reg_buf (int tid,
+				   struct reg_buffer_common *reg_buf,
+				   int za_regnum, int svg_regnum,
+				   int svcr_regnum);
 #endif /* NAT_AARCH64_SCALABLE_LINUX_PTRACE_H */
diff --git a/gdb/nat/aarch64-scalable-linux-sigcontext.h b/gdb/nat/aarch64-scalable-linux-sigcontext.h
index e0120e0..74407bd 100644
--- a/gdb/nat/aarch64-scalable-linux-sigcontext.h
+++ b/gdb/nat/aarch64-scalable-linux-sigcontext.h
@@ -22,8 +22,11 @@
 #ifndef NAT_AARCH64_SCALABLE_LINUX_SIGCONTEXT_H
 #define NAT_AARCH64_SCALABLE_LINUX_SIGCONTEXT_H
 
+#ifndef SVE_SIG_ZREGS_SIZE
+
 #define SVE_MAGIC	0x53564501
 
+
 struct sve_context {
 	struct _aarch64_ctx head;
 	__u16 vl;
@@ -132,7 +135,7 @@ struct sve_context {
 
 #define SVE_SIG_CONTEXT_SIZE(vq) (SVE_SIG_REGS_OFFSET + SVE_SIG_REGS_SIZE(vq))
 
-/* SVE/FP/SIMD state (NT_ARM_SVE) */
+/* SVE/FP/SIMD state (NT_ARM_SVE and NT_ARM_SSVE) */
 
 struct user_sve_header {
 	__u32 size; /* total meaningful regset content in bytes */
@@ -242,6 +245,7 @@ struct user_sve_header {
 	(SVE_PT_SVE_PREG_OFFSET(vq, SVE_NUM_PREGS) - \
 		SVE_PT_SVE_PREGS_OFFSET(vq))
 
+/* For streaming mode SVE (SSVE) FFR must be read and written as zero.  */
 #define SVE_PT_SVE_FFR_OFFSET(vq) \
 	__SVE_SIG_TO_PT(SVE_SIG_FFR_OFFSET(vq))
 
@@ -267,4 +271,55 @@ struct user_sve_header {
 		  SVE_PT_SVE_OFFSET + SVE_PT_SVE_SIZE(vq, flags)	\
 		: SVE_PT_FPSIMD_OFFSET + SVE_PT_FPSIMD_SIZE(vq, flags))
 
+#endif /* SVE_SIG_ZREGS_SIZE */
+
+/* Scalable Matrix Extensions (SME) definitions.  */
+
+/* Make sure we only define these if the kernel header doesn't.  */
+#ifndef ZA_PT_SIZE
+
+/* ZA state (NT_ARM_ZA) */
+struct user_za_header {
+	__u32 size; /* total meaningful regset content in bytes */
+	__u32 max_size; /* maximum possible size for this thread */
+	__u16 vl; /* current vector length */
+	__u16 max_vl; /* maximum possible vector length */
+	__u16 flags;
+	__u16 __reserved;
+};
+
+/* The remainder of the ZA state follows struct user_za_header.  The
+   total size of the ZA state (including header) depends on the
+   metadata in the header:  ZA_PT_SIZE(vq, flags) gives the total size
+   of the state in bytes, including the header.
+
+   Refer to arch/arm64/include/uapi/asm/sigcontext.h from the Linux kernel
+   for details of how to pass the correct "vq" argument to these macros.  */
+
+/* Offset from the start of struct user_za_header to the register data */
+#define ZA_PT_ZA_OFFSET						  \
+	((sizeof (struct user_za_header) + (__SVE_VQ_BYTES - 1))  \
+		/ __SVE_VQ_BYTES * __SVE_VQ_BYTES)
+
+/* The payload starts at offset ZA_PT_ZA_OFFSET, and is of size
+   ZA_PT_ZA_SIZE(vq, flags).
+
+   The ZA array is stored as a sequence of horizontal vectors ZAV of SVL/8
+   bytes each, starting from vector 0.
+
+   Additional data might be appended in the future.
+
+   The ZA matrix is represented in memory in an endianness-invariant layout
+   which differs from the layout used for the FPSIMD V-registers on big-endian
+   systems: see sigcontext.h for more explanation.  */
+
+#define ZA_PT_ZAV_OFFSET(vq, n)				\
+	(ZA_PT_ZA_OFFSET + ((vq * __SVE_VQ_BYTES) * n))
+
+#define ZA_PT_ZA_SIZE(vq) ((vq * __SVE_VQ_BYTES) * (vq * __SVE_VQ_BYTES))
+
+#define ZA_PT_SIZE(vq)			      \
+	(ZA_PT_ZA_OFFSET + ZA_PT_ZA_SIZE(vq))
+#endif /* ZA_PT_SIZE */
+
 #endif /* NAT_AARCH64_SCALABLE_LINUX_SIGCONTEXT_H */