[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

839

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

839

Yes, done. See comments for weird details.

[omjavaid]

This look good

[DavidSpickett]

Rebase

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

 //===-- GDBRemoteRegisterContext.h ------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 

 
 class GDBRemoteDynamicRegisterInfo final : public DynamicRegisterInfo {
 public:
   GDBRemoteDynamicRegisterInfo() : DynamicRegisterInfo() {}
 
   ~GDBRemoteDynamicRegisterInfo() override = default;
 
   void UpdateARM64SVERegistersInfos(uint64_t vg);
+  void UpdateARM64SMERegistersInfos(uint64_t vg);
 };
 
 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();
+  void AArch64SVEReconfigure();
+
+  void AArch64SMEReconfigure();
 
 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

 //===-- GDBRemoteRegisterContext.cpp --------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
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 SMEstate
   // 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)
+          if (do_reconfigure_arm64_sve) {
             AArch64SVEReconfigure();
+            AArch64SMEReconfigure();
+          }
+
+          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)
+          if (success && do_reconfigure_arm64_sve) {
             AArch64SVEReconfigure();
+            AArch64SMEReconfigure();
+            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;
-
+void GDBRemoteRegisterContext::AArch64SVEReconfigure() {
+  assert(m_reg_info_sp);
   const RegisterInfo *reg_info = m_reg_info_sp->GetRegisterInfo("vg");
   if (!reg_info)
-    return false;
+    return;
 
   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;
+    return;
 
   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;
+    return;
 
   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();
+void GDBRemoteRegisterContext::AArch64SMEReconfigure() {
+  assert(m_reg_info_sp);
+  const RegisterInfo *reg_info = m_reg_info_sp->GetRegisterInfo("svg");
+  // Target does not have SME, nothing for us to reconfigure.
+  if (!reg_info)
+    return;
 
-  return true;
+  uint64_t fail_value = LLDB_INVALID_ADDRESS;
+  uint32_t svg_reg_num = reg_info->kinds[eRegisterKindLLDB];
+  uint64_t svg_reg_value = ReadRegisterAsUnsigned(svg_reg_num, fail_value);
+
+  if (svg_reg_value == LLDB_INVALID_ADDRESS || svg_reg_value > 32)
+    return;
+
+  m_reg_info_sp->UpdateARM64SMERegistersInfos(svg_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());
 }
 
 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;
+  }
+
+  ConfigureOffsets();
+}

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

 //===-- ProcessGDBRemote.cpp ----------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
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)
+    if (reg_ctx_sp) {
       reg_ctx_sp->AArch64SVEReconfigure();
+      reg_ctx_sp->AArch64SMEReconfigure();
+      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

 //===-- DynamicRegisterInfo.cpp -------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
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() {