Enable AVX on boot if available

rafalmiel · Oct 4, 2024 · d6b7394 · d6b7394
1 parent a0bc2c9
commit d6b7394
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 1 deletion.
diff --git a/cykusz-rs/src/arch/x86_64/asm/boot.asm b/cykusz-rs/src/arch/x86_64/asm/boot.asm
@@ -13,6 +13,7 @@ extern setup_page_tables
 extern enable_paging
 extern __p4_table
 extern setup_SSE
+extern setup_AVX
 extern higher_half_start_ap
 
 section .apinit_trampoline
@@ -112,6 +113,7 @@ start:
     call setup_page_tables
     call enable_paging
     call setup_SSE
+    call setup_AVX
 
     lgdt [gdt64.pointer]
 

diff --git a/cykusz-rs/src/arch/x86_64/asm/sse.asm b/cykusz-rs/src/arch/x86_64/asm/sse.asm
@@ -1,6 +1,30 @@
 global setup_SSE
+global setup_AVX
 
 section .text
+setup_AVX:
+    push rax
+    push rcx
+    push rdx
+
+    ; check for AVX
+    mov rax, 0x1
+    cpuid
+    test ecx, 1<<28 ; check for AVX
+    jz .no_AVX
+    test ecx, 1<<26 ; check for XSAVE
+    jz .no_AVX
+
+    xor rcx, rcx
+    xgetbv      ; Load XCR0 register
+    or eax, 7   ; Set AVX, SSE, X87 bits
+    xsetbv      ; Save back to XCR0
+
+.no_AVX:
+    pop rdx
+    pop rcx
+    pop rax
+    ret
 ; Check for SSE and enable it. If it's not supported throw error "a".
 setup_SSE:
     ; check for SSE

diff --git a/cykusz-rs/src/arch/x86_64/task/mod.rs b/cykusz-rs/src/arch/x86_64/task/mod.rs
@@ -33,7 +33,7 @@ const KERN_STACK_ORDER: usize = 6;
 
 #[repr(align(16))]
 #[derive(Debug, Copy, Clone)]
-struct FpuState([u8; 512]);
+pub struct FpuState([u8; 512]);
 
 impl Default for FpuState {
     fn default() -> Self {