radv: Fix rayTracingPositionFetch with multiple geometies

The fix adds more indirections to avoid increasing register pressure by
tracking the primitive address.

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

src/amd/vulkan/bvh/bvh.h

@@ -59,6 +59,7 @@ struct radv_accel_struct_header {
 
    /* Everything after this gets updated/copied from the CPU. */
    uint32_t geometry_count;
+   uint32_t primitive_base_indices_offset;
    uint64_t instance_offset;
    uint64_t instance_count;
    uint32_t build_flags;

src/amd/vulkan/nir/radv_nir_lower_ray_queries.c

@@ -391,7 +391,9 @@ lower_rq_load(struct radv_device *device, nir_builder *b, nir_intrinsic_instr *i
    case nir_ray_query_value_intersection_triangle_vertex_positions: {
       nir_def *instance_node_addr = isec_load(b, intersection, instance_addr);
       nir_def *primitive_id = isec_load(b, intersection, primitive_id);
-      return radv_load_vertex_position(device, b, instance_node_addr, primitive_id, nir_intrinsic_column(instr));
+      nir_def *geometry_id = nir_iand_imm(b, isec_load(b, intersection, geometry_id_and_flags), 0xFFFFFF);
+      return radv_load_vertex_position(device, b, instance_node_addr, geometry_id, primitive_id,
+                                       nir_intrinsic_column(instr));
    }
    default:
       unreachable("Invalid nir_ray_query_value!");

src/amd/vulkan/nir/radv_nir_rt_common.c

@@ -313,14 +313,23 @@ nir_build_wto_matrix_load(nir_builder *b, nir_def *instance_addr, nir_def **out)
 }
 
 nir_def *
-radv_load_vertex_position(struct radv_device *device, nir_builder *b, nir_def *instance_addr, nir_def *primitive_id,
-                          uint32_t index)
+radv_load_vertex_position(struct radv_device *device, nir_builder *b, nir_def *instance_addr, nir_def *geometry_id,
+                          nir_def *primitive_id, uint32_t index)
 {
    nir_def *bvh_addr_id =
       nir_build_load_global(b, 1, 64, nir_iadd_imm(b, instance_addr, offsetof(struct radv_bvh_instance_node, bvh_ptr)));
    nir_def *bvh_addr = build_node_to_addr(device, b, bvh_addr_id, true);
 
-   nir_def *offset = nir_imul_imm(b, primitive_id, sizeof(struct radv_bvh_triangle_node));
+   nir_def *bvh_offset = nir_build_load_global(
+      b, 1, 32, nir_iadd_imm(b, instance_addr, offsetof(struct radv_bvh_instance_node, bvh_offset)));
+   nir_def *accel_struct = nir_isub(b, bvh_addr, nir_u2u64(b, bvh_offset));
+   nir_def *base_indices_offset = nir_build_load_global(
+      b, 1, 32,
+      nir_iadd_imm(b, accel_struct, offsetof(struct radv_accel_struct_header, primitive_base_indices_offset)));
+   nir_def *base_index_offset = nir_iadd(b, base_indices_offset, nir_imul_imm(b, geometry_id, sizeof(uint32_t)));
+   nir_def *base_index = nir_build_load_global(b, 1, 32, nir_iadd(b, accel_struct, nir_u2u64(b, base_index_offset)));
+
+   nir_def *offset = nir_imul_imm(b, nir_iadd(b, base_index, primitive_id), sizeof(struct radv_bvh_triangle_node));
    offset = nir_iadd_imm(b, offset, sizeof(struct radv_bvh_box32_node) + index * 3 * sizeof(float));
 
    return nir_build_load_global(b, 3, 32, nir_iadd(b, bvh_addr, nir_u2u64(b, offset)));

src/amd/vulkan/nir/radv_nir_rt_common.h

@@ -21,7 +21,7 @@ nir_def *nir_build_vec3_mat_mult(nir_builder *b, nir_def *vec, nir_def *matrix[]
 void nir_build_wto_matrix_load(nir_builder *b, nir_def *instance_addr, nir_def **out);
 
 nir_def *radv_load_vertex_position(struct radv_device *device, nir_builder *b, nir_def *instance_addr,
-                                   nir_def *primitive_id, uint32_t index);
+                                   nir_def *geometry_id, nir_def *primitive_id, uint32_t index);
 
 struct radv_ray_traversal_args;
 

src/amd/vulkan/nir/radv_nir_rt_shader.c

@@ -702,7 +702,9 @@ radv_lower_rt_instruction(nir_builder *b, nir_instr *instr, void *_data)
    case nir_intrinsic_load_ray_triangle_vertex_positions: {
       nir_def *instance_node_addr = nir_load_var(b, vars->instance_addr);
       nir_def *primitive_id = nir_load_var(b, vars->primitive_id);
-      ret = radv_load_vertex_position(vars->device, b, instance_node_addr, primitive_id, nir_intrinsic_column(intr));
+      nir_def *geometry_id = nir_iand_imm(b, nir_load_var(b, vars->geometry_id_and_flags), 0xFFFFFFF);
+      ret = radv_load_vertex_position(vars->device, b, instance_node_addr, geometry_id, primitive_id,
+                                      nir_intrinsic_column(intr));
       break;
    }
    default:

src/amd/vulkan/radv_acceleration_structure.c

@@ -46,6 +46,7 @@ static const uint32_t leaf_always_active_spv[] = {
 
 struct acceleration_structure_layout {
    uint32_t geometry_info_offset;
+   uint32_t primitive_base_indices_offset;
    uint32_t bvh_offset;
    uint32_t leaf_nodes_offset;
    uint32_t internal_nodes_offset;
@@ -94,6 +95,12 @@ radv_get_acceleration_structure_layout(struct radv_device *device, uint32_t leaf
       accel_struct->geometry_info_offset = offset;
       offset += sizeof(struct radv_accel_struct_geometry_info) * build_info->geometryCount;
    }
+
+   if (device->vk.enabled_features.rayTracingPositionFetch && geometry_type == VK_GEOMETRY_TYPE_TRIANGLES_KHR) {
+      accel_struct->primitive_base_indices_offset = offset;
+      offset += sizeof(uint32_t) * build_info->geometryCount;
+   }
+
    /* Parent links, which have to go directly before bvh_offset as we index them using negative
     * offsets from there. */
    offset += bvh_size / 64 * 4;
@@ -513,6 +520,7 @@ radv_init_header(VkCommandBuffer commandBuffer, const VkAccelerationStructureBui
 
    header.build_flags = build_info->flags;
    header.geometry_count = build_info->geometryCount;
+   header.primitive_base_indices_offset = layout.primitive_base_indices_offset;
 
    radv_update_buffer_cp(cmd_buffer, vk_acceleration_structure_get_va(dst) + base, (const char *)&header + base,
                          sizeof(header) - base);
@@ -538,6 +546,27 @@ radv_init_header(VkCommandBuffer commandBuffer, const VkAccelerationStructureBui
 
       free(geometry_infos);
    }
+
+   VkGeometryTypeKHR geometry_type = vk_get_as_geometry_type(build_info);
+   if (device->vk.enabled_features.rayTracingPositionFetch && geometry_type == VK_GEOMETRY_TYPE_TRIANGLES_KHR) {
+      uint32_t base_indices_size = sizeof(uint32_t) * build_info->geometryCount;
+      uint32_t *base_indices = malloc(base_indices_size);
+      if (!base_indices) {
+         vk_command_buffer_set_error(&cmd_buffer->vk, VK_ERROR_OUT_OF_HOST_MEMORY);
+         return;
+      }
+
+      uint32_t base_index = 0;
+      for (uint32_t i = 0; i < build_info->geometryCount; i++) {
+         base_indices[i] = base_index;
+         base_index += build_range_infos[i].primitiveCount;
+      }
+
+      radv_CmdUpdateBuffer(commandBuffer, vk_buffer_to_handle(dst->buffer),
+                           dst->offset + layout.primitive_base_indices_offset, base_indices_size, base_indices);
+
+      free(base_indices);
+   }
 }
 
 static void