renderer_vulkan: Resore color write dynamic state (#3453)

src/video_core/amdgpu/pixel_format.h

@@ -127,6 +127,14 @@ union CompMapping {
         };
     }
 
+    [[nodiscard]] u32 ApplyMask(u32 mask) const {
+        u32 swizzled_mask{};
+        for (u32 i = 0; i < 4; ++i) {
+            swizzled_mask |= ((mask >> i) & 1) << Map(i);
+        }
+        return swizzled_mask;
+    }
+
     [[nodiscard]] CompMapping Inverse() const {
         CompMapping result{};
         InverseSingle(result.r, CompSwizzle::Red);

src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp

@@ -168,6 +168,9 @@ GraphicsPipeline::GraphicsPipeline(
         dynamic_states.push_back(vk::DynamicState::eDepthBoundsTestEnable);
         dynamic_states.push_back(vk::DynamicState::eDepthBounds);
     }
+    if (instance.IsDynamicColorWriteMaskSupported()) {
+        dynamic_states.push_back(vk::DynamicState::eColorWriteMaskEXT);
+    }
     if (instance.IsVertexInputDynamicState()) {
         dynamic_states.push_back(vk::DynamicState::eVertexInputEXT);
     } else if (!vertex_bindings.empty()) {
@@ -288,7 +291,11 @@ GraphicsPipeline::GraphicsPipeline(
             .alphaBlendOp = control.separate_alpha_blend
                                 ? LiverpoolToVK::BlendOp(control.alpha_func)
                                 : color_blend,
-            .colorWriteMask = vk::ColorComponentFlags{key.write_masks[i]},
+            .colorWriteMask =
+                instance.IsDynamicColorWriteMaskSupported()
+                    ? vk::ColorComponentFlagBits::eR | vk::ColorComponentFlagBits::eG |
+                          vk::ColorComponentFlagBits::eB | vk::ColorComponentFlagBits::eA
+                    : key.write_masks[i],
         };
 
         // On GCN GPU there is an additional mask which allows to control color components exported

src/video_core/renderer_vulkan/vk_graphics_pipeline.h

@@ -38,7 +38,7 @@ struct GraphicsPipelineKey {
     u32 num_color_attachments;
     std::array<Shader::PsColorBuffer, Liverpool::NumColorBuffers> color_buffers;
     std::array<Liverpool::BlendControl, Liverpool::NumColorBuffers> blend_controls;
-    std::array<u32, Liverpool::NumColorBuffers> write_masks;
+    std::array<vk::ColorComponentFlags, Liverpool::NumColorBuffers> write_masks;
     Liverpool::ColorBufferMask cb_shader_mask;
     Liverpool::ColorControl::LogicOp logic_op;
     u32 num_samples;
@@ -80,6 +80,10 @@ public:
         return fetch_shader;
     }
 
+    auto GetWriteMasks() const {
+        return key.write_masks;
+    }
+
     u32 GetMrtMask() const {
         return key.mrt_mask;
     }

src/video_core/renderer_vulkan/vk_instance.cpp

@@ -255,6 +255,13 @@ bool Instance::CreateDevice() {
     // Optional
     maintenance_8 = add_extension(VK_KHR_MAINTENANCE_8_EXTENSION_NAME);
     depth_range_unrestricted = add_extension(VK_EXT_DEPTH_RANGE_UNRESTRICTED_EXTENSION_NAME);
+    dynamic_state_3 = add_extension(VK_EXT_EXTENDED_DYNAMIC_STATE_3_EXTENSION_NAME);
+    if (dynamic_state_3) {
+        dynamic_state_3_features =
+            feature_chain.get<vk::PhysicalDeviceExtendedDynamicState3FeaturesEXT>();
+        LOG_INFO(Render_Vulkan, "- extendedDynamicState3ColorWriteMask: {}",
+                 dynamic_state_3_features.extendedDynamicState3ColorWriteMask);
+    }
     robustness2 = add_extension(VK_EXT_ROBUSTNESS_2_EXTENSION_NAME);
     if (robustness2) {
         robustness2_features = feature_chain.get<vk::PhysicalDeviceRobustness2FeaturesEXT>();
@@ -419,6 +426,10 @@ bool Instance::CreateDevice() {
             .customBorderColors = true,
             .customBorderColorWithoutFormat = true,
         },
+        vk::PhysicalDeviceExtendedDynamicState3FeaturesEXT{
+            .extendedDynamicState3ColorWriteMask =
+                dynamic_state_3_features.extendedDynamicState3ColorWriteMask,
+        },
         vk::PhysicalDeviceDepthClipControlFeaturesEXT{
             .depthClipControl = true,
         },
@@ -494,6 +505,9 @@ bool Instance::CreateDevice() {
     if (!custom_border_color) {
         device_chain.unlink<vk::PhysicalDeviceCustomBorderColorFeaturesEXT>();
     }
+    if (!dynamic_state_3) {
+        device_chain.unlink<vk::PhysicalDeviceExtendedDynamicState3FeaturesEXT>();
+    }
     if (!depth_clip_control) {
         device_chain.unlink<vk::PhysicalDeviceDepthClipControlFeaturesEXT>();
     }

src/video_core/renderer_vulkan/vk_instance.h

@@ -139,6 +139,12 @@ public:
         return depth_range_unrestricted;
     }
 
+    /// Returns true when the extendedDynamicState3ColorWriteMask feature o
+    /// VK_EXT_extended_dynamic_state3 is supported.
+    bool IsDynamicColorWriteMaskSupported() const {
+        return dynamic_state_3 && dynamic_state_3_features.extendedDynamicState3ColorWriteMask;
+    }
+
     /// Returns true when VK_EXT_vertex_input_dynamic_state is supported.
     bool IsVertexInputDynamicState() const {
         return vertex_input_dynamic_state;
@@ -433,6 +439,7 @@ private:
     vk::PhysicalDeviceFeatures features;
     vk::PhysicalDeviceVulkan12Features vk12_features;
     vk::PhysicalDevicePortabilitySubsetFeaturesKHR portability_features;
+    vk::PhysicalDeviceExtendedDynamicState3FeaturesEXT dynamic_state_3_features;
     vk::PhysicalDeviceRobustness2FeaturesEXT robustness2_features;
     vk::PhysicalDeviceShaderAtomicFloat2FeaturesEXT shader_atomic_float2_features;
     vk::PhysicalDeviceWorkgroupMemoryExplicitLayoutFeaturesKHR
@@ -453,6 +460,7 @@ private:
     bool amd_shader_explicit_vertex_parameter{};
     bool depth_clip_control{};
     bool depth_clip_enable{};
+    bool dynamic_state_3{};
     bool depth_range_unrestricted{};
     bool vertex_input_dynamic_state{};
     bool robustness2{};

src/video_core/renderer_vulkan/vk_pipeline_cache.cpp

@@ -345,10 +345,8 @@ bool PipelineCache::RefreshGraphicsKey() {
                                              !col_buf.info.blend_bypass);
 
         // Apply swizzle to target mask
-        const auto& swizzle = key.color_buffers[cb].swizzle;
+        key.write_masks[cb] =
-        for (u32 i = 0; i < 4; ++i) {
+            vk::ColorComponentFlags{key.color_buffers[cb].swizzle.ApplyMask(target_mask)};
-            key.write_masks[cb] |= ((target_mask >> i) & 1) << swizzle.Map(i);
-        }
     }
 
     // Compile and bind shader stages

src/video_core/renderer_vulkan/vk_rasterizer.cpp

@@ -1097,6 +1097,7 @@ void Rasterizer::UpdateDynamicState(const GraphicsPipeline& pipeline) const {
 
     auto& dynamic_state = scheduler.GetDynamicState();
     dynamic_state.SetBlendConstants(liverpool->regs.blend_constants);
+    dynamic_state.SetColorWriteMasks(pipeline.GetWriteMasks());
 
     // Commit new dynamic state to the command buffer.
     dynamic_state.Commit(instance, scheduler.CommandBuffer());

src/video_core/renderer_vulkan/vk_scheduler.cpp

@@ -323,6 +323,12 @@ void DynamicState::Commit(const Instance& instance, const vk::CommandBuffer& cmd
         dirty_state.blend_constants = false;
         cmdbuf.setBlendConstants(blend_constants.data());
     }
+    if (dirty_state.color_write_masks) {
+        dirty_state.color_write_masks = false;
+        if (instance.IsDynamicColorWriteMaskSupported()) {
+            cmdbuf.setColorWriteMaskEXT(0, color_write_masks);
+        }
+    }
     if (dirty_state.line_width) {
         dirty_state.line_width = false;
         cmdbuf.setLineWidth(line_width);

src/video_core/renderer_vulkan/vk_scheduler.h

@@ -108,6 +108,7 @@ struct DynamicState {
         bool front_face : 1;
 
         bool blend_constants : 1;
+        bool color_write_masks : 1;
         bool line_width : 1;
     } dirty_state{};
 
@@ -143,6 +144,7 @@ struct DynamicState {
     vk::FrontFace front_face{};
 
     std::array<float, 4> blend_constants{};
+    ColorWriteMasks color_write_masks{};
     float line_width{};
 
     /// Commits the dynamic state to the provided command buffer.
@@ -306,6 +308,13 @@ struct DynamicState {
         }
     }
 
+    void SetColorWriteMasks(const ColorWriteMasks& color_write_masks_) {
+        if (!std::ranges::equal(color_write_masks, color_write_masks_)) {
+            color_write_masks = color_write_masks_;
+            dirty_state.color_write_masks = true;
+        }
+    }
+
     void SetLineWidth(const float width) {
         if (line_width != width) {
             line_width = width;