feat(kernel): lazy user FPU: off until first FP insn (illegal-insn trap -> zeroed regs, retry), save on dirty trap, restore on switch, off for threads without + kernel; spawntest per-thread FP check

crates/adapter-riscv/src/trap/fp.rs

@@ -0,0 +1,165 @@
+//! Floating point for user threads, lazily. A thread starts with the FPU
+//! off (`sstatus.FS`); its first floating-point instruction traps as an
+//! illegal instruction, and it is given zeroed registers and the FPU and
+//! retries. From then on its registers are saved when a trap finds them
+//! changed (FS dirty) and loaded whenever it is switched to; a thread
+//! without them runs with the FPU off, so it never sees another's. The
+//! kernel itself runs with the FPU off.
+
+use core::arch::asm;
+use domain::thread::{FpState, SavedContext};
+
+use super::TrapFrame;
+
+const FS: usize = 3 << 13;
+const FS_OFF: usize = 0;
+const FS_CLEAN: usize = 2 << 13;
+const FS_DIRTY: usize = 3 << 13;
+
+/// Turns the FPU off for the kernel; what the trap frame says the thread
+/// returns with is unchanged.
+pub(super) fn kernel_off() {
+    // SAFETY: clears sstatus.FS only.
+    unsafe { asm!("csrc sstatus, {0}", in(reg) FS) };
+}
+
+fn hardware_on() {
+    // SAFETY: sets sstatus.FS to (at least) clean, so the FP registers can
+    // be read and written; `kernel_off` turns it off again.
+    unsafe { asm!("csrs sstatus, {0}", in(reg) FS_CLEAN) };
+}
+
+/// On a trap from a thread: its registers into its context if it changed
+/// them, and the frame then says they are saved.
+pub(super) fn save_if_dirty(frame: &mut TrapFrame, ctx: &mut SavedContext) {
+    if frame.sstatus & FS != FS_DIRTY {
+        return;
+    }
+    if let Some(state) = ctx.fp_mut() {
+        hardware_on();
+        save(state);
+        frame.sstatus = (frame.sstatus & !FS) | FS_CLEAN;
+    }
+}
+
+/// Returning to a thread: its registers into the FPU, or the FPU off.
+pub(super) fn load_for(frame: &mut TrapFrame, ctx: &SavedContext) {
+    frame.sstatus &= !FS;
+    match ctx.fp() {
+        Some(state) => {
+            hardware_on();
+            load(state);
+            frame.sstatus |= FS_CLEAN;
+        }
+        None => frame.sstatus |= FS_OFF,
+    }
+}
+
+/// An illegal instruction from a thread with the FPU off and no registers
+/// yet: it gets them, zeroed, and retries. Whether that was it (otherwise
+/// the instruction is a real fault).
+pub(super) fn first_use(frame: &mut TrapFrame, ctx: &mut SavedContext) -> bool {
+    if frame.sstatus & FS != FS_OFF || ctx.fp().is_some() {
+        return false;
+    }
+    let state = ctx.enable_fp();
+    hardware_on();
+    load(state);
+    // sepc still names the instruction, so it runs again.
+    frame.sstatus = (frame.sstatus & !FS) | FS_CLEAN;
+    true
+}
+
+/// `$op` (fsd or fld) for every FP register, at `{base}` + 8 each.
+macro_rules! fp_registers {
+    ($op:literal) => {
+        concat!(
+            $op,
+            " f0, 0({base})\n",
+            $op,
+            " f1, 8({base})\n",
+            $op,
+            " f2, 16({base})\n",
+            $op,
+            " f3, 24({base})\n",
+            $op,
+            " f4, 32({base})\n",
+            $op,
+            " f5, 40({base})\n",
+            $op,
+            " f6, 48({base})\n",
+            $op,
+            " f7, 56({base})\n",
+            $op,
+            " f8, 64({base})\n",
+            $op,
+            " f9, 72({base})\n",
+            $op,
+            " f10, 80({base})\n",
+            $op,
+            " f11, 88({base})\n",
+            $op,
+            " f12, 96({base})\n",
+            $op,
+            " f13, 104({base})\n",
+            $op,
+            " f14, 112({base})\n",
+            $op,
+            " f15, 120({base})\n",
+            $op,
+            " f16, 128({base})\n",
+            $op,
+            " f17, 136({base})\n",
+            $op,
+            " f18, 144({base})\n",
+            $op,
+            " f19, 152({base})\n",
+            $op,
+            " f20, 160({base})\n",
+            $op,
+            " f21, 168({base})\n",
+            $op,
+            " f22, 176({base})\n",
+            $op,
+            " f23, 184({base})\n",
+            $op,
+            " f24, 192({base})\n",
+            $op,
+            " f25, 200({base})\n",
+            $op,
+            " f26, 208({base})\n",
+            $op,
+            " f27, 216({base})\n",
+            $op,
+            " f28, 224({base})\n",
+            $op,
+            " f29, 232({base})\n",
+            $op,
+            " f30, 240({base})\n",
+            $op,
+            " f31, 248({base})\n",
+        )
+    };
+}
+
+/// The FPU's registers into `state` (the FPU must be on).
+fn save(state: &mut FpState) {
+    let base = state.registers.as_mut_ptr();
+    let control: u64;
+    // SAFETY: `base` points at 32 u64s, written in place; frcsr only reads.
+    unsafe {
+        asm!(fp_registers!("fsd"), base = in(reg) base, options(nostack));
+        asm!("frcsr {0}", out(reg) control, options(nostack));
+    }
+    state.control = control;
+}
+
+/// `state` into the FPU's registers (the FPU must be on).
+fn load(state: &FpState) {
+    let base = state.registers.as_ptr();
+    // SAFETY: `base` points at 32 u64s, only read; fscsr writes fcsr.
+    unsafe {
+        asm!(fp_registers!("fld"), base = in(reg) base, options(nostack, readonly));
+        asm!("fscsr {0}", in(reg) state.control, options(nostack));
+    }
+}

crates/adapter-riscv/src/trap/mod.rs

@@ -1,4 +1,5 @@
 mod fast;
+mod fp;
 mod sched;
 mod vector;
 
@@ -17,6 +18,7 @@ const CAUSE_INTERRUPT_BIT: usize = 1 << 63;
 const CAUSE_S_SOFTWARE: usize = 1;
 const CAUSE_S_TIMER: usize = 5;
 const CAUSE_S_EXTERNAL: usize = 9;
+const CAUSE_ILLEGAL_INSTRUCTION: usize = 2;
 const CAUSE_ECALL_U: usize = 8;
 /// Set in `sstatus` when the trap came from supervisor mode.
 const SSTATUS_SPP: usize = 1 << 8;
@@ -29,6 +31,8 @@ pub fn init() {
         let addr = trap_vector as *const () as usize;
         core::arch::asm!("csrw stvec, {0}", in(reg) addr);
         core::arch::asm!("csrw sscratch, zero");
+        // The kernel never uses the FPU; user threads get it lazily.
+        core::arch::asm!("csrc sstatus, {0}", in(reg) 3usize << 13);
         // SIE stays clear; interrupts are enabled only by sret into userspace.
         let sie_bits = (1usize << 5) | (1usize << 9) | (1usize << 1);
         core::arch::asm!("csrs sie, {0}", in(reg) sie_bits);
@@ -90,6 +94,14 @@ extern "C" fn trap_handler_rust(frame: &mut TrapFrame) {
     }
     crate::smp::service_ipi();
     crate::smp::lock();
+    // The thread's FP registers are saved before anything can switch away.
+    if let Some(ctx) = kern()
+        .current_thread()
+        .and_then(|t| kern().thread_context_mut(t))
+    {
+        fp::save_if_dirty(frame, ctx);
+    }
+    fp::kernel_off();
     if scause & CAUSE_INTERRUPT_BIT != 0 {
         handle_interrupt(frame, scause & !CAUSE_INTERRUPT_BIT);
     } else {
@@ -163,6 +175,11 @@ fn handle_exception(frame: &mut TrapFrame, scause: usize) {
                 }
             }
         }
+        CAUSE_ILLEGAL_INSTRUCTION
+            if kern()
+                .current_thread()
+                .and_then(|t| kern().thread_context_mut(t))
+                .is_some_and(|ctx| fp::first_use(frame, ctx)) => {}
         _ => {
             let uart = crate::hw::uart::Uart::new();
             uart.write_str("FAULT: scause=");

crates/adapter-riscv/src/trap/sched.rs

@@ -14,6 +14,7 @@ fn load_context_to_frame(ctx: &domain::thread::SavedContext, frame: &mut TrapFra
     frame.sstatus = ctx.status();
     frame.sstatus &= !(1 << 8); // SPP: return to U-mode
     frame.sstatus |= 1 << 5; // SPIE
+    super::fp::load_for(frame, ctx);
 }
 
 fn apply_pending_result(frame: &mut TrapFrame, tid: types::id::ThreadId) {

crates/apps/spawntest/src/main.rs

@@ -173,6 +173,57 @@ fn encrypted_home(fs: &mut Storage, home_fs: &mut Storage) -> Result<(), &'stati
     Ok(())
 }
 
+/// Floating-point work that only comes out right if nobody else's
+/// registers get in: a sum over `ROUNDS`, yielding (so switching threads)
+/// all the way.
+fn float_work(step: f64) -> u64 {
+    const ROUNDS: u32 = 2000;
+    let mut sum = 0.0f64;
+    let mut x = core::hint::black_box(step);
+    for _ in 0..ROUNDS {
+        sum += x * 0.5;
+        x += step;
+        syscall::yield_now();
+    }
+    sum.to_bits()
+}
+
+/// What `float_work(step)` gives when nothing interferes.
+fn float_expected(step: f64) -> u64 {
+    let mut sum = 0.0f64;
+    let mut x = step;
+    for _ in 0..2000 {
+        sum += x * 0.5;
+        x += step;
+    }
+    sum.to_bits()
+}
+
+/// The other thread's answer, once it has one.
+static FLOAT_RESULT: core::sync::atomic::AtomicU64 = core::sync::atomic::AtomicU64::new(0);
+
+fn float_thread(_: &'static core::sync::atomic::AtomicU64) -> ! {
+    FLOAT_RESULT.store(float_work(0.25), core::sync::atomic::Ordering::Release);
+    syscall::exit_thread()
+}
+
+/// Two threads computing with floats at once each get their own registers
+/// (the kernel saves and restores them on every switch).
+fn floating_point_is_kept_per_thread() -> Result<(), &'static str> {
+    if !libkos::thread::spawn(&FLOAT_RESULT, float_thread) {
+        return Err("cannot start a thread");
+    }
+    let mine = float_work(1.75);
+    while FLOAT_RESULT.load(core::sync::atomic::Ordering::Acquire) == 0 {
+        syscall::yield_now();
+    }
+    let theirs = FLOAT_RESULT.load(core::sync::atomic::Ordering::Acquire);
+    if mine != float_expected(1.75) || theirs != float_expected(0.25) {
+        return Err("floating-point registers leaked between threads");
+    }
+    Ok(())
+}
+
 /// How many requests the `flaky` answering now has served.
 fn flaky() -> Result<usize, &'static str> {
     syscall::request(FLAKY, [flaky::op::WHO, 0, 0, 0]).map(|r| r[0]).map_err(|_| "flaky does not answer")
@@ -246,6 +297,9 @@ fn check() -> Result<(), &'static str> {
     revoke_stops_a_running_app(died)?;
     libkos::println!("spawntest: revoking stops a running app, allowing lets it run: ok");
 
+    floating_point_is_kept_per_thread()?;
+    libkos::println!("spawntest: floating point is kept per thread: ok");
+
     restarted_after_crashes()?;
     libkos::println!("spawntest: a crashed server is restarted, a crash loop backs off: ok");
 

crates/domain/src/thread/context.rs

@@ -1,9 +1,27 @@
 //! Saved register state of a thread.
 
+/// A thread's floating-point registers and their control word. A thread
+/// has none until its first floating-point instruction (see
+/// `SavedContext::enable_fp`): until then it cannot see another thread's.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct FpState {
+    pub registers: [u64; 32],
+    pub control: u64,
+}
+
+impl FpState {
+    pub const ZERO: Self = Self {
+        registers: [0; 32],
+        control: 0,
+    };
+}
+
 pub struct SavedContext {
     registers: [usize; 32],
     pc: usize,
     status: usize,
+    fp: FpState,
+    fp_enabled: bool,
 }
 
 impl SavedContext {
@@ -12,6 +30,8 @@ impl SavedContext {
             registers: [0; 32],
             pc: 0,
             status: 0,
+            fp: FpState::ZERO,
+            fp_enabled: false,
         }
     }
 
@@ -48,4 +68,39 @@ impl SavedContext {
     pub fn registers_mut(&mut self) -> &mut [usize; 32] {
         &mut self.registers
     }
+
+    /// Its floating-point registers, if it has used any.
+    pub fn fp(&self) -> Option<&FpState> {
+        self.fp_enabled.then_some(&self.fp)
+    }
+
+    pub fn fp_mut(&mut self) -> Option<&mut FpState> {
+        self.fp_enabled.then_some(&mut self.fp)
+    }
+
+    /// Gives it floating-point registers, all zero.
+    pub fn enable_fp(&mut self) -> &mut FpState {
+        self.fp = FpState::ZERO;
+        self.fp_enabled = true;
+        &mut self.fp
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn a_thread_has_no_fp_until_it_is_given_zeroed_registers() {
+        let mut c = SavedContext::empty();
+        assert!(c.fp().is_none() && c.fp_mut().is_none());
+        c.enable_fp().registers[3] = 7;
+        assert_eq!(c.fp().unwrap().registers[3], 7);
+        c.enable_fp();
+        assert_eq!(
+            c.fp(),
+            Some(&FpState::ZERO),
+            "enabling again starts from zero"
+        );
+    }
 }

crates/domain/src/thread/mod.rs

@@ -9,7 +9,7 @@ use crate::error::KernelError;
 mod context;
 mod table;
 
-pub use context::SavedContext;
+pub use context::{FpState, SavedContext};
 pub use table::{MAX_THREADS, ThreadTable};
 
 #[cfg(test)]