Vulkan: retire BLASes and the TLASes built from them safely (#7937)

A BLAS was destroyed immediately when its model was unloaded or when the
vertex heap grew and the BLAS had to be rebuilt, while the TLAS of a
frame still on the GPU could reference it by device address. Queue the
acceleration structure and its buffer on the deletion queue instead, the
same way ensureTlasCapacity already retires the TLAS. The deletion queue
is flushed after ray tracing shuts down, so shutdown defers too.

A built TLAS also kept pointing at retired BLASes for later frames: the
live shadow TLAS binding is only refreshed when a mission frame rebuilds
it, and buildTlas keeps the previous TLAS when it has no instances. With
ray-traced shadows on, flying a mission and then leaving and re-entering
the tech room (which frees every model on exit) bound such a TLAS on
every tech room frame. So when BLASes are retired, the built frame TLASes
are retired with them and the binding goes back to the empty fallback
TLAS until the next rebuild. onModelUnloaded reports whether it retired
anything, so the binding refresh only runs then.

Co-authored-by: Dustinspace217 <[email protected]>

code/graphics/vulkan/VulkanRaytracing.cpp

@@ -1,5 +1,7 @@
 #include "VulkanRaytracing.h"
 
+#include "VulkanDeletionQueue.h"
+
 namespace graphics::vulkan {
 
 namespace {
@@ -222,17 +224,36 @@ uint64_t VulkanRaytracingManager::makeKey(int pm_id, int submodel_num)
 
 void VulkanRaytracingManager::destroyBlasEntry(BlasEntry& entry)
 {
-	entry.accelStruct.reset();
+	// TLASes in frames still in flight hold this BLAS's device address, so defer like ensureTlasCapacity does
+	if (entry.accelStruct) {
+		getDeletionQueue()->queueAccelerationStructure(entry.accelStruct.release());
+	}
 
 	if (entry.buffer) {
-		m_device.destroyBuffer(entry.buffer);
+		getDeletionQueue()->queueBuffer(entry.buffer, entry.allocation);
 		entry.buffer = nullptr;
-	}
-	if (entry.allocation.isValid()) {
+		entry.allocation = VulkanAllocation();
+	} else if (entry.allocation.isValid()) {
 		m_memoryManager->freeAllocation(entry.allocation);
 	}
 }
 
+void VulkanRaytracingManager::retireBuiltTlases()
+{
+	// a built TLAS may still point at a retired BLAS, and screens that never rebuild the TLAS would keep binding it
+	for (auto& frame : m_frameTlas) {
+		if (frame.tlas) {
+			getDeletionQueue()->queueAccelerationStructure(frame.tlas.release());
+		}
+		if (frame.tlasBuffer) {
+			getDeletionQueue()->queueBuffer(frame.tlasBuffer, frame.tlasAllocation);
+			frame.tlasBuffer = nullptr;
+			frame.tlasAllocation = VulkanAllocation();
+		}
+		frame.tlasCapacity = 0;
+	}
+}
+
 vk::CommandBuffer VulkanRaytracingManager::beginSingleTimeCommands()
 {
 	vk::CommandBufferAllocateInfo allocInfo;

code/graphics/vulkan/VulkanRaytracing.h

@@ -93,8 +93,9 @@ public:
 	/**
 	 * @brief Purge all cached BLASes for a model and free their GPU resources.
 	 * Called from model_unload()/model_free_all() before the model is destroyed.
+	 * @return true if any BLAS was retired (the shadow TLAS binding must then be refreshed)
 	 */
-	void onModelUnloaded(int pm_id);
+	bool onModelUnloaded(int pm_id);
 
 	/**
 	 * @brief Rebuild the top-level acceleration structure for this frame from
@@ -172,6 +173,7 @@ private:
 	const BlasEntry* getOrBuildBlasEntry(int pm_id, int submodel_num);
 
 	void destroyBlasEntry(BlasEntry& entry);
+	void retireBuiltTlases();
 
 	bool createRawBuffer(vk::DeviceSize size,
 		vk::BufferUsageFlags usage,

code/graphics/vulkan/VulkanRaytracingBlas.cpp

@@ -182,6 +182,7 @@ const VulkanRaytracingManager::BlasEntry* VulkanRaytracingManager::getOrBuildBla
 		// address) -- destroy and fall through to rebuild.
 		destroyBlasEntry(it->second);
 		m_blasCache.erase(it);
+		retireBuiltTlases();
 	}
 
 	BlasEntry entry;
@@ -236,20 +237,27 @@ void VulkanRaytracingManager::onModelLoaded(int pm_id)
 	}
 }
 
-void VulkanRaytracingManager::onModelUnloaded(int pm_id)
+bool VulkanRaytracingManager::onModelUnloaded(int pm_id)
 {
 	if (!m_enabled) {
-		return;
+		return false;
 	}
 
+	bool retired = false;
 	for (auto it = m_blasCache.begin(); it != m_blasCache.end();) {
 		if (static_cast<int>(it->first >> 32) == pm_id) {
 			destroyBlasEntry(it->second);
 			it = m_blasCache.erase(it);
+			retired = true;
 		} else {
 			++it;
 		}
 	}
+
+	if (retired) {
+		retireBuiltTlases();
+	}
+	return retired;
 }
 
 } // namespace graphics::vulkan

code/graphics/vulkan/gr_vulkan.cpp

@@ -68,7 +68,10 @@ void vulkan_model_loaded(int pm_id)
 void vulkan_model_unloaded(int pm_id)
 {
 	if (auto* rt = getRaytracingManager()) {
-		rt->onModelUnloaded(pm_id);
+		if (rt->onModelUnloaded(pm_id)) {
+			getDescriptorManager()->setCurrentShadowTlas(rt->getTlasForShaderBinding());
+			getDrawManager()->invalidateGlobalSet();
+		}
 	}
 }