mesa/src/amd/vulkan/bvh/update.comp

/*
 * Copyright © 2023 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_memory_scope_semantics : require

layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;

#include "build_interface.h"

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

uint32_t fetch_parent_node(VOID_REF bvh, uint32_t node)
{
    uint64_t addr = bvh - node / 8 * 4 - 4;
    return DEREF(REF(uint32_t)(addr));
}

void main() {
    uint32_t bvh_offset = DEREF(args.src).bvh_offset;

    VOID_REF src_bvh = OFFSET(args.src, bvh_offset);
    VOID_REF dst_bvh = OFFSET(args.dst, bvh_offset);

    uint32_t leaf_node_size;
    if (args.geom_data.geometry_type == VK_GEOMETRY_TYPE_TRIANGLES_KHR)
        leaf_node_size = SIZEOF(radv_bvh_triangle_node);
    else if (args.geom_data.geometry_type == VK_GEOMETRY_TYPE_AABBS_KHR)
        leaf_node_size = SIZEOF(radv_bvh_aabb_node);
    else
        leaf_node_size = SIZEOF(radv_bvh_instance_node);

    uint32_t leaf_node_id = args.geom_data.first_id + gl_GlobalInvocationID.x;
    uint32_t first_leaf_offset = id_to_offset(RADV_BVH_ROOT_NODE) + SIZEOF(radv_bvh_box32_node);

    uint32_t dst_offset = leaf_node_id * leaf_node_size + first_leaf_offset;
    VOID_REF dst_ptr = OFFSET(dst_bvh, dst_offset);
    uint32_t src_offset = gl_GlobalInvocationID.x * args.geom_data.stride;

    radv_aabb bounds;
    bool is_active;
    if (args.geom_data.geometry_type == VK_GEOMETRY_TYPE_TRIANGLES_KHR) {
        is_active = build_triangle(bounds, dst_ptr, args.geom_data, gl_GlobalInvocationID.x);
    } else if (args.geom_data.geometry_type == VK_GEOMETRY_TYPE_AABBS_KHR) {
        VOID_REF src_ptr = OFFSET(args.geom_data.data, src_offset);
        is_active = build_aabb(bounds, src_ptr, dst_ptr, args.geom_data.geometry_id, gl_GlobalInvocationID.x);
    } else {
        VOID_REF src_ptr = OFFSET(args.geom_data.data, src_offset);
        /* arrayOfPointers */
        if (args.geom_data.stride == 8) {
            src_ptr = DEREF(REF(VOID_REF)(src_ptr));
        }

        is_active = build_instance(bounds, src_ptr, dst_ptr, gl_GlobalInvocationID.x);
    }

    if (!is_active)
        return;

    DEREF(INDEX(radv_aabb, args.leaf_bounds, leaf_node_id)) = bounds;
    memoryBarrier(gl_ScopeDevice,
        gl_StorageSemanticsBuffer,
        gl_SemanticsAcquireRelease | gl_SemanticsMakeAvailable | gl_SemanticsMakeVisible);

    uint32_t node_id = pack_node_id(dst_offset, 0);
    uint32_t parent_id = fetch_parent_node(src_bvh, node_id);
    uint32_t internal_nodes_offset = first_leaf_offset + args.leaf_node_count * leaf_node_size;
    while (parent_id != RADV_BVH_INVALID_NODE) {
        uint32_t offset = id_to_offset(parent_id);

        uint32_t parent_index = (offset - internal_nodes_offset) / SIZEOF(radv_bvh_box32_node) + 1;
        if (parent_id == RADV_BVH_ROOT_NODE)
            parent_index = 0;

        /* Make accesses to internal nodes in dst_bvh available and visible */
        memoryBarrier(gl_ScopeDevice,
                      gl_StorageSemanticsBuffer,
                      gl_SemanticsAcquireRelease | gl_SemanticsMakeAvailable | gl_SemanticsMakeVisible);

        radv_bvh_box32_node node = DEREF(REF(radv_bvh_box32_node)OFFSET(src_bvh, offset));
        uint32_t valid_child_count = 0;
        for (uint32_t i = 0; i < 4; ++valid_child_count, ++i)
            if (node.children[i] == RADV_BVH_INVALID_NODE)
                break;

        /* 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.
         */
        uint32_t ready_child_count = atomicAdd(
            DEREF(INDEX(uint32_t, args.internal_ready_count, parent_index)), 1, gl_ScopeDevice,
            gl_StorageSemanticsBuffer,
            gl_SemanticsAcquireRelease | gl_SemanticsMakeAvailable | gl_SemanticsMakeVisible);

        if (ready_child_count != valid_child_count - 1)
            break;

        for (uint32_t i = 0; i < valid_child_count; ++i) {
            uint32_t child_offset = id_to_offset(node.children[i]);
            if (child_offset == dst_offset)
                node.coords[i] = bounds;
            else if (child_offset >= internal_nodes_offset) {
                radv_aabb child_bounds = radv_aabb(vec3(INFINITY), vec3(-INFINITY));
                radv_bvh_box32_node child_node = DEREF(REF(radv_bvh_box32_node)OFFSET(dst_bvh, child_offset));
                for (uint32_t j = 0; j < 4; ++j) {
                    if (child_node.children[j] == RADV_BVH_INVALID_NODE)
                        break;
                    child_bounds.min = min(child_bounds.min, child_node.coords[j].min);
                    child_bounds.max = max(child_bounds.max, child_node.coords[j].max);
                }
                node.coords[i] = child_bounds;
            } else {
                uint32_t child_index = (child_offset - first_leaf_offset) / leaf_node_size;
                node.coords[i] = DEREF(INDEX(radv_aabb, args.leaf_bounds, child_index));
            }
        }

        DEREF(REF(radv_bvh_box32_node)OFFSET(dst_bvh, offset)) = node;

        if (parent_id == RADV_BVH_ROOT_NODE) {
            radv_aabb root_bounds = radv_aabb(vec3(INFINITY), vec3(-INFINITY));
            for (uint32_t i = 0; i < valid_child_count; ++i) {
                root_bounds.min = min(root_bounds.min, node.coords[i].min);
                root_bounds.max = max(root_bounds.max, node.coords[i].max);
            }
            DEREF(args.dst).aabb = root_bounds;
        }

        parent_id = fetch_parent_node(src_bvh, parent_id);
    }
}