ARM64EC: Emulate unaligned atomic access in non-JIT EC code

FEXCore/Source/Utils/ArchHelpers/Arm64.cpp

@@ -1952,8 +1952,8 @@ static uint64_t HandleAtomicLoadstoreExclusive(uintptr_t ProgramCounter, uint64_
 }
 
 [[nodiscard]]
-std::optional<int32_t>
-HandleUnalignedAccess(FEXCore::Core::InternalThreadState* Thread, UnalignedHandlerType HandleType, uintptr_t ProgramCounter, uint64_t* GPRs) {
+std::optional<int32_t> HandleUnalignedAccess(FEXCore::Core::InternalThreadState* Thread, UnalignedHandlerType HandleType,
+                                             uintptr_t ProgramCounter, uint64_t* GPRs, bool IsJIT) {
 #ifdef _M_ARM_64
   constexpr bool is_arm64 = true;
 #else
@@ -1978,7 +1978,7 @@ HandleUnalignedAccess(FEXCore::Core::InternalThreadState* Thread, UnalignedHandl
   uint32_t* StrictSplitLockMutex {CTX->Config.StrictInProcessSplitLocks ? &CTX->StrictSplitLockMutex : nullptr};
 
   // ParanoidTSO path doesn't modify any code.
-  if (HandleType == UnalignedHandlerType::Paranoid) [[unlikely]] {
+  if (HandleType == UnalignedHandlerType::Paranoid || !IsJIT) [[unlikely]] {
     if ((Instr & LDAXR_MASK) == LDAR_INST ||  // LDAR*
         (Instr & LDAXR_MASK) == LDAPR_INST) { // LDAPR*
       if (ArchHelpers::Arm64::HandleAtomicLoad(Instr, GPRs, 0)) {
@@ -2016,6 +2016,8 @@ HandleUnalignedAccess(FEXCore::Core::InternalThreadState* Thread, UnalignedHandl
         LogMan::Msg::EFmt("Unhandled JIT SIGBUS LDLUR*: PC: 0x{:x} Instruction: 0x{:08x}\n", ProgramCounter, PC[0]);
         return std::nullopt;
       }
+    } else if (!IsJIT) {
+      return 0;
     }
   }
 

FEXCore/include/FEXCore/Utils/ArchHelpers/Arm64.h

@@ -34,6 +34,6 @@ enum class UnalignedHandlerType {
  * by. FEXCore will return a positive or negative offset depending on internal handling.
  */
 [[nodiscard]]
-FEX_DEFAULT_VISIBILITY std::optional<int32_t>
-HandleUnalignedAccess(FEXCore::Core::InternalThreadState* Thread, UnalignedHandlerType HandleType, uintptr_t ProgramCounter, uint64_t* GPRs);
+FEX_DEFAULT_VISIBILITY std::optional<int32_t> HandleUnalignedAccess(
+  FEXCore::Core::InternalThreadState* Thread, UnalignedHandlerType HandleType, uintptr_t ProgramCounter, uint64_t* GPRs, bool IsJIT = true);
 } // namespace FEXCore::ArchHelpers::Arm64

Source/Windows/ARM64EC/Module.cpp

@@ -349,15 +349,11 @@ struct alignas(16) KiUserExceptionDispatcherStackLayout {
   uint64_t Redzone[2];
 };
 
-static bool HandleUnalignedAccess(ARM64_NT_CONTEXT& Context) {
+static bool HandleUnalignedAccess(ARM64_NT_CONTEXT& Context, bool IsJIT) {
   auto Thread = GetCPUArea().ThreadState();
-  if (!CTX->IsAddressInCodeBuffer(Thread, Context.Pc)) {
-    return false;
-  }
-
   FEXCORE_PROFILE_INSTANT_INCREMENT(Thread, AccumulatedSIGBUSCount, 1);
   const auto Result =
-    FEXCore::ArchHelpers::Arm64::HandleUnalignedAccess(Thread, HandlerConfig->GetUnalignedHandlerType(), Context.Pc, &Context.X0);
+    FEXCore::ArchHelpers::Arm64::HandleUnalignedAccess(Thread, HandlerConfig->GetUnalignedHandlerType(), Context.Pc, &Context.X0, IsJIT);
   Context.Pc += Result.value_or(0);
   return Result.has_value();
 }
@@ -761,16 +757,17 @@ bool ResetToConsistentStateImpl(EXCEPTION_RECORD* Exception, CONTEXT* GuestConte
     }
   }
 
-  if (!CTX->IsAddressInCodeBuffer(Thread, NativeContext->Pc) && !IsDispatcherAddress(NativeContext->Pc)) {
-    LogMan::Msg::DFmt("Passing through exception");
-    return false;
-  }
-
-  if (Exception->ExceptionCode == EXCEPTION_DATATYPE_MISALIGNMENT && Exception::HandleUnalignedAccess(*NativeContext)) {
+  bool IsJIT = CTX->IsAddressInCodeBuffer(Thread, NativeContext->Pc);
+  if (Exception->ExceptionCode == EXCEPTION_DATATYPE_MISALIGNMENT && Exception::HandleUnalignedAccess(*NativeContext, IsJIT)) {
     LogMan::Msg::DFmt("Handled unaligned atomic: new pc: {:X}", NativeContext->Pc);
     return true;
   }
 
+  if (!IsJIT && !IsDispatcherAddress(NativeContext->Pc)) {
+    LogMan::Msg::DFmt("Passing through exception");
+    return false;
+  }
+
   // The JIT (in CompileBlock) emits code to check the suspend doorbell at the start of every block, and run the following instruction if it is set:
   static constexpr uint32_t SuspendTrapMagic {0xD4395FC0}; // brk #0xCAFE
   if (Exception->ExceptionCode == EXCEPTION_ILLEGAL_INSTRUCTION && *reinterpret_cast<uint32_t*>(NativeContext->Pc) == SuspendTrapMagic) {