diff --git a/src/hyperlight_guest_bin/src/arch/amd64/context.rs b/src/hyperlight_guest_bin/src/arch/amd64/context.rs index ec1c337af..ceb01788d 100644 --- a/src/hyperlight_guest_bin/src/arch/amd64/context.rs +++ b/src/hyperlight_guest_bin/src/arch/amd64/context.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // Copyright 2025 The Hyperlight Authors. +use core::arch::global_asm; use super::machine::ExceptionInfo; #[repr(C)] @@ -8,98 +9,131 @@ use super::machine::ExceptionInfo; 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::() == 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::() + size_of::()).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::() == 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::() + size_of::()).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", +); diff --git a/src/hyperlight_guest_bin/src/arch/amd64/exception/entry.rs b/src/hyperlight_guest_bin/src/arch/amd64/exception/entry.rs index b1914d648..3957b2eed 100644 --- a/src/hyperlight_guest_bin/src/arch/amd64/exception/entry.rs +++ b/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),