radv: implement VK_KHR_dynamic_rendering_local_read

Signed-off-by: Samuel Pitoiset <samuel.pitoiset@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/27263>

src/amd/vulkan/meson.build

@@ -79,6 +79,7 @@ libradv_files = files(
   'nir/radv_nir_lower_view_index.c',
   'nir/radv_nir_lower_viewport_to_zero.c',
   'nir/radv_nir_lower_vs_inputs.c',
+  'nir/radv_nir_remap_color_attachment.c',
   'nir/radv_nir_rt_common.c',
   'nir/radv_nir_rt_shader.c',
   'winsys/null/radv_null_bo.c',

src/amd/vulkan/nir/radv_nir.h

@@ -71,6 +71,8 @@ bool radv_nir_lower_cooperative_matrix(nir_shader *shader, unsigned wave_size);
 
 bool radv_nir_lower_draw_id_to_zero(nir_shader *shader);
 
+bool radv_nir_remap_color_attachment(nir_shader *shader, const struct radv_graphics_state_key *gfx_state);
+
 #ifdef __cplusplus
 }
 #endif

src/amd/vulkan/nir/radv_nir_remap_color_attachment.c

@@ -0,0 +1,75 @@
+/*
+ * Copyright © 2024 Valve Corporation
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice (including the next
+ * paragraph) shall be included in all copies or substantial portions of the
+ * Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+ * IN THE SOFTWARE.
+ */
+
+#include "nir/nir.h"
+#include "nir/nir_builder.h"
+#include "radv_constants.h"
+#include "radv_nir.h"
+#include "radv_pipeline_graphics.h"
+#include "vk_graphics_state.h"
+
+static bool
+remap_color_attachment(nir_builder *b, nir_intrinsic_instr *intrin, void *state)
+{
+   const uint8_t *color_remap = (uint8_t *)state;
+
+   if (intrin->intrinsic != nir_intrinsic_store_output)
+      return false;
+
+   nir_io_semantics io_sem = nir_intrinsic_io_semantics(intrin);
+
+   if (io_sem.location < FRAG_RESULT_DATA0)
+      return false;
+
+   if (io_sem.dual_source_blend_index)
+      return false;
+
+   const unsigned location = io_sem.location - FRAG_RESULT_DATA0;
+   if (color_remap[location] == MESA_VK_ATTACHMENT_UNUSED) {
+      nir_instr_remove(&intrin->instr);
+      return false;
+   }
+
+   const unsigned new_location = FRAG_RESULT_DATA0 + color_remap[location];
+
+   io_sem.location = new_location;
+
+   nir_intrinsic_set_io_semantics(intrin, io_sem);
+
+   return true;
+}
+
+bool
+radv_nir_remap_color_attachment(nir_shader *shader, const struct radv_graphics_state_key *gfx_state)
+{
+   uint8_t color_remap[MAX_RTS];
+
+   /* Shader output locations to color attachment mappings. */
+   memset(color_remap, MESA_VK_ATTACHMENT_UNUSED, sizeof(color_remap));
+   for (uint32_t i = 0; i < MAX_RTS; i++) {
+      if (gfx_state->ps.epilog.color_map[i] != MESA_VK_ATTACHMENT_UNUSED)
+         color_remap[gfx_state->ps.epilog.color_map[i]] = i;
+   }
+
+   return nir_shader_intrinsics_pass(shader, remap_color_attachment, nir_metadata_all, &color_remap);
+}

src/amd/vulkan/radv_aco_shader_info.h

@@ -82,8 +82,7 @@ radv_aco_convert_ps_epilog_key(struct aco_ps_epilog_info *aco_info, const struct
    ASSIGN_FIELD(alpha_to_one);
 
    memcpy(aco_info->colors, radv_args->colors, sizeof(aco_info->colors));
-   for (uint32_t i = 0; i < MAX_RTS; i++)
-      aco_info->color_map[i] = i;
+   memcpy(aco_info->color_map, radv->color_map, sizeof(aco_info->color_map));
    aco_info->depth = radv_args->depth;
    aco_info->stencil = radv_args->stencil;
    aco_info->samplemask = radv_args->sample_mask;

src/amd/vulkan/radv_cmd_buffer.c

@@ -155,6 +155,11 @@ radv_bind_dynamic_state(struct radv_cmd_buffer *cmd_buffer, const struct radv_dy
       }
    }
 
+   if (memcmp(&dest->vk.cal.color_map, &src->vk.cal.color_map, sizeof(src->vk.cal.color_map))) {
+      typed_memcpy(dest->vk.cal.color_map, src->vk.cal.color_map, MAX_RTS);
+      cmd_buffer->state.dirty |= RADV_DYNAMIC_COLOR_ATTACHMENT_MAP;
+   }
+
 #define RADV_CMP_COPY(field, flag)                                                                                     \
    if (copy_mask & flag) {                                                                                             \
       if (dest->field != src->field) {                                                                                 \
@@ -4371,6 +4376,9 @@ lookup_ps_epilog(struct radv_cmd_buffer *cmd_buffer)
    const struct radv_dynamic_state *d = &cmd_buffer->state.dynamic;
    const struct radv_physical_device *pdev = radv_device_physical(device);
    struct radv_ps_epilog_state state = {0};
+   uint8_t color_remap[MAX_RTS];
+
+   memset(color_remap, MESA_VK_ATTACHMENT_UNUSED, sizeof(color_remap));
 
    state.color_attachment_count = render->color_att_count;
    for (unsigned i = 0; i < render->color_att_count; ++i) {
@@ -4391,6 +4399,10 @@ lookup_ps_epilog(struct radv_cmd_buffer *cmd_buffer)
           srcRGB == VK_BLEND_FACTOR_SRC_ALPHA_SATURATE || dstRGB == VK_BLEND_FACTOR_SRC_ALPHA_SATURATE ||
           srcRGB == VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA || dstRGB == VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA)
          state.need_src_alpha |= 1 << i;
+
+      state.color_attachment_mappings[i] = d->vk.cal.color_map[i];
+      if (state.color_attachment_mappings[i] != MESA_VK_ATTACHMENT_UNUSED)
+         color_remap[state.color_attachment_mappings[i]] = i;
    }
 
    state.mrt0_is_dual_src = radv_is_mrt0_dual_src(cmd_buffer);
@@ -4416,8 +4428,20 @@ lookup_ps_epilog(struct radv_cmd_buffer *cmd_buffer)
 
    struct radv_ps_epilog_key key = radv_generate_ps_epilog_key(device, &state);
 
+   /* Determine the actual colors written if outputs are remapped. */
+   uint32_t colors_written = 0;
+   for (uint32_t i = 0; i < MAX_RTS; i++) {
+      if (!((ps->info.ps.colors_written >> (i * 4)) & 0xf))
+         continue;
+
+      if (color_remap[i] == MESA_VK_ATTACHMENT_UNUSED)
+         continue;
+
+      colors_written |= 0xfu << (4 * color_remap[i]);
+   }
+
    /* Clear color attachments that aren't exported by the FS to match IO shader arguments. */
-   key.spi_shader_col_format &= ps->info.ps.colors_written;
+   key.spi_shader_col_format &= colors_written;
 
    return radv_shader_part_cache_get(device, &device->ps_epilogs, &cmd_buffer->ps_epilogs, &key);
 }
@@ -7924,6 +7948,23 @@ radv_CmdSetDepthBias2EXT(VkCommandBuffer commandBuffer, const VkDepthBiasInfoEXT
    state->dirty_dynamic |= RADV_CMD_DIRTY_DYNAMIC_DEPTH_BIAS;
 }
 
+VKAPI_ATTR void VKAPI_CALL
+radv_CmdSetRenderingAttachmentLocationsKHR(VkCommandBuffer commandBuffer,
+                                           const VkRenderingAttachmentLocationInfoKHR *pLocationInfo)
+{
+   VK_FROM_HANDLE(radv_cmd_buffer, cmd_buffer, commandBuffer);
+   struct radv_cmd_state *state = &cmd_buffer->state;
+
+   assume(pLocationInfo->colorAttachmentCount <= MESA_VK_MAX_COLOR_ATTACHMENTS);
+   for (uint32_t i = 0; i < pLocationInfo->colorAttachmentCount; i++) {
+      state->dynamic.vk.cal.color_map[i] = pLocationInfo->pColorAttachmentLocations[i] == VK_ATTACHMENT_UNUSED
+                                              ? MESA_VK_ATTACHMENT_UNUSED
+                                              : pLocationInfo->pColorAttachmentLocations[i];
+   }
+
+   state->dirty |= RADV_CMD_DIRTY_DYNAMIC_COLOR_ATTACHMENT_MAP;
+}
+
 VKAPI_ATTR void VKAPI_CALL
 radv_CmdExecuteCommands(VkCommandBuffer commandBuffer, uint32_t commandBufferCount, const VkCommandBuffer *pCmdBuffers)
 {
@@ -9558,7 +9599,7 @@ radv_emit_all_graphics_states(struct radv_cmd_buffer *cmd_buffer, const struct r
             (cmd_buffer->state.dirty_dynamic &
              (RADV_CMD_DIRTY_DYNAMIC_COLOR_WRITE_MASK | RADV_CMD_DIRTY_DYNAMIC_COLOR_BLEND_ENABLE |
               RADV_CMD_DIRTY_DYNAMIC_ALPHA_TO_COVERAGE_ENABLE | RADV_CMD_DIRTY_DYNAMIC_COLOR_BLEND_EQUATION |
-              RADV_CMD_DIRTY_DYNAMIC_ALPHA_TO_ONE_ENABLE))))) {
+              RADV_CMD_DIRTY_DYNAMIC_ALPHA_TO_ONE_ENABLE | RADV_CMD_DIRTY_DYNAMIC_COLOR_ATTACHMENT_MAP))))) {
          ps_epilog = lookup_ps_epilog(cmd_buffer);
          if (!ps_epilog) {
             vk_command_buffer_set_error(&cmd_buffer->vk, VK_ERROR_OUT_OF_HOST_MEMORY);

src/amd/vulkan/radv_cmd_buffer.h

@@ -74,7 +74,8 @@ enum radv_dynamic_state_bits {
    RADV_DYNAMIC_ATTACHMENT_FEEDBACK_LOOP_ENABLE = 1ull << 48,
    RADV_DYNAMIC_SAMPLE_LOCATIONS_ENABLE = 1ull << 49,
    RADV_DYNAMIC_ALPHA_TO_ONE_ENABLE = 1ull << 50,
-   RADV_DYNAMIC_ALL = (1ull << 51) - 1,
+   RADV_DYNAMIC_COLOR_ATTACHMENT_MAP = 1ull << 51,
+   RADV_DYNAMIC_ALL = (1ull << 52) - 1,
 };
 
 enum radv_cmd_dirty_dynamic_bits {
@@ -131,7 +132,8 @@ enum radv_cmd_dirty_dynamic_bits {
    RADV_CMD_DIRTY_DYNAMIC_ATTACHMENT_FEEDBACK_LOOP_ENABLE = 1ull << 48,
    RADV_CMD_DIRTY_DYNAMIC_SAMPLE_LOCATIONS_ENABLE = 1ull << 49,
    RADV_CMD_DIRTY_DYNAMIC_ALPHA_TO_ONE_ENABLE = 1ull << 50,
-   RADV_CMD_DIRTY_DYNAMIC_ALL = (1ull << 51) - 1,
+   RADV_CMD_DIRTY_DYNAMIC_COLOR_ATTACHMENT_MAP = 1ull << 51,
+   RADV_CMD_DIRTY_DYNAMIC_ALL = (1ull << 52) - 1,
 };
 
 enum radv_cmd_dirty_bits {

src/amd/vulkan/radv_pipeline_graphics.c

@@ -1053,6 +1053,10 @@ radv_pipeline_init_dynamic_state(const struct radv_device *device, struct radv_g
          uses_ds_feedback_loop ? (VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT) : VK_IMAGE_ASPECT_NONE;
    }
 
+   for (uint32_t i = 0; i < MAX_RTS; i++) {
+      dynamic->vk.cal.color_map[i] = state->cal ? state->cal->color_map[i] : i;
+   }
+
    pipeline->dynamic_state.mask = states;
 }
 
@@ -1150,6 +1154,7 @@ radv_remove_point_size(const struct radv_graphics_state_key *gfx_state, nir_shad
 static void
 radv_remove_color_exports(const struct radv_graphics_state_key *gfx_state, nir_shader *nir)
 {
+   uint8_t color_remap[MAX_RTS];
    bool fixup_derefs = false;
 
    /* Do not remove color exports when a PS epilog is used because the format isn't known and the color write mask can
@@ -1157,6 +1162,13 @@ radv_remove_color_exports(const struct radv_graphics_state_key *gfx_state, nir_s
    if (gfx_state->ps.has_epilog)
       return;
 
+   /* Shader output locations to color attachment mappings. */
+   memset(color_remap, MESA_VK_ATTACHMENT_UNUSED, sizeof(color_remap));
+   for (uint32_t i = 0; i < MAX_RTS; i++) {
+      if (gfx_state->ps.epilog.color_map[i] != MESA_VK_ATTACHMENT_UNUSED)
+         color_remap[gfx_state->ps.epilog.color_map[i]] = i;
+   }
+
    nir_foreach_shader_out_variable (var, nir) {
       int idx = var->data.location;
       idx -= FRAG_RESULT_DATA0;
@@ -1164,7 +1176,8 @@ radv_remove_color_exports(const struct radv_graphics_state_key *gfx_state, nir_s
       if (idx < 0)
          continue;
 
-      unsigned col_format = (gfx_state->ps.epilog.spi_shader_col_format >> (4 * idx)) & 0xf;
+      const uint8_t cb_idx = color_remap[idx];
+      unsigned col_format = (gfx_state->ps.epilog.spi_shader_col_format >> (4 * cb_idx)) & 0xf;
 
       if (col_format == V_028714_SPI_SHADER_ZERO) {
          /* Remove the color export if it's unused or in presence of holes. */
@@ -1609,12 +1622,15 @@ radv_generate_ps_epilog_key(const struct radv_device *device, const struct radv_
    struct radv_ps_epilog_key key;
 
    memset(&key, 0, sizeof(key));
+   memset(key.color_map, MESA_VK_ATTACHMENT_UNUSED, sizeof(key.color_map));
 
    for (unsigned i = 0; i < state->color_attachment_count; ++i) {
       unsigned cf;
+      unsigned cb_idx = state->color_attachment_mappings[i];
       VkFormat fmt = state->color_attachment_formats[i];
 
-      if (fmt == VK_FORMAT_UNDEFINED || !(state->color_write_mask & (0xfu << (i * 4)))) {
+      if (fmt == VK_FORMAT_UNDEFINED || !(state->color_write_mask & (0xfu << (i * 4))) ||
+          cb_idx == MESA_VK_ATTACHMENT_UNUSED) {
          cf = V_028714_SPI_SHADER_ZERO;
       } else {
          bool blend_enable = state->color_blend_enable & (0xfu << (i * 4));
@@ -1630,6 +1646,8 @@ radv_generate_ps_epilog_key(const struct radv_device *device, const struct radv_
       }
 
       col_format |= cf << (4 * i);
+
+      key.color_map[i] = state->color_attachment_mappings[i];
    }
 
    if (!(col_format & 0xf) && state->need_src_alpha & (1 << 0)) {
@@ -1637,12 +1655,14 @@ radv_generate_ps_epilog_key(const struct radv_device *device, const struct radv_
        * alpha coverage is enabled because the depth attachment needs it.
        */
       col_format |= V_028714_SPI_SHADER_32_AR;
+      key.color_map[0] = 0;
    }
 
    /* The output for dual source blending should have the same format as the first output. */
    if (state->mrt0_is_dual_src) {
       assert(!(col_format >> 4));
       col_format |= (col_format & 0xf) << 4;
+      key.color_map[1] = 1;
    }
 
    if (state->alpha_to_coverage_via_mrtz)
@@ -1722,6 +1742,10 @@ radv_pipeline_generate_ps_epilog_key(const struct radv_device *device, const str
    if (state->ms)
       ps_epilog.alpha_to_one = state->ms->alpha_to_one_enable;
 
+   for (uint32_t i = 0; i < MAX_RTS; i++) {
+      ps_epilog.color_attachment_mappings[i] = state->cal ? state->cal->color_map[i] : i;
+   }
+
    return radv_generate_ps_epilog_key(device, &ps_epilog);
 }
 
@@ -2535,6 +2559,9 @@ radv_graphics_shaders_compile(struct radv_device *device, struct vk_pipeline_cac
 
    if (stages[MESA_SHADER_FRAGMENT].nir) {
       radv_nir_lower_poly_line_smooth(stages[MESA_SHADER_FRAGMENT].nir, gfx_state);
+
+      if (!gfx_state->ps.has_epilog)
+         radv_nir_remap_color_attachment(stages[MESA_SHADER_FRAGMENT].nir, gfx_state);
    }
 
    radv_fill_shader_info(device, RADV_PIPELINE_GRAPHICS, gfx_state, stages, active_nir_stages);

src/amd/vulkan/radv_pipeline_graphics.h

@@ -555,6 +555,7 @@ struct radv_shader *radv_get_shader(struct radv_shader *const *shaders, gl_shade
 struct radv_ps_epilog_state {
    uint8_t color_attachment_count;
    VkFormat color_attachment_formats[MAX_RTS];
+   uint8_t color_attachment_mappings[MAX_RTS];
 
    uint32_t color_write_mask;
    uint32_t color_blend_enable;

src/amd/vulkan/radv_shader.h

@@ -96,6 +96,7 @@ struct radv_ps_epilog_key {
    uint8_t enable_mrt_output_nan_fixup;
 
    uint32_t colors_written;
+   uint8_t color_map[MAX_RTS];
    bool mrt0_is_dual_src;
    bool export_depth;
    bool export_stencil;

src/amd/vulkan/radv_shader_object.c

@@ -144,6 +144,9 @@ radv_shader_object_init_graphics(struct radv_shader_object *shader_obj, struct r
    if (pdev->info.gfx_level >= GFX11)
       gfx_state.ps.exports_mrtz_via_epilog = true;
 
+   for (uint32_t i = 0; i < MAX_RTS; i++)
+      gfx_state.ps.epilog.color_map[i] = i;
+
    struct radv_shader *shader = NULL;
    struct radv_shader_binary *binary = NULL;
 
@@ -420,6 +423,9 @@ radv_shader_object_create_linked(VkDevice _device, uint32_t createInfoCount, con
    if (pdev->info.gfx_level >= GFX11)
       gfx_state.ps.exports_mrtz_via_epilog = true;
 
+   for (uint32_t i = 0; i < MAX_RTS; i++)
+      gfx_state.ps.epilog.color_map[i] = i;
+
    for (unsigned i = 0; i < createInfoCount; i++) {
       const VkShaderCreateInfoEXT *pCreateInfo = &pCreateInfos[i];
       gl_shader_stage s = vk_to_mesa_shader_stage(pCreateInfo->stage);