Fix/xsave context save

src/hyperlight_guest_bin/src/arch/amd64/context.rs

@@ -1,105 +1,139 @@
 // SPDX-License-Identifier: Apache-2.0
 // Copyright 2025 The Hyperlight Authors.
 
+use core::arch::global_asm;
 use super::machine::ExceptionInfo;
 
 #[repr(C)]
 /// Saved context, pushed onto the stack by exception entry code
 pub struct Context {
     /// in order: ds, gs, fs, es
     pub segments: [u64; 4],
-    pub fxsave: [u8; 512],
+
+    /// Extended CPU state (xsave or fxsave area)
+    /// Size varies by CPU:
+    ///   512  bytes — fxsave (no AVX)
+    ///   576  bytes — xsave  (AVX)
+    ///   2688 bytes — xsave  (AVX-512)
+    /// Always allocated at maximum to keep struct layout fixed.
+    pub extended_state: [u8; 2688],
+
+    /// Actual bytes written by xsave/fxsave — set by save_context
+    pub extended_size: u64,
+
     /// no `rsp`, since the processor saved it
     /// `rax` is at the top, `r15` the bottom
     pub gprs: [u64; 15],
-    _padding: u64,
-}
-const _: () = assert!(size_of::<Context>() == 32 + 512 + 120 + 8);
-// The combination of the ExceptionInfo (pushed by the CPU) and the
-// register Context that we save to the stack must be 16byte aligned
-// before calling the hl_exception_handler as specified in the x86-64
-// ELF System V psABI specification, Section 3.2.2:
-//
-// https://gitlab.com/x86-psABIs/x86-64-ABI/-/jobs/artifacts/master/raw/x86-64-ABI/abi.pdf?job=build
-const _: () = assert!((size_of::<Context>() + size_of::<ExceptionInfo>()).is_multiple_of(16));
-
-// Defines `context_save` and `context_restore`
-macro_rules! save {
-    () => {
-        concat!(
-            // Save general-purpose registers
-            "    sub rsp, 8\n",
-            "    push rax\n",
-            "    push rbx\n",
-            "    push rcx\n",
-            "    push rdx\n",
-            "    push rsi\n",
-            "    push rdi\n",
-            "    push rbp\n",
-            "    push r8\n",
-            "    push r9\n",
-            "    push r10\n",
-            "    push r11\n",
-            "    push r12\n",
-            "    push r13\n",
-            "    push r14\n",
-            "    push r15\n",
-            // Save floating-point/SSE registers
-            // TODO: Don't do this unconditionally: get the exn
-            //       handlers compiled without sse
-            // TODO: Check if we ever generate code with ymm/zmm in
-            //       the handlers and save/restore those as well
-            "    sub rsp, 512\n",
-            "    mov rax, rsp\n",
-            "    fxsave [rax]\n",
-            // Save the rest of the segment registers
-            "    mov rax, es\n",
-            "    push rax\n",
-            "    mov rax, fs\n",
-            "    push rax\n",
-            "    mov rax, gs\n",
-            "    push rax\n",
-            "    mov rax, ds\n",
-            "    push rax\n",
-        )
-    };
-}
-pub(super) use save;
-
-macro_rules! restore {
-    () => {
-        concat!(
-            // Restore most segment registers
-            "    pop rax\n",
-            "    mov ds, rax\n",
-            "    pop rax\n",
-            "    mov gs, rax\n",
-            "    pop rax\n",
-            "    mov fs, rax\n",
-            "    pop rax\n",
-            "    mov es, rax\n",
-            // Restore floating-point/SSE registers
-            "    mov rax, rsp\n",
-            "    fxrstor [rax]\n",
-            "    add rsp, 512\n",
-            // Restore general-purpose registers
-            "    pop r15\n",
-            "    pop r14\n",
-            "    pop r13\n",
-            "    pop r12\n",
-            "    pop r11\n",
-            "    pop r10\n",
-            "    pop r9\n",
-            "    pop r8\n",
-            "    pop rbp\n",
-            "    pop rdi\n",
-            "    pop rsi\n",
-            "    pop rdx\n",
-            "    pop rcx\n",
-            "    pop rbx\n",
-            "    pop rax\n",
-            "    add rsp, 8\n",
-        )
-    };
+
+    _padding: [u64; 2],
 }
-pub(super) use restore;
+
+const _: () = assert!(size_of::<Context>() == 32 + 2688 + 8 + 120 + 16);
+
+// The combination of ExceptionInfo and Context must be 16-byte aligned
+// before calling hl_exception_handler as per x86-64 System V ABI.
+const _: () = assert!(
+    (size_of::<Context>() + size_of::<ExceptionInfo>()).is_multiple_of(16)
+);
+
+global_asm!(
+    ".global save_context",
+    "save_context:",
+    "    sub rsp, 8",
+    "    push rax",
+    "    push rbx",
+    "    push rcx",
+    "    push rdx",
+    "    push rsi",
+    "    push rdi",
+    "    push rbp",
+    "    push r8",
+    "    push r9",
+    "    push r10",
+    "    push r11",
+    "    push r12",
+    "    push r13",
+    "    push r14",
+    "    push r15",
+
+    // CPUID — get xsave area size
+    "    mov eax, 0xD",
+    "    xor ecx, ecx",
+    "    cpuid",
+    // ebx = required xsave size
+    "    sub rsp, rbx",
+    "    and rsp, -64",
+    "    push rbx",       // save size for restore
+
+    // check AVX support
+    "    mov eax, 1",
+    "    cpuid",
+    "    bt ecx, 28",
+    "    jnc use_fxsave",
+
+    "use_xsave:",
+    "    mov eax, 0x7",
+    "    xor edx, edx",
+    "    xsave [rsp]",
+    "    jmp save_done",
+
+    "use_fxsave:",
+    "    fxsave [rsp]",
+
+    "save_done:",
+    "    mov rax, es",
+    "    push rax",
+    "    mov rax, fs",
+    "    push rax",
+    "    mov rax, gs",
+    "    push rax",
+    "    mov rax, ds",
+    "    push rax",
+    "    ret",
+
+    ".global restore_context",
+    "restore_context:",
+    "    pop rax",
+    "    mov ds, rax",
+    "    pop rax",
+    "    mov gs, rax",
+    "    pop rax",
+    "    mov fs, rax",
+    "    pop rax",
+    "    mov es, rax",
+
+    "    pop rbx",        // restore saved size
+    "    mov eax, 1",
+    "    cpuid",
+    "    bt ecx, 28",
+    "    jnc use_fxrstor",
+
+    "use_xrstor:",
+    "    mov eax, 0x7",
+    "    xor edx, edx",
+    "    xrstor [rsp]",
+    "    jmp restore_done",
+
+    "use_fxrstor:",
+    "    fxrstor [rsp]",
+
+    "restore_done:",
+    "    add rsp, rbx",
+    "    pop r15",
+    "    pop r14",
+    "    pop r13",
+    "    pop r12",
+    "    pop r11",
+    "    pop r10",
+    "    pop r9",
+    "    pop r8",
+    "    pop rbp",
+    "    pop rdi",
+    "    pop rsi",
+    "    pop rdx",
+    "    pop rcx",
+    "    pop rbx",
+    "    pop rax",
+    "    add rsp, 8",
+    "    ret",
+);

src/hyperlight_guest_bin/src/arch/amd64/exception/entry.rs

@@ -8,7 +8,6 @@ use core::arch::{asm, global_asm};
 
 use hyperlight_common::outb::Exception;
 
-use super::super::context;
 use super::super::machine::{IDT, IdtEntry, IdtPointer, ProcCtrl};
 
 unsafe extern "C" {
@@ -58,8 +57,7 @@ macro_rules! generate_excp {
             "_do_excp",
             stringify!($num),
             ":\n",
-            context::save!(),
-            // rsi is the exception number.
+            "    call save_context\n",
             "    mov rsi, ",
             stringify!($num),
             "\n",
@@ -81,7 +79,7 @@ macro_rules! generate_excp {
             // For the ones that don't, we push a 0 to keep the
             // stack aligned.
             "   push 0\n",
-            context::save!(),
+            "    call save_context\n",
             // rsi is the exception number.
             "    mov rsi, ",
             stringify!($num),
@@ -100,13 +98,13 @@ macro_rules! generate_excp {
             "_do_excp",
             stringify!($num),
             ":\n",
-            context::save!(),
+            "    mov rdx, cr2\n",
+            "    call save_context\n",
             "    mov rsi, ",
             stringify!($num),
             "\n",
             // In a page fault exception, the cr2 register
             // contains the address that caused the page fault.
-            "    mov rdx, cr2\n",
             "    jmp _do_excp_common\n"
         )
     };
@@ -124,7 +122,7 @@ macro_rules! generate_exceptions {
             // stack pointer just before it was called.
             "    mov rdi, rsp\n",
             "    call {hl_exception_handler}\n",
-            context::restore!(),
+            "    call restore_context\n",
             "    add rsp, 8\n", // error code
             "    iretq\n",      // iretq is used to return from exception in x86_64
             generate_excp!(0, pusherrcode),