From d2a364c5c01eab1d921beb3b85f1c497f15614cb Mon Sep 17 00:00:00 2001
From: Wenyu Huang <huangwenyuu@outlook.com>
Date: Fri, 23 Aug 2024 07:12:17 -0400
Subject: [PATCH] vmm: Add support for enabling SVE in vm guests

This change enables SVE automatically if the host support SVE/SVE2.

Signed-off-by: Wenyu Huang <huangwenyuu@outlook.com>
---
 hypervisor/src/cpu.rs      |  9 +++++++++
 hypervisor/src/kvm/mod.rs  |  9 +++++++++
 vmm/src/cpu.rs             | 27 +++++++++++++++++++++++++--
 vmm/src/seccomp_filters.rs |  2 ++
 4 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/hypervisor/src/cpu.rs b/hypervisor/src/cpu.rs
index 985f89608..5b0581299 100644
--- a/hypervisor/src/cpu.rs
+++ b/hypervisor/src/cpu.rs
@@ -158,6 +158,11 @@ pub enum HypervisorCpuError {
     #[error("Failed to init vcpu: {0}")]
     VcpuInit(#[source] anyhow::Error),
     ///
+    /// Vcpu Finalize error
+    ///
+    #[error("Failed to finalize vcpu: {0}")]
+    VcpuFinalize(#[source] anyhow::Error),
+    ///
     /// Setting one reg error
     ///
     #[error("Failed to init vcpu: {0}")]
@@ -416,6 +421,10 @@ pub trait Vcpu: Send + Sync {
     ///
     #[cfg(target_arch = "aarch64")]
     fn vcpu_init(&self, kvi: &VcpuInit) -> Result<()>;
+
+    #[cfg(target_arch = "aarch64")]
+    fn vcpu_finalize(&self, feature: i32) -> Result<()>;
+
     ///
     /// Gets a list of the guest registers that are supported for the
     /// KVM_GET_ONE_REG/KVM_SET_ONE_REG calls.
diff --git a/hypervisor/src/kvm/mod.rs b/hypervisor/src/kvm/mod.rs
index fc87a7510..c636118d2 100644
--- a/hypervisor/src/kvm/mod.rs
+++ b/hypervisor/src/kvm/mod.rs
@@ -1881,6 +1881,15 @@ impl cpu::Vcpu for KvmVcpu {
             .map_err(|e| cpu::HypervisorCpuError::VcpuInit(e.into()))
     }
 
+    #[cfg(target_arch = "aarch64")]
+    fn vcpu_finalize(&self, feature: i32) -> cpu::Result<()> {
+        self.fd
+            .lock()
+            .unwrap()
+            .vcpu_finalize(&feature)
+            .map_err(|e| cpu::HypervisorCpuError::VcpuFinalize(e.into()))
+    }
+
     ///
     /// Gets a list of the guest registers that are supported for the
     /// KVM_GET_ONE_REG/KVM_SET_ONE_REG calls.
diff --git a/vmm/src/cpu.rs b/vmm/src/cpu.rs
index 7ca7764e1..f71212b48 100644
--- a/vmm/src/cpu.rs
+++ b/vmm/src/cpu.rs
@@ -140,6 +140,10 @@ pub enum Error {
     #[error("Error initialising vCPU: {0}")]
     VcpuArmInit(#[source] hypervisor::HypervisorCpuError),
 
+    #[cfg(target_arch = "aarch64")]
+    #[error("Error finalising vCPU: {0}")]
+    VcpuArmFinalize(#[source] hypervisor::HypervisorCpuError),
+
     #[error("Failed to join on vCPU threads: {0:?}")]
     ThreadCleanup(std::boxed::Box<dyn std::any::Any + std::marker::Send>),
 
@@ -415,8 +419,10 @@ impl Vcpu {
     /// Initializes an aarch64 specific vcpu for booting Linux.
     #[cfg(target_arch = "aarch64")]
     pub fn init(&self, vm: &Arc<dyn hypervisor::Vm>) -> Result<()> {
+        use std::arch::is_aarch64_feature_detected;
         let mut kvi: kvm_bindings::kvm_vcpu_init = kvm_bindings::kvm_vcpu_init::default();
-
+        let sve_supported =
+            is_aarch64_feature_detected!("sve") || is_aarch64_feature_detected!("sve2");
         // This reads back the kernel's preferred target type.
         vm.get_preferred_target(&mut kvi)
             .map_err(Error::VcpuArmPreferredTarget)?;
@@ -430,11 +436,28 @@ impl Vcpu {
         {
             kvi.features[0] |= 1 << kvm_bindings::KVM_ARM_VCPU_PMU_V3;
         }
+
+        if sve_supported
+            && vm
+                .as_any()
+                .downcast_ref::<hypervisor::kvm::KvmVm>()
+                .unwrap()
+                .check_extension(Cap::ArmSve)
+        {
+            kvi.features[0] |= 1 << kvm_bindings::KVM_ARM_VCPU_SVE;
+        }
+
         // Non-boot cpus are powered off initially.
         if self.id > 0 {
             kvi.features[0] |= 1 << kvm_bindings::KVM_ARM_VCPU_POWER_OFF;
         }
-        self.vcpu.vcpu_init(&kvi).map_err(Error::VcpuArmInit)
+        self.vcpu.vcpu_init(&kvi).map_err(Error::VcpuArmInit)?;
+        if sve_supported {
+            self.vcpu
+                .vcpu_finalize(kvm_bindings::KVM_ARM_VCPU_SVE as i32)
+                .map_err(Error::VcpuArmFinalize)?;
+        }
+        Ok(())
     }
 
     /// Runs the VCPU until it exits, returning the reason.
diff --git a/vmm/src/seccomp_filters.rs b/vmm/src/seccomp_filters.rs
index 4782960a4..3a6c75a98 100644
--- a/vmm/src/seccomp_filters.rs
+++ b/vmm/src/seccomp_filters.rs
@@ -427,12 +427,14 @@ fn create_vmm_ioctl_seccomp_rule_kvm() -> Result<Vec<SeccompRule>, BackendError>
     const KVM_ARM_PREFERRED_TARGET: u64 = 0x8020_aeaf;
     const KVM_ARM_VCPU_INIT: u64 = 0x4020_aeae;
     const KVM_SET_GUEST_DEBUG: u64 = 0x4208_ae9b;
+    const KVM_ARM_VCPU_FINALIZE: u64 = 0x4004_aec2;
 
     let common_rules = create_vmm_ioctl_seccomp_rule_common(HypervisorType::Kvm)?;
     let mut arch_rules = or![
         and![Cond::new(1, ArgLen::Dword, Eq, KVM_ARM_PREFERRED_TARGET,)?],
         and![Cond::new(1, ArgLen::Dword, Eq, KVM_ARM_VCPU_INIT,)?],
         and![Cond::new(1, ArgLen::Dword, Eq, KVM_SET_GUEST_DEBUG,)?],
+        and![Cond::new(1, ArgLen::Dword, Eq, KVM_ARM_VCPU_FINALIZE,)?],
     ];
     arch_rules.extend(common_rules);