Compiler Explorer

Source code

/*
template <typename T>
__global__ void entry0(T *a, int k) {

extern __shared__ T hen[];
T* top_array = (T*)(hen+(k * blockDim.x));
top_array[0] = a[0];
}

int main(){
    int *a = nullptr;
    entry0<<<1,1, 1024>>>(a, 1);
}
*/

// below is my all code
#include <cuda_runtime.h>
#include <stdio.h>
//#include <kat/on_device/shared_memory.cuh>
#include "device_launch_parameters.h"
#include <vector>
#include <iostream>
#include <math.h>
#include <thrust/limits.h>
#include <cfloat>
//#include "cudaMath.cuh"

#define BLOCK_SIZE 256
#define GRID_SIZE 128

template <typename T>
class cudaMath {
public:
    static T sum(T* in, int count);
    static T topk(T *in, int count, int k);
    static T max(T *in, int count);
    //static T reserve(const std::vector<T> &input);
};

#define CUDA_CHECK_ERROR(call)                                                   \
    do {                                                                         \
        cudaError_t err = call;                                                  \
        if (err != cudaSuccess) {                                                \
            std::cerr << "CUDA error in " << __FILE__ << " at line " << __LINE__ \
                      << ": " << cudaGetErrorString(err) << std::endl;           \
            std::exit(EXIT_FAILURE);                                             \
        }                                                                        \
    } while (0)

template <typename T>
__device__ void insert_value(T *array, int k, T data)
{
    for(int i=0; i<k; i++)
    {
        if(array[i] == data)
        {
            return;
        }
    }
    if(data < array[k-1])
    {
        return;
    }
    for(int i = k-2; i>=0; i--)
    {
        if(data > array[i])
        {
            array[i + 1] = array[i];
        }
        else
        {
            array[i+1] = data;
            return;
        }
    }
    array[0] = data;
}

template <typename T>
__global__ void topk_kernel(T *input, T *output, int length, int k)
{ 
    extern __shared__ unsigned char scheng[];
    T* heng = (T*)scheng;
    T* top_array = (T*)(heng+(k * blockDim.x));
    //auto heng = kat::shared_memory::dynamic::proxy<T>();
    //T* top_array = heng + threadIdx.x * k;

T lowest = (std::is_same<T, int>::value) ? INT_MIN : 
               (std::is_same<T, float>::value) ? -FLT_MAX : 
               (std::is_same<T, double>::value) ? -DBL_MAX : T();

for(int i = 0; i<k; i++)
    {
        top_array[i] = lowest;
    }

for(int idx = threadIdx.x + blockDim.x * blockIdx.x; idx < length; idx += gridDim.x * blockDim.x)
    {
        insert_value(top_array, k, input[idx]);
    }

// 将 top_array 写入共享内存
    for(int i =0; i<k; i++)
    {
        heng[k * threadIdx.x + i] = top_array[i]; 
    }
    __syncthreads();

for(int i = BLOCK_SIZE/2; i>=1; i/=2)
    {
        if(threadIdx.x < i)
        {
            for(int m=0; m<k; m++)
            {
                insert_value(top_array, k, heng[k *(threadIdx.x + i) + m]);
            }
        }
        __syncthreads();
        if(threadIdx.x < i)
        {
            for(int m=0; m<k; m++)
            {
                heng[k * threadIdx.x + m] = top_array[m];
            }
        }
        __syncthreads();
    }
    if(blockIdx.x * blockDim.x < length)
    {
        if(threadIdx.x == 0 )
        {
            for(int i =0; i < k; i++)
            {
                output[k * blockIdx.x + i] = heng[i];
            }
        }
    }
}

template <typename T>
T cudaMath<T>::topk(T *in, int count, int k) 
{
    T* d_in = nullptr;
    T* d_out = nullptr;
    T* _1_pass_result = nullptr;
    T h_out = 0;

CUDA_CHECK_ERROR(cudaMalloc((void**)&d_in, count * sizeof(T)));
    
    CUDA_CHECK_ERROR(cudaMalloc((void**)&d_out, sizeof(T)));

CUDA_CHECK_ERROR(cudaMalloc((void**)&_1_pass_result, (k * GRID_SIZE) *sizeof(T)));
    
    CUDA_CHECK_ERROR(cudaMemcpy(d_in, in, count * sizeof(T), cudaMemcpyHostToDevice));

CUDA_CHECK_ERROR(cudaMemcpy(d_out, &h_out, sizeof(T), cudaMemcpyHostToDevice));

size_t shared_mem_size = k * BLOCK_SIZE * sizeof(T);

cudaFuncSetAttribute(topk_kernel<T>, cudaFuncAttributeMaxDynamicSharedMemorySize, 65536);
    topk_kernel<<<GRID_SIZE, BLOCK_SIZE, shared_mem_size>>>(d_in, _1_pass_result, count, k);
    topk_kernel<<<1, BLOCK_SIZE, shared_mem_size>>>(_1_pass_result, d_out, k*GRID_SIZE, k);

CUDA_CHECK_ERROR(cudaGetLastError());
    CUDA_CHECK_ERROR(cudaDeviceSynchronize());

CUDA_CHECK_ERROR(cudaMemcpy(&h_out, d_out, sizeof(T), cudaMemcpyDeviceToHost));

CUDA_CHECK_ERROR(cudaFree(d_in));
    CUDA_CHECK_ERROR(cudaFree(d_out));
    CUDA_CHECK_ERROR(cudaFree(_1_pass_result));

return h_out;
}

template int cudaMath<int>::topk(int* in, int count, int k);
template float cudaMath<float>::topk(float* in, int count, int k);