[lldb][AArch64] Implement resizing of SME's ZA register

The size of ZA depends on the streaming vector length regardless
of the active mode. So in addition to vg (which reports the active
mode) we must send the client svg.

Otherwise the mechanics are the same as for non-streaming SVE.
Use the svg value to update the defined size of ZA, accounting
for the fact that ZA is not a single vector but a suqare matrix.

So if svg is 8, a single streaming vector would be 8*8 = 64 bytes.
ZA is that squared, so 64*64 = 4096 bytes.

Testing is included in a later patch.

Reviewed By: omjavaid

Differential Revision: https://reviews.llvm.org/D159504

Author: DavidSpickett

lldb/include/lldb/Target/DynamicRegisterInfo.h

@@ -93,6 +93,8 @@ class DynamicRegisterInfo {
     return llvm::iterator_range<reg_collection::const_iterator>(m_regs);
   }
 
+  void ConfigureOffsets();
+
 protected:
   // Classes that inherit from DynamicRegisterInfo can see and modify these
   typedef std::vector<lldb_private::RegisterSet> set_collection;
@@ -116,8 +118,6 @@ class DynamicRegisterInfo {
 
   void Finalize(const lldb_private::ArchSpec &arch);
 
-  void ConfigureOffsets();
-
   reg_collection m_regs;
   set_collection m_sets;
   set_reg_num_collection m_set_reg_nums;

lldb/source/Plugins/Process/gdb-remote/GDBRemoteRegisterContext.cpp

@@ -373,14 +373,14 @@ bool GDBRemoteRegisterContext::WriteRegisterBytes(const RegisterInfo *reg_info,
   if (dst == nullptr)
     return false;
 
-  // Code below is specific to AArch64 target in SVE state
+  // Code below is specific to AArch64 target in SVE or SME state
   // If vector granule (vg) register is being written then thread's
   // register context reconfiguration is triggered on success.
-  bool do_reconfigure_arm64_sve = false;
+  // We do not allow writes to SVG so it is not mentioned here.
   const ArchSpec &arch = process->GetTarget().GetArchitecture();
-  if (arch.IsValid() && arch.GetTriple().isAArch64())
-    if (strcmp(reg_info->name, "vg") == 0)
-      do_reconfigure_arm64_sve = true;
+  bool do_reconfigure_arm64_sve = arch.IsValid() &&
+                                  arch.GetTriple().isAArch64() &&
+                                  (strcmp(reg_info->name, "vg") == 0);
 
   if (data.CopyByteOrderedData(data_offset,                // src offset
                                reg_info->byte_size,        // src length
@@ -400,10 +400,10 @@ bool GDBRemoteRegisterContext::WriteRegisterBytes(const RegisterInfo *reg_info,
                 {m_reg_data.GetDataStart(), size_t(m_reg_data.GetByteSize())}))
 
         {
-          SetAllRegisterValid(false);
-
           if (do_reconfigure_arm64_sve)
-            AArch64SVEReconfigure();
+            AArch64Reconfigure();
+
+          InvalidateAllRegisters();
 
           return true;
         }
@@ -435,8 +435,10 @@ bool GDBRemoteRegisterContext::WriteRegisterBytes(const RegisterInfo *reg_info,
           // This is an actual register, write it
           success = SetPrimordialRegister(reg_info, gdb_comm);
 
-          if (success && do_reconfigure_arm64_sve)
-            AArch64SVEReconfigure();
+          if (success && do_reconfigure_arm64_sve) {
+            AArch64Reconfigure();
+            InvalidateAllRegisters();
+          }
         }
 
         // Check if writing this register will invalidate any other register
@@ -760,37 +762,47 @@ uint32_t GDBRemoteRegisterContext::ConvertRegisterKindToRegisterNumber(
   return m_reg_info_sp->ConvertRegisterKindToRegisterNumber(kind, num);
 }
 
-bool GDBRemoteRegisterContext::AArch64SVEReconfigure() {
-  if (!m_reg_info_sp)
-    return false;
-
-  const RegisterInfo *reg_info = m_reg_info_sp->GetRegisterInfo("vg");
-  if (!reg_info)
-    return false;
-
-  uint64_t fail_value = LLDB_INVALID_ADDRESS;
-  uint32_t vg_reg_num = reg_info->kinds[eRegisterKindLLDB];
-  uint64_t vg_reg_value = ReadRegisterAsUnsigned(vg_reg_num, fail_value);
+void GDBRemoteRegisterContext::AArch64Reconfigure() {
+  assert(m_reg_info_sp);
 
-  if (vg_reg_value == fail_value || vg_reg_value > 32)
-    return false;
-
-  reg_info = m_reg_info_sp->GetRegisterInfo("p0");
-  // Predicate registers have 1 bit per byte in the vector so their size is
-  // VL / 8. VG is in units of 8 bytes already, so if the size of p0 == VG
-  // already, we do not have to reconfigure.
-  if (!reg_info || vg_reg_value == reg_info->byte_size)
-    return false;
+  // Once we start to reconfigure registers, we cannot read any of them.
+  // So we must read VG and SVG up front.
 
-  m_reg_info_sp->UpdateARM64SVERegistersInfos(vg_reg_value);
-  // Make a heap based buffer that is big enough to store all registers
-  m_reg_data.SetData(std::make_shared<DataBufferHeap>(
-      m_reg_info_sp->GetRegisterDataByteSize(), 0));
-  m_reg_data.SetByteOrder(GetByteOrder());
+  const uint64_t fail_value = LLDB_INVALID_ADDRESS;
+  std::optional<uint64_t> vg_reg_value;
+  const RegisterInfo *vg_reg_info = m_reg_info_sp->GetRegisterInfo("vg");
+  if (vg_reg_info) {
+    uint32_t vg_reg_num = vg_reg_info->kinds[eRegisterKindLLDB];
+    uint64_t reg_value = ReadRegisterAsUnsigned(vg_reg_num, fail_value);
+    if (reg_value != fail_value && reg_value <= 32)
+      vg_reg_value = reg_value;
+  }
 
-  InvalidateAllRegisters();
+  std::optional<uint64_t> svg_reg_value;
+  const RegisterInfo *svg_reg_info = m_reg_info_sp->GetRegisterInfo("svg");
+  if (svg_reg_info) {
+    uint32_t svg_reg_num = svg_reg_info->kinds[eRegisterKindLLDB];
+    uint64_t reg_value = ReadRegisterAsUnsigned(svg_reg_num, fail_value);
+    if (reg_value != fail_value && reg_value <= 32)
+      svg_reg_value = reg_value;
+  }
 
-  return true;
+  if (vg_reg_value)
+    m_reg_info_sp->UpdateARM64SVERegistersInfos(*vg_reg_value);
+  if (svg_reg_value)
+    m_reg_info_sp->UpdateARM64SMERegistersInfos(*svg_reg_value);
+
+  // At this point if we have updated any registers, their offsets will all be
+  // invalid. If we did, we need to update them all.
+  if (vg_reg_value || svg_reg_value) {
+    m_reg_info_sp->ConfigureOffsets();
+    // From here we are able to read registers again.
+
+    // Make a heap based buffer that is big enough to store all registers
+    m_reg_data.SetData(std::make_shared<DataBufferHeap>(
+        m_reg_info_sp->GetRegisterDataByteSize(), 0));
+    m_reg_data.SetByteOrder(GetByteOrder());
+  }
 }
 
 void GDBRemoteDynamicRegisterInfo::UpdateARM64SVERegistersInfos(uint64_t vg) {
@@ -811,7 +823,14 @@ void GDBRemoteDynamicRegisterInfo::UpdateARM64SVERegistersInfos(uint64_t vg) {
     }
     reg.byte_offset = LLDB_INVALID_INDEX32;
   }
+}
 
-  // Re-calculate register offsets
-  ConfigureOffsets();
+void GDBRemoteDynamicRegisterInfo::UpdateARM64SMERegistersInfos(uint64_t svg) {
+  for (auto &reg : m_regs) {
+    if (strcmp(reg.name, "za") == 0) {
+      // ZA is a register with size (svg*8) * (svg*8). A square essentially.
+      reg.byte_size = (svg * 8) * (svg * 8);
+    }
+    reg.byte_offset = LLDB_INVALID_INDEX32;
+  }
 }

lldb/source/Plugins/Process/gdb-remote/GDBRemoteRegisterContext.h

@@ -39,6 +39,7 @@ class GDBRemoteDynamicRegisterInfo final : public DynamicRegisterInfo {
   ~GDBRemoteDynamicRegisterInfo() override = default;
 
   void UpdateARM64SVERegistersInfos(uint64_t vg);
+  void UpdateARM64SMERegistersInfos(uint64_t svg);
 };
 
 class GDBRemoteRegisterContext : public RegisterContext {
@@ -77,7 +78,8 @@ class GDBRemoteRegisterContext : public RegisterContext {
   uint32_t ConvertRegisterKindToRegisterNumber(lldb::RegisterKind kind,
                                                uint32_t num) override;
 
-  bool AArch64SVEReconfigure();
+  // Reconfigure variable sized registers for AArch64 SVE and SME.
+  void AArch64Reconfigure();
 
 protected:
   friend class ThreadGDBRemote;

lldb/source/Plugins/Process/gdb-remote/ProcessGDBRemote.cpp

@@ -1660,17 +1660,18 @@ ThreadSP ProcessGDBRemote::SetThreadStopInfo(
     gdb_thread->PrivateSetRegisterValue(lldb_regnum, buffer_sp->GetData());
   }
 
-  // AArch64 SVE specific code below calls AArch64SVEReconfigure to update
-  // SVE register sizes and offsets if value of VG register has changed
-  // since last stop.
+  // AArch64 SVE/SME specific code below updates SVE and ZA register sizes and
+  // offsets if value of VG or SVG registers has changed since last stop.
   const ArchSpec &arch = GetTarget().GetArchitecture();
   if (arch.IsValid() && arch.GetTriple().isAArch64()) {
     GDBRemoteRegisterContext *reg_ctx_sp =
         static_cast<GDBRemoteRegisterContext *>(
             gdb_thread->GetRegisterContext().get());
 
-    if (reg_ctx_sp)
-      reg_ctx_sp->AArch64SVEReconfigure();
+    if (reg_ctx_sp) {
+      reg_ctx_sp->AArch64Reconfigure();
+      reg_ctx_sp->InvalidateAllRegisters();
+    }
   }
 
   thread_sp->SetName(thread_name.empty() ? nullptr : thread_name.c_str());

lldb/source/Target/DynamicRegisterInfo.cpp

@@ -614,10 +614,11 @@ void DynamicRegisterInfo::Finalize(const ArchSpec &arch) {
   ConfigureOffsets();
 
   // Check if register info is reconfigurable
-  // AArch64 SVE register set has configurable register sizes
+  // AArch64 SVE register set has configurable register sizes, as does the ZA
+  // register that SME added (the streaming state of SME reuses the SVE state).
   if (arch.GetTriple().isAArch64()) {
     for (const auto &reg : m_regs) {
-      if (strcmp(reg.name, "vg") == 0) {
+      if ((strcmp(reg.name, "vg") == 0) || (strcmp(reg.name, "svg") == 0)) {
         m_is_reconfigurable = true;
         break;
       }