pvr, pco: implement VK_EXT_color_write_enable

Signed-off-by: Simon Perretta <simon.perretta@imgtec.com>
Acked-by: Erik Faye-Lund <erik.faye-lund@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/36412>

docs/features.txt

@@ -594,7 +594,7 @@ Khronos extensions that are not part of any Vulkan version:
   VK_EXT_border_color_swizzle                           DONE (anv, hasvk, hk, lvp, nvk, panvk, radv/gfx10+, tu, v3dv, vn)
   VK_EXT_buffer_device_address                          DONE (anv, hasvk, hk, nvk, panvk, radv, vn)
   VK_EXT_calibrated_timestamps                          DONE (anv, hasvk, hk, nvk, panvk/v10+, lvp, radv, vn, tu/a750+)
-  VK_EXT_color_write_enable                             DONE (anv, hasvk, hk, lvp, nvk, radv, tu, v3dv, vn)
+  VK_EXT_color_write_enable                             DONE (anv, hasvk, hk, lvp, nvk, pvr, radv, tu, v3dv, vn)
   VK_EXT_conditional_rendering                          DONE (anv, hasvk, lvp, nvk, radv, tu, vn)
   VK_EXT_conservative_rasterization                     DONE (anv, nvk, radv, vn, tu/a7xx+)
   VK_EXT_custom_border_color                            DONE (anv, hasvk, hk, lvp, nvk, panvk, radv, tu, v3dv, vn)

src/imagination/pco/pco_data.h

@@ -121,6 +121,7 @@ typedef struct _pco_fs_data {
    struct {
       bool alpha_to_one;
       bool sample_mask;
+      bool color_write_enable;
    } meta_present;
 } pco_fs_data;
 

src/imagination/pco/pco_nir.c

@@ -244,6 +244,8 @@ static void gather_fs_data(nir_shader *nir, pco_data *data)
    }
 
    data->fs.uses.fbfetch = nir->info.fs.uses_fbfetch_output;
+   data->fs.uses.fbfetch |= data->fs.meta_present.color_write_enable;
+
    data->fs.uses.early_frag = nir->info.fs.early_fragment_tests;
    data->fs.uses.sample_shading |= nir->info.fs.uses_sample_shading;
 }

src/imagination/pco/pco_nir_pvfio.c

@@ -1196,6 +1196,48 @@ lower_load_sample_mask(nir_builder *b, nir_instr *instr, void *cb_data)
    return smp_msk;
 }
 
+static nir_def *
+lower_color_write_enable(nir_builder *b, nir_instr *instr, UNUSED void *cb_data)
+{
+   nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
+
+   nir_src *input_src = &intr->src[0];
+   nir_def *input = input_src->ssa;
+   nir_def *offset = intr->src[1].ssa;
+
+   struct nir_io_semantics io_semantics = nir_intrinsic_io_semantics(intr);
+   unsigned color_write_index = io_semantics.location - FRAG_RESULT_DATA0;
+   io_semantics.fb_fetch_output = true;
+
+   b->cursor = nir_before_instr(&intr->instr);
+
+   /* TODO: nir op that returns bool based on whether a bit is set. */
+   /* TODO: define for 1 */
+   nir_def *color_write_enabled =
+      nir_ine_imm(b,
+                  nir_ubitfield_extract_imm(b,
+                                            nir_load_fs_meta_pco(b),
+                                            1 + color_write_index,
+                                            1),
+                  0);
+
+   nir_def *prev_input =
+      nir_load_output(b,
+                      input->num_components,
+                      input->bit_size,
+                      offset,
+                      .base = nir_intrinsic_base(intr),
+                      .range = nir_intrinsic_range(intr),
+                      .component = nir_intrinsic_component(intr),
+                      .dest_type = nir_intrinsic_src_type(intr),
+                      .io_semantics = io_semantics);
+
+   nir_src_rewrite(input_src,
+                   nir_bcsel(b, color_write_enabled, input, prev_input));
+
+   return NIR_LOWER_INSTR_PROGRESS;
+}
+
 /**
  * \brief Per-fragment output pass.
  *
@@ -1223,12 +1265,21 @@ bool pco_nir_pfo(nir_shader *shader, pco_fs_data *fs)
 
    bool progress = false;
 
+   /* TODO: instead of doing multiple passes, probably better to just cache all
+    * the stores
+    */
    if (fs->meta_present.alpha_to_one)
       progress |= nir_shader_lower_instructions(shader,
                                                 is_frag_color_out,
                                                 lower_alpha_to_one,
                                                 &state);
 
+   if (fs->meta_present.color_write_enable)
+      progress |= nir_shader_lower_instructions(shader,
+                                                is_frag_color_out,
+                                                lower_color_write_enable,
+                                                NULL);
+
    progress |= nir_shader_lower_instructions(shader, is_pfo, lower_pfo, &state);
    progress |= lower_isp_fb(&b, &state);
 

src/imagination/vulkan/pvr_cmd_buffer.c

@@ -3935,6 +3935,9 @@ static VkResult pvr_setup_descriptor_mappings(
 
             fs_meta |= cmd_buffer->vk.dynamic_graphics_state.ms.sample_mask
                        << 9;
+            fs_meta |=
+               cmd_buffer->vk.dynamic_graphics_state.cb.color_write_enables
+               << 1;
 
             struct pvr_suballoc_bo *fs_meta_bo;
             result = pvr_cmd_buffer_upload_general(cmd_buffer,
@@ -6275,7 +6278,9 @@ static VkResult pvr_validate_draw_state(struct pvr_cmd_buffer *cmd_buffer)
    state->dirty.fragment_descriptors |= state->dirty.gfx_desc_dirty;
 
    if (BITSET_TEST(dynamic_state->dirty, MESA_VK_DYNAMIC_CB_BLEND_CONSTANTS) ||
-       BITSET_TEST(dynamic_state->dirty, MESA_VK_DYNAMIC_RS_FRONT_FACE)) {
+       BITSET_TEST(dynamic_state->dirty, MESA_VK_DYNAMIC_RS_FRONT_FACE) ||
+       BITSET_TEST(dynamic_state->dirty,
+                   MESA_VK_DYNAMIC_CB_COLOR_WRITE_ENABLES)) {
       state->dirty.fragment_descriptors = true;
    }
 

src/imagination/vulkan/pvr_device.c

@@ -191,6 +191,7 @@ static void pvr_physical_device_get_supported_extensions(
       .KHR_uniform_buffer_standard_layout = true,
       .KHR_vertex_attribute_divisor = true,
       .KHR_zero_initialize_workgroup_memory = false,
+      .EXT_color_write_enable = true,
       .EXT_depth_clamp_zero_one = true,
       .EXT_external_memory_dma_buf = true,
       .EXT_host_query_reset = true,
@@ -279,6 +280,9 @@ static void pvr_physical_device_get_supported_features(
       /* Vulkan 1.2 / VK_KHR_uniform_buffer_standard_layout */
       .uniformBufferStandardLayout = true,
 
+      /* VK_EXT_color_write_enable */
+      .colorWriteEnable = true,
+
       /* Vulkan 1.2 / VK_EXT_host_query_reset */
       .hostQueryReset = true,
 

src/imagination/vulkan/pvr_pipeline.c

@@ -1848,6 +1848,7 @@ static void pvr_alloc_fs_sysvals(pco_data *data, nir_shader *nir)
 
    has_meta |= data->fs.meta_present.alpha_to_one;
    has_meta |= data->fs.meta_present.sample_mask;
+   has_meta |= data->fs.meta_present.color_write_enable;
    if (!has_meta)
       return;
 
@@ -2493,6 +2494,11 @@ pvr_preprocess_shader_data(pco_data *data,
       data->fs.rasterization_samples = state->ms->rasterization_samples;
       nir->info.fs.uses_sample_shading = state->ms->rasterization_samples >
                                          VK_SAMPLE_COUNT_1_BIT;
+      if (BITSET_TEST(state->dynamic, MESA_VK_DYNAMIC_CB_COLOR_WRITE_ENABLES) ||
+          (state->cb && state->cb->color_write_enables !=
+                           BITFIELD_MASK(MESA_VK_MAX_COLOR_ATTACHMENTS))) {
+         data->fs.meta_present.color_write_enable = true;
+      }
 
       /* TODO: push consts, dynamic state, etc. */
       break;