Merge pull request #4089 from bylaws/wineavx

ARM64EC: AVX register save/restore support under wine
FEX-Emu · Sep 30, 2024 · e9d34c7 · e9d34c7
2 parents c2f2c58 + b659701
commit e9d34c7
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 2 deletions.
diff --git a/Source/Windows/ARM64EC/Module.cpp b/Source/Windows/ARM64EC/Module.cpp
@@ -291,7 +291,13 @@ static void LoadStateFromECContext(FEXCore::Core::InternalThreadState* Thread, C
 
   if ((Context.ContextFlags & CONTEXT_FLOATING_POINT) == CONTEXT_FLOATING_POINT) {
     // Floating-point register state
-    CTX->SetXMMRegistersFromState(Thread, reinterpret_cast<const __uint128_t*>(Context.FltSave.XmmRegisters), nullptr);
+    if ((Context.ContextFlags & CONTEXT_XSTATE) == CONTEXT_XSTATE) {
+      const auto* Ymm = RtlLocateExtendedFeature(reinterpret_cast<CONTEXT_EX*>(&Context + 1), XSTATE_AVX, nullptr);
+      CTX->SetXMMRegistersFromState(Thread, reinterpret_cast<const __uint128_t*>(Context.FltSave.XmmRegisters),
+                                    reinterpret_cast<const __uint128_t*>(Ymm));
+    } else {
+      CTX->SetXMMRegistersFromState(Thread, reinterpret_cast<const __uint128_t*>(Context.FltSave.XmmRegisters), nullptr);
+    }
     memcpy(State.mm, Context.FltSave.FloatRegisters, sizeof(State.mm));
 
     State.FCW = Context.FltSave.ControlWord;
@@ -330,6 +336,11 @@ static ARM64_NT_CONTEXT StoreStateToPackedECContext(FEXCore::Core::InternalThrea
   ARM64_NT_CONTEXT ECContext {};
 
   ECContext.ContextFlags = CONTEXT_ARM64_FULL;
+  if (CPUFeatures->IsFeaturePresent(PF_AVX2_INSTRUCTIONS_AVAILABLE)) {
+    // This is a FEX extension and requires corresponding wine-side patches to be of use, however it is harmless to set
+    // even if those patches are not used.
+    ECContext.ContextFlags |= CONTEXT_ARM64_FEX_YMMSTATE;
+  }
 
   auto& State = Thread->CurrentFrame->State;
 
@@ -352,7 +363,7 @@ static ARM64_NT_CONTEXT StoreStateToPackedECContext(FEXCore::Core::InternalThrea
 
   ECContext.Pc = State.rip;
 
-  CTX->ReconstructXMMRegisters(Thread, reinterpret_cast<__uint128_t*>(&ECContext.V[0]), nullptr);
+  CTX->ReconstructXMMRegisters(Thread, reinterpret_cast<__uint128_t*>(&ECContext.V[0]), reinterpret_cast<__uint128_t*>(&ECContext.V[16]));
 
   ECContext.Lr = State.mm[0][0];
   ECContext.X6 = State.mm[1][0];

diff --git a/Source/Windows/include/winnt.h b/Source/Windows/include/winnt.h
@@ -122,8 +122,25 @@ typedef struct _IMAGE_ARM64EC_CODE_RANGE_ENTRY_POINT {
   ULONG EntryPoint;
 } IMAGE_ARM64EC_CODE_RANGE_ENTRY_POINT;
 
+typedef struct _CONTEXT_CHUNK {
+  LONG Offset;
+  ULONG Length;
+} CONTEXT_CHUNK, *PCONTEXT_CHUNK;
+
+typedef struct _CONTEXT_EX {
+  CONTEXT_CHUNK All;
+  CONTEXT_CHUNK Legacy;
+  CONTEXT_CHUNK XState;
+#ifdef _WIN64
+  ULONG64 align;
+#endif
+} CONTEXT_EX, *PCONTEXT_EX;
+
 NTSYSAPI DWORD WINAPI RtlRunOnceExecuteOnce(PRTL_RUN_ONCE, PRTL_RUN_ONCE_INIT_FN, PVOID, PVOID*);
 
+// This is a FEX extension, and requires corresponding wine patches
+#define CONTEXT_ARM64_FEX_YMMSTATE (CONTEXT_ARM64 | 0x00000040)
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/Source/Windows/include/winternl.h b/Source/Windows/include/winternl.h
@@ -4,6 +4,7 @@
 #pragma once
 
 #include_next <winternl.h>
+#include <winnt.h>
 
 #ifdef __cplusplus
 extern "C" {
@@ -468,6 +469,7 @@ void WINAPI RtlInitializeConditionVariable(RTL_CONDITION_VARIABLE*);
 NTSTATUS WINAPI RtlInitializeCriticalSection(RTL_CRITICAL_SECTION*);
 void WINAPI RtlInitializeSRWLock(RTL_SRWLOCK*);
 NTSTATUS WINAPI RtlLeaveCriticalSection(RTL_CRITICAL_SECTION*);
+void* WINAPI RtlLocateExtendedFeature(CONTEXT_EX*, ULONG, ULONG*);
 NTSTATUS WINAPI RtlMultiByteToUnicodeN(LPWSTR, DWORD, LPDWORD, LPCSTR, DWORD);
 NTSTATUS WINAPI RtlMultiByteToUnicodeSize(DWORD*, LPCSTR, ULONG);
 BOOL WINAPI RtlQueryPerformanceCounter(LARGE_INTEGER*);