From 98aad84d73c515bbe01821f1b4b541c882bcb92d Mon Sep 17 00:00:00 2001 From: Alyssa Rosenzweig Date: Sat, 12 Jul 2025 18:48:50 -0400 Subject: [PATCH] hk: push descriptor set addresses saves an indirection and sets us up for more goodness. Signed-off-by: Alyssa Rosenzweig Part-of: --- src/asahi/vulkan/hk_cmd_buffer.c | 16 +++++++-- src/asahi/vulkan/hk_nir_lower_descriptors.c | 18 +++------- src/asahi/vulkan/hk_shader.c | 40 ++++++++++++--------- src/asahi/vulkan/hk_shader.h | 1 + src/compiler/nir/nir_intrinsics.py | 5 +++ 5 files changed, 47 insertions(+), 33 deletions(-) diff --git a/src/asahi/vulkan/hk_cmd_buffer.c b/src/asahi/vulkan/hk_cmd_buffer.c index 92fab5452f3..9dd3353b089 100644 --- a/src/asahi/vulkan/hk_cmd_buffer.c +++ b/src/asahi/vulkan/hk_cmd_buffer.c @@ -677,12 +677,22 @@ hk_upload_usc_words(struct hk_cmd_buffer *cmd, struct hk_shader *s, struct agx_usc_builder b = agx_usc_builder(t.cpu, usc_size); - uint64_t root_ptr; + uint64_t root_ptr, set_ptr; - if (sw_stage == PIPE_SHADER_COMPUTE) + if (sw_stage == PIPE_SHADER_COMPUTE) { root_ptr = hk_cmd_buffer_upload_root(cmd, VK_PIPELINE_BIND_POINT_COMPUTE); - else + set_ptr = hk_pool_upload(cmd, cmd->state.cs.descriptors.root.sets, + sizeof(uint64_t) * s->info.set_count, 32); + } else { root_ptr = cmd->state.gfx.root; + set_ptr = hk_pool_upload(cmd, cmd->state.gfx.descriptors.root.sets, + sizeof(uint64_t) * s->info.set_count, 32); + } + + if (s->info.set_count) { + unsigned set_base = s->info.image_heap_uniform + 4; + agx_usc_uniform(&b, set_base, 4 * s->info.set_count, set_ptr); + } static_assert(offsetof(struct hk_root_descriptor_table, root_desc_addr) == 0, "self-reflective"); diff --git a/src/asahi/vulkan/hk_nir_lower_descriptors.c b/src/asahi/vulkan/hk_nir_lower_descriptors.c index b092e10d063..20eeaf7099b 100644 --- a/src/asahi/vulkan/hk_nir_lower_descriptors.c +++ b/src/asahi/vulkan/hk_nir_lower_descriptors.c @@ -84,16 +84,6 @@ lower_load_constant(nir_builder *b, nir_intrinsic_instr *load, return true; } -static nir_def * -load_descriptor_set_addr(nir_builder *b, uint32_t set, - UNUSED const struct lower_descriptors_ctx *ctx) -{ - uint32_t set_addr_offset = - hk_root_descriptor_offset(sets) + set * sizeof(uint64_t); - - return load_root(b, 1, 64, nir_imm_int(b, set_addr_offset), 8); -} - static nir_def * load_dynamic_buffer_start(nir_builder *b, uint32_t set, const struct lower_descriptors_ctx *ctx) @@ -152,8 +142,8 @@ load_descriptor(nir_builder *b, unsigned num_components, unsigned bit_size, } case VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK: { - nir_def *base_addr = nir_iadd_imm( - b, load_descriptor_set_addr(b, set, ctx), binding_layout->offset); + nir_def *base_addr = nir_iadd_imm(b, nir_load_descriptor_set_agx(b, set), + binding_layout->offset); assert(binding_layout->stride == 1); const uint32_t binding_size = binding_layout->array_size; @@ -177,7 +167,7 @@ load_descriptor(nir_builder *b, unsigned num_components, unsigned bit_size, desc_align_offset %= desc_align_mul; nir_def *desc; - nir_def *set_addr = load_descriptor_set_addr(b, set, ctx); + nir_def *set_addr = nir_load_descriptor_set_agx(b, set); desc = nir_load_global_constant_offset( b, num_components, bit_size, set_addr, desc_ubo_offset, .align_mul = desc_align_mul, .align_offset = desc_align_offset, @@ -715,7 +705,7 @@ lower_ssbo_resource_index(nir_builder *b, nir_intrinsic_instr *intrin, switch (binding_layout->type) { case VK_DESCRIPTOR_TYPE_MUTABLE_EXT: case VK_DESCRIPTOR_TYPE_STORAGE_BUFFER: { - nir_def *set_addr = load_descriptor_set_addr(b, set, ctx); + nir_def *set_addr = nir_load_descriptor_set_agx(b, set); binding_addr = nir_iadd_imm(b, set_addr, binding_layout->offset); binding_stride = binding_layout->stride; break; diff --git a/src/asahi/vulkan/hk_shader.c b/src/asahi/vulkan/hk_shader.c index 4237c5b0e51..946a80bc22d 100644 --- a/src/asahi/vulkan/hk_shader.c +++ b/src/asahi/vulkan/hk_shader.c @@ -944,6 +944,7 @@ hk_init_link_ht(struct hk_shader *shader, gl_shader_stage sw_stage) } struct fixed_uniforms { + unsigned sets; unsigned image_heap; unsigned root; }; @@ -952,7 +953,8 @@ static bool lower_uniforms(nir_builder *b, nir_intrinsic_instr *intr, void *data) { if (intr->intrinsic != nir_intrinsic_load_texture_handle_agx && - intr->intrinsic != nir_intrinsic_load_root_agx) + intr->intrinsic != nir_intrinsic_load_root_agx && + intr->intrinsic != nir_intrinsic_load_descriptor_set_agx) return false; b->cursor = nir_before_instr(&intr->instr); @@ -962,6 +964,9 @@ lower_uniforms(nir_builder *b, nir_intrinsic_instr *intr, void *data) if (intr->intrinsic == nir_intrinsic_load_texture_handle_agx) { rep = nir_bindless_image_agx(b, intr->src[0].ssa, .desc_set = ctx->image_heap); + } else if (intr->intrinsic == nir_intrinsic_load_descriptor_set_agx) { + unsigned s = nir_intrinsic_desc_set(intr); + rep = nir_load_preamble(b, 1, 64, .base = ctx->sets + (4 * s)); } else { rep = nir_load_preamble(b, 1, 64, .base = ctx->root); } @@ -1008,7 +1013,7 @@ hk_compile_nir(struct hk_device *dev, const VkAllocationCallbacks *pAllocator, const struct vk_pipeline_robustness_state *rs, const struct hk_fs_key *fs_key, enum hk_feature_key features, struct hk_shader *shader, gl_shader_stage sw_stage, bool hw, - nir_xfb_info *xfb_info) + nir_xfb_info *xfb_info, unsigned set_count) { unsigned nr_vbos = 0; @@ -1086,16 +1091,19 @@ hk_compile_nir(struct hk_device *dev, const VkAllocationCallbacks *pAllocator, } } - struct fixed_uniforms f = {.root = 0, .image_heap = 4}; + struct fixed_uniforms f = {.root = 0, .image_heap = 4, .sets = 8}; if (sw_stage == MESA_SHADER_FRAGMENT) { f.image_heap = AGX_ABI_FUNI_COUNT; f.root = AGX_ABI_FUNI_ROOT; + f.sets = AGX_ABI_FUNI_COUNT + 4; } else if (sw_stage == MESA_SHADER_VERTEX) { f.root = AGX_ABI_VUNI_COUNT_VK(nr_vbos); f.image_heap = f.root + 4; + f.sets = f.root + 8; } shader->info.image_heap_uniform = f.image_heap; + shader->info.set_count = set_count; /* XXX: rename */ NIR_PASS(_, nir, hk_lower_uvs_index, nr_vbos); @@ -1112,10 +1120,8 @@ hk_compile_nir(struct hk_device *dev, const VkAllocationCallbacks *pAllocator, #endif struct agx_shader_key backend_key = { - /* the image heap is always the last fixed uniform, so we can start - * preamble after that. - */ - .reserved_preamble = f.image_heap + 4, + /* Sets at the end */ + .reserved_preamble = f.image_heap + (4 * (set_count + 1)), .dev = agx_gather_device_key(&dev->dev), .no_stop = nir->info.stage == MESA_SHADER_FRAGMENT, @@ -1311,9 +1317,10 @@ hk_compile_shader(struct hk_device *dev, struct vk_shader_compile_info *info, for (unsigned v = 0; v < ARRAY_SIZE(variants); ++v) { if (variants[v].in) { - result = hk_compile_nir( - dev, pAllocator, variants[v].in, info->flags, info->robustness, - NULL, features, variants[v].out, sw_stage, true, NULL); + result = + hk_compile_nir(dev, pAllocator, variants[v].in, info->flags, + info->robustness, NULL, features, variants[v].out, + sw_stage, true, NULL, info->set_layout_count); if (result != VK_SUCCESS) { hk_api_shader_destroy(&dev->vk, &obj->vk, pAllocator); @@ -1379,9 +1386,10 @@ hk_compile_shader(struct hk_device *dev, struct vk_shader_compile_info *info, } /* hk_compile_nir takes ownership of the clone */ - result = hk_compile_nir(dev, pAllocator, clone, info->flags, - info->robustness, fs_key, features, shader, - sw_stage, hw, nir->xfb_info); + result = + hk_compile_nir(dev, pAllocator, clone, info->flags, + info->robustness, fs_key, features, shader, sw_stage, + hw, nir->xfb_info, info->set_layout_count); if (result != VK_SUCCESS) { hk_api_shader_destroy(&dev->vk, &obj->vk, pAllocator); ralloc_free(nir); @@ -1392,9 +1400,9 @@ hk_compile_shader(struct hk_device *dev, struct vk_shader_compile_info *info, struct hk_shader *shader = hk_only_variant(obj); /* hk_compile_nir takes ownership of nir */ - result = - hk_compile_nir(dev, pAllocator, nir, info->flags, info->robustness, - fs_key, features, shader, sw_stage, true, NULL); + result = hk_compile_nir(dev, pAllocator, nir, info->flags, + info->robustness, fs_key, features, shader, + sw_stage, true, NULL, info->set_layout_count); if (result != VK_SUCCESS) { hk_api_shader_destroy(&dev->vk, &obj->vk, pAllocator); return result; diff --git a/src/asahi/vulkan/hk_shader.h b/src/asahi/vulkan/hk_shader.h index fcade1cba1e..3112bd8af81 100644 --- a/src/asahi/vulkan/hk_shader.h +++ b/src/asahi/vulkan/hk_shader.h @@ -108,6 +108,7 @@ struct hk_shader_info { uint8_t clip_distance_array_size; uint8_t cull_distance_array_size; uint8_t image_heap_uniform; + uint8_t set_count; /* XXX: is there a less goofy way to do this? I really don't want dynamic * allocation here. diff --git a/src/compiler/nir/nir_intrinsics.py b/src/compiler/nir/nir_intrinsics.py index 01d3fdc08c7..653ec16fcda 100644 --- a/src/compiler/nir/nir_intrinsics.py +++ b/src/compiler/nir/nir_intrinsics.py @@ -2122,6 +2122,11 @@ intrinsic("load_texture_handle_agx", [1], 1, [], flags=[CAN_ELIMINATE, CAN_REORDER], bit_sizes=[32]) +# Load descriptor set address +intrinsic("load_descriptor_set_agx", [], 1, [DESC_SET], + flags=[CAN_ELIMINATE, CAN_REORDER], + bit_sizes=[64]) + # Given a bindless texture handle, load the address of the texture descriptor # described by that. This allows inspecting the descriptor from the shader. This # does not actually load the content of the descriptor, only the content of the