graphics: resize image heap tables in place when their descriptor size changes

src/graphics/host_gpu/renderer/pipeline/imageHeap.cpp

@@ -12,6 +12,7 @@
 #include "kernel/memory.h"
 
 #include <algorithm>
+#include <bit>
 #include <xxhash.h>
 
 namespace Libs::Graphics {
@@ -44,6 +45,20 @@ void Widen(size_t& begin, size_t& end, size_t first, size_t last) {
 	end              = empty ? last : std::max(end, last);
 }
 
+// Probes at first + k * step whose first dword lies within a table of size bytes.
+uint64_t Probes(uint64_t size, uint64_t first, uint64_t step) {
+	return size >= first + 4u ? (size - first - 4u) / step + 1u : 0u;
+}
+
+// Tables grow a little on most frames, so their slot words get room to double.
+std::unique_ptr<Buffer> SlotBuffer(RenderContext& context, uint64_t size) {
+	return std::make_unique<Buffer>(
+	    context.GetGraphics(), context.GetCommandScheduler(), MemoryUsage::DeviceLocal, 0,
+	    vk::BufferUsageFlagBits::eTransferDst | vk::BufferUsageFlagBits::eStorageBuffer |
+	        vk::BufferUsageFlagBits::eShaderDeviceAddress,
+	    std::bit_ceil(size + 4u));
+}
+
 // Guest bytes of a table that its tracker pages [begin, end) cover.
 std::pair<uint64_t, uint64_t> PageRange(uint64_t base, uint64_t size, size_t begin, size_t end) {
 	const auto origin = Common::AlignDown(base, TRACKER_PAGE_SIZE);
@@ -122,22 +137,17 @@ bool ImageHeap::Acquire(const IR::ImageHeapRequest& request, IR::ImageHeapView&
 	    ((request.base | request.first | request.step) & 3u) != 0u) {
 		return false;
 	}
-	auto memory = std::ranges::find_if(m_memories, [&](const auto& candidate) {
+	auto memory = std::ranges::find_if(
-		return candidate->base == request.base && candidate->size == size;
+	    m_memories, [&](const auto& candidate) { return candidate->base == request.base; });
-	});
+	// A table follows its descriptor's size, as the first program of a draw acquires it.
-	if (memory == m_memories.end()) {
+	if (memory == m_memories.end() || ((*memory)->size != size && (*memory)->draw != m_draw)) {
 		if (!m_context.IsMapped(request.base, size)) return false;
-		auto created   = std::make_unique<Memory>();
+		if (memory == m_memories.end()) {
-		created->base  = request.base;
+			std::scoped_lock lock {m_mutex};
-		created->size  = size;
+			memory = m_memories.insert(m_memories.end(),
-		created->words.resize(size / 4u);
+			                           std::make_unique<Memory>(Memory {.base = request.base}));
-		created->dead.resize(size / 4u);
+		}
-		created->dirty_pages.assign((Common::AlignUp(request.base + size, TRACKER_PAGE_SIZE) -
+		Resize(**memory, size);
-		                             Common::AlignDown(request.base, TRACKER_PAGE_SIZE)) /
-		                                TRACKER_PAGE_SIZE,
-		                            uint8_t {1});
-		std::scoped_lock lock {m_mutex};
-		memory = m_memories.insert(m_memories.end(), std::move(created));
 	}
 	// A draw refreshes a table and a view takes its changes once, so all programs of the draw see
 	// one snapshot and one set of kinds.
@@ -146,27 +156,28 @@ bool ImageHeap::Acquire(const IR::ImageHeapRequest& request, IR::ImageHeapView&
 		(*memory)->draw = m_draw;
 	}
 	auto&      views   = (*memory)->views;
+	const auto bytes   = (*memory)->size;
 	const auto profile = IR::ImageHeapResource(*request.root, 0u);
+	const auto probes  = Probes(size, request.first, request.step);
+	const auto limit   = request.count < probes ? request.count : UINT64_MAX;
 	auto       view    = std::ranges::find_if(views, [&](const auto& candidate) {
 		return candidate->step == request.step && candidate->first == request.first &&
-		       candidate->count == request.count && candidate->profile == profile;
+		       candidate->limit == limit && candidate->profile == profile;
 	});
 	if (view == views.end()) {
 		const auto epoch   = Image::heap_epoch.load(std::memory_order_relaxed);
+		const auto count   = std::min(limit, Probes(bytes, request.first, request.step));
 		auto       created = std::make_unique<View>(View {.memory        = **memory,
 		                                                  .step          = request.step,
 		                                                  .first         = request.first,
-		                                                  .count         = request.count,
+		                                                  .count         = count,
+		                                                  .limit         = limit,
 		                                                  .profile       = profile,
 		                                                  .resolve_epoch = epoch,
 		                                                  .layout_epoch  = epoch});
-		created->slots.resize(size / 4u + 1u);
+		created->slots.resize(bytes / 4u + 1u);
-		created->buffer = std::make_unique<Buffer>(
+		created->buffer = SlotBuffer(m_context, bytes);
-		    m_context.GetGraphics(), m_context.GetCommandScheduler(), MemoryUsage::DeviceLocal, 0,
+		UpdateProbes(*created, 0u, bytes);
-		    vk::BufferUsageFlagBits::eTransferDst | vk::BufferUsageFlagBits::eStorageBuffer |
-		        vk::BufferUsageFlagBits::eShaderDeviceAddress,
-		    size + 4u);
-		UpdateProbes(*created, 0u, size);
 		// The first upload also clears the slots no probe selects.
 		created->dirty_begin = 0;
 		created->dirty_end   = created->slots.size();
@@ -388,6 +399,43 @@ void ImageHeap::Retire(View& view) {
 	view.released.clear();
 }
 
+void ImageHeap::Resize(Memory& memory, uint64_t size) {
+	{
+		std::scoped_lock lock {m_mutex};
+		Unwatch(memory);
+		memory.size = size;
+		// The next refresh rereads the whole table.
+		memory.dirty_pages.assign((Common::AlignUp(memory.base + size, TRACKER_PAGE_SIZE) -
+		                           Common::AlignDown(memory.base, TRACKER_PAGE_SIZE)) /
+		                              TRACKER_PAGE_SIZE,
+		                          uint8_t {1});
+	}
+	memory.words.resize(size / 4u);
+	memory.dead.resize(size / 4u);
+	auto& scheduler = m_context.GetCommandScheduler();
+	for (auto& view: memory.views) {
+		const auto count = std::min(view->limit, Probes(size, view->first, view->step));
+		for (auto probe = count; probe < view->count; ++probe) {
+			Assign(*view, (view->first + probe * view->step) / 4u, 0u);
+		}
+		view->count       = count;
+		const auto slots  = view->slots.size();
+		view->slots.resize(size / 4u + 1u);
+		view->dirty_end   = std::min(view->dirty_end, view->slots.size());
+		view->dirty_begin = std::min(view->dirty_begin, view->dirty_end);
+		if (view->slots.size() * sizeof(uint32_t) > view->buffer->Size()) {
+			if (scheduler.Active()) {
+				scheduler.DeferOperation([buffer = std::move(view->buffer)] {});
+			}
+			view->buffer = SlotBuffer(m_context, size);
+			Widen(view->dirty_begin, view->dirty_end, 0u, view->slots.size());
+		} else if (slots < view->slots.size()) {
+			// Device words past the old end may hold stale slots.
+			Widen(view->dirty_begin, view->dirty_end, slots, view->slots.size());
+		}
+	}
+}
+
 void ImageHeap::Rebind(RenderExecutor& executor) {
 	KYTY_PROFILER_FUNCTION();
 	auto&      textures = m_context.GetTextureCache();
@@ -526,46 +574,26 @@ void ImageHeap::UnmapMemory(uint64_t address, uint64_t size) {
 	} else {
 		m_unmapped.push_back({address, size});
 	}
+	auto&           scheduler = m_context.GetCommandScheduler();
 	std::scoped_lock lock {m_mutex};
 	std::erase_if(m_memories, [&](const std::unique_ptr<Memory>& memory) {
 		if (address >= memory->base + memory->size || memory->base >= address + size) {
 			return false;
 		}
-		Drop(*memory);
+		Unwatch(*memory);
+		for (auto& view: memory->views) {
+			std::erase(m_active, view.get());
+			for (const auto& [value, index]: view->elements) {
+				Unreference(m_elements[index]);
+				view->released.push_back(index);
+			}
+			Retire(*view);
+			if (scheduler.Active()) {
+				scheduler.DeferOperation([buffer = std::move(view->buffer)] {});
+			}
+		}
 		return true;
 	});
 }
 
-void ImageHeap::RunGarbageCollector() {
-	// A table whose descriptor grew keeps its old snapshot until a frame passes without it.
-	std::scoped_lock      lock {m_mutex};
-	std::vector<uint64_t> acquired;
-	for (const auto& memory: m_memories) {
-		if (memory->draw > m_frame_draw) acquired.push_back(memory->base);
-	}
-	std::erase_if(m_memories, [&](const std::unique_ptr<Memory>& memory) {
-		const bool superseded = memory->draw <= m_frame_draw &&
-		                        std::ranges::find(acquired, memory->base) != acquired.end();
-		if (superseded) Drop(*memory);
-		return superseded;
-	});
-	m_frame_draw = m_draw;
-}
-
-void ImageHeap::Drop(Memory& memory) {
-	auto& scheduler = m_context.GetCommandScheduler();
-	Unwatch(memory);
-	for (auto& view: memory.views) {
-		std::erase(m_active, view.get());
-		for (const auto& [value, index]: view->elements) {
-			Unreference(m_elements[index]);
-			view->released.push_back(index);
-		}
-		Retire(*view);
-		if (scheduler.Active()) {
-			scheduler.DeferOperation([buffer = std::move(view->buffer)] {});
-		}
-	}
-}
-
 } // namespace Libs::Graphics

src/graphics/host_gpu/renderer/pipeline/imageHeap.h

@@ -22,10 +22,11 @@ class RenderExecutor;
 struct ImageHeapTestAccess;
 
 // Guest image tables that shaders index by GPU-computed offsets, kept resident in one bindless
-// sampled-image array. Each view of a table owns device slot words, one per table dword, that
+// sampled-image array. A table is keyed by its base and follows its descriptor's size. Each view
-// select an array element. CPU writes to a table fault in its snapshot, which each draw refreshes
+// of a table owns device slot words, one per table dword, that select an array element. CPU
-// once; a view takes the changes when a draw first acquires it. Changes to the images behind
+// writes to a table fault in its snapshot, which each draw refreshes once; a view takes the
-// elements bump Image::heap_epoch.
+// changes when a draw first acquires it. Changes to the images behind elements bump
+// Image::heap_epoch.
 class ImageHeap final: public ShaderRecompiler::IR::ImageHeapSource {
 public:
 	ImageHeap(RenderContext& context, PageManager& page_manager);
@@ -46,8 +47,6 @@ public:
 	// CPU write fault, on any thread.
 	void InvalidateMemory(uint64_t address, uint64_t size);
 	void UnmapMemory(uint64_t address, uint64_t size);
-	// Frame boundary: drops table snapshots that another at the same base replaced this frame.
-	void RunGarbageCollector();
 
 private:
 	using DescriptorValue = ShaderRecompiler::IR::DescriptorValue;
@@ -72,6 +71,7 @@ private:
 		uint64_t                            step  = 0;
 		uint64_t                            first = 0;
 		uint64_t                            count = 0;
+		uint64_t                            limit = 0; // Probes a key domain allows, if fewer.
 		ShaderRecompiler::IR::ImageResource profile;
 		std::vector<uint32_t>               slots; // Plus a trailing null slot.
 		std::vector<uint32_t>               kind_keys;
@@ -115,7 +115,7 @@ private:
 	Image*   ImageOf(const Element& element) const;
 	uint32_t Allocate();
 	void     Retire(View& view);
-	void     Drop(Memory& memory);
+	void     Resize(Memory& memory, uint64_t size);
 	void     WriteDescriptors();
 
 	RenderContext&                       m_context;
@@ -123,8 +123,7 @@ private:
 	std::mutex                           m_mutex; // Fault-visible Memory state and m_memories.
 	std::vector<std::unique_ptr<Memory>> m_memories;
 	std::vector<View*>                   m_active; // Acquired since the last commit.
-	uint64_t                             m_draw       = 1;
+	uint64_t                             m_draw = 1;
-	uint64_t                             m_frame_draw = 0; // Last draw before this frame.
 	std::vector<Element>                 m_elements;
 	std::vector<uint32_t>                m_free;
 	std::deque<std::pair<uint64_t, uint32_t>> m_retired; // Tick, element.

src/graphics/host_gpu/renderer/renderContext.cpp

@@ -148,7 +148,6 @@ void RenderContext::RunGarbageCollector() {
 		m_buffer_cache.ProcessFaultBuffer();
 	}
 	m_texture_cache.ProcessDownloadImages();
-	m_image_heap.RunGarbageCollector();
 	m_texture_cache.RunGarbageCollector();
 	m_buffer_cache.RunGarbageCollector();
 }

tests/ShaderRecompilerComputeTests.cpp

@@ -5511,16 +5511,33 @@ public:
                   ImageHeapTestAccess::ImageOf(heap, filtered.slots[240]),
               "a reactivated view kept an element whose image slot holds a new image");
 
-      // A table whose descriptor grew keeps its old snapshot until a frame passes without it.
+      // A table follows its descriptor's size and keeps its elements. Growing probes the new
+      // words, shrinking releases them, and a draw sees one size.
+      write(table_size / 32u, descriptor(40, Format::k8_8_8_8UNorm));
+      const auto elements = ImageHeapTestAccess::ElementCount(heap);
       const IR::ImageHeapRequest grown{table, table_size * 2u, 32u, 0u, table_size / 16u, &sampled};
-      for (uint32_t frame = 0; frame < 2u; ++frame) {
+      Require(name, "grown table",
-        Require(name, "grown table", heap.Acquire(grown, view), "a grown table was declined");
+              heap.Acquire(grown, view) && view.slot_count == table_size / 2u &&
-        draw();
+                  ImageHeapTestAccess::MemoryCount(heap) == 1u,
-        heap.RunGarbageCollector();
+              "a grown table was declined or kept apart from its old size");
-        heap.BeginDraw();
+      draw();
-        Require(name, "superseded table", ImageHeapTestAccess::MemoryCount(heap) == 2u - frame,
+      const auto added = ImageHeapTestAccess::ImageOf(heap, filtered.slots[table_size / 4u]);
-                "a snapshot was dropped in a frame that acquired it or kept after one without it");
+      Require(name, "grown descriptor",
-      }
+              added && ImageHeapTestAccess::ElementCount(heap) <= elements + 1u,
+              "a grown table missed its new descriptor or interned its old ones again");
+      Require(name, "shrunk table",
+              acquire(sampled) && view.slot_count == table_size / 4u &&
+                  filtered.slots.back() == 0u && textures.GetImage(added).heap_refs.load() == 0u,
+              "a shrunk table kept a descriptor past its end");
+      Require(name, "one size per draw",
+              heap.Acquire(grown, view) && view.slot_count == table_size / 4u,
+              "a draw saw two sizes of one table");
+      draw();
+      Require(name, "regrown table", heap.Acquire(grown, view), "a regrown table was declined");
+      draw();
+      Require(name, "regrown descriptor",
+              ImageHeapTestAccess::ImageOf(heap, filtered.slots[table_size / 4u]) == added,
+              "a table grown again after shrinking missed its descriptor");
 
       heap.UnmapMemory(table, table_size);
       Require(name, "unmapped table",