[lldb][AArch64] Add the tpidr2 TLS register that comes with SME

[DavidSpickett]

Summary

This changes the TLS regset to not only be dynamic in that it could
exist or not (though it always does) but also of a dynamic size.

If SME is present then the regset is 16 bytes and contains both tpidr
and tpidr2.

Testing is the same as tpidr. Write from assembly, read from lldb and
vice versa since we have no way to predict what its value should be
by just running a program.

[omjavaid]

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

269

GetTLS could be GetTLSBuffer similar to GetGPRBuffer/GetFPRBuffer above.

lldb/test/API/linux/aarch64/tls_registers/TestAArch64LinuxTLSRegisters.py

73

These three tests have a lot of commonalities may be merge them into one testing the whole logic. Doesn't look like we are getting much out of emitting three tests from this fairly basic test class.

lldb/test/API/linux/aarch64/tls_registers/main.c

38

It would be interesting to test reading/writing tpidr2 when SME is not enabled.

[DavidSpickett]

lldb/test/API/linux/aarch64/tls_registers/TestAArch64LinuxTLSRegisters.py

73

The tradeoff is execution time vs. a HWCAP check in the program file and a bunch of ifs in Python.

Let me see what I can do, but I'm leaning toward the implementation complexity outweighing the performance gained.

90

And speaking of words, let me change this skip reason to "not be present".

lldb/test/API/linux/aarch64/tls_registers/main.c

38

Not enabled, or not present? (I admit, these two words are used interchangeably in places)

Not enabled is actually the state here, as there's no SMTART used here. Architecture wise, I don't see anything to indicate it makes a difference if SME is active or not.

Not present is covered by test_tpidr2_no_sme.

[DavidSpickett]

Add "Buffer" to method names.

enabled -> present in skipped messages. I realised that "enabled" is ambiguous
does it mean enabled in the CPU or in the process, present is more clearly
meaning is it on the CPU at all.

[DavidSpickett]

Note the behaviour of tpidr2 on a system without SME.

[DavidSpickett]

lldb/test/API/linux/aarch64/tls_registers/TestAArch64LinuxTLSRegisters.py

73

I've looked into this. A major thing to note is that reading tpidr2 (or rather, the system register operands that represent it) on a system without SME will SIGILL. I tried this on a Mountain Jade system.

This means we cannot have the program file set both regardless and only test what we expect to be present. We still need the option to the program and still need to run it again each time we want to check a different register.

(I guess you could have an option for each register, but again we're trading complexity here)

That does mean that the program file is the same between the 3 tests, so you could merge them to only build once. That said I think the time saved rebuilding a small example is not much when put against the pile of if you would need to use to merge the 3 tests into one (and therefore have to unpick later).

I could merge test_tpidr2_no_sme into test_tpidr but again, I'd rather have the clear separation of it being its own test than add an if not self.isAArch64SME... in there.

If I'm addressing completely different aspects than you had considered here, let me know and I'll fix those.

[omjavaid]

This looks good just a minor suggestion inline.

lldb/test/API/linux/aarch64/tls_registers/TestAArch64LinuxTLSRegisters.py

84

Should we try to read/write both tpidr and tpidr2 here?

[DavidSpickett]

Thanks for the review, I'll try making the test a bit more flexible so we can read both registers at once for a bit more coverage.

lldb/test/API/linux/aarch64/tls_registers/TestAArch64LinuxTLSRegisters.py

84

We could but on a system with SME test_tpidr will also run. So you're checking both regardless.

I could change the test program so that it always sets tpidr, but only sets tpdir2 if told to. Then we could check both here just in case there's an issue reading one after the other. I'll give that a go.

[DavidSpickett]

Refactor the test so there are tests for:

• reading tpidr on non-SME systems
• reading tpidr and tpidr2 on SME systems
• tpidr2 not being present on non-SME systems

You could try to fold the last one into the first one but
it gets fiddly for not much saving.

[DavidSpickett]

Correct missing "tls" in test name.

[DavidSpickett]

Rebase and put this after the SSVE registers patch.

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

 //===-- NativeRegisterContextLinux_arm64.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
 //
 //===----------------------------------------------------------------------===//
 
protected:
 
   lldb::addr_t FixWatchpointHitAddress(lldb::addr_t hit_addr) override;
 
 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_tls_tpidr_is_valid;
 
   bool m_sve_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;
 
   bool m_refresh_hwdebug_info;
 
   struct user_pac_mask {
     uint64_t data_mask;
     uint64_t insn_mask;
   };
 
   struct user_pac_mask m_pac_mask;
 
   uint64_t m_mte_ctrl_reg;
   uint64_t m_sme_ctrl_reg;
 
-  uint64_t m_tls_tpidr_reg;
+  struct tls_regs {
+    uint64_t tpidr_reg;
+    // Only valid when SME is present.
+    uint64_t tpidr2_reg;
+  };
+
+  struct tls_regs m_tls_regs;
 
   bool IsGPR(unsigned reg) const;
 
   bool IsFPR(unsigned reg) const;
 
   Status ReadAllSVE();
 
   Status WriteAllSVE();
 
   Status ReadSVEHeader();
 
   Status WriteSVEHeader();
 
   Status ReadPAuthMask();
 
   Status ReadMTEControl();
 
   Status WriteMTEControl();
 
-  Status ReadTLSTPIDR();
+  Status ReadTLS();
 
-  Status WriteTLSTPIDR();
+  Status WriteTLS();
 
   // SVCR is a pseudo register and we do not allow writes to it.
   Status ReadSMEControl();
 
   bool IsSVE(unsigned reg) const;
   bool IsPAuth(unsigned reg) const;
   bool IsMTE(unsigned reg) const;
   bool IsTLS(unsigned reg) const;
   bool IsSME(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 *GetPACMask() { return &m_pac_mask; }
 
   void *GetMTEControl() { return &m_mte_ctrl_reg; }
 
-  void *GetTLSTPIDR() { return &m_tls_tpidr_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();
 
   size_t GetMTEControlSize() { return sizeof(m_mte_ctrl_reg); }
 
-  size_t GetTLSTPIDRSize() { return sizeof(m_tls_tpidr_reg); }
-
   void *GetSMEControl() { return &m_sme_ctrl_reg; }
 
   size_t GetSMEControlSize() { return sizeof(m_sme_ctrl_reg); }
 
+  size_t GetTLSBufferSize() { return m_tls_size; }
+
   llvm::Error ReadHardwareDebugInfo() override;
 
   llvm::Error WriteHardwareDebugRegs(DREGType hwbType) override;
 
   uint32_t CalculateFprOffset(const RegisterInfo *reg_info) const;
 
   RegisterInfoPOSIX_arm64 &GetRegisterInfo() const;
 

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

 //===-- NativeRegisterContextLinux_arm64.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
 //
 //===----------------------------------------------------------------------===//
 
: NativeRegisterContextRegisterInfo(native_thread,
                                         register_info_up.release()),
       NativeRegisterContextLinux(native_thread) {
   ::memset(&m_fpr, 0, sizeof(m_fpr));
   ::memset(&m_gpr_arm64, 0, sizeof(m_gpr_arm64));
   ::memset(&m_hwp_regs, 0, sizeof(m_hwp_regs));
   ::memset(&m_hbp_regs, 0, sizeof(m_hbp_regs));
   ::memset(&m_sve_header, 0, sizeof(m_sve_header));
   ::memset(&m_pac_mask, 0, sizeof(m_pac_mask));
+  ::memset(&m_tls_regs, 0, sizeof(m_tls_regs));
 
   m_mte_ctrl_reg = 0;
-  m_tls_tpidr_reg = 0;
   m_sme_ctrl_reg = 0;
 
   // 16 is just a maximum value, query hardware for actual watchpoint count
   m_max_hwp_supported = 16;
   m_max_hbp_supported = 16;
 
   m_refresh_hwdebug_info = true;
 
   m_gpr_is_valid = false;
   m_fpu_is_valid = false;
   m_sve_buffer_is_valid = false;
   m_sve_header_is_valid = false;
   m_pac_mask_is_valid = false;
   m_mte_ctrl_is_valid = false;
-  m_tls_tpidr_is_valid = false;
+  m_tls_is_valid = false;
+
+  // SME adds the tpidr2 register
+  m_tls_size = GetRegisterInfo().IsSSVEEnabled() ? sizeof(m_tls_regs)
+                                                 : sizeof(m_tls_regs.tpidr_reg);
 
   if (GetRegisterInfo().IsSVEEnabled() || GetRegisterInfo().IsSSVEEnabled())
     m_sve_state = SVEState::Unknown;
   else
     m_sve_state = SVEState::Disabled;
 }
 
 RegisterInfoPOSIX_arm64 &
if (m_sve_state == SVEState::Disabled) {
           sve_reg_num = reg_info->value_regs[0];
         offset = CalculateSVEOffset(GetRegisterInfoAtIndex(sve_reg_num));
       }
 
       assert(offset < GetSVEBufferSize());
       src = (uint8_t *)GetSVEBuffer() + offset;
     }
   } else if (IsTLS(reg)) {
-    error = ReadTLSTPIDR();
+    error = ReadTLS();
     if (error.Fail())
       return error;
 
     offset = reg_info->byte_offset - GetRegisterInfo().GetTLSOffset();
-    assert(offset < GetTLSTPIDRSize());
-    src = (uint8_t *)GetTLSTPIDR() + offset;
+    assert(offset < GetTLSBufferSize());
+    src = (uint8_t *)GetTLSBuffer() + offset;
   } else if (IsSVE(reg)) {
     if (m_sve_state == SVEState::Disabled || m_sve_state == SVEState::Unknown)
       return Status("SVE disabled or not supported");
 
     if (GetRegisterInfo().IsSVERegVG(reg)) {
       sve_vg = GetSVERegVG();
       src = (uint8_t *)&sve_vg;
     } else {
if (IsGPR(reg)) {
 
     offset = reg_info->byte_offset - GetRegisterInfo().GetMTEOffset();
     assert(offset < GetMTEControlSize());
     dst = (uint8_t *)GetMTEControl() + offset;
     ::memcpy(dst, reg_value.GetBytes(), reg_info->byte_size);
 
     return WriteMTEControl();
   } else if (IsTLS(reg)) {
-    error = ReadTLSTPIDR();
+    error = ReadTLS();
     if (error.Fail())
       return error;
 
     offset = reg_info->byte_offset - GetRegisterInfo().GetTLSOffset();
-    assert(offset < GetTLSTPIDRSize());
-    dst = (uint8_t *)GetTLSTPIDR() + offset;
+    assert(offset < GetTLSBufferSize());
+    dst = (uint8_t *)GetTLSBuffer() + offset;
     ::memcpy(dst, reg_value.GetBytes(), reg_info->byte_size);
 
-    return WriteTLSTPIDR();
+    return WriteTLS();
   } else if (IsSME(reg)) {
     return Status("Writing to SVCR is not supported.");
   }
 
   return Status("Failed to write register value");
 }
 
 Status NativeRegisterContextLinux_arm64::ReadAllRegisterValues(
Status NativeRegisterContextLinux_arm64::ReadAllRegisterValues(
 
   if (GetRegisterInfo().IsMTEEnabled()) {
     reg_data_byte_size += GetMTEControlSize();
     error = ReadMTEControl();
     if (error.Fail())
       return error;
   }
 
-  // tpidr is always present but there will be more in future.
-  reg_data_byte_size += GetTLSTPIDRSize();
-  error = ReadTLSTPIDR();
+  // tpidr is always present but tpidr2 depends on SME.
+  reg_data_byte_size += GetTLSBufferSize();
+  error = ReadTLS();
   if (error.Fail())
     return error;
 
   data_sp.reset(new DataBufferHeap(reg_data_byte_size, 0));
   uint8_t *dst = data_sp->GetBytes();
 
   ::memcpy(dst, GetGPRBuffer(), GetGPRBufferSize());
   dst += GetGPRBufferSize();
 
   if (GetRegisterInfo().IsSVEEnabled() || GetRegisterInfo().IsSSVEEnabled()) {
     *dst = static_cast<uint8_t>(m_sve_state);
     dst += sizeof(m_sve_state);
     ::memcpy(dst, GetSVEBuffer(), GetSVEBufferSize());
     dst += GetSVEBufferSize();
   } else {
     ::memcpy(dst, GetFPRBuffer(), GetFPRSize());
     dst += GetFPRSize();
   }
 
   if (GetRegisterInfo().IsMTEEnabled())
     ::memcpy(dst, GetMTEControl(), GetMTEControlSize());
 
-  ::memcpy(dst, GetTLSTPIDR(), GetTLSTPIDRSize());
+  ::memcpy(dst, GetTLSBuffer(), GetTLSBufferSize());
 
   return error;
 }
 
 Status NativeRegisterContextLinux_arm64::WriteAllRegisterValues(
     const lldb::DataBufferSP &data_sp) {
   // AArch64 register data must contain GPRs, either FPR or SVE registers
   // (which can be streaming or non-streaming) and optional MTE register.

 
 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_pac_mask_is_valid = false;
   m_mte_ctrl_is_valid = false;
-  m_tls_tpidr_is_valid = false;
+  m_tls_is_valid = false;
 
   // Update SVE 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::WriteMTEControl() {
   ioVec.iov_base = GetMTEControl();
   ioVec.iov_len = GetMTEControlSize();
 
   m_mte_ctrl_is_valid = false;
 
   return WriteRegisterSet(&ioVec, GetMTEControlSize(), NT_ARM_TAGGED_ADDR_CTRL);
 }
 
-Status NativeRegisterContextLinux_arm64::ReadTLSTPIDR() {
+Status NativeRegisterContextLinux_arm64::ReadTLS() {
   Status error;
 
-  if (m_tls_tpidr_is_valid)
+  if (m_tls_is_valid)
     return error;
 
   struct iovec ioVec;
-  ioVec.iov_base = GetTLSTPIDR();
-  ioVec.iov_len = GetTLSTPIDRSize();
+  ioVec.iov_base = GetTLSBuffer();
+  ioVec.iov_len = GetTLSBufferSize();
 
-  error = ReadRegisterSet(&ioVec, GetTLSTPIDRSize(), NT_ARM_TLS);
+  error = ReadRegisterSet(&ioVec, GetTLSBufferSize(), NT_ARM_TLS);
 
   if (error.Success())
-    m_tls_tpidr_is_valid = true;
+    m_tls_is_valid = true;
 
   return error;
 }
 
-Status NativeRegisterContextLinux_arm64::WriteTLSTPIDR() {
+Status NativeRegisterContextLinux_arm64::WriteTLS() {
   Status error;
 
-  error = ReadTLSTPIDR();
+  error = ReadTLS();
   if (error.Fail())
     return error;
 
   struct iovec ioVec;
-  ioVec.iov_base = GetTLSTPIDR();
-  ioVec.iov_len = GetTLSTPIDRSize();
+  ioVec.iov_base = GetTLSBuffer();
+  ioVec.iov_len = GetTLSBufferSize();
 
-  m_tls_tpidr_is_valid = false;
+  m_tls_is_valid = false;
 
-  return WriteRegisterSet(&ioVec, GetTLSTPIDRSize(), NT_ARM_TLS);
+  return WriteRegisterSet(&ioVec, GetTLSBufferSize(), NT_ARM_TLS);
 }
 
 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.

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

 //===-- RegisterInfoPOSIX_arm64.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
 //
 //===----------------------------------------------------------------------===//
 
public:
   size_t GetRegisterSetCount() const override;
 
   size_t GetRegisterSetFromRegisterIndex(uint32_t reg_index) const override;
 
   void AddRegSetPAuth();
 
   void AddRegSetMTE();
 
-  void AddRegSetTLS();
+  void AddRegSetTLS(bool has_tpidr2);
 
   void AddRegSetSME();
 
   uint32_t ConfigureVectorLength(uint32_t sve_vq);
 
   bool VectorSizeIsValid(uint32_t vq) {
     // coverity[unsigned_compare]
     if (vq >= eVectorQuadwordAArch64 && vq <= eVectorQuadwordAArch64SVEMax)

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

 //===-- RegisterInfoPOSIX_arm64.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
 //
 //===---------------------------------------------------------------------===//
 

 
 static lldb_private::RegisterInfo g_register_infos_pauth[] = {
     DEFINE_EXTENSION_REG(data_mask), DEFINE_EXTENSION_REG(code_mask)};
 
 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)};
+    DEFINE_EXTENSION_REG(tpidr),
+    // Only present when SME is present
+    DEFINE_EXTENSION_REG(tpidr2)};
 
 static lldb_private::RegisterInfo g_register_infos_sme[] = {
     DEFINE_EXTENSION_REG(svcr)};
 
 // 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,
-  k_num_tls_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[] = {
static const lldb_private::RegisterSet g_reg_sets_arm64[k_num_register_sets] = {
      g_sve_regnums_arm64}};
 
 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};
 
-static const lldb_private::RegisterSet g_reg_set_tls_arm64 = {
-    "Thread Local Storage Registers", "tls", k_num_tls_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) {
if (m_opt_regsets.AnySet(eRegsetMaskDynamic)) {
       llvm::copy(reg_sets_ref, std::back_inserter(m_dynamic_reg_sets));
 
       if (m_opt_regsets.AllSet(eRegsetMaskPAuth))
         AddRegSetPAuth();
 
       if (m_opt_regsets.AllSet(eRegsetMaskMTE))
         AddRegSetMTE();
 
-      // tpidr is always present, but in future there will be others so this is
-      // done as a dynamic set.
-      AddRegSetTLS();
+      // The TLS set always contains tpidr but only has tpidr2 when SME is
+      // present.
+      AddRegSetTLS(m_opt_regsets.AllSet(eRegsetMaskSSVE));
 
       if (m_opt_regsets.AllSet(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();
void RegisterInfoPOSIX_arm64::AddRegSetMTE() {
   m_dynamic_reg_infos[mte_regnum].kinds[lldb::eRegisterKindLLDB] = mte_regnum;
 
   m_per_regset_regnum_range[m_register_set_count] =
       std::make_pair(mte_regnum, mte_regnum + 1);
   m_dynamic_reg_sets.push_back(g_reg_set_mte_arm64);
   m_dynamic_reg_sets.back().registers = m_mte_regnum_collection.data();
 }
 
-void RegisterInfoPOSIX_arm64::AddRegSetTLS() {
+void RegisterInfoPOSIX_arm64::AddRegSetTLS(bool has_tpidr2) {
   uint32_t tls_regnum = m_dynamic_reg_infos.size();
-  m_tls_regnum_collection.push_back(tls_regnum);
-  m_dynamic_reg_infos.push_back(g_register_infos_tls[0]);
-  m_dynamic_reg_infos[tls_regnum].byte_offset =
-      m_dynamic_reg_infos[tls_regnum - 1].byte_offset +
-      m_dynamic_reg_infos[tls_regnum - 1].byte_size;
-  m_dynamic_reg_infos[tls_regnum].kinds[lldb::eRegisterKindLLDB] = tls_regnum;
+  uint32_t num_regs = has_tpidr2 ? 2 : 1;
+  for (uint32_t i = 0; i < num_regs; i++) {
+    m_tls_regnum_collection.push_back(tls_regnum + i);
+    m_dynamic_reg_infos.push_back(g_register_infos_tls[i]);
+    m_dynamic_reg_infos[tls_regnum + i].byte_offset =
+        m_dynamic_reg_infos[tls_regnum + i - 1].byte_offset +
+        m_dynamic_reg_infos[tls_regnum + i - 1].byte_size;
+    m_dynamic_reg_infos[tls_regnum + i].kinds[lldb::eRegisterKindLLDB] =
+        tls_regnum + i;
+  }
 
   m_per_regset_regnum_range[m_register_set_count] =
-      std::make_pair(tls_regnum, tls_regnum + 1);
-  m_dynamic_reg_sets.push_back(g_reg_set_tls_arm64);
+      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();
 }
 
 void RegisterInfoPOSIX_arm64::AddRegSetSME() {
   uint32_t sme_regnum = m_dynamic_reg_infos.size();
   m_sme_regnum_collection.push_back(sme_regnum);
   m_dynamic_reg_infos.push_back(g_register_infos_sme[0]);
   m_dynamic_reg_infos[sme_regnum].byte_offset =

lldb/test/API/linux/aarch64/tls_register/main.c

-#include <stdbool.h>
-#include <stdint.h>
-
-int main() {
-  // Save tpidr to restore later.
-  uint64_t tpidr = 0;
-  __asm__ volatile("mrs %0, tpidr_el0" : "=r"(tpidr));
-
-  // Set a pattern for lldb to find.
-  uint64_t pattern = 0x1122334455667788;
-  __asm__ volatile("msr tpidr_el0, %0" ::"r"(pattern));
-
-  // Set break point at this line.
-  // lldb will now set its own pattern for us to find.
-
-  uint64_t new_tpidr = pattern;
-  __asm__ volatile("mrs %0, tpidr_el0" : "=r"(new_tpidr));
-  volatile bool tpidr_was_set = new_tpidr == 0x8877665544332211;
-
-  // Restore the original.
-  __asm__ volatile("msr tpidr_el0, %0" ::"r"(tpidr));
-
-  return 0; // Set break point 2 at this line.
-}

lldb/test/API/linux/aarch64/tls_registers/Makefile

 C_SOURCES := main.c
 
 include Makefile.rules

lldb/test/API/linux/aarch64/tls_registers/TestAArch64LinuxTLSRegisters.py

 """
-Test lldb's ability to read and write the AArch64 TLS register tpidr.
+Test lldb's ability to read and write the AArch64 TLS registers.
 """
 
 import lldb
 from lldbsuite.test.decorators import *
 from lldbsuite.test.lldbtest import *
 from lldbsuite.test import lldbutil
 
 
-class AArch64LinuxTLSRegister(TestBase):
+class AArch64LinuxTLSRegisters(TestBase):
     NO_DEBUG_INFO_TESTCASE = True
 
-    @skipUnlessArch("aarch64")
-    @skipUnlessPlatform(["linux"])
-    def test_tls(self):
+    def setup(self, register_name):
         self.build()
         self.runCmd("file " + self.getBuildArtifact("a.out"), CURRENT_EXECUTABLE_SET)
 
         lldbutil.run_break_set_by_file_and_line(
             self,
             "main.c",
             line_number("main.c", "// Set break point at this line."),
             num_expected_locations=1,
         )
 
         lldbutil.run_break_set_by_file_and_line(
             self,
             "main.c",
             line_number("main.c", "// Set break point 2 at this line."),
             num_expected_locations=1,
         )
 
+        self.runCmd("settings set target.run-args {}".format(
+            {"tpidr": 1, "tpidr2": 2}[register_name]))
         self.runCmd("run", RUN_SUCCEEDED)
 
         if self.process().GetState() == lldb.eStateExited:
             self.fail("Test program failed to run.")
 
         self.expect(
             "thread list",
             STOPPED_DUE_TO_BREAKPOINT,
             substrs=["stopped", "stop reason = breakpoint"],
         )
 
+    def check_tls_reg(self, register_name):
+        self.setup(register_name)
+
         # Since we can't predict what the value will be, the program has set
         # a target value for us to find.
 
         regs = self.thread().GetSelectedFrame().GetRegisters()
         tls_regs = regs.GetFirstValueByName("Thread Local Storage Registers")
         self.assertTrue(tls_regs.IsValid(), "No TLS registers found.")
-        tpidr = tls_regs.GetChildMemberWithName("tpidr")
-        self.assertTrue(tpidr.IsValid(), "No tpidr register found.")
+        tls_reg = tls_regs.GetChildMemberWithName(register_name)
+        self.assertTrue(tls_reg.IsValid(), "{} register not found.".format(
+            register_name))
 
-        self.assertEqual(tpidr.GetValueAsUnsigned(), 0x1122334455667788)
+        self.assertEqual(tls_reg.GetValueAsUnsigned(), 0x1122334455667788)
 
         # Set our own value for the program to find.
-        self.expect("register write tpidr 0x{:x}".format(0x8877665544332211))
+        self.expect("register write {} 0x{:x}".format(
+                    register_name, 0x8877665544332211))
         self.expect("continue")
 
         self.expect(
             "thread list",
             STOPPED_DUE_TO_BREAKPOINT,
             substrs=["stopped", "stop reason = breakpoint"],
         )
 
-        self.expect("p tpidr_was_set", substrs=["true"])
+        self.expect("p reg_was_set", substrs=["true"])
+
+    @skipUnlessArch("aarch64")
+    @skipUnlessPlatform(["linux"])
+    def test_tpidr(self):
+        self.check_tls_reg("tpidr")
+
+    @skipUnlessArch("aarch64")
+    @skipUnlessPlatform(["linux"])
+    def test_tpidr2(self):
+        if not self.isAArch64SME():
+            self.skipTest("SME must present.")
+
+        self.check_tls_reg("tpidr2")
+
+    @skipUnlessArch("aarch64")
+    @skipUnlessPlatform(["linux"])
+    def test_tpidr2_no_sme(self):
+        if self.isAArch64SME():
+            self.skipTest("SME must not be present.")
+
+        self.setup("tpidr")
+
+        regs = self.thread().GetSelectedFrame().GetRegisters()
+        tls_regs = regs.GetFirstValueByName("Thread Local Storage Registers")
+        self.assertTrue(tls_regs.IsValid(), "No TLS registers found.")
+        tls_reg = tls_regs.GetChildMemberWithName("tpidr2")
+        self.assertFalse(tls_reg.IsValid(),
+            "tpdir2 should not be present without SME")

lldb/test/API/linux/aarch64/tls_registers/main.c

+#include <stdbool.h>
+#include <stdint.h>
+
+uint64_t get_tpidr(void) {
+  uint64_t tpidr = 0;
+  __asm__ volatile("mrs %0, tpidr_el0" : "=r"(tpidr));
+  return tpidr;
+}
+
+uint64_t get_tpidr2(void) {
+  uint64_t tpidr2 = 0;
+  // S3_3_C13_C0_5 means tpidr2, and will work with older tools.
+  __asm__ volatile("mrs %0, S3_3_C13_C0_5" : "=r"(tpidr2));
+  return tpidr2;
+}
+
+void set_tpidr(uint64_t value) {
+  __asm__ volatile("msr tpidr_el0, %0" ::"r"(value));
+}
+
+void set_tpidr2(uint64_t value) {
+  __asm__ volatile("msr S3_3_C13_C0_5, %0" ::"r"(value));
+}
+
+int main(int argc, char *argv[]) {
+  uint64_t (*getter)(void) = get_tpidr;
+  void (*setter)(uint64_t) = set_tpidr;
+
+  // Expect just the register number.
+  if (argc != 2)
+    return 1;
+
+  // Note that accessing tpidr2 on a system without it will SIGILL. That is why
+  // we have this option instead of trying to set both regardless.
+  switch (argv[1][0]) {
+  case '1':
+    break;
+  case '2':
+    getter = get_tpidr2;
+    setter = set_tpidr2;
+    break;
+  default:
+    return 1;
+  }
+
+  uint64_t original = getter();
+  uint64_t pattern = 0x1122334455667788;
+  setter(pattern);
+
+  // Set break point at this line.
+  // lldb will now set its own pattern for us to find.
+
+  uint64_t new_value = getter();
+  volatile bool reg_was_set = new_value == 0x8877665544332211;
+
+  setter(original);
+
+  return 0; // Set break point 2 at this line.
+}