From e5556e16a7029f583c47688d4e7f4db020e7425a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Thomas=20Lov=C3=A9n?= <thomasloven@gmail.com>
Date: Sat, 22 Jan 2022 22:26:03 +0100
Subject: [PATCH] Enable SSE

---
 src/kernel/cpu/cpu.c        |  1 +
 src/kernel/cpu/sse.S        | 30 ++++++++++++++++++++++++++++++
 src/kernel/include/cpu.h    |  7 ++++++-
 src/kernel/include/proc.h   |  1 +
 src/kernel/proc/process.c   |  5 +++++
 src/kernel/proc/scheduler.c |  3 +++
 6 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100644 src/kernel/cpu/sse.S

diff --git a/src/kernel/cpu/cpu.c b/src/kernel/cpu/cpu.c
index b603230..fc51435 100644
--- a/src/kernel/cpu/cpu.c
+++ b/src/kernel/cpu/cpu.c
@@ -27,4 +27,5 @@ void cpu_init() {
   apic_init();
   ioapic_init();
   timer_init();
+  sse_init();
 }
\ No newline at end of file
diff --git a/src/kernel/cpu/sse.S b/src/kernel/cpu/sse.S
new file mode 100644
index 0000000..80efa5d
--- /dev/null
+++ b/src/kernel/cpu/sse.S
@@ -0,0 +1,30 @@
+.intel_syntax noprefix
+
+.global sse_init
+sse_init:
+  //; Enable bits 9 and 10 in cr4
+  //; 9: Enable SSE instructions
+  //; 10: Enable SSE exceptions
+  mov rax, cr4
+  or rax, 1<<9 | 1<<10
+  mov cr4, rax
+
+  //; Set bit 1 and unset bit 2 in cr0
+  //; 1: Coprocessor monitoring
+  //; 2: Coprocessor emulation
+  mov rax, cr0
+  or rax, 1<<1
+  and rax, ~(1<<2)
+  mov cr0, rax
+
+  ret
+
+.global sse_save
+sse_save:
+  fxsave [rdi]
+  ret
+
+.global sse_restore
+sse_restore:
+  fxrstor [rdi]
+  ret
diff --git a/src/kernel/include/cpu.h b/src/kernel/include/cpu.h
index 546ed62..9862e2c 100644
--- a/src/kernel/include/cpu.h
+++ b/src/kernel/include/cpu.h
@@ -98,4 +98,9 @@ void irq_unmask(int irq);
 void ioapic_init();
 
 // cpu/timer.c
-void timer_init();
\ No newline at end of file
+void timer_init();
+
+// cpu/sse.S
+void sse_init();
+void sse_save(void *addr);
+void sse_restore(void *addr);
\ No newline at end of file
diff --git a/src/kernel/include/proc.h b/src/kernel/include/proc.h
index 63c1892..73b7351 100644
--- a/src/kernel/include/proc.h
+++ b/src/kernel/include/proc.h
@@ -7,6 +7,7 @@ struct process {
   uint64_t state;
   uint64_t P4;
   struct process *q_next;
+  void *sse;
   uint8_t stack[];
 };
 
diff --git a/src/kernel/proc/process.c b/src/kernel/proc/process.c
index b064b82..cf5f62c 100644
--- a/src/kernel/proc/process.c
+++ b/src/kernel/proc/process.c
@@ -1,6 +1,7 @@
 #include <stdint.h>
 #include <proc.h>
 #include <memory.h>
+#include <stdlib.h>
 
 struct swtch_stack {
   uint64_t RBP;
@@ -24,6 +25,10 @@ struct process *new_process(void (*function)(void)) {
   p->q_next = 0;
   p->P4 = new_P4();
 
+  // Storage for SSE registers must be 16 byte aligned
+  void *sse = malloc(512+15);
+  p->sse = (void *) (((uintptr_t)sse + 15) & ~0xF);
+
   struct swtch_stack *stck = p->stack_ptr;
   stck->RBP = (uint64_t)&stck->RBP2;
   stck->ret = (uint64_t)function;
diff --git a/src/kernel/proc/scheduler.c b/src/kernel/proc/scheduler.c
index 785fc87..57533c6 100644
--- a/src/kernel/proc/scheduler.c
+++ b/src/kernel/proc/scheduler.c
@@ -32,8 +32,11 @@ void scheduler() {
 
     current_proc = new;
     write_cr3(new->P4);
+    sse_restore(current_proc->sse);
+
     switch_stack(&scheduler_proc->stack_ptr, &new->stack_ptr);
 
+    sse_save(current_proc->sse);
     scheduler_insert(current_proc);
     current_proc = 0;
   }