Save and restore the additional XMM registers in OSR too.

lib/runtime-core/image-loading-linux-x86-64.s

@@ -41,6 +41,30 @@ add $8, %rsp
 movq (%rsp), %xmm7
 add $8, %rsp
 
+movq (%rsp), %xmm8
+add $8, %rsp
+
+movq (%rsp), %xmm9
+add $8, %rsp
+
+movq (%rsp), %xmm10
+add $8, %rsp
+
+movq (%rsp), %xmm11
+add $8, %rsp
+
+movq (%rsp), %xmm12
+add $8, %rsp
+
+movq (%rsp), %xmm13
+add $8, %rsp
+
+movq (%rsp), %xmm14
+add $8, %rsp
+
+movq (%rsp), %xmm15
+add $8, %rsp
+
 popq %rbp
 popq %rax
 popq %rbx

lib/runtime-core/image-loading-macos-x86-64.s

@@ -41,6 +41,30 @@ add $8, %rsp
 movq (%rsp), %xmm7
 add $8, %rsp
 
+movq (%rsp), %xmm8
+add $8, %rsp
+
+movq (%rsp), %xmm9
+add $8, %rsp
+
+movq (%rsp), %xmm10
+add $8, %rsp
+
+movq (%rsp), %xmm11
+add $8, %rsp
+
+movq (%rsp), %xmm12
+add $8, %rsp
+
+movq (%rsp), %xmm13
+add $8, %rsp
+
+movq (%rsp), %xmm14
+add $8, %rsp
+
+movq (%rsp), %xmm15
+add $8, %rsp
+
 popq %rbp
 popq %rax
 popq %rbx

lib/runtime-core/src/fault.rs

@@ -207,7 +207,7 @@ pub fn allocate_and_run<R, F: FnOnce() -> R>(size: usize, f: F) -> R {
 
         // NOTE: Keep this consistent with `image-loading-*.s`.
         stack[end_offset - 4 - 10] = &mut ctx as *mut Context<F, R> as usize as u64; // rdi
-        const NUM_SAVED_REGISTERS: usize = 23;
+        const NUM_SAVED_REGISTERS: usize = 31;
         let stack_begin = stack
             .as_mut_ptr()
             .offset((end_offset - 4 - NUM_SAVED_REGISTERS) as isize);
@@ -347,7 +347,7 @@ unsafe fn install_sighandler() {
 pub struct FaultInfo {
     pub faulting_addr: *const c_void,
     pub ip: *const c_void,
-    pub known_registers: [Option<u64>; 24],
+    pub known_registers: [Option<u64>; 32],
 }
 
 #[cfg(all(target_os = "linux", target_arch = "x86_64"))]
@@ -378,7 +378,7 @@ pub unsafe fn get_fault_info(siginfo: *const c_void, ucontext: *const c_void) ->
     let gregs = &(*ucontext).uc_mcontext.gregs;
     let fpregs = &*(*ucontext).uc_mcontext.fpregs;
 
-    let mut known_registers: [Option<u64>; 24] = [None; 24];
+    let mut known_registers: [Option<u64>; 32] = [None; 32];
     known_registers[X64Register::GPR(GPR::R15).to_index().0] = Some(gregs[REG_R15 as usize] as _);
     known_registers[X64Register::GPR(GPR::R14).to_index().0] = Some(gregs[REG_R14 as usize] as _);
     known_registers[X64Register::GPR(GPR::R13).to_index().0] = Some(gregs[REG_R13 as usize] as _);
@@ -405,6 +405,14 @@ pub unsafe fn get_fault_info(siginfo: *const c_void, ucontext: *const c_void) ->
     known_registers[X64Register::XMM(XMM::XMM5).to_index().0] = Some(read_xmm(&fpregs._xmm[5]));
     known_registers[X64Register::XMM(XMM::XMM6).to_index().0] = Some(read_xmm(&fpregs._xmm[6]));
     known_registers[X64Register::XMM(XMM::XMM7).to_index().0] = Some(read_xmm(&fpregs._xmm[7]));
+    known_registers[X64Register::XMM(XMM::XMM8).to_index().0] = Some(read_xmm(&fpregs._xmm[8]));
+    known_registers[X64Register::XMM(XMM::XMM9).to_index().0] = Some(read_xmm(&fpregs._xmm[9]));
+    known_registers[X64Register::XMM(XMM::XMM10).to_index().0] = Some(read_xmm(&fpregs._xmm[10]));
+    known_registers[X64Register::XMM(XMM::XMM11).to_index().0] = Some(read_xmm(&fpregs._xmm[11]));
+    known_registers[X64Register::XMM(XMM::XMM12).to_index().0] = Some(read_xmm(&fpregs._xmm[12]));
+    known_registers[X64Register::XMM(XMM::XMM13).to_index().0] = Some(read_xmm(&fpregs._xmm[13]));
+    known_registers[X64Register::XMM(XMM::XMM14).to_index().0] = Some(read_xmm(&fpregs._xmm[14]));
+    known_registers[X64Register::XMM(XMM::XMM15).to_index().0] = Some(read_xmm(&fpregs._xmm[15]));
 
     FaultInfo {
         faulting_addr: si_addr as usize as _,
@@ -458,8 +466,17 @@ pub unsafe fn get_fault_info(siginfo: *const c_void, ucontext: *const c_void) ->
     }
     #[repr(C)]
     struct fpstate {
-        _unused: [u8; 168],
+        _cwd: u16,
-        xmm: [[u64; 2]; 8],
+        _swd: u16,
+        _ftw: u16,
+        _fop: u16,
+        _rip: u64,
+        _rdp: u64,
+        _mxcsr: u32,
+        _mxcr_mask: u32,
+        _st: [[u16; 8]; 8],
+        xmm: [[u64; 2]; 16],
+        _padding: [u32; 24],
     }
     #[allow(dead_code)]
     #[repr(C)]
@@ -476,7 +493,7 @@ pub unsafe fn get_fault_info(siginfo: *const c_void, ucontext: *const c_void) ->
     let ss = &(*(*ucontext).uc_mcontext).ss;
     let fs = &(*(*ucontext).uc_mcontext).fs;
 
-    let mut known_registers: [Option<u64>; 24] = [None; 24];
+    let mut known_registers: [Option<u64>; 32] = [None; 32];
 
     known_registers[X64Register::GPR(GPR::R15).to_index().0] = Some(ss.r15);
     known_registers[X64Register::GPR(GPR::R14).to_index().0] = Some(ss.r14);
@@ -504,6 +521,14 @@ pub unsafe fn get_fault_info(siginfo: *const c_void, ucontext: *const c_void) ->
     known_registers[X64Register::XMM(XMM::XMM5).to_index().0] = Some(fs.xmm[5][0]);
     known_registers[X64Register::XMM(XMM::XMM6).to_index().0] = Some(fs.xmm[6][0]);
     known_registers[X64Register::XMM(XMM::XMM7).to_index().0] = Some(fs.xmm[7][0]);
+    known_registers[X64Register::XMM(XMM::XMM8).to_index().0] = Some(fs.xmm[8][0]);
+    known_registers[X64Register::XMM(XMM::XMM9).to_index().0] = Some(fs.xmm[9][0]);
+    known_registers[X64Register::XMM(XMM::XMM10).to_index().0] = Some(fs.xmm[10][0]);
+    known_registers[X64Register::XMM(XMM::XMM11).to_index().0] = Some(fs.xmm[11][0]);
+    known_registers[X64Register::XMM(XMM::XMM12).to_index().0] = Some(fs.xmm[12][0]);
+    known_registers[X64Register::XMM(XMM::XMM13).to_index().0] = Some(fs.xmm[13][0]);
+    known_registers[X64Register::XMM(XMM::XMM14).to_index().0] = Some(fs.xmm[14][0]);
+    known_registers[X64Register::XMM(XMM::XMM15).to_index().0] = Some(fs.xmm[15][0]);
 
     FaultInfo {
         faulting_addr: si_addr,

lib/runtime-core/src/state.rs

@@ -811,10 +811,10 @@ pub mod x64 {
     pub unsafe fn read_stack<'a, I: Iterator<Item = &'a CodeVersion>, F: Fn() -> I + 'a>(
         versions: F,
         mut stack: *const u64,
-        initially_known_registers: [Option<u64>; 24],
+        initially_known_registers: [Option<u64>; 32],
         mut initial_address: Option<u64>,
     ) -> ExecutionStateImage {
-        let mut known_registers: [Option<u64>; 24] = initially_known_registers;
+        let mut known_registers: [Option<u64>; 32] = initially_known_registers;
         let mut results: Vec<WasmFunctionStateDump> = vec![];
         let mut was_baseline = true;