anv: Track leaf block offset map

Track where is each leaf_id encoded in final BVH.
It's a map of leaf_id == final_bvh_offset. This will help us to navigate
the BVH layout in update pass.

Leaf block offset will give us : Leaf id -> bvh block
and parent-child map can be used for: bvh_block -> parent offset.

Signed-off-by: Sagar Ghuge <sagar.ghuge@intel.com>

src/intel/vulkan/bvh/anv_build_interface.h

@@ -34,6 +34,7 @@ struct encode_args {
    uint32_t geometry_type;
 
    VOID_REF parent_child_map;
+   VOID_REF leaf_block_offset_map;
 };
 
 struct header_args {

src/intel/vulkan/bvh/anv_bvh.h

@@ -333,6 +333,8 @@ struct anv_instance_leaf {
 | For a BLAS, nothing here      |
 |-------------------------------| bvh_layout.parent_child_map_offset
 | Parent - child map            |
+|-------------------------------| bvh_layout.leaf_block_map_offset
+| Leaf block offset map         |
 |-------------------------------|
 | padding to align to           |
 | 64 bytes boundary             | bvh_layout.size
@@ -355,6 +357,9 @@ struct bvh_layout {
     * */
    uint64_t parent_child_map_offset;
 
+   /* This map stores BVH block index for each leaf id (IR ID) */
+   uint64_t leaf_block_map_offset;
+
    /* Total size = bvh_offset + leaves + internal_nodes (assuming there's no
     * internal node collpased)
     */

src/intel/vulkan/bvh/encode.comp

@@ -491,6 +491,13 @@ main()
                child_aabb = DEREF(REF(vk_ir_node)NODE_OFFSET(child)).aabb;
 
             uint32_t type = ir_id_to_type(child);
+	    if (child != VK_BVH_INVALID_NODE &&
+                (type == vk_ir_node_triangle || type == vk_ir_node_aabb)) {
+               uint32_t ir_offset = ir_id_to_offset(child);
+               uint32_t leaf_id = ir_offset / intermediate_leaf_node_size;
+               DEREF(INDEX(uint32_t, args.leaf_block_offset_map, leaf_id)) = child_block;
+            }
+
             /* Track each children's parent in the map. */
             if (child != VK_BVH_INVALID_NODE && type != vk_ir_node_instance) {
                uint32_t pcm_val = 0;

src/intel/vulkan/genX_acceleration_structure.c

@@ -286,6 +286,10 @@ get_bvh_layout(const struct vk_acceleration_structure_build_state *state,
    layout->parent_child_map_offset = offset;
    offset += parent_child_map_size;
 
+   uint64_t leaf_block_offset_size = (internal_count + leaf_count) * sizeof(uint32_t);
+   layout->leaf_block_map_offset = offset;
+   offset += leaf_block_offset_size;
+
    layout->size = align64(offset, 64);
 }
 
@@ -427,6 +431,8 @@ anv_encode_as(VkCommandBuffer commandBuffer, const struct vk_acceleration_struct
                               bvh_layout.instance_leaves_offset,
       .parent_child_map = vk_acceleration_structure_get_va(dst) +
                           bvh_layout.parent_child_map_offset,
+      .leaf_block_offset_map = vk_acceleration_structure_get_va(dst) +
+                               bvh_layout.leaf_block_map_offset,
    };
    anv_bvh_build_set_args(commandBuffer, &args, sizeof(args));