feat: Request permission for Intel AMX

Intel AMX (Advanced Matrix Extensions) was introduced in Intel Sapphire
Rapids to accelerate deep learning and AI workloads. Since it requires
a larger area to save its state, the TILEDATA feature is disabled by
default. We request permission for it by default because it can be
disabled via CPU template. Otherwise, kernels prior to v6.4 have a bug
where KVM_GET_SUPPORTED_CPUID returns an inconsistent state of TILECFG
enabled but TILEDATA disabled by default, causing guest's #GP fault on
xsetbv instruction.

Signed-off-by: Takahiro Itazuri <[email protected]>

Author: zulinx86

src/vmm/Cargo.toml

@@ -9,7 +9,7 @@ license = "Apache-2.0"
 bench = false
 
 [dependencies]
-acpi_tables = { path = "../acpi-tables" } 
+acpi_tables = { path = "../acpi-tables" }
 aes-gcm =  { version = "0.10.1", default-features = false, features = ["aes"] }
 arrayvec = { version = "0.7.6", optional = true }
 aws-lc-rs = { version = "1.12.4", features = ["bindgen"] }

src/vmm/src/vstate/kvm.rs

@@ -6,7 +6,11 @@ use kvm_bindings::KVM_API_VERSION;
 use kvm_bindings::{CpuId, MsrList, KVM_MAX_CPUID_ENTRIES};
 use kvm_ioctls::Kvm as KvmFd;
 use serde::{Deserialize, Serialize};
+#[cfg(target_arch = "x86_64")]
+use vmm_sys_util::syscall::SyscallReturnCode;
 
+#[cfg(target_arch = "x86_64")]
+use crate::arch::x86_64::gen::arch_prctl;
 use crate::cpu_config::templates::KvmCapability;
 use crate::vstate::memory::{GuestMemory, GuestMemoryMmap};
 
@@ -25,8 +29,14 @@ pub enum KvmError {
     #[cfg(target_arch = "x86_64")]
     /// Failed to get supported cpuid: {0}
     GetSupportedCpuId(kvm_ioctls::Error),
+    #[cfg(target_arch = "x86_64")]
+    /// Failed to get supported XSTATE features: {0}
+    GetSupportedXstateFeatures(std::io::Error),
     /// The number of configured slots is bigger than the maximum reported by KVM
     NotEnoughMemorySlots,
+    #[cfg(target_arch = "x86_64")]
+    /// Failed to request permission for a XSTATE feature ({0}): {1}
+    RequestXstateFeatures(u32, std::io::Error),
 }
 
 /// Struct with kvm fd and kvm associated paramenters.
@@ -73,6 +83,15 @@ impl Kvm {
 
         #[cfg(target_arch = "x86_64")]
         {
+            // Request permission for Intel AMX (Advanced Matrix Extensions) TILEDATA.
+            //
+            // Unless requested, on kernels prior to v6.4, KVM_GET_SUPPORTED_CPUID returns an
+            // inconsistent state where TILECFG is set but TILEDATA isn't. Such a half-enabled state
+            // causes guest crash during boot because a guest calls XSETBV instruction with all
+            // XSAVE feature bits enumerated on CPUID and XSETBV only accepts either of both Intel
+            // AMX bits enabled or disabled; otherwise resulting in general protection fault.
+            Self::request_xstate_feature_permission(arch_prctl::ARCH_XCOMP_TILEDATA)?;
+
             let supported_cpuid = kvm_fd
                 .get_supported_cpuid(KVM_MAX_CPUID_ENTRIES)
                 .map_err(KvmError::GetSupportedCpuId)?;
@@ -86,6 +105,65 @@ impl Kvm {
         }
     }
 
+    /// Request permission for a dynamic XSTATE features.
+    ///
+    /// Some XSTATE features are not permitted by default, because they require a larger area
+    /// to save their states than the traditional 4096-byte area. Instead, the permission for them
+    /// can be requested via arch_prctl().
+    /// https://github.com/torvalds/linux/blob/master/Documentation/arch/x86/xstate.rst
+    ///
+    /// We request permission for them by default if available in order to retrieve the correct
+    /// supported feature set via KVM_GET_SUPPORTED_CPUID.
+    /// https://docs.kernel.org/virt/kvm/api.html#kvm-get-supported-cpuid
+    ///
+    /// Note that such requested features can be disabled by a CPU template and no memory allocation
+    /// to save their states happens here immediately.
+    #[cfg(target_arch = "x86_64")]
+    fn request_xstate_feature_permission(xfeature: u32) -> Result<(), KvmError> {
+        // Get the supported dynamic xstate features.
+        let mut supported_xfeatures: libc::c_ulong = 0;
+        // SAFETY: Safe because the third input (`addr`) is a valid `c_ulong`` pointer.
+        // https://man7.org/linux/man-pages/man2/arch_prctl.2.html
+        SyscallReturnCode(unsafe {
+            libc::syscall(
+                libc::SYS_arch_prctl,
+                arch_prctl::ARCH_GET_XCOMP_SUPP,
+                &mut supported_xfeatures as *mut libc::c_ulong,
+            )
+        })
+        .into_empty_result()
+        .or_else(|err| {
+            if err.raw_os_error() == Some(libc::EINVAL) {
+                // EINVAL is returned if the dynamic XSTATE feature enabling is not supported (e.g.
+                // kernel version prior to v5.17).
+                // https://github.com/torvalds/linux/commit/980fe2fddcff21937c93532b4597c8ea450346c1
+                //
+                // `supported_xfeatures` remains 0 here, so will skip permission request.
+                Ok(())
+            } else {
+                Err(err)
+            }
+        })
+        .map_err(KvmError::GetSupportedXstateFeatures)?;
+
+        // Request permission for the given XSTATE feature only if available
+        let xfeature_mask: libc::c_ulong = 1u64 << xfeature;
+        if (supported_xfeatures & xfeature_mask) == xfeature_mask {
+            // SAFETY: Safe because all inputs are valid as `c_ulong`` values.
+            SyscallReturnCode(unsafe {
+                libc::syscall(
+                    libc::SYS_arch_prctl,
+                    arch_prctl::ARCH_REQ_XCOMP_GUEST_PERM,
+                    xfeature,
+                )
+            })
+            .into_empty_result()
+            .map_err(|err| KvmError::RequestXstateFeatures(xfeature, err))?;
+        }
+
+        Ok(())
+    }
+
     /// Msrs needed to be saved on snapshot creation.
     #[cfg(target_arch = "x86_64")]
     pub fn msrs_to_save(&self) -> Result<MsrList, crate::arch::x86_64::msr::MsrError> {
@@ -215,4 +293,54 @@ pub(crate) mod tests {
             .iter()
             .any(|c| *c == kvm_bindings::KVM_CAP_IOEVENTFD));
     }
+
+    #[cfg(target_arch = "x86_64")]
+    #[test]
+    fn test_request_xstate_feature_permission() {
+        // Test request_xstate_feature_permission() for Intel AMX.
+        Kvm::request_xstate_feature_permission(arch_prctl::ARCH_XCOMP_TILEDATA).unwrap();
+
+        let mut supported_xfeatures: libc::c_ulong = 0;
+        // SAFETY: Safe because the third input (`addr`) is a valid `c_ulong` pointer.
+        match SyscallReturnCode(unsafe {
+            libc::syscall(
+                libc::SYS_arch_prctl,
+                arch_prctl::ARCH_GET_XCOMP_SUPP,
+                &mut supported_xfeatures as *mut libc::c_ulong,
+            )
+        })
+        .into_empty_result()
+        {
+            Ok(_) => {} // Continue this test
+            Err(err) if err.raw_os_error() == Some(libc::EINVAL) => {
+                // Dynamic XSTATE feature enabling is not supported in the first place, so nothing
+                // to test on this kernel version.
+                return;
+            }
+            Err(err) => panic!("Unexpected error: {}", err),
+        };
+
+        // If Intel AMX is not supported, nothing to test on this CPU.
+        let intel_amx_feature_mask: libc::c_ulong = 1u64 << arch_prctl::ARCH_XCOMP_TILEDATA;
+        if supported_xfeatures & intel_amx_feature_mask != intel_amx_feature_mask {
+            return;
+        }
+
+        let mut permitted_xfeatures: libc::c_ulong = 0;
+        // SAFETY: Safe because the third input (`addr`) is a valid `c_ulong` pointer.
+        SyscallReturnCode(unsafe {
+            libc::syscall(
+                libc::SYS_arch_prctl,
+                arch_prctl::ARCH_GET_XCOMP_GUEST_PERM,
+                &mut permitted_xfeatures as *mut libc::c_ulong,
+            )
+        })
+        .into_empty_result()
+        .unwrap();
+
+        assert_eq!(
+            permitted_xfeatures & intel_amx_feature_mask,
+            intel_amx_feature_mask
+        );
+    }
 }