From 8e2869fa4114622a0f63c439c2dec110ed231c97 Mon Sep 17 00:00:00 2001 From: Samuel Pitoiset Date: Tue, 14 Apr 2026 17:21:07 +0200 Subject: [PATCH] vulkan: add an option to lower SHADER_RECORD_INDEX to non-uniform Applications are required to set NonUniform if the resource is arrayed, but with VK_DESCRIPTOR_MAPPING_SOURCE_HEAP_WITH_SHADER_RECORD_INDEX_EXT, the resource is non-arrayed in the shader. So, it's technically not required to set it. Although, the offset can vary per-lane and NonUniform is implicit. Backport-to: 26.1 Signed-off-by: Samuel Pitoiset Part-of: --- src/amd/vulkan/radv_shader.c | 2 +- .../runtime/vk_nir_lower_descriptor_heaps.c | 64 +++++++++++++++---- .../runtime/vk_nir_lower_descriptor_heaps.h | 6 ++ src/vulkan/runtime/vk_pipeline.c | 2 +- src/vulkan/runtime/vk_shader.c | 2 +- 5 files changed, 59 insertions(+), 17 deletions(-) diff --git a/src/amd/vulkan/radv_shader.c b/src/amd/vulkan/radv_shader.c index 0e3c0647779..691ef4a2d86 100644 --- a/src/amd/vulkan/radv_shader.c +++ b/src/amd/vulkan/radv_shader.c @@ -800,7 +800,7 @@ radv_shader_spirv_to_nir(struct radv_device *device, struct radv_shader_stage *s if (stage->key.descriptor_heap) { progress = false; - NIR_PASS(progress, nir, vk_nir_lower_descriptor_heaps, stage->layout.mapping, &embedded_samplers); + NIR_PASS(progress, nir, vk_nir_lower_descriptor_heaps, stage->layout.mapping, NULL, &embedded_samplers); if (progress) { NIR_PASS(_, nir, nir_remove_dead_variables, nir_var_uniform | nir_var_image, NULL); NIR_PASS(_, nir, nir_opt_dce); diff --git a/src/vulkan/runtime/vk_nir_lower_descriptor_heaps.c b/src/vulkan/runtime/vk_nir_lower_descriptor_heaps.c index e25852e63df..e1fdf75d999 100644 --- a/src/vulkan/runtime/vk_nir_lower_descriptor_heaps.c +++ b/src/vulkan/runtime/vk_nir_lower_descriptor_heaps.c @@ -162,6 +162,8 @@ vk_hash_descriptor_heap_mappings( struct heap_mapping_ctx { const VkShaderDescriptorSetAndBindingMappingInfoEXT *info; + const vk_nir_lower_descriptor_heaps_options *options; + /* Map from vk_sampler_state to indices */ struct hash_table *sampler_idx_map; }; @@ -242,15 +244,33 @@ unpack_combined_image_sampler(nir_builder *b, nir_def *combined, return nir_ubitfield_extract_imm(b, combined, 0, 20); } +static bool +is_mapping_implicitly_non_uniform(struct heap_mapping_ctx *ctx, + const VkDescriptorSetAndBindingMappingEXT *mapping, + nir_def *index) +{ + /* Non-arrayed resources backed by HEAP_WITH_SHADER_RECORD_INDEX can be + * implicitly non-uniform: different lanes in a subgroup may have different + * shader record indices (and thus different heap entries) with no + * descriptor indexing in the shader to annotate it. + */ + return mapping->source == VK_DESCRIPTOR_MAPPING_SOURCE_HEAP_WITH_SHADER_RECORD_INDEX_EXT && + index == NULL && ctx->options && ctx->options->lower_shader_record_index_to_non_uniform; +} + static nir_def * vk_build_descriptor_heap_offset(nir_builder *b, + struct heap_mapping_ctx *ctx, const VkDescriptorSetAndBindingMappingEXT *mapping, nir_resource_type resource_type, uint32_t binding, nir_def *index, - bool is_sampler) + bool is_sampler, bool *non_uniform_out) { assert(util_is_power_of_two_nonzero(resource_type)); + if (non_uniform_out != NULL) + *non_uniform_out = is_mapping_implicitly_non_uniform(ctx, mapping, index); + if (index == NULL) index = nir_imm_int(b, 0); @@ -601,7 +621,8 @@ build_buffer_addr_for_deref(nir_builder *b, nir_def *root_addr, /* The cursor is not where you left it when this function returns. */ static nir_def * build_deref_heap_offset(nir_builder *b, nir_deref_instr *deref, - bool is_sampler, struct heap_mapping_ctx *ctx) + bool is_sampler, struct heap_mapping_ctx *ctx, + bool *non_uniform_out) { uint32_t set, binding; nir_resource_type resource_type; @@ -619,11 +640,9 @@ build_deref_heap_offset(nir_builder *b, nir_deref_instr *deref, b->cursor = nir_before_instr(&deref->instr); - if (index == NULL) - index = nir_imm_int(b, 0); - - return vk_build_descriptor_heap_offset(b, mapping, resource_type, - binding, index, is_sampler); + return vk_build_descriptor_heap_offset(b, ctx, mapping, resource_type, + binding, index, is_sampler, + non_uniform_out); } else { nir_deref_instr *root_cast = deref_get_root_cast(deref); if (root_cast == NULL) @@ -671,10 +690,13 @@ lower_heaps_tex(nir_builder *b, nir_tex_instr *tex, assert(texture != NULL); { - nir_def *heap_offset = build_deref_heap_offset(b, texture, false, ctx); + bool texture_non_uniform = false; + nir_def *heap_offset = build_deref_heap_offset(b, texture, false, ctx, + &texture_non_uniform); if (heap_offset != NULL) { nir_src_rewrite(&tex->src[texture_src_idx].src, heap_offset); tex->src[texture_src_idx].src_type = nir_tex_src_texture_heap_offset; + tex->texture_non_uniform |= texture_non_uniform; progress = true; } } @@ -689,10 +711,13 @@ lower_heaps_tex(nir_builder *b, nir_tex_instr *tex, const VkSamplerCreateInfo *embedded_sampler = get_deref_embedded_sampler(sampler, ctx); if (embedded_sampler == NULL) { - nir_def *heap_offset = build_deref_heap_offset(b, sampler, true, ctx); + bool sampler_non_uniform = false; + nir_def *heap_offset = build_deref_heap_offset(b, sampler, true, ctx, + &sampler_non_uniform); if (heap_offset != NULL) { nir_src_rewrite(&tex->src[sampler_src_idx].src, heap_offset); tex->src[sampler_src_idx].src_type = nir_tex_src_sampler_heap_offset; + tex->sampler_non_uniform |= sampler_non_uniform; progress = true; } } else { @@ -719,7 +744,9 @@ lower_heaps_image(nir_builder *b, nir_intrinsic_instr *intrin, struct heap_mapping_ctx *ctx, bool deref) { nir_deref_instr *image = nir_src_as_deref(intrin->src[0]); - nir_def *heap_offset = build_deref_heap_offset(b, image, false, ctx); + bool is_non_uniform = false; + nir_def *heap_offset = build_deref_heap_offset(b, image, false, ctx, + &is_non_uniform); if (heap_offset == NULL) return false; @@ -729,6 +756,11 @@ lower_heaps_image(nir_builder *b, nir_intrinsic_instr *intrin, nir_src_rewrite(&intrin->src[0], heap_offset); } + if (is_non_uniform) { + nir_intrinsic_set_access(intrin, + nir_intrinsic_access(intrin) | ACCESS_NON_UNIFORM); + } + return true; } @@ -790,9 +822,10 @@ try_lower_heaps_deref_access(nir_builder *b, nir_intrinsic_instr *intrin, b->cursor = nir_before_instr(&desc_load->instr); nir_def *heap_offset = - vk_build_descriptor_heap_offset(b, mapping, resource_type, binding, + vk_build_descriptor_heap_offset(b, ctx, mapping, resource_type, binding, NULL /* index */, - false /* is_sampler */); + false /* is_sampler */, + NULL /* non_uniform_out */); /* This moves the cursor */ heap_offset = build_buffer_addr_for_deref(b, heap_offset, deref, @@ -922,8 +955,9 @@ lower_heaps_load_descriptor(nir_builder *b, nir_intrinsic_instr *desc_load, /* Everything else is an offset */ nir_def *heap_offset = - vk_build_descriptor_heap_offset(b, mapping, resource_type, binding, - index, false /* is_sampler */); + vk_build_descriptor_heap_offset(b, ctx, mapping, resource_type, binding, + index, false /* is_sampler */, + NULL /* non_uniform_out */); nir_def *desc = nir_load_heap_descriptor(b, desc_load->def.num_components, desc_load->def.bit_size, heap_offset, @@ -1001,10 +1035,12 @@ bool vk_nir_lower_descriptor_heaps( nir_shader *nir, const VkShaderDescriptorSetAndBindingMappingInfoEXT *mapping, + const vk_nir_lower_descriptor_heaps_options *options, struct vk_sampler_state_array *embedded_samplers_out) { struct heap_mapping_ctx ctx = { .info = mapping, + .options = options, .sampler_idx_map = _mesa_hash_table_create(NULL, hash_sampler, samplers_equal), }; diff --git a/src/vulkan/runtime/vk_nir_lower_descriptor_heaps.h b/src/vulkan/runtime/vk_nir_lower_descriptor_heaps.h index 37fc05aba47..6b344dc2d65 100644 --- a/src/vulkan/runtime/vk_nir_lower_descriptor_heaps.h +++ b/src/vulkan/runtime/vk_nir_lower_descriptor_heaps.h @@ -67,9 +67,15 @@ vk_sampler_state_array_finish(struct vk_sampler_state_array *arr) free(arr->samplers); } +typedef struct vk_nir_lower_descriptor_heaps_options { + /* Whether to lower non-arrayed resources backed by SRI to non-uniform. */ + bool lower_shader_record_index_to_non_uniform : 1; +} vk_nir_lower_descriptor_heaps_options; + bool vk_nir_lower_descriptor_heaps( nir_shader *nir, const VkShaderDescriptorSetAndBindingMappingInfoEXT *mapping, + const vk_nir_lower_descriptor_heaps_options *options, struct vk_sampler_state_array *embedded_samplers_out); #endif diff --git a/src/vulkan/runtime/vk_pipeline.c b/src/vulkan/runtime/vk_pipeline.c index db98ebf6144..c4acd592ed9 100644 --- a/src/vulkan/runtime/vk_pipeline.c +++ b/src/vulkan/runtime/vk_pipeline.c @@ -1010,7 +1010,7 @@ vk_pipeline_precompile_shader(struct vk_device *device, struct vk_sampler_state_array embedded_samplers; bool heaps_progress = false; NIR_PASS(heaps_progress, nir, vk_nir_lower_descriptor_heaps, - desc_map, &embedded_samplers); + desc_map, NULL, &embedded_samplers); if (heaps_progress) { NIR_PASS(_, nir, nir_remove_dead_variables, nir_var_uniform | nir_var_image, NULL); diff --git a/src/vulkan/runtime/vk_shader.c b/src/vulkan/runtime/vk_shader.c index f78d06e76f2..d3e9aad084e 100644 --- a/src/vulkan/runtime/vk_shader.c +++ b/src/vulkan/runtime/vk_shader.c @@ -283,7 +283,7 @@ vk_shader_to_nir(struct vk_device *device, bool heaps_progress = false; NIR_PASS(heaps_progress, nir, vk_nir_lower_descriptor_heaps, - desc_map, embedded_samplers_out); + desc_map, NULL, embedded_samplers_out); if (heaps_progress) { NIR_PASS(_, nir, nir_remove_dead_variables, nir_var_uniform | nir_var_image, NULL);