From 8b1bccb79bd58c70ad5e7cf4a8fafb99b8df59b7 Mon Sep 17 00:00:00 2001
From: Bananymous <bananymousosq@gmail.com>
Date: Mon, 31 Jul 2023 22:28:57 +0300
Subject: [PATCH] Kernel: store/load sse/mmx/fpu state on isr/irq/syscall

I haven't tested this yet but should be fine. This will be optimized
to only save state from threads that are using it
---
 kernel/arch/x86_64/IDT.cpp     | 8 ++++++++
 kernel/include/kernel/Thread.h | 5 +++++
 kernel/kernel/Syscall.cpp      | 4 ++++
 3 files changed, 17 insertions(+)

diff --git a/kernel/arch/x86_64/IDT.cpp b/kernel/arch/x86_64/IDT.cpp
index b0a4eeb7cf..0d9199f6e6 100644
--- a/kernel/arch/x86_64/IDT.cpp
+++ b/kernel/arch/x86_64/IDT.cpp
@@ -138,6 +138,8 @@ namespace IDT
 
 	extern "C" void cpp_isr_handler(uint64_t isr, uint64_t error, Kernel::InterruptStack& interrupt_stack, const Registers* regs)
 	{
+		Kernel::Thread::current().save_sse();
+
 		pid_t tid = Kernel::Scheduler::current_tid();
 		pid_t pid = tid ? Kernel::Process::current().pid() : 0;
 
@@ -205,10 +207,14 @@ namespace IDT
 		default:
 			break;
 		}
+
+		Kernel::Thread::current().load_sse();
 	}
 
 	extern "C" void cpp_irq_handler(uint64_t irq, Kernel::InterruptStack& interrupt_stack)
 	{
+		Kernel::Thread::current().save_sse();
+
 		if (Kernel::Scheduler::current_tid())
 		{
 			Kernel::Thread::current().set_return_rsp(interrupt_stack.rsp);
@@ -242,6 +248,8 @@ namespace IDT
 		default:
 			break;
 		}
+
+		Kernel::Thread::current().load_sse();
 	}
 
 	static void flush_idt()
diff --git a/kernel/include/kernel/Thread.h b/kernel/include/kernel/Thread.h
index 715bd66a28..1de1b3eac9 100644
--- a/kernel/include/kernel/Thread.h
+++ b/kernel/include/kernel/Thread.h
@@ -82,6 +82,9 @@ namespace Kernel
 
 		bool is_userspace() const { return m_is_userspace; }
 
+		void save_sse() { asm volatile("fxsave %0" :: "m"(m_sse_storage)); }
+		void load_sse() { asm volatile("fxrstor %0" :: "m"(m_sse_storage)); }
+
 	private:
 		Thread(pid_t tid, Process*);
 		void on_exit();
@@ -111,6 +114,8 @@ namespace Kernel
 
 		uint64_t m_terminate_blockers { 0 };
 
+		alignas(16) uint8_t m_sse_storage[512];
+
 		friend class TerminateBlocker;
 		friend class Scheduler;
 	};
diff --git a/kernel/kernel/Syscall.cpp b/kernel/kernel/Syscall.cpp
index 2b8afff98e..f83f2fcfc3 100644
--- a/kernel/kernel/Syscall.cpp
+++ b/kernel/kernel/Syscall.cpp
@@ -30,6 +30,8 @@ namespace Kernel
 			return 0;
 		}
 
+		Thread::current().save_sse();
+
 		asm volatile("sti");
 
 		(void)arg1;
@@ -183,6 +185,8 @@ namespace Kernel
 			break;
 		}
 
+		Thread::current().load_sse();
+
 		if (ret.is_error())
 			return -ret.error().get_error_code();
 		return ret.value();