//
// Copyright (C) 2009-2021 Intel Corporation
//
// SPDX-License-Identifier: MIT
//
//

module traversal_shader;

kernel_module morton_kernels ("traversal_shader.cl")
{
    links lsc_intrinsics;

    kernel TS_primrefs_from_instances       < kernelFunction = "TS_primrefs_from_instances" >;
    kernel TS_primrefs_from_instances_indirect       < kernelFunction = "TS_primrefs_from_instances_indirect" >;
    kernel TS_primrefs_from_instances_ptrs  < kernelFunction = "TS_primrefs_from_instances_pointers" >;
    kernel TS_primrefs_from_instances_ptrs_indirect  < kernelFunction = "TS_primrefs_from_instances_pointers_indirect" >;
    kernel TS_update_instance_leaves        < kernelFunction = "TS_update_instance_leaves" >;
    kernel TS_Refit_per_one_startpoint_sg   < kernelFunction = "TS_Refit_per_one_startpoint_sg" >;
    kernel TS_fixup_leaves                  < kernelFunction = "TS_fixup_leaves" >;
}

struct MKTSBuildArgs
{
    qword build_globals;
    qword bvh_buffer;
    qword instance_descs;
    qword build_primref_buffer;
    qword aabb_buffer;
    qword is_procedural_buffer;
    qword leaf_creation_index_buffer;
    dword aabb_stride;
    dword num_instances;
    dword leaf_creation_index_stride;
};

const BUILD_PRIMREFS_GROUPSIZE = 16;


metakernel TS_build_primrefs( MKTSBuildArgs build_state, dword allowUpdate )
{
    define num_groups((build_state.num_instances + BUILD_PRIMREFS_GROUPSIZE - 1) / BUILD_PRIMREFS_GROUPSIZE);
    dispatch TS_primrefs_from_instances(num_groups, 1, 1) args(
        build_state.build_globals,
        build_state.bvh_buffer,
        build_state.instance_descs,
        build_state.num_instances,
        build_state.build_primref_buffer,
        build_state.aabb_buffer,
        build_state.is_procedural_buffer,
        build_state.aabb_stride,
        allowUpdate
    );

}

metakernel TS_build_primrefs_indirect(MKTSBuildArgs build_state, qword indirectBuildRangeInfo, dword allowUpdate)
{
    define num_groups  REG0;
    define groupsize_1 REG1; // groupsize - 1
    define C_4         REG2;

    // init with primitiveCount
    num_groups = load_dword(indirectBuildRangeInfo);
    groupsize_1 = 15; // BUILD_PRIMREFS_GROUPSIZE - 1
    C_4 = 4;          // log_2(BUILD_PRIMREFS_GROUPSIZE)

    num_groups = num_groups + groupsize_1;
    num_groups = num_groups >> C_4; // num_groups / BUILD_PRIMREFS_GROUPSIZE;

    DISPATCHDIM_X = num_groups.lo;
    DISPATCHDIM_Y = 1;
    DISPATCHDIM_Z = 1;

    dispatch_indirect TS_primrefs_from_instances_indirect args(
        build_state.build_globals,
        build_state.bvh_buffer,
        build_state.instance_descs,
        build_state.build_primref_buffer,
        build_state.aabb_buffer,
        build_state.is_procedural_buffer,
        build_state.aabb_stride,
        allowUpdate,
        indirectBuildRangeInfo
    );

}

metakernel TS_build_primrefs_array_of_pointers( MKTSBuildArgs build_state, dword allowUpdate )
{
    define num_groups((build_state.num_instances + BUILD_PRIMREFS_GROUPSIZE - 1) / BUILD_PRIMREFS_GROUPSIZE);
    dispatch TS_primrefs_from_instances_ptrs(num_groups, 1, 1) args(
        build_state.build_globals,
        build_state.bvh_buffer,
        build_state.instance_descs,
        build_state.num_instances,
        build_state.build_primref_buffer,
        build_state.aabb_buffer,
        build_state.is_procedural_buffer,
        build_state.aabb_stride,
        allowUpdate
    );
}

metakernel
TS_build_primrefs_array_of_pointers_indirect(MKTSBuildArgs build_state, qword indirectBuildRangeInfo, dword allowUpdate)
{
    define num_groups  REG0;
    define groupsize_1 REG1; // groupsize - 1
    define C_4         REG2;

    // init with primitiveCount
    num_groups = load_dword(indirectBuildRangeInfo);
    groupsize_1 = 15; // BUILD_PRIMREFS_GROUPSIZE - 1
    C_4 = 4;          // log_2(BUILD_PRIMREFS_GROUPSIZE)

    num_groups = num_groups + groupsize_1;
    num_groups = num_groups >> C_4; // num_groups / BUILD_PRIMREFS_GROUPSIZE;

    DISPATCHDIM_X = num_groups.lo;
    DISPATCHDIM_Y = 1;
    DISPATCHDIM_Z = 1;

    dispatch_indirect TS_primrefs_from_instances_ptrs_indirect args(
        build_state.build_globals,
        build_state.bvh_buffer,
        build_state.instance_descs,
        build_state.build_primref_buffer,
        build_state.aabb_buffer,
        build_state.is_procedural_buffer,
        build_state.aabb_stride,
        allowUpdate,
        indirectBuildRangeInfo
    );
}




const UPDATE_INSTANCE_LEAVES_GROUPSIZE = 16;

struct MKTSUpdateArgs
{
    qword bvh_buffer;
    qword instance_descs;
    qword instance_descs_ptrs;
    qword aabb_buffer;
    qword is_procedural_buffer;
    qword refit_scratch;
    dword aabb_stride;
    dword num_instances;
};

metakernel TS_update_instance_leaves( MKTSUpdateArgs update_state )
{
    define num_groups((update_state.num_instances + UPDATE_INSTANCE_LEAVES_GROUPSIZE - 1) / UPDATE_INSTANCE_LEAVES_GROUPSIZE);
    dispatch TS_update_instance_leaves(num_groups, 1, 1) args(
        update_state.bvh_buffer,
        update_state.instance_descs,
        update_state.instance_descs_ptrs,
        update_state.refit_scratch,
        update_state.aabb_buffer,
        update_state.is_procedural_buffer,
        update_state.aabb_stride
    );
}

metakernel TS_update_instance_leaves_indirect( MKTSUpdateArgs update_state, qword indirectBuildRangeInfo )
{
    define num_groups  REG0;
    define groupsize_1 REG1; // groupsize - 1
    define C_4         REG2;

    // init with primitiveCount
    num_groups = load_dword(indirectBuildRangeInfo);
    groupsize_1 = 15; // UPDATE_INSTANCE_LEAVES_GROUPSIZE - 1
    C_4 = 4;          // log_2(UPDATE_INSTANCE_LEAVES_GROUPSIZE)

    num_groups = num_groups + groupsize_1;
    num_groups = num_groups >> C_4; // num_groups / UPDATE_INSTANCE_LEAVES_GROUPSIZE;

    DISPATCHDIM_X = num_groups.lo;
    DISPATCHDIM_Y = 1;
    DISPATCHDIM_Z = 1;

    // need to add indirect offset?
    dispatch_indirect TS_update_instance_leaves args(
        update_state.bvh_buffer,
        update_state.instance_descs,
        update_state.instance_descs_ptrs,
        update_state.refit_scratch,
        update_state.aabb_buffer,
        update_state.is_procedural_buffer,
        update_state.aabb_stride
    );
}

metakernel TS_refit(MKTSUpdateArgs update_state, qword bvh_inner_nodes_start_value, qword bvh_inner_nodes_end )
{
    REG0 = bvh_inner_nodes_start_value;
    REG1.lo = load_dword(bvh_inner_nodes_end);
    REG1.hi = 0;
    REG2 = REG1 - REG0;

    DISPATCHDIM_X = REG2.lo;
    DISPATCHDIM_Y = 1;
    DISPATCHDIM_Z = 1;

    dispatch_indirect TS_Refit_per_one_startpoint_sg
    args(
        update_state.bvh_buffer,
        update_state.refit_scratch,
        update_state.is_procedural_buffer
    );
}


const FIXUP_LEAVES_NODES_PER_GROUP = 2;

metakernel TS_fixup_leaves(MKTSBuildArgs build_state, qword bvh_inner_nodes_start_value, qword bvh_inner_nodes_end )
{
    define ONE REG3;

    ONE = 1;
    REG0 = bvh_inner_nodes_start_value;
    REG1.lo = load_dword(bvh_inner_nodes_end);
    REG1.hi = 0;
    REG2 = REG1 - REG0;
    REG2 = REG2 + ONE;
    REG2 = REG2 >> ONE;

    DISPATCHDIM_X = REG2.lo;
    DISPATCHDIM_Y = 1;
    DISPATCHDIM_Z = 1;

    dispatch_indirect TS_fixup_leaves
        args(
            build_state.bvh_buffer,
            build_state.leaf_creation_index_buffer,
            build_state.build_primref_buffer,
            build_state.leaf_creation_index_stride
        );

}