Merge pull request #10304 from Ondra97/codex/fix-agc-dissolved-staging-owners

fix(agc): keep buffer owners when a staged address space dissolves

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

@@ -5563,7 +5563,10 @@ void GuestBufferMemory::resolveOverlays() {
         if (regions[index].overlay) members.push_back(index);
     }
     if (members.empty()) return;
-    std::vector<bool> drop(regions.size(), space == nullptr);
+    std::vector<bool> drop(regions.size(), false);
+    if (space == nullptr) {
+        for (const auto member : members) drop[member] = true;
+    }
     std::sort(members.begin(), members.end(), [&](std::size_t left, std::size_t right) { return regions[left].begin < regions[right].begin; });
     for (std::size_t first = 0; space != nullptr && first < members.size();) {
         Region cluster{regions[members[first]].begin, regions[members[first]].end, true, {}, nullptr};

core/libs/prx/libSceAgcDriver/tests/BdaResources.cpp

@@ -813,6 +813,69 @@ void failedJoinWrittenViewTests(const Context& context, const BdaTestAccess& acc
     ::operator delete(block, std::align_val_t{2 * piece});
 }
 
+void failedJoinDissolvedOverlayTests(const Context& context, const BdaTestAccess& access) {
+    constexpr std::size_t piece = 65536;
+    constexpr std::size_t large = 3 * piece;
+    constexpr std::uint32_t stored = 0x0b5e55ed;
+    auto* block = static_cast<std::uint8_t*>(::operator new(2 * large, std::align_val_t{2 * piece}));
+    std::memset(block, 0x5a, 2 * large);
+    const auto first = reinterpret_cast<std::uintptr_t>(block);
+    const auto second = first + large;
+    const auto written = second + 128;
+    auto importing = context;
+    importing.hostImportAlignment = piece;
+    ClearHostImports(importing.device);
+    SetHostImportWindow(2 * piece);
+    SetHostImportBudget(2 * large);
+    {
+        GuestAllocations::Mutation mutation;
+        mutation.Add(block, large, true, false, true);
+        mutation.Add(block + large, large, true, false, true);
+    }
+    Require(HostImportFor(importing, first, 16) != nullptr && HostImportFor(importing, second, 16) != nullptr && HostImportLiveBytes() == 2 * large, "the failed join overlay fixture did not import its two ranges apart");
+    access.limitImports(large);
+    {
+        GuestBufferMemory memory(importing);
+        memory.AddReadable(second - 16, 32);
+        memory.Upload(false);
+        Require(ExistingHostImportFor(importing, first, 16) == nullptr && ExistingHostImportFor(importing, second, 16) == nullptr, "the join over the import limit did not retire the imports it replaced");
+        memory.WriteBack();
+    }
+    SetHostImportBudget(2 * large);
+    Require(HostImportFor(importing, first, 16) != nullptr && HostImportFor(importing, second, 16) != nullptr, "the ranges of the failed join were not imported apart again");
+    {
+        std::lock_guard gpu(AgcDriver::GuestMemory::GpuMutex());
+        Recorder recorder(importing);
+        recorder.Activate();
+        GuestBufferMemory memory(importing);
+        memory.AllowDeviceStaging();
+        memory.AcquireRegistered();
+        memory.AddWritable(written, 32);
+        memory.AddWritable(second - 16, 32);
+        memory.Upload(true);
+        std::uint32_t adjustment = 0;
+        const auto view = memory.Descriptor(written, 32, adjustment);
+        std::memcpy(access.bytes(view.buffer).data() + view.offset + adjustment, &stored, sizeof(stored));
+        std::uint32_t crossing = 0;
+        const auto spanning = memory.Descriptor(second - 16, 32, crossing);
+        Require(spanning.range >= 32, "the view across the two ranges of the failed join has no owner covering it");
+        memory.WriteBack();
+    }
+    std::uint32_t kept = 0;
+    std::memcpy(&kept, block + large + 128, sizeof(kept));
+    Require(kept == stored, "a write through a view staged in a space that a crossing view dissolved was lost");
+    access.limitImports(std::nullopt);
+    ClearHostImports(importing.device);
+    {
+        GuestAllocations::Mutation mutation;
+        mutation.Remove(block);
+        mutation.Remove(block + large);
+    }
+    SetHostImportBudget(0);
+    SetHostImportWindow(0);
+    ::operator delete(block, std::align_val_t{2 * piece});
+}
+
 // One registered range of `cells` import windows of `window` bytes, with the window and the import budget set for the
 // test and reset after it.
 class WindowedRange {
@@ -4612,6 +4675,7 @@ void RunBdaResourceTests(const Context& context, const BdaTestAccess& access) {
     importCoverageHintBudgetTests(context, access);
     failedImportJoinProbeBudgetTests(context, access);
     failedJoinWrittenViewTests(context, access);
+    failedJoinDissolvedOverlayTests(context, access);
     importWholeRangeOverWindowTests(context);
     importMappedFailureGenerationTests(context);
     importStaleProbeTests(context, access);

core/libs/prx/libSceAgcDriver/tests/Recorder.cpp

@@ -4059,10 +4059,11 @@ bool bdaStagingTests(const Device& device, Recorder& recorder) {
         return false;
     }
     constexpr std::size_t bytes = 4u << 20u;
+    constexpr std::size_t allocationBytes = bytes + 65536;
 #ifdef _WIN32
-    void* block = VirtualAlloc(nullptr, bytes, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+    void* block = VirtualAlloc(nullptr, allocationBytes, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
 #else
-    void* block = std::aligned_alloc(65536, bytes);
+    void* block = std::aligned_alloc(65536, allocationBytes);
 #endif
     Require(block != nullptr, "cannot allocate the test block");
     std::memset(block, 0x11, bytes);
@@ -4152,6 +4153,51 @@ bool bdaStagingTests(const Device& device, Recorder& recorder) {
         recorder.Sync();
         memory.WriteBack();
     }
+    std::memset(block, 0xff, allocationBytes);
+    {
+        std::cout << "checking a crossing write after an address-based staging candidate\n";
+        GuestBufferMemory memory(context);
+        memory.AllowDeviceStaging();
+        memory.AcquireRegistered();
+        memory.AddWritable(element, elementBytes, false, false);
+        const auto crossing = address + bytes - 16;
+        const auto dissolved = AddressSpaceCounters().dissolvedOverlap;
+        memory.AddWritable(crossing, 32, false, false);
+        Require(AddressSpaceCounters().dissolvedOverlap == dissolved + 1, "the crossing write did not dissolve its cached address space");
+        const auto ordinary = address + bytes + 4096;
+        memory.AddReadable(ordinary, 64);
+        memory.Upload(true);
+        std::uint32_t elementAdjustment = 0;
+        std::uint32_t crossingAdjustment = 0;
+        std::uint32_t ordinaryAdjustment = 0;
+        const auto elementView = memory.Descriptor(element, elementBytes, elementAdjustment);
+        const auto crossingView = memory.Descriptor(crossing, 32, crossingAdjustment);
+        const auto ordinaryView = memory.Descriptor(ordinary, 64, ordinaryAdjustment);
+        auto uploaded = std::make_shared<Buffer>(context, 96, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT);
+        auto downloaded = std::make_shared<Buffer>(context, 64, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT);
+        const auto source = uploaded->Bytes();
+        std::fill(source.begin(), source.end(), std::byte{0});
+        TextureDetiler repack(context);
+        const auto commands = recorder.Commands();
+        RecordMemoryBarrier(context, commands, VK_PIPELINE_STAGE_HOST_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_HOST_WRITE_BIT | VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT);
+        repack.DispatchRepack(commands, uploaded->Handle(), 0, elementView.buffer, elementView.offset + elementAdjustment, 64);
+        repack.DispatchRepack(commands, uploaded->Handle(), 64, crossingView.buffer, crossingView.offset + crossingAdjustment, 32);
+        repack.DispatchRepack(commands, ordinaryView.buffer, ordinaryView.offset + ordinaryAdjustment, downloaded->Handle(), 0, 64);
+        RecordMemoryBarrier(context, commands, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_HOST_BIT | VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_HOST_READ_BIT | VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_TRANSFER_READ_BIT);
+        memory.RecordCopyBacks(recorder);
+        recorder.Sync();
+        memory.MarkDirectWrites();
+        memory.WriteBack();
+        const auto contents = downloaded->Bytes().first(64);
+        Require(std::all_of(contents.begin(), contents.end(), [](auto value) { return value == std::byte{0xff}; }), "a plain read lost its GPU owner when a crossing write dissolved the staged address space");
+        for (std::size_t index = 0; index < allocationBytes / 4; ++index) {
+            const auto byte = index * 4;
+            const bool written = (byte >= element - address && byte < element - address + 64) || (byte >= bytes - 16 && byte < bytes + 16);
+            const auto expected = written ? 0u : 0xffffffffu;
+            Require(words[index] == expected, "dissolving a staged address space lost a write or changed bytes outside its views");
+        }
+        std::cout << "a crossing write dissolves the staged space while keeping the heap, written views and ordinary read GPU-backed\n";
+    }
     {
         GuestAllocations::Mutation mutation;
         mutation.Remove(block);