diff --git a/src/amd/vulkan/bvh/bvh.h b/src/amd/vulkan/bvh/bvh.h index 2b87ec47664..d15ad5f276b 100644 --- a/src/amd/vulkan/bvh/bvh.h +++ b/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; diff --git a/src/amd/vulkan/nir/radv_nir_lower_ray_queries.c b/src/amd/vulkan/nir/radv_nir_lower_ray_queries.c index 5600571e3c5..cc6b4c832d6 100644 --- a/src/amd/vulkan/nir/radv_nir_lower_ray_queries.c +++ b/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!"); diff --git a/src/amd/vulkan/nir/radv_nir_rt_common.c b/src/amd/vulkan/nir/radv_nir_rt_common.c index cd98516169f..e5bc4ea05d8 100644 --- a/src/amd/vulkan/nir/radv_nir_rt_common.c +++ b/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))); diff --git a/src/amd/vulkan/nir/radv_nir_rt_common.h b/src/amd/vulkan/nir/radv_nir_rt_common.h index f4acde364b2..66021a9d0ba 100644 --- a/src/amd/vulkan/nir/radv_nir_rt_common.h +++ b/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; diff --git a/src/amd/vulkan/nir/radv_nir_rt_shader.c b/src/amd/vulkan/nir/radv_nir_rt_shader.c index d1a3a656689..d066665e671 100644 --- a/src/amd/vulkan/nir/radv_nir_rt_shader.c +++ b/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: diff --git a/src/amd/vulkan/radv_acceleration_structure.c b/src/amd/vulkan/radv_acceleration_structure.c index 71aa2cd8f19..1bc7b6b3173 100644 --- a/src/amd/vulkan/radv_acceleration_structure.c +++ b/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