RenderPassTrace

#include "Raytracer.hlsli"

cbuffer Params : register(b0)
{
    float4 ResolutionParams;
    float4 BoundsMin;
    float4 BoundsMax;
}

RWStructuredBuffer<Ray> Rays : register(u0);
RWTexture2D<float4> Output: register(u1);
RWStructuredBuffer<float4> VertexData: register(u2);
RWStructuredBuffer<MemoryNode> VertexNodes: register(u3);
RWStructuredBuffer<uint> IndexData: register(u4);
RWStructuredBuffer<MemoryNode> IndexNodes: register(u5);
RWStructuredBuffer<GeometryNode> Geometries: register(u6);
RWStructuredBuffer<InstanceNode> Instances: register(u7);
RWStructuredBuffer<BVHNode> ASTreeData: register(u8);
RWStructuredBuffer<MemoryNode> ASTreeNodes: register(u9);
RWStructuredBuffer<uint> ASIndexData: register(u10);
RWStructuredBuffer<MemoryNode> ASIndexNodes: register(u11);
RWStructuredBuffer<float4> WoopData: register(u12);
RWStructuredBuffer<MemoryNode> WoopNodes: register(u13);

bool IntersectRayTriangle(float4 origin, float4 direction, float4 v0, float4 v1, float4 v2, out float distance, out float u, out float v)
{
    float4 e1 = v1 - v0;
    float4 e2 = v2 - v0;
    float4 pvec = float4(cross(direction.xyz, e2.xyz), 0.0f);
    float det = dot(e1, pvec);
    if (det > -1e-6 && det < 1e-6)
    {
        return false;
    }
    float inv_det = 1.0f / det;

    float4 tvec = origin - v0;
    u = dot(tvec, pvec) * inv_det;
    if (u < 0.0f || u > 1.0f)
    {
        return false;
    }

    float4 qvec = float4(cross(tvec.xyz, e1.xyz), 0.0f);
    v = dot(direction, qvec) * inv_det;
    if (v < 0.0f || u + v > 1.0f)
    {
        return false;
    }

    distance = dot(e2, qvec) * inv_det;
    return (distance > 0.0f);
}

#define BVH_STACK_SIZE 64

[numthreads(32, 32, 1)]
void RenderPass(uint3 GI : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID)
{
    uint width = asuint(ResolutionParams.x);
    uint height = asuint(ResolutionParams.y);

    int id = DTid.y * width + DTid.x;

    Ray r = Rays[id];

    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;
    bool hit = false;

    float4 temp = float4(0.0f, 0.0f, 0.0f, 1.0f);

    InstanceNode instance = Instances[0];

    int i;

    // Traversal (use for for testing)
    [loop] for (i = 0; i < 1000; i++)
    //while (node_id != 0xFFFFFFFF) // This is significantly slower on RDNA 2, not sure about others
    {
        temp.x += 0.001f;

        [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 (stack_ptr == meshbvh_stack_ptr)
                {
                    meshbvh_stack_ptr = -1;
                    o = r.Origin;
                    d = r.Direction;
                    inv = r.Inverse;
                    oinv = o * inv;
                }

                node_id = stack[stack_ptr];
                stack_ptr--;
            }
            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;
                    }
                }
            }
        }
        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;
        }
        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;
                    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];

                    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;

                                //id = prims_ids[n];
                            }
                        }
                    }
                }
            }

            if (stack_ptr == meshbvh_stack_ptr)
            {
                meshbvh_stack_ptr = -1;
                o = r.Origin;
                d = r.Direction;
                inv = r.Inverse;
                oinv = o * inv;
            }

            node_id = stack[stack_ptr];
            stack_ptr--;
        }

        if (node_id == 0xFFFFFFFF)
        {
            break;
        }
    }

    //Output[DTid.xy] = float4(temp.x, temp.x * 0.1f, temp.x * 0.01f, 1.0f);

    Output[DTid.xy] = float4(bU, bV, temp.x, temp.w);
}