fix(agc): import GPU-mapped guest memory when it is mapped

core/libs/prx/libSceAgcDriver/Execution/include/VulkanDevice.hpp

@@ -9,6 +9,7 @@
 #include "prx/libSceAgcDriver/Execution/include/Recipe.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/Draw.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/Recorder.hpp"
+#include "prx/libc/include/GuestAllocations.hpp"
 #include <functional>
 #include <memory>
 #include <optional>
@@ -47,6 +48,7 @@ class VulkanDevice {
     // three-step form below so the GPU wait happens without the mutex.
     void WaitIdle();
     void PrepareForReplacement();
+    bool ImportGuestMemory(const GuestAllocations::Mapped& ranges, std::uint64_t generation, bool adoptDevice);
     // Sends recorded work to the GPU without waiting for it. With `reapFirst` it first retires batches
     // that already finished, so the in-flight list stays short (APS5_NO_OPPORTUNISTIC_REAP=1 skips
     // that). A reap runs completion actions, and a write-back can wait for a later batch under the

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Lifecycle.cpp

@@ -4,6 +4,8 @@
 #include "prx/libSceAgcDriver/Execution/include/WorkerSampler.hpp"
 #include "prx/libSceAgcDriver/Execution/include/ProfileOutput.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/GuestBufferMemory.hpp"
+#include "prx/libc/include/GuestAllocations.hpp"
+#include <cstdio>
 
 namespace AgcDriver::DriverDetail {
 
@@ -13,6 +15,7 @@ Driver& Driver::Get() {
 }
 
 Driver::~Driver() {
+    GuestAllocations::GuestAllocationsSetGpuMapObserver_nid_postfix(nullptr);
     stop();
 }
 
@@ -57,6 +60,16 @@ Driver::Driver() {
         stop();
         throw;
     }
+    GuestAllocations::GuestAllocationsSetGpuMapObserver_nid_postfix([](const GuestAllocations::Mapped& ranges, std::uint64_t generation) {
+        try {
+            const auto current = Driver::Get().device.Load();
+            if (current == nullptr || current->ImportGuestMemory(ranges, generation, false)) return;
+            std::lock_guard lock(GuestMemory::GpuMutex());
+            if (Driver::Get().device.Load() == current) current->ImportGuestMemory(ranges, generation, true);
+        } catch (const std::exception& error) {
+            std::fprintf(stderr, "[gpu] import of GPU memory at its mapping failed: %s\n", error.what());
+        }
+    });
 }
 
 void Driver::WaitIdle() {

core/libs/prx/libSceAgcDriver/Execution/src/VulkanDevice.cpp

@@ -1381,6 +1381,11 @@ VulkanDevice::VulkanDevice(const PresentationWindow* window) : state(std::make_u
 
 VulkanDevice::~VulkanDevice() = default;
 
+bool VulkanDevice::ImportGuestMemory(const GuestAllocations::Mapped& ranges, std::uint64_t generation, bool adoptDevice) {
+    if (adoptDevice) GuestMemory::AssertGpuLockHeld("VulkanDevice::ImportGuestMemory");
+    return Graphics::ImportMappedRanges(graphicsContext(), ranges, generation, adoptDevice);
+}
+
 void VulkanDevice::PrepareForReplacement() {
     GuestMemory::AssertGpuLockHeld("VulkanDevice::PrepareForReplacement");
     require(state->recorder != nullptr && Graphics::Recorder::Active() == state->recorder.get(), "device replacement requires its active recorder");

core/libs/prx/libSceAgcDriver/Graphics/include/GuestBufferMemory.hpp

@@ -55,6 +55,7 @@ void SetImportWatch(const Context& context, ImportWatch watch);
 // (alignment and budget permitting), or null. Bytes at `address` are at `address - import->base` in
 // the import's buffer.
 const HostImport* HostImportFor(const Context& context, std::uint64_t address, std::size_t bytes);
+bool ImportMappedRanges(const Context& context, const GuestAllocations::Mapped& ranges, std::uint64_t generation, bool adoptDevice);
 // Whether an existing import covers [address, address + bytes), without reconciling the imports
 // with the registry or making one (HostImportFor may take a registry lease): a hint for choices
 // made outside the device lock (a sampled texture's path, a dispatch's pre-sync); the path taken

core/libs/prx/libSceAgcDriver/Graphics/src/GuestBufferMemory.cpp

@@ -143,6 +143,7 @@ struct HostImports {
     PFN_vkDestroyBuffer destroyBuffer = nullptr;
     PFN_vkFreeMemory freeMemory = nullptr;
     std::map<std::uint64_t, HostImport> imports;
+    std::uint64_t liveBytes = 0;
     std::set<std::uint64_t> failed;
     // Registry generation the imports were last reconciled with.
     std::uint64_t refreshedGeneration = 0;
@@ -200,6 +201,7 @@ void retireImport(VkDevice device, HostImports& state, std::map<std::uint64_t, H
     auto holder = std::make_shared<RetiredImport>(state.device, state.destroyBuffer, state.freeMemory, it->second);
     if (auto* recorder = Recorder::Active(); recorder != nullptr && !recorder->Idle()) recorder->Keep(std::move(holder));
     ++state.epoch;
+    state.liveBytes -= it->second.bytes;
     state.imports.erase(it);
 }
 
@@ -537,8 +539,7 @@ const HostImport* importAllocation(const Context& context, HostImports& state, s
         const char* value = std::getenv("APS5_HOST_IMPORT_MIB");
         return (value ? std::strtoull(value, nullptr, 10) : 6144ull) << 20u;
     }();
-    std::uint64_t live = 0;
-    for (const auto& [address, existing] : state.imports) live += existing.bytes;
+    const auto live = state.liveBytes;
     if (live + bytes > budget) {
         // A refused import turns every later use of the range into CPU copies, so say so.
         static std::uint64_t refused = 0, refusedBytes = 0;
@@ -635,9 +636,8 @@ const HostImport* importAllocation(const Context& context, HostImports& state, s
     static std::uint64_t importedBytes = 0;
     importedBytes += bytes;
     static const bool trace = std::getenv("APS5_TRACE_HOST_IMPORT") != nullptr;
-    std::uint64_t liveBytes = bytes;
-    for (const auto& [address, existing] : state.imports) liveBytes += existing.bytes;
-    if (trace) std::fprintf(stderr, "[gpu] host import of 0x%llx+0x%llx ok (%zu live, %.1f MiB live, %.1f MiB ever)\n", static_cast<unsigned long long>(base), static_cast<unsigned long long>(bytes), state.imports.size() + 1, liveBytes / 1048576.0, importedBytes / 1048576.0);
+    state.liveBytes += bytes;
+    if (trace) std::fprintf(stderr, "[gpu] host import of 0x%llx+0x%llx ok (%zu live, %.1f MiB live, %.1f MiB ever)\n", static_cast<unsigned long long>(base), static_cast<unsigned long long>(bytes), state.imports.size() + 1, state.liveBytes / 1048576.0, importedBytes / 1048576.0);
     return &state.imports.emplace(base, entry).first->second;
 }
 
@@ -669,6 +669,7 @@ void refreshImports(const Context& context, HostImports& state, const GuestAlloc
 #endif
         }
         state.imports.clear();
+        state.liveBytes = 0;
         state.failed.clear();
         state.device = context.device;
         state.destroyBuffer = context.Function<PFN_vkDestroyBuffer>("vkDestroyBuffer");
@@ -1492,6 +1493,7 @@ void ClearHostImports(VkDevice device) {
 #endif
     }
     state.imports.clear();
+    state.liveBytes = 0;
     state.failed.clear();
     state.device = VK_NULL_HANDLE;
     state.refreshedGeneration = 0;
@@ -1682,6 +1684,32 @@ const HostImport* HostImportFor(const Context& context, std::uint64_t address, s
     return nullptr;
 }
 
+bool ImportMappedRanges(const Context& context, const GuestAllocations::Mapped& ranges, std::uint64_t generation, bool adoptDevice) {
+    if (context.hostImportAlignment == 0) return true;
+    ensurePinWaiter();
+    auto& state = Imports();
+    std::lock_guard lock(state.mutex);
+    if (!adoptDevice && (state.device != context.device || state.watchDevice != context.device)) return false;
+    const bool adopting = state.device != context.device;
+    GuestAllocations::Lease current;
+    if (adopting) {
+        current = GuestAllocations::GuestAllocationsAcquire_nid_postfix();
+        refreshImports(context, state, current);
+    } else if (state.refreshedGeneration > generation) {
+        return true;
+    }
+    for (const auto& mapped : ranges) {
+        const auto range = mapped.lock();
+        if (range == nullptr || (adopting && leasedRangeOwner(current, range->address) != range)) continue;
+        try {
+            if (importAllocation(context, state, range->address, range->bytes, GuestAllocations::Lease{range}) == nullptr && GuestAllocations::GuestAllocationsGeneration_nid_postfix() != generation) state.failed.erase(range->address);
+        } catch (const std::exception& error) {
+            std::fprintf(stderr, "[gpu] import of GPU memory mapped at 0x%llx+0x%zx failed: %s\n", static_cast<unsigned long long>(range->address), range->bytes, error.what());
+        }
+    }
+    return true;
+}
+
 bool RegisteredReadableCovers(std::uint64_t address, std::size_t bytes) {
     if (bytes == 0 || bytes > std::numeric_limits<std::uint64_t>::max() - address) return false;
     const auto lease = GuestAllocations::GuestAllocationsAcquire_nid_postfix();

core/libs/prx/libSceAgcDriver/tests/execution/HostImportAtMap.cpp

@@ -0,0 +1,182 @@
+#include "prx/libSceAgcDriver/Execution/include/GuestMemory.hpp"
+#include "prx/libSceAgcDriver/Execution/include/VulkanDevice.hpp"
+#include "prx/libSceAgcDriver/Graphics/include/GuestBufferMemory.hpp"
+#include "prx/libc/include/GuestAllocations.hpp"
+#include "VulkanTestDevice.hpp"
+#ifdef _WIN32
+#ifndef NOMINMAX
+#define NOMINMAX
+#endif
+#include <windows.h>
+#endif
+#include <array>
+#include <thread>
+#include <atomic>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <iostream>
+#include <mutex>
+
+namespace {
+
+using AgcDriver::Graphics::Require;
+
+constexpr std::size_t BlockBytes = 65536;
+constexpr std::size_t AllocationBytes = BlockBytes * 4;
+
+class GuestAllocation {
+public:
+    GuestAllocation() {
+#ifdef _WIN32
+        memory = static_cast<std::uint8_t*>(VirtualAlloc(nullptr, AllocationBytes, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE));
+#else
+        memory = static_cast<std::uint8_t*>(std::aligned_alloc(BlockBytes, AllocationBytes));
+#endif
+        Require(memory != nullptr, "host import at map: cannot allocate the guest memory");
+        std::memset(memory, 0, AllocationBytes);
+        Register();
+    }
+
+    ~GuestAllocation() {
+        if (registered) Unregister();
+#ifdef _WIN32
+        VirtualFree(memory, 0, MEM_RELEASE);
+#else
+        std::free(memory);
+#endif
+    }
+
+    GuestAllocation(const GuestAllocation&) = delete;
+    GuestAllocation& operator=(const GuestAllocation&) = delete;
+
+    std::uint64_t Address() const { return static_cast<std::uint64_t>(reinterpret_cast<std::uintptr_t>(memory)); }
+
+    void Register() {
+        GuestAllocations::Mutation().Add(memory, AllocationBytes, true, true, true);
+        registered = true;
+    }
+
+    void Unregister() {
+        GuestAllocations::Mutation().Remove(memory);
+        registered = false;
+    }
+
+    GuestAllocations::Mapped Ranges() const {
+        GuestAllocations::Mapped ranges;
+        for (const auto& range : GuestAllocations::GuestAllocationsAcquire_nid_postfix()) {
+            if (range->address == Address()) ranges.push_back(range);
+        }
+        Require(ranges.size() == 1, "host import at map: the allocation is not registered");
+        return ranges;
+    }
+
+private:
+    std::uint8_t* memory = nullptr;
+    bool registered = false;
+};
+
+void Reconcile(AgcDriver::VulkanDevice& device, const GuestAllocation& probe) {
+    std::lock_guard lock(AgcDriver::GuestMemory::GpuMutex());
+    const std::array<std::uint32_t, 4> pattern{5u, 6u, 7u, 8u};
+    Require(device.FillBuffer(probe.Address(), BlockBytes, pattern), "host import at map: the probe fill was not imported");
+    device.WaitIdle();
+}
+
+void Import(AgcDriver::VulkanDevice& device, const GuestAllocations::Mapped& ranges, std::uint64_t generation = GuestAllocations::GuestAllocationsGeneration_nid_postfix()) {
+    if (device.ImportGuestMemory(ranges, generation, false)) return;
+    std::lock_guard lock(AgcDriver::GuestMemory::GpuMutex());
+    Require(device.ImportGuestMemory(ranges, generation, true), "host import at map: the import under the GPU lock refused");
+}
+
+bool Imported(VkDevice device, std::uint64_t address, std::size_t bytes) {
+    AgcDriver::Graphics::Context probe{};
+    probe.device = device;
+    probe.hostImportAlignment = 1;
+    return AgcDriver::Graphics::HostImportCovers(probe, address, bytes);
+}
+
+}
+
+int main() {
+    try {
+#ifdef _WIN32
+        _putenv_s("APS5_PIN_WAIT_MS", "2000");
+#else
+        setenv("APS5_PIN_WAIT_MS", "2000", 1);
+#endif
+        auto device = OpenVulkanTestDevice();
+        if (!device) return VulkanTestSkipped;
+        GuestAllocation first;
+        Import(*device, first.Ranges());
+        const bool importedFirst = Imported(device->Device(), first.Address(), AllocationBytes);
+        first.Unregister();
+        Require(!Imported(device->Device(), first.Address(), AllocationBytes), "host import at map: unmapping an allocation imported before any lookup kept its import");
+        GuestAllocation probe;
+        bool filled = false;
+        {
+            std::lock_guard lock(AgcDriver::GuestMemory::GpuMutex());
+            const std::array<std::uint32_t, 4> pattern{1u, 2u, 3u, 4u};
+            filled = device->FillBuffer(probe.Address(), BlockBytes, pattern);
+            device->WaitIdle();
+        }
+        if (!filled) {
+            std::printf("skipped, the device does not import guest memory\n");
+            return VulkanTestSkipped;
+        }
+        Require(importedFirst, "host import at map: a mapping before any lookup was not imported");
+        GuestAllocation guest;
+        const auto handle = device->Device();
+        Require(!Imported(handle, guest.Address(), AllocationBytes), "host import at map: the allocation was imported before anything asked for it");
+        Import(*device, guest.Ranges());
+        Require(Imported(handle, guest.Address(), AllocationBytes), "host import at map: the mapped allocation was not imported");
+
+        GuestAllocation stale;
+        Import(*device, stale.Ranges());
+        stale.Unregister();
+        GuestAllocation mapped;
+        Import(*device, mapped.Ranges());
+        Require(Imported(handle, mapped.Address(), AllocationBytes), "host import at map: a mapping made after another changed was not imported");
+
+        GuestAllocation remapped;
+        Import(*device, remapped.Ranges());
+        remapped.Unregister();
+        remapped.Register();
+        Import(*device, remapped.Ranges());
+        Require(Imported(handle, remapped.Address(), AllocationBytes), "host import at map: a range mapped again was not imported");
+        GuestAllocation gone;
+        const auto goneRanges = gone.Ranges();
+        gone.Unregister();
+        Import(*device, goneRanges);
+        Require(!Imported(handle, gone.Address(), AllocationBytes), "host import at map: a mapping unregistered before its import was imported");
+        gone.Register();
+
+        GuestAllocation late;
+        const auto lateRanges = late.Ranges();
+        const auto lateAt = GuestAllocations::GuestAllocationsGeneration_nid_postfix();
+        GuestAllocation newer;
+        Reconcile(*device, probe);
+        Import(*device, lateRanges, lateAt);
+        Require(!Imported(handle, late.Address(), AllocationBytes), "host import at map: an import older than the last reconcile was made");
+
+        GuestAllocation contended;
+        std::atomic<bool> held{false};
+        std::atomic<bool> release{false};
+        std::thread holder([&] {
+            std::lock_guard lock(AgcDriver::GuestMemory::GpuMutex());
+            held = true;
+            while (!release) std::this_thread::yield();
+        });
+        while (!held) std::this_thread::yield();
+        const bool unlocked = device->ImportGuestMemory(contended.Ranges(), GuestAllocations::GuestAllocationsGeneration_nid_postfix(), false);
+        release = true;
+        holder.join();
+        Require(unlocked && Imported(handle, contended.Address(), AllocationBytes), "host import at map: a mapping on a device that already imports needed the GPU lock");
+        std::puts("host import at map tests passed");
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        return 1;
+    }
+}

core/libs/prx/libc/include/GuestAllocations.hpp

@@ -22,6 +22,7 @@ struct Range {
 };
 
 using Lease = std::vector<std::shared_ptr<const Range>>;
+using Mapped = std::vector<std::weak_ptr<const Range>>;
 
 extern "C" {
 void* GuestAllocationsBegin_nid_postfix();
@@ -54,6 +55,8 @@ void GuestAllocationsInvalidate_nid_postfix(std::uintptr_t address, std::size_t
 // nothing to wait for counts against the same bound as a plain spin). Without one the mutation spins
 // until the lease is dropped by another thread.
 void GuestAllocationsSetPinWaiter_nid_postfix(bool (*callback)(std::uintptr_t address, std::size_t bytes));
+void GuestAllocationsSetGpuMapObserver_nid_postfix(void (*callback)(const Mapped& ranges, std::uint64_t generation));
+void GuestAllocationsNoteGpuMapping_nid_postfix(void* mutation, const void* pointer, std::size_t bytes);
 }
 
 class Mutation {
@@ -76,6 +79,7 @@ class Mutation {
     void Remove(const void* pointer) { GuestAllocationsRemove_nid_postfix(handle, pointer); }
     void Unmap(const void* pointer, std::size_t bytes, const std::function<void(const void*, std::size_t, const void*, bool)>& apply) { GuestAllocationsUnmap_nid_postfix(handle, pointer, bytes, apply); }
     void Protect(const void* pointer, std::size_t bytes, bool readable, bool writable, bool gpu, const std::function<void()>& apply) { GuestAllocationsProtect_nid_postfix(handle, pointer, bytes, readable, writable, gpu, apply); }
+    void NoteGpuMapping(const void* pointer, std::size_t bytes) { GuestAllocationsNoteGpuMapping_nid_postfix(handle, pointer, bytes); }
 
 private:
     void* handle;