From 620f9ee71a7aa908db8a0734fd310538309bbe93 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Sat, 23 Jul 2022 13:31:48 +0200
Subject: [PATCH] [NTOS] Add MxCsr handling to trap handlers

---
 ntoskrnl/ke/amd64/trap.S           | 7 +++++++
 sdk/include/asm/ksamd64.template.h | 2 +-
 sdk/include/asm/trapamd64.inc      | 9 +++++++++
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/ntoskrnl/ke/amd64/trap.S b/ntoskrnl/ke/amd64/trap.S
index 9d694518d0f..63dacf8166b 100644
--- a/ntoskrnl/ke/amd64/trap.S
+++ b/ntoskrnl/ke/amd64/trap.S
@@ -822,6 +822,10 @@ PUBLIC KiSystemCallEntry64
     mov ds, ax
     mov es, ax
 
+    /* Save MCXSR and set kernel value */
+    stmxcsr [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_MxCsr]
+    ldmxcsr gs:[PcMxCsr]
+
 #if DBG
     /* Check IRQL */
     mov rax, cr8
@@ -861,6 +865,9 @@ no_user_apc_pending:
     /* Disable interrupts for return */
     cli
 
+    /* Restore MCXSR */
+    ldmxcsr [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_MxCsr]
+
     /* Restore old trap frame */
     mov rcx, gs:[PcCurrentThread]
     mov rdx, [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_TrapFrame]
diff --git a/sdk/include/asm/ksamd64.template.h b/sdk/include/asm/ksamd64.template.h
index 1072df28b34..6fcb43149ea 100644
--- a/sdk/include/asm/ksamd64.template.h
+++ b/sdk/include/asm/ksamd64.template.h
@@ -586,7 +586,7 @@ OFFSET(PcIrql, KPCR, Irql),
 OFFSET(PcStallScaleFactor, KPCR, StallScaleFactor),
 OFFSET(PcHalReserved, KPCR, HalReserved),
 //OFFSET(PcPrcb, KPCR, Prcb),
-//OFFSET(PcMxCsr, KPCR, MxCsr),
+OFFSET(PcMxCsr, KIPCR, Prcb.MxCsr),
 //OFFSET(PcNumber, KPCR, Number),
 //OFFSET(PcInterruptRequest, KPCR, InterruptRequest),
 //OFFSET(PcIdleHalt, KPCR, IdleHalt),
diff --git a/sdk/include/asm/trapamd64.inc b/sdk/include/asm/trapamd64.inc
index eeb804ec970..d5d51dc50a6 100644
--- a/sdk/include/asm/trapamd64.inc
+++ b/sdk/include/asm/trapamd64.inc
@@ -97,6 +97,9 @@ MACRO(EnterTrap, Flags)
         mov [rbp + KTRAP_FRAME_SegGs], gs
     endif
 
+    /* Save MCXSR */
+    stmxcsr [rbp + KTRAP_FRAME_MxCsr]
+
 #if DBG
     mov ecx, MSR_GS_BASE
     rdmsr
@@ -116,6 +119,9 @@ MACRO(EnterTrap, Flags)
     mov es, ax
     swapgs
 
+    /* Load kernel MXCSR */
+    ldmxcsr gs:[PcMxCsr]
+
 #if DBG
     /* Check IRQL */
     mov rax, cr8
@@ -255,6 +261,9 @@ kernel_mode_return:
         movdqa xmm5, [rbp + KTRAP_FRAME_Xmm5]
     endif
 
+    /* Restore MCXSR */
+    ldmxcsr [rbp + KTRAP_FRAME_MxCsr]
+
     /* Restore rbp */
     mov rbp, [rbp + KTRAP_FRAME_Rbp]