Add a Cargo Feature for Enabling SSE

src/avx.s

@@ -0,0 +1,17 @@
+.section .boot, "awx"
+.intel_syntax noprefix
+.code64
+
+# This asm file enables AVX support before the OS starts.
+# AVX is not a requirement for an OS to boot.
+# This file should be loaded after stage 3 and just before stage 4.
+enable_avx:
+    push rax
+    push rcx
+    xor rcx, rcx
+    xgetbv
+    or eax, 7
+    xsetbv
+    pop rcx
+    pop rax
+    ret

src/main.rs

@@ -36,7 +36,9 @@ include!(concat!(env!("OUT_DIR"), "/bootloader_config.rs"));
 global_asm!(include_str!("stage_1.s"));
 global_asm!(include_str!("stage_2.s"));
 global_asm!(include_str!("e820.s"));
+global_asm!(include_str!("sse.s"));
 global_asm!(include_str!("stage_3.s"));
+global_asm!(include_str!("avx.s"));
 
 #[cfg(feature = "vga_320x200")]
 global_asm!(include_str!("video_mode/vga_320x200.s"));
@@ -87,6 +89,9 @@ extern "C" {
 
 #[no_mangle]
 pub unsafe extern "C" fn stage_4() -> ! {
+if is_x86_feature_detected!("avx") {
+enable_avx();
+}
     // Set stack segment
     asm!("mov bx, 0x0
           mov ss, bx" ::: "bx" : "intel");

src/sse.s

@@ -0,0 +1,53 @@
+.section .boot, "awx"
+.intel_syntax noprefix
+.code32
+
+# Checks for SSE support and enables it. Should be loaded after stage 2.
+no_sse_msg: .asciz "This system does not support SSE"
+no_sse2_msg: .asciz "This system does not support SSE 2"
+no_xsave_msg: .asciz "This system does not support XSAVE"
+
+# As a part of the implementation of x86-64, AMD demands a minimum amount of SSE support.
+# This function will fail if SSE, SSE2 and XSAVE support are not found together.
+enable_sse:
+    mov eax, 0x1
+    cpuid
+    test edx, 1 << 25
+    jz .no_sse
+    mov eax, 0x1
+    cpuid
+     test edx, 1 << 26
+    jz .no_sse2
+    mov eax, 0x1
+    cpuid
+     test ecx, 1 << 26
+    jz .no_xsave
+    mov eax, cr0
+    # clear coprocessor emulation CR0.EM
+    and ax, 0xFFFB
+    # set coprocessor monitoring  CR0.MP
+    or ax, 0x2
+    mov cr0, eax
+    mov eax, cr4
+    # set CR4.OSFXSR and CR4.OSXMMEXCPT at the same time
+    or ax, 3 << 9
+    mov cr4, eax
+    ret
+
+.no_sse:
+    lea si, [no_sse_msg]
+    call real_mode_println
+.no_sse_spin:
+jmp .no_sse_spin
+
+.no_sse2:
+    lea si, [no_sse2_msg]
+    call real_mode_println
+.no_sse2_spin:
+jmp .no_sse2_spin
+
+.no_xsave:
+    lea si, [no_xsave_msg]
+    call real_mode_println
+.no_xsave_spin:
+jmp .no_xsave_spin

src/stage_2.s

@@ -95,7 +95,7 @@ enter_protected_mode_again:
     mov eax, cr0
     or al, 1    # set protected mode bit
     mov cr0, eax
-
+    call enable_sse
     push 0x8
     lea eax, [stage_3]
     push eax