[lldb][AArch64] Add SME's Array Storage (ZA) register

[DavidSpickett]

Summary

Note: This requires later commits for ZA to function properly,
it is split for ease of review. Testing is also in a later patch.

The "Matrix" part of the Scalable Matrix Extension is a new register
"ZA". You can think of this as a square matrix made up of scalable rows,
where each row is one scalable vector long. However it is not made
of the existing scalable vector registers, it is its own register.
Meaning that the size of ZA is the vector length in bytes * the vector
length in bytes.

https://developer.arm.com/documentation/ddi0616/latest/

It uses the streaming vector length, even when streaming mode itself
is not active. For this reason, it's register data header always
includes the streaming vector length.

Due to it's size I've changed kMaxRegisterByteSize to the maximum
possible ZA size and kTypicalRegisterByteSize will be the maximum
possible scalable vector size. Therefore ZA transactions will cause heap
allocations, and non ZA registers will perform exactly as before.

ZA can be enabled and disabled independently of streaming mode. The way
this works in ptrace is different to SVE versus streaming SVE. Writing
NT_ARM_ZA without register data disables ZA, writing NT_ARM_ZA with
register data enables ZA (LLDB will only support the latter, and only
because it's convenient for us to do so).

https://kernel.org/doc/html/v6.2/arm64/sme.html

LLDB does not handle registers that can appear and dissappear at
runtime. Rather than add complexity to implement that, LLDB will
show a block of 0s when ZA is disabled.

The alternative is not only updating the vector lengths every stop,
but every register definition. It's possible but I'm not sure it's worth
pursuing.

Users should refer to the SVCR register (added in later patches)
for the final word on whether ZA is active or not.

Writing to "VG" during streaming mode will change the size of the
streaming sve registers and ZA. LLDB will not attempt to preserve
register values in this case, we'll just read back the undefined
content the kernel shows. This is in line with, as stated, the
kernel ABIs and the prospective software ABIs look like.

ZA is defined as a vector of size SVL*SVL, so the display in lldb
is very basic. A giant block of values. This is no worse than SVE,
just larger. There is scope to improve this but that can wait
until we see some use cases.

[DavidSpickett]

Rebase

[omjavaid]

This looks good overall thanks for doing the patch split makes the review way less overwhelming.

lldb/source/Plugins/Process/Linux/NativeRegisterContextLinux_arm64.cpp

342

In case of ZA inactive can we avoid having to transfer these zeros over gdb protocol and construct this register on the user side without even doing the transfer?

[DavidSpickett]

lldb/source/Plugins/Process/Linux/NativeRegisterContextLinux_arm64.cpp

342

You could but you'd still need a way to signal to the user side that ZA is in fact inactive. Most obvious way is to send a single 0 and know that if you get that but the streaming vector length is provided too, we must extend the value.

However, that bumps up against a lot of checks in lldb's client side. I know there's at least one that fails if the data received is < the assumed size of the register (for more we just truncate it seems). Bypassing these for 1 target specific register is effort and more importantly, potentially hiding issues. The register size check was actually very useful developing ZA support, so I wouldn't like to subvert it.

GDB also doesn't do anything like this, so we would be compatible with a gdbserver that sent us all the data but an lldb-server wouldn't be compatible with a gdb because it only sends a bit of the data.

You're right that there's scope to compress the data here but it seems better done at a protocol level or by adding some kind of new way to describe registers overall.

[DavidSpickett]

Put ZA in the SME register set instead.

Once other patches are applied you end up with:

Scalable Matrix Extension Registers:
      svcr = 0x0000000000000000
       svg = 0x0000000000000004
        za = {0x00 <...>}

[omjavaid]

lldb/source/Plugins/Process/Linux/NativeRegisterContextLinux_arm64.cpp

566

ZA -> SME?

lldb/source/Plugins/Process/Utility/RegisterInfoPOSIX_arm64.h

110–114

ConfigureVectorLength -> ConfigureVectorLengthSVE?

[omjavaid]

lldb/source/Plugins/Process/Linux/NativeRegisterContextLinux_arm64.cpp

596

can you explain why in case of ZA inactive shouldnt we fill the buffer with zeros here as well ?

[DavidSpickett]

lldb/source/Plugins/Process/Linux/NativeRegisterContextLinux_arm64.cpp

596

Start state: ZA is disabled, reading ptrace gives you just the header.

What we must do to restore that state is to write the header back with no register data. So that's why we don't insert 0s here or trust the currently cached za buffer.

Unlike SVE, there's no flag to say make this active. The presence of the register data is that flag.

...which I will put in a comment as well.

[DavidSpickett]

Address review comments.

[DavidSpickett]

lldb/source/Plugins/Process/Linux/NativeRegisterContextLinux_arm64.cpp

596

Comment is a bit below this.

[omjavaid]

This is LGTM. Thanks for your patience.

[DavidSpickett]

Rebase

lldb/include/lldb/Utility/RegisterValue.h

 
 class RegisterValue {
 public:
   enum {
     // What we can reasonably put on the stack, big enough to support up to 256
     // byte AArch64 SVE.
     kTypicalRegisterByteSize = 256u,
     // Anything else we'll heap allocate storage for it.
-    kMaxRegisterByteSize = kTypicalRegisterByteSize,
+    // 256x256 to support 256 byte AArch64 SME's array storage (ZA) register.
+    // Which is a square of vector length x vector length.
+    kMaxRegisterByteSize = 256u * 256u,
   };
 
   typedef llvm::SmallVector<uint8_t, kTypicalRegisterByteSize> BytesContainer;
 
   enum Type {
     eTypeInvalid,
     eTypeUInt8,
     eTypeUInt16,

lldb/source/Plugins/Process/Linux/NativeRegisterContextLinux_arm64.h

 
 private:
   bool m_gpr_is_valid;
   bool m_fpu_is_valid;
   bool m_sve_buffer_is_valid;
   bool m_mte_ctrl_is_valid;
 
   bool m_sve_header_is_valid;
+  bool m_za_buffer_is_valid;
+  bool m_za_header_is_valid;
   bool m_pac_mask_is_valid;
   bool m_tls_is_valid;
   size_t m_tls_size;
 
   struct user_pt_regs m_gpr_arm64; // 64-bit general purpose registers.
 
   RegisterInfoPOSIX_arm64::FPU
       m_fpr; // floating-point registers including extended register sets.
 
   SVEState m_sve_state;
   struct sve::user_sve_header m_sve_header;
   std::vector<uint8_t> m_sve_ptrace_payload;
 
+  sve::user_za_header m_za_header;
+  std::vector<uint8_t> m_za_ptrace_payload;
+
   bool m_refresh_hwdebug_info;
 
   struct user_pac_mask {
     uint64_t data_mask;
     uint64_t insn_mask;
   };
 
   struct user_pac_mask m_pac_mask;
private:
   Status ReadMTEControl();
 
   Status WriteMTEControl();
 
   Status ReadTLS();
 
   Status WriteTLS();
 
+  Status ReadZAHeader();
+
+  Status ReadZA();
+
+  Status WriteZA();
+
+  // No WriteZAHeader because writing only the header will disable ZA.
+  // Instead use WriteZA and ensure you have the correct ZA buffer size set
+  // beforehand if you wish to disable it.
+
   bool IsSVE(unsigned reg) const;
+  bool IsSME(unsigned reg) const;
   bool IsPAuth(unsigned reg) const;
   bool IsMTE(unsigned reg) const;
   bool IsTLS(unsigned reg) const;
 
   uint64_t GetSVERegVG() { return m_sve_header.vl / 8; }
 
   void SetSVERegVG(uint64_t vg) { m_sve_header.vl = vg * 8; }
 
   void *GetSVEHeader() { return &m_sve_header; }
 
+  void *GetZAHeader() { return &m_za_header; }
+
+  size_t GetZAHeaderSize() { return sizeof(m_za_header); }
+
   void *GetPACMask() { return &m_pac_mask; }
 
   void *GetMTEControl() { return &m_mte_ctrl_reg; }
 
   void *GetTLSBuffer() { return &m_tls_regs; }
 
   void *GetSVEBuffer() { return m_sve_ptrace_payload.data(); }
 
   size_t GetSVEHeaderSize() { return sizeof(m_sve_header); }
 
   size_t GetPACMaskSize() { return sizeof(m_pac_mask); }
 
   size_t GetSVEBufferSize() { return m_sve_ptrace_payload.size(); }
 
   unsigned GetSVERegSet();
 
+  void *GetZABuffer() { return m_za_ptrace_payload.data(); };
+
+  size_t GetZABufferSize() { return m_za_ptrace_payload.size(); }
+
   size_t GetMTEControlSize() { return sizeof(m_mte_ctrl_reg); }
 
   size_t GetTLSBufferSize() { return m_tls_size; }
 
   llvm::Error ReadHardwareDebugInfo() override;
 
   llvm::Error WriteHardwareDebugRegs(DREGType hwbType) override;
 

lldb/source/Plugins/Process/Linux/NativeRegisterContextLinux_arm64.cpp

 #define NT_ARM_SVE 0x405 /* ARM Scalable Vector Extension */
 #endif
 
 #ifndef NT_ARM_SSVE
 #define NT_ARM_SSVE                                                            \
   0x40b /* ARM Scalable Matrix Extension, Streaming SVE mode */
 #endif
 
+#ifndef NT_ARM_ZA
+#define NT_ARM_ZA 0x40c /* ARM Scalable Matrix Extension, Array Storage */
+#endif
+
 #ifndef NT_ARM_PAC_MASK
 #define NT_ARM_PAC_MASK 0x406 /* Pointer authentication code masks */
 #endif
 
 #ifndef NT_ARM_TAGGED_ADDR_CTRL
 #define NT_ARM_TAGGED_ADDR_CTRL 0x409 /* Tagged address control register */
 #endif
 
if (NativeProcessLinux::PtraceWrapper(PTRACE_GETREGSET,
       regset = NT_ARM_SSVE;
       if (NativeProcessLinux::PtraceWrapper(PTRACE_GETREGSET,
                                             native_thread.GetID(), &regset,
                                             &ioVec, sizeof(sve_header))
               .Success())
         opt_regsets.Set(RegisterInfoPOSIX_arm64::eRegsetMaskSSVE);
     }
 
+    sve::user_za_header za_header;
+    ioVec.iov_base = &za_header;
+    ioVec.iov_len = sizeof(za_header);
+    regset = NT_ARM_ZA;
+    if (NativeProcessLinux::PtraceWrapper(PTRACE_GETREGSET,
+                                          native_thread.GetID(), &regset,
+                                          &ioVec, sizeof(za_header))
+            .Success())
+      opt_regsets.Set(RegisterInfoPOSIX_arm64::eRegsetMaskZA);
+
     NativeProcessLinux &process = native_thread.GetProcess();
 
     std::optional<uint64_t> auxv_at_hwcap =
         process.GetAuxValue(AuxVector::AUXV_AT_HWCAP);
     if (auxv_at_hwcap && (*auxv_at_hwcap & HWCAP_PACA))
       opt_regsets.Set(RegisterInfoPOSIX_arm64::eRegsetMaskPAuth);
 
     std::optional<uint64_t> auxv_at_hwcap2 =
NativeRegisterContextLinux_arm64::ReadRegister(const RegisterInfo *reg_info,
   } else if (IsMTE(reg)) {
     error = ReadMTEControl();
     if (error.Fail())
       return error;
 
     offset = reg_info->byte_offset - GetRegisterInfo().GetMTEOffset();
     assert(offset < GetMTEControlSize());
     src = (uint8_t *)GetMTEControl() + offset;
+  } else if (IsSME(reg)) {
+    error = ReadZAHeader();
+    if (error.Fail())
+      return error;
+
+    // If there is only a header and no registers, ZA is inactive. Read as 0
+    // in this case.
+    if (m_za_header.size == sizeof(m_za_header)) {
+      // This will get reconfigured/reset later, so we are safe to use it.
+      // ZA is a square of VL * VL and the ptrace buffer also includes the
+      // header itself.
+      m_za_ptrace_payload.resize(((m_za_header.vl) * (m_za_header.vl)) +
+                                 GetZAHeaderSize());
+      std::fill(m_za_ptrace_payload.begin(), m_za_ptrace_payload.end(), 0);
+    } else {
+      // ZA is active, read the real register.
+      error = ReadZA();
+      if (error.Fail())
+        return error;
+    }
+
+    // ZA is part of the SME set but uses a seperate member buffer for storage.
+    // Therefore its effective byte offset is always 0 even if it isn't 0 within
+    // the SME register set.
+    src = (uint8_t *)GetZABuffer() + GetZAHeaderSize();
   } else
     return Status("failed - register wasn't recognized to be a GPR or an FPR, "
                   "write strategy unknown");
 
   reg_value.SetFromMemoryData(*reg_info, src, reg_info->byte_size,
                               eByteOrderLittle, error);
 
   return error;
else {
 
         if (sve::vl_valid(vg_value * 8)) {
           if (m_sve_header_is_valid && vg_value == GetSVERegVG())
             return error;
 
           SetSVERegVG(vg_value);
 
           error = WriteSVEHeader();
-          if (error.Success())
+          if (error.Success()) {
+            // Changing VG during streaming mode also changes the size of ZA.
+            if (m_sve_state == SVEState::Streaming)
+              m_za_header_is_valid = false;
             ConfigureRegisterContext();
+          }
 
           if (m_sve_header_is_valid && vg_value == GetSVERegVG())
             return error;
         }
 
         return Status("SVE vector length update failed.");
       }
 
if (error.Fail())
       return error;
 
     offset = reg_info->byte_offset - GetRegisterInfo().GetTLSOffset();
     assert(offset < GetTLSBufferSize());
     dst = (uint8_t *)GetTLSBuffer() + offset;
     ::memcpy(dst, reg_value.GetBytes(), reg_info->byte_size);
 
     return WriteTLS();
+  } else if (IsSME(reg)) {
+    error = ReadZA();
+    if (error.Fail())
+      return error;
+
+    // ZA is part of the SME set but not stored with the other SME registers.
+    // So its byte offset is effectively always 0.
+    dst = (uint8_t *)GetZABuffer() + GetZAHeaderSize();
+    ::memcpy(dst, reg_value.GetBytes(), reg_info->byte_size);
+
+    // While this is writing a header that contains a vector length, the only
+    // way to change that is via the vg register. So here we assume the length
+    // will always be the current length and no reconfigure is needed.
+    return WriteZA();
   }
 
   return Status("Failed to write register value");
 }
 
 enum RegisterSetType : uint32_t {
   GPR,
   SVE, // Used for SVE and SSVE.
   FPR, // When there is no SVE, or SVE in FPSIMD mode.
+  // Pointer authentication registers are read only, so not included here.
   MTE,
   TLS,
+  SME, // ZA only, SVCR and SVG are pseudo registers.
 };
 
 static uint8_t *AddRegisterSetType(uint8_t *dst,
                                    RegisterSetType register_set_type) {
   *(reinterpret_cast<uint32_t *>(dst)) = register_set_type;
   return dst + sizeof(uint32_t);
 }
 

 Status
 NativeRegisterContextLinux_arm64::CacheAllRegisters(uint32_t &cached_size) {
   Status error;
   cached_size = sizeof(RegisterSetType) + GetGPRBufferSize();
   error = ReadGPR();
   if (error.Fail())
     return error;
 
+  // Here this means, does the system have ZA, not whether it is active.
+  if (GetRegisterInfo().IsZAEnabled()) {
+    error = ReadZAHeader();
+    if (error.Fail())
+      return error;
+    // Use header size here because the buffer may contain fake data when ZA is
+    // disabled. We do not want to write this fake data (all 0s) because this
+    // would tell the kernel that we want ZA to become active. Which is the
+    // opposite of what we want in the case where it is currently inactive.
+    cached_size += sizeof(RegisterSetType) + m_za_header.size;
+    // For the same reason, we need to force it to be re-read so that it will
+    // always contain the real header.
+    m_za_buffer_is_valid = false;
+    error = ReadZA();
+    if (error.Fail())
+      return error;
+  }
+
   // If SVE is enabled we need not copy FPR separately.
   if (GetRegisterInfo().IsSVEEnabled() || GetRegisterInfo().IsSSVEEnabled()) {
     // Store mode and register data.
     cached_size +=
         sizeof(RegisterSetType) + sizeof(m_sve_state) + GetSVEBufferSize();
     error = ReadAllSVE();
   } else {
     cached_size += sizeof(RegisterSetType) + GetFPRSize();
if (error.Fail())
     return error;
 
   data_sp.reset(new DataBufferHeap(reg_data_byte_size, 0));
   uint8_t *dst = data_sp->GetBytes();
 
   dst = AddSavedRegisters(dst, RegisterSetType::GPR, GetGPRBuffer(),
                           GetGPRBufferSize());
 
+  // Streaming SVE and the ZA register both use the streaming vector length.
+  // When you change this, the kernel will invalidate parts of the process
+  // state. Therefore we need a specific order of restoration for each mode, if
+  // we also have ZA to restore.
+  //
+  // Streaming mode enabled, ZA enabled:
+  // * Write streaming registers. This sets SVCR.SM and clears SVCR.ZA.
+  // * Write ZA, this set SVCR.ZA. The register data we provide is written to
+  // ZA.
+  // * Result is SVCR.SM and SVCR.ZA set, with the expected data in both
+  //   register sets.
+  //
+  // Streaming mode disabled, ZA enabled:
+  // * Write ZA. This sets SVCR.ZA, and the ZA content. In the majority of cases
+  //   the streaming vector length is changing, so the thread is converted into
+  //   an FPSIMD thread if it is not already one. This also clears SVCR.SM.
+  // * Write SVE registers, which also clears SVCR.SM but most importantly, puts
+  //   us into full SVE mode instead of FPSIMD mode (where the registers are
+  //   actually the 128 bit Neon registers).
+  // * Result is we have SVCR.SM = 0, SVCR.ZA = 1 and the expected register
+  //   state.
+  //
+  // Restoring in different orders leads to things like the SVE registers being
+  // truncated due to the FPSIMD mode and ZA being disabled or filled with 0s
+  // (disabled and 0s looks the same from inside lldb since we fake the value
+  // when it's disabled).
+  //
+  // For more information on this, look up the uses of the relevant NT_ARM_
+  // constants and the functions vec_set_vector_length, sve_set_common and
+  // za_set in the Linux Kernel.
+
+  if ((m_sve_state != SVEState::Streaming) && GetRegisterInfo().IsZAEnabled()) {
+    // Use the header size not the buffer size, as we may be using the buffer
+    // for fake data, which we do not want to write out.
+    assert(m_za_header.size <= GetZABufferSize());
+    dst = AddSavedRegisters(dst, RegisterSetType::SME, GetZABuffer(),
+                            m_za_header.size);
+  }
+
   if (GetRegisterInfo().IsSVEEnabled() || GetRegisterInfo().IsSSVEEnabled()) {
     dst = AddRegisterSetType(dst, RegisterSetType::SVE);
     *(reinterpret_cast<SVEState *>(dst)) = m_sve_state;
     dst += sizeof(m_sve_state);
     dst = AddSavedRegistersData(dst, GetSVEBuffer(), GetSVEBufferSize());
   } else {
     dst = AddSavedRegisters(dst, RegisterSetType::FPR, GetFPRBuffer(),
                             GetFPRSize());
   }
 
+  if ((m_sve_state == SVEState::Streaming) && GetRegisterInfo().IsZAEnabled()) {
+    assert(m_za_header.size <= GetZABufferSize());
+    dst = AddSavedRegisters(dst, RegisterSetType::SME, GetZABuffer(),
+                            m_za_header.size);
+  }
+
   if (GetRegisterInfo().IsMTEEnabled()) {
     dst = AddSavedRegisters(dst, RegisterSetType::MTE, GetMTEControl(),
                             GetMTEControlSize());
   }
 
   dst = AddSavedRegisters(dst, RegisterSetType::TLS, GetTLSBuffer(),
                           GetTLSBufferSize());
 
case RegisterSetType::SVE:
         return error;
       }
       m_sve_header_is_valid = true;
       error = WriteSVEHeader();
       if (error.Fail())
         return error;
 
       // SVE header has been written configure SVE vector length if needed.
+      // This could change ZA data too, but that will be restored again later
+      // anyway.
       ConfigureRegisterContext();
 
       // Write header and register data, incrementing src this time.
       error = RestoreRegisters(
           GetSVEBuffer(), &src, GetSVEBufferSize(), m_sve_buffer_is_valid,
           std::bind(&NativeRegisterContextLinux_arm64::WriteAllSVE, this));
       break;
     case RegisterSetType::FPR:
       error = RestoreRegisters(
           GetFPRBuffer(), &src, GetFPRSize(), m_fpu_is_valid,
           std::bind(&NativeRegisterContextLinux_arm64::WriteFPR, this));
       break;
     case RegisterSetType::MTE:
       error = RestoreRegisters(
           GetMTEControl(), &src, GetMTEControlSize(), m_mte_ctrl_is_valid,
           std::bind(&NativeRegisterContextLinux_arm64::WriteMTEControl, this));
       break;
     case RegisterSetType::TLS:
       error = RestoreRegisters(
           GetTLSBuffer(), &src, GetTLSBufferSize(), m_tls_is_valid,
           std::bind(&NativeRegisterContextLinux_arm64::WriteTLS, this));
       break;
+    case RegisterSetType::SME:
+      // To enable or disable ZA you write the regset with or without register
+      // data. The kernel detects this by looking at the ioVec's length, not the
+      // ZA header size you pass in. Therefore we must write header and register
+      // data (if present) in one go every time. Read the header only first just
+      // to get the size.
+      ::memcpy(GetZAHeader(), src, GetZAHeaderSize());
+      // Read the header and register data. Can't use the buffer size here, it
+      // may be incorrect due to being filled with dummy data previously. Resize
+      // this so WriteZA uses the correct size.
+      m_za_ptrace_payload.resize(m_za_header.size);
+      ::memcpy(GetZABuffer(), src, GetZABufferSize());
+      m_za_buffer_is_valid = true;
+
+      error = WriteZA();
+      if (error.Fail())
+        return error;
+
+      // Update size of ZA, which resizes the ptrace payload potentially
+      // trashing our copy of the data we just wrote.
+      ConfigureRegisterContext();
+
+      // ZA buffer now has proper size, read back the data we wrote above, from
+      // ptrace.
+      error = ReadZA();
+      src += GetZABufferSize();
+      break;
     }
 
     if (error.Fail())
       return error;
   }
 
   return error;
 }
if (GetRegisterInfo().GetRegisterSetFromRegisterIndex(reg) ==
     return true;
   return false;
 }
 
 bool NativeRegisterContextLinux_arm64::IsSVE(unsigned reg) const {
   return GetRegisterInfo().IsSVEReg(reg);
 }
 
+bool NativeRegisterContextLinux_arm64::IsSME(unsigned reg) const {
+  return GetRegisterInfo().IsSMEReg(reg);
+}
+
 bool NativeRegisterContextLinux_arm64::IsPAuth(unsigned reg) const {
   return GetRegisterInfo().IsPAuthReg(reg);
 }
 
 bool NativeRegisterContextLinux_arm64::IsMTE(unsigned reg) const {
   return GetRegisterInfo().IsMTEReg(reg);
 }
 
Status NativeRegisterContextLinux_arm64::WriteFPR() {
   return WriteRegisterSet(&ioVec, GetFPRSize(), NT_FPREGSET);
 }
 
 void NativeRegisterContextLinux_arm64::InvalidateAllRegisters() {
   m_gpr_is_valid = false;
   m_fpu_is_valid = false;
   m_sve_buffer_is_valid = false;
   m_sve_header_is_valid = false;
+  m_za_buffer_is_valid = false;
+  m_za_header_is_valid = false;
   m_pac_mask_is_valid = false;
   m_mte_ctrl_is_valid = false;
   m_tls_is_valid = false;
 
-  // Update SVE registers in case there is change in configuration.
+  // Update SVE and ZA registers in case there is change in configuration.
   ConfigureRegisterContext();
 }
 
 unsigned NativeRegisterContextLinux_arm64::GetSVERegSet() {
   return m_sve_state == SVEState::Streaming ? NT_ARM_SSVE : NT_ARM_SVE;
 }
 
 Status NativeRegisterContextLinux_arm64::ReadSVEHeader() {
Status NativeRegisterContextLinux_arm64::WriteTLS() {
   ioVec.iov_base = GetTLSBuffer();
   ioVec.iov_len = GetTLSBufferSize();
 
   m_tls_is_valid = false;
 
   return WriteRegisterSet(&ioVec, GetTLSBufferSize(), NT_ARM_TLS);
 }
 
+Status NativeRegisterContextLinux_arm64::ReadZAHeader() {
+  Status error;
+
+  if (m_za_header_is_valid)
+    return error;
+
+  struct iovec ioVec;
+  ioVec.iov_base = GetZAHeader();
+  ioVec.iov_len = GetZAHeaderSize();
+
+  error = ReadRegisterSet(&ioVec, GetZAHeaderSize(), NT_ARM_ZA);
+
+  if (error.Success())
+    m_za_header_is_valid = true;
+
+  return error;
+}
+
+Status NativeRegisterContextLinux_arm64::ReadZA() {
+  Status error;
+
+  if (m_za_buffer_is_valid)
+    return error;
+
+  struct iovec ioVec;
+  ioVec.iov_base = GetZABuffer();
+  ioVec.iov_len = GetZABufferSize();
+
+  error = ReadRegisterSet(&ioVec, GetZABufferSize(), NT_ARM_ZA);
+
+  if (error.Success())
+    m_za_buffer_is_valid = true;
+
+  return error;
+}
+
+Status NativeRegisterContextLinux_arm64::WriteZA() {
+  // Note that because the ZA ptrace payload contains the header also, this
+  // method will write both. This is done because writing only the header
+  // will disable ZA, even if .size in the header is correct for an enabled ZA.
+  Status error;
+
+  error = ReadZA();
+  if (error.Fail())
+    return error;
+
+  struct iovec ioVec;
+  ioVec.iov_base = GetZABuffer();
+  ioVec.iov_len = GetZABufferSize();
+
+  m_za_buffer_is_valid = false;
+  m_za_header_is_valid = false;
+
+  return WriteRegisterSet(&ioVec, GetZABufferSize(), NT_ARM_ZA);
+}
+
 void NativeRegisterContextLinux_arm64::ConfigureRegisterContext() {
   // ConfigureRegisterContext gets called from InvalidateAllRegisters
   // on every stop and configures SVE vector length and whether we are in
   // streaming SVE mode.
   // If m_sve_state is set to SVEState::Disabled on first stop, code below will
   // be deemed non operational for the lifetime of current process.
   if (!m_sve_header_is_valid && m_sve_state != SVEState::Disabled) {
     // If we have SVE we may also have the SVE streaming mode that SME added.
if (!(error.Success() && ((m_sve_header.flags & sve::ptrace_regs_mask) ==
       } else {
         m_sve_state = SVEState::Disabled;
       }
     }
 
     if (m_sve_state == SVEState::Full || m_sve_state == SVEState::FPSIMD ||
         m_sve_state == SVEState::Streaming) {
       // On every stop we configure SVE vector length by calling
-      // ConfigureVectorLength regardless of current SVEState of this thread.
+      // ConfigureVectorLengthSVE regardless of current SVEState of this thread.
       uint32_t vq = RegisterInfoPOSIX_arm64::eVectorQuadwordAArch64SVE;
       if (sve::vl_valid(m_sve_header.vl))
         vq = sve::vq_from_vl(m_sve_header.vl);
 
-      GetRegisterInfo().ConfigureVectorLength(vq);
+      GetRegisterInfo().ConfigureVectorLengthSVE(vq);
       m_sve_ptrace_payload.resize(sve::PTraceSize(vq, sve::ptrace_regs_sve));
     }
   }
+
+  if (!m_za_header_is_valid) {
+    Status error = ReadZAHeader();
+    if (error.Success()) {
+      uint32_t vq = RegisterInfoPOSIX_arm64::eVectorQuadwordAArch64SVE;
+      if (sve::vl_valid(m_za_header.vl))
+        vq = sve::vq_from_vl(m_za_header.vl);
+
+      GetRegisterInfo().ConfigureVectorLengthZA(vq);
+      m_za_ptrace_payload.resize(m_za_header.size);
+      m_za_buffer_is_valid = false;
+    }
+  }
 }
 
 uint32_t NativeRegisterContextLinux_arm64::CalculateFprOffset(
     const RegisterInfo *reg_info) const {
   return reg_info->byte_offset - GetGPRSize();
 }
 
 uint32_t NativeRegisterContextLinux_arm64::CalculateSVEOffset(
uint32_t NativeRegisterContextLinux_arm64::CalculateSVEOffset(
   }
   return sve_reg_offset;
 }
 
 std::vector<uint32_t> NativeRegisterContextLinux_arm64::GetExpeditedRegisters(
     ExpeditedRegs expType) const {
   std::vector<uint32_t> expedited_reg_nums =
       NativeRegisterContext::GetExpeditedRegisters(expType);
+  // SVE, non-streaming vector length.
   if (m_sve_state == SVEState::FPSIMD || m_sve_state == SVEState::Full)
     expedited_reg_nums.push_back(GetRegisterInfo().GetRegNumSVEVG());
 
   return expedited_reg_nums;
 }
 
 llvm::Expected<NativeRegisterContextLinux::MemoryTaggingDetails>
 NativeRegisterContextLinux_arm64::GetMemoryTaggingDetails(int32_t type) {

lldb/source/Plugins/Process/Utility/LinuxPTraceDefines_arm64sve.h

struct user_sve_header {
   uint32_t size;     /* total meaningful regset content in bytes */
   uint32_t max_size; /* maxmium possible size for this thread */
   uint16_t vl;       /* current vector length */
   uint16_t max_vl;   /* maximum possible vector length */
   uint16_t flags;
   uint16_t reserved;
 };
 
+using user_za_header = user_sve_header;
+
 /* Definitions for user_sve_header.flags: */
 const uint16_t ptrace_regs_mask = 1 << 0;
 const uint16_t ptrace_regs_fpsimd = 0;
 const uint16_t ptrace_regs_sve = ptrace_regs_mask;
 
 /*
  * The remainder of the SVE state follows struct user_sve_header.  The
  * total size of the SVE state (including header) depends on the

lldb/source/Plugins/Process/Utility/RegisterContextPOSIX_arm64.h

protected:
 
   bool IsGPR(unsigned reg);
 
   bool IsFPR(unsigned reg);
 
   size_t GetFPUSize() { return sizeof(RegisterInfoPOSIX_arm64::FPU); }
 
   bool IsSVE(unsigned reg) const;
+  bool IsSME(unsigned reg) const;
   bool IsPAuth(unsigned reg) const;
   bool IsTLS(unsigned reg) const;
 
   bool IsSVEZ(unsigned reg) const { return m_register_info_up->IsSVEZReg(reg); }
   bool IsSVEP(unsigned reg) const { return m_register_info_up->IsSVEPReg(reg); }
   bool IsSVEVG(unsigned reg) const {
     return m_register_info_up->IsSVERegVG(reg);
   }

lldb/source/Plugins/Process/Utility/RegisterContextPOSIX_arm64.cpp

if (m_register_info_up->GetRegisterSetFromRegisterIndex(reg) ==
     return true;
   return false;
 }
 
 bool RegisterContextPOSIX_arm64::IsSVE(unsigned reg) const {
   return m_register_info_up->IsSVEReg(reg);
 }
 
+bool RegisterContextPOSIX_arm64::IsSME(unsigned reg) const {
+  return m_register_info_up->IsSMEReg(reg);
+}
+
 bool RegisterContextPOSIX_arm64::IsPAuth(unsigned reg) const {
   return m_register_info_up->IsPAuthReg(reg);
 }
 
 bool RegisterContextPOSIX_arm64::IsTLS(unsigned reg) const {
   return m_register_info_up->IsTLSReg(reg);
 }
 

lldb/source/Plugins/Process/Utility/RegisterInfoPOSIX_arm64.h

public:
   // AArch64 register set mask value
   enum {
     eRegsetMaskDefault = 0,
     eRegsetMaskSVE = 1,
     eRegsetMaskSSVE = 2,
     eRegsetMaskPAuth = 4,
     eRegsetMaskMTE = 8,
     eRegsetMaskTLS = 16,
+    eRegsetMaskZA = 32,
     eRegsetMaskDynamic = ~1,
   };
 
   // AArch64 Register set FP/SIMD feature configuration
   enum {
     eVectorQuadwordAArch64,
     eVectorQuadwordAArch64SVE,
     eVectorQuadwordAArch64SVEMax = 256
public:
   size_t GetRegisterSetFromRegisterIndex(uint32_t reg_index) const override;
 
   void AddRegSetPAuth();
 
   void AddRegSetMTE();
 
   void AddRegSetTLS(bool has_tpidr2);
 
-  uint32_t ConfigureVectorLength(uint32_t sve_vq);
+  void AddRegSetSME();
+
+  uint32_t ConfigureVectorLengthSVE(uint32_t sve_vq);
+
+  void ConfigureVectorLengthZA(uint32_t za_vq);
 
   bool VectorSizeIsValid(uint32_t vq) {
     // coverity[unsigned_compare]
     if (vq >= eVectorQuadwordAArch64 && vq <= eVectorQuadwordAArch64SVEMax)
       return true;
     return false;
   }
 
   bool IsSVEEnabled() const { return m_opt_regsets.AnySet(eRegsetMaskSVE); }
   bool IsSSVEEnabled() const { return m_opt_regsets.AnySet(eRegsetMaskSSVE); }
+  bool IsZAEnabled() const { return m_opt_regsets.AnySet(eRegsetMaskZA); }
   bool IsPAuthEnabled() const { return m_opt_regsets.AnySet(eRegsetMaskPAuth); }
   bool IsMTEEnabled() const { return m_opt_regsets.AnySet(eRegsetMaskMTE); }
   bool IsTLSEnabled() const { return m_opt_regsets.AnySet(eRegsetMaskTLS); }
 
   bool IsSVEReg(unsigned reg) const;
   bool IsSVEZReg(unsigned reg) const;
   bool IsSVEPReg(unsigned reg) const;
   bool IsSVERegVG(unsigned reg) const;
   bool IsPAuthReg(unsigned reg) const;
   bool IsMTEReg(unsigned reg) const;
   bool IsTLSReg(unsigned reg) const;
+  bool IsSMEReg(unsigned reg) const;
 
   uint32_t GetRegNumSVEZ0() const;
   uint32_t GetRegNumSVEFFR() const;
   uint32_t GetRegNumFPCR() const;
   uint32_t GetRegNumFPSR() const;
   uint32_t GetRegNumSVEVG() const;
   uint32_t GetPAuthOffset() const;
   uint32_t GetMTEOffset() const;
   uint32_t GetTLSOffset() const;
+  uint32_t GetSMEOffset() const;
 
 private:
   typedef std::map<uint32_t, std::vector<lldb_private::RegisterInfo>>
       per_vq_register_infos;
 
   per_vq_register_infos m_per_vq_reg_infos;
 
   uint32_t m_vector_reg_vq = eVectorQuadwordAArch64;
+  uint32_t m_za_reg_vq = eVectorQuadwordAArch64;
 
+  // In normal operation this is const. Only when SVE or SME registers change
+  // size is it either replaced or the content modified.
   const lldb_private::RegisterInfo *m_register_info_p;
   uint32_t m_register_info_count;
 
   const lldb_private::RegisterSet *m_register_set_p;
   uint32_t m_register_set_count;
 
   // Contains pair of [start, end] register numbers of a register set with start
   // and end included.
   std::map<uint32_t, std::pair<uint32_t, uint32_t>> m_per_regset_regnum_range;
 
   lldb_private::Flags m_opt_regsets;
 
   std::vector<lldb_private::RegisterInfo> m_dynamic_reg_infos;
   std::vector<lldb_private::RegisterSet> m_dynamic_reg_sets;
 
   std::vector<uint32_t> pauth_regnum_collection;
   std::vector<uint32_t> m_mte_regnum_collection;
   std::vector<uint32_t> m_tls_regnum_collection;
+  std::vector<uint32_t> m_sme_regnum_collection;
 };
 
 #endif

lldb/source/Plugins/Process/Utility/RegisterInfoPOSIX_arm64.cpp

 static lldb_private::RegisterInfo g_register_infos_mte[] = {
     DEFINE_EXTENSION_REG(mte_ctrl)};
 
 static lldb_private::RegisterInfo g_register_infos_tls[] = {
     DEFINE_EXTENSION_REG(tpidr),
     // Only present when SME is present
     DEFINE_EXTENSION_REG(tpidr2)};
 
+static lldb_private::RegisterInfo g_register_infos_sme[] =
+    // 16 is a default size we will change later.
+    {{"za", nullptr, 16, 0, lldb::eEncodingVector, lldb::eFormatVectorOfUInt8,
+      KIND_ALL_INVALID, nullptr, nullptr, nullptr}};
+
 // Number of register sets provided by this context.
 enum {
   k_num_gpr_registers = gpr_w28 - gpr_x0 + 1,
   k_num_fpr_registers = fpu_fpcr - fpu_v0 + 1,
   k_num_sve_registers = sve_ffr - sve_vg + 1,
   k_num_mte_register = 1,
   // Number of TLS registers is dynamic so it is not listed here.
   k_num_pauth_register = 2,
+  k_num_sme_register = 1,
   k_num_register_sets_default = 2,
   k_num_register_sets = 3
 };
 
 // ARM64 general purpose registers.
 static const uint32_t g_gpr_regnums_arm64[] = {
     gpr_x0,  gpr_x1,   gpr_x2,  gpr_x3,
     gpr_x4,  gpr_x5,   gpr_x6,  gpr_x7,

 static const lldb_private::RegisterSet g_reg_set_pauth_arm64 = {
     "Pointer Authentication Registers", "pauth", k_num_pauth_register, nullptr};
 
 static const lldb_private::RegisterSet g_reg_set_mte_arm64 = {
     "MTE Control Register", "mte", k_num_mte_register, nullptr};
 
 // The size of the TLS set is dynamic, so not listed here.
 
+static const lldb_private::RegisterSet g_reg_set_sme_arm64 = {
+    "Scalable Matrix Extension Registers", "sme", k_num_sme_register, nullptr};
+
 RegisterInfoPOSIX_arm64::RegisterInfoPOSIX_arm64(
     const lldb_private::ArchSpec &target_arch, lldb_private::Flags opt_regsets)
     : lldb_private::RegisterInfoAndSetInterface(target_arch),
       m_opt_regsets(opt_regsets) {
   switch (target_arch.GetMachine()) {
   case llvm::Triple::aarch64:
   case llvm::Triple::aarch64_32: {
     m_register_set_p = g_reg_sets_arm64;
if (m_opt_regsets.AnySet(eRegsetMaskDynamic)) {
 
       if (m_opt_regsets.AllSet(eRegsetMaskMTE))
         AddRegSetMTE();
 
       // The TLS set always contains tpidr but only has tpidr2 when SME is
       // present.
       AddRegSetTLS(m_opt_regsets.AllSet(eRegsetMaskSSVE));
 
+      if (m_opt_regsets.AnySet(eRegsetMaskSSVE))
+        AddRegSetSME();
+
       m_register_info_count = m_dynamic_reg_infos.size();
       m_register_info_p = m_dynamic_reg_infos.data();
       m_register_set_p = m_dynamic_reg_sets.data();
       m_register_set_count = m_dynamic_reg_sets.size();
     }
     break;
   }
   default:
void RegisterInfoPOSIX_arm64::AddRegSetTLS(bool has_tpidr2) {
 
   m_per_regset_regnum_range[m_register_set_count] =
       std::make_pair(tls_regnum, m_dynamic_reg_infos.size());
   m_dynamic_reg_sets.push_back(
       {"Thread Local Storage Registers", "tls", num_regs, nullptr});
   m_dynamic_reg_sets.back().registers = m_tls_regnum_collection.data();
 }
 
-uint32_t RegisterInfoPOSIX_arm64::ConfigureVectorLength(uint32_t sve_vq) {
+void RegisterInfoPOSIX_arm64::AddRegSetSME() {
+  uint32_t sme_regnum = m_dynamic_reg_infos.size();
+  for (uint32_t i = 0; i < k_num_sme_register; i++) {
+    m_sme_regnum_collection.push_back(sme_regnum + i);
+    m_dynamic_reg_infos.push_back(g_register_infos_sme[i]);
+    m_dynamic_reg_infos[sme_regnum + i].byte_offset =
+        m_dynamic_reg_infos[sme_regnum + i - 1].byte_offset +
+        m_dynamic_reg_infos[sme_regnum + i - 1].byte_size;
+    m_dynamic_reg_infos[sme_regnum + i].kinds[lldb::eRegisterKindLLDB] =
+        sme_regnum + i;
+  }
+
+  m_per_regset_regnum_range[m_register_set_count] =
+      std::make_pair(sme_regnum, m_dynamic_reg_infos.size());
+  m_dynamic_reg_sets.push_back(g_reg_set_sme_arm64);
+  m_dynamic_reg_sets.back().registers = m_sme_regnum_collection.data();
+}
+
+uint32_t RegisterInfoPOSIX_arm64::ConfigureVectorLengthSVE(uint32_t sve_vq) {
   // sve_vq contains SVE Quad vector length in context of AArch64 SVE.
   // SVE register infos if enabled cannot be disabled by selecting sve_vq = 0.
   // Also if an invalid or previously set vector length is passed to this
   // function then it will exit immediately with previously set vector length.
   if (!VectorSizeIsValid(sve_vq) || m_vector_reg_vq == sve_vq)
     return m_vector_reg_vq;
 
   // We cannot enable AArch64 only mode if SVE was enabled.
if (reg_info_ref.empty()) {
 
     m_per_vq_reg_infos[sve_vq] = reg_info_ref;
   }
 
   m_register_info_p = m_per_vq_reg_infos[sve_vq].data();
   return m_vector_reg_vq;
 }
 
+void RegisterInfoPOSIX_arm64::ConfigureVectorLengthZA(uint32_t za_vq) {
+  if (!VectorSizeIsValid(za_vq) || m_za_reg_vq == za_vq)
+    return;
+
+  m_za_reg_vq = za_vq;
+
+  // For SVE changes, we replace m_register_info_p completely. ZA is in a
+  // dynamic set and is just 1 register so we make an exception to const here.
+  lldb_private::RegisterInfo *non_const_reginfo =
+      const_cast<lldb_private::RegisterInfo *>(m_register_info_p);
+  non_const_reginfo[m_sme_regnum_collection[0]].byte_size =
+      (za_vq * 16) * (za_vq * 16);
+}
+
 bool RegisterInfoPOSIX_arm64::IsSVEReg(unsigned reg) const {
   if (m_vector_reg_vq > eVectorQuadwordAArch64)
     return (sve_vg <= reg && reg <= sve_ffr);
   else
     return false;
 }
 
 bool RegisterInfoPOSIX_arm64::IsSVEZReg(unsigned reg) const {

 bool RegisterInfoPOSIX_arm64::IsMTEReg(unsigned reg) const {
   return llvm::is_contained(m_mte_regnum_collection, reg);
 }
 
 bool RegisterInfoPOSIX_arm64::IsTLSReg(unsigned reg) const {
   return llvm::is_contained(m_tls_regnum_collection, reg);
 }
 
+bool RegisterInfoPOSIX_arm64::IsSMEReg(unsigned reg) const {
+  return llvm::is_contained(m_sme_regnum_collection, reg);
+}
+
 uint32_t RegisterInfoPOSIX_arm64::GetRegNumSVEZ0() const { return sve_z0; }
 
 uint32_t RegisterInfoPOSIX_arm64::GetRegNumSVEFFR() const { return sve_ffr; }
 
 uint32_t RegisterInfoPOSIX_arm64::GetRegNumFPCR() const { return fpu_fpcr; }
 
 uint32_t RegisterInfoPOSIX_arm64::GetRegNumFPSR() const { return fpu_fpsr; }
 
 uint32_t RegisterInfoPOSIX_arm64::GetRegNumSVEVG() const { return sve_vg; }
 
 uint32_t RegisterInfoPOSIX_arm64::GetPAuthOffset() const {
   return m_register_info_p[pauth_regnum_collection[0]].byte_offset;
 }
 
 uint32_t RegisterInfoPOSIX_arm64::GetMTEOffset() const {
   return m_register_info_p[m_mte_regnum_collection[0]].byte_offset;
 }
 
 uint32_t RegisterInfoPOSIX_arm64::GetTLSOffset() const {
   return m_register_info_p[m_tls_regnum_collection[0]].byte_offset;
 }
+
+uint32_t RegisterInfoPOSIX_arm64::GetSMEOffset() const {
+  return m_register_info_p[m_sme_regnum_collection[0]].byte_offset;
+}

lldb/source/Plugins/Process/elf-core/RegisterContextPOSIXCore_arm64.cpp

if (m_sve_data.GetByteSize() > sizeof(sve::user_sve_header)) {
     if (!sve::vl_valid(m_sve_vector_length)) {
       m_sve_state = SVEState::Disabled;
       m_sve_vector_length = 0;
     }
   } else
     m_sve_state = SVEState::Disabled;
 
   if (m_sve_state != SVEState::Disabled)
-    m_register_info_up->ConfigureVectorLength(
+    m_register_info_up->ConfigureVectorLengthSVE(
         sve::vq_from_vl(m_sve_vector_length));
 }
 
 uint32_t RegisterContextCorePOSIX_arm64::CalculateSVEOffset(
     const RegisterInfo *reg_info) {
   // Start of Z0 data is after GPRs plus 8 bytes of vg register
   uint32_t sve_reg_offset = LLDB_INVALID_INDEX32;
   if (m_sve_state == SVEState::FPSIMD) {

lldb/source/Plugins/Process/elf-core/RegisterUtilities.h

constexpr RegsetDesc FPR_Desc[] = {
     {llvm::Triple::NetBSD, llvm::Triple::x86_64, NETBSD::AMD64::NT_FPREGS},
     {llvm::Triple::OpenBSD, llvm::Triple::UnknownArch, OPENBSD::NT_FPREGS},
 };
 
 constexpr RegsetDesc AARCH64_SVE_Desc[] = {
     {llvm::Triple::Linux, llvm::Triple::aarch64, llvm::ELF::NT_ARM_SVE},
 };
 
+constexpr RegsetDesc AARCH64_ZA_Desc[] = {
+    {llvm::Triple::Linux, llvm::Triple::aarch64, llvm::ELF::NT_ARM_ZA},
+};
+
 constexpr RegsetDesc AARCH64_PAC_Desc[] = {
     {llvm::Triple::Linux, llvm::Triple::aarch64, llvm::ELF::NT_ARM_PAC_MASK},
 };
 
 constexpr RegsetDesc AARCH64_TLS_Desc[] = {
     {llvm::Triple::Linux, llvm::Triple::aarch64, llvm::ELF::NT_ARM_TLS},
 };