mesa/src/vulkan/runtime/bvh/lbvh_generate_ir.comp

/*
 * Copyright © 2022 Bas Nieuwenhuizen
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice (including the next
 * paragraph) shall be included in all copies or substantial portions of the
 * Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 */

#version 460

#include "vk_build_interface.h"

layout(local_size_x_id = SUBGROUP_SIZE_ID, local_size_y = 1, local_size_z = 1) in;

TYPE(lbvh_node_info, 4);

layout(push_constant) uniform CONSTS
{
   lbvh_generate_ir_args args;
};

void
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;
   vk_aabb previous_bounds;
   previous_bounds.min = vec3(INFINITY);
   previous_bounds.max = vec3(-INFINITY);
   uint32_t previous_flags;

   for (;;) {
      uint32_t count = 0;

      /* Check if all children have been processed. As this is an atomic the last path coming from
       * a child will pass here, while earlier paths break.
       */
      count = atomicAdd(
         DEREF(INDEX(lbvh_node_info, args.node_info, idx)).path_count, 1, gl_ScopeDevice,
         gl_StorageSemanticsBuffer,
         gl_SemanticsAcquireRelease | gl_SemanticsMakeAvailable | gl_SemanticsMakeVisible);
      if (count != 2)
         break;

      /* 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 = atomicAdd(DEREF(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);

      REF(vk_ir_box_node) node = REF(vk_ir_box_node)(OFFSET(args.bvh, current_offset));
      vk_aabb bounds = previous_bounds;
      uint32_t flags = VK_BVH_BOX_FLAG_ONLY_OPAQUE | VK_BVH_BOX_FLAG_NO_OPAQUE;

      lbvh_node_info info = DEREF(INDEX(lbvh_node_info, args.node_info, idx));

      uint32_t children[2] = info.children;

      /* Try using the cached previous_bounds instead of fetching the bounds twice. */
      int32_t previous_child_index = -1;
      if (previous_id == children[0])
         previous_child_index = 0;
      else if (previous_id == children[1])
         previous_child_index = 1;

      if (previous_child_index == -1) {
         if (children[0] != VK_BVH_INVALID_NODE) {
            uint32_t child_offset = ir_id_to_offset(children[0]);
            REF(vk_ir_node) child = REF(vk_ir_node)(OFFSET(args.bvh, child_offset));
            vk_aabb child_bounds = DEREF(child).aabb;
            bounds.min = min(bounds.min, child_bounds.min);
            bounds.max = max(bounds.max, child_bounds.max);

            if (VK_BUILD_FLAG(VK_BUILD_FLAG_PROPAGATE_CULL_FLAGS))
               flags &= fetch_child_flags(args.bvh, children[0]);
         }
         previous_child_index = 0;
      } else {
            if (VK_BUILD_FLAG(VK_BUILD_FLAG_PROPAGATE_CULL_FLAGS))
               flags &= previous_flags;
      }

      /* Fetch the non-cached child */
      if (children[1 - previous_child_index] != VK_BVH_INVALID_NODE) {
         uint32_t child_offset = ir_id_to_offset(children[1 - previous_child_index]);
         REF(vk_ir_node) child = REF(vk_ir_node)(OFFSET(args.bvh, child_offset));
         vk_aabb child_bounds = DEREF(child).aabb;
         bounds.min = min(bounds.min, child_bounds.min);
         bounds.max = max(bounds.max, child_bounds.max);

         if (VK_BUILD_FLAG(VK_BUILD_FLAG_PROPAGATE_CULL_FLAGS))
            flags &= fetch_child_flags(args.bvh, children[1 - previous_child_index]);
      }


      vk_ir_box_node node_value;

      node_value.base.aabb = bounds;
      node_value.bvh_offset = VK_UNKNOWN_BVH_OFFSET;
      node_value.children = children;

      if (VK_BUILD_FLAG(VK_BUILD_FLAG_PROPAGATE_CULL_FLAGS))
         node_value.flags = flags;

      DEREF(node) = node_value;

      if (info.parent == VK_BVH_INVALID_NODE)
         break;

      idx = info.parent & ~LBVH_RIGHT_CHILD_BIT;

      DEREF(INDEX(lbvh_node_info, args.node_info, idx))
         .children[(info.parent >> LBVH_RIGHT_CHILD_BIT_SHIFT) & 1] = current_id;

      previous_id = current_id;
      previous_bounds = bounds;
      previous_flags = flags;

      memoryBarrierBuffer();
   }
}