xref: /aosp_15_r20/external/mesa3d/src/amd/vulkan/bvh/update.comp (revision 6104692788411f58d303aa86923a9ff6ecaded22)
1/*
2 * Copyright © 2023 Valve Corporation
3 *
4 * SPDX-License-Identifier: MIT
5 */
6
7#version 460
8
9#extension GL_GOOGLE_include_directive : require
10
11#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
12#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
13#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
14#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require
15#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
16#extension GL_EXT_scalar_block_layout : require
17#extension GL_EXT_buffer_reference : require
18#extension GL_EXT_buffer_reference2 : require
19#extension GL_KHR_memory_scope_semantics : require
20
21layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
22
23#include "build_interface.h"
24
25layout(push_constant) uniform CONSTS {
26    update_args args;
27};
28
29uint32_t fetch_parent_node(VOID_REF bvh, uint32_t node)
30{
31    uint64_t addr = bvh - node / 8 * 4 - 4;
32    return DEREF(REF(uint32_t)(addr));
33}
34
35void main() {
36    uint32_t bvh_offset = DEREF(args.src).bvh_offset;
37
38    VOID_REF src_bvh = OFFSET(args.src, bvh_offset);
39    VOID_REF dst_bvh = OFFSET(args.dst, bvh_offset);
40
41    uint32_t leaf_node_size;
42    if (args.geom_data.geometry_type == VK_GEOMETRY_TYPE_TRIANGLES_KHR)
43        leaf_node_size = SIZEOF(radv_bvh_triangle_node);
44    else if (args.geom_data.geometry_type == VK_GEOMETRY_TYPE_AABBS_KHR)
45        leaf_node_size = SIZEOF(radv_bvh_aabb_node);
46    else
47        leaf_node_size = SIZEOF(radv_bvh_instance_node);
48
49    uint32_t leaf_node_id = args.geom_data.first_id + gl_GlobalInvocationID.x;
50    uint32_t first_leaf_offset = id_to_offset(RADV_BVH_ROOT_NODE) + SIZEOF(radv_bvh_box32_node);
51
52    uint32_t dst_offset = leaf_node_id * leaf_node_size + first_leaf_offset;
53    VOID_REF dst_ptr = OFFSET(dst_bvh, dst_offset);
54    uint32_t src_offset = gl_GlobalInvocationID.x * args.geom_data.stride;
55
56    radv_aabb bounds;
57    bool is_active;
58    if (args.geom_data.geometry_type == VK_GEOMETRY_TYPE_TRIANGLES_KHR) {
59        is_active = build_triangle(bounds, dst_ptr, args.geom_data, gl_GlobalInvocationID.x);
60    } else {
61        VOID_REF src_ptr = OFFSET(args.geom_data.data, src_offset);
62        is_active = build_aabb(bounds, src_ptr, dst_ptr, args.geom_data.geometry_id, gl_GlobalInvocationID.x);
63    }
64
65    if (!is_active)
66        return;
67
68    DEREF(INDEX(radv_aabb, args.leaf_bounds, leaf_node_id)) = bounds;
69    memoryBarrier(gl_ScopeDevice,
70        gl_StorageSemanticsBuffer,
71        gl_SemanticsAcquireRelease | gl_SemanticsMakeAvailable | gl_SemanticsMakeVisible);
72
73    uint32_t node_id = pack_node_id(dst_offset, 0);
74    uint32_t parent_id = fetch_parent_node(src_bvh, node_id);
75    uint32_t internal_nodes_offset = first_leaf_offset + args.leaf_node_count * leaf_node_size;
76    while (parent_id != RADV_BVH_INVALID_NODE) {
77        uint32_t offset = id_to_offset(parent_id);
78
79        uint32_t parent_index = (offset - internal_nodes_offset) / SIZEOF(radv_bvh_box32_node) + 1;
80        if (parent_id == RADV_BVH_ROOT_NODE)
81            parent_index = 0;
82
83        /* Make accesses to internal nodes in dst_bvh available and visible */
84        memoryBarrier(gl_ScopeDevice,
85                      gl_StorageSemanticsBuffer,
86                      gl_SemanticsAcquireRelease | gl_SemanticsMakeAvailable | gl_SemanticsMakeVisible);
87
88        REF(radv_bvh_box32_node) src_node = REF(radv_bvh_box32_node)OFFSET(src_bvh, offset);
89        REF(radv_bvh_box32_node) dst_node = REF(radv_bvh_box32_node)OFFSET(dst_bvh, offset);
90        uint32_t children[4];
91        for (uint32_t i = 0; i < 4; ++i)
92            children[i] = DEREF(src_node).children[i];
93
94        uint32_t valid_child_count = 0;
95        for (uint32_t i = 0; i < 4; ++valid_child_count, ++i)
96            if (children[i] == RADV_BVH_INVALID_NODE)
97                break;
98
99        /* Check if all children have been processed. As this is an atomic the last path coming from
100         * a child will pass here, while earlier paths break.
101         */
102        uint32_t ready_child_count = atomicAdd(
103            DEREF(INDEX(uint32_t, args.internal_ready_count, parent_index)), 1, gl_ScopeDevice,
104            gl_StorageSemanticsBuffer,
105            gl_SemanticsAcquireRelease | gl_SemanticsMakeAvailable | gl_SemanticsMakeVisible);
106
107        if (ready_child_count != valid_child_count - 1)
108            break;
109
110        for (uint32_t i = 0; i < 4; ++i)
111            DEREF(dst_node).children[i] = children[i];
112
113        for (uint32_t i = 0; i < valid_child_count; ++i) {
114            uint32_t child_offset = id_to_offset(children[i]);
115            radv_aabb child_bounds;
116            if (child_offset == dst_offset)
117                child_bounds = bounds;
118            else if (child_offset >= internal_nodes_offset) {
119                child_bounds = radv_aabb(vec3(INFINITY), vec3(-INFINITY));
120                REF(radv_bvh_box32_node) child_node = REF(radv_bvh_box32_node)OFFSET(dst_bvh, child_offset);
121                for (uint32_t j = 0; j < 4; ++j) {
122                    if (DEREF(child_node).children[j] == RADV_BVH_INVALID_NODE)
123                        break;
124                    child_bounds.min = min(child_bounds.min, DEREF(child_node).coords[j].min);
125                    child_bounds.max = max(child_bounds.max, DEREF(child_node).coords[j].max);
126                }
127            } else {
128                uint32_t child_index = (child_offset - first_leaf_offset) / leaf_node_size;
129                child_bounds = DEREF(INDEX(radv_aabb, args.leaf_bounds, child_index));
130            }
131
132            DEREF(dst_node).coords[i] = child_bounds;
133        }
134
135        if (parent_id == RADV_BVH_ROOT_NODE) {
136            radv_aabb root_bounds = radv_aabb(vec3(INFINITY), vec3(-INFINITY));
137            for (uint32_t i = 0; i < valid_child_count; ++i) {
138                radv_aabb bounds = DEREF(dst_node).coords[i];
139                root_bounds.min = min(root_bounds.min, bounds.min);
140                root_bounds.max = max(root_bounds.max, bounds.max);
141            }
142            DEREF(args.dst).aabb = root_bounds;
143        }
144
145        parent_id = fetch_parent_node(src_bvh, parent_id);
146    }
147}
148