fix(aarch64/vcpu): correct handling of KVM_REG_ARM64_SVE_VLS

KVM_REG_ARM64_SVE_VLS needs to be treated differently from
other SVE registers. It needs to be set before vcpu is
finalized.

Signed-off-by: Egor Lazarchuk <yegorlz@amazon.co.uk>

Author: ShadowCurse

src/vmm/src/arch/aarch64/regs.rs

@@ -109,6 +109,10 @@ arm64_sys_reg!(ID_AA64MMFR2_EL1, 3, 0, 0, 7, 2);
 // EL0 Virtual Timer Registers
 arm64_sys_reg!(KVM_REG_ARM_TIMER_CNT, 3, 3, 14, 3, 2);
 
+/// Vector lengths pseudo-register
+pub const KVM_REG_ARM64_SVE_VLS: u64 =
+    KVM_REG_ARM64 | KVM_REG_ARM64_SVE as u64 | KVM_REG_SIZE_U512 | 0xffff;
+
 /// Different aarch64 registers sizes
 #[derive(Debug)]
 pub enum RegSize {

src/vmm/src/vstate/vcpu/aarch64.rs

@@ -12,7 +12,7 @@ use kvm_ioctls::*;
 use serde::{Deserialize, Serialize};
 
 use crate::arch::aarch64::regs::{
-    arm64_core_reg_id, offset__of, Aarch64RegisterVec, KVM_REG_ARM_TIMER_CNT,
+    arm64_core_reg_id, offset__of, Aarch64RegisterVec, KVM_REG_ARM64_SVE_VLS, KVM_REG_ARM_TIMER_CNT,
 };
 use crate::arch::aarch64::vcpu::{
     get_all_registers, get_all_registers_ids, get_mpidr, get_mpstate, get_registers, set_mpstate,
@@ -190,12 +190,28 @@ impl KvmVcpu {
             Some(kvi) => kvi,
             None => Self::default_kvi(vm_fd, self.index)?,
         };
+        self.kvi = state.kvi;
 
         self.init_vcpu(&kvi)?;
+
+        // If KVM_REG_ARM64_SVE_VLS is present it needs to
+        // be set before vcpu is finalized.
+        if let Some(sve_vls_reg) = state
+            .regs
+            .iter()
+            .find(|reg| reg.id == KVM_REG_ARM64_SVE_VLS)
+        {
+            set_register(&self.fd, sve_vls_reg).map_err(KvmVcpuError::RestoreState)?;
+        }
+
         self.finilize_vcpu(&kvi)?;
-        self.kvi = state.kvi;
 
-        for reg in state.regs.iter() {
+        // KVM_REG_ARM64_SVE_VLS needs to be skipped after vcpu is finalized.
+        for reg in state
+            .regs
+            .iter()
+            .filter(|reg| reg.id != KVM_REG_ARM64_SVE_VLS)
+        {
             set_register(&self.fd, reg).map_err(KvmVcpuError::RestoreState)?;
         }
         set_mpstate(&self.fd, state.mp_state).map_err(KvmVcpuError::RestoreState)?;