feat: Intel AMX support

CHANGELOG.md

@@ -10,6 +10,9 @@ and this project adheres to
 
 ### Added
 
+- [#5065](https://github.com/firecracker-microvm/firecracker/pull/5065) Added
+  support for Intel AMX (Advanced Matrix Extensions).
+
 ### Changed
 
 ### Deprecated

resources/seccomp/x86_64-unknown-linux-musl.json

@@ -1078,6 +1078,18 @@
                     }
                 ]
             },
+            {
+                "syscall": "ioctl",
+                "args": [
+                    {
+                        "index": 1,
+                        "type": "dword",
+                        "op": "eq",
+                        "val": 2415963855,
+                        "comment": "KVM_GET_XSAVE2"
+                    }
+                ]
+            },
             {
                 "syscall": "ioctl",
                 "args": [

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"] }
@@ -22,8 +22,8 @@ displaydoc = "0.2.5"
 event-manager = "0.4.0"
 gdbstub = { version = "0.7.3", optional = true }
 gdbstub_arch = { version = "0.3.1", optional = true }
-kvm-bindings = { version = "0.11.0", features = ["fam-wrappers", "serde"] }
-kvm-ioctls = "0.20.0"
+kvm-bindings = { version = "0.11.1", features = ["fam-wrappers", "serde"] }
+kvm-ioctls = "0.21.0"
 libc = "0.2.170"
 linux-loader = "0.13.0"
 log = { version = "0.4.26", features = ["std", "serde"] }

src/vmm/src/arch/x86_64/gen/arch_prctl.rs

@@ -0,0 +1,43 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+// automatically generated by tools/bindgen.sh
+
+#![allow(
+    non_camel_case_types,
+    non_upper_case_globals,
+    dead_code,
+    non_snake_case,
+    clippy::ptr_as_ptr,
+    clippy::undocumented_unsafe_blocks,
+    missing_debug_implementations,
+    clippy::tests_outside_test_module
+)]
+
+pub const ARCH_SET_GS: u32 = 4097;
+pub const ARCH_SET_FS: u32 = 4098;
+pub const ARCH_GET_FS: u32 = 4099;
+pub const ARCH_GET_GS: u32 = 4100;
+pub const ARCH_GET_CPUID: u32 = 4113;
+pub const ARCH_SET_CPUID: u32 = 4114;
+pub const ARCH_GET_XCOMP_SUPP: u32 = 4129;
+pub const ARCH_GET_XCOMP_PERM: u32 = 4130;
+pub const ARCH_REQ_XCOMP_PERM: u32 = 4131;
+pub const ARCH_GET_XCOMP_GUEST_PERM: u32 = 4132;
+pub const ARCH_REQ_XCOMP_GUEST_PERM: u32 = 4133;
+pub const ARCH_XCOMP_TILECFG: u32 = 17;
+pub const ARCH_XCOMP_TILEDATA: u32 = 18;
+pub const ARCH_MAP_VDSO_X32: u32 = 8193;
+pub const ARCH_MAP_VDSO_32: u32 = 8194;
+pub const ARCH_MAP_VDSO_64: u32 = 8195;
+pub const ARCH_GET_UNTAG_MASK: u32 = 16385;
+pub const ARCH_ENABLE_TAGGED_ADDR: u32 = 16386;
+pub const ARCH_GET_MAX_TAG_BITS: u32 = 16387;
+pub const ARCH_FORCE_TAGGED_SVA: u32 = 16388;
+pub const ARCH_SHSTK_ENABLE: u32 = 20481;
+pub const ARCH_SHSTK_DISABLE: u32 = 20482;
+pub const ARCH_SHSTK_LOCK: u32 = 20483;
+pub const ARCH_SHSTK_UNLOCK: u32 = 20484;
+pub const ARCH_SHSTK_STATUS: u32 = 20485;
+pub const ARCH_SHSTK_SHSTK: u32 = 1;
+pub const ARCH_SHSTK_WRSS: u32 = 2;

src/vmm/src/arch/x86_64/gen/mod.rs

@@ -5,9 +5,9 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the THIRD-PARTY file.
 
+pub mod arch_prctl;
 pub mod hyperv;
 pub mod hyperv_tlfs;
 pub mod mpspec;
-
 pub mod msr_index;
 pub mod perf_event;

src/vmm/src/arch/x86_64/mod.rs

@@ -19,6 +19,8 @@ mod mptable;
 pub mod msr;
 /// Logic for configuring x86_64 registers.
 pub mod regs;
+/// Logic for configuring XSTATE features.
+pub mod xstate;
 
 #[allow(missing_docs)]
 pub mod gen;

src/vmm/src/arch/x86_64/xstate.rs

@@ -0,0 +1,136 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use vmm_sys_util::syscall::SyscallReturnCode;
+
+use crate::arch::x86_64::gen::arch_prctl;
+
+const INTEL_AMX_MASK: u64 = 1u64 << arch_prctl::ARCH_XCOMP_TILEDATA;
+
+/// Errors assocaited with x86_64's dynamic XSAVE state features.
+#[derive(Debug, thiserror::Error, displaydoc::Display)]
+pub enum XstateError {
+    /// Failed to get supported XSTATE features: {0}
+    GetSupportedFeatures(std::io::Error),
+    /// Failed to request permission for XSTATE feature ({0}): {1}
+    RequestFeaturePermission(u32, std::io::Error),
+}
+
+/// Request permission for all dynamic XSTATE features.
+///
+/// Some XSTATE features are not permitted by default, because they may require a larger area to
+/// save their states than the tranditional 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
+///
+/// Firecracker requests permission for them by default if available in order to retrieve the
+/// full supported feature set via KVM_GET_SUPPORTED_CPUID.
+/// https://docs.kernel.org/virt/kvm/api.html#kvm-get-supported-cpuid
+///
+/// Note that requested features can be masked by a CPU template.
+pub fn request_dynamic_xstate_features() -> Result<(), XstateError> {
+    let supported_xfeatures =
+        match get_supported_xfeatures().map_err(XstateError::GetSupportedFeatures)? {
+            Some(supported_xfeatures) => supported_xfeatures,
+            // Exit early if dynamic XSTATE feature enabling is not supported on the kernel.
+            None => return Ok(()),
+        };
+
+    // Intel AMX's 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.
+    if supported_xfeatures & INTEL_AMX_MASK == INTEL_AMX_MASK {
+        request_xfeature_permission(arch_prctl::ARCH_XCOMP_TILEDATA).map_err(|err| {
+            XstateError::RequestFeaturePermission(arch_prctl::ARCH_XCOMP_TILEDATA, err)
+        })?;
+    }
+
+    Ok(())
+}
+
+/// Get supported XSTATE features
+///
+/// Returns Ok(None) if dynamic XSTATE feature enabling is not supported.
+fn get_supported_xfeatures() -> Result<Option<u64>, std::io::Error> {
+    let mut supported_xfeatures: u64 = 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
+    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(()) => Ok(Some(supported_xfeatures)),
+        // 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
+        Err(err) if err.raw_os_error() == Some(libc::EINVAL) => Ok(None),
+        Err(err) => Err(err),
+    }
+}
+
+/// Request permission for a dynamic XSTATE feature.
+///
+/// This should be called after `get_supported_xfeatures()` that also checks that dynamic XSTATE
+/// feature enabling is supported.
+fn request_xfeature_permission(xfeature: u32) -> Result<(), std::io::Error> {
+    // SAFETY: Safe because the third input (`addr`) is a valid `c_ulong` value.
+    // https://man7.org/linux/man-pages/man2/arch_prctl.2.html
+    SyscallReturnCode(unsafe {
+        libc::syscall(
+            libc::SYS_arch_prctl,
+            arch_prctl::ARCH_REQ_XCOMP_GUEST_PERM as libc::c_ulong,
+            xfeature as libc::c_ulong,
+        )
+    })
+    .into_empty_result()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // Get permitted XSTATE features.
+    fn get_permitted_xstate_features() -> Result<u64, std::io::Error> {
+        let mut permitted_xfeatures: u64 = 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_GUEST_PERM,
+                &mut permitted_xfeatures as *mut libc::c_ulong,
+            )
+        })
+        .into_empty_result()
+        {
+            Ok(()) => Ok(permitted_xfeatures),
+            Err(err) => Err(err),
+        }
+    }
+
+    #[test]
+    fn test_request_xstate_feature_permission() {
+        request_dynamic_xstate_features().unwrap();
+
+        let supported_xfeatures = match get_supported_xfeatures().unwrap() {
+            Some(supported_xfeatures) => supported_xfeatures,
+            // Nothing to test if dynamic XSTATE feature enabling is not supported on the kernel.
+            None => return,
+        };
+
+        // Check each dynamic feature is enabled. (currently only Intel AMX TILEDATA)
+        if supported_xfeatures & INTEL_AMX_MASK == INTEL_AMX_MASK {
+            let permitted_xfeatures = get_permitted_xstate_features().unwrap();
+            assert_eq!(permitted_xfeatures & INTEL_AMX_MASK, INTEL_AMX_MASK);
+        }
+    }
+}

src/vmm/src/vstate/kvm.rs

@@ -7,13 +7,15 @@ 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 crate::arch::x86_64::xstate::{request_dynamic_xstate_features, XstateError};
 use crate::cpu_config::templates::KvmCapability;
 use crate::vstate::memory::{GuestMemory, GuestMemoryMmap};
 
 /// Errors associated with the wrappers over KVM ioctls.
 /// Needs `rustfmt::skip` to make multiline comments work
 #[rustfmt::skip]
-#[derive(Debug, PartialEq, Eq, thiserror::Error, displaydoc::Display)]
+#[derive(Debug, thiserror::Error, displaydoc::Display)]
 pub enum KvmError {
     /// The host kernel reports an invalid KVM API version: {0}
     ApiVersion(i32),
@@ -27,6 +29,9 @@ pub enum KvmError {
     GetSupportedCpuId(kvm_ioctls::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 dynamic XSTATE features: {0}
+    XstateFeatures(XstateError),
 }
 
 /// Struct with kvm fd and kvm associated paramenters.
@@ -73,6 +78,8 @@ impl Kvm {
 
         #[cfg(target_arch = "x86_64")]
         {
+            request_dynamic_xstate_features().map_err(KvmError::XstateFeatures)?;
+
             let supported_cpuid = kvm_fd
                 .get_supported_cpuid(KVM_MAX_CPUID_ENTRIES)
                 .map_err(KvmError::GetSupportedCpuId)?;