lavapipe: Fix ray tracing position fetch with multiple geometries

The previous code didn't take into account, that the primitive index is
set back to 0 for each geometry.

Reviewed-by: Mike Blumenkrantz <michael.blumenkrantz@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/34496>

src/gallium/frontends/lavapipe/lvp_ray_tracing_pipeline.c

@@ -266,6 +266,7 @@ struct lvp_ray_tracing_state {
    nir_variable *tmax;
 
    nir_variable *instance_addr;
+   nir_variable *primitive_addr;
    nir_variable *primitive_id;
    nir_variable *geometry_id_and_flags;
    nir_variable *hit_kind;
@@ -312,8 +313,10 @@ lvp_ray_tracing_pipeline_compiler_get_stack_size(
 }
 
 static void
-lvp_ray_tracing_state_init(nir_shader *nir, struct lvp_ray_tracing_state *state)
+lvp_ray_tracing_state_init(nir_shader *nir, struct lvp_ray_tracing_pipeline_compiler *compiler)
 {
+   struct lvp_ray_tracing_state *state = &compiler->state;
+
    state->bvh_base = nir_variable_create(nir, nir_var_shader_temp, glsl_uint64_t_type(), "bvh_base");
    state->flags = nir_variable_create(nir, nir_var_shader_temp, glsl_uint_type(), "flags");
    state->cull_mask = nir_variable_create(nir, nir_var_shader_temp, glsl_uint_type(), "cull_mask");
@@ -338,6 +341,9 @@ lvp_ray_tracing_state_init(nir_shader *nir, struct lvp_ray_tracing_state *state)
    state->accept = nir_variable_create(nir, nir_var_shader_temp, glsl_bool_type(), "accept");
    state->terminate = nir_variable_create(nir, nir_var_shader_temp, glsl_bool_type(), "terminate");
    state->opaque = nir_variable_create(nir, nir_var_shader_temp, glsl_bool_type(), "opaque");
+
+   if (compiler->pipeline->device->vk.enabled_features.rayTracingPositionFetch)
+      state->primitive_addr = nir_variable_create(nir, nir_var_shader_temp, glsl_uint64_t_type(), "primitive_addr");
 }
 
 static void
@@ -655,6 +661,12 @@ lvp_handle_triangle_intersection(nir_builder *b,
    nir_store_var(b, state->hit_kind,
                  nir_bcsel(b, intersection->frontface, nir_imm_int(b, 0xFE), nir_imm_int(b, 0xFF)), 0x1);
 
+   nir_def *prev_primitive_addr = NULL;
+   if (state->primitive_addr) {
+      prev_primitive_addr = nir_load_var(b, state->primitive_addr);
+      nir_store_var(b, state->primitive_addr, intersection->base.node_addr, 0x1);
+   }         
+
    nir_store_scratch(b, intersection->barycentrics, barycentrics_offset);
 
    nir_def *geometry_id = nir_iand_imm(b, intersection->base.geometry_id_and_flags, 0xfffffff);
@@ -702,6 +714,8 @@ lvp_handle_triangle_intersection(nir_builder *b,
       nir_store_var(b, state->geometry_id_and_flags, prev_geometry_id_and_flags, 0x1);
       nir_store_var(b, state->hit_kind, prev_hit_kind, 0x1);
       nir_store_scratch(b, prev_barycentrics, barycentrics_offset);
+      if (state->primitive_addr)
+         nir_store_var(b, state->primitive_addr, prev_primitive_addr, 0x1);
    }
    nir_pop_if(b, NULL);
 }
@@ -993,8 +1007,7 @@ lvp_lower_ray_tracing_instr(nir_builder *b, nir_instr *instr, void *data)
    }
    case nir_intrinsic_load_ray_triangle_vertex_positions: {
       def = lvp_load_vertex_position(
-         b, nir_load_var(b, state->instance_addr), nir_load_var(b, state->primitive_id),
-         nir_intrinsic_column(intr));
+         b, nir_load_var(b, state->primitive_addr), nir_intrinsic_column(intr));
       break;
    }
    /* Internal system values */
@@ -1041,7 +1054,7 @@ lvp_compile_ray_tracing_pipeline(struct lvp_pipeline *pipeline,
       .pipeline = pipeline,
       .flags = vk_rt_pipeline_create_flags(create_info),
    };
-   lvp_ray_tracing_state_init(b->shader, &compiler.state);
+   lvp_ray_tracing_state_init(b->shader, &compiler);
    compiler.functions = _mesa_pointer_hash_table_create(NULL);
 
    nir_def *launch_id = nir_load_ray_launch_id(b);

src/gallium/frontends/lavapipe/nir/lvp_nir.h

@@ -14,8 +14,7 @@ nir_def *lvp_mul_vec3_mat(nir_builder *b, nir_def *vec, nir_def *matrix[], bool
 
 void lvp_load_wto_matrix(nir_builder *b, nir_def *instance_addr, nir_def **node_data, nir_def **out);
 
-nir_def *lvp_load_vertex_position(nir_builder *b, nir_def *instance_addr,
-                                  nir_def *primitive_id, uint32_t index);
+nir_def *lvp_load_vertex_position(nir_builder *b, nir_def *primitive_addr, uint32_t index);
 
 struct lvp_ray_traversal_args;
 

src/gallium/frontends/lavapipe/nir/lvp_nir_lower_ray_queries.c

@@ -121,6 +121,7 @@ struct ray_query_traversal_vars {
 };
 
 struct ray_query_intersection_vars {
+   rq_variable *primitive_addr;
    rq_variable *primitive_id;
    rq_variable *geometry_id_and_flags;
    rq_variable *instance_addr;
@@ -182,6 +183,8 @@ init_ray_query_intersection_vars(void *ctx, nir_shader *shader, unsigned array_l
 
    const struct glsl_type *vec2_type = glsl_vector_type(GLSL_TYPE_FLOAT, 2);
 
+   result.primitive_addr =
+      rq_variable_create(ctx, shader, array_length, glsl_uint64_t_type(), VAR_NAME("_primitive_addr"));
    result.primitive_id =
       rq_variable_create(ctx, shader, array_length, glsl_uint_type(), VAR_NAME("_primitive_id"));
    result.geometry_id_and_flags = rq_variable_create(ctx, shader, array_length, glsl_uint_type(),
@@ -249,6 +252,7 @@ lower_ray_query(nir_shader *shader, nir_variable *ray_query, struct hash_table *
 static void
 copy_candidate_to_closest(nir_builder *b, nir_def *index, struct ray_query_vars *vars)
 {
+   rq_copy_var(b, index, vars->closest.primitive_addr, vars->candidate.primitive_addr, 0x1);
    rq_copy_var(b, index, vars->closest.barycentrics, vars->candidate.barycentrics, 0x3);
    rq_copy_var(b, index, vars->closest.geometry_id_and_flags, vars->candidate.geometry_id_and_flags,
                0x1);
@@ -431,8 +435,7 @@ lower_rq_load(nir_builder *b, nir_def *index, nir_intrinsic_instr *instr,
       return rq_load_var(b, index, vars->origin);
    case nir_ray_query_value_intersection_triangle_vertex_positions:
       return lvp_load_vertex_position(
-         b, rq_load_var(b, index, intersection->instance_addr),
-         rq_load_var(b, index, intersection->primitive_id), column);
+         b, rq_load_var(b, index, intersection->primitive_addr), column);
    default:
       unreachable("Invalid nir_ray_query_value!");
    }
@@ -474,6 +477,7 @@ handle_candidate_triangle(nir_builder *b, struct lvp_triangle_intersection *inte
    nir_def *index = data->index;
 
    rq_store_var(b, index, vars->candidate.barycentrics, intersection->barycentrics, 3);
+   rq_store_var(b, index, vars->candidate.primitive_addr, intersection->base.node_addr, 1);
    rq_store_var(b, index, vars->candidate.primitive_id, intersection->base.primitive_id, 1);
    rq_store_var(b, index, vars->candidate.geometry_id_and_flags,
                 intersection->base.geometry_id_and_flags, 1);

src/gallium/frontends/lavapipe/nir/lvp_nir_ray_tracing.c

@@ -56,20 +56,9 @@ lvp_load_wto_matrix(nir_builder *b, nir_def *instance_addr, nir_def **node_data,
 }
 
 nir_def *
-lvp_load_vertex_position(nir_builder *b, nir_def *instance_addr, nir_def *primitive_id,
-                         uint32_t index)
+lvp_load_vertex_position(nir_builder *b, nir_def *primitive_addr, uint32_t index)
 {
-   nir_def *bvh_addr = nir_build_load_global(
-      b, 1, 64, nir_iadd_imm(b, instance_addr, offsetof(struct lvp_bvh_instance_node, bvh_ptr)));
-
-   nir_def *leaf_nodes_offset = nir_build_load_global(
-      b, 1, 32, nir_iadd_imm(b, bvh_addr, offsetof(struct lvp_bvh_header, leaf_nodes_offset)));
-
-   nir_def *offset = nir_imul_imm(b, primitive_id, sizeof(struct lvp_bvh_triangle_node));
-   offset = nir_iadd(b, offset, leaf_nodes_offset);
-   offset = nir_iadd_imm(b, offset, index * 3 * sizeof(float));
-
-   return nir_build_load_global(b, 3, 32, nir_iadd(b, bvh_addr, nir_u2u64(b, offset)));
+   return nir_build_load_global(b, 3, 32, nir_iadd_imm(b, primitive_addr, index * 3 * sizeof(float)));
 }
 
 static nir_def *