radv/gfx9: add support for 3d images to blit 2d paths

This add support for a 3D image reading path to the blit 2d paths,
like I did for the clear paths.

Fixes: e38685cc62e 'Revert "radv: disable support for VEGA for now."'
Reviewed-by: Bas Nieuwenhuizen <[email protected]>
Tested-by: Alex Smith <[email protected]>
Signed-off-by: Dave Airlie <[email protected]>

2 files changed, 97 insertions, 23 deletions

diff --git a/src/amd/vulkan/radv_meta_blit2d.c b/src/amd/vulkan/radv_meta_blit2d.c
index 30f58abb5f8..40d49f454dd 100644
--- a/src/amd/vulkan/radv_meta_blit2d.c
+++ b/src/amd/vulkan/radv_meta_blit2d.c
@@ -30,6 +30,7 @@
 
 enum blit2d_src_type {
 	BLIT2D_SRC_TYPE_IMAGE,
+	BLIT2D_SRC_TYPE_IMAGE_3D,
 	BLIT2D_SRC_TYPE_BUFFER,
 	BLIT2D_NUM_SRC_TYPES,
 };
@@ -41,6 +42,8 @@ create_iview(struct radv_cmd_buffer *cmd_buffer,
               VkImageAspectFlagBits aspects)
 {
 	VkFormat format;
+	VkImageViewType view_type = cmd_buffer->device->physical_device->rad_info.chip_class < GFX9 ? VK_IMAGE_VIEW_TYPE_2D :
+		radv_meta_get_view_type(surf->image);
 
 	if (depth_format)
 		format = depth_format;
@@ -51,7 +54,7 @@ create_iview(struct radv_cmd_buffer *cmd_buffer,
 			     &(VkImageViewCreateInfo) {
 				     .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
 					     .image = radv_image_to_handle(surf->image),
-					     .viewType = VK_IMAGE_VIEW_TYPE_2D,
+					     .viewType = view_type,
 					     .format = format,
 					     .subresourceRange = {
 					     .aspectMask = aspects,
@@ -126,6 +129,12 @@ blit2d_bind_src(struct radv_cmd_buffer *cmd_buffer,
 	} else {
 		create_iview(cmd_buffer, src_img, &tmp->iview, depth_format, aspects);
 
+		if (src_type == BLIT2D_SRC_TYPE_IMAGE_3D)
+			radv_CmdPushConstants(radv_cmd_buffer_to_handle(cmd_buffer),
+					      device->meta_state.blit2d.p_layouts[src_type],
+					      VK_SHADER_STAGE_FRAGMENT_BIT, 16, 4,
+					      &src_img->layer);
+
 		radv_meta_push_descriptor_set(cmd_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
 					      device->meta_state.blit2d.p_layouts[src_type],
 					      0, /* set */
@@ -341,8 +350,10 @@ radv_meta_blit2d(struct radv_cmd_buffer *cmd_buffer,
 		 unsigned num_rects,
 		 struct radv_meta_blit2d_rect *rects)
 {
+	bool use_3d = cmd_buffer->device->physical_device->rad_info.chip_class >= GFX9 &&
+		(src_img && src_img->image->type == VK_IMAGE_TYPE_3D);
 	enum blit2d_src_type src_type = src_buf ? BLIT2D_SRC_TYPE_BUFFER :
-						  BLIT2D_SRC_TYPE_IMAGE;
+		use_3d ? BLIT2D_SRC_TYPE_IMAGE_3D : BLIT2D_SRC_TYPE_IMAGE;
 	radv_meta_blit2d_normal_dst(cmd_buffer, src_img, src_buf, dst,
 				    num_rects, rects, src_type);
 }
@@ -407,29 +418,46 @@ build_nir_vertex_shader(void)
 
 typedef nir_ssa_def* (*texel_fetch_build_func)(struct nir_builder *,
                                                struct radv_device *,
-                                               nir_ssa_def *);
+                                               nir_ssa_def *, bool);
 
 static nir_ssa_def *
 build_nir_texel_fetch(struct nir_builder *b, struct radv_device *device,
-                      nir_ssa_def *tex_pos)
+                      nir_ssa_def *tex_pos, bool is_3d)
 {
+	enum glsl_sampler_dim dim = is_3d ? GLSL_SAMPLER_DIM_3D : GLSL_SAMPLER_DIM_2D;
 	const struct glsl_type *sampler_type =
-		glsl_sampler_type(GLSL_SAMPLER_DIM_2D, false, false, GLSL_TYPE_UINT);
+		glsl_sampler_type(dim, false, false, GLSL_TYPE_UINT);
 	nir_variable *sampler = nir_variable_create(b->shader, nir_var_uniform,
 						    sampler_type, "s_tex");
 	sampler->data.descriptor_set = 0;
 	sampler->data.binding = 0;
 
+	nir_ssa_def *tex_pos_3d = NULL;
+	if (is_3d) {
+		nir_intrinsic_instr *layer = nir_intrinsic_instr_create(b->shader, nir_intrinsic_load_push_constant);
+		nir_intrinsic_set_base(layer, 16);
+		nir_intrinsic_set_range(layer, 4);
+		layer->src[0] = nir_src_for_ssa(nir_imm_int(b, 0));
+		layer->num_components = 1;
+		nir_ssa_dest_init(&layer->instr, &layer->dest, 1, 32, "layer");
+		nir_builder_instr_insert(b, &layer->instr);
+
+		nir_ssa_def *chans[3];
+		chans[0] = nir_channel(b, tex_pos, 0);
+		chans[1] = nir_channel(b, tex_pos, 1);
+		chans[2] = &layer->dest.ssa;
+		tex_pos_3d = nir_vec(b, chans, 3);
+	}
 	nir_tex_instr *tex = nir_tex_instr_create(b->shader, 2);
-	tex->sampler_dim = GLSL_SAMPLER_DIM_2D;
+	tex->sampler_dim = dim;
 	tex->op = nir_texop_txf;
 	tex->src[0].src_type = nir_tex_src_coord;
-	tex->src[0].src = nir_src_for_ssa(tex_pos);
+	tex->src[0].src = nir_src_for_ssa(is_3d ? tex_pos_3d : tex_pos);
 	tex->src[1].src_type = nir_tex_src_lod;
 	tex->src[1].src = nir_src_for_ssa(nir_imm_int(b, 0));
 	tex->dest_type = nir_type_uint;
 	tex->is_array = false;
-	tex->coord_components = 2;
+	tex->coord_components = is_3d ? 3 : 2;
 	tex->texture = nir_deref_var_create(tex, sampler);
 	tex->sampler = NULL;
 
@@ -442,7 +470,7 @@ build_nir_texel_fetch(struct nir_builder *b, struct radv_device *device,
 
 static nir_ssa_def *
 build_nir_buffer_fetch(struct nir_builder *b, struct radv_device *device,
-                      nir_ssa_def *tex_pos)
+		       nir_ssa_def *tex_pos, bool is_3d)
 {
 	const struct glsl_type *sampler_type =
 		glsl_sampler_type(GLSL_SAMPLER_DIM_BUF, false, false, GLSL_TYPE_UINT);
@@ -490,7 +518,7 @@ static const VkPipelineVertexInputStateCreateInfo normal_vi_create_info = {
 
 static nir_shader *
 build_nir_copy_fragment_shader(struct radv_device *device,
-                               texel_fetch_build_func txf_func, const char* name)
+                               texel_fetch_build_func txf_func, const char* name, bool is_3d)
 {
 	const struct glsl_type *vec4 = glsl_vec4_type();
 	const struct glsl_type *vec2 = glsl_vector_type(GLSL_TYPE_FLOAT, 2);
@@ -511,7 +539,7 @@ build_nir_copy_fragment_shader(struct radv_device *device,
 	unsigned swiz[4] = { 0, 1 };
 	nir_ssa_def *tex_pos = nir_swizzle(&b, pos_int, swiz, 2, false);
 
-	nir_ssa_def *color = txf_func(&b, device, tex_pos);
+	nir_ssa_def *color = txf_func(&b, device, tex_pos, is_3d);
 	nir_store_var(&b, color_out, color, 0xf);
 
 	return b.shader;
@@ -519,7 +547,7 @@ build_nir_copy_fragment_shader(struct radv_device *device,
 
 static nir_shader *
 build_nir_copy_fragment_shader_depth(struct radv_device *device,
-				     texel_fetch_build_func txf_func, const char* name)
+				     texel_fetch_build_func txf_func, const char* name, bool is_3d)
 {
 	const struct glsl_type *vec4 = glsl_vec4_type();
 	const struct glsl_type *vec2 = glsl_vector_type(GLSL_TYPE_FLOAT, 2);
@@ -540,7 +568,7 @@ build_nir_copy_fragment_shader_depth(struct radv_device *device,
 	unsigned swiz[4] = { 0, 1 };
 	nir_ssa_def *tex_pos = nir_swizzle(&b, pos_int, swiz, 2, false);
 
-	nir_ssa_def *color = txf_func(&b, device, tex_pos);
+	nir_ssa_def *color = txf_func(&b, device, tex_pos, is_3d);
 	nir_store_var(&b, color_out, color, 0x1);
 
 	return b.shader;
@@ -548,7 +576,7 @@ build_nir_copy_fragment_shader_depth(struct radv_device *device,
 
 static nir_shader *
 build_nir_copy_fragment_shader_stencil(struct radv_device *device,
-				       texel_fetch_build_func txf_func, const char* name)
+				       texel_fetch_build_func txf_func, const char* name, bool is_3d)
 {
 	const struct glsl_type *vec4 = glsl_vec4_type();
 	const struct glsl_type *vec2 = glsl_vector_type(GLSL_TYPE_FLOAT, 2);
@@ -569,7 +597,7 @@ build_nir_copy_fragment_shader_stencil(struct radv_device *device,
 	unsigned swiz[4] = { 0, 1 };
 	nir_ssa_def *tex_pos = nir_swizzle(&b, pos_int, swiz, 2, false);
 
-	nir_ssa_def *color = txf_func(&b, device, tex_pos);
+	nir_ssa_def *color = txf_func(&b, device, tex_pos, is_3d);
 	nir_store_var(&b, color_out, color, 0x1);
 
 	return b.shader;
@@ -629,6 +657,10 @@ blit2d_init_color_pipeline(struct radv_device *device,
 		src_func = build_nir_texel_fetch;
 		name = "meta_blit2d_image_fs";
 		break;
+	case BLIT2D_SRC_TYPE_IMAGE_3D:
+		src_func = build_nir_texel_fetch;
+		name = "meta_blit3d_image_fs";
+		break;
 	case BLIT2D_SRC_TYPE_BUFFER:
 		src_func = build_nir_buffer_fetch;
 		name = "meta_blit2d_buffer_fs";
@@ -642,7 +674,7 @@ blit2d_init_color_pipeline(struct radv_device *device,
 	struct radv_shader_module fs = { .nir = NULL };
 
 
-	fs.nir = build_nir_copy_fragment_shader(device, src_func, name);
+	fs.nir = build_nir_copy_fragment_shader(device, src_func, name, src_type == BLIT2D_SRC_TYPE_IMAGE_3D);
 	vi_create_info = &normal_vi_create_info;
 
 	struct radv_shader_module vs = {
@@ -788,6 +820,10 @@ blit2d_init_depth_only_pipeline(struct radv_device *device,
 		src_func = build_nir_texel_fetch;
 		name = "meta_blit2d_depth_image_fs";
 		break;
+	case BLIT2D_SRC_TYPE_IMAGE_3D:
+		src_func = build_nir_texel_fetch;
+		name = "meta_blit3d_depth_image_fs";
+		break;
 	case BLIT2D_SRC_TYPE_BUFFER:
 		src_func = build_nir_buffer_fetch;
 		name = "meta_blit2d_depth_buffer_fs";
@@ -800,7 +836,7 @@ blit2d_init_depth_only_pipeline(struct radv_device *device,
 	const VkPipelineVertexInputStateCreateInfo *vi_create_info;
 	struct radv_shader_module fs = { .nir = NULL };
 
-	fs.nir = build_nir_copy_fragment_shader_depth(device, src_func, name);
+	fs.nir = build_nir_copy_fragment_shader_depth(device, src_func, name, src_type == BLIT2D_SRC_TYPE_IMAGE_3D);
 	vi_create_info = &normal_vi_create_info;
 
 	struct radv_shader_module vs = {
@@ -943,6 +979,10 @@ blit2d_init_stencil_only_pipeline(struct radv_device *device,
 		src_func = build_nir_texel_fetch;
 		name = "meta_blit2d_stencil_image_fs";
 		break;
+	case BLIT2D_SRC_TYPE_IMAGE_3D:
+		src_func = build_nir_texel_fetch;
+		name = "meta_blit3d_stencil_image_fs";
+		break;
 	case BLIT2D_SRC_TYPE_BUFFER:
 		src_func = build_nir_buffer_fetch;
 		name = "meta_blit2d_stencil_buffer_fs";
@@ -955,7 +995,7 @@ blit2d_init_stencil_only_pipeline(struct radv_device *device,
 	const VkPipelineVertexInputStateCreateInfo *vi_create_info;
 	struct radv_shader_module fs = { .nir = NULL };
 
-	fs.nir = build_nir_copy_fragment_shader_stencil(device, src_func, name);
+	fs.nir = build_nir_copy_fragment_shader_stencil(device, src_func, name, src_type == BLIT2D_SRC_TYPE_IMAGE_3D);
 	vi_create_info = &normal_vi_create_info;
 
 	struct radv_shader_module vs = {
@@ -1120,6 +1160,7 @@ VkResult
 radv_device_init_meta_blit2d_state(struct radv_device *device)
 {
 	VkResult result;
+	bool create_3d = device->physical_device->rad_info.chip_class >= GFX9;
 
 	const VkPushConstantRange push_constant_ranges[] = {
 		{VK_SHADER_STAGE_VERTEX_BIT, 0, 16},
@@ -1155,6 +1196,37 @@ radv_device_init_meta_blit2d_state(struct radv_device *device)
 	if (result != VK_SUCCESS)
 		goto fail;
 
+	if (create_3d) {
+		result = radv_CreateDescriptorSetLayout(radv_device_to_handle(device),
+							&(VkDescriptorSetLayoutCreateInfo) {
+								.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
+									.flags = VK_DESCRIPTOR_SET_LAYOUT_CREATE_PUSH_DESCRIPTOR_BIT_KHR,
+									.bindingCount = 1,
+									.pBindings = (VkDescriptorSetLayoutBinding[]) {
+									{
+										.binding = 0,
+										.descriptorType = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
+										.descriptorCount = 1,
+										.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT,
+										.pImmutableSamplers = NULL
+									},
+								}
+							}, &device->meta_state.alloc, &device->meta_state.blit2d.ds_layouts[BLIT2D_SRC_TYPE_IMAGE_3D]);
+		if (result != VK_SUCCESS)
+			goto fail;
+		result = radv_CreatePipelineLayout(radv_device_to_handle(device),
+						   &(VkPipelineLayoutCreateInfo) {
+							   .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
+								   .setLayoutCount = 1,
+								   .pSetLayouts = &device->meta_state.blit2d.ds_layouts[BLIT2D_SRC_TYPE_IMAGE_3D],
+								   .pushConstantRangeCount = 2,
+								   .pPushConstantRanges = push_constant_ranges,
+								   },
+						   &device->meta_state.alloc, &device->meta_state.blit2d.p_layouts[BLIT2D_SRC_TYPE_IMAGE_3D]);
+		if (result != VK_SUCCESS)
+			goto fail;
+	}
+
 	result = radv_CreateDescriptorSetLayout(radv_device_to_handle(device),
 						&(VkDescriptorSetLayoutCreateInfo) {
 							.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
@@ -1187,6 +1259,8 @@ radv_device_init_meta_blit2d_state(struct radv_device *device)
 		goto fail;
 
 	for (unsigned src = 0; src < BLIT2D_NUM_SRC_TYPES; src++) {
+		if (src == BLIT2D_SRC_TYPE_IMAGE_3D && !create_3d)
+			continue;
 		for (unsigned j = 0; j < ARRAY_SIZE(pipeline_formats); ++j) {
 			result = blit2d_init_color_pipeline(device, src, pipeline_formats[j]);
 			if (result != VK_SUCCESS)
diff --git a/src/amd/vulkan/radv_private.h b/src/amd/vulkan/radv_private.h
index 9b626379ee0..47ae0342ce2 100644
--- a/src/amd/vulkan/radv_private.h
+++ b/src/amd/vulkan/radv_private.h
@@ -401,15 +401,15 @@ struct radv_meta_state {
 	struct {
 		VkRenderPass render_passes[NUM_META_FS_KEYS];
 
-		VkPipelineLayout p_layouts[2];
-		VkDescriptorSetLayout ds_layouts[2];
-		VkPipeline pipelines[2][NUM_META_FS_KEYS];
+		VkPipelineLayout p_layouts[3];
+		VkDescriptorSetLayout ds_layouts[3];
+		VkPipeline pipelines[3][NUM_META_FS_KEYS];
 
 		VkRenderPass depth_only_rp;
-		VkPipeline depth_only_pipeline[2];
+		VkPipeline depth_only_pipeline[3];
 
 		VkRenderPass stencil_only_rp;
-		VkPipeline stencil_only_pipeline[2];
+		VkPipeline stencil_only_pipeline[3];
 	} blit2d;
 
 	struct {