renderer: rewrite buffer lookup around page index

src/graphics/host_gpu/renderer/cache/bufferCache.cpp

@@ -56,11 +56,32 @@ struct BufferCache::DownloadRange {
 	uint64_t offset  = 0;
 };
 
+void BufferCache::Register(BufferId id) {
+	const auto& buffer = m_slot_buffers[id];
+	PageTable::PageRange pages {};
+	EXIT_IF(!PageTable::TryGetPageRange(buffer.CpuAddress(), buffer.Size(), pages));
+	for (size_t page = pages.first; page < pages.last_exclusive; ++page) {
+		m_page_table[page] = id;
+	}
+	const auto [it, inserted] = m_buffers.emplace(buffer.CpuAddress(), id);
+	(void)it;
+	EXIT_IF(!inserted);
+	m_total_used_memory += buffer.Size();
+}
+
 void BufferCache::Unregister(BufferId id) {
 	auto& buffer = m_slot_buffers[id];
 	if (buffer.is_deleted) {
 		return;
 	}
+	PageTable::PageRange pages {};
+	EXIT_IF(!PageTable::TryGetPageRange(buffer.CpuAddress(), buffer.Size(), pages));
+	for (size_t page = pages.first; page < pages.last_exclusive; ++page) {
+		auto* owner = m_page_table.Find(page);
+		if (owner != nullptr && *owner == id) {
+			*owner = {};
+		}
+	}
 	const auto found = m_buffers.find(buffer.CpuAddress());
 	EXIT_IF(found == m_buffers.end() || found->second != id);
 	m_buffers.erase(found);
@@ -252,15 +273,15 @@ void BufferCache::ReadMemoryOnGpu(uint64_t vaddr, uint64_t size, bool is_write)
 		    m_memory_tracker.ValidateGpuDirtyPages(m_gpu_modified_ranges, address, bytes,
 		                                           "memory invalidation");
 	    },
-	    [&](uint64_t address, uint64_t bytes) noexcept {
+		[&](uint64_t address, uint64_t bytes) noexcept {
 		    for (const auto range: m_gpu_modified_ranges.Intersections(address, bytes)) {
 			    for (uint64_t copied = 0; copied < range.size;) {
 				    const auto copy_address = range.address + copied;
-				    auto       owner        = m_buffers.upper_bound(copy_address);
+				    const auto* owner = m_page_table.Find(copy_address >> PageTable::kPageBits);
-				    if (owner == m_buffers.begin()) {
+				    if (owner == nullptr || !*owner) {
 					    EXIT("BufferCache: invalidation readback has no buffer owner\n");
 				    }
-				    auto& buffer = m_slot_buffers[std::prev(owner)->second];
+				    auto& buffer = m_slot_buffers[*owner];
 				    if (!buffer.IsInBounds(copy_address, 1)) {
 					    EXIT("BufferCache: invalidation readback is outside its buffer owner\n");
 				    }
@@ -369,45 +390,35 @@ BufferId BufferCache::FindBuffer(uint64_t vaddr, uint64_t size) {
 	if (size == 0 || vaddr >= TRACKER_ADDRESS_SIZE || size > TRACKER_ADDRESS_SIZE - vaddr) {
 		EXIT("BufferCache: invalid buffer discovery request\n");
 	}
+	const auto* owner = m_page_table.Find(vaddr >> PageTable::kPageBits);
+	if (owner != nullptr && *owner) {
+		auto& buffer = m_slot_buffers[*owner];
+		if (buffer.IsInBounds(vaddr, size)) {
+			buffer.tick_accessed_last = m_gc_tick;
+			return *owner;
+		}
+	}
+	return CreateBuffer(vaddr, size);
+}
+
+BufferId BufferCache::CreateBuffer(uint64_t vaddr, uint64_t size) {
 	auto& command = m_scheduler.Current();
 	EXIT_IF(command.IsInvalid() || command.IsExecute());
 	auto       begin = vaddr & ~(CACHING_PAGE_SIZE - 1);
 	auto       end   = (vaddr + size + CACHING_PAGE_SIZE - 1) & ~(CACHING_PAGE_SIZE - 1);
-	const auto next  = m_buffers.upper_bound(vaddr);
+	auto       first = m_buffers.lower_bound(begin);
-	if (next != m_buffers.begin()) {
-		const auto owner  = std::prev(next);
-		auto&      buffer = m_slot_buffers[owner->second];
-		if (buffer.IsInBounds(vaddr, size)) {
-			buffer.tick_accessed_last = m_gc_tick;
-			return owner->second;
-		}
-	}
-
-	auto merge = [&](const Buffer& buffer) {
-		const auto buffer_begin = buffer.CpuAddress();
-		const auto buffer_end   = buffer_begin + buffer.Size();
-		if (begin >= buffer_end || buffer_begin >= end) {
-			return false;
-		}
-		begin = std::min(begin, buffer_begin);
-		end   = std::max(end, buffer_end);
-		return true;
-	};
-	std::vector<BufferId> overlaps;
-	auto                  first = m_buffers.lower_bound(begin);
 	if (first != m_buffers.begin()) {
 		const auto previous = std::prev(first);
-		if (merge(m_slot_buffers[previous->second])) {
+		if (const auto& buffer = m_slot_buffers[previous->second];
+		    buffer.CpuAddress() + buffer.Size() > begin) {
 			first = previous;
 		}
 	}
-	for (auto candidate = first; candidate != m_buffers.end(); ++candidate) {
+	auto last = first;
-		if (candidate->first >= end) {
+	for (; last != m_buffers.end() && last->first < end; ++last) {
-			break;
+		const auto& buffer = m_slot_buffers[last->second];
-		}
+		begin              = std::min(begin, buffer.CpuAddress());
-		if (merge(m_slot_buffers[candidate->second])) {
+		end                = std::max(end, buffer.CpuAddress() + buffer.Size());
-			overlaps.push_back(candidate->second);
-		}
 	}
 
 	const auto id = m_slot_buffers.insert(m_graphics, m_scheduler, MemoryUsage::DeviceLocal, begin,
@@ -416,17 +427,14 @@ BufferId BufferCache::FindBuffer(uint64_t vaddr, uint64_t size) {
 	buffer.tick_accessed_last = m_gc_tick;
 	SetVulkanObjectNameF(m_graphics.device, buffer.Handle(),
 	                     "Kyty.GameBuffer[guest=0x{:016x} size=0x{:x}]", begin, end - begin);
-	for (const auto overlap: overlaps) {
+	for (auto overlap = first; overlap != last;) {
-		const auto& old = m_slot_buffers[overlap];
+		const auto current = overlap++;
+		const auto old_id  = current->second;
+		const auto& old    = m_slot_buffers[old_id];
 		buffer.CopyFrom(command, old, 0, old.CpuAddress() - begin, old.Size());
+		DeleteBuffer(old_id);
 	}
-	for (const auto overlap: overlaps) {
+	Register(id);
-		DeleteBuffer(overlap);
-	}
-	m_total_used_memory += buffer.Size();
-	const auto [it, inserted] = m_buffers.emplace(begin, id);
-	(void)it;
-	EXIT_IF(!inserted);
 	return id;
 }
 
@@ -544,13 +552,12 @@ std::pair<Buffer*, uint64_t> BufferCache::ObtainBufferForImage(uint64_t vaddr, u
 		EXIT("BufferCache: invalid image source\n");
 	}
 	auto find_owner = [&]() -> Buffer* {
-		auto owner = m_buffers.upper_bound(vaddr);
+		const auto* owner = m_page_table.Find(vaddr >> PageTable::kPageBits);
-		if (owner == m_buffers.begin()) {
+		if (owner == nullptr || !*owner) {
 			return nullptr;
 		}
-		--owner;
+		auto& buffer = m_slot_buffers[*owner];
-		auto* buffer = m_slot_buffers.try_get(owner->second);
+		return buffer.IsInBounds(vaddr, size) ? &buffer : nullptr;
-		return buffer != nullptr && buffer->IsInBounds(vaddr, size) ? buffer : nullptr;
 	};
 
 	{
@@ -565,14 +572,7 @@ std::pair<Buffer*, uint64_t> BufferCache::ObtainBufferForImage(uint64_t vaddr, u
 			if (!IsRegionRegistered(vaddr, size)) {
 				EXIT("BufferCache: GPU-dirty image source has no native buffer\n");
 			}
-			const auto id       = FindBuffer(vaddr, size);
+			owner = &m_slot_buffers[FindBuffer(vaddr, size)];
-			auto&      buffer   = m_slot_buffers[id];
-			owner               = &buffer;
-			const auto resolved = m_buffers.find(buffer.CpuAddress());
-			if (resolved == m_buffers.end() || resolved->second != id ||
-			    !buffer.IsInBounds(vaddr, size)) {
-				EXIT("BufferCache: merged image source does not contain the requested range\n");
-			}
 		}
 		if (owner != nullptr && !cpu_modified && (!gpu_modified || has_dirty_buffer_source)) {
 			owner->tick_accessed_last = m_gc_tick;

src/graphics/host_gpu/renderer/cache/bufferCache.h

@@ -6,6 +6,7 @@
 #include "common/slotVector.h"
 #include "graphics/host_gpu/memoryTracker.h"
 #include "graphics/host_gpu/rangeSet.h"
+#include "graphics/host_gpu/renderer/cache/multiLevelPageTable.h"
 #include "graphics/host_gpu/renderer/cache/streamBuffer.h"
 
 #include <map>
@@ -66,12 +67,16 @@ private:
 
 	struct DownloadCopy;
 	struct DownloadRange;
+	using PageTable = MultiLevelPageTable<BufferId, 14, 40, 16>;
+	static_assert(CACHING_PAGE_SIZE == (uint64_t {1} << PageTable::kPageBits));
 	static constexpr uint64_t               DOWNLOAD_ALIGNMENT = 64;
 	[[nodiscard]] static constexpr uint64_t AlignDownload(uint64_t size) noexcept {
 		return (size + DOWNLOAD_ALIGNMENT - 1) & ~(DOWNLOAD_ALIGNMENT - 1);
 	}
 	[[nodiscard]] static std::pair<uint64_t, uint64_t> DownloadEnvelope(const DownloadCopy& copy);
 	void WriteDataBuffer(Buffer& buffer, uint64_t address, const void* source, uint64_t size);
+	[[nodiscard]] BufferId CreateBuffer(uint64_t vaddr, uint64_t size);
+	void                   Register(BufferId id);
 	void Unregister(BufferId id);
 	void DeleteBuffer(BufferId id);
 	[[nodiscard]] bool SynchronizeBuffer(Buffer& buffer, uint64_t vaddr, uint64_t size,
@@ -91,6 +96,7 @@ private:
 	Buffer                       m_gds_buffer;
 	Common::SlotVector<Buffer>   m_slot_buffers;
 	std::map<uint64_t, BufferId> m_buffers;
+	PageTable                    m_page_table;
 	RangeSet                     m_gpu_modified_ranges;
 	MemoryTracker                m_memory_tracker;
 	StreamBuffer                 m_staging_buffer;

tests/ShaderRecompilerComputeTests.cpp

@@ -155,6 +155,12 @@ struct BufferCacheTestAccess {
                                          uint64_t address, uint64_t size) {
     return cache.SynchronizeBufferFromImage(buffer, address, size);
   }
+
+  static BufferId PageOwner(const BufferCache &cache, uint64_t address) {
+    const auto *owner = cache.m_page_table.Find(
+        address >> BufferCache::PageTable::kPageBits);
+    return owner == nullptr ? BufferId{} : *owner;
+  }
 };
 
 struct StreamBufferTestAccess {
@@ -2962,13 +2968,23 @@ public:
       constexpr uint64_t index_span = 3 * index_page;
       const uint64_t index_begin = base + index_offset;
       Require(name, "registered-range empty lookup",
-              !cache.IsRegionRegistered(index_begin, index_span),
+              !cache.IsRegionRegistered(index_begin, index_span) &&
+                  !BufferCacheTestAccess::PageOwner(cache, index_begin),
               "empty Buffer cache reported a registered range");
       const auto index_left =
           cache.FindBuffer(index_begin + 0x100, sizeof(uint32_t));
       const auto index_right = cache.FindBuffer(
           index_begin + 2 * index_page + 0x100, sizeof(uint32_t));
-      Require(name, "registered-range owners", index_left && index_right,
+      Require(name, "registered-range owners",
+              index_left && index_right &&
+                  cache.FindBuffer(index_begin + 0x200, sizeof(uint32_t)) ==
+                      index_left &&
+                  BufferCacheTestAccess::PageOwner(cache, index_begin) ==
+                      index_left &&
+                  !BufferCacheTestAccess::PageOwner(cache,
+                                                    index_begin + index_page) &&
+                  BufferCacheTestAccess::PageOwner(
+                      cache, index_begin + 2 * index_page) == index_right,
               "failed to create disjoint Buffer index owners");
       Require(
           name, "registered-range boundaries",
@@ -2986,8 +3002,23 @@ public:
           cache.FindBuffer(index_begin + index_page - 1, index_page + 2);
       Require(name, "registered-range merge",
               index_bridge && cache.IsRegionRegistered(index_begin + index_page,
-                                                       index_page),
+                                                       index_page) &&
+                  BufferCacheTestAccess::PageOwner(cache, index_begin) ==
+                      index_bridge &&
+                  BufferCacheTestAccess::PageOwner(
+                      cache, index_begin + index_page) == index_bridge &&
+                  BufferCacheTestAccess::PageOwner(
+                      cache, index_begin + 2 * index_page) == index_bridge,
               "bridging acquisition did not publish its merged Buffer range");
+      scheduler.FinishCurrent();
+      constexpr uint64_t recycled_offset = index_offset + 4 * index_page;
+      const auto recycled = cache.FindBuffer(base + recycled_offset + 0x100,
+                                             sizeof(uint32_t));
+      Require(name, "slot generation reuse",
+              (recycled.index == index_left.index && recycled != index_left) ||
+                  (recycled.index == index_right.index &&
+                   recycled != index_right),
+              "retired Buffer slot was not reused with a new generation");
       const auto [resolved_left, resolved_left_offset] = cache.ObtainBuffer(
           index_begin + 0x100, sizeof(uint32_t), true, false, index_left);
       const auto [resolved_right, resolved_right_offset] =
@@ -2999,8 +3030,16 @@ public:
                   resolved_left_offset != resolved_right_offset,
               "saved descriptor IDs did not resolve through the merged owner");
       cache.UnmapMemory(index_begin, index_span);
+      cache.UnmapMemory(base + recycled_offset, index_page);
       Require(name, "registered-range removal",
-              !cache.IsRegionRegistered(index_begin, index_span),
+              !cache.IsRegionRegistered(index_begin, index_span) &&
+                  !BufferCacheTestAccess::PageOwner(cache, index_begin) &&
+                  !BufferCacheTestAccess::PageOwner(
+                      cache, index_begin + index_page) &&
+                  !BufferCacheTestAccess::PageOwner(
+                      cache, index_begin + 2 * index_page) &&
+                  !BufferCacheTestAccess::PageOwner(cache,
+                                                    base + recycled_offset),
               "unmapped Buffer owner remained in the registered-range index");
 
       MarkGpuWrite(base + first_offset, sizeof(first_value));