From 1ecd24f5d6f22c660e9900963d2a40ac7050f6d0 Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Sun, 2 Oct 2022 18:53:54 +0200
Subject: [PATCH] Kernel: Add SSE support (enable SSE on boot and save context
 (user tasks only) on task switch)

---
 apps/Makefile                   |  2 +-
 kernel/include/thread/Task.h    | 10 +++++++++-
 kernel/src/init/Init.cpp        |  4 ++++
 kernel/src/main.asm             | 13 ++++++++++++-
 kernel/src/thread/Scheduler.cpp |  8 ++++++--
 kernel/src/thread/Task.cpp      | 17 +++++++++++++++++
 libs/libc/Makefile              |  2 +-
 7 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/apps/Makefile b/apps/Makefile
index 61df706d..3a3d1bcc 100644
--- a/apps/Makefile
+++ b/apps/Makefile
@@ -6,7 +6,7 @@ APPS_BIN := $(APPS_DIR)/bin
 
 REAL_APPS := $(patsubst %, $(APPS_BIN)/%, $(APPS))
 
-CFLAGS := -Wall -Wextra -Werror -mgeneral-regs-only -Os
+CFLAGS := -Wall -Wextra -Werror -Os
 
 $(APPS_BIN)/%: $(APPS_SRC)/%.c
 	@mkdir -p $(@D)
diff --git a/kernel/include/thread/Task.h b/kernel/include/thread/Task.h
index 4fb0ccca..b197e46e 100644
--- a/kernel/include/thread/Task.h
+++ b/kernel/include/thread/Task.h
@@ -25,7 +25,15 @@ struct Task
     TaskState state;
 
     uint64_t cpu_time = 0;
+
+    char floating_region[512] __attribute__((aligned(16)));
+    bool floating_saved = false;
+
+    bool is_user_task();
 };
 
 void set_context_from_task(Task& task, Context* ctx);
-void get_context_to_task(Task& task, Context* ctx);
\ No newline at end of file
+void get_context_to_task(Task& task, Context* ctx);
+
+void task_save_floating(Task& task);
+void task_restore_floating(Task& task);
\ No newline at end of file
diff --git a/kernel/src/init/Init.cpp b/kernel/src/init/Init.cpp
index fe84bdb4..15085901 100644
--- a/kernel/src/init/Init.cpp
+++ b/kernel/src/init/Init.cpp
@@ -32,11 +32,15 @@ void Init::disable_smp()
     return;
 }
 
+extern "C" void asm_enable_sse();
+
 void Init::early_init()
 {
     Interrupts::disable();
     asm volatile("cld");
 
+    asm_enable_sse();
+
     framebuffer0.init((void*)bootboot.fb_ptr, bootboot.fb_type, bootboot.fb_scanline, bootboot.fb_width,
                       bootboot.fb_height);
 
diff --git a/kernel/src/main.asm b/kernel/src/main.asm
index 7c5c15f8..04e8bf2d 100644
--- a/kernel/src/main.asm
+++ b/kernel/src/main.asm
@@ -14,4 +14,15 @@ idle_task_function:
     sti
 .idle:
     hlt
-    jmp .idle
\ No newline at end of file
+    jmp .idle
+
+global asm_enable_sse
+asm_enable_sse:
+    mov rax, cr0
+    and ax, 0xFFFB
+    or ax, 0x2
+    mov cr0, rax
+    mov rax, cr4
+    or ax, 3 << 9
+    mov cr4, rax
+    ret
\ No newline at end of file
diff --git a/kernel/src/thread/Scheduler.cpp b/kernel/src/thread/Scheduler.cpp
index 59e6c127..b630e033 100644
--- a/kernel/src/thread/Scheduler.cpp
+++ b/kernel/src/thread/Scheduler.cpp
@@ -272,16 +272,20 @@ void Scheduler::task_yield(Context* context)
         sched_current_task = sched_current_task->next_task;
         if (sched_current_task->state == sched_current_task->Running)
         {
+            if (sched_current_task->id != original_task->id)
+            {
+                if (!was_idle && original_task->is_user_task()) { task_save_floating(*original_task); }
+                if (sched_current_task->is_user_task()) { task_restore_floating(*sched_current_task); }
+            }
             sched_current_task->task_time = 20;
             set_context_from_task(*sched_current_task, context);
-            Interrupts::enable();
             return;
         }
     } while (sched_current_task != original_task);
+    if (!was_idle && original_task->is_user_task()) task_save_floating(*original_task);
     sched_current_task = &idle_task;
     sched_current_task->task_time = frequency;
     if (!was_idle) { set_context_from_task(*sched_current_task, context); }
-    Interrupts::enable();
     return;
 }
 
diff --git a/kernel/src/thread/Task.cpp b/kernel/src/thread/Task.cpp
index 202b7de9..17226a60 100644
--- a/kernel/src/thread/Task.cpp
+++ b/kernel/src/thread/Task.cpp
@@ -9,4 +9,21 @@ void set_context_from_task(Task& task, Context* ctx)
 void get_context_to_task(Task& task, Context* ctx)
 {
     memcpy(&task.regs, ctx, sizeof(Context));
+}
+
+void task_save_floating(Task& task)
+{
+    task.floating_saved = true;
+    asm volatile("fxsave (%0)" : : "r"(&task.floating_region));
+}
+
+void task_restore_floating(Task& task)
+{
+    if (!task.floating_saved) return;
+    asm volatile("fxrstor (%0)" : : "r"(&task.floating_region));
+}
+
+bool Task::is_user_task()
+{
+    return (regs.cs & 3) > 0;
 }
\ No newline at end of file
diff --git a/libs/libc/Makefile b/libs/libc/Makefile
index 54a4067c..3d3a08ed 100644
--- a/libs/libc/Makefile
+++ b/libs/libc/Makefile
@@ -5,7 +5,7 @@ LIBC_BIN := $(LIBC_DIR)/bin
 
 DESTDIR ?= $(LUNA_BASE)/usr/lib
 
-CFLAGS := -Wall -Wextra -Werror -Os -nostdlib -fno-omit-frame-pointer -mno-mmx -mno-sse -mno-sse2
+CFLAGS := -Wall -Wextra -Werror -Os -nostdlib -fno-omit-frame-pointer
 CXXFLAGS := -fno-rtti -fno-exceptions
 ASMFLAGS := -felf64