From 8e6d49252404d5ad85f327fb6938ae8cc8fe8a7d Mon Sep 17 00:00:00 2001
From: psucien <bad_cast@protonmail.com>
Date: Thu, 30 May 2024 12:55:37 +0200
Subject: [PATCH] video_core: renderer_vulkan: added color mask support

---
 .../renderer_vulkan/vk_graphics_pipeline.cpp        |  5 +++++
 src/video_core/renderer_vulkan/vk_instance.cpp      | 13 +++++++++++++
 src/video_core/renderer_vulkan/vk_instance.h        |  6 ++++++
 src/video_core/renderer_vulkan/vk_rasterizer.cpp    | 13 +++++++++++++
 4 files changed, 37 insertions(+)

diff --git a/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp b/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
index 44bb61bf..31255a63 100644
--- a/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
+++ b/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
@@ -113,6 +113,11 @@ GraphicsPipeline::GraphicsPipeline(const Instance& instance_, Scheduler& schedul
         vk::DynamicState::eBlendConstants,
     };
 
+    if (instance.IsColorWriteEnableSupported()) {
+        dynamic_states.push_back(vk::DynamicState::eColorWriteEnableEXT);
+        dynamic_states.push_back(vk::DynamicState::eColorWriteMaskEXT);
+    }
+
     const vk::PipelineDynamicStateCreateInfo dynamic_info = {
         .dynamicStateCount = static_cast<u32>(dynamic_states.size()),
         .pDynamicStates = dynamic_states.data(),
diff --git a/src/video_core/renderer_vulkan/vk_instance.cpp b/src/video_core/renderer_vulkan/vk_instance.cpp
index d35d35d1..578a8654 100644
--- a/src/video_core/renderer_vulkan/vk_instance.cpp
+++ b/src/video_core/renderer_vulkan/vk_instance.cpp
@@ -115,6 +115,7 @@ bool Instance::CreateDevice() {
         vk::PhysicalDeviceCustomBorderColorFeaturesEXT, vk::PhysicalDeviceIndexTypeUint8FeaturesEXT,
         vk::PhysicalDeviceFragmentShaderInterlockFeaturesEXT,
         vk::PhysicalDevicePipelineCreationCacheControlFeaturesEXT,
+        vk::PhysicalDeviceColorWriteEnableFeaturesEXT,
         vk::PhysicalDeviceFragmentShaderBarycentricFeaturesKHR>();
     const vk::StructureChain properties_chain =
         physical_device.getProperties2<vk::PhysicalDeviceProperties2,
@@ -150,6 +151,8 @@ bool Instance::CreateDevice() {
     tooling_info = add_extension(VK_EXT_TOOLING_INFO_EXTENSION_NAME);
     custom_border_color = add_extension(VK_EXT_CUSTOM_BORDER_COLOR_EXTENSION_NAME);
     index_type_uint8 = add_extension(VK_KHR_INDEX_TYPE_UINT8_EXTENSION_NAME);
+    color_write_en = add_extension(VK_EXT_COLOR_WRITE_ENABLE_EXTENSION_NAME);
+    add_extension(VK_EXT_EXTENDED_DYNAMIC_STATE_3_EXTENSION_NAME);
     add_extension(VK_KHR_PUSH_DESCRIPTOR_EXTENSION_NAME);
     add_extension(VK_KHR_MAINTENANCE_4_EXTENSION_NAME);
 
@@ -216,12 +219,22 @@ bool Instance::CreateDevice() {
         vk::PhysicalDeviceIndexTypeUint8FeaturesEXT{
             .indexTypeUint8 = true,
         },
+        vk::PhysicalDeviceColorWriteEnableFeaturesEXT{
+            .colorWriteEnable = true,
+        },
+        vk::PhysicalDeviceExtendedDynamicState3FeaturesEXT{
+            .extendedDynamicState3ColorWriteMask = true,
+        },
     };
 
     if (!index_type_uint8) {
         device_chain.unlink<vk::PhysicalDeviceIndexTypeUint8FeaturesEXT>();
     }
 
+    if (!color_write_en) {
+        device_chain.unlink<vk::PhysicalDeviceColorWriteEnableFeaturesEXT>();
+    }
+
     try {
         device = physical_device.createDeviceUnique(device_chain.get());
     } catch (vk::ExtensionNotPresentError& err) {
diff --git a/src/video_core/renderer_vulkan/vk_instance.h b/src/video_core/renderer_vulkan/vk_instance.h
index 5e1be6ce..e4ee9aa4 100644
--- a/src/video_core/renderer_vulkan/vk_instance.h
+++ b/src/video_core/renderer_vulkan/vk_instance.h
@@ -111,6 +111,11 @@ public:
         return external_memory_host;
     }
 
+    /// Returns true when VK_EXT_color_write_enable is supported
+    bool IsColorWriteEnableSupported() const {
+        return color_write_en;
+    }
+
     /// Returns the vendor ID of the physical device
     u32 GetVendorID() const {
         return properties.vendorID;
@@ -218,6 +223,7 @@ private:
     bool fragment_shader_barycentric{};
     bool shader_stencil_export{};
     bool external_memory_host{};
+    bool color_write_en{};
     u64 min_imported_host_pointer_alignment{};
     bool tooling_info{};
     bool debug_utils_supported{};
diff --git a/src/video_core/renderer_vulkan/vk_rasterizer.cpp b/src/video_core/renderer_vulkan/vk_rasterizer.cpp
index 5bdb443d..fa973ed8 100644
--- a/src/video_core/renderer_vulkan/vk_rasterizer.cpp
+++ b/src/video_core/renderer_vulkan/vk_rasterizer.cpp
@@ -138,6 +138,19 @@ void Rasterizer::UpdateDynamicState() {
     auto& regs = liverpool->regs;
     const auto cmdbuf = scheduler.CommandBuffer();
     cmdbuf.setBlendConstants(&regs.blend_constants.red);
+
+    if (instance.IsColorWriteEnableSupported()) {
+        std::array<VkBool32, Liverpool::NumColorBuffers> write_en{};
+        std::array<vk::ColorComponentFlags, Liverpool::NumColorBuffers> write_mask{};
+        for (int col_buf_idx = 0; col_buf_idx < Liverpool::NumColorBuffers; ++col_buf_idx) {
+            const auto mask = regs.color_target_mask.raw >> (col_buf_idx * 4);
+            write_en[col_buf_idx] = mask ? vk::True : vk::False;
+            write_mask[col_buf_idx] = vk::ColorComponentFlags{mask};
+        }
+
+        cmdbuf.setColorWriteEnableEXT(write_en);
+        cmdbuf.setColorWriteMaskEXT(0, write_mask);
+    }
 }
 
 void Rasterizer::UpdateViewportScissorState() {