radv/bvh: Add custom leaf node builder

This custom builder implements fine-grained instance node bounds
calculation by looking at all AABBs at tree depth 2.

Shaves off 0.3ms in the start scene for Indiana Jones: The Great Circle
on Deck (roughly 29.1ms->28.7ms).

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/32797>

src/amd/vulkan/bvh/leaf.comp

@@ -0,0 +1,98 @@
+/*
+ * Copyright © 2024 Valve Corporation
+ *
+ * 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_shader_explicit_arithmetic_types_float16 : require
+#extension GL_EXT_scalar_block_layout : require
+#extension GL_EXT_buffer_reference : require
+#extension GL_EXT_buffer_reference2 : require
+#extension GL_KHR_shader_subgroup_vote : require
+#extension GL_KHR_shader_subgroup_arithmetic : require
+#extension GL_KHR_shader_subgroup_ballot : require
+
+#include "build_helpers.h"
+
+mat3 abs(mat3 in_mat) {
+    return mat3(abs(in_mat[0]), abs(in_mat[1]), abs(in_mat[2]));
+}
+
+void
+update_instance_aabb(inout vk_aabb instance_aabb, vk_aabb blas_aabb, mat3x4 transform)
+{
+    /* https://zeux.io/2010/10/17/aabb-from-obb-with-component-wise-abs */
+    vec3 blas_aabb_center = (blas_aabb.max + blas_aabb.min) / 2.0;
+    vec3 blas_aabb_extent = (blas_aabb.max - blas_aabb.min) / 2.0;
+
+    blas_aabb_center = vec4(blas_aabb_center, 1.0f) * transform;
+    blas_aabb_extent = blas_aabb_extent * abs(mat3(transform));
+
+    instance_aabb.min = min(instance_aabb.min, blas_aabb_center - blas_aabb_extent);
+    instance_aabb.max = max(instance_aabb.max, blas_aabb_center + blas_aabb_extent);
+}
+
+vk_aabb
+calculate_fine_instance_node_bounds(VOID_REF bvh, mat3x4 transform)
+{
+    REF(radv_accel_struct_header) header = REF(radv_accel_struct_header)(bvh);
+    bvh = OFFSET(bvh, DEREF(header).bvh_offset);
+    REF(radv_bvh_box32_node) root = REF(radv_bvh_box32_node)(bvh);
+
+    uint32_t children[4] = DEREF(root).children;
+
+    vk_aabb aabb;
+    aabb.min = vec3(INFINITY);
+    aabb.max = vec3(-INFINITY);
+    if (children[0] == 0xFFFFFFFF) {
+        aabb.min = vec3(NAN);
+        aabb.max = vec3(NAN);
+    }
+
+    for (uint32_t child_idx = 0; child_idx < 4; ++child_idx) {
+        if (children[child_idx] == 0xFFFFFFFF)
+            break;
+        if (id_to_type(children[child_idx]) != radv_bvh_node_box32) {
+            update_instance_aabb(aabb, DEREF(root).coords[child_idx], transform);
+            continue;
+        }
+
+        radv_bvh_box32_node node = DEREF(REF(radv_bvh_box32_node)OFFSET(bvh, id_to_offset(children[child_idx])));
+
+        for (uint32_t grandchild_idx = 0; grandchild_idx < 4; ++grandchild_idx) {
+            if (node.children[grandchild_idx] == 0xFFFFFFFF)
+                break;
+            update_instance_aabb(aabb, node.coords[grandchild_idx], transform);
+        }
+    }
+
+    return aabb;
+}
+
+#define CALCULATE_FINE_INSTANCE_NODE_BOUNDS calculate_fine_instance_node_bounds
+
+#include "leaf.h"

src/amd/vulkan/bvh/meson.build

@@ -28,6 +28,16 @@ bvh_shaders = [
     'update',
     [],
   ],
+  [
+    'leaf.comp',
+    'radv_leaf',
+    ['ALWAYS_ACTIVE=0'],
+  ],
+  [
+    'leaf.comp',
+    'radv_leaf_always_active',
+    ['ALWAYS_ACTIVE=1'],
+  ],
 ]
 
 bvh_include_dir = dir_source_root + '/src/amd/vulkan/bvh'

src/amd/vulkan/radv_acceleration_structure.c

@@ -36,6 +36,14 @@ static const uint32_t update_spv[] = {
 #include "bvh/update.spv.h"
 };
 
+static const uint32_t leaf_spv[] = {
+#include "bvh/radv_leaf.spv.h"
+};
+
+static const uint32_t leaf_always_active_spv[] = {
+#include "bvh/radv_leaf_always_active.spv.h"
+};
+
 struct acceleration_structure_layout {
    uint32_t geometry_info_offset;
    uint32_t bvh_offset;
@@ -649,6 +657,10 @@ static const struct vk_acceleration_structure_build_ops build_ops = {
    .init_update_scratch = radv_init_update_scratch,
    .update_bind_pipeline[0] = radv_update_bind_pipeline,
    .update_as[0] = radv_update_as,
+   .leaf_spirv_override = leaf_spv,
+   .leaf_spirv_override_size = sizeof(leaf_spv),
+   .leaf_always_active_spirv_override = leaf_always_active_spv,
+   .leaf_always_active_spirv_override_size = sizeof(leaf_always_active_spv),
 };
 
 static void