radv: Add conversion shader for internal nodes

Reviewed-by: Konstantin Seurer <konstantin.seurer@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/18769>

src/amd/vulkan/bvh/build_helpers.h

@@ -379,4 +379,23 @@ calculate_node_bounds(VOID_REF bvh, uint32_t id)
    return aabb;
 }
 
+AABB
+load_aabb(REF(radv_ir_node) node) {
+   AABB bounds;
+   bounds.min.x = DEREF(node).aabb[0][0];
+   bounds.min.y = DEREF(node).aabb[0][1];
+   bounds.min.z = DEREF(node).aabb[0][2];
+   bounds.max.x = DEREF(node).aabb[1][0];
+   bounds.max.y = DEREF(node).aabb[1][1];
+   bounds.max.z = DEREF(node).aabb[1][2];
+   return bounds;
+}
+
+float
+aabb_surface_area(AABB aabb)
+{
+   vec3 diagonal = aabb.max - aabb.min;
+   return 2 * diagonal.x * diagonal.y + 2 * diagonal.y * diagonal.z + 2 * diagonal.x * diagonal.z;
+}
+
 #endif

src/amd/vulkan/bvh/build_interface.h

@@ -75,4 +75,12 @@ struct copy_args {
    uint32_t mode;
 };
 
+struct convert_internal_args {
+   VOID_REF intermediate_bvh;
+   VOID_REF output_bvh;
+   uint32_t leaf_node_count;
+   uint32_t internal_node_count;
+   uint32_t geometry_type;
+};
+
 #endif

src/amd/vulkan/bvh/converter_internal.comp

@@ -0,0 +1,174 @@
+/*
+ * Copyright © 2022 Friedrich Vock
+ *
+ * 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
+
+#extension GL_GOOGLE_include_directive : require
+
+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require
+#extension GL_EXT_scalar_block_layout : require
+#extension GL_EXT_buffer_reference : require
+#extension GL_EXT_buffer_reference2 : require
+
+layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
+
+#include "build_helpers.h"
+#include "build_interface.h"
+
+layout(push_constant) uniform CONSTS {
+   convert_internal_args args;
+};
+
+void
+main()
+{
+   uint32_t global_id = gl_GlobalInvocationID.x;
+
+   uint32_t intermediate_leaf_node_size;
+   uint32_t output_leaf_node_size;
+   switch (args.geometry_type) {
+   case VK_GEOMETRY_TYPE_TRIANGLES_KHR:
+      intermediate_leaf_node_size = SIZEOF(radv_ir_triangle_node);
+      output_leaf_node_size = SIZEOF(radv_bvh_triangle_node);
+      break;
+   case VK_GEOMETRY_TYPE_AABBS_KHR:
+      intermediate_leaf_node_size = SIZEOF(radv_ir_aabb_node);
+      output_leaf_node_size = SIZEOF(radv_bvh_aabb_node);
+      break;
+   default: /* instances */
+      intermediate_leaf_node_size = SIZEOF(radv_ir_instance_node);
+      output_leaf_node_size = SIZEOF(radv_bvh_instance_node);
+      break;
+   }
+
+   uint32_t intermediate_leaf_nodes_size = args.leaf_node_count * intermediate_leaf_node_size;
+   uint32_t dst_leaf_offset =
+      align(SIZEOF(radv_accel_struct_header), 64) + SIZEOF(radv_bvh_box32_node);
+   uint32_t dst_internal_offset = dst_leaf_offset + args.leaf_node_count * output_leaf_node_size;
+
+   REF(radv_ir_box_node) intermediate_internal_nodes =
+      REF(radv_ir_box_node)OFFSET(args.intermediate_bvh, intermediate_leaf_nodes_size);
+   radv_ir_box_node src =
+      DEREF(INDEX(radv_ir_box_node, intermediate_internal_nodes, global_id));
+
+   REF(radv_bvh_box32_node) out_internal_nodes =
+      REF(radv_bvh_box32_node) OFFSET(args.output_bvh, dst_internal_offset);
+   REF(radv_bvh_box32_node) dst_node = INDEX(radv_bvh_box32_node, out_internal_nodes, global_id);
+
+   if (global_id == args.internal_node_count - 1) {
+      dst_node = REF(radv_bvh_box32_node)
+         OFFSET(args.output_bvh, id_to_offset(RADV_BVH_ROOT_NODE));
+   }
+
+   uint32_t found_child_count = 0;
+   uint32_t children[4] = {NULL_NODE_ID, NULL_NODE_ID, NULL_NODE_ID, NULL_NODE_ID};
+
+   for (uint32_t i = 0; i < 2; ++i)
+      if (src.children[i] != NULL_NODE_ID)
+         children[found_child_count++] = src.children[i];
+
+   while (found_child_count < 4) {
+      uint32_t collapsed_child_index;
+      float largest_surface_area = 0.0f;
+
+      for (uint32_t i = 0; i < found_child_count; ++i) {
+         if (ir_id_to_type(children[i]) != radv_ir_node_internal)
+            continue;
+
+         AABB bounds =
+            load_aabb(REF(radv_ir_node)OFFSET(args.intermediate_bvh,
+                                              ir_id_to_offset(children[i])));
+
+         float surface_area = aabb_surface_area(bounds);
+         if (surface_area > largest_surface_area) {
+            largest_surface_area = surface_area;
+            collapsed_child_index = i;
+         }
+      }
+
+      if (largest_surface_area > 0.0f) {
+         REF(radv_ir_box_node) child_node =
+            REF(radv_ir_box_node)OFFSET(args.intermediate_bvh,
+                                        ir_id_to_offset(children[collapsed_child_index]));
+         uint32_t grandchildren[2] = DEREF(child_node).children;
+         uint32_t valid_grandchild_count = 0;
+
+         if (grandchildren[1] != NULL_NODE_ID)
+            ++valid_grandchild_count;
+
+         if (grandchildren[0] != NULL_NODE_ID)
+            ++valid_grandchild_count;
+         else
+            grandchildren[0] = grandchildren[1];
+
+         if (valid_grandchild_count > 1)
+            children[found_child_count++] = grandchildren[1];
+
+         if (valid_grandchild_count > 0)
+            children[collapsed_child_index] = grandchildren[0];
+      } else
+         break;
+   }
+
+   for (uint32_t i = 0; i < found_child_count; ++i) {
+      uint32_t type = ir_id_to_type(children[i]);
+      uint32_t offset = ir_id_to_offset(children[i]);
+      uint32_t dst_offset;
+
+      if (offset < intermediate_leaf_nodes_size) {
+         uint32_t child_index = offset / intermediate_leaf_node_size;
+         dst_offset = dst_leaf_offset + child_index * output_leaf_node_size;
+      } else {
+         uint32_t offset_in_internal_nodes = offset - intermediate_leaf_nodes_size;
+         uint32_t child_index = offset_in_internal_nodes / SIZEOF(radv_ir_box_node);
+         dst_offset = dst_internal_offset + child_index * SIZEOF(radv_bvh_box32_node);
+      }
+
+      AABB child_aabb =
+         load_aabb(REF(radv_ir_node)OFFSET(args.intermediate_bvh, offset));
+
+      DEREF(dst_node).coords[i][0][0] = child_aabb.min.x;
+      DEREF(dst_node).coords[i][0][1] = child_aabb.min.y;
+      DEREF(dst_node).coords[i][0][2] = child_aabb.min.z;
+      DEREF(dst_node).coords[i][1][0] = child_aabb.max.x;
+      DEREF(dst_node).coords[i][1][1] = child_aabb.max.y;
+      DEREF(dst_node).coords[i][1][2] = child_aabb.max.z;
+      children[i] = pack_node_id(dst_offset, ir_type_to_bvh_type(type));
+   }
+
+   for (uint i = found_child_count; i < 4; ++i) {
+      for (uint vec = 0; vec < 2; ++vec)
+         for (uint comp = 0; comp < 3; ++comp)
+            DEREF(dst_node).coords[i][vec][comp] = NAN;
+   }
+
+   DEREF(dst_node).children = children;
+
+   if (global_id == args.internal_node_count - 1) {
+      REF(radv_accel_struct_header) header = REF(radv_accel_struct_header)(args.output_bvh);
+      DEREF(header).aabb = src.base.aabb;
+   }
+}

src/amd/vulkan/bvh/meson.build

@@ -23,6 +23,7 @@ bvh_shaders = [
   'lbvh_internal.comp',
   'leaf.comp',
   'morton.comp',
+  'converter_internal.comp',
 ]
 
 bvh_include_dir = meson.source_root() + '/src/amd/vulkan/bvh'