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

module rebraid;

kernel init_scratch             < source="bvh_rebraid.cl", kernelFunction="rebraid_init_scratch"                        >
kernel chase_instance_ptrs      < source="bvh_rebraid.cl", kernelFunction="rebraid_chase_instance_pointers"             >
kernel calc_aabb                < source="bvh_rebraid.cl", kernelFunction="rebraid_computeAABB_DXR_instances"           >
kernel calc_aabb_indirect                < source="bvh_rebraid.cl", kernelFunction="rebraid_computeAABB_DXR_instances_indirect"           >
kernel calc_aabb_ptr            < source="bvh_rebraid.cl", kernelFunction="rebraid_computeAABB_DXR_instances_pointers"  >
kernel calc_aabb_ptr_indirect            < source="bvh_rebraid.cl", kernelFunction="rebraid_computeAABB_DXR_instances_pointers_indirect"  >
kernel count_splits             < source="bvh_rebraid.cl", kernelFunction="rebraid_count_splits"          >
kernel count_splits_SG          < source="bvh_rebraid.cl", kernelFunction="rebraid_count_splits_SG"       >
kernel count_splits_SG_indirect          < source="bvh_rebraid.cl", kernelFunction="rebraid_count_splits_SG_indirect"       >
kernel build_primrefs           < source="bvh_rebraid.cl", kernelFunction="rebraid_build_primrefs"        >
kernel build_primrefs_indirect           < source="bvh_rebraid.cl", kernelFunction="rebraid_build_primrefs_indirect"        >

//kernel ISA_TEST < source="bvh_rebraid.cl", kernelFunction="ISA_TEST" >
//kernel DEBUG_PRINT < source="bvh_rebraid.cl", kernelFunction="DEBUG_PRINT" >


const PRIMREF_GROUP_SIZE = 256;

const COUNT_SPLITS_GROUP_SIZE = 16;

struct MKRebraidArgs
{
  qword bvh_buffer;
  qword primref_buffer;
  qword global_buffer;
  qword instances_buffer;
  qword rebraid_scratch;
  qword flat_instances_buffer;
  dword num_instances;
  dword num_extra_primrefs;
};

metakernel rebraid(
    MKRebraidArgs Args
  )
{
  dispatch init_scratch(1,1,1) args( Args.rebraid_scratch );
  dispatch calc_aabb(Args.num_instances,1,1) args( Args.bvh_buffer, Args.instances_buffer );
  control( wait_idle );

  //define num_count_groups ((Args.num_instances + (COUNT_SPLITS_GROUP_SIZE-1)) / COUNT_SPLITS_GROUP_SIZE);
  //dispatch count_splits(num_count_groups,1,1) args( Args.bvh_buffer, Args.instances_buffer, Args.rebraid_scratch, Args.num_instances );
  
  dispatch count_splits_SG(Args.num_instances,1,1) args( Args.bvh_buffer, Args.instances_buffer, Args.rebraid_scratch );
  control( wait_idle );

  define num_primref_groups ((Args.num_instances + (PRIMREF_GROUP_SIZE-1)) / PRIMREF_GROUP_SIZE);

  dispatch build_primrefs(num_primref_groups,1,1) args( Args.global_buffer, Args.bvh_buffer, Args.instances_buffer, Args.rebraid_scratch, Args.primref_buffer, Args.num_extra_primrefs, Args.num_instances );
  control( wait_idle );

  //dispatch DEBUG_PRINT(1,1,1) args( Args.global_buffer, Args.instances_buffer, Args.rebraid_scratch, Args.primref_buffer, Args.num_extra_primrefs, Args.num_instances );
}

metakernel rebraid_indirect(MKRebraidArgs Args, qword indirectBuildRangeInfo)
{

    dispatch init_scratch(1, 1, 1) args(Args.rebraid_scratch);

    define num_groups  REG0;
    num_groups = load_dword(indirectBuildRangeInfo);
    DISPATCHDIM_X = num_groups.lo;
    DISPATCHDIM_Y = 1;
    DISPATCHDIM_Z = 1;

    dispatch_indirect calc_aabb_indirect args(Args.bvh_buffer, Args.instances_buffer, indirectBuildRangeInfo);
    control(wait_idle);

    dispatch_indirect count_splits_SG_indirect
        args(Args.bvh_buffer, Args.instances_buffer, Args.rebraid_scratch, indirectBuildRangeInfo);

    define groupsize_1 REG1; // groupsize - 1
    define C_8         REG2;

    groupsize_1 = 255; // PRIMREF_GROUP_SIZE - 1
    C_8 = 8;          // log_2(PRIMREF_GROUP_SIZE)

    num_groups = num_groups + groupsize_1;
    num_groups = num_groups >> C_8; // num_groups / PRIMREF_GROUP_SIZE;
    DISPATCHDIM_X = num_groups.lo;

    control(wait_idle);

    dispatch_indirect build_primrefs_indirect args(
        Args.global_buffer,
        Args.bvh_buffer,
        Args.instances_buffer,
        Args.rebraid_scratch,
        Args.primref_buffer,
        indirectBuildRangeInfo,
        Args.num_extra_primrefs);
    control(wait_idle);
}

metakernel rebraid_ptrs(
    MKRebraidArgs Args
  )
{
  dispatch init_scratch(1,1,1) args( Args.rebraid_scratch );
  dispatch chase_instance_ptrs( Args.num_instances, 1, 1)   args( Args.instances_buffer, Args.flat_instances_buffer );
  dispatch calc_aabb_ptr(Args.num_instances,1,1) args( Args.bvh_buffer, Args.instances_buffer );
  control( wait_idle );

  //define num_count_groups ((Args.num_instances + (COUNT_SPLITS_GROUP_SIZE-1)) / COUNT_SPLITS_GROUP_SIZE);
  //dispatch count_splits(num_count_groups,1,1) args( Args.bvh_buffer, Args.instances_buffer, Args.rebraid_scratch );
  
  dispatch count_splits_SG(Args.num_instances,1,1) args( Args.bvh_buffer, Args.flat_instances_buffer, Args.rebraid_scratch );
  control( wait_idle );

  define num_primref_groups ((Args.num_instances + (PRIMREF_GROUP_SIZE-1)) / PRIMREF_GROUP_SIZE);


  dispatch build_primrefs(num_primref_groups,1,1) args( Args.global_buffer, Args.bvh_buffer, Args.flat_instances_buffer, Args.rebraid_scratch,  Args.primref_buffer, Args.num_extra_primrefs, Args.num_instances );
  control( wait_idle );

}

metakernel rebraid_ptrs_indirect(MKRebraidArgs Args, qword indirectBuildRangeInfo)
{
    dispatch init_scratch(1, 1, 1) args(Args.rebraid_scratch);

    define num_groups  REG0;
    num_groups = load_dword(indirectBuildRangeInfo);
    DISPATCHDIM_X = num_groups.lo;
    DISPATCHDIM_Y = 1;
    DISPATCHDIM_Z = 1;

    dispatch_indirect chase_instance_ptrs
             args(Args.instances_buffer, Args.flat_instances_buffer, indirectBuildRangeInfo);
    dispatch_indirect calc_aabb_ptr_indirect args(Args.bvh_buffer, Args.instances_buffer, indirectBuildRangeInfo);
    control(wait_idle);

    dispatch_indirect count_splits_SG_indirect
        args(Args.bvh_buffer, Args.flat_instances_buffer, Args.rebraid_scratch, indirectBuildRangeInfo);

    define groupsize_1 REG1; // groupsize - 1
    define C_8         REG2;

    groupsize_1 = 255; // PRIMREF_GROUP_SIZE - 1
    C_8 = 8;          // log_2(PRIMREF_GROUP_SIZE)

    num_groups = num_groups + groupsize_1;
    num_groups = num_groups >> C_8; // num_groups / PRIMREF_GROUP_SIZE;
    DISPATCHDIM_X = num_groups.lo;

    control(wait_idle);

    dispatch_indirect build_primrefs_indirect args(
        Args.global_buffer,
        Args.bvh_buffer,
        Args.flat_instances_buffer,
        Args.rebraid_scratch,
        Args.primref_buffer,
        Args.num_extra_primrefs,
        indirectBuildRangeInfo,
        Args.num_instances);
    control(wait_idle);
}