/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of NVIDIA CORPORATION nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/* Perform second step of bisection algorithm for large matrices for
 * intervals that contained after the first step more than one eigenvalue
 */

#ifndef _BISECT_KERNEL_LARGE_MULTI_H_
#define _BISECT_KERNEL_LARGE_MULTI_H_

#include <cooperative_groups.h>

namespace cg = cooperative_groups;
// includes, project
#include "config.h"
#include "util.h"

// additional kernel
#include "bisect_util.cu"

////////////////////////////////////////////////////////////////////////////////
//! Perform second step of bisection algorithm for large matrices for
//! intervals that after the first step contained more than one eigenvalue
//! @param  g_d  diagonal elements of symmetric, tridiagonal matrix
//! @param  g_s  superdiagonal elements of symmetric, tridiagonal matrix
//! @param  n    matrix size
//! @param  blocks_mult  start addresses of blocks of intervals that are
//!                      processed by one block of threads, each of the
//!                      intervals contains more than one eigenvalue
//! @param  blocks_mult_sum  total number of eigenvalues / singleton intervals
//!                          in one block of intervals
//! @param  g_left  left limits of intervals
//! @param  g_right  right limits of intervals
//! @param  g_left_count  number of eigenvalues less than left limits
//! @param  g_right_count  number of eigenvalues less than right limits
//! @param  g_lambda  final eigenvalue
//! @param  g_pos  index of eigenvalue (in ascending order)
//! @param  precision  desired precision of eigenvalues
////////////////////////////////////////////////////////////////////////////////
__global__ void bisectKernelLarge_MultIntervals(float             *g_d,
                                                float             *g_s,
                                                const unsigned int n,
                                                unsigned int      *blocks_mult,
                                                unsigned int      *blocks_mult_sum,
                                                float             *g_left,
                                                float             *g_right,
                                                unsigned int      *g_left_count,
                                                unsigned int      *g_right_count,
                                                float             *g_lambda,
                                                unsigned int      *g_pos,
                                                float              precision)
{
    // Handle to thread block group
    cg::thread_block   cta = cg::this_thread_block();
    const unsigned int tid = threadIdx.x;

    // left and right limits of interval
    __shared__ float s_left[2 * MAX_THREADS_BLOCK];
    __shared__ float s_right[2 * MAX_THREADS_BLOCK];

    // number of eigenvalues smaller than interval limits
    __shared__ unsigned int s_left_count[2 * MAX_THREADS_BLOCK];
    __shared__ unsigned int s_right_count[2 * MAX_THREADS_BLOCK];

    // helper array for chunk compaction of second chunk
    __shared__ unsigned int s_compaction_list[2 * MAX_THREADS_BLOCK + 1];
    // compaction list helper for exclusive scan
    unsigned int *s_compaction_list_exc = s_compaction_list + 1;

    // flag if all threads are converged
    __shared__ unsigned int all_threads_converged;
    // number of active threads
    __shared__ unsigned int num_threads_active;
    // number of threads to employ for compaction
    __shared__ unsigned int num_threads_compaction;
    // flag if second chunk has to be compacted
    __shared__ unsigned int compact_second_chunk;

    // parameters of block of intervals processed by this block of threads
    __shared__ unsigned int c_block_start;
    __shared__ unsigned int c_block_end;
    __shared__ unsigned int c_block_offset_output;

    // midpoint of currently active interval of the thread
    float mid = 0.0f;
    // number of eigenvalues smaller than \a mid
    unsigned int mid_count = 0;
    // current interval parameter
    float        left;
    float        right;
    unsigned int left_count;
    unsigned int right_count;
    // helper for compaction, keep track which threads have a second child
    unsigned int is_active_second = 0;

    // initialize common start conditions
    if (0 == tid) {
        c_block_start         = blocks_mult[blockIdx.x];
        c_block_end           = blocks_mult[blockIdx.x + 1];
        c_block_offset_output = blocks_mult_sum[blockIdx.x];

        num_threads_active     = c_block_end - c_block_start;
        s_compaction_list[0]   = 0;
        num_threads_compaction = ceilPow2(num_threads_active);

        all_threads_converged = 1;
        compact_second_chunk  = 0;
    }

    cg::sync(cta);

    // read data into shared memory
    if (tid < num_threads_active) {
        s_left[tid]        = g_left[c_block_start + tid];
        s_right[tid]       = g_right[c_block_start + tid];
        s_left_count[tid]  = g_left_count[c_block_start + tid];
        s_right_count[tid] = g_right_count[c_block_start + tid];
    }

    cg::sync(cta);

    // do until all threads converged
    while (true) {
        // for (int iter=0; iter < 0; iter++) {

        // subdivide interval if currently active and not already converged
        subdivideActiveInterval(tid,
                                s_left,
                                s_right,
                                s_left_count,
                                s_right_count,
                                num_threads_active,
                                left,
                                right,
                                left_count,
                                right_count,
                                mid,
                                all_threads_converged);

        cg::sync(cta);

        // stop if all eigenvalues have been found
        if (1 == all_threads_converged) {
            break;
        }

        // compute number of eigenvalues smaller than mid for active and not
        // converged intervals, use all threads for loading data from gmem and
        // s_left and s_right as scratch space to store the data load from gmem
        // in shared memory
        mid_count = computeNumSmallerEigenvalsLarge(
            g_d, g_s, n, mid, tid, num_threads_active, s_left, s_right, (left == right), cta);

        cg::sync(cta);

        if (tid < num_threads_active) {
            // store intervals
            if (left != right) {
                storeNonEmptyIntervals(tid,
                                       num_threads_active,
                                       s_left,
                                       s_right,
                                       s_left_count,
                                       s_right_count,
                                       left,
                                       mid,
                                       right,
                                       left_count,
                                       mid_count,
                                       right_count,
                                       precision,
                                       compact_second_chunk,
                                       s_compaction_list_exc,
                                       is_active_second);
            }
            else {
                storeIntervalConverged(s_left,
                                       s_right,
                                       s_left_count,
                                       s_right_count,
                                       left,
                                       mid,
                                       right,
                                       left_count,
                                       mid_count,
                                       right_count,
                                       s_compaction_list_exc,
                                       compact_second_chunk,
                                       num_threads_active,
                                       is_active_second);
            }
        }

        cg::sync(cta);

        // compact second chunk of intervals if any of the threads generated
        // two child intervals
        if (1 == compact_second_chunk) {
            createIndicesCompaction(s_compaction_list_exc, num_threads_compaction, cta);

            compactIntervals(s_left,
                             s_right,
                             s_left_count,
                             s_right_count,
                             mid,
                             right,
                             mid_count,
                             right_count,
                             s_compaction_list,
                             num_threads_active,
                             is_active_second);
        }

        cg::sync(cta);

        // update state variables
        if (0 == tid) {
            num_threads_active += s_compaction_list[num_threads_active];
            num_threads_compaction = ceilPow2(num_threads_active);

            compact_second_chunk  = 0;
            all_threads_converged = 1;
        }

        cg::sync(cta);

        // clear
        s_compaction_list_exc[threadIdx.x]              = 0;
        s_compaction_list_exc[threadIdx.x + blockDim.x] = 0;

        cg::sync(cta);

    } // end until all threads converged

    // write data back to global memory
    if (tid < num_threads_active) {
        unsigned int addr = c_block_offset_output + tid;

        g_lambda[addr] = s_left[tid];
        g_pos[addr]    = s_right_count[tid];
    }
}

#endif // #ifndef _BISECT_KERNEL_LARGE_MULTI_H_
