code: def _get_sorting_kernel(kernel_id, block_size):
      """
      Builds kernels used for sorting inputs. There are several kernels here
      corresponding to the steps in the algorithm. The algorithm works by
      determining the sorted position for each input item. This is done with
      a bucket sort algorithm, where each word_id is a bucket. The first step
      determines the size of each bucket (number of occurences of each word_id).
      Next, a prefix some is computed over the list of bucket sizes to find
      where each bucket will be placed in the output buffer. Finally, each thread
      places it's index into the correct sorted position based on the bucket
      start index (computed from the prefix sum) and that thread's offset into
      the bucket (which is taken from the output of the atomic add done in the
      first step.)
  
      Arguments:
          kernel_id (Integer): Which step to build the kernel for [0, 4]
          block_size (Integer): Number of threads per block for the prefix sum
              kernels.
      """
      code = r"""
  #define THREADS %(threads)s
  #define STORE_BLOCKSUM %(store_blocksum)s
  __global__ void sort_inputs0(
          int* inputs, int* index_buffer, int* offset_buffer, int* word_counts, const int vocab_size,
          const int input_length)
  {
      const int tid = threadIdx.x + (blockDim.x * blockIdx.x);
      int word_id;
  
      if(tid < input_length)
      {
          word_id = inputs[tid];
          offset_buffer[tid] = atomicAdd(&word_counts[word_id], 1);
      }
  }
  
  __device__ void scan(int* buffer, int* blocksum, int global_length)
  {
      const int tid = (threadIdx.x << 1) + 1;
      const int gid = ((threadIdx.x + (blockIdx.x * blockDim.x)) << 1) + 1;
  
      __shared__ int local_counts[THREADS * 2];
      local_counts[tid] = buffer[gid];
      local_counts[tid - 1] = buffer[gid - 1];
  
      #pragma unroll
      for(int skip = 1; skip <= THREADS; skip <<= 1)
      {
          int mask = (skip << 1) - 1;
          if((tid & mask) == mask)
          {
              local_counts[tid] += local_counts[tid - skip];
          }
  
          __syncthreads();
      }
  
      if(tid == (THREADS * 2 - 1))
      {
  #if STORE_BLOCKSUM
          blocksum[blockIdx.x] = local_counts[tid];
  #endif
          local_counts[tid] = 0;
      }
  
      #pragma unroll
      for(int skip = THREADS; skip > 0; skip >>= 1)
      {
          int mask = (skip << 1) - 1;
          if((tid & mask) == mask)
          {
              int temp = local_counts[tid - skip];
              local_counts[tid - skip] = local_counts[tid];
              local_counts[tid] += temp;
          }
  
          __syncthreads();
      }
  
      if(gid < global_length)
      {
          buffer[gid] = local_counts[tid];
          buffer[gid - 1] = local_counts[tid - 1];
      }
  }
  
  __global__ void sort_inputs1(
          int* inputs, int* index_buffer, int* offset_buffer, int* word_counts, const int vocab_size,
          const int input_length)
  {
      scan(word_counts, word_counts + vocab_size, vocab_size);
  }
  
  __global__ void sort_inputs2(
          int* inputs, int* index_buffer, int* offset_buffer, int* word_counts, const int vocab_size,
          const int input_length)
  {
      scan(word_counts + vocab_size, 0, blockDim.x);
  }
  
  __global__ void sort_inputs3(
          int* inputs, int* index_buffer, int* offset_buffer, int* word_counts, const int vocab_size,
          const int input_length)
  {
      const int gid = (threadIdx.x + (blockIdx.x * blockDim.x)) << 1;
  
      if(gid < vocab_size)
      {
          word_counts[gid] += word_counts[vocab_size + blockIdx.x];
          word_counts[gid + 1] += word_counts[vocab_size + blockIdx.x];
      }
  }
  
  __global__ void sort_inputs4(
          int* inputs, int* index_buffer, int* offset_buffer, int* word_counts, const int vocab_size,
          const int input_length)
  {
      const int tid = threadIdx.x + (blockDim.x * blockIdx.x);
      int word_id;
  
      if(tid < input_length)
      {
          word_id = inputs[tid];
          int sorted_position = word_counts[word_id] + offset_buffer[tid];
          index_buffer[sorted_position] = tid;
      }
  }
  """
      code = code % {
          "threads": block_size,
          "store_blocksum": (1 if kernel_id == 1 else 0),
      }
      module = SourceModule(code, options=["--use_fast_math"])
  
      function_name = "sort_inputs" + native_str(kernel_id)
      kernel = module.get_function(function_name)
      kernel.prepare("PPPPII")
      kernel.name = "sort_inputs"
      return kernel
  
imports: from future.utils import native_str
  from pycuda.tools import context_dependent_memoize
  from neon.backends.cuda_templates import _ew_types
  from neon.backends.util.source_module import SourceModule
label_desc: Use image node_id 2 for calling a global function with example usage: _get_sorting_kernel(kernel_id, block_size) and returns: kernel