[LLDB] AArch64 SVE restore SVE registers after expression

[omjavaid]

Summary

This patch fixes register save/restore on expression call to also include SVE registers.

This will fix expression calls like:

re re p1

<Register Value P1 before expression>

p <var-name or function call>

re re p1

<Register Value P1 after expression>

In above example register P1 should remain the same before and after the expression evaluation.

[DavidSpickett]

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

451

Do we need PAC registers here too? (and in write values)

543

I'm a bit dubious about these size checks as a way to tell if things are present, should this one have GetRegisterInfo().IsSVEEnabled() also?

I'm thinking what if PAC plus MTE plus <random future ext> adds enough registers to exceed the size of the SVE header, but the process doesn't in fact have SVE.

You could do the size checks in terms of bytes left to read from the data_sp, that might simplify it some too.

587

This could go wrong if data_sp contains GPR+FPR+SVE data and somehow the register info thinks MTE is enabled. Not sure if that's a realistic thing to handle here.

If you did this in terms of bytes left you could check that there are bytes left to read from data_sp by this point.

(it's a shame the data_sp we get isn't some richer structure we could ask for offsets into)

lldb/test/API/commands/register/register/aarch64_sve_registers/rw_access_static_config/TestSVERegisters.py

175

regiter -> register

[omjavaid]

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

451

We implement PAC as read only so we dont wanna save them either.

543

We are not sure if SVE was enabled or not when we restore. So size checks is the only way we can figure out if a register checkpoint was created when SVE was enabled. This works for now but if we have to stack another register set which is greater than size of SVE header then we will have to put information about enabled register sets inside the checkpoint maybe in a bitmap describing which register set is enabled and look for those at predefined offsets. I thought this wasn't needed for now and will over complicate already complicated variable sized checkpoint.

587

FPR and SVE cannot live together when SVE is present FPR registers are replicated by first 128bits of SVE Z registers so we will either save FPR or SVE data not both.

reg_data_min_size can either be FPR + GPR or SVE + GPR thats why we calculate it again when data has SVE.

lldb/test/API/commands/register/register/aarch64_sve_registers/rw_access_static_config/TestSVERegisters.py

175

Ack.

[omjavaid]

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

543

Just to shed more light on this if SVE is note enabled we just have GPR + FPR + MTE predictable from size. As MTE is just 8 bytes we can be sure if checkpoint has data beyond FPR which is greater than GPR + FPR + SVE_HEADER that means we atleast have SVE header present. Then we check if its a valid header using !sve_vl_valid(m_sve_header.vl)) check and write it to configure SVE registers size before writing actual registers.
If we dont have register data beyond greater than GPR+FPR+MTE then we can safely say there is no SVE present at all.

[rankov]

I have tested this patch with my setup and it fixes the issues described. Thanks.

[DavidSpickett]

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

451

Cool, please add a comment in both places to that effect.

543

Right, of course because PAC is read only we only care about MTE and we know that the SVE header is > 8 bytes. Sounds good to me.

Another comment with that explanation please.

587

Got it.

[DavidSpickett]

Hopefully in the future we can get some kind of smarter container for this register data in future but this looks good for now with a few extra comments for the logic of it. (and so we know where to make changes later)

[omjavaid]

This update addresses review comments.

[DavidSpickett]

LGTM with the clang-format warning fixed.

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

private:
   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 *GetSVEBuffer();
+  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(); }
 
   size_t GetMTEControlSize() { return sizeof(m_mte_ctrl_reg); }

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

 
 #ifndef NT_ARM_TAGGED_ADDR_CTRL
 #define NT_ARM_TAGGED_ADDR_CTRL 0x409 /* Tagged address control register */
 #endif
 
 #define HWCAP_PACA (1 << 30)
 #define HWCAP2_MTE (1 << 18)
 
-#define REG_CONTEXT_SIZE (GetGPRSize() + GetFPRSize())
-
 using namespace lldb;
 using namespace lldb_private;
 using namespace lldb_private::process_linux;
 
 std::unique_ptr<NativeRegisterContextLinux>
 NativeRegisterContextLinux::CreateHostNativeRegisterContextLinux(
     const ArchSpec &target_arch, NativeThreadLinux &native_thread) {
   switch (target_arch.GetMachine()) {
if (IsGPR(reg)) {
     ::memcpy(dst, reg_value.GetBytes(), reg_info->byte_size);
 
     return WriteMTEControl();
   }
 
   return Status("Failed to write register value");
 }
 
 Status NativeRegisterContextLinux_arm64::ReadAllRegisterValues(
     lldb::DataBufferSP &data_sp) {
   Status error;
 
-  data_sp.reset(new DataBufferHeap(REG_CONTEXT_SIZE, 0));
-
+  // Read all register values if not done already and calculate total size of
+  // register set data. We store all register values in data_sp by copying full
+  // PTrace data that corresponds to register sets enabled by current register
+  // context.
+  uint32_t reg_data_byte_size = GetGPRBufferSize();
   error = ReadGPR();
   if (error.Fail())
     return error;
 
+  // If SVE is enabled we need not copy FPR separately.
+  if (GetRegisterInfo().IsSVEEnabled()) {
+    reg_data_byte_size += GetSVEBufferSize();
+    error = ReadAllSVE();
+  } else {
+    reg_data_byte_size += GetFPRSize();
     error = ReadFPR();
+  }
+  if (error.Fail())
+    return error;
+
+  if (GetRegisterInfo().IsMTEEnabled()) {
+    reg_data_byte_size += GetMTEControlSize();
+    error = ReadMTEControl();
     if (error.Fail())
       return error;
+  }
 
+  data_sp.reset(new DataBufferHeap(reg_data_byte_size, 0));
   uint8_t *dst = data_sp->GetBytes();
-  ::memcpy(dst, GetGPRBuffer(), GetGPRSize());
-  dst += GetGPRSize();
+
+  ::memcpy(dst, GetGPRBuffer(), GetGPRBufferSize());
+  dst += GetGPRBufferSize();
+
+  if (GetRegisterInfo().IsSVEEnabled()) {
+    ::memcpy(dst, GetSVEBuffer(), GetSVEBufferSize());
+    dst += GetSVEBufferSize();
+  } else {
     ::memcpy(dst, GetFPRBuffer(), GetFPRSize());
+    dst += GetFPRSize();
+  }
+
+  if (GetRegisterInfo().IsMTEEnabled())
+    ::memcpy(dst, GetMTEControl(), GetMTEControlSize());
 
   return error;
 }
 
 Status NativeRegisterContextLinux_arm64::WriteAllRegisterValues(
     const lldb::DataBufferSP &data_sp) {
   Status error;
 
   if (!data_sp) {
     error.SetErrorStringWithFormat(
         "NativeRegisterContextLinux_arm64::%s invalid data_sp provided",
         __FUNCTION__);
     return error;
   }
 
-  if (data_sp->GetByteSize() != REG_CONTEXT_SIZE) {
-    error.SetErrorStringWithFormat(
-        "NativeRegisterContextLinux_arm64::%s data_sp contained mismatched "
-        "data size, expected %" PRIu64 ", actual %" PRIu64,
-        __FUNCTION__, REG_CONTEXT_SIZE, data_sp->GetByteSize());
-    return error;
-  }
-
   uint8_t *src = data_sp->GetBytes();
   if (src == nullptr) {
     error.SetErrorStringWithFormat("NativeRegisterContextLinux_arm64::%s "
                                    "DataBuffer::GetBytes() returned a null "
                                    "pointer",
                                    __FUNCTION__);
     return error;
   }
-  ::memcpy(GetGPRBuffer(), src, GetRegisterInfoInterface().GetGPRSize());
+
+  uint64_t reg_data_min_size = GetGPRBufferSize() + GetFPRSize();
+  if (data_sp->GetByteSize() < reg_data_min_size) {
+    error.SetErrorStringWithFormat(
+        "NativeRegisterContextLinux_arm64::%s data_sp contained insufficient "
+        "register data bytes, expected at least %" PRIu64 ", actual %" PRIu64,
+        __FUNCTION__, reg_data_min_size, data_sp->GetByteSize());
+    return error;
+  }
+
+  // Register data starts with GPRs
+  ::memcpy(GetGPRBuffer(), src, GetGPRBufferSize());
+  m_gpr_is_valid = true;
 
   error = WriteGPR();
   if (error.Fail())
     return error;
 
-  src += GetRegisterInfoInterface().GetGPRSize();
-  ::memcpy(GetFPRBuffer(), src, GetFPRSize());
+  src += GetGPRBufferSize();
+
+  // Verify if register data may contain SVE register values.
+  bool contains_sve_reg_data =
+      (data_sp->GetByteSize() > (reg_data_min_size + GetSVEHeaderSize()));
+
+  if (contains_sve_reg_data) {
+    // We have SVE register data first write SVE header.
+    ::memcpy(GetSVEHeader(), src, GetSVEHeaderSize());
+    if (!sve_vl_valid(m_sve_header.vl)) {
+      m_sve_header_is_valid = false;
+      error.SetErrorStringWithFormat("NativeRegisterContextLinux_arm64::%s "
+                                     "Invalid SVE header in data_sp",
+                                     __FUNCTION__);
+      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.
+    ConfigureRegisterContext();
 
+    // Make sure data_sp contains sufficient data to write all SVE registers.
+    reg_data_min_size = GetGPRBufferSize() + GetSVEBufferSize();
+    if (data_sp->GetByteSize() < reg_data_min_size) {
+      error.SetErrorStringWithFormat(
+          "NativeRegisterContextLinux_arm64::%s data_sp contained insufficient "
+          "register data bytes, expected %" PRIu64 ", actual %" PRIu64,
+          __FUNCTION__, reg_data_min_size, data_sp->GetByteSize());
+      return error;
+    }
+
+    ::memcpy(GetSVEBuffer(), src, GetSVEBufferSize());
+    m_sve_buffer_is_valid = true;
+    error = WriteAllSVE();
+    src += GetSVEBufferSize();
+  } else {
+    ::memcpy(GetFPRBuffer(), src, GetFPRSize());
+    m_fpu_is_valid = true;
     error = WriteFPR();
+    src += GetFPRSize();
+  }
+
   if (error.Fail())
     return error;
 
+  if (GetRegisterInfo().IsMTEEnabled() && data_sp->GetByteSize() > reg_data_min_size) {
+    ::memcpy(GetMTEControl(), src, GetMTEControlSize());
+    m_mte_ctrl_is_valid = true;
+    error = WriteMTEControl();
+  }
+
   return error;
 }
 
 bool NativeRegisterContextLinux_arm64::IsGPR(unsigned reg) const {
   if (GetRegisterInfo().GetRegisterSetFromRegisterIndex(reg) ==
       RegisterInfoPOSIX_arm64::GPRegSet)
     return true;
   return false;
uint32_t NativeRegisterContextLinux_arm64::CalculateSVEOffset(
   } else if (m_sve_state == SVEState::Full) {
     uint32_t sve_z0_offset = GetGPRSize() + 16;
     sve_reg_offset =
         sve::SigRegsOffset() + reg_info->byte_offset - sve_z0_offset;
   }
   return sve_reg_offset;
 }
 
-void *NativeRegisterContextLinux_arm64::GetSVEBuffer() {
-  if (m_sve_state == SVEState::FPSIMD)
-    return m_sve_ptrace_payload.data() + sve::ptrace_fpsimd_offset;
-
-  return m_sve_ptrace_payload.data();
-}
-
 std::vector<uint32_t> NativeRegisterContextLinux_arm64::GetExpeditedRegisters(
     ExpeditedRegs expType) const {
   std::vector<uint32_t> expedited_reg_nums =
       NativeRegisterContext::GetExpeditedRegisters(expType);
   if (m_sve_state == SVEState::FPSIMD || m_sve_state == SVEState::Full)
     expedited_reg_nums.push_back(GetRegisterInfo().GetRegNumSVEVG());
 
   return expedited_reg_nums;

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

public:
   bool VectorSizeIsValid(uint32_t vq) {
     if (vq >= eVectorQuadwordAArch64 && vq <= eVectorQuadwordAArch64SVEMax)
       return true;
     return false;
   }
 
   bool IsSVEEnabled() const { return m_opt_regsets.AnySet(eRegsetMaskSVE); }
   bool IsPAuthEnabled() const { return m_opt_regsets.AnySet(eRegsetMaskPAuth); }
+  bool IsMTEEnabled() const { return m_opt_regsets.AnySet(eRegsetMaskMTE); }
 
   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;
 

lldb/test/API/commands/register/register/aarch64_sve_registers/rw_access_static_config/TestSVERegisters.py

class RegisterCommandsTestCase(TestBase):
 
     def check_sve_register_size(self, set, name, expected):
         reg_value = set.GetChildMemberWithName(name)
         self.assertTrue(reg_value.IsValid(),
                         'Verify we have a register named "%s"' % (name))
         self.assertEqual(reg_value.GetByteSize(), expected,
                          'Verify "%s" == %i' % (name, expected))
 
+    def check_sve_regs_read(self, z_reg_size):
+        p_reg_size = int(z_reg_size / 8)
+
+        for i in range(32):
+            z_regs_value = '{' + \
+                ' '.join('0x{:02x}'.format(i + 1)
+                         for _ in range(z_reg_size)) + '}'
+            self.expect("register read " + 'z%i' %
+                        (i), substrs=[z_regs_value])
+
+        p_value_bytes = ['0xff', '0x55', '0x11', '0x01', '0x00']
+        for i in range(16):
+            p_regs_value = '{' + \
+                ' '.join(p_value_bytes[i % 5] for _ in range(p_reg_size)) + '}'
+            self.expect("register read " + 'p%i' % (i), substrs=[p_regs_value])
+
+        self.expect("register read ffr", substrs=[p_regs_value])
+
+    def check_sve_regs_read_after_write(self, z_reg_size):
+        p_reg_size = int(z_reg_size / 8)
+
+        z_regs_value = '{' + \
+            ' '.join(('0x9d' for _ in range(z_reg_size))) + '}'
+
+        p_regs_value = '{' + \
+            ' '.join(('0xee' for _ in range(p_reg_size))) + '}'
+
+        for i in range(32):
+            self.runCmd('register write ' + 'z%i' %
+                        (i) + " '" + z_regs_value + "'")
+
+        for i in range(32):
+            self.expect("register read " + 'z%i' % (i), substrs=[z_regs_value])
+
+        for i in range(16):
+            self.runCmd('register write ' + 'p%i' %
+                        (i) + " '" + p_regs_value + "'")
+
+        for i in range(16):
+            self.expect("register read " + 'p%i' % (i), substrs=[p_regs_value])
+
+        self.runCmd('register write ' + 'ffr ' + "'" + p_regs_value + "'")
+
+        self.expect("register read " + 'ffr', substrs=[p_regs_value])
+
     mydir = TestBase.compute_mydir(__file__)
 
     @no_debug_info_test
     @skipIf(archs=no_match(["aarch64"]))
     @skipIf(oslist=no_match(['linux']))
     def test_sve_registers_configuration(self):
         """Test AArch64 SVE registers size configuration."""
         self.build()
def test_sve_registers_read_write(self):
 
         vg_reg = sve_registers.GetChildMemberWithName("vg")
 
         vg_reg_value = sve_registers.GetChildMemberWithName(
             "vg").GetValueAsUnsigned()
 
         z_reg_size = vg_reg_value * 8
 
-        p_reg_size = int(z_reg_size / 8)
-
-        for i in range(32):
-            z_regs_value = '{' + \
-                ' '.join('0x{:02x}'.format(i + 1)
-                         for _ in range(z_reg_size)) + '}'
-            self.expect("register read " + 'z%i' %
-                        (i), substrs=[z_regs_value])
+        self.check_sve_regs_read(z_reg_size)
 
-        p_value_bytes = ['0xff', '0x55', '0x11', '0x01', '0x00']
-        for i in range(16):
-            p_regs_value = '{' + \
-                ' '.join(p_value_bytes[i % 5] for _ in range(p_reg_size)) + '}'
-            self.expect("register read " + 'p%i' % (i), substrs=[p_regs_value])
+        # Evaluate simple expression and print function expr_eval_func address.
+        self.expect("p expr_eval_func", substrs=["= 0x"])
 
-        self.expect("register read ffr", substrs=[p_regs_value])
+        # Evaluate expression call function expr_eval_func.
+        self.expect_expr("expr_eval_func()",
+                         result_type="int", result_value="1")
 
-        z_regs_value = '{' + \
-            ' '.join(('0x9d' for _ in range(z_reg_size))) + '}'
+        # We called a jitted function above which must not have changed SVE
+        # vector length or regiter values.
+        self.check_sve_regs_read(z_reg_size)
 
-        p_regs_value = '{' + \
-            ' '.join(('0xee' for _ in range(p_reg_size))) + '}'
-
-        for i in range(32):
-            self.runCmd('register write ' + 'z%i' %
-                        (i) + " '" + z_regs_value + "'")
-
-        for i in range(32):
-            self.expect("register read " + 'z%i' % (i), substrs=[z_regs_value])
-
-        for i in range(16):
-            self.runCmd('register write ' + 'p%i' %
-                        (i) + " '" + p_regs_value + "'")
-
-        for i in range(16):
-            self.expect("register read " + 'p%i' % (i), substrs=[p_regs_value])
-
-        self.runCmd('register write ' + 'ffr ' + "'" + p_regs_value + "'")
-
-        self.expect("register read " + 'ffr', substrs=[p_regs_value])
+        self.check_sve_regs_read_after_write(z_reg_size)

lldb/test/API/commands/register/register/aarch64_sve_registers/rw_access_static_config/main.c

-int main() {
+#include <sys/prctl.h>
+
+void write_sve_regs () {
   asm volatile("setffr\n\t");
   asm volatile("ptrue p0.b\n\t");
   asm volatile("ptrue p1.h\n\t");
   asm volatile("ptrue p2.s\n\t");
   asm volatile("ptrue p3.d\n\t");
   asm volatile("pfalse p4.b\n\t");
   asm volatile("ptrue p5.b\n\t");
   asm volatile("ptrue p6.h\n\t");
void write_sve_regs () {
   asm volatile("cpy  z24.b, p0/z, #25\n\t");
   asm volatile("cpy  z25.b, p5/z, #26\n\t");
   asm volatile("cpy  z26.b, p10/z, #27\n\t");
   asm volatile("cpy  z27.b, p15/z, #28\n\t");
   asm volatile("cpy  z28.b, p0/z, #29\n\t");
   asm volatile("cpy  z29.b, p5/z, #30\n\t");
   asm volatile("cpy  z30.b, p10/z, #31\n\t");
   asm volatile("cpy  z31.b, p15/z, #32\n\t");
+}
+
+// This function will be called using jitted expression call. We change vector
+// length and write SVE registers. Our program context should restore to
+// orignal vector length and register values after expression evaluation.
+int expr_eval_func() {
+  prctl(PR_SVE_SET_VL, 8 * 2);
+  write_sve_regs();
+  prctl(PR_SVE_SET_VL, 8 * 4);
+  write_sve_regs();
+  return 1;
+}
+
+int main() {
+  write_sve_regs();
   return 0; // Set a break point here.
 }