radv/meta: fuse depth/stencil aspects copy with the GFX path

Depth/stencil copies on graphics are twice as fast now.

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

src/amd/vulkan/meta/radv_meta.h

@@ -79,6 +79,7 @@ enum radv_meta_object_key_type {
    RADV_META_OBJECT_KEY_BLIT2D_COLOR,
    RADV_META_OBJECT_KEY_BLIT2D_DEPTH,
    RADV_META_OBJECT_KEY_BLIT2D_STENCIL,
+   RADV_META_OBJECT_KEY_BLIT2D_DEPTH_STENCIL,
    RADV_META_OBJECT_KEY_FILL_MEMORY,
    RADV_META_OBJECT_KEY_COPY_MEMORY,
    RADV_META_OBJECT_KEY_COPY_IMAGE_TO_BUFFER,

src/amd/vulkan/meta/radv_meta_blit2d.c

@@ -30,6 +30,10 @@ static VkResult get_stencil_only_pipeline(struct radv_device *device, enum blit2
                                           uint32_t log2_samples, VkPipeline *pipeline_out,
                                           VkPipelineLayout *layout_out);
 
+static VkResult get_depth_stencil_pipeline(struct radv_device *device, enum blit2d_src_type src_type,
+                                           uint32_t log2_samples, VkPipeline *pipeline_out,
+                                           VkPipelineLayout *layout_out);
+
 static void
 create_iview(struct radv_cmd_buffer *cmd_buffer, struct radv_meta_blit2d_surf *surf, struct radv_image_view *iview,
              VkFormat depth_format, VkImageAspectFlagBits aspects)
@@ -237,26 +241,22 @@ radv_gfx_copy_image(struct radv_cmd_buffer *cmd_buffer, struct radv_meta_blit2d_
                          .extent = (VkExtent2D){extent->width, extent->height},
                       });
 
-   u_foreach_bit (i, dst->aspect_mask) {
-      unsigned dst_aspect_mask = 1u << i;
-      unsigned src_aspect_mask = dst_aspect_mask;
    VkFormat src_format = src->format;
    VkFormat dst_format = dst->format;
 
-      if (dst_aspect_mask == VK_IMAGE_ASPECT_STENCIL_BIT) {
+   if (dst->aspect_mask == (VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT)) {
+      dst_format = dst->image->vk.format;
+   } else if (dst->aspect_mask == VK_IMAGE_ASPECT_STENCIL_BIT) {
       dst_format = vk_format_stencil_only(dst->image->vk.format);
       src_format = dst_format;
-      } else if (dst_aspect_mask == VK_IMAGE_ASPECT_DEPTH_BIT) {
+   } else if (dst->aspect_mask == VK_IMAGE_ASPECT_DEPTH_BIT) {
       dst_format = vk_format_depth_only(dst->image->vk.format);
       src_format = dst_format;
-      } else {
-         src_aspect_mask = src->aspect_mask;
    }
 
-      /* Adjust the aspect for color to depth/stencil image copies. */
+   /* Adjust the formats for color to depth/stencil image copies. */
    if (vk_format_is_color(src->image->vk.format) && vk_format_is_depth_or_stencil(dst->image->vk.format)) {
       assert(src->aspect_mask == VK_IMAGE_ASPECT_COLOR_BIT);
-         src_aspect_mask = src->aspect_mask;
       src_format = src->format;
    } else if (vk_format_is_depth_or_stencil(src->image->vk.format) && vk_format_is_color(dst->image->vk.format)) {
       if (src->aspect_mask == VK_IMAGE_ASPECT_STENCIL_BIT) {
@@ -268,7 +268,7 @@ radv_gfx_copy_image(struct radv_cmd_buffer *cmd_buffer, struct radv_meta_blit2d_
    }
 
    struct radv_image_view dst_iview;
-      create_iview(cmd_buffer, dst, &dst_iview, dst_format, dst_aspect_mask);
+   create_iview(cmd_buffer, dst, &dst_iview, dst_format, dst->aspect_mask);
 
    const VkRenderingAttachmentInfo att_info = {
       .sType = VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO,
@@ -289,23 +289,34 @@ radv_gfx_copy_image(struct radv_cmd_buffer *cmd_buffer, struct radv_meta_blit2d_
       .layerCount = 1,
    };
 
-      if (dst_aspect_mask == VK_IMAGE_ASPECT_COLOR_BIT || dst_aspect_mask == VK_IMAGE_ASPECT_PLANE_0_BIT ||
+   if (dst->aspect_mask == VK_IMAGE_ASPECT_COLOR_BIT || dst->aspect_mask == VK_IMAGE_ASPECT_PLANE_0_BIT ||
-          dst_aspect_mask == VK_IMAGE_ASPECT_PLANE_1_BIT || dst_aspect_mask == VK_IMAGE_ASPECT_PLANE_2_BIT) {
+       dst->aspect_mask == VK_IMAGE_ASPECT_PLANE_1_BIT || dst->aspect_mask == VK_IMAGE_ASPECT_PLANE_2_BIT) {
       result = get_color_pipeline(device, src_type, dst_format, log2_samples, &pipeline, &layout);
       if (result != VK_SUCCESS) {
          vk_command_buffer_set_error(&cmd_buffer->vk, result);
-            goto fail;
+         goto fail_pipeline;
       }
 
       rendering_info.colorAttachmentCount = 1;
       rendering_info.pColorAttachments = &att_info;
 
       radv_CmdBindPipeline(radv_cmd_buffer_to_handle(cmd_buffer), VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline);
-      } else if (dst_aspect_mask == VK_IMAGE_ASPECT_DEPTH_BIT) {
+   } else if (dst->aspect_mask == (VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT)) {
+      result = get_depth_stencil_pipeline(device, src_type, log2_samples, &pipeline, &layout);
+      if (result != VK_SUCCESS) {
+         vk_command_buffer_set_error(&cmd_buffer->vk, result);
+         goto fail_pipeline;
+      }
+
+      rendering_info.pDepthAttachment = &att_info;
+      rendering_info.pStencilAttachment = &att_info;
+
+      radv_CmdBindPipeline(radv_cmd_buffer_to_handle(cmd_buffer), VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline);
+   } else if (dst->aspect_mask == VK_IMAGE_ASPECT_DEPTH_BIT) {
       result = get_depth_only_pipeline(device, src_type, log2_samples, &pipeline, &layout);
       if (result != VK_SUCCESS) {
          vk_command_buffer_set_error(&cmd_buffer->vk, result);
-            goto fail;
+         goto fail_pipeline;
       }
 
       rendering_info.pDepthAttachment = &att_info,
@@ -313,11 +324,11 @@ radv_gfx_copy_image(struct radv_cmd_buffer *cmd_buffer, struct radv_meta_blit2d_
 
       radv_CmdBindPipeline(radv_cmd_buffer_to_handle(cmd_buffer), VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline);
 
-      } else if (dst_aspect_mask == VK_IMAGE_ASPECT_STENCIL_BIT) {
+   } else if (dst->aspect_mask == VK_IMAGE_ASPECT_STENCIL_BIT) {
       result = get_stencil_only_pipeline(device, src_type, log2_samples, &pipeline, &layout);
       if (result != VK_SUCCESS) {
          vk_command_buffer_set_error(&cmd_buffer->vk, result);
-            goto fail;
+         goto fail_pipeline;
       }
 
       rendering_info.pDepthAttachment = (dst->image->vk.aspects & VK_IMAGE_ASPECT_DEPTH_BIT) ? &att_info : NULL,
@@ -358,8 +369,34 @@ radv_gfx_copy_image(struct radv_cmd_buffer *cmd_buffer, struct radv_meta_blit2d_
       radv_CmdPushConstants2(radv_cmd_buffer_to_handle(cmd_buffer), &pc_info_fs);
    }
 
-      struct radv_image_view src_iview;
+   struct radv_image_view src_iview, src_iview_depth, src_iview_stencil;
-      create_iview(cmd_buffer, src, &src_iview, src_format, src_aspect_mask);
+
+   if (dst->aspect_mask == (VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT)) {
+      VkFormat depth_format = vk_format_depth_only(dst->image->vk.format);
+      VkFormat stencil_format = vk_format_stencil_only(dst->image->vk.format);
+
+      create_iview(cmd_buffer, src, &src_iview_depth, depth_format, VK_IMAGE_ASPECT_DEPTH_BIT);
+      create_iview(cmd_buffer, src, &src_iview_stencil, stencil_format, VK_IMAGE_ASPECT_STENCIL_BIT);
+
+      radv_meta_bind_descriptors(
+         cmd_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, layout, 2,
+         (VkDescriptorGetInfoEXT[]){{.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_GET_INFO_EXT,
+                                     .type = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
+                                     .data.pSampledImage =
+                                        (VkDescriptorImageInfo[]){
+                                           {.sampler = VK_NULL_HANDLE,
+                                            .imageView = radv_image_view_to_handle(&src_iview_depth),
+                                            .imageLayout = VK_IMAGE_LAYOUT_GENERAL},
+                                        }},
+                                    {.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_GET_INFO_EXT,
+                                     .type = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
+                                     .data.pSampledImage = (VkDescriptorImageInfo[]){
+                                        {.sampler = VK_NULL_HANDLE,
+                                         .imageView = radv_image_view_to_handle(&src_iview_stencil),
+                                         .imageLayout = VK_IMAGE_LAYOUT_GENERAL},
+                                     }}});
+   } else {
+      create_iview(cmd_buffer, src, &src_iview, src_format, src->aspect_mask);
 
       radv_meta_bind_descriptors(cmd_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, layout, 1,
                                  (VkDescriptorGetInfoEXT[]){{.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_GET_INFO_EXT,
@@ -371,6 +408,7 @@ radv_gfx_copy_image(struct radv_cmd_buffer *cmd_buffer, struct radv_meta_blit2d_
                                                                    .imageLayout = VK_IMAGE_LAYOUT_GENERAL,
                                                                 },
                                                              }}});
+   }
 
    radv_CmdBeginRendering(radv_cmd_buffer_to_handle(cmd_buffer), &rendering_info);
 
@@ -382,10 +420,14 @@ radv_gfx_copy_image(struct radv_cmd_buffer *cmd_buffer, struct radv_meta_blit2d_
 
    radv_CmdEndRendering2KHR(radv_cmd_buffer_to_handle(cmd_buffer), &end_info);
 
-   fail:
+fail_pipeline:
+   if (dst->aspect_mask == (VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT)) {
+      radv_image_view_finish(&src_iview_depth);
+      radv_image_view_finish(&src_iview_stencil);
+   } else {
       radv_image_view_finish(&src_iview);
-      radv_image_view_finish(&dst_iview);
    }
+   radv_image_view_finish(&dst_iview);
 }
 
 struct radv_blit2d_key {
@@ -405,18 +447,27 @@ create_layout(struct radv_device *device, int idx, VkPipelineLayout *layout_out)
    key.type = RADV_META_OBJECT_KEY_BLIT2D;
    key.index = idx;
 
-   const VkDescriptorSetLayoutBinding binding = {
+   const VkDescriptorSetLayoutBinding bindings[] = {
+      {
          .binding = 0,
          .descriptorType = desc_type,
          .descriptorCount = 1,
          .stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT,
+      },
+      {
+         .binding = 1,
+         .descriptorType = desc_type,
+         .descriptorCount = 1,
+         .stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT,
+      },
+
    };
 
    const VkDescriptorSetLayoutCreateInfo desc_info = {
       .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
       .flags = VK_DESCRIPTOR_SET_LAYOUT_CREATE_PUSH_DESCRIPTOR_BIT,
-      .bindingCount = 1,
+      .bindingCount = 2,
-      .pBindings = &binding,
+      .pBindings = bindings,
    };
 
    const VkPushConstantRange pc_range = {
@@ -886,3 +937,153 @@ get_stencil_only_pipeline(struct radv_device *device, enum blit2d_src_type src_t
    ralloc_free(fs_module);
    return result;
 }
+
+static VkResult
+get_depth_stencil_pipeline(struct radv_device *device, enum blit2d_src_type src_type, uint32_t log2_samples,
+                           VkPipeline *pipeline_out, VkPipelineLayout *layout_out)
+{
+   struct radv_blit2d_ds_key key;
+   const char *name;
+   VkResult result;
+
+   result = create_layout(device, src_type, layout_out);
+   if (result != VK_SUCCESS)
+      return result;
+
+   memset(&key, 0, sizeof(key));
+   key.type = RADV_META_OBJECT_KEY_BLIT2D_DEPTH_STENCIL;
+   key.src_type = src_type;
+   key.log2_samples = log2_samples;
+
+   VkPipeline pipeline_from_cache = vk_meta_lookup_pipeline(&device->meta_state.device, &key, sizeof(key));
+   if (pipeline_from_cache != VK_NULL_HANDLE) {
+      *pipeline_out = pipeline_from_cache;
+      return VK_SUCCESS;
+   }
+
+   radv_meta_nir_texel_fetch_build_func src_func;
+   switch (src_type) {
+   case BLIT2D_SRC_TYPE_IMAGE:
+      src_func = radv_meta_nir_build_blit2d_texel_fetch;
+      name = "meta_blit2d_depth_stencil_image_fs";
+      break;
+   case BLIT2D_SRC_TYPE_IMAGE_3D:
+      src_func = radv_meta_nir_build_blit2d_texel_fetch;
+      name = "meta_blit3d_depth_stencil_image_fs";
+      break;
+   default:
+      UNREACHABLE("unknown blit src type\n");
+      break;
+   }
+
+   nir_shader *vs_module = radv_meta_nir_build_blit2d_vertex_shader(device);
+   nir_shader *fs_module = radv_meta_nir_build_blit2d_copy_fragment_shader_depth_stencil(
+      device, src_func, name, src_type == BLIT2D_SRC_TYPE_IMAGE_3D, log2_samples > 0);
+
+   const VkGraphicsPipelineCreateInfo pipeline_create_info = {
+      .sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO,
+      .stageCount = 2,
+      .pStages =
+         (VkPipelineShaderStageCreateInfo[]){
+            {.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
+             .stage = VK_SHADER_STAGE_VERTEX_BIT,
+             .module = vk_shader_module_handle_from_nir(vs_module),
+             .pName = "main",
+             .pSpecializationInfo = NULL},
+            {.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
+             .stage = VK_SHADER_STAGE_FRAGMENT_BIT,
+             .module = vk_shader_module_handle_from_nir(fs_module),
+             .pName = "main",
+             .pSpecializationInfo = NULL},
+         },
+      .pVertexInputState =
+         &(VkPipelineVertexInputStateCreateInfo){
+            .sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
+            .vertexBindingDescriptionCount = 0,
+            .vertexAttributeDescriptionCount = 0,
+         },
+      .pInputAssemblyState =
+         &(VkPipelineInputAssemblyStateCreateInfo){
+            .sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
+            .topology = VK_PRIMITIVE_TOPOLOGY_META_RECT_LIST_MESA,
+            .primitiveRestartEnable = false,
+         },
+      .pViewportState =
+         &(VkPipelineViewportStateCreateInfo){
+            .sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO,
+            .viewportCount = 1,
+            .scissorCount = 1,
+         },
+      .pRasterizationState =
+         &(VkPipelineRasterizationStateCreateInfo){.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO,
+                                                   .rasterizerDiscardEnable = false,
+                                                   .polygonMode = VK_POLYGON_MODE_FILL,
+                                                   .cullMode = VK_CULL_MODE_NONE,
+                                                   .frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE,
+                                                   .depthBiasConstantFactor = 0.0f,
+                                                   .depthBiasClamp = 0.0f,
+                                                   .depthBiasSlopeFactor = 0.0f,
+                                                   .lineWidth = 1.0f},
+      .pMultisampleState =
+         &(VkPipelineMultisampleStateCreateInfo){
+            .sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO,
+            .rasterizationSamples = 1 << log2_samples,
+            .sampleShadingEnable = false,
+            .pSampleMask = (VkSampleMask[]){UINT32_MAX},
+         },
+      .pColorBlendState =
+         &(VkPipelineColorBlendStateCreateInfo){
+            .sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO,
+            .attachmentCount = 0,
+            .pAttachments = NULL,
+            .blendConstants = {0.0f, 0.0f, 0.0f, 0.0f},
+         },
+      .pDepthStencilState =
+         &(VkPipelineDepthStencilStateCreateInfo){
+            .sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO,
+            .depthTestEnable = true,
+            .depthWriteEnable = true,
+            .stencilTestEnable = true,
+            .front = {.failOp = VK_STENCIL_OP_REPLACE,
+                      .passOp = VK_STENCIL_OP_REPLACE,
+                      .depthFailOp = VK_STENCIL_OP_REPLACE,
+                      .compareOp = VK_COMPARE_OP_ALWAYS,
+                      .compareMask = 0xff,
+                      .writeMask = 0xff,
+                      .reference = 0},
+            .back = {.failOp = VK_STENCIL_OP_REPLACE,
+                     .passOp = VK_STENCIL_OP_REPLACE,
+                     .depthFailOp = VK_STENCIL_OP_REPLACE,
+                     .compareOp = VK_COMPARE_OP_ALWAYS,
+                     .compareMask = 0xff,
+                     .writeMask = 0xff,
+                     .reference = 0},
+            .depthCompareOp = VK_COMPARE_OP_ALWAYS,
+            .minDepthBounds = 0.0f,
+            .maxDepthBounds = 1.0f,
+         },
+      .pDynamicState =
+         &(VkPipelineDynamicStateCreateInfo){
+            .sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO,
+            .dynamicStateCount = 2,
+            .pDynamicStates =
+               (VkDynamicState[]){
+                  VK_DYNAMIC_STATE_VIEWPORT,
+                  VK_DYNAMIC_STATE_SCISSOR,
+               },
+         },
+      .layout = *layout_out,
+   };
+
+   struct vk_meta_rendering_info render = {
+      .depth_attachment_format = VK_FORMAT_D32_SFLOAT,
+      .stencil_attachment_format = VK_FORMAT_S8_UINT,
+   };
+
+   result = vk_meta_create_graphics_pipeline(&device->vk, &device->meta_state.device, &pipeline_create_info, &render,
+                                             &key, sizeof(key), pipeline_out);
+
+   ralloc_free(vs_module);
+   ralloc_free(fs_module);
+   return result;
+}

src/amd/vulkan/nir/radv_meta_nir.c

@@ -484,6 +484,39 @@ radv_meta_nir_build_blit2d_copy_fragment_shader_stencil(struct radv_device *devi
    return b.shader;
 }
 
+nir_shader *
+radv_meta_nir_build_blit2d_copy_fragment_shader_depth_stencil(struct radv_device *device,
+                                                              radv_meta_nir_texel_fetch_build_func txf_func,
+                                                              const char *name, bool is_3d, bool is_multisampled)
+{
+   const struct glsl_type *vec4 = glsl_vec4_type();
+   const struct glsl_type *vec2 = glsl_vector_type(GLSL_TYPE_FLOAT, 2);
+   nir_builder b = radv_meta_nir_init_shader(device, MESA_SHADER_FRAGMENT, "%s", name);
+
+   nir_variable *tex_pos_in = nir_variable_create(b.shader, nir_var_shader_in, vec2, "v_tex_pos");
+   tex_pos_in->data.location = VARYING_SLOT_VAR0;
+
+   nir_def *pos_int = nir_f2i32(&b, nir_load_var(&b, tex_pos_in));
+   nir_def *tex_pos = nir_trim_vector(&b, pos_int, 2);
+
+   /* Depth */
+   nir_variable *depth_out = nir_variable_create(b.shader, nir_var_shader_out, vec4, "f_depth");
+   depth_out->data.location = FRAG_RESULT_DEPTH;
+
+   nir_def *depth = txf_func(&b, 0, tex_pos, is_3d, is_multisampled);
+   nir_store_var(&b, depth_out, depth, 0x1);
+
+   /* Stencil */
+   nir_variable *stencil_out = nir_variable_create(b.shader, nir_var_shader_out, vec4, "f_stencil");
+   stencil_out->data.location = FRAG_RESULT_STENCIL;
+
+   nir_def *stencil = txf_func(&b, 1, tex_pos, is_3d, is_multisampled);
+   nir_store_var(&b, stencil_out, stencil, 0x1);
+
+   b.shader->info.fs.uses_sample_shading = is_multisampled;
+   return b.shader;
+}
+
 nir_shader *
 radv_meta_nir_build_itob_compute_shader(struct radv_device *dev, bool is_3d)
 {

src/amd/vulkan/nir/radv_meta_nir.h

@@ -64,6 +64,10 @@ nir_shader *radv_meta_nir_build_blit2d_copy_fragment_shader_depth(struct radv_de
 nir_shader *radv_meta_nir_build_blit2d_copy_fragment_shader_stencil(struct radv_device *device,
                                                                     radv_meta_nir_texel_fetch_build_func txf_func,
                                                                     const char *name, bool is_3d, bool is_multisampled);
+nir_shader *radv_meta_nir_build_blit2d_copy_fragment_shader_depth_stencil(struct radv_device *device,
+                                                                          radv_meta_nir_texel_fetch_build_func txf_func,
+                                                                          const char *name, bool is_3d,
+                                                                          bool is_multisampled);
 
 void radv_meta_nir_build_clear_color_shaders(struct radv_device *dev, struct nir_shader **out_vs,
                                              struct nir_shader **out_fs, uint32_t frag_output);