radv: call vk_nir_lower_descriptor_heaps()

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

src/amd/vulkan/radv_pipeline.c

@@ -228,6 +228,13 @@ radv_pipeline_stage_init(VkPipelineCreateFlags2 pipeline_flags, const VkPipeline
    vk_pipeline_hash_shader_stage(pipeline_flags, sinfo, NULL, out_stage->shader_blake3);
 }
 
+void
+radv_pipeline_stage_finish(struct radv_shader_stage *stage)
+{
+   ralloc_free(stage->nir);
+   vk_sampler_state_array_finish(&stage->layout.embedded_samplers);
+}
+
 void
 radv_shader_layout_init(const struct radv_pipeline_layout *pipeline_layout, mesa_shader_stage stage,
                         struct radv_shader_layout *layout)

src/amd/vulkan/radv_pipeline.h

@@ -94,6 +94,8 @@ void radv_pipeline_stage_init(VkPipelineCreateFlags2 pipeline_flags, const VkPip
 void radv_shader_layout_init(const struct radv_pipeline_layout *pipeline_layout, mesa_shader_stage stage,
                              struct radv_shader_layout *layout);
 
+void radv_pipeline_stage_finish(struct radv_shader_stage *stage);
+
 void radv_postprocess_nir(struct radv_device *device, const struct radv_graphics_state_key *gfx_state,
                           struct radv_shader_stage *stage);
 

src/amd/vulkan/radv_pipeline_compute.c

@@ -226,7 +226,7 @@ radv_compute_pipeline_compile(const VkComputePipelineCreateInfo *pCreateInfo, st
    if (radv_can_dump_shader_stats(device, cs_stage.nir)) {
       radv_dump_shader_stats(device, &pipeline->base, pipeline->base.shaders[MESA_SHADER_COMPUTE], stderr);
    }
-   ralloc_free(cs_stage.nir);
+   radv_pipeline_stage_finish(&cs_stage);
 
 done:
    pipeline_feedback.duration = os_time_get_nano() - pipeline_start;

src/amd/vulkan/radv_pipeline_graphics.c

@@ -2969,7 +2969,7 @@ radv_graphics_pipeline_state_finish(struct radv_device *device, struct radv_grap
 
    if (gfx_state->stages) {
       for (uint32_t i = 0; i < MESA_VULKAN_SHADER_STAGES; i++)
-         ralloc_free(gfx_state->stages[i].nir);
+         radv_pipeline_stage_finish(&gfx_state->stages[i]);
 
       ralloc_free(gfx_state->stages[MESA_SHADER_GEOMETRY].gs_copy_shader);
 

src/amd/vulkan/radv_pipeline_rt.c

@@ -951,7 +951,7 @@ radv_rt_compile_shaders(struct radv_device *device, struct vk_pipeline_cache *ca
 
 cleanup:
    for (uint32_t i = 0; i < pCreateInfo->stageCount; i++)
-      ralloc_free(stages[i].nir);
+      radv_pipeline_stage_finish(&stages[i]);
    free(stages);
    return result;
 }

src/amd/vulkan/radv_shader.c

@@ -40,6 +40,8 @@
 #include "vk_debug_report.h"
 #include "vk_format.h"
 #include "vk_nir.h"
+#include "vk_nir_lower_descriptor_heaps.h"
+#include "vk_sampler.h"
 #include "vk_nir_convert_ycbcr.h"
 #include "vk_semaphore.h"
 #include "vk_sync.h"
@@ -448,11 +450,12 @@ ycbcr_conversion_lookup(const void *data, uint32_t set, uint32_t binding, uint32
 }
 
 nir_shader *
-radv_shader_spirv_to_nir(struct radv_device *device, const struct radv_shader_stage *stage,
+radv_shader_spirv_to_nir(struct radv_device *device, struct radv_shader_stage *stage,
                          const struct radv_spirv_to_nir_options *options, bool is_internal)
 {
    const struct radv_physical_device *pdev = radv_device_physical(device);
    struct radv_instance *instance = radv_physical_device_instance(pdev);
+   struct vk_sampler_state_array embedded_samplers;
    nir_shader *nir;
    bool progress;
 
@@ -775,6 +778,29 @@ radv_shader_spirv_to_nir(struct radv_device *device, const struct radv_shader_st
    if (lower_flrp != 0)
       NIR_PASS(progress, nir, nir_lower_flrp, lower_flrp, false /* always precise */);
 
+   if (stage->key.descriptor_heap) {
+      progress = false;
+      NIR_PASS(progress, nir, vk_nir_lower_descriptor_heaps, stage->layout.mapping, &embedded_samplers);
+      if (progress) {
+         NIR_PASS(_, nir, nir_remove_dead_variables, nir_var_uniform | nir_var_image, NULL);
+         NIR_PASS(_, nir, nir_opt_dce);
+
+         if (embedded_samplers.sampler_count > 0) {
+            /* Copy embedded samplers to the shader layout for easier uses. */
+            stage->layout.embedded_samplers.samplers =
+               malloc(embedded_samplers.sampler_count * sizeof(struct vk_sampler_state));
+            if (!stage->layout.embedded_samplers.samplers)
+               return NULL;
+
+            memcpy(stage->layout.embedded_samplers.samplers, embedded_samplers.samplers,
+                   embedded_samplers.sampler_count * sizeof(*embedded_samplers.samplers));
+            stage->layout.embedded_samplers.sampler_count = embedded_samplers.sampler_count;
+
+            vk_sampler_state_array_finish(&embedded_samplers);
+         }
+      }
+   }
+
    NIR_PASS(_, nir, nir_lower_explicit_io, nir_var_mem_push_const, nir_address_format_32bit_offset);
 
    NIR_PASS(_, nir, nir_lower_explicit_io, nir_var_mem_ubo | nir_var_mem_ssbo,

src/amd/vulkan/radv_shader.h

@@ -21,6 +21,7 @@
 #include "radv_constants.h"
 #include "radv_shader_args.h"
 #include "radv_shader_info.h"
+#include "vk_nir_lower_descriptor_heaps.h"
 #include "vk_pipeline_cache.h"
 
 #include "aco_shader_info.h"
@@ -265,6 +266,8 @@ struct radv_shader_layout {
    bool independent_sets;
 
    const VkShaderDescriptorSetAndBindingMappingInfoEXT *mapping;
+
+   struct vk_sampler_state_array embedded_samplers;
 };
 
 struct radv_shader_stage {
@@ -514,7 +517,7 @@ void radv_optimize_nir_algebraic(nir_shader *shader, bool opt_offsets, bool opt_
 
 struct radv_shader_stage;
 
-nir_shader *radv_shader_spirv_to_nir(struct radv_device *device, const struct radv_shader_stage *stage,
+nir_shader *radv_shader_spirv_to_nir(struct radv_device *device, struct radv_shader_stage *stage,
                                      const struct radv_spirv_to_nir_options *options, bool is_internal);
 
 void radv_init_shader_arenas(struct radv_device *device);