From c1d875d841ad2692b7792bd10813bcf253723d27 Mon Sep 17 00:00:00 2001
From: kd-11 <karokidii@gmail.com>
Date: Mon, 12 Jun 2023 23:31:13 +0300
Subject: [PATCH] vk: Handle VK_ERROR_FRAGMENTATION when allocating descriptor
 pools

---
 rpcs3/Emu/RSX/VK/VKGSRender.cpp          |  6 ++++++
 rpcs3/Emu/RSX/VK/VKGSRender.h            |  3 +++
 rpcs3/Emu/RSX/VK/VKHelpers.cpp           |  8 ++++++++
 rpcs3/Emu/RSX/VK/vkutils/descriptors.cpp | 24 +++++++++++++++++++++---
 rpcs3/Emu/RSX/VK/vkutils/shared.cpp      |  3 +++
 5 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/rpcs3/Emu/RSX/VK/VKGSRender.cpp b/rpcs3/Emu/RSX/VK/VKGSRender.cpp
index 1470c1e6e9..e337165980 100644
--- a/rpcs3/Emu/RSX/VK/VKGSRender.cpp
+++ b/rpcs3/Emu/RSX/VK/VKGSRender.cpp
@@ -1202,6 +1202,12 @@ bool VKGSRender::on_vram_exhausted(rsx::problem_severity severity)
 	return any_cache_relieved;
 }
 
+void VKGSRender::on_descriptor_pool_fragmentation()
+{
+	// Just flush everything. Unless the hardware is very deficient, this should happen very rarely.
+	flush_command_queue(true, true);
+}
+
 void VKGSRender::notify_tile_unbound(u32 tile)
 {
 	//TODO: Handle texture writeback
diff --git a/rpcs3/Emu/RSX/VK/VKGSRender.h b/rpcs3/Emu/RSX/VK/VKGSRender.h
index 44be389ab8..b8de04eb23 100644
--- a/rpcs3/Emu/RSX/VK/VKGSRender.h
+++ b/rpcs3/Emu/RSX/VK/VKGSRender.h
@@ -263,6 +263,9 @@ public:
 	// External callback to handle out of video memory problems
 	bool on_vram_exhausted(rsx::problem_severity severity);
 
+	// Handle pool creation failure due to fragmentation
+	void on_descriptor_pool_fragmentation();
+
 	// Conditional rendering
 	void begin_conditional_rendering(const std::vector<rsx::reports::occlusion_query_info*>& sources) override;
 	void end_conditional_rendering() override;
diff --git a/rpcs3/Emu/RSX/VK/VKHelpers.cpp b/rpcs3/Emu/RSX/VK/VKHelpers.cpp
index 78b297b694..f6d3ea417e 100644
--- a/rpcs3/Emu/RSX/VK/VKHelpers.cpp
+++ b/rpcs3/Emu/RSX/VK/VKHelpers.cpp
@@ -268,4 +268,12 @@ namespace vk
 
 		renderer->emergency_query_cleanup(&cmd);
 	}
+
+	void on_descriptor_pool_fragmentation()
+	{
+		if (auto vkthr = dynamic_cast<VKGSRender*>(rsx::get_current_renderer()))
+		{
+			vkthr->on_descriptor_pool_fragmentation();
+		}
+	}
 }
diff --git a/rpcs3/Emu/RSX/VK/vkutils/descriptors.cpp b/rpcs3/Emu/RSX/VK/vkutils/descriptors.cpp
index 3f80e1854d..d9d4101eee 100644
--- a/rpcs3/Emu/RSX/VK/vkutils/descriptors.cpp
+++ b/rpcs3/Emu/RSX/VK/vkutils/descriptors.cpp
@@ -4,6 +4,9 @@
 
 namespace vk
 {
+	// Error handler callback
+	extern void on_descriptor_pool_fragmentation();
+
 	namespace descriptors
 	{
 		class dispatch_manager
@@ -111,6 +114,7 @@ namespace vk
 		ensure(max_sets > 16);
 
 		m_create_info_pool_sizes = pool_sizes;
+
 		for (auto& size : m_create_info_pool_sizes)
 		{
 			ensure(size.descriptorCount < 128); // Sanity check. Remove before commit.
@@ -221,8 +225,6 @@ namespace vk
 			vk::get_gc()->dispose(cleanup_obj);
 		}
 
-		std::lock_guard lock(m_subpool_lock);
-
 		m_current_subpool_offset = 0;
 		m_current_subpool_index = umax;
 
@@ -238,7 +240,23 @@ namespace vk
 		if (m_current_subpool_index == umax)
 		{
 			VkDescriptorPool subpool = VK_NULL_HANDLE;
-			CHECK_RESULT(vkCreateDescriptorPool(*m_owner, &m_create_info, nullptr, &subpool));
+
+			// Only attempt recovery once. Can be bumped up if we have a more complex setup in future.
+			int retries = 1;
+
+			while (VkResult result = vkCreateDescriptorPool(*m_owner, &m_create_info, nullptr, &subpool))
+			{
+				if (retries-- && (result == VK_ERROR_FRAGMENTATION_EXT))
+				{
+					rsx_log.warning("Descriptor pool creation failed with fragmentation error. Will attempt to recover.");
+					vk::on_descriptor_pool_fragmentation();
+					continue;
+				}
+
+				vk::die_with_error(result);
+			}
+
+			std::lock_guard lock(m_subpool_lock);
 
 			m_device_subpools.push_back(
 			{
diff --git a/rpcs3/Emu/RSX/VK/vkutils/shared.cpp b/rpcs3/Emu/RSX/VK/vkutils/shared.cpp
index aab9e789ff..e524c62983 100644
--- a/rpcs3/Emu/RSX/VK/vkutils/shared.cpp
+++ b/rpcs3/Emu/RSX/VK/vkutils/shared.cpp
@@ -96,6 +96,9 @@ namespace vk
 		case VK_ERROR_INVALID_EXTERNAL_HANDLE_KHR:
 			error_message = "Invalid external handle (VK_ERROR_INVALID_EXTERNAL_HANDLE_KHR)";
 			break;
+		case VK_ERROR_FRAGMENTATION_EXT:
+			error_message = "Descriptor pool creation failed (VK_ERROR_FRAGMENTATION)";
+			break;
 		default:
 			error_message = fmt::format("Unknown Code (%Xh, %d)%s", static_cast<s32>(error_code), static_cast<s32>(error_code), src_loc{line, col, file, func});
 			break;