4 files changed, 27 insertions, 11 deletions

diff --git a/src/amd/vulkan/radv_debug.c b/src/amd/vulkan/radv_debug.c
index 97ccfac694f..08ee54c9765 100644
--- a/src/amd/vulkan/radv_debug.c
+++ b/src/amd/vulkan/radv_debug.c
@@ -495,8 +495,14 @@ radv_dump_shader(struct radv_pipeline *pipeline,
 	if (!shader)
 		return;
 
-	fprintf(f, "%s:\n%s\n\n", radv_get_shader_name(shader, stage),
-		shader->disasm_string);
+	fprintf(f, "%s:\n\n", radv_get_shader_name(shader, stage));
+
+	if (shader->nir) {
+		fprintf(f, "NIR:\n");
+		nir_print_shader(shader->nir, f);
+	}
+
+	fprintf(stderr, "DISASM:\n%s\n", shader->disasm_string);
 
 	radv_shader_dump_stats(pipeline->device, shader, stage, f);
 }
diff --git a/src/amd/vulkan/radv_pipeline.c b/src/amd/vulkan/radv_pipeline.c
index 228fdf4ab07..5800f297074 100644
--- a/src/amd/vulkan/radv_pipeline.c
+++ b/src/amd/vulkan/radv_pipeline.c
@@ -139,7 +139,7 @@ radv_pipeline_compile(struct radv_pipeline *pipeline,
 		return NULL;
 
 	if (!variant) {
-		variant = radv_shader_variant_create(pipeline->device, nir,
+		variant = radv_shader_variant_create(pipeline->device, module, nir,
 						     layout, key, &code,
 						     &code_size);
 	}
@@ -163,7 +163,7 @@ radv_pipeline_compile(struct radv_pipeline *pipeline,
 
 		free(gs_copy_code);
 	}
-	if (!module->nir)
+	if (!module->nir && !pipeline->device->trace_bo)
 		ralloc_free(nir);
 
 	if (variant)
@@ -270,7 +270,7 @@ radv_tess_pipeline_compile(struct radv_pipeline *pipeline,
 	nir_lower_tes_patch_vertices(tes_nir,
 				     tcs_nir->info.tess.tcs_vertices_out);
 
-	tes_variant = radv_shader_variant_create(pipeline->device, tes_nir,
+	tes_variant = radv_shader_variant_create(pipeline->device, tes_module, tes_nir,
 						 layout, &tes_key, &tes_code,
 						 &tes_code_size);
 
@@ -282,14 +282,14 @@ radv_tess_pipeline_compile(struct radv_pipeline *pipeline,
 
 	radv_hash_shader(tcs_sha1, tcs_module, tcs_entrypoint, tcs_spec_info, layout, &tcs_key, 0);
 
-	tcs_variant = radv_shader_variant_create(pipeline->device, tcs_nir,
+	tcs_variant = radv_shader_variant_create(pipeline->device, tcs_module, tcs_nir,
 						 layout, &tcs_key, &tcs_code,
 						 &tcs_code_size);
 
-	if (!tes_module->nir)
+	if (!tes_module->nir && !pipeline->device->trace_bo)
 		ralloc_free(tes_nir);
 
-	if (!tcs_module->nir)
+	if (!tcs_module->nir && !pipeline->device->trace_bo)
 		ralloc_free(tcs_nir);
 
 	if (tes_variant)
diff --git a/src/amd/vulkan/radv_shader.c b/src/amd/vulkan/radv_shader.c
index ca0ad2d5758..349e291b06a 100644
--- a/src/amd/vulkan/radv_shader.c
+++ b/src/amd/vulkan/radv_shader.c
@@ -377,6 +377,7 @@ radv_fill_shader_variant(struct radv_device *device,
 
 static struct radv_shader_variant *
 shader_variant_create(struct radv_device *device,
+		      struct radv_shader_module *module,
 		      struct nir_shader *shader,
 		      gl_shader_stage stage,
 		      struct ac_nir_compiler_options *options,
@@ -430,6 +431,9 @@ shader_variant_create(struct radv_device *device,
 
 	if (device->trace_bo) {
 		variant->disasm_string = binary.disasm_string;
+		if (!gs_copy_shader && !module->nir) {
+			variant->nir = shader;
+		}
 	} else {
 		free(binary.disasm_string);
 	}
@@ -439,6 +443,7 @@ shader_variant_create(struct radv_device *device,
 
 struct radv_shader_variant *
 radv_shader_variant_create(struct radv_device *device,
+			   struct radv_shader_module *module,
 			   struct nir_shader *shader,
 			   struct radv_pipeline_layout *layout,
 			   const struct ac_shader_variant_key *key,
@@ -454,7 +459,7 @@ radv_shader_variant_create(struct radv_device *device,
 	options.unsafe_math = !!(device->debug_flags & RADV_DEBUG_UNSAFE_MATH);
 	options.supports_spill = device->llvm_supports_spill;
 
-	return shader_variant_create(device, shader, shader->stage,
+	return shader_variant_create(device, module, shader, shader->stage,
 				     &options, false, code_out, code_size_out);
 }
 
@@ -469,7 +474,7 @@ radv_create_gs_copy_shader(struct radv_device *device,
 
 	options.key.has_multiview_view_index = multiview;
 
-	return shader_variant_create(device, shader, MESA_SHADER_VERTEX,
+	return shader_variant_create(device, NULL, shader, MESA_SHADER_VERTEX,
 				     &options, true, code_out, code_size_out);
 }
 
@@ -484,6 +489,7 @@ radv_shader_variant_destroy(struct radv_device *device,
 	list_del(&variant->slab_list);
 	mtx_unlock(&device->shader_slab_mutex);
 
+	ralloc_free(variant->nir);
 	free(variant->disasm_string);
 	free(variant);
 }
diff --git a/src/amd/vulkan/radv_shader.h b/src/amd/vulkan/radv_shader.h
index 467557e5fc3..7ad38bf5792 100644
--- a/src/amd/vulkan/radv_shader.h
+++ b/src/amd/vulkan/radv_shader.h
@@ -49,7 +49,10 @@ struct radv_shader_variant {
 	struct ac_shader_variant_info info;
 	unsigned rsrc1;
 	unsigned rsrc2;
-	char *disasm_string; /* debug only */
+
+	/* debug only */
+	struct nir_shader *nir;
+	char *disasm_string;
 
 	struct list_head slab_list;
 };
@@ -78,6 +81,7 @@ radv_destroy_shader_slabs(struct radv_device *device);
 
 struct radv_shader_variant *
 radv_shader_variant_create(struct radv_device *device,
+			   struct radv_shader_module *module,
 			   struct nir_shader *shader,
 			   struct radv_pipeline_layout *layout,
 			   const struct ac_shader_variant_key *key,