/* 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.
 */

////////////////////////////////////////////////////////////////////////////////
//
//  cdpAdvancedQuicksort.cu
//
//  Implementation of a parallel quicksort in CUDA. It comes in
//  several parts:
//
//  1. A small-set insertion sort. We do this on any set with <=32 elements
//  2. A partitioning kernel, which - given a pivot - separates an input
//     array into elements <=pivot, and >pivot. Two quicksorts will then
//     be launched to resolve each of these.
//  3. A quicksort co-ordinator, which figures out what kernels to launch
//     and when.
//
////////////////////////////////////////////////////////////////////////////////
#include <cooperative_groups.h>
#include <thrust/device_vector.h>
#include <thrust/random.h>

namespace cg = cooperative_groups;

#include <helper_cuda.h>
#include <helper_string.h>

#include "cdpQuicksort.h"

////////////////////////////////////////////////////////////////////////////////
// Inline PTX call to return index of highest non-zero bit in a word
////////////////////////////////////////////////////////////////////////////////
static __device__ __forceinline__ unsigned int __qsflo(unsigned int word)
{
    unsigned int ret;
    asm volatile("bfind.u32 %0, %1;" : "=r"(ret) : "r"(word));
    return ret;
}

////////////////////////////////////////////////////////////////////////////////
//
//  ringbufAlloc
//
//  Allocates from a ringbuffer. Allows for not failing when we run out
//  of stack for tracking the offset counts for each sort subsection.
//
//  We use the atomicMax trick to allow out-of-order retirement. If we
//  hit the size limit on the ringbuffer, then we spin-wait for people
//  to complete.
//
////////////////////////////////////////////////////////////////////////////////
template <typename T> static __device__ T *ringbufAlloc(qsortRingbuf *ringbuf)
{
    // Wait for there to be space in the ring buffer. We'll retry only a fixed
    // number of times and then fail, to avoid an out-of-memory deadlock.
    unsigned int loop = 10000;

    while (((ringbuf->head - ringbuf->tail) >= ringbuf->stacksize) && (loop-- > 0))
        ;

    if (loop == 0)
        return NULL;

    // Note that the element includes a little index book-keeping, for freeing
    // later.
    unsigned int index = atomicAdd((unsigned int *)&ringbuf->head, 1);
    T           *ret   = (T *)(ringbuf->stackbase) + (index & (ringbuf->stacksize - 1));
    ret->index         = index;

    return ret;
}

////////////////////////////////////////////////////////////////////////////////
//
//  ringBufFree
//
//  Releases an element from the ring buffer. If every element is released
//  up to and including this one, we can advance the tail to indicate that
//  space is now available.
//
////////////////////////////////////////////////////////////////////////////////
template <typename T> static __device__ void ringbufFree(qsortRingbuf *ringbuf, T *data)
{
    unsigned int index = data->index; // Non-wrapped index to free
    unsigned int count = atomicAdd((unsigned int *)&(ringbuf->count), 1) + 1;
    unsigned int max   = atomicMax((unsigned int *)&(ringbuf->max), index + 1);

    // Update the tail if need be. Note we update "max" to be the new value in
    // ringbuf->max
    if (max < (index + 1))
        max = index + 1;

    if (max == count)
        atomicMax((unsigned int *)&(ringbuf->tail), count);
}

////////////////////////////////////////////////////////////////////////////////
//
//  qsort_warp
//
//  Simplest possible implementation, does a per-warp quicksort with no
//  inter-warp
//  communication. This has a high atomic issue rate, but the rest should
//  actually
//  be fairly quick because of low work per thread.
//
//  A warp finds its section of the data, then writes all data <pivot to one
//  buffer and all data >pivot to the other. Atomics are used to get a unique
//  section of the buffer.
//
//  Obvious optimisation: do multiple chunks per warp, to increase in-flight
//  loads
//  and cover the instruction overhead.
//
////////////////////////////////////////////////////////////////////////////////
__global__ void qsort_warp(unsigned        *indata,
                           unsigned        *outdata,
                           unsigned int     offset,
                           unsigned int     len,
                           qsortAtomicData *atomicData,
                           qsortRingbuf    *atomicDataStack,
                           unsigned int     source_is_indata,
                           unsigned int     depth)
{
    // Handle to thread block group
    cg::thread_block cta = cg::this_thread_block();
    // Find my data offset, based on warp ID
    unsigned int thread_id = threadIdx.x + (blockIdx.x << QSORT_BLOCKSIZE_SHIFT);
    // unsigned int warp_id = threadIdx.x >> 5;   // Used for debug only
    unsigned int lane_id = threadIdx.x & (warpSize - 1);

    // Exit if I'm outside the range of sort to be done
    if (thread_id >= len)
        return;

    //
    // First part of the algorithm. Each warp counts the number of elements that
    // are
    // greater/less than the pivot.
    //
    // When a warp knows its count, it updates an atomic counter.
    //

    // Read in the data and the pivot. Arbitrary pivot selection for now.
    unsigned pivot = indata[offset + len / 2];
    unsigned data  = indata[offset + thread_id];

    // Count how many are <= and how many are > pivot.
    // If all are <= pivot then we adjust the comparison
    // because otherwise the sort will move nothing and
    // we'll iterate forever.
    cg::coalesced_group active  = cg::coalesced_threads();
    unsigned int        greater = (data > pivot);
    unsigned int        gt_mask = active.ballot(greater);

    if (gt_mask == 0) {
        greater = (data >= pivot);
        gt_mask = active.ballot(greater); // Must re-ballot for adjusted comparator
    }

    unsigned int lt_mask  = active.ballot(!greater);
    unsigned int gt_count = __popc(gt_mask);
    unsigned int lt_count = __popc(lt_mask);

    // Atomically adjust the lt_ and gt_offsets by this amount. Only one thread
    // need do this. Share the result using shfl
    unsigned int lt_offset, gt_offset;

    if (lane_id == 0) {
        if (lt_count > 0)
            lt_offset = atomicAdd((unsigned int *)&atomicData->lt_offset, lt_count);

        if (gt_count > 0)
            gt_offset = len - (atomicAdd((unsigned int *)&atomicData->gt_offset, gt_count) + gt_count);
    }

    lt_offset = active.shfl((int)lt_offset, 0); // Everyone pulls the offsets from lane 0
    gt_offset = active.shfl((int)gt_offset, 0);

    // Now compute my own personal offset within this. I need to know how many
    // threads with a lane ID less than mine are going to write to the same buffer
    // as me. We can use popc to implement a single-operation warp scan in this
    // case.
    unsigned lane_mask_lt;
    asm("mov.u32 %0, %%lanemask_lt;" : "=r"(lane_mask_lt));
    unsigned int my_mask   = greater ? gt_mask : lt_mask;
    unsigned int my_offset = __popc(my_mask & lane_mask_lt);

    // Move data.
    my_offset += greater ? gt_offset : lt_offset;
    outdata[offset + my_offset] = data;

    // Count up if we're the last warp in. If so, then Kepler will launch the next
    // set of sorts directly from here.
    if (lane_id == 0) {
        // Count "elements written". If I wrote the last one, then trigger the next
        // qsorts
        unsigned int mycount = lt_count + gt_count;

        if (atomicAdd((unsigned int *)&atomicData->sorted_count, mycount) + mycount == len) {
            // We're the last warp to do any sorting. Therefore it's up to us to
            // launch the next stage.
            unsigned int lt_len = atomicData->lt_offset;
            unsigned int gt_len = atomicData->gt_offset;

            cudaStream_t lstream, rstream;
            cudaStreamCreateWithFlags(&lstream, cudaStreamNonBlocking);
            cudaStreamCreateWithFlags(&rstream, cudaStreamNonBlocking);

            // Begin by freeing our atomicData storage. It's better for the ringbuffer
            // algorithm
            // if we free when we're done, rather than re-using (makes for less
            // fragmentation).
            ringbufFree<qsortAtomicData>(atomicDataStack, atomicData);

            // Exceptional case: if "lt_len" is zero, then all values in the batch
            // are equal. We are then done (may need to copy into correct buffer,
            // though)
            if (lt_len == 0) {
                if (source_is_indata)
                    cudaMemcpyAsync(indata + offset,
                                    outdata + offset,
                                    gt_len * sizeof(unsigned),
                                    cudaMemcpyDeviceToDevice,
                                    lstream);

                return;
            }

            // Start with lower half first
            if (lt_len > BITONICSORT_LEN) {
                // If we've exceeded maximum depth, fall through to backup
                // big_bitonicsort
                if (depth >= QSORT_MAXDEPTH) {
                    // The final bitonic stage sorts in-place in "outdata". We therefore
                    // re-use "indata" as the out-of-range tracking buffer. For (2^n)+1
                    // elements we need (2^(n+1)) bytes of oor buffer. The backup qsort
                    // buffer is at least this large when sizeof(QTYPE) >= 2.
                    big_bitonicsort<<<1, BITONICSORT_LEN, 0, lstream>>>(
                        outdata, source_is_indata ? indata : outdata, indata, offset, lt_len);
                }
                else {
                    // Launch another quicksort. We need to allocate more storage for the
                    // atomic data.
                    if ((atomicData = ringbufAlloc<qsortAtomicData>(atomicDataStack)) == NULL)
                        printf("Stack-allocation error. Failing left child launch.\n");
                    else {
                        atomicData->lt_offset = atomicData->gt_offset = atomicData->sorted_count = 0;
                        unsigned int numblocks =
                            (unsigned int)(lt_len + (QSORT_BLOCKSIZE - 1)) >> QSORT_BLOCKSIZE_SHIFT;
                        qsort_warp<<<numblocks, QSORT_BLOCKSIZE, 0, lstream>>>(
                            outdata, indata, offset, lt_len, atomicData, atomicDataStack, !source_is_indata, depth + 1);
                    }
                }
            }
            else if (lt_len > 1) {
                // Final stage uses a bitonic sort instead. It's important to
                // make sure the final stage ends up in the correct (original) buffer.
                // We launch the smallest power-of-2 number of threads that we can.
                unsigned int bitonic_len = 1 << (__qsflo(lt_len - 1U) + 1);
                bitonicsort<<<1, bitonic_len, 0, lstream>>>(
                    outdata, source_is_indata ? indata : outdata, offset, lt_len);
            }
            // Finally, if we sorted just one single element, we must still make
            // sure that it winds up in the correct place.
            else if (source_is_indata && (lt_len == 1))
                indata[offset] = outdata[offset];

            if (cudaPeekAtLastError() != cudaSuccess)
                printf("Left-side launch fail: %s\n", cudaGetErrorString(cudaGetLastError()));

            // Now the upper half.
            if (gt_len > BITONICSORT_LEN) {
                // If we've exceeded maximum depth, fall through to backup
                // big_bitonicsort
                if (depth >= QSORT_MAXDEPTH)
                    big_bitonicsort<<<1, BITONICSORT_LEN, 0, rstream>>>(
                        outdata, source_is_indata ? indata : outdata, indata, offset + lt_len, gt_len);
                else {
                    // Allocate new atomic storage for this launch
                    if ((atomicData = ringbufAlloc<qsortAtomicData>(atomicDataStack)) == NULL)
                        printf("Stack allocation error! Failing right-side launch.\n");
                    else {
                        atomicData->lt_offset = atomicData->gt_offset = atomicData->sorted_count = 0;
                        unsigned int numblocks =
                            (unsigned int)(gt_len + (QSORT_BLOCKSIZE - 1)) >> QSORT_BLOCKSIZE_SHIFT;
                        qsort_warp<<<numblocks, QSORT_BLOCKSIZE, 0, rstream>>>(outdata,
                                                                               indata,
                                                                               offset + lt_len,
                                                                               gt_len,
                                                                               atomicData,
                                                                               atomicDataStack,
                                                                               !source_is_indata,
                                                                               depth + 1);
                    }
                }
            }
            else if (gt_len > 1) {
                unsigned int bitonic_len = 1 << (__qsflo(gt_len - 1U) + 1);
                bitonicsort<<<1, bitonic_len, 0, rstream>>>(
                    outdata, source_is_indata ? indata : outdata, offset + lt_len, gt_len);
            }
            else if (source_is_indata && (gt_len == 1))
                indata[offset + lt_len] = outdata[offset + lt_len];

            if (cudaPeekAtLastError() != cudaSuccess)
                printf("Right-side launch fail: %s\n", cudaGetErrorString(cudaGetLastError()));
        }
    }
}

////////////////////////////////////////////////////////////////////////////////
//
//  run_quicksort
//
//  Host-side code to run the Kepler version of quicksort. It's pretty
//  simple, because all launch control is handled on the device via CDP.
//
//  All parallel quicksorts require an equal-sized scratch buffer. This
//  must be passed in ahead of time.
//
//  Returns the time elapsed for the sort.
//
////////////////////////////////////////////////////////////////////////////////
float run_quicksort_cdp(unsigned *gpudata, unsigned *scratchdata, unsigned int count, cudaStream_t stream)
{
    unsigned int stacksize = QSORT_STACK_ELEMS;

    // This is the stack, for atomic tracking of each sort's status
    qsortAtomicData *gpustack;
    checkCudaErrors(cudaMalloc((void **)&gpustack, stacksize * sizeof(qsortAtomicData)));
    checkCudaErrors(cudaMemset(gpustack, 0, sizeof(qsortAtomicData))); // Only need set first entry to 0

    // Create the memory ringbuffer used for handling the stack.
    // Initialise everything to where it needs to be.
    qsortRingbuf  buf;
    qsortRingbuf *ringbuf;
    checkCudaErrors(cudaMalloc((void **)&ringbuf, sizeof(qsortRingbuf)));
    buf.head      = 1; // We start with one allocation
    buf.tail      = 0;
    buf.count     = 0;
    buf.max       = 0;
    buf.stacksize = stacksize;
    buf.stackbase = gpustack;
    checkCudaErrors(cudaMemcpy(ringbuf, &buf, sizeof(buf), cudaMemcpyHostToDevice));

    // Timing events...
    cudaEvent_t ev1, ev2;
    checkCudaErrors(cudaEventCreate(&ev1));
    checkCudaErrors(cudaEventCreate(&ev2));
    checkCudaErrors(cudaEventRecord(ev1));

    // Now we trivially launch the qsort kernel
    if (count > BITONICSORT_LEN) {
        unsigned int numblocks = (unsigned int)(count + (QSORT_BLOCKSIZE - 1)) >> QSORT_BLOCKSIZE_SHIFT;
        qsort_warp<<<numblocks, QSORT_BLOCKSIZE, 0, stream>>>(
            gpudata, scratchdata, 0U, count, gpustack, ringbuf, true, 0);
    }
    else {
        bitonicsort<<<1, BITONICSORT_LEN>>>(gpudata, gpudata, 0, count);
    }

    checkCudaErrors(cudaGetLastError());
    checkCudaErrors(cudaEventRecord(ev2));
    checkCudaErrors(cudaDeviceSynchronize());

    float elapse = 0.0f;

    if (cudaPeekAtLastError() != cudaSuccess)
        printf("Launch failure: %s\n", cudaGetErrorString(cudaGetLastError()));
    else
        checkCudaErrors(cudaEventElapsedTime(&elapse, ev1, ev2));

    // Sanity check that the stack allocator is doing the right thing
    checkCudaErrors(cudaMemcpy(&buf, ringbuf, sizeof(*ringbuf), cudaMemcpyDeviceToHost));

    if (count > BITONICSORT_LEN && buf.head != buf.tail) {
        printf("Stack allocation error!\nRingbuf:\n");
        printf("\t head = %u\n", buf.head);
        printf("\t tail = %u\n", buf.tail);
        printf("\tcount = %u\n", buf.count);
        printf("\t  max = %u\n", buf.max);
    }

    // Release our stack data once we're done
    checkCudaErrors(cudaFree(ringbuf));
    checkCudaErrors(cudaFree(gpustack));

    return elapse;
}

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
int run_qsort(unsigned int size, int seed, int debug, int loop, int verbose)
{
    if (seed > 0)
        srand(seed);

    // Create and set up our test
    unsigned *gpudata, *scratchdata;
    checkCudaErrors(cudaMalloc((void **)&gpudata, size * sizeof(unsigned)));
    checkCudaErrors(cudaMalloc((void **)&scratchdata, size * sizeof(unsigned)));

    // Create CPU data.
    unsigned    *data = new unsigned[size];
    unsigned int min  = loop ? loop : size;
    unsigned int max  = size;
    loop              = (loop == 0) ? 1 : loop;

    for (size = min; size <= max; size += loop) {
        if (verbose)
            printf(" Input: ");

        for (unsigned int i = 0; i < size; i++) {
            // Build data 8 bits at a time
            data[i]   = 0;
            char *ptr = (char *)&(data[i]);

            for (unsigned j = 0; j < sizeof(unsigned); j++) {
                // Easy-to-read data in debug mode
                if (debug) {
                    *ptr++ = (char)(rand() % 10);
                    break;
                }

                *ptr++ = (char)(rand() & 255);
            }

            if (verbose) {
                if (i && !(i % 32))
                    printf("\n        ");

                printf("%u ", data[i]);
            }
        }

        if (verbose)
            printf("\n");

        checkCudaErrors(cudaMemcpy(gpudata, data, size * sizeof(unsigned), cudaMemcpyHostToDevice));

        // So we're now populated and ready to go! We size our launch as
        // blocks of up to BLOCKSIZE threads, and appropriate grid size.
        // One thread is launched per element.
        float elapse;
        elapse = run_quicksort_cdp(gpudata, scratchdata, size, NULL);

        // run_bitonicsort<SORTTYPE>(gpudata, scratchdata, size, verbose);
        checkCudaErrors(cudaDeviceSynchronize());

        // Copy back the data and verify correct sort
        checkCudaErrors(cudaMemcpy(data, gpudata, size * sizeof(unsigned), cudaMemcpyDeviceToHost));

        if (verbose) {
            printf("Output: ");

            for (unsigned int i = 0; i < size; i++) {
                if (i && !(i % 32))
                    printf("\n        ");

                printf("%u ", data[i]);
            }

            printf("\n");
        }

        unsigned int check;

        for (check = 1; check < size; check++) {
            if (data[check] < data[check - 1]) {
                printf("FAILED at element: %d\n", check);
                break;
            }
        }

        if (check != size) {
            printf("    cdpAdvancedQuicksort FAILED\n");
            exit(EXIT_FAILURE);
        }
        else
            printf("    cdpAdvancedQuicksort PASSED\n");

        // Display the time between event recordings
        printf("Sorted %u elems in %.3f ms (%.3f Melems/sec)\n", size, elapse, (float)size / (elapse * 1000.0f));
        fflush(stdout);
    }

    // Release everything and we're done
    checkCudaErrors(cudaFree(scratchdata));
    checkCudaErrors(cudaFree(gpudata));
    delete (data);
    return 0;
}

static void usage()
{
    printf("Syntax: cdpAdvancedQuicksort [-size=<num>] [-seed=<num>] [-debug] "
           "[-loop-step=<num>] [-verbose]\n");
    printf("If loop_step is non-zero, will run from 1->array_len in steps of "
           "loop_step\n");
}

// Host side entry
int main(int argc, char *argv[])
{
    int          size    = 1000000;
    unsigned int seed    = 0;
    int          debug   = 0;
    int          loop    = 0;
    int          verbose = 0;

    if (checkCmdLineFlag(argc, (const char **)argv, "help") || checkCmdLineFlag(argc, (const char **)argv, "h")) {
        usage();
        printf("&&&& cdpAdvancedQuicksort WAIVED\n");
        exit(EXIT_WAIVED);
    }

    if (checkCmdLineFlag(argc, (const char **)argv, "size")) {
        size = getCmdLineArgumentInt(argc, (const char **)argv, "size");
    }

    if (checkCmdLineFlag(argc, (const char **)argv, "seed")) {
        seed = getCmdLineArgumentInt(argc, (const char **)argv, "seed");
    }

    if (checkCmdLineFlag(argc, (const char **)argv, "loop-step")) {
        loop = getCmdLineArgumentInt(argc, (const char **)argv, "loop-step");
    }

    if (checkCmdLineFlag(argc, (const char **)argv, "debug")) {
        debug = 1;
    }

    if (checkCmdLineFlag(argc, (const char **)argv, "verbose")) {
        verbose = 1;
    }

    // Get device properties
    int            cuda_device = findCudaDevice(argc, (const char **)argv);
    cudaDeviceProp properties;
    checkCudaErrors(cudaGetDeviceProperties(&properties, cuda_device));
    int cdpCapable = (properties.major == 3 && properties.minor >= 5) || properties.major >= 4;

    printf("GPU device %s has compute capabilities (SM %d.%d)\n", properties.name, properties.major, properties.minor);

    if (!cdpCapable) {
        printf("cdpAdvancedQuicksort requires SM 3.5 or higher to use CUDA Dynamic "
               "Parallelism.  Exiting...\n");
        exit(EXIT_WAIVED);
    }

    printf("Running qsort on %d elements with seed %d, on %s\n", size, seed, properties.name);

    // Increase the pending kernel launch limit
    cudaDeviceSetLimit(cudaLimitDevRuntimePendingLaunchCount, 4096);

    run_qsort(size, seed, debug, loop, verbose);

    exit(EXIT_SUCCESS);
}
