anv: Track parent-child map for BVH update

This map stores parent BVH offset for each of their children. This will
help us to walk the BVH layout later in the update pass.

Since we are tracking block indexes, even with 2^32 large BVH size, we
can have 2^26 max indices (each block 64B wide) that leaves us 6 bits in
which we can track child slot index occupancies in parent.

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

src/intel/vulkan/bvh/anv_build_interface.h

@@ -32,6 +32,8 @@ struct encode_args {
 
    uint32_t leaf_node_count;
    uint32_t geometry_type;
+
+   VOID_REF parent_child_map;
 };
 
 struct header_args {

src/intel/vulkan/bvh/anv_bvh.h

@@ -327,13 +327,12 @@ struct anv_instance_leaf {
 | start with root node,         |
 | followed by interleaving      |
 | internal nodes and leaves     |
-|-------------------------------|
-| padding to align to           |
-| 64 bytes boundary             |
 |-------------------------------| bvh_layout.instance_leaves_offset
 | For a TLAS, the pointers      |
 | to all anv_instance_leaves    |
 | For a BLAS, nothing here      |
+|-------------------------------| bvh_layout.parent_child_map_offset
+| Parent - child map            |
 |-------------------------------|
 | padding to align to           |
 | 64 bytes boundary             | bvh_layout.size
@@ -346,13 +345,20 @@ struct bvh_layout {
     */
    uint64_t bvh_offset;
 
+   /* This tracks pointers to all anv_instance_leaves for BLAS. */
+   uint64_t instance_leaves_offset;
+
+   /* This map stores parent BVH offset for each child
+    *
+    * Lower 26bits - parent block index
+    * upper 6bits  - parent child slot index
+    * */
+   uint64_t parent_child_map_offset;
+
    /* Total size = bvh_offset + leaves + internal_nodes (assuming there's no
     * internal node collpased)
     */
    uint64_t size;
-
-   /* This tracks pointers to all anv_instance_leaves for BLAS. */
-   uint64_t instance_leaves_offset;
 };
 
 #endif

src/intel/vulkan/bvh/encode.comp

@@ -359,6 +359,7 @@ main()
       /* Tracks BLOCK where the next children should be encoded. */
       DEREF(args.header).dst_node_offset = 1;
       DEREF(header).instance_count = 0;
+      DEREF(INDEX(uint32_t, args.parent_child_map, 0)) = VK_NULL_BVH_OFFSET;
    }
 
    IR_NODE children[6] = {VK_BVH_INVALID_NODE, VK_BVH_INVALID_NODE,
@@ -490,9 +491,17 @@ 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_internal)
-               encode_leaf_node(type, child, BLOCK_OFFSET(child_block),
-                                header);
+            /* Track each children's parent in the map. */
+            if (child != VK_BVH_INVALID_NODE && type != vk_ir_node_instance) {
+               uint32_t pcm_val = 0;
+               pcm_val = (cluster.idx << 26) | internal_node_block;
+               DEREF(INDEX(uint32_t, args.parent_child_map, child_block)) = pcm_val;
+            }
+
+            if (child != VK_BVH_INVALID_NODE && type != vk_ir_node_internal) {
+               encode_leaf_node(type, child,
+                              BLOCK_OFFSET(child_block), header);
+            }
          }
 
          BLOCK child_block_offset =

src/intel/vulkan/genX_acceleration_structure.c

@@ -273,7 +273,6 @@ get_bvh_layout(const struct vk_acceleration_structure_build_state *state,
       UNREACHABLE("Unknown VkGeometryTypeKHR");
    }
 
-   offset = align64(offset, 64);
    layout->instance_leaves_offset = offset;
 
    /* For a TLAS, we store the address of anv_instance_leaf after header
@@ -283,6 +282,10 @@ get_bvh_layout(const struct vk_acceleration_structure_build_state *state,
       offset += leaf_count * sizeof(uint64_t);
    }
 
+   uint64_t parent_child_map_size = (internal_count + leaf_count) * sizeof(uint32_t);
+   layout->parent_child_map_offset = offset;
+   offset += parent_child_map_size;
+
    layout->size = align64(offset, 64);
 }
 
@@ -422,6 +425,8 @@ anv_encode_as(VkCommandBuffer commandBuffer, const struct vk_acceleration_struct
       .geometry_type = geometry_type,
       .instance_leaves_addr = vk_acceleration_structure_get_va(dst) +
                               bvh_layout.instance_leaves_offset,
+      .parent_child_map = vk_acceleration_structure_get_va(dst) +
+                          bvh_layout.parent_child_map_offset,
    };
    anv_bvh_build_set_args(commandBuffer, &args, sizeof(args));