Traversal bvh2

///////////////////////////////////////////////////////////////////////////////////////////////////
//
// TraversalBvh2.hlsli
//
// Implements ray traversal through multi-level BVH-2 acceleration structure.
//
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef __TRAVERSAL_BVH2__HLSLI__
#define __TRAVERSAL_BVH2__HLSLI__

#include "../Raytracer.hlsli"

// Definition to compact funciton parameters into TraceRayCompute funciton (BVH-2 variant)
#define TRACE_RAY_PARAMS RWStructuredBuffer<GeometryNode> Geometries,\
                         RWStructuredBuffer<InstanceNode> Instances,\
                         RWStructuredBuffer<MemoryNode> ASTreeNodes,\
                         RWStructuredBuffer<BVH2Node> ASTreeData,\
                         RWStructuredBuffer<MemoryNode> ASIndexNodes,\
                         RWStructuredBuffer<uint> ASIndexData,\
                         RWStructuredBuffer<MemoryNode> WoopNodes,\
                         RWStructuredBuffer<float4> WoopData

// Definition to compact argument passing into TraceRayCompute function (BVH-2 variant)
#define TRACE_RAY_ARGS  Geometries,\
                        Instances,\
                        ASTreeNodes,\
                        ASTreeData,\
                        ASIndexNodes,\
                        ASIndexData,\
                        WoopNodes,\
                        WoopData

/// <summary>
/// Performs ray traversal through acceleration structure for single ray.
///
/// This function performs traversal through binary Bounding Volume Hierarchy (BVH-2). Result of
/// this funciton is represented by barycentric coordinates, primitive ID, geometry ID and distance
/// along the ray to hitpoint.
/// </summary>
/// <param name="r">Ray to trace.</param>
/// <param name="Geometries">Buffer of GeometryNode - holds all definition for geometries in the scene</param>
/// <param name="Instances">Buffer of InstanceNode - holds all geometry instances definitions in the scene</param>
/// <param name="ASTreeNodes">Buffer of memory nodes - each representing single BVH node data definition/offsets</param>
/// <param name="ASTreeData">Buffer of BVH nodes - BVH node data</param>
/// <param name="ASIndexNodes">Buffer of memory nodes - each representing single BVH index data definition/offsets</param>
/// <param name="ASIndexData">Buffer of BVH indices - BVH index data</param>
/// <param name="WoopNodes">Buffer of memory nodes - each representing definition/offsets into data buffer holding woop triangle data</param>
/// <param name="WoopData">Buffer of woop triangle data - Woop triangle geometry data</param>
/// <returns>
/// 4-component vector, where:
///   1st component has packed U/V barycentric coordinates (see PackFloat2/UnpackFloat2)
///   2nd component distance along the ray to hit
///   3rd component primitive ID (unsigned int)
///   4th component geometry ID (unsigned int)
/// </returns>
float4 TraceRayCompute(Ray r, TRACE_RAY_PARAMS)
{
    float4 o = r.Origin;
    float4 d = r.Direction;
    float4 inv = r.Inverse;
    float4 oinv = o * inv;

    uint node_id = 0;
    uint stack[BVH_STACK_SIZE];
    uint stack_ptr = 0;
    stack[stack_ptr] = 0xFFFFFFFF;
    int meshbvh_stack_ptr = -1;

    float tmin = 0.0f;
    float tmax = 10000.0f;
    float bU = 0.0f;
    float bV = 0.0f;
    float dist = tmax;
    uint prim_id = 0;
    uint geometryID = 0;
    bool hit = false;

    InstanceNode instance = Instances[0];

    int i;

    // Traversal (use for for testing)
    [loop]
    for (i = 0; i < 1024; i++)
    //while (node_id != 0xFFFFFFFF) // This is significantly slower on RDNA 2, not sure about others
    {
        // When interior node was intersected
        [branch]
        if (ASTreeData[node_id].PrimitiveCount == 0)
        {
            // Fetch children bounding boxes
            float4 n0xy = ASTreeData[node_id].LXY;
            float4 n1xy = ASTreeData[node_id].RXY;
            float4 nz = ASTreeData[node_id].LRZ;

            // Test against child AABBs
            float c0lox = n0xy.x * inv.x - oinv.x;
            float c0hix = n0xy.y * inv.x - oinv.x;
            float c0loy = n0xy.z * inv.y - oinv.y;
            float c0hiy = n0xy.w * inv.y - oinv.y;
            float c0loz = nz.x * inv.z - oinv.z;
            float c0hiz = nz.y * inv.z - oinv.z;
            float c1loz = nz.z * inv.z - oinv.z;
            float c1hiz = nz.w * inv.z - oinv.z;
            float c0min = max(max(min(c0lox, c0hix), min(c0loy, c0hiy)), max(min(c0loz, c0hiz), tmin));
            float c0max = min(min(max(c0lox, c0hix), max(c0loy, c0hiy)), min(max(c0loz, c0hiz), tmax));
            float c1lox = n1xy.x * inv.x - oinv.x;
            float c1hix = n1xy.y * inv.x - oinv.x;
            float c1loy = n1xy.z * inv.y - oinv.y;
            float c1hiy = n1xy.w * inv.y - oinv.y;
            float c1min = max(max(min(c1lox, c1hix), min(c1loy, c1hiy)), max(min(c1loz, c1hiz), tmin));
            float c1max = min(min(max(c1lox, c1hix), max(c1loy, c1hiy)), min(max(c1loz, c1hiz), tmax));

            bool traverseChild0 = (c0max >= c0min);
            bool traverseChild1 = (c1max >= c1min);

            // If no children was hit, get node from stack
            if (!traverseChild0 && !traverseChild1)
            {
                // If we're in BLAS and we're on entrypoint, then reset the ray as the traversal will
                // continue in TLAS
                if (stack_ptr == meshbvh_stack_ptr)
                {
                    meshbvh_stack_ptr = -1;
                    o = r.Origin;
                    d = r.Direction;
                    inv = r.Inverse;
                    oinv = o * inv;
                }

                // Pop from stack
                node_id = stack[stack_ptr];
                stack_ptr--;
            }
            // If either (or both children) were it - traverse next on (along the ray direction),
            // and if both were hit, push the other one to stack
            else if (traverseChild0 || traverseChild1)
            {
                uint first_child = node_id + 1;
                uint second_child = ASTreeData[node_id].PrimitiveOffset;
                node_id = (traverseChild0) ? first_child : second_child;

                if (traverseChild0 && traverseChild1)
                {
                    if (c1min < c0min)
                    {
                        node_id = second_child;
                        stack_ptr++;
                        stack[stack_ptr] = first_child;
                    }
                    else
                    {
                        stack_ptr++;
                        stack[stack_ptr] = second_child;
                    }
                }
            }
        }
        // When leaf node was intersected - and we're in TLAS
        //
        // Setup the new ray origin/direction for traversing BLAS and enter the BLAS, record
        // entry-point to BLAS into stack so we know when we will finish BLAS traversal
        else if (ASTreeData[node_id].PrimitiveCount == -1)
        {
            meshbvh_stack_ptr = stack_ptr;

            uint blas_offset = ASTreeData[node_id].PrimitiveOffset;
            uint instance_index = ASIndexData[blas_offset];
            instance = Instances[instance_index];

            node_id = ASTreeNodes[Geometries[instance.GeometryNode].BVHNode + 1].Offset / 64;

            o = mul(r.Origin, instance.TransformInverse);
            d = mul(r.Direction, instance.TransformInverse);
            inv = rcp(d);
            oinv = o * inv;
        }
        // When leaf node was intersected - and we're in BLAS
        //
        // Perform ray-triangle intersection test for all triangles in leaf node
        else
        {
            if (ASTreeData[node_id].PrimitiveCount > 0)
            {
                GeometryNode geom = Geometries[instance.GeometryNode];
                //MemoryNode vbo = VertexNodes[geom.VertexBufferNode];
                //MemoryNode ibo = IndexNodes[geom.IndexBufferNode];
                MemoryNode wbo = WoopNodes[geom.WoopBufferNode];

                uint index_offset = ASIndexNodes[ASTreeData[node_id].PrimitiveOffset].Offset / 4;

                for (uint j = 0; j < ASTreeData[node_id].PrimitiveCount; j++)
                {
                    // Don't trash cache by reading index through it
                    uint tri_idx = ASIndexData[ASTreeData[node_id].PrimitiveOffset + j] * 3;

                    // Fetch data for Woop's triangle
                    float4 r = WoopData[wbo.Offset / 16 + tri_idx + 0];
                    float4 p = WoopData[wbo.Offset / 16 + tri_idx + 1];
                    float4 q = WoopData[wbo.Offset / 16 + tri_idx + 2];

                    // Perform intersection
                    float o_z = r.w - o.x * r.x - o.y * r.y - o.z * r.z;
                    float i_z = 1.0f / (d.x * r.x + d.y * r.y + d.z * r.z);
                    float t = o_z * i_z;

                    if (t > tmin && t < tmax)
                    {
                        float o_x = p.w + o.x * p.x + o.y * p.y + o.z * p.z;
                        float d_x = d.x * p.x + d.y * p.y + d.z * p.z;
                        float u = o_x + t * d_x;

                        if (u >= 0.0f && u <= 1.0f)
                        {
                            float o_y = q.w + o.x * q.x + o.y * q.y + o.z * q.z;
                            float d_y = d.x * q.x + d.y * q.y + d.z * q.z;
                            float v = o_y + t * d_y;

                            if (v >= 0.0f && u + v <= 1.0f)
                            {
                                tmax = t;
                                bU = u;
                                bV = v;
                                hit = true;

                                geometryID = instance.GeometryNode;
                                prim_id = tri_idx;
                            }
                        }
                    }
                }
            }

            // If we're in BLAS and we're on entrypoint, then reset the ray as the traversal will
            // continue in TLAS
            if (stack_ptr == meshbvh_stack_ptr)
            {
                meshbvh_stack_ptr = -1;
                o = r.Origin;
                d = r.Direction;
                inv = r.Inverse;
                oinv = o * inv;
            }

            // Pop from stack
            node_id = stack[stack_ptr];
            stack_ptr--;
        }

        // If we hit entrypoint into the traversal, terminate the traversal as we are finished
        if (node_id == 0xFFFFFFFF)
        {
            break;
        }
    }

    return float4(PackFloat2(bU, bV), tmax, asfloat(prim_id), asfloat(geometryID));
}

#endif