From 5c25d82f598433353fc96122afe15cf9bb0aa650 Mon Sep 17 00:00:00 2001 From: Ruben Hakobyan Date: Thu, 21 May 2026 19:04:20 -0700 Subject: [PATCH] hypervisor, vmm: Fix aarch64 SVE register save/restore When SVE is enabled, KVM replaces the FPSIMD V-registers with wider SVE Z-registers. Attempting to access the old FPSIMD offsets returns EINVAL. Fix by classifying each register from KVM_GET_REG_LIST as core, system, or extended. Extended registers (currently SVE only) are saved as generic `ExtendedReg` entries split into `pre_finalize_regs` (registers like SVE VLS that must be written before `vcpu_finalize`) and `extended_regs`. FPSIMD registers are only accessed when SVE is absent. Unrecognized register families error immediately so future extensions like SME fail clearly rather than silently losing state. The snapshot is deserialized before vCPU init to make pre-finalize register state available for the init -> VLS -> finalize ordering required by KVM. Signed-off-by: Ruben Hakobyan --- hypervisor/src/arch/aarch64/mod.rs | 8 +++ hypervisor/src/cpu.rs | 20 ++++++ hypervisor/src/kvm/aarch64/mod.rs | 38 +++++++--- hypervisor/src/kvm/mod.rs | 109 +++++++++++++++++++++-------- hypervisor/src/lib.rs | 12 ++++ vmm/src/cpu.rs | 23 ++++-- 6 files changed, 165 insertions(+), 45 deletions(-) diff --git a/hypervisor/src/arch/aarch64/mod.rs b/hypervisor/src/arch/aarch64/mod.rs index 9d19f44bf..53ed17c15 100644 --- a/hypervisor/src/arch/aarch64/mod.rs +++ b/hypervisor/src/arch/aarch64/mod.rs @@ -4,3 +4,11 @@ pub mod gic; pub mod regs; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Serialize, Deserialize)] +pub struct ExtendedReg { + pub id: u64, + pub data: Vec, +} diff --git a/hypervisor/src/cpu.rs b/hypervisor/src/cpu.rs index 044c81a2e..aaf741129 100644 --- a/hypervisor/src/cpu.rs +++ b/hypervisor/src/cpu.rs @@ -18,6 +18,8 @@ use {anyhow::anyhow, vm_memory::GuestAddress}; use crate::RegList; #[cfg(target_arch = "aarch64")] use crate::VcpuInit; +#[cfg(target_arch = "aarch64")] +use crate::arch::aarch64::ExtendedReg; #[cfg(target_arch = "x86_64")] use crate::arch::x86::{CpuIdEntry, FpuState, LapicState, MsrEntry, SpecialRegisters}; #[cfg(feature = "tdx")] @@ -221,6 +223,16 @@ pub enum HypervisorCpuError { #[error("Failed to set aarch64 core register")] SetAarchCoreRegister(#[source] anyhow::Error), /// + /// Getting extended register error + /// + #[error("Failed to get extended register")] + GetExtendedRegister(#[source] anyhow::Error), + /// + /// Setting extended register error + /// + #[error("Failed to set extended register")] + SetExtendedRegister(#[source] anyhow::Error), + /// /// Getting RISC-V 64-bit core register error /// #[error("Failed to get riscv64 core register")] @@ -460,6 +472,14 @@ pub trait Vcpu: Send + Sync { #[cfg(target_arch = "aarch64")] fn vcpu_finalize(&self, feature: i32) -> Result<()>; /// + /// Sets pre-finalize registers (e.g. SVE VLS). + /// Must be called after vcpu_init and before vcpu_finalize. + /// + #[cfg(target_arch = "aarch64")] + fn set_pre_finalize_regs(&self, _regs: &[ExtendedReg]) -> Result<()> { + Ok(()) + } + /// /// Gets the features that have been finalized for a given CPU. /// #[cfg(target_arch = "aarch64")] diff --git a/hypervisor/src/kvm/aarch64/mod.rs b/hypervisor/src/kvm/aarch64/mod.rs index a94ed55f1..156d79684 100644 --- a/hypervisor/src/kvm/aarch64/mod.rs +++ b/hypervisor/src/kvm/aarch64/mod.rs @@ -11,8 +11,9 @@ pub mod gic; use kvm_bindings::{ - KVM_REG_ARM_COPROC_MASK, KVM_REG_ARM_CORE, KVM_REG_SIZE_MASK, KVM_REG_SIZE_U32, - KVM_REG_SIZE_U64, kvm_mp_state, kvm_one_reg, kvm_regs, + KVM_REG_ARM_COPROC_MASK, KVM_REG_ARM_CORE, KVM_REG_ARM64, KVM_REG_ARM64_SVE, KVM_REG_SIZE_MASK, + KVM_REG_SIZE_SHIFT, KVM_REG_SIZE_U32, KVM_REG_SIZE_U64, KVM_REG_SIZE_U128, KVM_REG_SIZE_U256, + KVM_REG_SIZE_U512, KVM_REG_SIZE_U1024, KVM_REG_SIZE_U2048, kvm_mp_state, kvm_one_reg, kvm_regs, }; pub use kvm_ioctls::{Cap, Kvm}; use serde::{Deserialize, Serialize}; @@ -70,15 +71,26 @@ pub fn is_system_register(regid: u64) -> bool { } let size = regid & KVM_REG_SIZE_MASK; - - assert!( - !(size != KVM_REG_SIZE_U32 && size != KVM_REG_SIZE_U64), - "Unexpected register size for system register {size}" - ); - - true + match size { + KVM_REG_SIZE_U32 | KVM_REG_SIZE_U64 => true, + KVM_REG_SIZE_U128 | KVM_REG_SIZE_U256 | KVM_REG_SIZE_U512 | KVM_REG_SIZE_U1024 + | KVM_REG_SIZE_U2048 => false, + _ => unreachable!("Unexpected register size {size:#x} for register id {regid:#x}"), + } } +pub fn is_sve_register(regid: u64) -> bool { + (regid & KVM_REG_ARM_COPROC_MASK as u64) == KVM_REG_ARM64_SVE as u64 +} + +pub fn reg_size(regid: u64) -> usize { + let shift = ((regid & KVM_REG_SIZE_MASK) >> KVM_REG_SIZE_SHIFT as u64) as u32; + 1usize << shift +} + +pub const KVM_ARM64_SVE_VLS_REGID: u64 = + KVM_REG_ARM64 | KVM_REG_SIZE_U512 | KVM_REG_ARM64_SVE as u64 | 0xffff; + pub fn check_required_kvm_extensions(kvm: &Kvm) -> KvmResult<()> { macro_rules! check_extension { ($cap:expr) => { @@ -100,9 +112,17 @@ pub fn check_required_kvm_extensions(kvm: &Kvm) -> KvmResult<()> { Ok(()) } +pub use crate::arch::aarch64::ExtendedReg; + +pub const PRE_FINALIZE_IDS: &[u64] = &[KVM_ARM64_SVE_VLS_REGID]; + #[derive(Clone, Default, Serialize, Deserialize)] pub struct VcpuKvmState { pub mp_state: kvm_mp_state, pub core_regs: kvm_regs, pub sys_regs: Vec, + #[serde(default)] + pub pre_finalize_regs: Vec, + #[serde(default)] + pub extended_regs: Vec, } diff --git a/hypervisor/src/kvm/mod.rs b/hypervisor/src/kvm/mod.rs index a7fa74ca6..4ed7c0ec7 100644 --- a/hypervisor/src/kvm/mod.rs +++ b/hypervisor/src/kvm/mod.rs @@ -42,7 +42,10 @@ use vmm_sys_util::eventfd::EventFd; #[cfg(target_arch = "aarch64")] use crate::aarch64::gic::KvmGicV3Its; #[cfg(target_arch = "aarch64")] -pub use crate::aarch64::{VcpuKvmState, check_required_kvm_extensions, is_system_register}; +pub use crate::aarch64::{ + ExtendedReg, KVM_ARM64_SVE_VLS_REGID, PRE_FINALIZE_IDS, VcpuKvmState, + check_required_kvm_extensions, is_sve_register, is_system_register, reg_size, +}; #[cfg(target_arch = "aarch64")] use crate::arch::aarch64::gic::{Vgic, VgicConfig}; #[cfg(target_arch = "riscv64")] @@ -113,10 +116,10 @@ pub use kvm_bindings::{ }; #[cfg(target_arch = "aarch64")] use kvm_bindings::{ - KVM_GUESTDBG_USE_HW, KVM_NR_SPSR, KVM_REG_ARM_CORE, KVM_REG_ARM64, KVM_REG_ARM64_SYSREG, - KVM_REG_ARM64_SYSREG_CRM_MASK, KVM_REG_ARM64_SYSREG_CRN_MASK, KVM_REG_ARM64_SYSREG_OP0_MASK, - KVM_REG_ARM64_SYSREG_OP1_MASK, KVM_REG_ARM64_SYSREG_OP2_MASK, KVM_REG_SIZE_U32, - KVM_REG_SIZE_U64, KVM_REG_SIZE_U128, kvm_regs, user_pt_regs, + KVM_GUESTDBG_USE_HW, KVM_NR_SPSR, KVM_REG_ARM_COPROC_MASK, KVM_REG_ARM_CORE, KVM_REG_ARM64, + KVM_REG_ARM64_SYSREG, KVM_REG_ARM64_SYSREG_CRM_MASK, KVM_REG_ARM64_SYSREG_CRN_MASK, + KVM_REG_ARM64_SYSREG_OP0_MASK, KVM_REG_ARM64_SYSREG_OP1_MASK, KVM_REG_ARM64_SYSREG_OP2_MASK, + KVM_REG_SIZE_U32, KVM_REG_SIZE_U64, KVM_REG_SIZE_U128, kvm_regs, user_pt_regs, }; #[cfg(target_arch = "riscv64")] use kvm_bindings::{KVM_REG_RISCV_CORE, kvm_riscv_core}; @@ -1895,7 +1898,6 @@ impl cpu::Vcpu for KvmVcpu { off += std::mem::size_of::(); } - self.get_fpsimd_regs(&mut state)?; Ok(state.into()) } @@ -2049,7 +2051,6 @@ impl cpu::Vcpu for KvmVcpu { off += std::mem::size_of::(); } - self.set_fpsimd_regs(&kvm_regs_state)?; Ok(()) } @@ -2636,6 +2637,16 @@ impl cpu::Vcpu for KvmVcpu { .map_err(|e| cpu::HypervisorCpuError::VcpuFinalize(e.into())) } + #[cfg(target_arch = "aarch64")] + fn set_pre_finalize_regs(&self, regs: &[ExtendedReg]) -> cpu::Result<()> { + for reg in regs { + self.fd + .set_one_reg(reg.id, ®.data) + .map_err(|e| cpu::HypervisorCpuError::SetExtendedRegister(e.into()))?; + } + Ok(()) + } + #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))] /// /// Gets a list of the guest registers that are supported for the @@ -2912,42 +2923,64 @@ impl cpu::Vcpu for KvmVcpu { mp_state: self.get_mp_state()?.into(), ..Default::default() }; + // Get core registers state.core_regs = self.get_regs()?.into(); - // Get systerm register // Call KVM_GET_REG_LIST to get all registers available to the guest. - // For ArmV8 there are around 500 registers. - let mut sys_regs: Vec = Vec::new(); + // ARM64_REGS_MAX in kvm-bindings caps this at 500. let mut reg_list = kvm_bindings::RegList::new(500).unwrap(); self.fd .get_reg_list(&mut reg_list) .map_err(|e| cpu::HypervisorCpuError::GetRegList(e.into()))?; - // At this point reg_list should contain: core registers and system - // registers. - // The register list contains the number of registers and their ids. We - // will be needing to call KVM_GET_ONE_REG on each id in order to save - // all of them. We carve out from the list the core registers which are - // represented in the kernel by kvm_regs structure and for which we can - // calculate the id based on the offset in the structure. - reg_list.retain(|regid| is_system_register(*regid)); + let mut sys_regs: Vec = Vec::new(); + let mut pre_finalize_regs: Vec = Vec::new(); + let mut extended_regs: Vec = Vec::new(); + let mut has_sve = false; - // Now, for the rest of the registers left in the previously fetched - // register list, we are simply calling KVM_GET_ONE_REG. - let indices = reg_list.as_slice(); - for index in indices.iter() { - let mut bytes = [0_u8; 8]; - self.fd - .get_one_reg(*index, &mut bytes) - .map_err(|e| cpu::HypervisorCpuError::GetSysRegister(e.into()))?; - sys_regs.push(kvm_bindings::kvm_one_reg { - id: *index, - addr: u64::from_le_bytes(bytes), - }); + for regid in reg_list.as_slice().iter().copied() { + if (regid & KVM_REG_ARM_COPROC_MASK as u64) == KVM_REG_ARM_CORE as u64 { + // Handled by get_regs() above. + } else if is_system_register(regid) { + let mut bytes = [0_u8; 8]; + self.fd + .get_one_reg(regid, &mut bytes) + .map_err(|e| cpu::HypervisorCpuError::GetSysRegister(e.into()))?; + sys_regs.push(kvm_bindings::kvm_one_reg { + id: regid, + addr: u64::from_le_bytes(bytes), + }); + } else if is_sve_register(regid) { + let size = reg_size(regid); + let mut bytes = vec![0u8; size]; + self.fd + .get_one_reg(regid, &mut bytes) + .map_err(|e| cpu::HypervisorCpuError::GetExtendedRegister(e.into()))?; + has_sve = true; + let reg = ExtendedReg { + id: regid, + data: bytes, + }; + if PRE_FINALIZE_IDS.contains(®id) { + pre_finalize_regs.push(reg); + } else { + extended_regs.push(reg); + } + } else { + return Err(cpu::HypervisorCpuError::GetExtendedRegister(anyhow!( + "Unsupported register family: {regid:#x}" + ))); + } + } + + if !has_sve { + self.get_fpsimd_regs(&mut state.core_regs)?; } state.sys_regs = sys_regs; + state.pre_finalize_regs = pre_finalize_regs; + state.extended_regs = extended_regs; Ok(state.into()) } @@ -3116,8 +3149,24 @@ impl cpu::Vcpu for KvmVcpu { #[cfg(target_arch = "aarch64")] fn set_state(&self, state: &CpuState) -> cpu::Result<()> { let state: VcpuKvmState = state.clone().into(); + // Set core registers self.set_regs(&state.core_regs.into())?; + + let mut has_sve = state + .pre_finalize_regs + .iter() + .any(|r| is_sve_register(r.id)); + for reg in &state.extended_regs { + has_sve |= is_sve_register(reg.id); + self.fd + .set_one_reg(reg.id, ®.data) + .map_err(|e| cpu::HypervisorCpuError::SetExtendedRegister(e.into()))?; + } + if !has_sve { + self.set_fpsimd_regs(&state.core_regs)?; + } + // Set system registers for reg in &state.sys_regs { self.fd diff --git a/hypervisor/src/lib.rs b/hypervisor/src/lib.rs index f224e7217..6800cc66e 100644 --- a/hypervisor/src/lib.rs +++ b/hypervisor/src/lib.rs @@ -161,6 +161,18 @@ pub enum CpuState { Mshv(mshv::VcpuMshvState), } +#[cfg(target_arch = "aarch64")] +impl CpuState { + pub fn pre_finalize_regs(&self) -> &[arch::aarch64::ExtendedReg] { + match self { + #[cfg(feature = "kvm")] + CpuState::Kvm(state) => &state.pre_finalize_regs, + #[cfg(feature = "mshv")] + CpuState::Mshv(_) => &[], + } + } +} + #[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)] #[cfg(target_arch = "x86_64")] pub enum ClockData { diff --git a/vmm/src/cpu.rs b/vmm/src/cpu.rs index ec281b945..6a905d837 100644 --- a/vmm/src/cpu.rs +++ b/vmm/src/cpu.rs @@ -152,6 +152,10 @@ pub enum Error { #[error("Error finalising vCPU")] VcpuArmFinalize(#[source] hypervisor::HypervisorCpuError), + #[cfg(target_arch = "aarch64")] + #[error("Error setting pre-finalize registers")] + VcpuSetPreFinalizeRegs(#[source] hypervisor::HypervisorCpuError), + #[cfg(target_arch = "aarch64")] #[error("Error initialising GICR base address")] VcpuSetGicrBaseAddr(#[source] hypervisor::HypervisorCpuError), @@ -991,15 +995,22 @@ impl CpuManager { )?; if let Some(snapshot) = snapshot { - #[cfg(target_arch = "aarch64")] - { - vcpu.init(self.vm.as_ref())?; - vcpu.finalize_sve()?; - } - let state: CpuState = snapshot.to_state().map_err(|e| { Error::VcpuCreate(anyhow!("Could not get vCPU state from snapshot {e:?}")) })?; + + #[cfg(target_arch = "aarch64")] + { + vcpu.init(self.vm.as_ref())?; + let pre_finalize = state.pre_finalize_regs(); + if !pre_finalize.is_empty() { + vcpu.vcpu + .set_pre_finalize_regs(pre_finalize) + .map_err(Error::VcpuSetPreFinalizeRegs)?; + } + vcpu.finalize_sve()?; + } + vcpu.vcpu .set_state(&state) .map_err(|e| Error::VcpuCreate(anyhow!("Could not set the vCPU state {e:?}")))?;