feat: Support Intel AMX

Intel AMX was introduced in Intel Sapphire Rapids and stands for
Advanced Matrix Extensions that accelerates deep learning and artificial
intelligence workloads. Since it requires larger area to save its state,
the TILEDATA feature is disabled by default. We enable 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 enable XSTATE features ({0:#b}): {1}
+    RequestXstateFeatures(u32, std::io::Error),
 }
 
 /// Struct with kvm fd and kvm associated paramenters.
@@ -73,6 +83,8 @@ impl Kvm {
 
         #[cfg(target_arch = "x86_64")]
         {
+            Self::enable_intel_amx()?;
+
             let supported_cpuid = kvm_fd
                 .get_supported_cpuid(KVM_MAX_CPUID_ENTRIES)
                 .map_err(KvmError::GetSupportedCpuId)?;
@@ -86,6 +98,74 @@ impl Kvm {
         }
     }
 
+    #[cfg(target_arch = "x86_64")]
+    // XSTATE feature mask for Intel AMX.
+    const INTEL_AMX_XCOMP_MASK: libc::c_ulong =
+        (1u64 << arch_prctl::ARCH_XCOMP_TILECFG) | (1u64 << arch_prctl::ARCH_XCOMP_TILEDATA);
+
+    /// Enable Intel AMX if available.
+    ///
+    /// Intel AMX (Advanced Matrix Extensions) is an instruction set for AI workloads that was
+    /// introduced in Intel Sapphire Rapids (*7i.metal). Since it requires larger area to save the
+    /// state, it is disabled by default.
+    /// https://github.com/torvalds/linux/blob/master/Documentation/arch/x86/xstate.rst
+    ///
+    /// We enable it by default but can be disabled by CPU template; otherwise,
+    /// KVM_GET_SUPPORTED_CPUID returns a inconsistent state where TILECFG is enabled but TILEDATA
+    /// is disabled, causing guest's #GP fault on xsetbv due to the lack of sanity check.
+    /// https://lore.kernel.org/all/[email protected]/
+    ///
+    /// Dynamically-enabled feature bits need to be requested with arch_prctl() before calling
+    /// KVM_GET_SUPPORTED_CPUID. Feature bits that have not been requested are excluded from the
+    /// result of KVM_GET_SUPPORTED_CPUID.
+    /// https://docs.kernel.org/virt/kvm/api.html
+    ///
+    /// Note that no memory allocation to save Intel AMX state happens here immediately.
+    #[cfg(target_arch = "x86_64")]
+    fn enable_intel_amx() -> Result<(), KvmError> {
+        // Get the supported xstate features.
+        let mut supported_xfeatures: libc::c_ulong = 0;
+        // SAFETY: Safe because the second input (`op`) might not be valid for unsupported kernels
+        // but EINVAL is handled later, and 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| {
+            // EINVAL is returned if ARCH_GET_XCOMP_SUPP is not supported (e.g. kernel versions
+            // prior to v5.17).
+            // https://github.com/torvalds/linux/commit/980fe2fddcff21937c93532b4597c8ea450346c1
+            if err.raw_os_error() == Some(libc::EINVAL) {
+                Ok(())
+            } else {
+                Err(err)
+            }
+        })
+        .map_err(KvmError::GetSupportedXstateFeatures)?;
+
+        // Enable Intel AMX if supported.
+        if (supported_xfeatures & Self::INTEL_AMX_XCOMP_MASK) == Self::INTEL_AMX_XCOMP_MASK {
+            // SAFETY: Safe because ARCH_REQ_XCOMP_GUEST_PERM is supported if ARCH_GET_XCOMP_SUPP is
+            // supported and it has been confirmed that ARCH_XCOMP_TILEDATA is supported.
+            SyscallReturnCode(unsafe {
+                libc::syscall(
+                    libc::SYS_arch_prctl,
+                    arch_prctl::ARCH_REQ_XCOMP_GUEST_PERM,
+                    arch_prctl::ARCH_XCOMP_TILEDATA,
+                )
+            })
+            .into_empty_result()
+            .map_err(|err| KvmError::RequestXstateFeatures(arch_prctl::ARCH_XCOMP_TILEDATA, 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 +295,99 @@ pub(crate) mod tests {
             .iter()
             .any(|c| *c == kvm_bindings::KVM_CAP_IOEVENTFD));
     }
+
+    #[cfg(target_arch = "x86_64")]
+    mod x86_64 {
+        use super::*;
+        use crate::arch::x86_64::cpu_model::CpuModel;
+
+        #[derive(PartialEq, PartialOrd)]
+        struct KernelVersion(u32, u32);
+
+        impl KernelVersion {
+            fn current() -> Self {
+                let version_str = std::fs::read_to_string("/proc/sys/kernel/osrelease").unwrap();
+                let mut parts = version_str.trim().split('.');
+
+                let major = parts.next().unwrap().parse::<u32>().unwrap();
+                let minor = parts.next().unwrap().parse::<u32>().unwrap();
+
+                KernelVersion(major, minor)
+            }
+        }
+
+        #[derive(PartialEq)]
+        enum Vendor {
+            Intel,
+            Amd,
+        }
+
+        impl Vendor {
+            fn new() -> Self {
+                let vendor_id = Self::get_vendor_id_str();
+                match vendor_id.as_str() {
+                    "GenuineIntel" => Vendor::Intel,
+                    "AuthenticAMD" => Vendor::Amd,
+                    _ => panic!("Unknown vendor_id: {}", vendor_id),
+                }
+            }
+
+            fn get_vendor_id_str() -> String {
+                let cpuinfo = std::fs::read_to_string("/proc/cpuinfo").unwrap();
+
+                for line in cpuinfo.lines() {
+                    if line.starts_with("vendor_id") {
+                        return line
+                            .split(':')
+                            .nth(1)
+                            .map(|s| s.trim().to_string())
+                            .unwrap();
+                    }
+                }
+                panic!("`vendor_id` not found in /proc/cpuinfo");
+            }
+        }
+
+        #[cfg(target_arch = "x86_64")]
+        #[test]
+        fn test_enable_intel_amx() {
+            Kvm::enable_intel_amx().unwrap();
+
+            // ARCH_{REQ,GET}_XCOMP_GUEST_PERM were added in kernel v5.17.
+            //  https://github.com/torvalds/linux/commit/980fe2fddcff21937c93532b4597c8ea450346c1
+            let supported_version = KernelVersion(5, 17);
+            let current_version = KernelVersion::current();
+
+            if current_version >= supported_version {
+                let mut permitted_xfeatures: libc::c_ulong = 0;
+                // SAFETY: Safe because the second input (`op`) should be valid on kernel v5.17+,
+                // and 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();
+
+                // Intel AMX is available only on Intel processors now.
+                let vendor = Vendor::new();
+
+                // Intel AMX is introduced in Intel Sapphire Rapids (CPUID.01H:EAX = 0x000806f8).
+                let supported_cpu = CpuModel::from(&0x000806f8);
+                let current_cpu = CpuModel::get_cpu_model();
+
+                if current_cpu >= supported_cpu && vendor == Vendor::Intel {
+                    assert_eq!(
+                        permitted_xfeatures & Kvm::INTEL_AMX_XCOMP_MASK,
+                        Kvm::INTEL_AMX_XCOMP_MASK
+                    );
+                } else {
+                    assert_eq!(permitted_xfeatures & Kvm::INTEL_AMX_XCOMP_MASK, 0);
+                }
+            }
+        }
+    }
 }