MultiLevelBVH Traversal

struct TopLevelNode
{
    unsigned int bvh_index;
    unsigned int pad[3];
    float4 mat0;
    float4 mat1;
    float4 mat2;
    float4 mat3;
};

// Transform vector with a 4x4 matrix.
float4 mul4x1(float4 vec, float4 mat0, float4 mat1, float4 mat2, float4 mat3)
{
    float4 out;

    // Assume vec.w = 1 to save a multiplication
    out.x = mat0.x * vec.x + mat1.x * vec.y + mat2.x * vec.z + mat3.x;
    out.y = mat0.y * vec.x + mat1.y * vec.y + mat2.y * vec.z + mat3.y;
    out.z = mat0.z * vec.x + mat1.z * vec.y + mat2.z * vec.z + mat3.z;
    out.w = 1.0f;

    return out;
}

// Transform vector with a 3x3 rotation matrix.
// This function ignores the 4 row/col of the matrix.
float4 mul3x1(float4 vec, float4 mat0, float4 mat1, float4 mat2)
{
    float4 out;

    out.x = mat0.x * vec.x + mat1.x * vec.y + mat2.x * vec.z;
    out.y = mat0.y * vec.x + mat1.y * vec.y + mat2.y * vec.z;
    out.z = mat0.z * vec.x + mat1.z * vec.y + mat2.z * vec.z;
    out.w = 0.0f;

    return out;
}

float4 TraceRay_MultiLevel(float4 ray_origin,
    float4 ray_direction,
    __global float4* triangles,
    __global struct BVHStackNode* nodes,
    __global unsigned int* indices,
    __global struct TopLevelNode* toplevelData)
{
    float4 o = ray_origin;
    float4 d = ray_direction;
    float4 inv = native_recip(d);
    float4 oinv = o * inv;

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

    float tmin = 0.0f;
    float tmax = 10000.0f;
    float bU = 0.0f;
    float bV = 0.0f;
    float dist = tmax;

    float leaf = 0.0f;
    int id = -1;

    do
    {
        __global struct BVHStackNode* n = &(nodes[node_id]);
        int prim_count = n->prim_count;

        // Interior node
        if (prim_count == 0)
        {
            // Fetch children bounding boxes
            float4 n0xy = n->n0xy;
            float4 n1xy = n->n1xy;
            float4 nz = n->nz;

            // 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 = ray_origin;
                    d = ray_direction;
                    inv = native_recip(d);
                    oinv = o * inv;
                }

                node_id = stack[stack_ptr];
                stack_ptr--;
            }
            else if (traverseChild0 || traverseChild1)
            {
                unsigned int first_child = node_id + 1;
                unsigned int second_child = n->second_child;
                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 (prim_count < 0)
        {
            meshbvh_stack_ptr = stack_ptr;

            unsigned int bottom_level_offset = n->prim_offset;
            __global unsigned int *bottom_level_ids = &indices[bottom_level_offset];

            int bottom_levels_count = -prim_count;

            int bvh_offset = bottom_level_ids[0];
            int bvh_index = toplevelData[bvh_offset].bvh_index;

            node_id = bvh_index;

            o = mul4x1(ray_origin, toplevelData[bvh_offset].mat0, toplevelData[bvh_offset].mat1, toplevelData[bvh_offset].mat2, toplevelData[bvh_offset].mat3);
            d = mul3x1(ray_direction, toplevelData[bvh_offset].mat0, toplevelData[bvh_offset].mat1, toplevelData[bvh_offset].mat2);
            inv = native_recip(d);
            oinv = o * inv;
        }
        // Leaf node
        else
        {
            unsigned int prim_offset = n->prim_offset;
            __global unsigned int *prims_ids = &indices[prim_offset];

            for (int n = 0; n < prim_count; n++)
            {
                // Don't trash cache by reading index through it
                int tri_idx = prims_ids[n] * 3;
                float4 r = triangles[tri_idx + 0];
                float4 p = triangles[tri_idx + 1];
                float4 q = triangles[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;
                            id = prims_ids[n];
                        }
                    }
                }
            }

            if (stack_ptr == meshbvh_stack_ptr)
            {
                meshbvh_stack_ptr = -1;
                o = ray_origin;
                d = ray_direction;
                inv = native_recip(d);
                oinv = o * inv;
            }

            node_id = stack[stack_ptr];
            stack_ptr--;
        }
    }
    while (node_id != 0xffffffff);

    return (float4)(bU, bV, tmax, as_float(id));
}