Merge pull request #5660 from Sonicadvance1/183

Wow64: Spin loop on atomic with WFE

FEXCore/include/FEXCore/Utils/SpinWaitLock.h

@@ -144,6 +144,21 @@ static inline void WaitPred(T* Futex, T ComparisonValue) {
   }
 }
 
+template<typename T, typename Pred>
+static inline void WaitBitMaskPred(T* Futex, T BitMask, T ComparisonValue, Pred Predicate) {
+  auto AtomicFutex = std::atomic_ref<T>(*Futex);
+  T Result = AtomicFutex.load();
+
+  while (!Predicate(Result & BitMask, ComparisonValue)) {
+    Result = LoadExclusive(Futex);
+    if (Predicate(Result & BitMask, ComparisonValue)) {
+      return;
+    }
+
+    Result = WFELoadAtomic(Futex);
+  }
+}
+
 template<typename T, typename TT>
 static inline bool Wait(T* Futex, TT ExpectedValue, const std::chrono::nanoseconds& Timeout) {
   auto AtomicFutex = std::atomic_ref<T>(*Futex);
@@ -213,6 +228,16 @@ static inline void WaitPred(T* Futex, T ComparisonValue) {
   }
 }
 
+template<typename T, typename Pred>
+static inline void WaitBitMaskPred(T* Futex, T BitMask, T ComparisonValue, Pred Predicate) {
+  auto AtomicFutex = std::atomic_ref<T>(*Futex);
+  T Result = AtomicFutex.load();
+
+  while (!Predicate(Result & BitMask, ComparisonValue)) {
+    Result = AtomicFutex.load();
+  }
+}
+
 template<typename T, typename TT>
 static inline bool Wait(T* Futex, TT ExpectedValue, const std::chrono::nanoseconds& Timeout) {
   auto AtomicFutex = std::atomic_ref<T>(*Futex);

Source/Windows/WOW64/Module.cpp

@@ -95,6 +95,10 @@ struct TLS {
     return reinterpret_cast<std::atomic<uint32_t>&>(TEB->TlsSlots[FEXCore::ToUnderlying(Slot::CONTROL_WORD)]);
   }
 
+  uint32_t* ControlWordAddress() const {
+    return reinterpret_cast<uint32_t*>(&TEB->TlsSlots[FEXCore::ToUnderlying(Slot::CONTROL_WORD)]);
+  }
+
   CONTEXT*& EntryContext() const {
     return reinterpret_cast<CONTEXT*&>(TEB->TlsSlots[FEXCore::ToUnderlying(Slot::ENTRY_CONTEXT)]);
   }
@@ -822,8 +826,7 @@ NTSTATUS BTCpuSuspendLocalThread(HANDLE Thread, ULONG* Count) {
   }
 
   // Spin until the JIT is interrupted
-  while (TLS.ControlWord().load() & ControlBits::IN_JIT)
-    ;
+  FEXCore::Utils::SpinWaitLock::WaitBitMaskPred(TLS.ControlWordAddress(), ControlBits::IN_JIT, 0U, std::equal_to<>());
 
   // The JIT has now been interrupted and the context stored in the thread's CPU area is up-to-date
   if (Err = NtSuspendThread(*ThreadDup, Count); Err) {