Merge pull request #4378 from bylaws/volmd

Implement PE volatile metadata support

FEXCore/Source/Interface/Config/Config.json.in

@@ -420,6 +420,14 @@
           "Should work without issues in most cases."
         ]
       },
+      "VolatileMetadata": {
+        "Type": "bool",
+        "Default": "true",
+        "Desc": [
+          "Use volatile metadata in PE files to inform TSO instructions when available.",
+	  "When metadata is unavailable falls back to the currently enabled TSO options."
+        ]
+      },
       "X87ReducedPrecision": {
         "Type": "bool",
         "Default": "false",

FEXCore/Source/Interface/Context/Context.h

@@ -186,6 +186,10 @@ public:
 
   void AddThunkTrampolineIRHandler(uintptr_t Entrypoint, uintptr_t GuestThunkEntrypoint) override;
 
+  void AddForceTSOInformation(const IntervalList<uint64_t>& ValidRanges, fextl::set<uint64_t>&& Instructions) override;
+
+  void RemoveForceTSOInformation(uint64_t Address, uint64_t Size) override;
+
 public:
   friend class FEXCore::HLE::SyscallHandler;
 #ifdef JIT_ARM64
@@ -373,5 +377,7 @@ private:
   std::shared_mutex CustomIRMutex;
   std::atomic<bool> HasCustomIRHandlers {};
   fextl::unordered_map<uint64_t, std::tuple<CustomIREntrypointHandler, void*, void*>> CustomIRHandlers;
+  IntervalList<uint64_t> ForceTSOValidRanges; // The ranges for which ForceTSOInstructions has populated data
+  fextl::set<uint64_t> ForceTSOInstructions;
 };
 } // namespace FEXCore::Context

FEXCore/Source/Interface/Core/Core.cpp

@@ -565,6 +565,16 @@ ContextImpl::GenerateIR(FEXCore::Core::InternalThreadState* Thread, uint64_t Gue
 
     for (size_t j = 0; j < CodeBlocks->size(); ++j) {
       const FEXCore::Frontend::Decoder::DecodedBlocks& Block = CodeBlocks->at(j);
+
+      bool BlockInForceTSOValidRange = false;
+      auto InstForceTSOIt = ForceTSOInstructions.end();
+      if (ForceTSOValidRanges.Contains({Block.Entry, Block.Entry + Block.Size})) {
+        if (auto It = ForceTSOInstructions.lower_bound(Block.Entry); *It < Block.Entry + Block.Size) {
+          InstForceTSOIt = It;
+          BlockInForceTSOValidRange = true;
+        }
+      }
+
       // Set the block entry point
       Thread->OpDispatcher->SetNewBlockIfChanged(Block.Entry);
 
@@ -581,6 +591,7 @@ ContextImpl::GenerateIR(FEXCore::Core::InternalThreadState* Thread, uint64_t Gue
       }
 
       for (size_t i = 0; i < InstsInBlock; ++i) {
+        uint64_t InstAddress = Block.Entry + BlockInstructionsLength;
         const FEXCore::X86Tables::X86InstInfo* TableInfo {nullptr};
         const FEXCore::X86Tables::DecodedInst* DecodedInfo {nullptr};
 
@@ -603,7 +614,7 @@ ContextImpl::GenerateIR(FEXCore::Core::InternalThreadState* Thread, uint64_t Gue
         Thread->OpDispatcher->FlushRegisterCache(true);
 
         if (ExtendedDebugInfo || Thread->OpDispatcher->CanHaveSideEffects(TableInfo, DecodedInfo)) {
-          Thread->OpDispatcher->_GuestOpcode(Block.Entry + BlockInstructionsLength - GuestRIP);
+          Thread->OpDispatcher->_GuestOpcode(InstAddress - GuestRIP);
         }
 
         if (Config.SMCChecks == FEXCore::Config::CONFIG_SMC_FULL) {
@@ -620,7 +631,7 @@ ContextImpl::GenerateIR(FEXCore::Core::InternalThreadState* Thread, uint64_t Gue
 
           Thread->OpDispatcher->SetCurrentCodeBlock(CodeWasChangedBlock);
           Thread->OpDispatcher->_ThreadRemoveCodeEntry();
-          Thread->OpDispatcher->ExitFunction(Thread->OpDispatcher->_EntrypointOffset(GPRSize, Block.Entry + BlockInstructionsLength - GuestRIP));
+          Thread->OpDispatcher->ExitFunction(Thread->OpDispatcher->_EntrypointOffset(GPRSize, InstAddress - GuestRIP));
 
           auto NextOpBlock = Thread->OpDispatcher->CreateNewCodeBlockAfter(CurrentBlock);
 
@@ -632,17 +643,27 @@ ContextImpl::GenerateIR(FEXCore::Core::InternalThreadState* Thread, uint64_t Gue
           auto Fn = TableInfo->OpcodeDispatcher;
           Thread->OpDispatcher->ResetHandledLock();
           Thread->OpDispatcher->ResetDecodeFailure();
+          IR::ForceTSOMode ForceTSO =
+            BlockInForceTSOValidRange ?
+              (InstForceTSOIt != ForceTSOInstructions.end() && *InstForceTSOIt == InstAddress ? IR::ForceTSOMode::ForceEnabled :
+                                                                                                IR::ForceTSOMode::ForceDisabled) :
+              IR::ForceTSOMode::NoOverride;
+          Thread->OpDispatcher->SetForceTSO(ForceTSO);
           std::invoke(Fn, Thread->OpDispatcher, DecodedInfo);
           if (Thread->OpDispatcher->HadDecodeFailure()) {
             HadDispatchError = true;
           } else {
             if (Thread->OpDispatcher->HasHandledLock() != IsLocked) {
               HadDispatchError = true;
-              LogMan::Msg::EFmt("Missing LOCK HANDLER at 0x{:x}{{'{}'}}", Block.Entry + BlockInstructionsLength, TableInfo->Name ?: "UND");
+              LogMan::Msg::EFmt("Missing LOCK HANDLER at 0x{:x}{{'{}'}}", InstAddress, TableInfo->Name ?: "UND");
             }
             BlockInstructionsLength += DecodedInfo->InstSize;
             TotalInstructionsLength += DecodedInfo->InstSize;
             ++TotalInstructions;
+
+            // Walk InstForceTSOIt forward past the handled instruction
+            InstForceTSOIt =
+              std::find_if(InstForceTSOIt, ForceTSOInstructions.end(), [&](auto Val) { return Val >= Block.Entry + BlockInstructionsLength; });
           }
         } else {
           // Invalid instruction
@@ -971,6 +992,19 @@ void ContextImpl::AddThunkTrampolineIRHandler(uintptr_t Entrypoint, uintptr_t Gu
   }
 }
 
+void ContextImpl::AddForceTSOInformation(const IntervalList<uint64_t>& ValidRanges, fextl::set<uint64_t>&& Instructions) {
+  LogMan::Throw::AFmt(CodeInvalidationMutex.try_lock() == false, "CodeInvalidationMutex needs to be unique_locked here");
+  ForceTSOValidRanges.Insert(ValidRanges);
+  ForceTSOInstructions.merge(Instructions);
+}
+
+void ContextImpl::RemoveForceTSOInformation(uint64_t Address, uint64_t Size) {
+  LogMan::Throw::AFmt(CodeInvalidationMutex.try_lock() == false, "CodeInvalidationMutex needs to be unique_locked here");
+
+  ForceTSOValidRanges.Remove({Address, Address + Size});
+  ForceTSOInstructions.erase(ForceTSOInstructions.lower_bound(Address), ForceTSOInstructions.upper_bound(Address + Size));
+}
+
 void ContextImpl::RemoveCustomIREntrypoint(uintptr_t Entrypoint) {
   LOGMAN_THROW_A_FMT(Config.Is64BitMode || !(Entrypoint >> 32), "64-bit Entrypoint in 32-bit mode {:x}", Entrypoint);
 

FEXCore/Source/Interface/Core/OpcodeDispatcher.h

@@ -47,6 +47,12 @@ enum class BTAction {
   BTComplement,
 };
 
+enum class ForceTSOMode {
+  NoOverride,
+  ForceDisabled,
+  ForceEnabled,
+};
+
 struct LoadSourceOptions {
   // Alignment of the load in bytes. iInvalid signifies opsize aligned.
   IR::OpSize Align = OpSize::iInvalid;
@@ -261,6 +267,13 @@ public:
     return HandledLock;
   }
 
+  void SetForceTSO(ForceTSOMode Mode) {
+    ForceTSO = Mode;
+  }
+  ForceTSOMode GetForceTSO() const {
+    return ForceTSO;
+  }
+
   void SetDumpIR(bool DumpIR) {
     ShouldDump = DumpIR;
   }
@@ -1319,6 +1332,7 @@ private:
   bool HandledLock {false};
   bool DecodeFailure {false};
   bool NeedsBlockEnd {false};
+  ForceTSOMode ForceTSO {ForceTSOMode::NoOverride};
   // Used during new op bringup
   bool ShouldDump {false};
 
@@ -2346,7 +2360,11 @@ private:
   IROp_IRHeader* CurrentHeader {};
 
   bool IsTSOEnabled(FEXCore::IR::RegisterClassType Class) {
-    if (Class == FPRClass) {
+    if (ForceTSO == ForceTSOMode::ForceEnabled) {
+      return true;
+    } else if (ForceTSO == ForceTSOMode::ForceDisabled) {
+      return false;
+    } else if (Class == FPRClass) {
       return CTX->IsVectorAtomicTSOEnabled();
     } else {
       return CTX->IsAtomicTSOEnabled();

FEXCore/include/FEXCore/Core/Context.h

@@ -7,6 +7,7 @@
 #include <FEXCore/Core/CPUID.h>
 #include <FEXCore/IR/IR.h>
 #include <FEXCore/Utils/CompilerDefs.h>
+#include <FEXCore/Utils/IntervalList.h>
 #include <FEXCore/fextl/memory.h>
 #include <FEXCore/fextl/set.h>
 #include <FEXCore/fextl/string.h>
@@ -216,6 +217,15 @@ public:
    */
   FEX_DEFAULT_VISIBILITY virtual void AddThunkTrampolineIRHandler(uintptr_t Entrypoint, uintptr_t GuestThunkEntrypoint) = 0;
 
+  /**
+   * @brief Adds additional per-instruction granularity TSO enable/disable information for the given range.
+   *
+   * @param ValidRanges The set of address ranges covered by this information
+   * @param Instructions The set of instruction addresses within the given ranges for which TSO should be enabled
+   */
+  FEX_DEFAULT_VISIBILITY virtual void AddForceTSOInformation(const IntervalList<uint64_t>& ValidRanges, fextl::set<uint64_t>&& Instructions) = 0;
+
+  FEX_DEFAULT_VISIBILITY virtual void RemoveForceTSOInformation(uint64_t Address, uint64_t Size) = 0;
 private:
 };
 

FEXCore/include/FEXCore/Utils/IntervalList.h

@@ -6,6 +6,7 @@
 
 #include <FEXCore/fextl/vector.h>
 
+namespace FEXCore {
 template<typename SizeType>
 class IntervalList {
 public:
@@ -66,6 +67,12 @@ public:
     FirstIt->End = End;
   }
 
+  void Insert(const IntervalList<SizeType>& Other) {
+    for (const auto& Interval : Other.Intervals) {
+      Insert(Interval);
+    }
+  }
+
   void Remove(Interval Entry) {
     if (Entry.Offset == Entry.End) {
       return;
@@ -121,7 +128,7 @@ public:
     Intervals.erase(EraseStartIt, EraseEndIt);
   }
 
-  QueryResult Query(SizeType Offset) {
+  QueryResult Query(SizeType Offset) const {
     const auto It = std::upper_bound(Intervals.begin(), Intervals.end(), Offset, [](const auto& LHS, const auto& RHS) {
       return LHS < RHS.End;
     }); // Lowest offset interval that (maybe) overlaps with the query offset
@@ -135,11 +142,21 @@ public:
     }
   }
 
-  bool Intersect(Interval Entry) {
+  bool Intersect(Interval Entry) const {
     const auto It = std::upper_bound(Intervals.begin(), Intervals.end(), Entry, [](const auto& LHS, const auto& RHS) {
       return LHS.Offset < RHS.End;
     }); // Lowest offset interval that (maybe) overlaps with the query offset
 
     return It != Intervals.end() && It->Offset < Entry.End;
   }
+
+  bool Contains(Interval Entry) const {
+    const auto It = std::upper_bound(Intervals.begin(), Intervals.end(), Entry, [](const auto& LHS, const auto& RHS) {
+      return LHS.Offset < RHS.End;
+    }); // Lowest offset interval that (maybe) overlaps with the query offset
+
+    return It != Intervals.end() && It->Offset <= Entry.Offset && It->End >= Entry.End;
+  }
 };
+
+} // namespace FEXCore

Source/Tools/FEXConfig/main.qml

@@ -601,6 +601,11 @@ ApplicationWindow {
                         text: qsTr("Enable non-tearing split-lock atomics")
                         config: "StrictInProcessSplitLocks"
                     }
+                    ConfigCheckBox {
+                        topPadding: 4
+                        text: qsTr("Use PE volatile metadata for ARM64EC")
+                        config: "VolatileMetadata"
+                    }
                 }
             }
 

Source/Windows/ARM64EC/Module.cpp

@@ -24,6 +24,7 @@ $end_info$
 #include <FEXCore/Utils/ArchHelpers/Arm64.h>
 #include <FEXCore/Utils/MathUtils.h>
 #include <FEXCore/Utils/TypeDefines.h>
+#include <FEXCore/Utils/SignalScopeGuards.h>
 
 #include "Common/ArgumentLoader.h"
 #include "Common/Config.h"
@@ -243,6 +244,55 @@ void InitSyscalls() {
   FillNtDllLUTs(NtDll);
   PatchCallChecker();
 }
+
+void LoadImageVolatileMetadata(uint64_t Address) {
+  const auto Module = reinterpret_cast<HMODULE>(Address);
+  IMAGE_NT_HEADERS* Nt = RtlImageNtHeader(Module);
+  uint64_t EndAddress = Address + Nt->OptionalHeader.SizeOfImage;
+  ULONG Size;
+  const auto* LoadConfig =
+    reinterpret_cast<_IMAGE_LOAD_CONFIG_DIRECTORY64*>(RtlImageDirectoryEntryToData(Module, true, IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG, &Size));
+  if (!LoadConfig || LoadConfig->Size <= offsetof(_IMAGE_LOAD_CONFIG_DIRECTORY64, VolatileMetadataPointer)) {
+    return;
+  }
+
+  if (LoadConfig->VolatileMetadataPointer < Address || LoadConfig->VolatileMetadataPointer + sizeof(IMAGE_VOLATILE_METADATA) >= EndAddress) {
+    return;
+  }
+
+  const auto* VolatileMetadata = reinterpret_cast<IMAGE_VOLATILE_METADATA*>(LoadConfig->VolatileMetadataPointer);
+  if (!VolatileMetadata || Address + VolatileMetadata->VolatileAccessTable + VolatileMetadata->VolatileAccessTableSize >= EndAddress ||
+      Address + VolatileMetadata->VolatileInfoRangeTable + VolatileMetadata->VolatileInfoRangeTableSize >= EndAddress) {
+    return;
+  }
+
+  fextl::set<uint64_t> VolatileInstructions;
+  const auto* VolatileAccessTableBegin = reinterpret_cast<IMAGE_VOLATILE_RVA_METADATA*>(Address + VolatileMetadata->VolatileAccessTable);
+  const auto* VolatileAccessTableEnd =
+    VolatileAccessTableBegin + (VolatileMetadata->VolatileAccessTableSize / sizeof(IMAGE_VOLATILE_RVA_METADATA));
+  for (auto It = VolatileAccessTableBegin; It != VolatileAccessTableEnd; It++) {
+    VolatileInstructions.emplace(Address + It->Rva);
+  }
+
+  FEXCore::IntervalList<uint64_t> VolatileValidRanges;
+  const auto* VolatileInfoRangeTableBegin = reinterpret_cast<IMAGE_VOLATILE_RANGE_METADATA*>(Address + VolatileMetadata->VolatileInfoRangeTable);
+  const auto* VolatileInfoRangeTableEnd =
+    VolatileInfoRangeTableBegin + (VolatileMetadata->VolatileInfoRangeTableSize / sizeof(IMAGE_VOLATILE_RANGE_METADATA));
+  for (auto It = VolatileInfoRangeTableBegin; It != VolatileInfoRangeTableEnd; It++) {
+    VolatileValidRanges.Insert({Address + It->Rva, Address + It->Rva + It->Size});
+  }
+
+  std::scoped_lock Lock(CTX->GetCodeInvalidationMutex());
+  CTX->AddForceTSOInformation(VolatileValidRanges, std::move(VolatileInstructions));
+}
+
+void HandleImageMap(uint64_t Address) {
+  FEX_CONFIG_OPT(VolatileMetadata, VOLATILEMETADATA);
+  if (VolatileMetadata) {
+    LoadImageVolatileMetadata(Address);
+  }
+  InvalidationTracker->HandleImageMap(Address);
+}
 } // namespace
 
 namespace Exception {
@@ -563,7 +613,7 @@ NTSTATUS ProcessInit() {
   InvalidationTracker.emplace(*CTX, Threads);
 
   auto MainModule = reinterpret_cast<__TEB*>(NtCurrentTeb())->Peb->ImageBaseAddress;
-  InvalidationTracker->HandleImageMap(reinterpret_cast<uint64_t>(MainModule));
+  HandleImageMap(reinterpret_cast<uint64_t>(MainModule));
 
   CPUFeatures.emplace(*CTX);
 
@@ -745,8 +795,12 @@ NTSTATUS NotifyMapViewOfSection(void* Unk1, void* Address, void* Unk2, SIZE_T Si
     return STATUS_SUCCESS;
   }
 
-  std::scoped_lock Lock(ThreadCreationMutex);
-  InvalidationTracker->HandleImageMap(reinterpret_cast<uint64_t>(Address));
+  {
+    std::scoped_lock Lock(ThreadCreationMutex);
+    HandleImageMap(reinterpret_cast<uint64_t>(Address));
+  }
+
+
   return STATUS_SUCCESS;
 }
 
@@ -760,7 +814,11 @@ void NotifyUnmapViewOfSection(void* Address, BOOL After, NTSTATUS Status) {
   }
 
   std::scoped_lock Lock(ThreadCreationMutex);
-  InvalidationTracker->InvalidateContainingSection(reinterpret_cast<uint64_t>(Address), true);
+  auto [Start, Size] = InvalidationTracker->InvalidateContainingSection(reinterpret_cast<uint64_t>(Address), true);
+  if (Size) {
+    std::scoped_lock Lock(CTX->GetCodeInvalidationMutex());
+    CTX->RemoveForceTSOInformation(Start, Size);
+  }
 }
 
 void FlushInstructionCacheHeavy(const void* Address, SIZE_T Size) {

Source/Windows/Common/InvalidationTracker.cpp

@@ -50,10 +50,10 @@ void InvalidationTracker::HandleImageMap(uint64_t Address) {
   }
 }
 
-void InvalidationTracker::InvalidateContainingSection(uint64_t Address, bool Free) {
+InvalidationTracker::InvalidateContainingSectionResult InvalidationTracker::InvalidateContainingSection(uint64_t Address, bool Free) {
   MEMORY_BASIC_INFORMATION Info;
   if (NtQueryVirtualMemory(NtCurrentProcess(), reinterpret_cast<void*>(Address), MemoryBasicInformation, &Info, sizeof(Info), nullptr)) {
-    return;
+    return {Address, 0};
   }
 
   const auto SectionBase = reinterpret_cast<uint64_t>(Info.AllocationBase);
@@ -75,6 +75,8 @@ void InvalidationTracker::InvalidateContainingSection(uint64_t Address, bool Fre
     std::scoped_lock Lock(RWXIntervalsLock);
     RWXIntervals.Remove({SectionBase, SectionBase + SectionSize});
   }
+
+  return {SectionBase, SectionSize};
 }
 
 void InvalidationTracker::InvalidateAlignedInterval(uint64_t Address, uint64_t Size, bool Free) {

Source/Windows/Common/InvalidationTracker.h

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 #pragma once
 
-#include "IntervalList.h"
+#include <FEXCore/Utils/IntervalList.h>
 #include <mutex>
 #include <unordered_map>
 
@@ -22,13 +22,17 @@ public:
   InvalidationTracker(FEXCore::Context::Context& CTX, const std::unordered_map<DWORD, FEXCore::Core::InternalThreadState*>& Threads);
   void HandleMemoryProtectionNotification(uint64_t Address, uint64_t Size, ULONG Prot);
   void HandleImageMap(uint64_t Address);
-  void InvalidateContainingSection(uint64_t Address, bool Free);
+  struct InvalidateContainingSectionResult {
+    uint64_t SectionStart;
+    uint64_t SectionSize;
+  };
+  InvalidateContainingSectionResult InvalidateContainingSection(uint64_t Address, bool Free);
   void InvalidateAlignedInterval(uint64_t Address, uint64_t Size, bool Free);
   void ReprotectRWXIntervals(uint64_t Address, uint64_t Size);
   bool HandleRWXAccessViolation(uint64_t FaultAddress);
 
 private:
-  IntervalList<uint64_t> RWXIntervals;
+  FEXCore::IntervalList<uint64_t> RWXIntervals;
   std::mutex RWXIntervalsLock;
   FEXCore::Context::Context& CTX;
   const std::unordered_map<DWORD, FEXCore::Core::InternalThreadState*>& Threads;

Source/Windows/Common/OvercommitTracker.h

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 #pragma once
 
-#include "IntervalList.h"
+#include <FEXCore/Utils/IntervalList.h>
 #include <thread>
 #include <shared_mutex>
 
@@ -13,7 +13,7 @@ namespace FEX::Windows {
 class OvercommitTracker {
 private:
   bool IsWine;
-  IntervalList<uint64_t> OvercommitIntervals;
+  FEXCore::IntervalList<uint64_t> OvercommitIntervals;
   std::shared_mutex OvercommitIntervalsMutex;
 
 public:

Source/Windows/include/winnt.h

@@ -136,6 +136,25 @@ typedef struct _CONTEXT_EX {
 #endif
 } CONTEXT_EX, *PCONTEXT_EX;
 
+// From process hacker
+typedef struct _IMAGE_VOLATILE_METADATA {
+  ULONG Size;
+  ULONG Version;
+  ULONG VolatileAccessTable;
+  ULONG VolatileAccessTableSize;
+  ULONG VolatileInfoRangeTable;
+  ULONG VolatileInfoRangeTableSize;
+} IMAGE_VOLATILE_METADATA, *PIMAGE_VOLATILE_METADATA;
+
+typedef struct _IMAGE_VOLATILE_RVA_METADATA {
+  ULONG Rva;
+} IMAGE_VOLATILE_RVA_METADATA, *PIMAGE_VOLATILE_RVA_METADATA;
+
+typedef struct _IMAGE_VOLATILE_RANGE_METADATA {
+  ULONG Rva;
+  ULONG Size;
+} IMAGE_VOLATILE_RANGE_METADATA, *PIMAGE_VOLATILE_RANGE_METADATA;
+
 NTSYSAPI DWORD WINAPI RtlRunOnceExecuteOnce(PRTL_RUN_ONCE, PRTL_RUN_ONCE_INIT_FN, PVOID, PVOID*);
 
 // This is a FEX extension, and requires corresponding wine patches