[CUDA] Top_k问题

欢迎访问小站，阅读原文 http://www.yandong.org/archives/759

Top_K问题是从一组数据中选出最大（最小）的K个数据。

问题分析

使用双调排序网络对数据进行排序，然后选择前k个数据。其中面临的问题是BLOCK之间的通信，可以使用global memory和原子操作，或者返回host，进行调度。

为了进行BLOCK之间的通信，选用返回host调度的方式。

函数bitonic_half_cleaner 和bitonic_merge是比较网络，使用CUDA加速。有N个数据，则启用N/2个线程进行比较。然后在host中调度，采用非递归版本的双调排序网络。

代码

#include <stdio.h>#include <stdlib.h>#include <cuda_runtime.h>#include <cuda.h>#include <time.h>#define RAND_MAX 0x7fff#define imin(a,b) (a<b?a:b)#define TYPE intconst int N = 4096;const int threadsPerBlock = 256;const int blocksPerGrid = N / threadsPerBlock ;/*len为区块长度，每个区块分配len/2个线程*/__global__ void bitonic_half_cleaner( TYPE *a, int len) {
	/*每个线程处理两个数据，所以只有N/2个线程，然后要映射到N个数据上，f(tid)=base*len+offset
	*f(tid)即为数组中的下标
	*/
	int index = threadIdx.x + blockIdx.x * blockDim.x;//线程号
	int offset= index%(len/2); /*区块内偏移*/
	int base = index/(len/2); /*所在的区块号*/

	index = base*len +offset;/*线程对应的数据的下标*/
	int temp;
	int target=index + len/2;
	if(a[index] > a[target])
	{
		temp = a[index];
		a[index]=a[target];
		a[target] =temp;
	}}/*len为区块长度，每个区块分配len/2个线程*/__global__ void bitonic_merge( TYPE *a, int len) {
	/*每个线程处理两个数据，所以只有N/2个线程，然后要映射到N个数据上，f(tid)=base*len+offset
	*f(tid)即为数组中的下标
	*/
	int index = threadIdx.x + blockIdx.x * blockDim.x;//线程号
	int offset= index%(len/2); /*区块内偏移*/
	int base = index/(len/2); /*所在的区块号*/

	int source = base*len +offset;/*线程对应的数据的下标*/
	int temp;
	int target=base*len+len-offset-1;
	if(a[source] >  a[target])
	{
		temp = a[source];
		a[source]=a[target];
		a[target] =temp;
	}}int main( void ) {
    TYPE   *a;
    TYPE   *dev_a;

    // allocate memory on the cpu side
    a = (TYPE*)malloc( N*sizeof(TYPE) );
	srand(0);

    // allocate the memory on the GPU
    cudaMalloc( (void**)&dev_a,  N*sizeof(TYPE) ) ;

    // fill in the host memory with data
    for (int i=0; i<N; i++) {
		a[i]=(rand()+i/100);
		printf("%d ",a[i]);
    }
	printf("\n\n\n %d\n", 3%1);

    // copy the arrays 'a' and 'b' to the GPU
    cudaMemcpy( dev_a, a, N*sizeof(TYPE),cudaMemcpyHostToDevice ) ;
	int BLOCKS=(N/2)/threadsPerBlock;
	printf("B %d T %d\n", BLOCKS, threadsPerBlock);

	/*双调排序，非递归，使用CUDA加速比较网络部分*/
	for(int i=2;i<=N;i=i<<1)
	{
		printf("Merge i=%d\n",i);
		bitonic_merge<<<BLOCKS,threadsPerBlock>>>(dev_a, i);
		for(int j=(i>>1);j>=2;j=j>>1)
		{
			printf("i=%d, j=%d\n",i ,j);
			bitonic_half_cleaner<<<BLOCKS,threadsPerBlock>>>(dev_a, j);
		}
	}
	

    // copy the array 'c' back from the GPU to the CPU
   cudaMemcpy( a, dev_a, N*sizeof(TYPE), cudaMemcpyDeviceToHost ) ;

	for (int i=0; i<N; i++) {
		printf("%d ",a[i]);
    }
	

    // free memory on the gpu side
    cudaFree( dev_a ) ;

    // free memory on the cpu side
    free( a );}