fix(agc): reclaim idle buffers on allocation OOM

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

@@ -61,12 +61,13 @@ class BufferPool {
     // The size a buffer for `bytes` is created with: its size class, or `bytes` itself when large.
     static std::size_t Capacity(std::size_t bytes, VkMemoryPropertyFlags properties);
     static VkBufferUsageFlags Usage(VkBufferUsageFlags usage, VkMemoryPropertyFlags properties);
+    VkDeviceMemory AllocateMemory(const VkMemoryAllocateInfo& allocation, const char* operation);
     std::optional<BufferAllocation> Take(std::size_t bytes, VkBufferUsageFlags usage, VkMemoryPropertyFlags properties = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
     void Put(const BufferAllocation& allocation) noexcept;
     static bool SlabEligible(std::size_t capacity, VkDeviceSize alignment, VkDeviceSize size, VkDeviceSize atom);
     static VkDeviceSize SlabBlockBytes(std::size_t capacity);
     SlabSlot TakeSlot(const Context& context, std::uint32_t memoryType, std::size_t capacity, bool addressable);
-    void PutSlot(VkDeviceMemory memory, VkDeviceSize offset) noexcept;
+    VkDeviceSize PutSlot(VkDeviceMemory memory, VkDeviceSize offset) noexcept;
     std::size_t SlabBlocks();
 
 private:
@@ -117,7 +118,8 @@ class BufferPool {
     static bool DeviceTierEnabled();
     static bool DeviceTiered(VkMemoryPropertyFlags properties);
     static std::size_t NextCapacity(std::size_t capacity);
-    void destroy(const BufferAllocation& allocation) noexcept;
+    VkDeviceSize destroy(const BufferAllocation& allocation) noexcept;
+    VkDeviceSize trim();
     // Moves the tier's least recently used slot (the oldest front of its lists) to `evicted`; the
     // caller destroys those after releasing the mutex, so builds taking buffers on other threads
     // do not wait behind the Vulkan destroy calls. Nothing changes when the vector cannot grow.

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

@@ -41,6 +41,30 @@ BufferPool::~BufferPool() {
     for (auto& [memory, block] : slabBlocks) freeBlock(*block);
 }
 
+VkDeviceMemory BufferPool::AllocateMemory(const VkMemoryAllocateInfo& allocation, const char* operation) {
+    VkDeviceMemory memory = VK_NULL_HANDLE;
+    auto result = allocateMemory(device, &allocation, nullptr, &memory);
+    const auto initialResult = result;
+    unsigned attempts = 1;
+    VkDeviceSize reclaimed = 0;
+    if (result == VK_ERROR_OUT_OF_DEVICE_MEMORY || result == VK_ERROR_OUT_OF_HOST_MEMORY) {
+        reclaimed = trim();
+        if (reclaimed != 0) {
+            memory = VK_NULL_HANDLE;
+            result = allocateMemory(device, &allocation, nullptr, &memory);
+            ++attempts;
+        }
+    }
+    if (result == VK_SUCCESS) {
+        if (attempts == 2) std::fprintf(stderr, "[bufferpool] %s recovered: allocation=%llu reclaimed=%llu initialResult=%d attempts=%u\n", operation, static_cast<unsigned long long>(allocation.allocationSize), static_cast<unsigned long long>(reclaimed), static_cast<int>(initialResult), attempts);
+        return memory;
+    }
+    char details[512];
+    std::snprintf(details, sizeof(details), "%s: allocation=%llu memoryType=%u reclaimed=%llu initialResult=%d attempts=%u", operation, static_cast<unsigned long long>(allocation.allocationSize), allocation.memoryTypeIndex, static_cast<unsigned long long>(reclaimed), static_cast<int>(initialResult), attempts);
+    Check(result, details);
+    return VK_NULL_HANDLE;
+}
+
 bool BufferPool::SlabEligible(std::size_t capacity, VkDeviceSize alignment, VkDeviceSize size, VkDeviceSize atom) {
     if (capacity >= classLimit || size > capacity || !std::has_single_bit(capacity)) return false;
     const auto alignmentOk = alignment == 0 || (std::has_single_bit(alignment) && alignment <= capacity);
@@ -80,7 +104,7 @@ SlabSlot BufferPool::TakeSlot(const Context& context, std::uint32_t memoryType,
     if (addressable) allocation.pNext = &flags;
     allocation.allocationSize = blockBytes;
     allocation.memoryTypeIndex = memoryType;
-    Check(allocateMemory(device, &allocation, nullptr, &block->memory), "vkAllocateMemory buffer slab");
+    block->memory = AllocateMemory(allocation, "vkAllocateMemory buffer slab");
     if ((context.memory.memoryTypes[memoryType].propertyFlags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) != 0) {
         void* mapping = nullptr;
         const auto mapped = mapMemory(device, block->memory, 0, VK_WHOLE_SIZE, 0, &mapping);
@@ -105,26 +129,27 @@ void BufferPool::freeBlock(SlabBlock& block) noexcept {
     freeMemory(device, block.memory, nullptr);
 }
 
-void BufferPool::PutSlot(VkDeviceMemory memory, VkDeviceSize offset) noexcept {
+VkDeviceSize BufferPool::PutSlot(VkDeviceMemory memory, VkDeviceSize offset) noexcept {
     std::unique_ptr<SlabBlock> released;
     {
         std::lock_guard lock(slabMutex);
         const auto found = slabBlocks.find(memory);
-        if (found == slabBlocks.end()) return;
+        if (found == slabBlocks.end()) return 0;
         auto& block = *found->second;
         auto& slab = slabs[block.slab];
         if (block.free.empty()) slab.available.push_back(&block);
         block.free.push_back(static_cast<std::uint32_t>(offset / block.slotBytes));
-        if (--block.used != 0) return;
+        if (--block.used != 0) return 0;
         if (slab.emptyBlocks == 0) {
             ++slab.emptyBlocks;
-            return;
+            return 0;
         }
         slab.available.erase(std::find(slab.available.begin(), slab.available.end(), &block));
         released = std::move(found->second);
         slabBlocks.erase(found);
     }
     freeBlock(*released);
+    return SlabBlockBytes(released->slotBytes);
 }
 
 std::size_t BufferPool::SlabBlocks() {
@@ -150,16 +175,16 @@ VkDeviceSize BufferPool::DeviceBudget(const VkPhysicalDeviceMemoryProperties& me
     return std::max<VkDeviceSize>(VkDeviceSize{512} << 20u, largest / 8u);
 }
 
-void BufferPool::destroy(const BufferAllocation& allocation) noexcept {
+VkDeviceSize BufferPool::destroy(const BufferAllocation& allocation) noexcept {
     if (allocation.slab) {
         destroyBuffer(device, allocation.buffer, nullptr);
-        PutSlot(allocation.memory, allocation.offset);
-        return;
+        return PutSlot(allocation.memory, allocation.offset);
     }
     // Device-local allocations (see DeviceBuffer) are never mapped.
     if (allocation.mapping != nullptr) unmap(device, allocation.memory);
     destroyBuffer(device, allocation.buffer, nullptr);
     freeMemory(device, allocation.memory, nullptr);
+    return allocation.allocationBytes;
 }
 
 bool BufferPool::DeviceTiered(VkMemoryPropertyFlags properties) {
@@ -240,6 +265,43 @@ std::size_t BufferPool::MaxSlots() {
     return slots;
 }
 
+VkDeviceSize BufferPool::trim() {
+    Tier retiredSmall;
+    Tier retiredLarge;
+    Tier retiredDevice;
+    {
+        std::lock_guard lock(mutex);
+        for (const auto& [tier, retired] : {std::pair{&smallTier, &retiredSmall}, std::pair{&largeTier, &retiredLarge}, std::pair{&deviceTier, &retiredDevice}}) {
+            tier->free.swap(retired->free);
+            tier->evictions += tier->slots;
+            tier->retainedBytes = 0;
+            tier->slots = 0;
+        }
+    }
+    VkDeviceSize bytes = 0;
+    for (const auto* tier : {&retiredSmall, &retiredLarge, &retiredDevice}) {
+        for (const auto& [key, slots] : tier->free) {
+            for (const auto& slot : slots) bytes += destroy(slot.allocation);
+        }
+    }
+    for (;;) {
+        std::unique_ptr<SlabBlock> released;
+        {
+            std::lock_guard lock(slabMutex);
+            const auto found = std::find_if(slabBlocks.begin(), slabBlocks.end(), [](const auto& entry) { return entry.second->used == 0; });
+            if (found == slabBlocks.end()) break;
+            auto& slab = slabs.at(found->second->slab);
+            slab.available.erase(std::find(slab.available.begin(), slab.available.end(), found->second.get()));
+            --slab.emptyBlocks;
+            released = std::move(found->second);
+            slabBlocks.erase(found);
+        }
+        freeBlock(*released);
+        bytes += SlabBlockBytes(released->slotBytes);
+    }
+    return bytes;
+}
+
 void BufferPool::Put(const BufferAllocation& allocation) noexcept {
     // Evicted allocations are destroyed after the mutex is released (see evictOldest). The vector
     // may throw on growth; a Put that cannot retain simply destroys, as noexcept requires.

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

@@ -61,7 +61,7 @@ Buffer::Buffer(const Context& context, std::size_t size, VkBufferUsageFlags usag
             ready = true;
             return;
         }
-        Check(context.Function<PFN_vkAllocateMemory>("vkAllocateMemory")(context.device, &allocation, nullptr, &memory), "vkAllocateMemory buffer");
+        memory = cache->AllocateMemory(allocation, "vkAllocateMemory buffer");
         Check(context.Function<PFN_vkBindBufferMemory>("vkBindBufferMemory")(context.device, buffer, memory, 0), "vkBindBufferMemory");
         initializeAddress(usage);
         if ((properties & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) != 0) Check(context.Function<PFN_vkMapMemory>("vkMapMemory")(context.device, memory, 0, VK_WHOLE_SIZE, 0, &mapping), "vkMapMemory");
@@ -131,7 +131,7 @@ DeviceBuffer::DeviceBuffer(const Context& context, std::size_t size, VkBufferUsa
         allocation.allocationSize = requirements.size;
         allocationBytes = requirements.size;
         allocation.memoryTypeIndex = context.MemoryType(requirements.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
-        Check(context.Function<PFN_vkAllocateMemory>("vkAllocateMemory")(context.device, &allocation, nullptr, &memory), "vkAllocateMemory device buffer");
+        memory = cache->AllocateMemory(allocation, "vkAllocateMemory device buffer");
         Check(context.Function<PFN_vkBindBufferMemory>("vkBindBufferMemory")(context.device, buffer, memory, 0), "vkBindBufferMemory device");
     } catch (...) {
         release();