vulkan: Handle inactive primitives with LBVH builds

cc: mesa-stable

Reviewed-by: Natalie Vock <natalie.vock@gmx.de>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/39378>

src/vulkan/runtime/bvh/lbvh_generate_ir.comp

@@ -39,6 +39,10 @@ main(void)
 {
    uint32_t global_id = gl_GlobalInvocationID.x;
 
+   uint32_t internal_node_count = max(DEREF(args.header).active_leaf_count, 2) - 1;
+   if (global_id >= internal_node_count)
+      return;
+
    uint32_t idx = global_id;
 
    uint32_t previous_id = VK_BVH_INVALID_NODE;
@@ -63,8 +67,7 @@ main(void)
       /* We allocate nodes on demand with the atomic here to ensure children come before their
        * parents, which is a requirement of the encoder.
        */
-      uint32_t dst_idx =
+      uint32_t dst_idx = atomicAdd(DEREF(args.header).ir_internal_node_count, 1);
-         atomicAdd(DEREF(REF(vk_ir_header)(args.header)).ir_internal_node_count, 1);
 
       uint32_t current_offset = args.internal_node_base + dst_idx * SIZEOF(vk_ir_box_node);
       uint32_t current_id = pack_ir_node_id(current_offset, vk_ir_node_internal);

src/vulkan/runtime/bvh/lbvh_main.comp

@@ -35,9 +35,9 @@ layout(push_constant) uniform CONSTS
 };
 
 int32_t
-longest_common_prefix(int32_t i, uint32_t key_i, int32_t j)
+longest_common_prefix(int32_t i, uint32_t key_i, int32_t j, uint32_t active_leaf_count)
 {
-   if (j < 0 || j >= args.id_count)
+   if (j < 0 || j >= active_leaf_count)
       return -1;
 
    uint32_t key_j = DEREF(INDEX(key_id_pair, args.src_ids, j)).key;
@@ -78,12 +78,14 @@ longest_common_prefix(int32_t i, uint32_t key_i, int32_t j)
 void
 main()
 {
-   if (args.id_count <= 1) {
+   uint32_t active_leaf_count = DEREF(args.header).active_leaf_count;
+
+   if (active_leaf_count <= 1) {
       REF(lbvh_node_info) dst = REF(lbvh_node_info)(args.node_info);
       DEREF(dst).parent = VK_BVH_INVALID_NODE;
       DEREF(dst).path_count = 2;
       DEREF(dst).children[0] =
-         args.id_count == 1 ? DEREF(INDEX(key_id_pair, args.src_ids, 0)).id : VK_BVH_INVALID_NODE;
+         active_leaf_count == 1 ? DEREF(INDEX(key_id_pair, args.src_ids, 0)).id : VK_BVH_INVALID_NODE;
       DEREF(dst).children[1] = VK_BVH_INVALID_NODE;
       return;
    }
@@ -91,8 +93,8 @@ main()
    int32_t id = int32_t(gl_GlobalInvocationID.x);
    uint32_t id_key = DEREF(INDEX(key_id_pair, args.src_ids, id)).key;
 
-   int32_t left_lcp = longest_common_prefix(id, id_key, id - 1);
+   int32_t left_lcp = longest_common_prefix(id, id_key, id - 1, active_leaf_count);
-   int32_t right_lcp = longest_common_prefix(id, id_key, id + 1);
+   int32_t right_lcp = longest_common_prefix(id, id_key, id + 1, active_leaf_count);
    int32_t dir = right_lcp > left_lcp ? 1 : -1;
    int32_t lcp_min = min(left_lcp, right_lcp);
 
@@ -100,13 +102,13 @@ main()
     * this subtree is going to own.
     */
    int32_t lmax = 128;
-   while (longest_common_prefix(id, id_key, id + dir * lmax) > lcp_min) {
+   while (longest_common_prefix(id, id_key, id + dir * lmax, active_leaf_count) > lcp_min) {
       lmax *= 2;
    }
 
    int32_t length = 0;
    for (int32_t t = lmax / 2; t >= 1; t /= 2) {
-      if (longest_common_prefix(id, id_key, id + (length + t) * dir) > lcp_min)
+      if (longest_common_prefix(id, id_key, id + (length + t) * dir, active_leaf_count) > lcp_min)
          length += t;
    }
    int32_t other_end = id + length * dir;
@@ -114,11 +116,11 @@ main()
    /* The number of bits in the prefix that is the same for all elements in the
     * range.
     */
-   int32_t lcp_node = longest_common_prefix(id, id_key, other_end);
+   int32_t lcp_node = longest_common_prefix(id, id_key, other_end, active_leaf_count);
    int32_t child_range = 0;
    for (int32_t diff = 2; diff < 2 * length; diff *= 2) {
       int32_t t = DIV_ROUND_UP(length, diff);
-      if (longest_common_prefix(id, id_key, id + (child_range + t) * dir) > lcp_node)
+      if (longest_common_prefix(id, id_key, id + (child_range + t) * dir, active_leaf_count) > lcp_node)
          child_range += t;
    }
 

src/vulkan/runtime/bvh/vk_build_interface.h

@@ -81,14 +81,14 @@ struct lbvh_main_args {
    VOID_REF bvh;
    REF(key_id_pair) src_ids;
    VOID_REF node_info;
-   uint32_t id_count;
+   REF(vk_ir_header) header;
    uint32_t internal_node_base;
 };
 
 struct lbvh_generate_ir_args {
    VOID_REF bvh;
    VOID_REF node_info;
-   VOID_REF header;
+   REF(vk_ir_header) header;
    uint32_t internal_node_base;
 };
 

src/vulkan/runtime/vk_acceleration_structure.c

@@ -961,7 +961,7 @@ lbvh_build_internal(VkCommandBuffer commandBuffer,
          .bvh = pInfos[i].scratchData.deviceAddress + bvh_states[i].vk.scratch.ir_offset,
          .src_ids = pInfos[i].scratchData.deviceAddress + src_scratch_offset,
          .node_info = pInfos[i].scratchData.deviceAddress + bvh_states[i].vk.scratch.lbvh_node_offset,
-         .id_count = bvh_states[i].vk.leaf_node_count,
+         .header = pInfos[i].scratchData.deviceAddress + bvh_states[i].vk.scratch.header_offset,
          .internal_node_base = bvh_states[i].vk.scratch.internal_node_offset - bvh_states[i].vk.scratch.ir_offset,
       };