/* ========================================================================== *
 *  The following code is an adaptation of the bitonic sort implementation
 *  found in the CUDA Toolkit Samples.
 *
 *  Author: Firas Safadi (fsafadi@ulg.ac.be)
 *  Original code by NVIDIA
 * ========================================================================== */

#include "sort.h"

////////////////////////////////////////////////////////////////////////////////
// Shortcut definition
////////////////////////////////////////////////////////////////////////////////
typedef unsigned int uint;

////////////////////////////////////////////////////////////////////////////////
// CUDA sorting networks
////////////////////////////////////////////////////////////////////////////////

extern "C" uint bitonicSort(
  uint *d_SrcKey,
  real *d_SrcVal,
  uint batchSize,
  uint arrayLength,
  uint dir
  );

//Enables maximum occupancy
#define SHARED_SIZE_LIMIT 1024U
#define GLOBAL_MERGE_WIDTH (SHARED_SIZE_LIMIT / 2)

//Map to single instructions on G8x / G9x / G100
#define    UMUL(a, b) __umul24((a), (b))
#define UMAD(a, b, c) ( UMUL((a), (b)) + (c) )

__device__ inline void Comparator(
  uint &keyA,
  real &valA,
  uint &keyB,
  real &valB,
  uint dir
  )
{
  uint t;
  real u;

  if ((valA > valB) == dir)
  {
    t = keyA;
    keyA = keyB;
    keyB = t;
    u = valA;
    valA = valB;
    valB = u;
  }
}

////////////////////////////////////////////////////////////////////////////////
// Monolithic bitonic sort kernel for short arrays fitting into shared memory
////////////////////////////////////////////////////////////////////////////////
__global__ void bitonicSortShared(
  uint *d_SrcKey,
  real *d_SrcVal,
  uint arrayLength,
  uint dir
  )
{
  //Shared memory storage for one or more short vectors
  __shared__ uint s_key[SHARED_SIZE_LIMIT];
  __shared__ real s_val[SHARED_SIZE_LIMIT];

  //Offset to the beginning of subbatch and load data
  d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;
  d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;
  s_key[threadIdx.x + 0] = d_SrcKey[0];
  s_val[threadIdx.x + 0] = d_SrcVal[0];
  s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];
  s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];

  for (uint size = 2; size < arrayLength; size <<= 1)
  {
    //Bitonic merge
    uint ddd = dir ^ ((threadIdx.x & (size / 2)) != 0);

    for (uint stride = size / 2; stride > 0; stride >>= 1)
    {
      __syncthreads();
      uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));
      Comparator(
        s_key[pos + 0], s_val[pos + 0],
        s_key[pos + stride], s_val[pos + stride],
        ddd
        );
    }
  }

  //ddd == dir for the last bitonic merge step
  {
    for (uint stride = arrayLength / 2; stride > 0; stride >>= 1)
    {
      __syncthreads();
      uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));
      Comparator(
        s_key[pos + 0], s_val[pos + 0],
        s_key[pos + stride], s_val[pos + stride],
        dir
        );
    }
  }

  __syncthreads();
  d_SrcKey[0] = s_key[threadIdx.x + 0];
  d_SrcVal[0] = s_val[threadIdx.x + 0];
  d_SrcKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];
  d_SrcVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];
}

////////////////////////////////////////////////////////////////////////////////
// Bitonic sort kernel for large arrays (not fitting into shared memory)
////////////////////////////////////////////////////////////////////////////////
//Bottom-level bitonic sort
//Almost the same as bitonicSortShared with the exception of
//even / odd subarrays being sorted in opposite directions
//Bitonic merge accepts both
//Ascending | descending or descending | ascending sorted pairs
__global__ void bitonicSortShared1(
  uint *d_SrcKey,
  real *d_SrcVal
  )
{
  //Shared memory storage for current subarray
  __shared__ uint s_key[SHARED_SIZE_LIMIT];
  __shared__ real s_val[SHARED_SIZE_LIMIT];

  //Offset to the beginning of subarray and load data
  d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;
  d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;
  s_key[threadIdx.x + 0] = d_SrcKey[0];
  s_val[threadIdx.x + 0] = d_SrcVal[0];
  s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];
  s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];

  for (uint size = 2; size < SHARED_SIZE_LIMIT; size <<= 1)
  {
    //Bitonic merge
    uint ddd = (threadIdx.x & (size / 2)) != 0;

    for (uint stride = size / 2; stride > 0; stride >>= 1)
    {
      __syncthreads();
      uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));
      Comparator(
        s_key[pos + 0], s_val[pos + 0],
        s_key[pos + stride], s_val[pos + stride],
        ddd
        );
    }
  }

  //Odd / even arrays of SHARED_SIZE_LIMIT elements
  //sorted in opposite directions
  uint ddd = blockIdx.x & 1;
  {
    for (uint stride = SHARED_SIZE_LIMIT / 2; stride > 0; stride >>= 1)
    {
      __syncthreads();
      uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));
      Comparator(
        s_key[pos + 0], s_val[pos + 0],
        s_key[pos + stride], s_val[pos + stride],
        ddd
        );
    }
  }


  __syncthreads();
  d_SrcKey[0] = s_key[threadIdx.x + 0];
  d_SrcVal[0] = s_val[threadIdx.x + 0];
  d_SrcKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];
  d_SrcVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];
}

//Bitonic merge iteration for stride >= SHARED_SIZE_LIMIT
__global__ void bitonicMergeGlobal(
  uint *d_SrcKey,
  real *d_SrcVal,
  uint arrayLength,
  uint size,
  uint stride,
  uint dir
  )
{
  uint global_comparatorI = blockIdx.x * blockDim.x + threadIdx.x;
  uint        comparatorI = global_comparatorI & (arrayLength / 2 - 1);

  //Bitonic merge
  uint ddd = dir ^ ((comparatorI & (size / 2)) != 0);
  uint pos = 2 * global_comparatorI - (global_comparatorI & (stride - 1));

  uint keyA = d_SrcKey[pos + 0];
  real valA = d_SrcVal[pos + 0];
  uint keyB = d_SrcKey[pos + stride];
  real valB = d_SrcVal[pos + stride];

  Comparator(
    keyA, valA,
    keyB, valB,
    ddd
    );

  d_SrcKey[pos + 0] = keyA;
  d_SrcVal[pos + 0] = valA;
  d_SrcKey[pos + stride] = keyB;
  d_SrcVal[pos + stride] = valB;
}

//Combined bitonic merge steps for
//size > SHARED_SIZE_LIMIT and stride = [1 .. SHARED_SIZE_LIMIT / 2]
__global__ void bitonicMergeShared(
  uint *d_SrcKey,
  real *d_SrcVal,
  uint arrayLength,
  uint size,
  uint dir
  )
{
  //Shared memory storage for current subarray
  __shared__ uint s_key[SHARED_SIZE_LIMIT];
  __shared__ real s_val[SHARED_SIZE_LIMIT];

  d_SrcKey += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;
  d_SrcVal += blockIdx.x * SHARED_SIZE_LIMIT + threadIdx.x;
  s_key[threadIdx.x + 0] = d_SrcKey[0];
  s_val[threadIdx.x + 0] = d_SrcVal[0];
  s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcKey[(SHARED_SIZE_LIMIT / 2)];
  s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)] = d_SrcVal[(SHARED_SIZE_LIMIT / 2)];

  //Bitonic merge
  uint comparatorI = UMAD(blockIdx.x, blockDim.x, threadIdx.x) & ((arrayLength / 2) - 1);
  uint ddd = dir ^ ((comparatorI & (size / 2)) != 0);

  for (uint stride = SHARED_SIZE_LIMIT / 2; stride > 0; stride >>= 1)
  {
    __syncthreads();
    uint pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));
    Comparator(
      s_key[pos + 0], s_val[pos + 0],
      s_key[pos + stride], s_val[pos + stride],
      ddd
      );
  }

  __syncthreads();
  d_SrcKey[0] = s_key[threadIdx.x + 0];
  d_SrcVal[0] = s_val[threadIdx.x + 0];
  d_SrcKey[(SHARED_SIZE_LIMIT / 2)] = s_key[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];
  d_SrcVal[(SHARED_SIZE_LIMIT / 2)] = s_val[threadIdx.x + (SHARED_SIZE_LIMIT / 2)];
}

////////////////////////////////////////////////////////////////////////////////
// Interface function
////////////////////////////////////////////////////////////////////////////////
//Helper function (also used by odd-even merge sort)
extern "C" uint factorRadix2(uint *log2L, uint L)
{
  if (!L)
  {
    *log2L = 0;
    return 0;
  }
  else
  {
    for (*log2L = 0; (L & 1) == 0; L >>= 1, *log2L++);

    return L;
  }
}

extern "C" uint bitonicSort(
  uint *d_SrcKey,
  real *d_SrcVal,
  uint batchSize,
  uint arrayLength,
  uint dir
  )
{
  //Nothing to sort
  if (arrayLength < 2)
    return 0;

  //Only power-of-two array lengths are supported by this implementation
  uint log2L;
  uint factorizationRemainder = factorRadix2(&log2L, arrayLength);

  if (factorizationRemainder != 1)
    error("Array length must be a power of 2");

  dir = (dir != 0);

  uint  blockCount = batchSize * arrayLength / SHARED_SIZE_LIMIT;
  uint threadCount = SHARED_SIZE_LIMIT / 2;

  if (arrayLength <= SHARED_SIZE_LIMIT)
  {
    if ((batchSize * arrayLength) % SHARED_SIZE_LIMIT != 0)
      error("Batch size must be greater than or equal to shared memory size limit");

    AsyncKernelLaunch(bitonicSortShared, blockCount, threadCount, d_SrcKey, d_SrcVal, arrayLength, dir);
  }
  else
  {
    AsyncKernelLaunch(bitonicSortShared1, blockCount, threadCount, d_SrcKey, d_SrcVal);

    for (uint size = 2 * SHARED_SIZE_LIMIT; size <= arrayLength; size <<= 1)
      for (unsigned stride = size / 2; stride > 0; stride >>= 1)
        if (stride >= SHARED_SIZE_LIMIT)
        {
          AsyncKernelLaunch(bitonicMergeGlobal, (batchSize * arrayLength) / GLOBAL_MERGE_WIDTH, GLOBAL_MERGE_WIDTH / 2, d_SrcKey, d_SrcVal, arrayLength, size, stride, dir);
        }
        else
        {
          AsyncKernelLaunch(bitonicMergeShared, blockCount, threadCount, d_SrcKey, d_SrcVal, arrayLength, size, dir);
          break;
        }
  }

  return threadCount;
}

void bitonic_sort(real* val, nat* key, nat batchSize, nat arrayLength)
{
  bitonicSort(key, val, batchSize, arrayLength, 0);
}
