From a1e1e1c2d83665d089a314094ca8e7ff26dcd433 Mon Sep 17 00:00:00 2001
From: Drew Galbraith <dgalbraith33@gmail.com>
Date: Mon, 20 Nov 2023 16:40:07 -0800
Subject: [PATCH] [Zion] Enable SSE instructions at the start of boot.

These aren't ready to be used yet as we need to save them on task
switch.
---
 init-dbg.sh         | 11 +++++++++++
 zion/CMakeLists.txt |  1 +
 zion/common/cpu.cpp | 30 ++++++++++++++++++++++++++++++
 zion/common/cpu.h   |  3 +++
 zion/zion.cpp       |  3 +++
 5 files changed, 48 insertions(+)
 create mode 100644 zion/common/cpu.cpp
 create mode 100644 zion/common/cpu.h

diff --git a/init-dbg.sh b/init-dbg.sh
index cf10141..952efde 100755
--- a/init-dbg.sh
+++ b/init-dbg.sh
@@ -1,8 +1,19 @@
 #! /bin/bash
 
+set -e
+
 CWD="$(pwd)"
 BIN=$CWD/toolchain/local/bin
 GCC=$BIN/x86_64-pc-acadia-gcc
 AR=$BIN/x86_64-pc-acadia-ar
 SYSROOT=$CWD/sysroot
 cmake -B builddbg/ -G Ninja -D CMAKE_CXX_COMPILER=${GCC} -D CMAKE_ASM-ATT_COMPILER=${GCC} -D CMAKE_AR=${AR} -D CMAKE_BUILD_TYPE=Debug -D CMAKE_INSTALL_PREFIX=${SYSROOT}
+
+pushd yunq
+virtualenv venv
+source venv/bin/activate
+pip install -r requirements.txt
+deactivate
+popd
+
+echo "Set up environment properly."
diff --git a/zion/CMakeLists.txt b/zion/CMakeLists.txt
index 1e30851..fa2e22e 100644
--- a/zion/CMakeLists.txt
+++ b/zion/CMakeLists.txt
@@ -2,6 +2,7 @@ add_executable(zion
   boot/acpi.cpp
   boot/boot_info.cpp
   capability/capability_table.cpp
+  common/cpu.cpp
   common/gdt.cpp
   common/load_gdt.s
   common/msr.cpp
diff --git a/zion/common/cpu.cpp b/zion/common/cpu.cpp
new file mode 100644
index 0000000..16b4768
--- /dev/null
+++ b/zion/common/cpu.cpp
@@ -0,0 +1,30 @@
+#include "common/cpu.h"
+
+#include <cpuid.h>
+
+#include "debug/debug.h"
+
+void ProbeCpuAndEnableFeatures() {
+  dbgln("CPUID");
+  uint32_t eax, ebx, ecx, edx;
+  __get_cpuid(1, &eax, &ebx, &ecx, &edx);
+
+  if (!(edx & (0x3 << 25))) {
+    panic("SSE & SSE2 not available.");
+  }
+
+  if (!(ecx & (0x1 | (0x1 << 9)))) {
+    panic("SSE3, SSSE3 not available.");
+  }
+
+  dbgln("Setting SSE");
+  asm volatile(
+      "mov %%cr0, %%rax;"
+      "and $0xFFFB, %%ax;"  // Clear EM
+      "or $0x2, %%ax;"      // Set MP
+      "mov %%rax, %%cr0;"
+      "mov %%cr4, %%rax;"
+      "or $0x600, %%ax;"  // Set OSFXSR, OSXMMEXCPT
+      "mov %%rax, %%cr4;" ::
+          : "rax");
+}
diff --git a/zion/common/cpu.h b/zion/common/cpu.h
new file mode 100644
index 0000000..5df6f84
--- /dev/null
+++ b/zion/common/cpu.h
@@ -0,0 +1,3 @@
+#pragma once
+
+void ProbeCpuAndEnableFeatures();
diff --git a/zion/zion.cpp b/zion/zion.cpp
index c96599e..c5e04e4 100644
--- a/zion/zion.cpp
+++ b/zion/zion.cpp
@@ -1,6 +1,7 @@
 #include <stdint.h>
 
 #include "boot/acpi.h"
+#include "common/cpu.h"
 #include "common/gdt.h"
 #include "debug/debug.h"
 #include "interrupt/apic.h"
@@ -22,6 +23,8 @@ extern "C" void zion() {
   InitGdt();
   InitIdt();
 
+  ProbeCpuAndEnableFeatures();
+
   dbgln("[boot] Init Physical Memory Manager.");
   phys_mem::InitBootstrapPageAllocation();