x86_64::task: FPU, SSE

Signed-off-by: Andy-Python-Programmer <[email protected]>

Author: Andy-Python-Programmer

src/aero_kernel/src/arch/x86_64/controlregs.rs

@@ -185,6 +185,73 @@ bitflags::bitflags! {
     }
 }
 
+bitflags::bitflags! {
+    pub struct MxCsr: u32 {
+        const INVALID_OPERATION = 1 << 0;
+        const DENORMAL = 1 << 1;
+        const DIVIDE_BY_ZERO = 1 << 2;
+        const OVERFLOW = 1 << 3;
+        const UNDERFLOW = 1 << 4;
+        const PRECISION = 1 << 5;
+        const DENORMALS_ARE_ZEROS = 1 << 6;
+        const INVALID_OPERATION_MASK = 1 << 7;
+        const DENORMAL_MASK = 1 << 8;
+        const DIVIDE_BY_ZERO_MASK = 1 << 9;
+        const OVERFLOW_MASK = 1 << 10;
+        const UNDERFLOW_MASK = 1 << 11;
+        const PRECISION_MASK = 1 << 12;
+        const ROUNDING_CONTROL_NEGATIVE = 1 << 13;
+        const ROUNDING_CONTROL_POSITIVE = 1 << 14;
+        const ROUNDING_CONTROL_ZERO = 3 << 13;
+        const FLUSH_TO_ZERO = 1 << 15;
+    }
+}
+
+bitflags::bitflags! {
+    /// Configuration flags of the XCr0 register.
+    ///
+    /// For MPX, [`BNDREG`](XCr0Flags::BNDREG) and [`BNDCSR`](XCr0Flags::BNDCSR) must be set/unset simultaneously.
+    /// For AVX-512, [`OPMASK`](XCr0Flags::OPMASK), [`ZMM_HI256`](XCr0Flags::ZMM_HI256), and [`HI16_ZMM`](XCr0Flags::HI16_ZMM)
+    /// must be set/unset simultaneously.
+    #[repr(transparent)]
+    pub struct XCr0Flags: u64 {
+        /// Enables using the x87 FPU state
+        /// with `XSAVE`/`XRSTOR`.
+        ///
+        /// Must be set.
+        const X87 = 1;
+        /// Enables using MXCSR and the XMM registers
+        /// with `XSAVE`/`XRSTOR`.
+        ///
+        /// Must be set if [`AVX`](XCr0Flags::AVX) is set.
+        const SSE = 1 << 1;
+        /// Enables AVX instructions and using the upper halves of the AVX registers
+        /// with `XSAVE`/`XRSTOR`.
+        const AVX = 1 << 2;
+        /// Enables MPX instructions and using the BND0-BND3 bound registers
+        /// with `XSAVE`/`XRSTOR` (Intel Only).
+        const BNDREG = 1 << 3;
+        /// Enables MPX instructions and using the BNDCFGU and BNDSTATUS registers
+        /// with `XSAVE`/`XRSTOR` (Intel Only).
+        const BNDCSR = 1 << 4;
+        /// Enables AVX-512 instructions and using the K0-K7 mask registers
+        /// with `XSAVE`/`XRSTOR` (Intel Only).
+        const OPMASK = 1 << 5;
+        /// Enables AVX-512 instructions and using the upper halves of the lower ZMM registers
+        /// with `XSAVE`/`XRSTOR` (Intel Only).
+        const ZMM_HI256 = 1 << 6;
+        /// Enables AVX-512 instructions and using the upper ZMM registers
+        /// with `XSAVE`/`XRSTOR` (Intel Only).
+        const HI16_ZMM = 1 << 7;
+        /// Enables using the PKRU register
+        /// with `XSAVE`/`XRSTOR`.
+        const MPK = 1 << 9;
+        /// Enables Lightweight Profiling extensions and managing LWP state
+        /// with `XSAVE`/`XRSTOR` (AMD Only).
+        const LWP = 1 << 62;
+    }
+}
+
 /// Returns the current value of the RFLAGS register.
 pub fn read_rflags() -> RFlags {
     let value: u64;
@@ -230,6 +297,35 @@ pub fn read_cr0() -> Cr0Flags {
     Cr0Flags::from_bits_truncate(value) // Get the flags from the bits.
 }
 
+pub fn read_xcr0() -> XCr0Flags {
+    let (low, high): (u32, u32);
+
+    unsafe {
+        asm!(
+            "xgetbv",
+            in("ecx") 0,
+            out("rax") low, out("rdx") high,
+            options(nomem, nostack, preserves_flags),
+        );
+    }
+
+    XCr0Flags::from_bits_truncate((high as u64) << 32 | (low as u64))
+}
+
+pub unsafe fn write_xcr0(value: XCr0Flags) {
+    let low = value.bits() as u32;
+    let high = (value.bits() >> 32) as u32;
+
+    unsafe {
+        asm!(
+            "xsetbv",
+            in("ecx") 0,
+            in("rax") low, in("rdx") high,
+            options(nomem, nostack, preserves_flags),
+        );
+    }
+}
+
 /// Write the given set of CR4 flags.
 ///
 /// ## Safety
@@ -273,3 +369,17 @@ pub fn read_cr2() -> VirtAddr {
         VirtAddr::new(value)
     }
 }
+
+pub fn read_mxcsr() -> MxCsr {
+    let mut mxcsr: u32 = 0;
+    unsafe {
+        asm!("stmxcsr [{}]", in(reg) &mut mxcsr, options(nostack, preserves_flags));
+    }
+    MxCsr::from_bits_truncate(mxcsr)
+}
+
+pub fn write_mxcsr(value: MxCsr) {
+    unsafe {
+        asm!("ldmxcsr [{}]", in(reg) &value, options(nostack, readonly));
+    }
+}

src/aero_kernel/src/arch/x86_64/interrupts/exceptions.rs

@@ -58,10 +58,25 @@ interrupt_exception!(fn protection() => "Protection Fault");
 interrupt_exception!(fn fpu_fault() => "FPU floating point fault");
 interrupt_exception!(fn alignment_check() => "Alignment check fault");
 interrupt_exception!(fn machine_check() => "Machine check fault");
-interrupt_exception!(fn simd() => "SIMD floating point fault");
 interrupt_exception!(fn virtualization() => "Virtualization fault");
 interrupt_exception!(fn security() => "Security exception");
 
+pub fn simd(stack: &mut InterruptErrorStack) {
+    unwind::prepare_panic();
+
+    log::error!("EXCEPTION: SIMD floating point fault");
+    log::error!("Stack: {:#x?}", stack);
+    log::error!("MXCSR: {:?}", controlregs::read_mxcsr());
+
+    unwind::unwind_stack_trace();
+
+    unsafe {
+        loop {
+            super::halt();
+        }
+    }
+}
+
 pub fn invalid_opcode(stack: &mut InterruptErrorStack) {
     // Catch SYSENTER on AMD CPUs.
     //

src/aero_kernel/src/arch/x86_64/mod.rs

@@ -235,17 +235,27 @@ pub fn enable_acpi() {
     aml::get_subsystem().enable_acpi(INTERRUPT_CONTROLLER.method() as _);
 }
 
+fn enable_xsave() {
+    use controlregs::XCr0Flags;
+
+    // Enable XSAVE and x{get,set}bv
+    let mut cr4 = controlregs::read_cr4();
+    cr4.insert(controlregs::Cr4Flags::OSXSAVE);
+    unsafe { controlregs::write_cr4(cr4) }
+
+    let mut xcr0 = controlregs::read_xcr0();
+    xcr0.insert(XCr0Flags::X87 | XCr0Flags::SSE);
+    unsafe { controlregs::write_xcr0(xcr0) }
+}
+
 pub fn init_cpu() {
     unsafe {
         // Enable the no-execute page protection feature.
         io::wrmsr(io::IA32_EFER, io::rdmsr(io::IA32_EFER) | 1 << 11);
 
-        // Check if SSE is supported. SSE support is a requirement for running Aero.
-        let has_sse = CpuId::new()
-            .get_feature_info()
-            .map_or(false, |i| i.has_sse());
+        let features = CpuId::new().get_feature_info().unwrap();
 
-        assert!(has_sse);
+        assert!(features.has_sse());
 
         {
             let mut cr0 = controlregs::read_cr0();
@@ -264,5 +274,8 @@ pub fn init_cpu() {
 
             controlregs::write_cr4(cr4);
         }
+
+        assert!(features.has_xsave(), "init: xsave not supported!");
+        enable_xsave();
     }
 }

src/aero_kernel/src/arch/x86_64/task.rs

@@ -35,17 +35,20 @@
 use alloc::alloc::alloc_zeroed;
 
 use aero_syscall::{MMapFlags, MMapProt};
+use alloc::boxed::Box;
 use alloc::vec::Vec;
+use raw_cpuid::CpuId;
 
 use core::alloc::Layout;
 use core::ptr::Unique;
 
-use crate::arch::interrupts::InterruptErrorStack;
+use crate::arch::controlregs::MxCsr;
 use crate::fs::cache::DirCacheItem;
 use crate::mem::paging::*;
 use crate::syscall::ExecArgs;
 use crate::userland::vm::Vm;
 use crate::utils::StackHelper;
+use crate::{arch::interrupts::InterruptErrorStack, mem::alloc_boxed_buffer};
 
 use super::{controlregs, io};
 
@@ -121,6 +124,8 @@ pub struct ArchTask {
 
     fs_base: VirtAddr,
     gs_base: VirtAddr,
+
+    fpu_storage: Option<Box<[u8]>>,
 }
 
 impl ArchTask {
@@ -136,6 +141,8 @@ impl ArchTask {
 
             fs_base: VirtAddr::zero(),
             gs_base: VirtAddr::zero(),
+
+            fpu_storage: None,
         }
     }
 
@@ -178,6 +185,8 @@ impl ArchTask {
 
             fs_base: VirtAddr::zero(),
             gs_base: VirtAddr::zero(),
+
+            fpu_storage: None,
         }
     }
 
@@ -221,6 +230,9 @@ impl ArchTask {
         context.rip = fork_init as _;
         context.cr3 = address_space.cr3().start_address().as_u64();
 
+        let mut fpu_storage = alloc_boxed_buffer::<u8>(xsave_size() as usize);
+        fpu_storage.copy_from_slice(self.fpu_storage.as_ref().unwrap());
+
         Ok(Self {
             context: unsafe { Unique::new_unchecked(context) },
             context_switch_rsp: VirtAddr::new(switch_stack as u64),
@@ -230,6 +242,8 @@ impl ArchTask {
             // The FS and GS bases are inherited from the parent process.
             fs_base: self.fs_base.clone(),
             gs_base: self.gs_base.clone(),
+
+            fpu_storage: Some(fpu_storage),
         })
     }
 
@@ -272,6 +286,9 @@ impl ArchTask {
         context.rip = fork_init as u64;
         context.cr3 = new_address_space.cr3().start_address().as_u64();
 
+        let mut fpu_storage = alloc_boxed_buffer::<u8>(xsave_size() as usize);
+        fpu_storage.copy_from_slice(self.fpu_storage.as_ref().unwrap());
+
         Ok(Self {
             context: unsafe { Unique::new_unchecked(context) },
             context_switch_rsp: VirtAddr::new(switch_stack as u64),
@@ -281,6 +298,8 @@ impl ArchTask {
             // The FS and GS bases are inherited from the parent process.
             fs_base: self.fs_base.clone(),
             gs_base: self.gs_base.clone(),
+
+            fpu_storage: Some(fpu_storage),
         })
     }
 
@@ -324,6 +343,33 @@ impl ArchTask {
         self.fs_base = VirtAddr::zero();
         self.gs_base = VirtAddr::zero();
 
+        let mut fpu_storage = alloc_boxed_buffer::<u8>(xsave_size() as usize);
+
+        unsafe {
+            xrstor(&fpu_storage);
+
+            // The x87 FPU control word is set to 0x37f (default), which masks all
+            // floating-point exceptions, sets rounding to nearest, and sets the x87
+            // FPU precision to 64 bits (as documented in Intel SDM volume 1 section
+            // 8.1.5).
+            const DEFAULT_FPU_CWORD: u16 = 0x37f;
+            asm!("fldcw [{}]", in(reg) &DEFAULT_FPU_CWORD, options(nomem));
+
+            // Set the default MXCSR value at reset as documented in Intel SDM volume 2A.
+            controlregs::write_mxcsr(
+                MxCsr::INVALID_OPERATION_MASK
+                    | MxCsr::DENORMAL_MASK
+                    | MxCsr::DIVIDE_BY_ZERO_MASK
+                    | MxCsr::OVERFLOW_MASK
+                    | MxCsr::UNDERFLOW_MASK
+                    | MxCsr::PRECISION_MASK,
+            );
+
+            xsave(&mut fpu_storage);
+        }
+
+        self.fpu_storage = Some(fpu_storage);
+
         extern "C" {
             fn jump_userland_exec(stack: VirtAddr, rip: VirtAddr, rflags: u64);
         }
@@ -473,6 +519,28 @@ impl ArchTask {
     }
 }
 
+fn xsave_size() -> u32 {
+    static XSAVE_SIZE: Option<u32> = None;
+    XSAVE_SIZE.unwrap_or_else(|| {
+        CpuId::new()
+            .get_extended_state_info()
+            .expect("xsave: cpuid extended state info unavailable")
+            .xsave_size()
+    })
+}
+
+fn xsave(fpu: &mut Box<[u8]>) {
+    unsafe {
+        asm!("xsave [{}]", in(reg) fpu.as_ptr(), in("eax") 0xffffffffu32, in("edx") 0xffffffffu32)
+    }
+}
+
+fn xrstor(fpu: &Box<[u8]>) {
+    unsafe {
+        asm!("xrstor [{}]", in(reg) fpu.as_ptr(), in("eax") 0xffffffffu32, in("edx") 0xffffffffu32);
+    }
+}
+
 /// Check out the module level documentation for more information.
 pub fn arch_task_spinup(from: &mut ArchTask, to: &ArchTask) {
     extern "C" {
@@ -491,6 +559,14 @@ pub fn arch_task_spinup(from: &mut ArchTask, to: &ArchTask) {
         // update the swap GS target to point to the new GS base.
         io::wrmsr(io::IA32_KERNEL_GSBASE, to.gs_base.as_u64());
 
+        if let Some(fpu) = from.fpu_storage.as_mut() {
+            xsave(fpu);
+        }
+
+        if let Some(fpu) = to.fpu_storage.as_ref() {
+            xrstor(fpu);
+        }
+
         task_spinup(&mut from.context, to.context.as_ref());
     }
 }