tu: Implement VK_EXT_color_write_enable

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/16936>

docs/features.txt

@@ -532,7 +532,7 @@ Khronos extensions that are not part of any Vulkan version:
   VK_EXT_border_color_swizzle                           DONE (anv, lvp)
   VK_EXT_buffer_device_address                          DONE (radv)
   VK_EXT_calibrated_timestamps                          DONE (anv, lvp, radv)
-  VK_EXT_color_write_enable                             DONE (anv, lvp, radv, v3dv)
+  VK_EXT_color_write_enable                             DONE (anv, lvp, radv, tu, v3dv)
   VK_EXT_conditional_rendering                          DONE (anv, lvp, radv, tu)
   VK_EXT_conservative_rasterization                     DONE (anv/gen9+, radv)
   VK_EXT_custom_border_color                            DONE (anv, lvp, radv, tu, v3dv)

src/freedreno/vulkan/tu_cmd_buffer.c

@@ -2855,6 +2855,22 @@ tu_CmdSetLineStippleEXT(VkCommandBuffer commandBuffer,
    tu_stub();
 }
 
+VKAPI_ATTR void VKAPI_CALL
+tu_CmdSetColorWriteEnableEXT(VkCommandBuffer commandBuffer, uint32_t attachmentCount,
+                             const VkBool32 *pColorWriteEnables)
+{
+   TU_FROM_HANDLE(tu_cmd_buffer, cmd, commandBuffer);
+   uint32_t color_write_enable = 0;
+
+   for (unsigned i = 0; i < attachmentCount; i++) {
+      if (pColorWriteEnables[i])
+         color_write_enable |= BIT(i);
+   }
+
+   cmd->state.color_write_enable = color_write_enable;
+   cmd->state.dirty |= TU_CMD_DIRTY_BLEND;
+}
+
 static void
 tu_flush_for_access(struct tu_cache_state *cache,
                     enum tu_cmd_access_mask src_mask,
@@ -3819,6 +3835,10 @@ tu6_calculate_lrz_state(struct tu_cmd_buffer *cmd,
        cmd->state.logic_op_enabled && cmd->state.rop_reads_dst)
       gras_lrz_cntl.lrz_write = false;
 
+   if ((cmd->state.pipeline->dynamic_state_mask & BIT(TU_DYNAMIC_STATE_COLOR_WRITE_ENABLE)) &&
+        cmd->state.color_write_enable != MASK(cmd->state.pipeline->num_rts))
+      gras_lrz_cntl.lrz_write = false;
+
    /* LRZ is disabled until it is cleared, which means that one "wrong"
     * depth test or shader could disable LRZ until depth buffer is cleared.
     */
@@ -4053,6 +4073,10 @@ tu6_emit_blend(struct tu_cs *cs, struct tu_cmd_buffer *cmd)
    struct tu_pipeline *pipeline = cmd->state.pipeline;
    uint32_t color_write_enable = cmd->state.pipeline_color_write_enable;
 
+   if (pipeline->dynamic_state_mask &
+       BIT(TU_DYNAMIC_STATE_COLOR_WRITE_ENABLE))
+      color_write_enable &= cmd->state.color_write_enable;
+
    for (unsigned i = 0; i < pipeline->num_rts; i++) {
       tu_cs_emit_pkt4(cs, REG_A6XX_RB_MRT_CONTROL(i), 2);
       if (color_write_enable & BIT(i)) {
@@ -4069,7 +4093,8 @@ tu6_emit_blend(struct tu_cs *cs, struct tu_cmd_buffer *cmd)
 
    uint32_t blend_enable_mask =
       (cmd->state.logic_op_enabled && cmd->state.rop_reads_dst) ?
-      color_write_enable : cmd->state.pipeline_blend_enable;
+      color_write_enable : (cmd->state.pipeline_blend_enable &
+                            cmd->state.color_write_enable);
 
    tu_cs_emit_pkt4(cs, REG_A6XX_SP_BLEND_CNTL, 1);
    tu_cs_emit(cs, cmd->state.sp_blend_cntl |

src/freedreno/vulkan/tu_device.c

@@ -232,6 +232,7 @@ get_device_extensions(const struct tu_physical_device *device,
       .IMG_filter_cubic = device->info->a6xx.has_tex_filter_cubic,
       .VALVE_mutable_descriptor_type = true,
       .EXT_image_2d_view_of_3d = true,
+      .EXT_color_write_enable = true,
    };
 }
 
@@ -886,6 +887,12 @@ tu_GetPhysicalDeviceFeatures2(VkPhysicalDevice physicalDevice,
          features->sampler2DViewOf3D = true;
          break;
       }
+      case VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COLOR_WRITE_ENABLE_FEATURES_EXT: {
+         VkPhysicalDeviceColorWriteEnableFeaturesEXT *features =
+            (VkPhysicalDeviceColorWriteEnableFeaturesEXT *)ext;
+         features->colorWriteEnable = true;
+         break;
+      }
 
       default:
          break;

src/freedreno/vulkan/tu_pipeline.c

@@ -2155,6 +2155,17 @@ tu6_emit_rb_mrt_controls(struct tu_pipeline *pipeline,
                          bool *rop_reads_dst,
                          uint32_t *color_bandwidth_per_sample)
 {
+   const VkPipelineColorWriteCreateInfoEXT *color_info =
+      vk_find_struct_const(blend_info->pNext,
+                           PIPELINE_COLOR_WRITE_CREATE_INFO_EXT);
+
+   /* The static state is ignored if it's dynamic. In that case assume
+    * everything is enabled and then the appropriate registers will be zero'd
+    * dynamically.
+    */
+   if (pipeline->dynamic_state_mask & BIT(TU_DYNAMIC_STATE_COLOR_WRITE_ENABLE))
+      color_info = NULL;
+
    *rop_reads_dst = false;
    *color_bandwidth_per_sample = 0;
 
@@ -2174,7 +2185,8 @@ tu6_emit_rb_mrt_controls(struct tu_pipeline *pipeline,
 
       uint32_t rb_mrt_control = 0;
       uint32_t rb_mrt_blend_control = 0;
-      if (format != VK_FORMAT_UNDEFINED) {
+      if (format != VK_FORMAT_UNDEFINED &&
+          (!color_info || color_info->pColorWriteEnables[i])) {
          const bool has_alpha = vk_format_has_alpha(format);
 
          rb_mrt_control =
@@ -3149,6 +3161,17 @@ tu_pipeline_builder_parse_dynamic(struct tu_pipeline_builder *builder,
          pipeline->dynamic_state_mask |= BIT(TU_DYNAMIC_STATE_BLEND);
          pipeline->dynamic_state_mask |= BIT(TU_DYNAMIC_STATE_LOGIC_OP);
          break;
+      case VK_DYNAMIC_STATE_COLOR_WRITE_ENABLE_EXT:
+         pipeline->sp_blend_cntl_mask &= ~A6XX_SP_BLEND_CNTL_ENABLE_BLEND__MASK;
+         pipeline->rb_blend_cntl_mask &= ~A6XX_RB_BLEND_CNTL_ENABLE_BLEND__MASK;
+         pipeline->dynamic_state_mask |= BIT(TU_DYNAMIC_STATE_BLEND);
+
+         /* Dynamic color write enable doesn't directly change any of the
+          * registers, but it causes us to make some of the registers 0, so we
+          * set this dynamic state instead of making the register dynamic.
+          */
+         pipeline->dynamic_state_mask |= BIT(TU_DYNAMIC_STATE_COLOR_WRITE_ENABLE);
+         break;
       default:
          assert(!"unsupported dynamic state");
          break;
@@ -3615,7 +3638,8 @@ tu_pipeline_builder_parse_multisample_and_color_blend(
       VkFormat format = builder->color_attachment_formats[i];
       unsigned mask = MASK(vk_format_get_nr_components(format));
       if (format != VK_FORMAT_UNDEFINED &&
-          (blendAttachment.colorWriteMask & mask) != mask) {
+          ((blendAttachment.colorWriteMask & mask) != mask ||
+           !(pipeline->color_write_enable & BIT(i)))) {
          pipeline->lrz.force_disable_mask |= TU_LRZ_FORCE_DISABLE_WRITE;
       }
    }

src/freedreno/vulkan/tu_private.h

@@ -701,6 +701,7 @@ enum tu_dynamic_state
    TU_DYNAMIC_STATE_PRIMITIVE_TOPOLOGY = TU_DYNAMIC_STATE_COUNT,
    TU_DYNAMIC_STATE_PRIMITIVE_RESTART_ENABLE,
    TU_DYNAMIC_STATE_LOGIC_OP,
+   TU_DYNAMIC_STATE_COLOR_WRITE_ENABLE,
    /* re-use the line width enum as it uses GRAS_SU_CNTL: */
    TU_DYNAMIC_STATE_GRAS_SU_CNTL = VK_DYNAMIC_STATE_LINE_WIDTH,
 };
@@ -1169,6 +1170,7 @@ struct tu_cmd_state
    uint32_t rb_mrt_control_rop;
    uint32_t rb_blend_cntl, sp_blend_cntl;
    uint32_t pipeline_color_write_enable, pipeline_blend_enable;
+   uint32_t color_write_enable;
    bool logic_op_enabled;
    bool rop_reads_dst;
    enum pc_di_primtype primtype;