Save floating point unit registers

src/boot/traps.S

@@ -4,10 +4,10 @@
 .global trap_entry
 
 trap_entry:
-    # Create space on the stack for 32 registers (32 * 8 = 256 bytes)
-    addi sp, sp, -256
+    # Allocate stack space for GPRs (256 bytes) + FPU registers (512 bytes)
+    addi sp, sp, -768  # Total: 256 (GPRs) + 512 (FPU)
 
-    # Save all General Purpose Registers (GPRs)
+    # --- Save GPRs ---
     sd ra, 0(sp)
     sd gp, 8(sp)
     sd tp, 16(sp)
@@ -39,19 +39,108 @@ trap_entry:
     sd t5, 224(sp)
     sd t6, 232(sp)
 
-    # Save mepc to the stack (using offset 240)
+    # --- Save mepc ---
     csrr t0, mepc
     sd t0, 240(sp)
 
-    # Call trap handler
-    mv a0, sp # put the stack pointer to the first arg
+    # --- Check FPU Status (mstatus.FS) ---
+    csrr t1, mstatus
+    li t2, 0x6000
+    and t1, t1, t2
+    bnez t1, save_fpu
+
+    # If no FPU, skip to handler
+    j call_handler
+
+save_fpu:
+    # Save FPU registers (f0-f31) starting at offset 256
+    fsd f0, 256(sp)
+    fsd f1, 264(sp)
+    fsd f2, 272(sp)
+    fsd f3, 280(sp)
+    fsd f4, 288(sp)
+    fsd f5, 296(sp)
+    fsd f6, 304(sp)
+    fsd f7, 312(sp)
+    fsd f8, 320(sp)
+    fsd f9, 328(sp)
+    fsd f10, 336(sp)
+    fsd f11, 344(sp)
+    fsd f12, 352(sp)
+    fsd f13, 360(sp)
+    fsd f14, 368(sp)
+    fsd f15, 376(sp)
+    fsd f16, 384(sp)
+    fsd f17, 392(sp)
+    fsd f18, 400(sp)
+    fsd f19, 408(sp)
+    fsd f20, 416(sp)
+    fsd f21, 424(sp)
+    fsd f22, 432(sp)
+    fsd f23, 440(sp)
+    fsd f24, 448(sp)
+    fsd f25, 456(sp)
+    fsd f26, 464(sp)
+    fsd f27, 472(sp)
+    fsd f28, 480(sp)
+    fsd f29, 488(sp)
+    fsd f30, 496(sp)
+    fsd f31, 504(sp)
+
+call_handler:
+    # Pass stack pointer to C handler
+    mv a0, sp
     call handle_trap
 
-    # restore mepc
+    # --- Restore mepc ---
     ld t0, 240(sp)
     csrw mepc, t0
 
-    # Restore all GPRs
+    # --- Restore FPU (if it was saved) ---
+    csrr t1, mstatus
+    li t2, 0x6000
+    and t1, t1, t2
+    bnez t1, restore_fpu
+
+    j restore_gpr
+
+restore_fpu:
+    # Restore FPU registers
+    fld f0, 256(sp)
+    fld f1, 264(sp)
+    fld f2, 272(sp)
+    fld f3, 280(sp)
+    fld f4, 288(sp)
+    fld f5, 296(sp)
+    fld f6, 304(sp)
+    fld f7, 312(sp)
+    fld f8, 320(sp)
+    fld f9, 328(sp)
+    fld f10, 336(sp)
+    fld f11, 344(sp)
+    fld f12, 352(sp)
+    fld f13, 360(sp)
+    fld f14, 368(sp)
+    fld f15, 376(sp)
+    fld f16, 384(sp)
+    fld f17, 392(sp)
+    fld f18, 400(sp)
+    fld f19, 408(sp)
+    fld f20, 416(sp)
+    fld f21, 424(sp)
+    fld f22, 432(sp)
+    fld f23, 440(sp)
+    fld f24, 448(sp)
+    fld f25, 456(sp)
+    fld f26, 464(sp)
+    fld f27, 472(sp)
+    fld f28, 480(sp)
+    fld f29, 488(sp)
+    fld f30, 496(sp)
+    fld f31, 504(sp)
+
+restore_gpr:
+    # Restore GPRs
     ld ra, 0(sp)
     ld gp, 8(sp)
     ld tp, 16(sp)
@@ -83,8 +172,6 @@ trap_entry:
     ld t5, 224(sp)
     ld t6, 232(sp)
 
-    # Shrink the stack back
-    addi sp, sp, 256
-
-    # Return from Machine-mode trap
+    # Restore stack pointer
+    addi sp, sp, 768
     mret