radv/rt: move stack_sizes into radv_ray_tracing_module

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/21667>

src/amd/vulkan/radv_pipeline.c

@@ -133,10 +133,6 @@ radv_pipeline_destroy(struct radv_device *device, struct radv_pipeline *pipeline
          radv_shader_part_unref(device, graphics_pipeline->ps_epilog);
 
       vk_free(&device->vk.alloc, graphics_pipeline->state_data);
-   } else if (pipeline->type == RADV_PIPELINE_RAY_TRACING) {
-      struct radv_ray_tracing_pipeline *rt_pipeline = radv_pipeline_to_ray_tracing(pipeline);
-
-      free(rt_pipeline->stack_sizes);
    } else if (pipeline->type == RADV_PIPELINE_LIBRARY) {
       struct radv_library_pipeline *library_pipeline = radv_pipeline_to_library(pipeline);
 
@@ -3514,7 +3510,7 @@ radv_graphics_pipeline_compile(struct radv_graphics_pipeline *pipeline,
 
    bool found_in_application_cache = true;
    if (!skip_shaders_cache &&
-       radv_create_shaders_from_pipeline_cache(device, cache, hash, &pipeline->base, NULL, NULL,
+       radv_create_shaders_from_pipeline_cache(device, cache, hash, &pipeline->base, NULL, 0,
                                                &found_in_application_cache)) {
       if (found_in_application_cache)
          pipeline_feedback.flags |= VK_PIPELINE_CREATION_FEEDBACK_APPLICATION_PIPELINE_CACHE_HIT_BIT;
@@ -5381,8 +5377,7 @@ radv_compute_pipeline_compile(struct radv_compute_pipeline *pipeline,
                               const VkPipelineShaderStageCreateInfo *pStage,
                               const VkPipelineCreateFlags flags, const uint8_t *custom_hash,
                               const VkPipelineCreationFeedbackCreateInfo *creation_feedback,
-                              struct radv_pipeline_shader_stack_size **stack_sizes,
+                              struct radv_ray_tracing_module *rt_groups, uint32_t num_rt_groups)
-                              uint32_t *num_stack_sizes)
 {
    struct radv_shader_binary *binaries[MESA_VULKAN_SHADER_STAGES] = {NULL};
    unsigned char hash[20];
@@ -5409,8 +5404,8 @@ radv_compute_pipeline_compile(struct radv_compute_pipeline *pipeline,
 
    bool found_in_application_cache = true;
    if (!keep_executable_info &&
-       radv_create_shaders_from_pipeline_cache(device, cache, hash, &pipeline->base, stack_sizes,
+       radv_create_shaders_from_pipeline_cache(device, cache, hash, &pipeline->base, rt_groups,
-                                               num_stack_sizes, &found_in_application_cache)) {
+                                               num_rt_groups, &found_in_application_cache)) {
       if (found_in_application_cache)
          pipeline_feedback.flags |=
             VK_PIPELINE_CREATION_FEEDBACK_APPLICATION_PIPELINE_CACHE_HIT_BIT;
@@ -5476,9 +5471,8 @@ radv_compute_pipeline_compile(struct radv_compute_pipeline *pipeline,
    }
 
    if (!keep_executable_info) {
-      radv_pipeline_cache_insert_shaders(device, cache, hash, &pipeline->base, binaries,
+      radv_pipeline_cache_insert_shaders(device, cache, hash, &pipeline->base, binaries, rt_groups,
-                                         stack_sizes ? *stack_sizes : NULL,
+                                         num_rt_groups);
-                                         num_stack_sizes ? *num_stack_sizes : 0);
    }
 
    free(binaries[MESA_SHADER_COMPUTE]);
@@ -5529,7 +5523,7 @@ radv_compute_pipeline_create(VkDevice _device, VkPipelineCache _cache,
 
    result = radv_compute_pipeline_compile(pipeline, pipeline_layout, device, cache, &key,
                                           &pCreateInfo->stage, pCreateInfo->flags, NULL,
-                                          creation_feedback, NULL, NULL);
+                                          creation_feedback, NULL, 0);
    if (result != VK_SUCCESS) {
       radv_pipeline_destroy(device, &pipeline->base, pAllocator);
       return result;

src/amd/vulkan/radv_pipeline_cache.c

@@ -323,10 +323,11 @@ radv_pipeline_cache_add_entry(struct radv_pipeline_cache *cache, struct cache_en
 }
 
 bool
-radv_create_shaders_from_pipeline_cache(
+radv_create_shaders_from_pipeline_cache(struct radv_device *device,
-   struct radv_device *device, struct radv_pipeline_cache *cache, const unsigned char *sha1,
+                                        struct radv_pipeline_cache *cache,
-   struct radv_pipeline *pipeline, struct radv_pipeline_shader_stack_size **stack_sizes,
+                                        const unsigned char *sha1, struct radv_pipeline *pipeline,
-   uint32_t *num_stack_sizes, bool *found_in_application_cache)
+                                        struct radv_ray_tracing_module *rt_groups,
+                                        uint32_t num_rt_groups, bool *found_in_application_cache)
 {
    struct cache_entry *entry;
 
@@ -402,17 +403,11 @@ radv_create_shaders_from_pipeline_cache(
       pipeline->shaders[MESA_SHADER_COMPUTE] = NULL;
    }
 
-   if (num_stack_sizes) {
+   assert(num_rt_groups == entry->num_stack_sizes);
-      *num_stack_sizes = entry->num_stack_sizes;
+   for (int i = 0; i < num_rt_groups; ++i) {
-      if (entry->num_stack_sizes) {
+      memcpy(&rt_groups[i].stack_size, p, sizeof(struct radv_pipeline_shader_stack_size));
-         *stack_sizes = malloc(entry->num_stack_sizes * sizeof(**stack_sizes));
-         memcpy(*stack_sizes, p, entry->num_stack_sizes * sizeof(**stack_sizes));
-      }
-   } else {
-      assert(!entry->num_stack_sizes);
    }
-
+   p += entry->num_stack_sizes * sizeof(struct radv_pipeline_shader_stack_size);
-   p += entry->num_stack_sizes * sizeof(**stack_sizes);
 
    if (device->instance->debug_flags & RADV_DEBUG_NO_MEMORY_CACHE && cache == device->mem_cache)
       vk_free(&cache->alloc, entry);
@@ -431,8 +426,8 @@ void
 radv_pipeline_cache_insert_shaders(struct radv_device *device, struct radv_pipeline_cache *cache,
                                    const unsigned char *sha1, struct radv_pipeline *pipeline,
                                    struct radv_shader_binary *const *binaries,
-                                   const struct radv_pipeline_shader_stack_size *stack_sizes,
+                                   const struct radv_ray_tracing_module *rt_groups,
-                                   uint32_t num_stack_sizes)
+                                   uint32_t num_rt_groups)
 {
    if (!cache)
       cache = device->mem_cache;
@@ -462,7 +457,7 @@ radv_pipeline_cache_insert_shaders(struct radv_device *device, struct radv_pipel
       return;
    }
 
-   size_t size = sizeof(*entry) + sizeof(*stack_sizes) * num_stack_sizes;
+   size_t size = sizeof(*entry) + sizeof(struct radv_pipeline_shader_stack_size) * num_rt_groups;
    for (int i = 0; i < MESA_VULKAN_SHADER_STAGES; ++i)
       if (pipeline->shaders[i])
          size += binaries[i]->total_size;
@@ -490,11 +485,11 @@ radv_pipeline_cache_insert_shaders(struct radv_device *device, struct radv_pipel
       p += binaries[i]->total_size;
    }
 
-   if (num_stack_sizes) {
+   for (int i = 0; i < num_rt_groups; ++i) {
-      memcpy(p, stack_sizes, sizeof(*stack_sizes) * num_stack_sizes);
+      memcpy(p, &rt_groups->stack_size, sizeof(struct radv_pipeline_shader_stack_size));
-      p += sizeof(*stack_sizes) * num_stack_sizes;
+      p += sizeof(struct radv_pipeline_shader_stack_size);
    }
-   entry->num_stack_sizes = num_stack_sizes;
+   entry->num_stack_sizes = num_rt_groups;
 
    // Make valgrind happy by filling the alignment hole at the end.
    assert(p == (char *)entry + size_without_align);

src/amd/vulkan/radv_pipeline_rt.c

@@ -354,7 +354,7 @@ radv_rt_pipeline_has_dynamic_stack_size(const VkRayTracingPipelineCreateInfoKHR
 
 static unsigned
 compute_rt_stack_size(const VkRayTracingPipelineCreateInfoKHR *pCreateInfo,
-                      const struct radv_pipeline_shader_stack_size *stack_sizes)
+                      const struct radv_ray_tracing_module *groups)
 {
    if (radv_rt_pipeline_has_dynamic_stack_size(pCreateInfo))
       return -1u;
@@ -366,11 +366,11 @@ compute_rt_stack_size(const VkRayTracingPipelineCreateInfoKHR *pCreateInfo,
    unsigned non_recursive_size = 0;
 
    for (unsigned i = 0; i < pCreateInfo->groupCount; ++i) {
-      non_recursive_size = MAX2(stack_sizes[i].non_recursive_size, non_recursive_size);
+      non_recursive_size = MAX2(groups[i].stack_size.non_recursive_size, non_recursive_size);
 
       const VkRayTracingShaderGroupCreateInfoKHR *group_info = &pCreateInfo->pGroups[i];
       uint32_t shader_id = VK_SHADER_UNUSED_KHR;
-      unsigned size = stack_sizes[i].recursive_size;
+      unsigned size = groups[i].stack_size.recursive_size;
 
       switch (group_info->type) {
       case VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR:
@@ -482,9 +482,9 @@ radv_rt_pipeline_create(VkDevice _device, VkPipelineCache _cache,
 
    /* First check if we can get things from the cache before we take the expensive step of
     * generating the nir. */
-   result = radv_compute_pipeline_compile(&rt_pipeline->base, pipeline_layout, device, cache,
+   result = radv_compute_pipeline_compile(&rt_pipeline->base, pipeline_layout, device, cache, &key,
-                                          &key, &stage, flags, hash, creation_feedback,
+                                          &stage, flags, hash, creation_feedback,
-                                          &rt_pipeline->stack_sizes, &rt_pipeline->group_count);
+                                          rt_pipeline->groups, rt_pipeline->group_count);
 
    if (result != VK_SUCCESS && result != VK_PIPELINE_COMPILE_REQUIRED)
       goto pipeline_fail;
@@ -493,26 +493,18 @@ radv_rt_pipeline_create(VkDevice _device, VkPipelineCache _cache,
       if (pCreateInfo->flags & VK_PIPELINE_CREATE_FAIL_ON_PIPELINE_COMPILE_REQUIRED_BIT)
          goto pipeline_fail;
 
-      rt_pipeline->stack_sizes =
+      shader = create_rt_shader(device, &local_create_info, rt_pipeline->groups, &key);
-         calloc(sizeof(*rt_pipeline->stack_sizes), local_create_info.groupCount);
-      if (!rt_pipeline->stack_sizes) {
-         result = VK_ERROR_OUT_OF_HOST_MEMORY;
-         goto pipeline_fail;
-      }
-
-      shader = create_rt_shader(device, &local_create_info, rt_pipeline->stack_sizes,
-                                rt_pipeline->groups, &key);
       module.nir = shader;
       result = radv_compute_pipeline_compile(
-         &rt_pipeline->base, pipeline_layout, device, cache, &key, &stage, pCreateInfo->flags,
+         &rt_pipeline->base, pipeline_layout, device, cache, &key, &stage, pCreateInfo->flags, hash,
-         hash, creation_feedback, &rt_pipeline->stack_sizes, &rt_pipeline->group_count);
+         creation_feedback, rt_pipeline->groups, rt_pipeline->group_count);
       if (result != VK_SUCCESS)
          goto shader_fail;
    }
 
    radv_compute_pipeline_init(&rt_pipeline->base, pipeline_layout);
 
-   rt_pipeline->stack_size = compute_rt_stack_size(pCreateInfo, rt_pipeline->stack_sizes);
+   rt_pipeline->stack_size = compute_rt_stack_size(pCreateInfo, rt_pipeline->groups);
 
    *pPipeline = radv_pipeline_to_handle(&rt_pipeline->base.base);
 
@@ -598,7 +590,8 @@ radv_GetRayTracingShaderGroupStackSizeKHR(VkDevice device, VkPipeline _pipeline,
 {
    RADV_FROM_HANDLE(radv_pipeline, pipeline, _pipeline);
    struct radv_ray_tracing_pipeline *rt_pipeline = radv_pipeline_to_ray_tracing(pipeline);
-   const struct radv_pipeline_shader_stack_size *stack_size = &rt_pipeline->stack_sizes[group];
+   const struct radv_pipeline_shader_stack_size *stack_size =
+      &rt_pipeline->groups[group].stack_size;
 
    if (groupShader == VK_SHADER_GROUP_SHADER_ANY_HIT_KHR ||
        groupShader == VK_SHADER_GROUP_SHADER_INTERSECTION_KHR)

src/amd/vulkan/radv_private.h

@@ -412,13 +412,15 @@ struct radv_pipeline_shader_stack_size;
 
 bool radv_create_shaders_from_pipeline_cache(
    struct radv_device *device, struct radv_pipeline_cache *cache, const unsigned char *sha1,
-   struct radv_pipeline *pipeline, struct radv_pipeline_shader_stack_size **stack_sizes,
+   struct radv_pipeline *pipeline, struct radv_ray_tracing_module *rt_groups,
-   uint32_t *num_stack_sizes, bool *found_in_application_cache);
+   uint32_t num_rt_groups, bool *found_in_application_cache);
 
-void radv_pipeline_cache_insert_shaders(
+void radv_pipeline_cache_insert_shaders(struct radv_device *device,
-   struct radv_device *device, struct radv_pipeline_cache *cache, const unsigned char *sha1,
+                                        struct radv_pipeline_cache *cache,
-   struct radv_pipeline *pipeline, struct radv_shader_binary *const *binaries,
+                                        const unsigned char *sha1, struct radv_pipeline *pipeline,
-   const struct radv_pipeline_shader_stack_size *stack_sizes, uint32_t num_stack_sizes);
+                                        struct radv_shader_binary *const *binaries,
+                                        const struct radv_ray_tracing_module *rt_groups,
+                                        uint32_t num_rt_groups);
 
 enum radv_blit_ds_layout {
    RADV_BLIT_DS_LAYOUT_TILE_ENABLE,
@@ -2206,6 +2208,7 @@ struct radv_compute_pipeline {
 
 struct radv_ray_tracing_module {
    struct radv_pipeline_group_handle handle;
+   struct radv_pipeline_shader_stack_size stack_size;
 };
 
 struct radv_library_pipeline {
@@ -2239,7 +2242,6 @@ struct radv_graphics_lib_pipeline {
 struct radv_ray_tracing_pipeline {
    struct radv_compute_pipeline base;
 
-   struct radv_pipeline_shader_stack_size *stack_sizes;
    uint32_t group_count;
    uint32_t stack_size;
    struct radv_ray_tracing_module groups[];

src/amd/vulkan/radv_rt_shader.c

@@ -124,7 +124,7 @@ struct rt_variables {
    nir_variable *ahit_terminate;
 
    /* Array of stack size struct for recording the max stack size for each group. */
-   struct radv_pipeline_shader_stack_size *stack_sizes;
+   struct radv_ray_tracing_module *groups;
    unsigned stage_idx;
 };
 
@@ -135,19 +135,18 @@ reserve_stack_size(struct rt_variables *vars, uint32_t size)
       const VkRayTracingShaderGroupCreateInfoKHR *group = vars->create_info->pGroups + group_idx;
 
       if (vars->stage_idx == group->generalShader || vars->stage_idx == group->closestHitShader)
-         vars->stack_sizes[group_idx].recursive_size =
+         vars->groups[group_idx].stack_size.recursive_size =
-            MAX2(vars->stack_sizes[group_idx].recursive_size, size);
+            MAX2(vars->groups[group_idx].stack_size.recursive_size, size);
 
       if (vars->stage_idx == group->anyHitShader || vars->stage_idx == group->intersectionShader)
-         vars->stack_sizes[group_idx].non_recursive_size =
+         vars->groups[group_idx].stack_size.non_recursive_size =
-            MAX2(vars->stack_sizes[group_idx].non_recursive_size, size);
+            MAX2(vars->groups[group_idx].stack_size.non_recursive_size, size);
    }
 }
 
 static struct rt_variables
 create_rt_variables(nir_shader *shader, const VkRayTracingPipelineCreateInfoKHR *create_info,
-                    struct radv_pipeline_shader_stack_size *stack_sizes,
+                    struct radv_ray_tracing_module *groups, const struct radv_pipeline_key *key)
-                    const struct radv_pipeline_key *key)
 {
    struct rt_variables vars = {
       .create_info = create_info,
@@ -193,7 +192,7 @@ create_rt_variables(nir_shader *shader, const VkRayTracingPipelineCreateInfoKHR
    vars.ahit_terminate =
       nir_variable_create(shader, nir_var_shader_temp, glsl_bool_type(), "ahit_terminate");
 
-   vars.stack_sizes = stack_sizes;
+   vars.groups = groups;
    return vars;
 }
 
@@ -231,7 +230,7 @@ map_rt_variables(struct hash_table *var_remap, struct rt_variables *src,
    _mesa_hash_table_insert(var_remap, src->ahit_accept, dst->ahit_accept);
    _mesa_hash_table_insert(var_remap, src->ahit_terminate, dst->ahit_terminate);
 
-   src->stack_sizes = dst->stack_sizes;
+   src->groups = dst->groups;
    src->stage_idx = dst->stage_idx;
 }
 
@@ -828,7 +827,7 @@ insert_rt_case(nir_builder *b, nir_shader *shader, struct rt_variables *vars, ni
    nir_opt_dead_cf(shader);
 
    struct rt_variables src_vars =
-      create_rt_variables(shader, vars->create_info, vars->stack_sizes, vars->key);
+      create_rt_variables(shader, vars->create_info, vars->groups, vars->key);
    map_rt_variables(var_remap, &src_vars, vars);
 
    NIR_PASS_V(shader, lower_rt_instructions, &src_vars, call_idx_base);
@@ -1379,9 +1378,7 @@ load_stack_entry(nir_builder *b, nir_ssa_def *index, const struct radv_ray_trave
 static nir_shader *
 build_traversal_shader(struct radv_device *device,
                        const VkRayTracingPipelineCreateInfoKHR *pCreateInfo,
-                       struct radv_pipeline_shader_stack_size *stack_sizes,
+                       struct radv_ray_tracing_module *groups, const struct radv_pipeline_key *key)
-                       const struct radv_ray_tracing_module *groups,
-                       const struct radv_pipeline_key *key)
 {
    /* Create the traversal shader as an intersection shader to prevent validation failures due to
     * invalid variable modes.*/
@@ -1391,7 +1388,7 @@ build_traversal_shader(struct radv_device *device,
    b.shader->info.workgroup_size[1] = device->physical_device->rt_wave_size == 64 ? 8 : 4;
    b.shader->info.shared_size =
       device->physical_device->rt_wave_size * MAX_STACK_ENTRY_COUNT * sizeof(uint32_t);
-   struct rt_variables vars = create_rt_variables(b.shader, pCreateInfo, stack_sizes, key);
+   struct rt_variables vars = create_rt_variables(b.shader, pCreateInfo, groups, key);
 
    /* Register storage for hit attributes */
    nir_variable *hit_attribs[RADV_MAX_HIT_ATTRIB_SIZE / sizeof(uint32_t)];
@@ -1578,8 +1575,7 @@ move_rt_instructions(nir_shader *shader)
 
 nir_shader *
 create_rt_shader(struct radv_device *device, const VkRayTracingPipelineCreateInfoKHR *pCreateInfo,
-                 struct radv_pipeline_shader_stack_size *stack_sizes,
+                 struct radv_ray_tracing_module *groups, const struct radv_pipeline_key *key)
-                 const struct radv_ray_tracing_module *groups, const struct radv_pipeline_key *key)
 {
    nir_builder b = radv_meta_init_shader(device, MESA_SHADER_COMPUTE, "rt_combined");
    b.shader->info.internal = false;
@@ -1587,7 +1583,7 @@ create_rt_shader(struct radv_device *device, const VkRayTracingPipelineCreateInf
    b.shader->info.workgroup_size[1] = device->physical_device->rt_wave_size == 64 ? 8 : 4;
    b.shader->info.shared_size = device->physical_device->rt_wave_size * RADV_MAX_HIT_ATTRIB_SIZE;
 
-   struct rt_variables vars = create_rt_variables(b.shader, pCreateInfo, stack_sizes, key);
+   struct rt_variables vars = create_rt_variables(b.shader, pCreateInfo, groups, key);
    load_sbt_entry(&b, &vars, nir_imm_int(&b, 0), SBT_RAYGEN, SBT_GENERAL_IDX);
    nir_store_var(&b, vars.stack_ptr, nir_load_rt_dynamic_callable_stack_base_amd(&b), 0x1);
 
@@ -1611,7 +1607,7 @@ create_rt_shader(struct radv_device *device, const VkRayTracingPipelineCreateInf
    nir_ssa_def *idx = nir_load_var(&b, vars.idx);
 
    /* Insert traversal shader */
-   nir_shader *traversal = build_traversal_shader(device, pCreateInfo, stack_sizes, groups, key);
+   nir_shader *traversal = build_traversal_shader(device, pCreateInfo, groups, key);
    b.shader->info.shared_size = MAX2(b.shader->info.shared_size, traversal->info.shared_size);
    assert(b.shader->info.shared_size <= 32768);
    insert_rt_case(&b, traversal, &vars, idx, 0, 1);

src/amd/vulkan/radv_shader.h

@@ -565,7 +565,7 @@ VkResult radv_compute_pipeline_compile(
    const struct radv_pipeline_key *pipeline_key, const VkPipelineShaderStageCreateInfo *pStage,
    const VkPipelineCreateFlags flags, const uint8_t *custom_hash,
    const VkPipelineCreationFeedbackCreateInfo *creation_feedback,
-   struct radv_pipeline_shader_stack_size **stack_sizes, uint32_t *num_stack_sizes);
+   struct radv_ray_tracing_module *rt_groups, uint32_t num_rt_groups);
 
 struct radv_shader_args;
 
@@ -750,8 +750,7 @@ bool radv_lower_fs_intrinsics(nir_shader *nir, const struct radv_pipeline_stage
 
 nir_shader *create_rt_shader(struct radv_device *device,
                              const VkRayTracingPipelineCreateInfoKHR *pCreateInfo,
-                             struct radv_pipeline_shader_stack_size *stack_sizes,
+                             struct radv_ray_tracing_module *groups,
-                             const struct radv_ray_tracing_module *groups,
                              const struct radv_pipeline_key *key);
 
 #endif