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

[DavidSpickett]

Summary

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.

[DavidSpickett]

Rebase

[omjavaid]

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

42

svg

[DavidSpickett]

Rebase.

[omjavaid]

This apparently requires a rebase now.

[DavidSpickett]

Rebase

[omjavaid]

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

838

can we combine UpdateARM64SVERegistersInfos and UpdateARM64SMERegistersInfos to avoid back to back calls of ConfigureOffsets

[DavidSpickett]

Combine reconfigure into one function. The tricky thing is we have to read
svg and vg up front because we can't read registers while we're messing
with register offsets.

Otherwise, the result is a bit clearer and it also avoids a second heap
allocation for the storage.

[DavidSpickett]

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

838

Yes, done. See comments for weird details.

[omjavaid]

This look good

[DavidSpickett]

Rebase

lldb/include/lldb/Target/DynamicRegisterInfo.h

public:
   const lldb_private::RegisterInfo *
   GetRegisterInfo(llvm::StringRef reg_name) const;
 
   typedef std::vector<lldb_private::RegisterInfo> reg_collection;
   llvm::iterator_range<reg_collection::const_iterator> registers() const {
     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;
   typedef std::vector<uint32_t> reg_num_collection;
   typedef std::vector<reg_num_collection> set_reg_num_collection;
   typedef std::vector<lldb_private::ConstString> name_collection;
   typedef std::map<uint32_t, reg_num_collection> reg_to_regs_map;
   typedef std::map<uint32_t, uint32_t> reg_offset_map;
 
   llvm::Expected<uint32_t> ByteOffsetFromSlice(uint32_t index,
                                                llvm::StringRef slice_str,
                                                lldb::ByteOrder byte_order);
   llvm::Expected<uint32_t> ByteOffsetFromComposite(
       uint32_t index, lldb_private::StructuredData::Array &composite_reg_list,
       lldb::ByteOrder byte_order);
   llvm::Expected<uint32_t> ByteOffsetFromRegInfoDict(
       uint32_t index, lldb_private::StructuredData::Dictionary &reg_info_dict,
       lldb::ByteOrder byte_order);
 
   void MoveFrom(DynamicRegisterInfo &&info);
 
   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;
   name_collection m_set_names;
   reg_to_regs_map m_value_regs_map;
   reg_to_regs_map m_invalidate_regs_map;
   reg_offset_map m_value_reg_offset_map;
   size_t m_reg_data_byte_size = 0u; // The number of bytes required to store

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

 
 class GDBRemoteDynamicRegisterInfo final : public DynamicRegisterInfo {
 public:
   GDBRemoteDynamicRegisterInfo() : DynamicRegisterInfo() {}
 
   ~GDBRemoteDynamicRegisterInfo() override = default;
 
   void UpdateARM64SVERegistersInfos(uint64_t vg);
+  void UpdateARM64SMERegistersInfos(uint64_t svg);
 };
 
 class GDBRemoteRegisterContext : public RegisterContext {
 public:
   GDBRemoteRegisterContext(ThreadGDBRemote &thread, uint32_t concrete_frame_idx,
                            GDBRemoteDynamicRegisterInfoSP reg_info_sp,
                            bool read_all_at_once, bool write_all_at_once);
 
public:
   bool ReadAllRegisterValues(RegisterCheckpoint &reg_checkpoint) override;
 
   bool
   WriteAllRegisterValues(const RegisterCheckpoint &reg_checkpoint) override;
 
   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;
 
   bool ReadRegisterBytes(const RegisterInfo *reg_info);
 
   bool WriteRegisterBytes(const RegisterInfo *reg_info, DataExtractor &data,
                           uint32_t data_offset);

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

bool GDBRemoteRegisterContext::WriteRegisterBytes(const RegisterInfo *reg_info,
 
   // Grab a pointer to where we are going to put this register
   uint8_t *dst = const_cast<uint8_t *>(
       m_reg_data.PeekData(reg_info->byte_offset, reg_info->byte_size));
 
   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
                                dst,                        // dst
                                reg_info->byte_size,        // dst length
                                m_reg_data.GetByteOrder())) // dst byte order
   {
     GDBRemoteClientBase::Lock lock(gdb_comm);
     if (lock) {
       if (m_write_all_at_once) {
         // Invalidate all register values
         InvalidateIfNeeded(true);
 
         // Set all registers in one packet
         if (gdb_comm.WriteAllRegisters(
                 m_thread.GetProtocolID(),
                 {m_reg_data.GetDataStart(), size_t(m_reg_data.GetByteSize())}))
 
         {
-          SetAllRegisterValid(false);
-
           if (do_reconfigure_arm64_sve)
-            AArch64SVEReconfigure();
+            AArch64Reconfigure();
+
+          InvalidateAllRegisters();
 
           return true;
         }
       } else {
         bool success = true;
 
         if (reg_info->value_regs) {
           // This register is part of another register. In this case we read
if (lock) {
               success = false;
             else
               success = SetPrimordialRegister(value_reg_info, gdb_comm);
           }
         } else {
           // 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
         // values? If so, invalidate them
         if (reg_info->invalidate_regs) {
           for (uint32_t idx = 0, reg = reg_info->invalidate_regs[0];
                reg != LLDB_INVALID_REGNUM;
                reg = reg_info->invalidate_regs[++idx])
bool GDBRemoteRegisterContext::WriteAllRegisterValues(
   return false;
 }
 
 uint32_t GDBRemoteRegisterContext::ConvertRegisterKindToRegisterNumber(
     lldb::RegisterKind kind, uint32_t num) {
   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;
+void GDBRemoteRegisterContext::AArch64Reconfigure() {
+  assert(m_reg_info_sp);
 
-  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);
-
-  if (vg_reg_value == fail_value || vg_reg_value > 32)
-    return false;
+  // Once we start to reconfigure registers, we cannot read any of them.
+  // So we must read VG and SVG up front.
 
-  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;
+  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;
+  }
+
+  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;
+  }
+
+  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.
 
-  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());
-
-  InvalidateAllRegisters();
-
-  return true;
+  }
 }
 
 void GDBRemoteDynamicRegisterInfo::UpdateARM64SVERegistersInfos(uint64_t vg) {
   // SVE Z register size is vg x 8 bytes.
   uint32_t z_reg_byte_size = vg * 8;
 
   // SVE vector length has changed, accordingly set size of Z, P and FFR
   // registers. Also invalidate register offsets it will be recalculated
   // after SVE register size update.
   for (auto &reg : m_regs) {
     if (reg.value_regs == nullptr) {
       if (reg.name[0] == 'z' && isdigit(reg.name[1]))
         reg.byte_size = z_reg_byte_size;
       else if (reg.name[0] == 'p' && isdigit(reg.name[1]))
         reg.byte_size = vg;
       else if (strcmp(reg.name, "ffr") == 0)
         reg.byte_size = 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/ProcessGDBRemote.cpp

for (const auto &pair : expedited_register_map) {
     WritableDataBufferSP buffer_sp(
         new DataBufferHeap(reg_value_extractor.GetStringRef().size() / 2, 0));
     reg_value_extractor.GetHexBytes(buffer_sp->GetData(), '\xcc');
     uint32_t lldb_regnum = gdb_reg_ctx_sp->ConvertRegisterKindToRegisterNumber(
         eRegisterKindProcessPlugin, pair.first);
     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());
 
   gdb_thread->SetThreadDispatchQAddr(thread_dispatch_qaddr);
   // Check if the GDB server was able to provide the queue name, kind and serial
   // number
   if (queue_vars_valid)

lldb/source/Target/DynamicRegisterInfo.cpp

void DynamicRegisterInfo::Finalize(const ArchSpec &arch) {
   }
 
   // At this stage call ConfigureOffsets to calculate register offsets for
   // targets supporting dynamic offset calculation. It also calculates
   // total byte size of register data.
   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;
       }
     }
   }
 }
 
 void DynamicRegisterInfo::ConfigureOffsets() {