fix(vmm): Support kernel >= 5.16 and < 5.17 for xfeature enabling

The Linux kernel introduced some arch_prctl() APIs (ARCH_GET_XCOMP_SUPP,
ARCH_GET_XCOMP_PERM and ARCH_REQ_XCOMP_PERM) to support dynamic XSTATE
feature enabling for "userspace applications" in 5.16 [1]. To support
that for "KVM guest", v5.17 introduced additional arch_prctl() APIs
(ARCH_GET_XCOMP_GUEST_PERM and ARCH_REQ_XCOMP_GUEST_PERM) [2]. Note that
ARCH_XCOMP_GUEST_SUPP does not exist, so regardless of userspace
application or KVM guest, ARCH_GET_XCOMP_SUPP is used to retrieve
supported features. As a result, kernel versions >= 5.16 and < 5.17
report supported features via ARCH_GET_XCOMP_SUPP but
ARCH_REQ_XCOMP_GUEST_PERM fail with EINVAL.

[1]: torvalds/linux@db8268d
[2]: torvalds/linux@980fe2f
Signed-off-by: Takahiro Itazuri <itazur@amazon.com>

Author: zulinx86

CHANGELOG.md

@@ -27,6 +27,10 @@ and this project adheres to
   typo in Swagger definition of `MmdsConfig`, where the property `imds_compat`
   was spelled as `imds_comat`. This caused auto-generated clients to create bad
   requests.
+- [#5447](https://github.com/firecracker-microvm/firecracker/pull/5447): Fixed
+  Intel AMX enabling for kernels that support dynamic XSTATE features for
+  userspace applications but not for KVM guests (e.g. kernel versions >= 5.16
+  and < 5.17).
 
 ## [1.13.0]
 

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

@@ -4,6 +4,7 @@
 use vmm_sys_util::syscall::SyscallReturnCode;
 
 use crate::arch::x86_64::generated::arch_prctl;
+use crate::logger::info;
 
 const INTEL_AMX_MASK: u64 = 1u64 << arch_prctl::ARCH_XCOMP_TILEDATA;
 
@@ -72,28 +73,55 @@ fn get_supported_xfeatures() -> Result<Option<u64>, std::io::Error> {
     {
         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).
+        // version prior to v5.16).
         // https://github.com/torvalds/linux/commit/db8268df0983adc2bb1fb48c9e5f7bfbb5f617f3
-        Err(err) if err.raw_os_error() == Some(libc::EINVAL) => Ok(None),
+        Err(err) if err.raw_os_error() == Some(libc::EINVAL) => {
+            info!("Dynamic XSTATE feature enabling is not supported.");
+            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.
+/// This should be called after `get_supported_xfeatures()` that retrieves supported dynamic XSTATE
+/// features.
+///
+/// Returns Ok(()) if the permission request succeeded or dynamic XSTATE feature enabling for
+/// "guest" is not 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 {
+    match 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()
+    {
+        Ok(()) => Ok(()),
+        // EINVAL is returned if the dynamic XSTATE feature enabling for "guest" is not supported
+        // although that for "userspace application" is supported (e.g. kernel versions >= 5.16 and
+        // < 5.17).
+        // https://github.com/torvalds/linux/commit/980fe2fddcff21937c93532b4597c8ea450346c1
+        //
+        // Note that XFEATURE_MASK_XTILE (= XFEATURE_MASK_XTILE_DATA | XFEATURE_MASK_XTILE_CFG) was
+        // also added to KVM_SUPPORTED_XCR0 in kernel v5.17. KVM_SUPPORTED_XCR0 is used to
+        // initialize the guest-supported XCR0. Thus, KVM_GET_SUPPORTED_CPUID doesn't
+        // return AMX-half-enabled state, where XTILE_CFG is set but XTILE_DATA is unset, on such
+        // kernels.
+        // https://github.com/torvalds/linux/commit/86aff7a4799286635efd94dab17b513544703cad
+        // https://github.com/torvalds/linux/blame/f443e374ae131c168a065ea1748feac6b2e76613/arch/x86/kvm/x86.c#L8850-L8853
+        // https://github.com/firecracker-microvm/firecracker/pull/5065
+        Err(err) if err.raw_os_error() == Some(libc::EINVAL) => {
+            info!("Dynamic XSTATE feature enabling is not supported for guest.");
+            Ok(())
+        }
+        Err(err) => Err(err),
+    }
 }
 
 #[cfg(test)]