From edff3dfa4a42ae230157d9312df74d488b6db4fa Mon Sep 17 00:00:00 2001 From: Billy Laws Date: Mon, 22 Dec 2025 16:04:44 +0000 Subject: [PATCH] ImageTracker: Load AOT images For now, a simple directory structure with filename based mapping is used: LOCALAPPDATA/cache/
/{, } All cache blobs are loaded once at the start for simplicity, this will also be useful in the future as metadata validation for settings etc can be done just once. --- Source/Windows/Common/ImageTracker.cpp | 81 ++++++++++++++++++++++++++ Source/Windows/Common/ImageTracker.h | 7 +++ Source/Windows/include/winternl.h | 2 + 3 files changed, 90 insertions(+) diff --git a/Source/Windows/Common/ImageTracker.cpp b/Source/Windows/Common/ImageTracker.cpp index 0076b3b18..740e652df 100644 --- a/Source/Windows/Common/ImageTracker.cpp +++ b/Source/Windows/Common/ImageTracker.cpp @@ -167,6 +167,12 @@ FEXCore::ExecutableFileSectionInfo ImageTracker::HandleImageMap(std::string_view Writer->AppendSetMainExecutable(ImageInfo->Info); CTX.SetCodeMapWriter(std::move(Writer)); } + LoadAOTImages(*ImageInfo); + } + + auto AOTImage = AOTImages.find(ID); + if (AOTImage != AOTImages.end()) { + CTX.GetCodeCache().LoadData(nullptr, AOTImage->second.Data, ImageInfo->SectionInfo); } } @@ -209,4 +215,79 @@ int ImageTracker::OpenCodeMapFile() { } return _sopen(ActiveCodeMapPath.c_str(), O_CREAT | O_TRUNC | O_WRONLY | O_APPEND, _SH_DENYRW, 0644); } + +void ImageTracker::LoadAOTImages(MappedImageInfo& ImageInfo) { + const auto AnsiPath = + fmt::format("\\??\\{}cache\\{}", FEX::Config::GetCacheDirectory(), FEXCore::CodeMap::GetBaseFilename(ImageInfo.Info, false)); + + // Iterate over all files in the given executable's cache directory, mapping them into memory for future use. + // Each cache file name matches the unique ID (as returned by FEXCore::CodeMap::GetBaseFilename) of the image it corresponds to. + ScopedUnicodeString NtPath(AnsiPath.c_str()); + + OBJECT_ATTRIBUTES DirAttr; + InitializeObjectAttributes(&DirAttr, &*NtPath, OBJ_CASE_INSENSITIVE, NULL, NULL); + + ScopedHandle DirHandle; + IO_STATUS_BLOCK IOSB; + if (!NT_SUCCESS(NtOpenFile(&*DirHandle, FILE_LIST_DIRECTORY | SYNCHRONIZE, &DirAttr, &IOSB, FILE_SHARE_READ | FILE_SHARE_WRITE, + FILE_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT))) { + return; + } + + std::array DirBuffer; + bool FirstScan = true; + auto QueryDir = [&]() { + NTSTATUS Status = NtQueryDirectoryFile(*DirHandle, nullptr, nullptr, nullptr, &IOSB, DirBuffer.data(), DirBuffer.size(), + FileBothDirectoryInformation, FALSE, nullptr, FirstScan); + if (FirstScan) { + FirstScan = false; + } + return NT_SUCCESS(Status); + }; + + while (QueryDir()) { + auto* Info = reinterpret_cast(DirBuffer.data()); + + while (true) { + UNICODE_STRING CurrentFileName; + CurrentFileName.Buffer = Info->FileName; + CurrentFileName.Length = static_cast(Info->FileNameLength); + CurrentFileName.MaximumLength = CurrentFileName.Length; + + bool Skip = (Info->FileAttributes & FILE_ATTRIBUTE_DIRECTORY) || (Info->FileNameLength == 2 && Info->FileName[0] == '.') || + (Info->FileNameLength == 4 && Info->FileName[0] == '.' && Info->FileName[1] == '.'); + + if (!Skip) { + OBJECT_ATTRIBUTES FileAttr; + InitializeObjectAttributes(&FileAttr, &CurrentFileName, OBJ_CASE_INSENSITIVE, *DirHandle, nullptr); + + ScopedHandle FileHandle; + if (NT_SUCCESS(NtOpenFile(&*FileHandle, GENERIC_READ | SYNCHRONIZE, &FileAttr, &IOSB, FILE_SHARE_READ, FILE_SYNCHRONOUS_IO_NONALERT))) { + ScopedHandle SectionHandle; + if (NT_SUCCESS(NtCreateSection(&*SectionHandle, SECTION_MAP_EXECUTE | SECTION_MAP_READ, nullptr, nullptr, PAGE_EXECUTE_READ, + SEC_COMMIT, *FileHandle))) { + void* LoadAddress = nullptr; + SIZE_T MappedSize = 0; + if (NT_SUCCESS(NtMapViewOfSection(*SectionHandle, NtCurrentProcess(), &LoadAddress, 0, 0, nullptr, &MappedSize, ViewUnmap, + MEM_RESERVE | MEM_TOP_DOWN, PAGE_EXECUTE_READ))) { + fextl::string UniqueId; + ULONG AnsiLength = 0; + RtlUnicodeToMultiByteSize(&AnsiLength, Info->FileName, Info->FileNameLength); + UniqueId.resize(AnsiLength); + RtlUnicodeToMultiByteN(UniqueId.data(), AnsiLength, NULL, Info->FileName, Info->FileNameLength); + + AOTImages[UniqueId] = {.Data = static_cast(LoadAddress)}; + LogMan::Msg::IFmt("Loaded cache: {}", UniqueId); + } + } + } + } + + if (Info->NextEntryOffset == 0) { + break; + } + Info = reinterpret_cast(reinterpret_cast(Info) + Info->NextEntryOffset); + } + } +} } // namespace FEX::Windows diff --git a/Source/Windows/Common/ImageTracker.h b/Source/Windows/Common/ImageTracker.h index 9588861af..8744d5049 100644 --- a/Source/Windows/Common/ImageTracker.h +++ b/Source/Windows/Common/ImageTracker.h @@ -55,6 +55,12 @@ private: fextl::unordered_map Relocations); }; + struct AOTImageInfo { + std::byte* Data; + }; + + void LoadAOTImages(MappedImageInfo& Info); + FEXCore::Context::Context& CTX; FEX_CONFIG_OPT(ExtendedVolatileMetadataConfig, EXTENDEDVOLATILEMETADATA); @@ -62,6 +68,7 @@ private: std::shared_mutex ImagesLock; std::map MappedImages; + std::map AOTImages; std::string ActiveCodeMapPath; bool IsGeneratingCache; diff --git a/Source/Windows/include/winternl.h b/Source/Windows/include/winternl.h index 4a47000a1..01f8e50fb 100644 --- a/Source/Windows/include/winternl.h +++ b/Source/Windows/include/winternl.h @@ -519,6 +519,8 @@ NTSYSAPI NTSTATUS WINAPI NtMapViewOfSection(HANDLE, HANDLE, PVOID*, ULONG_PTR, S NTSTATUS WINAPI NtOpenKeyEx(PHANDLE, ACCESS_MASK, const OBJECT_ATTRIBUTES*, ULONG); NTSTATUS WINAPI NtProtectVirtualMemory(HANDLE, PVOID*, SIZE_T*, ULONG, ULONG*); NTSTATUS WINAPI NtQueryAttributesFile(const OBJECT_ATTRIBUTES*, FILE_BASIC_INFORMATION*); +NTSTATUS WINAPI NtQueryDirectoryFile(HANDLE, HANDLE, PIO_APC_ROUTINE, PVOID, PIO_STATUS_BLOCK, PVOID, ULONG, FILE_INFORMATION_CLASS, + BOOLEAN, PUNICODE_STRING, BOOLEAN); NTSTATUS WINAPI NtQueryValueKey(HANDLE, const UNICODE_STRING*, KEY_VALUE_INFORMATION_CLASS, void*, DWORD, DWORD*); NTSTATUS WINAPI NtQueryVirtualMemory(HANDLE, LPCVOID, MEMORY_INFORMATION_CLASS, PVOID, SIZE_T, SIZE_T*); NTSTATUS WINAPI NtReadFile(HANDLE, HANDLE, PIO_APC_ROUTINE, PVOID, PIO_STATUS_BLOCK, PVOID, ULONG, PLARGE_INTEGER, PULONG);