Fix/xsave context save

[JM00NJ]

Problem

The previous implementation used fxsave/fxrstor unconditionally,
which only preserves x87 and SSE state (fixed 512 bytes). Any AVX
(ymm) or AVX-512 (zmm) registers in use at the time of an exception
are silently lost — the guest resumes with corrupted FPU state.

Why global_asm! instead of Rust

The core issue is that the required save area size is not known at
compile time — it depends on the CPU's XSAVE capabilities, determined
at runtime via CPUID leaf 0xD (EBX).

Rust's type system requires struct fields to have a fixed size at
compile time. This means we cannot dynamically size the save area
inside the Context struct using safe Rust — any attempt to do so
would require heap allocation or dynamic dispatch, neither of which
is appropriate in a bare-metal exception handler that runs before
the allocator is available.

global_asm! gives us direct control over the stack layout:

• CPUID leaf 0xD is called to determine the exact save area size
• sub rsp, rbx opens exactly the right amount of stack space
• and rsp, -64 aligns to the 64-byte boundary required by xsave
• The size is pushed onto the stack for restore to use later

This way the save area lives on the stack with the correct size for
the current CPU, without touching the heap or requiring a fixed-size
struct field.

Changes

• context.rs: replaced fxsave/fxrstor macro approach with
  save_context/restore_context functions in global_asm!.
  CPUID selects xsave (AVX) or fxsave (no AVX) at runtime.
  Context.extended_state sized to 2688 bytes (AVX-512 maximum)
  to keep the struct layout fixed for the Rust handler.

• entry.rs: updated exception entry stubs to call
  save_context/restore_context instead of the old inline macros.

Fixes TODO in context.rs

// 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

Why entry.rs was updated

The old entry.rs used context::save!() and context::restore!()
inline macros, which expanded directly into the exception handler
assembly string via concat!(). These macros were tightly coupled
to the fixed fxsave layout in context.rs.

With save_context/restore_context now living as standalone
functions in global_asm!, the exception stubs simply call them:

call save_context   ; instead of context::save!() inline expansion
call restore_context

This also moves the page fault handler to read cr2 before calling
save_context, since CPUID inside save_context does not modify
cr2 but a nested fault could — reading it first is the safe
approach.

The use super::super::context; import was removed as it is no
longer needed.

[JM00NJ]

Verified the generated assembly with objdump on Linux
(x86_64-unknown-linux-gnu, debug build):

0000000000000000 <save_context>:
   0:	48 83 ec 08          	sub    $0x8,%rsp
   4:	50                   	push   %rax
   5:	53                   	push   %rbx
   6:	51                   	push   %rcx
   7:	52                   	push   %rdx
   8:	56                   	push   %rsi
   9:	57                   	push   %rdi
   a:	55                   	push   %rbp
   b:	41 50                	push   %r8
   d:	41 51                	push   %r9
   f:	41 52                	push   %r10
  11:	41 53                	push   %r11
  13:	41 54                	push   %r12
  15:	41 55                	push   %r13
  17:	41 56                	push   %r14
  19:	41 57                	push   %r15
  1b:	b8 0d 00 00 00       	mov    $0xd,%eax
  20:	31 c9                	xor    %ecx,%ecx
  22:	0f a2                	cpuid
  24:	48 29 dc             	sub    %rbx,%rsp
  27:	48 83 e4 c0          	and    $0xffffffffffffffc0,%rsp
  2b:	53                   	push   %rbx
  2c:	b8 01 00 00 00       	mov    $0x1,%eax
  31:	0f a2                	cpuid
  33:	0f ba e1 1c          	bt     $0x1c,%ecx
  37:	73 0d                	jae    46 <use_fxsave>

0000000000000039 <use_xsave>:
  39:	b8 07 00 00 00       	mov    $0x7,%eax
  3e:	31 d2                	xor    %edx,%edx
  40:	0f ae 24 24          	xsave  (%rsp)
  44:	eb 04                	jmp    4a <save_done>

0000000000000046 <use_fxsave>:
  46:	0f ae 04 24          	fxsave (%rsp)

000000000000004a <save_done>:
  4a:	48 8c c0             	mov    %es,%rax
  4d:	50                   	push   %rax
  4e:	48 8c e0             	mov    %fs,%rax
  51:	50                   	push   %rax
  52:	48 8c e8             	mov    %gs,%rax
  55:	50                   	push   %rax
  56:	48 8c d8             	mov    %ds,%rax
  59:	50                   	push   %rax
  5a:	c3                   	ret

000000000000005b <restore_context>:
  5b:	58                   	pop    %rax
  5c:	48 8e d8             	mov    %rax,%ds
  5f:	58                   	pop    %rax
  60:	48 8e e8             	mov    %rax,%gs
  63:	58                   	pop    %rax
  64:	48 8e e0             	mov    %rax,%fs
  67:	58                   	pop    %rax
  68:	48 8e c0             	mov    %rax,%es
  6b:	5b                   	pop    %rbx
  6c:	b8 01 00 00 00       	mov    $0x1,%eax
  71:	0f a2                	cpuid
  73:	0f ba e1 1c          	bt     $0x1c,%ecx
  77:	73 0d                	jae    86 <use_fxrstor>

0000000000000079 <use_xrstor>:
  79:	b8 07 00 00 00       	mov    $0x7,%eax
  7e:	31 d2                	xor    %edx,%edx
  80:	0f ae 2c 24          	xrstor (%rsp)
  84:	eb 04                	jmp    8a <restore_done>

0000000000000086 <use_fxrstor>:
  86:	0f ae 0c 24          	fxrstor (%rsp)

000000000000008a <restore_done>:
  8a:	48 01 dc             	add    %rbx,%rsp
  8d:	41 5f                	pop    %r15
  8f:	41 5e                	pop    %r14
  91:	41 5d                	pop    %r13
  93:	41 5c                	pop    %r12
  95:	41 5b                	pop    %r11
  97:	41 5a                	pop    %r10
  99:	41 59                	pop    %r9

CPUID leaf 0xD selects the xsave area size, AVX bit (ecx bit 28)
selects xsave vs fxsave, and the size is pushed/popped for restore.

[jsturtevant]

We don't currently compile code for the guest with AVX enabled. It's something we've discussed enabling but it's been lower on the priority list. Are you seeing this manifest as a bug or have a use case?

There would also be some performance concerns during context switching and saving off all the extra registers #711

[JM00NJ]

This PR aims to fix the silent loss of the upper 128 bits
of ymm registers when AVX is in use. It's written entirely
in global_asm! so there's no Rust overhead - and on AVX
CPUs xsave uses lazy saving, which can actually be faster
than an unconditional fxsave in practice.

If the direction isn't clear yet, I'm happy to leave this
open or close it as a draft — but I believe it's a net
improvement over the current code either way.

[jsturtevant]

This PR aims to fix the silent loss of the upper 128 bits of ymm registers when AVX is in use.

We don't use AVX, so i am not sure we get any benifits? I agree this would be useful if we did.

As per performance It's not rust vs ASM concerned about but bytes we are having to save. The benchmarks might be able to give us a sense if this is something to be concerned about in practice. @syntactically might have some thoughts too

[ludfjig]

hi @JM00NJ would you mind explaining the issue you're trying to fix? not only are guests complied with avx disabled, but in addition avx is disabled in the guest vm. The guest can enable it if they wishes, but avx is part of xsave so it will be cleared on restore()

[JM00NJ]

Thanks for the clarification.I wasn't aware that AVX is
also disabled at the VM level and cleared on restore; that
changes things significantly.

The motivation was to guard against silent ymm corruption
if AVX were ever enabled, but if the restore path already
clears AVX state, this PR doesn't add much value in the
current setup.

Happy to close this if the team decides it's not needed
at this stage. If AVX support moves up the priority list
in the future, this could serve as a starting point.