Enable hardware floats and save on every context switch

src/main.rs

@@ -86,7 +86,7 @@ mod start;
 mod tasking;
 mod virtual_memory;
 
-use core::{ptr, slice};
+use core::{arch::asm, ptr, slice};
 
 use bootinfo::BOOTINFO;
 use cast::usize;
@@ -105,7 +105,10 @@ use tar_no_std::TarArchiveRef;
 use tasking::{SleepReason, TASKING};
 use virtual_memory::{ACTIVE_SPACE, KERNEL_SPACE};
 use x86_64::{
-    registers::rflags::{self, RFlags},
+    registers::{
+        control::{Cr0, Cr0Flags, Cr4, Cr4Flags},
+        rflags::{self, RFlags},
+    },
     structures::paging::{Page, PageTableFlags},
     VirtAddr,
 };
@@ -129,6 +132,19 @@ pub fn main() {
     unsafe {
         rflags::write(rflags_data);
     }
+    // Enable the FPU
+    unsafe {
+        Cr0::update(|cr0| {
+            cr0.remove(Cr0Flags::EMULATE_COPROCESSOR);
+            cr0.insert(Cr0Flags::MONITOR_COPROCESSOR);
+        });
+        Cr4::update(|cr4| {
+            cr4.insert(Cr4Flags::OSFXSR);
+            cr4.insert(Cr4Flags::OSXMMEXCPT_ENABLE);
+        });
+        asm!("FNINIT");
+        tasking::store_initial_fpu_state();
+    }
     gdt::init();
     Lazy::force(&PHYSICAL_MEMORY);
     Lazy::force(&KERNEL_SPACE);

src/tasking.rs

@@ -1,19 +1,14 @@
 use crate::{
-    gdt,
-    interrupts::{send_ipc_to, REGISTERD_PIDS},
-    pit, println,
-    virtual_memory::{ASpaceMutex, AddressSpace, PagingError, ACTIVE_SPACE, KERNEL_SPACE},
+    gdt, interrupts::{send_ipc_to, REGISTERD_PIDS}, pit, println, virtual_memory::{ASpaceMutex, AddressSpace, PagingError, ACTIVE_SPACE, KERNEL_SPACE}
 };
 use alloc::{boxed::Box, collections::VecDeque, vec, vec::Vec};
 use cast::{u64, usize};
 use core::{
-    arch::naked_asm,
-    ptr::{addr_of, addr_of_mut},
-    sync::atomic::{AtomicBool, Ordering},
+    arch::{asm, naked_asm}, ptr::{addr_of, addr_of_mut}, sync::atomic::{AtomicBool, Ordering}
 };
 use crossbeam_queue::SegQueue;
 use slab::Slab;
-use spin::{Lazy, Mutex, RwLock};
+use spin::{Lazy, Mutex, Once, RwLock};
 use x86_64::{
     structures::paging::{Page, PageTableFlags},
     VirtAddr,
@@ -88,6 +83,36 @@ pub enum SleepReason {
     TimeSleep,
 }
 
+static INITIAL_FPU_STATE: Once<FpuState> = Once::new();
+
+pub fn store_initial_fpu_state() {
+    let mut state = FpuState::new();
+    state.save();
+    INITIAL_FPU_STATE.call_once(|| state);
+}
+
+#[repr(align(16))]
+#[derive(Debug, Clone, Copy)]
+struct FpuState {
+    state: [u8; 512]
+}
+
+impl FpuState {
+    fn new() -> Self {
+        Self {
+            state: [0; 512],
+        }
+    }
+
+    fn save(&mut self) {
+        unsafe { asm!("fxsave [rax]", in("rax") &mut self.state) };
+    }
+
+    fn load(&self) {
+        unsafe { asm!("fxrstor [rax]", in("rax") &self.state) };
+    }
+}
+
 #[derive(Debug)]
 pub struct Process {
     address_space: Option<AddressSpace>,
@@ -176,6 +201,7 @@ impl Process {
             kernel_stack,
             sleeping: RwLock::new(sleeping),
             tls_values: Mutex::new(vec![0; self.num_tls]),
+            fpu_state: Mutex::new(*INITIAL_FPU_STATE.get().unwrap()),
         };
         let idx = self.threads.write().insert(thread);
         self.num_threads += 1;
@@ -194,6 +220,7 @@ pub struct Thread {
     kernel_esp: *mut usize,
     kernel_esp_top: VirtAddr,
     tls_values: Mutex<Vec<u64>>,
+    fpu_state: Mutex<FpuState>,
 }
 
 pub struct IpcMessage {
@@ -351,9 +378,11 @@ impl Tasking {
                 if self.current_thread(|thread| thread.sleeping.read().is_none()) {
                     self.ready_to_run.lock().push_back((current_pid, self.current_tid().unwrap()));
                 }
+                self.current_thread(|thread| thread.fpu_state.lock().save());
                 let curr_stack = self.current_thread_mut(|thread| addr_of_mut!(thread.kernel_esp));
                 *self.current_pid.write() = Some(next_process_pid);
                 *self.current_tid.write() = Some(next_process_tid);
+                self.current_thread(|thread| thread.fpu_state.lock().load());
                 let kernel_esp = self.current_thread(|thread| {
                     gdt::set_tss_stack(thread.kernel_esp_top);
                     thread.kernel_esp
@@ -437,6 +466,7 @@ impl Tasking {
                             .expect("Non-current process has active page table")
                             .activate();
                     });
+                self.current_thread(|thread| thread.fpu_state.lock().load());
                 let kernel_esp = self.current_thread_mut(|thread| {
                     gdt::set_tss_stack(thread.kernel_esp_top);
                     thread.kernel_esp
@@ -493,6 +523,7 @@ impl Tasking {
                 }
                 *self.current_pid.write() = Some(next_process_pid);
                 *self.current_tid.write() = Some(next_process_tid);
+                self.current_thread(|thread| thread.fpu_state.lock().load());
                 let kernel_esp = self.current_thread_mut(|thread| {
                     gdt::set_tss_stack(thread.kernel_esp_top);
                     thread.kernel_esp