From 81313cef428de42709825f8452163f7b6d07d7ad Mon Sep 17 00:00:00 2001
From: Curle <curle@gemwire.uk>
Date: Fri, 25 Sep 2020 16:48:20 +0100
Subject: [PATCH] Refactor CPU preparation into cpu.c, work on enabling AVX and
 SSE

---
 chroma/inc/kernel/chroma.h |  3 +++
 chroma/kernel.c            |  6 ++---
 chroma/system/cpu.c        | 46 ++++++++++++++++++++++++++++++++++++++
 chroma/system/rw.c         |  2 +-
 4 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/chroma/inc/kernel/chroma.h b/chroma/inc/kernel/chroma.h
index 8a2f7b7..7ae00e8 100644
--- a/chroma/inc/kernel/chroma.h
+++ b/chroma/inc/kernel/chroma.h
@@ -51,6 +51,9 @@ size_t MemorySize;
 void DrawPixel(uint32_t x, uint32_t y, uint32_t color);
 void FillScreen(uint32_t color);
 
+void SetupExtensions();
+void PrepareCPU();
+
 void WriteString(const char* string);
 void WriteChar(const char character);
 
diff --git a/chroma/kernel.c b/chroma/kernel.c
index d60261e..0e5765b 100644
--- a/chroma/kernel.c
+++ b/chroma/kernel.c
@@ -34,15 +34,13 @@ int Main(void) {
 
     WriteStringWithFont("Initty Testing");
 
-    SetupInitialGDT();
-    SetupIDT();
-    InitInterrupts();
+    PrepareCPU();
 
     PCIEnumerate();
 
     InitMemoryManager();
 
-    DrawSplash();
+    //DrawSplash();
 
     InitPaging();
 
diff --git a/chroma/system/cpu.c b/chroma/system/cpu.c
index a66a924..dc02170 100644
--- a/chroma/system/cpu.c
+++ b/chroma/system/cpu.c
@@ -60,6 +60,52 @@ static void RefreshCS() {
 }
 
 
+void PrepareCPU() {
+
+    SetupInitialGDT();
+    SetupIDT();
+
+    SetupExtensions();
+
+    InitInterrupts();
+
+}
+
+void SetupExtensions() {
+
+    // Enable SSE
+    size_t CR0 = ReadControlRegister(0);
+
+    CR0 &= ~(1 << 2);
+    CR0 |= 1;
+
+    WriteControlRegister(0, CR0);
+
+    // Enable OSXSAVE and gang
+    size_t CR4 = ReadControlRegister(4);
+    CR4 |= (1 << 9);
+    CR4 |= (1 << 10);
+    CR4 |= (1 << 18);
+
+    WriteControlRegister(4, CR4);
+
+    // Enable AVX (and AVX-512 in future)
+    
+    CR0 = ReadExtendedControlRegister(0);
+    SerialPrintf("XCR0 is currently %x.\n", CR0);
+    CR0 |= (1 << 0);
+    CR0 |= (1 << 1);
+    CR0 |= (1 << 2);
+
+    /*CR0 |= (1 << 5);
+    CR0 |= (1 << 6);
+    CR0 |= (1 << 7);*/
+
+    SerialPrintf("About to write xcr0: %x\n", CR0);
+    WriteExtendedControlRegister(0, CR0);
+}
+
+
 void SetupInitialGDT() {
     DESC_TBL GDTData = {0};
     size_t TSSBase = (uint64_t) (&TSSEntries);
diff --git a/chroma/system/rw.c b/chroma/system/rw.c
index d1b2ff6..86c9402 100644
--- a/chroma/system/rw.c
+++ b/chroma/system/rw.c
@@ -189,7 +189,7 @@ size_t ReadExtendedControlRegister(size_t XCRX) {
 //TODO: make this just have an assert as returning data for a write to catch
 // errors that shoudlunt really happen doesent make alot of sense
 size_t WriteExtendedControlRegister(size_t XCRX, size_t Data){
-    uint32_t DataLow = Data & 0xffffffff;
+    uint32_t DataLow =   Data & 0x00000000ffffffff;
     uint32_t DataHigh = (Data & 0xffffffff00000000) >> 32;
     __asm__ __volatile__("xsetbv" : : "a" (DataLow), "c" (XCRX), "d" (DataHigh) :);
     return Data;