preemptive_threads/
context_full.rs

1/// Full CPU context including FPU/SSE/AVX state
2#[repr(C, align(64))]
3pub struct FullThreadContext {
4    /// General purpose registers
5    pub rsp: u64,
6    pub rbp: u64,
7    pub rbx: u64,
8    pub r12: u64,
9    pub r13: u64,
10    pub r14: u64,
11    pub r15: u64,
12    pub rflags: u64,
13
14    /// FPU/SSE/AVX state (512 bytes for XSAVE area)
15    pub xsave_area: [u8; 512],
16
17    /// MXCSR register for SSE state
18    pub mxcsr: u32,
19    pub mxcsr_mask: u32,
20}
21
22impl Default for FullThreadContext {
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28impl FullThreadContext {
29    pub const fn new() -> Self {
30        Self {
31            rsp: 0,
32            rbp: 0,
33            rbx: 0,
34            r12: 0,
35            r13: 0,
36            r14: 0,
37            r15: 0,
38            rflags: 0x202, // Default RFLAGS with interrupts enabled
39            xsave_area: [0; 512],
40            mxcsr: 0x1F80, // Default MXCSR
41            mxcsr_mask: 0xFFFF,
42        }
43    }
44}
45
46/// Check CPU features at runtime
47pub fn check_cpu_features() -> CpuFeatures {
48    let mut features = CpuFeatures::default();
49
50    unsafe {
51        // Check for XSAVE support
52        let result = core::arch::x86_64::__cpuid_count(1, 0);
53        features.xsave = (result.ecx & (1 << 26)) != 0;
54        features.avx = (result.ecx & (1 << 28)) != 0;
55        features.fma = (result.ecx & (1 << 12)) != 0;
56
57        // Check for AVX2
58        let result = core::arch::x86_64::__cpuid_count(7, 0);
59        features.avx2 = (result.ebx & (1 << 5)) != 0;
60
61        // Check for AVX-512
62        features.avx512f = (result.ebx & (1 << 16)) != 0;
63
64        // Get XSAVE feature mask
65        if features.xsave {
66            let result = core::arch::x86_64::__cpuid_count(0xD, 0);
67            features.xsave_mask = ((result.edx as u64) << 32) | (result.eax as u64);
68        }
69    }
70
71    features
72}
73
74#[derive(Default)]
75pub struct CpuFeatures {
76    pub xsave: bool,
77    pub avx: bool,
78    pub avx2: bool,
79    pub avx512f: bool,
80    pub fma: bool,
81    pub xsave_mask: u64,
82}
83
84pub static mut CPU_FEATURES: CpuFeatures = CpuFeatures {
85    xsave: false,
86    avx: false,
87    avx2: false,
88    avx512f: false,
89    fma: false,
90    xsave_mask: 0,
91};
92
93/// Initialize CPU feature detection
94pub fn init_cpu_features() {
95    unsafe {
96        CPU_FEATURES = check_cpu_features();
97    }
98}
99
100/// Full context switch with all CPU state
101/// 
102/// # Safety
103/// Caller must ensure valid thread context pointers and proper synchronization.
104#[cfg(target_arch = "x86_64")]
105#[unsafe(naked)]
106#[no_mangle]
107pub unsafe extern "C" fn switch_context_full(
108    from: *mut FullThreadContext,
109    to: *const FullThreadContext,
110) {
111    core::arch::naked_asm!(
112        "
113        # Save callee-saved registers
114        push rbp
115        push rbx
116        push r12
117        push r13
118        push r14
119        push r15
120        
121        # Save RFLAGS
122        pushfq
123        pop rax
124        mov [rdi + 56], rax    # Save RFLAGS to context
125        
126        # Save stack pointer
127        mov [rdi + 0], rsp
128        mov [rdi + 8], rbp
129        mov [rdi + 16], rbx
130        mov [rdi + 24], r12
131        mov [rdi + 32], r13
132        mov [rdi + 40], r14
133        mov [rdi + 48], r15
134        
135        # For now, just use FXSAVE/FXRSTOR which is always available
136        # Save FPU/SSE state
137        fxsave [rdi + 64]
138        
139        # Restore FPU/SSE state
140        fxrstor [rsi + 64]
141        
142        # Restore general registers
143        mov rsp, [rsi + 0]
144        mov rbp, [rsi + 8]
145        mov rbx, [rsi + 16]
146        mov r12, [rsi + 24]
147        mov r13, [rsi + 32]
148        mov r14, [rsi + 40]
149        mov r15, [rsi + 48]
150        
151        # Restore RFLAGS
152        mov rax, [rsi + 56]
153        push rax
154        popfq
155        
156        # Restore callee-saved registers from stack
157        pop r15
158        pop r14
159        pop r13
160        pop r12
161        pop rbx
162        pop rbp
163        
164        ret
165        "
166    );
167}
168
169/// Simple context switch (enhanced version with proper RFLAGS handling)
170/// 
171/// # Safety
172/// Caller must ensure valid thread context pointers and proper synchronization.
173#[cfg(target_arch = "x86_64")]
174#[unsafe(naked)]
175#[no_mangle]
176pub unsafe extern "C" fn switch_context_simple(
177    from: *mut crate::thread::ThreadContext,
178    to: *const crate::thread::ThreadContext,
179) {
180    core::arch::naked_asm!(
181        "
182        # Save callee-saved registers in correct order
183        push rbp
184        push rbx
185        push r12
186        push r13
187        push r14
188        push r15
189        
190        # Save RFLAGS
191        pushfq
192        
193        # Save current context
194        mov [rdi + 0], rsp
195        mov [rdi + 8], rbp
196        mov [rdi + 16], rbx
197        mov [rdi + 24], r12
198        mov [rdi + 32], r13
199        mov [rdi + 40], r14
200        mov [rdi + 48], r15
201        
202        # Switch to new context
203        mov rsp, [rsi + 0]
204        mov rbp, [rsi + 8]
205        mov rbx, [rsi + 16]
206        mov r12, [rsi + 24]
207        mov r13, [rsi + 32]
208        mov r14, [rsi + 40]
209        mov r15, [rsi + 48]
210        
211        # Restore RFLAGS
212        popfq
213        
214        # Restore callee-saved registers in reverse order
215        pop r15
216        pop r14
217        pop r13
218        pop r12
219        pop rbx
220        pop rbp
221        
222        ret
223        "
224    );
225}
preemptive_threads/context_full.rs

preemptive_threads/
context_full.rs
