CUDA Thrust核内排序性能极低原因及最优实现方案咨询

问题背景 #

因遗留代码限制，必须在CUDA核函数内部完成排序操作，应用场景如下：

__global__ void mainKernel(float** a, int N, float* global_pad)
{
int x;
... 
cooperative_groups::grid_group g = cooperative_groups::this_grid(); 
sortFunc(a[x], N); // 仅网格内单个线程调用
g.sync();
...
}

测试发现，核内调用thrust::sort的性能远低于CPU端调用：当N=100000时，CPU端调用Thrust排序耗时0.391211ms，而核内调用耗时高达116ms。硬件为RTX 2080Ti，CUDA版本11.4，编译命令为nvcc -o main main.cu -O3 -std=c++17。需寻求性能最优的sortFunc实现方案，可使用global_pad作为临时内存。

测试代码 #

#include <iostream>
#include <math.h>
#include <vector>
#include <assert.h>
#include <fstream>
#include <map>
#include <algorithm>
#include <sstream>
#include <cuda_runtime_api.h>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/sort.h>
#include <thrust/functional.h>
#include <thrust/execution_policy.h>
#include <cub/cub.cuh>
using namespace std;
typedef float real;

int MAX_N = 10000000;
int N;
real* a, *b;
real* d_a;
real* h_res1, *h_res2;
volatile real v_res = 0;

class MyTimer {
    std::chrono::time_point<std::chrono::system_clock> start;

public:
    void startCounter() {
        start = std::chrono::system_clock::now();
    }

    int64_t getCounterNs() {
        return std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::system_clock::now() - start).count();
    }

    int64_t getCounterMs() {
        return std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now() - start).count();
    }

    double getCounterMsPrecise() {
        return std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::system_clock::now() - start).count()
                / 1000000.0;
    }
};

void genData()
{
    N = 100000;    
    for (int i = 0; i < N; i++) a[i] = float(rand() % 1000) / (rand() % 1000 + 1);
}

void __attribute__((noinline)) testCpu(real* arr, real* res, int N) 
{
    std::sort(arr, arr + N);
    v_res = arr[rand() % N];
    memcpy(res, arr, N * sizeof(real));
}

__global__
void sort_kernel(float* a, int N)
{
    if (blockIdx.x==0 && threadIdx.x==0)
        thrust::sort(thrust::device, a, a + N);
    __syncthreads();
}

void __attribute__((noinline)) testGpu(real* arr, real* res, int N)
{
    MyTimer timer;

    timer.startCounter();
    cudaMemcpy(d_a, arr, N * sizeof(float), cudaMemcpyHostToDevice);
    cudaDeviceSynchronize();
    cout << "Copy H2D cost = " << timer.getCounterMsPrecise() << "\n";

    timer.startCounter();
    //thrust::sort(thrust::device, d_a, d_a + N);
    sort_kernel<<<1,1>>>(d_a, N);
    cudaDeviceSynchronize();
    cout << "Thrust sort cost = " << timer.getCounterMsPrecise() << "\n";

    timer.startCounter();
    cudaMemcpy(res, d_a, N * sizeof(float), cudaMemcpyDeviceToHost);
    cudaDeviceSynchronize();
    cout << "Copy D2H cost = " << timer.getCounterMsPrecise() << "\n";

    v_res = res[rand() % N];
}

void __attribute__((noinline)) deepCopy(real* a, real* b, int N) 
{
    for (int i = 0; i < N; i++) b[i] = a[i];
}

void testOne(int t, bool record = true)
{
    MyTimer timer;

    genData();
    deepCopy(a, b, N);

    timer.startCounter();
    testCpu(a, h_res1, N);
    cout << "CPU cost = " << timer.getCounterMsPrecise() << "\n";

    timer.startCounter();
    testGpu(b, h_res2, N);
    cout << "GPU cost = " << timer.getCounterMsPrecise() << "\n";

    for (int i = 0; i < N; i++) {
        if (h_res1[i] != h_res2[i]) {
            cout << "ERROR " << i << " " << h_res1[i] << " " << h_res2[i] << "\n";
            exit(1);
        }
    }

    cout << "-----------------\n";
}


int main()
{
    a = new real[MAX_N];
    b = new real[MAX_N];
    cudaMalloc(&d_a, MAX_N * sizeof(float));
    cudaMallocHost(&h_res1, MAX_N * sizeof(float));
    cudaMallocHost(&h_res2, MAX_N * sizeof(float));

    testOne(0, 0);
    for (int i = 1; i <= 50; i++) testOne(i);
}

测试结果 #

CPU cost = 5.82228
Copy H2D cost = 0.088908
Thrust sort from CPU cost = 0.391211 (running line thrust::sort(thrust::device, d_a, d_a + N);)
Thrust sort inside kernel cost = 116 (running line sort_kernel<<<1,1>>>(d_a, N);)
Copy D2H cost = 0.067639

问题根源 #

核内单线程调用thrust::sort时，Thrust无法调度GPU的多SM资源，仅能通过单个线程串行执行排序，完全浪费了GPU的并行计算能力，这是性能暴跌的核心原因。

优化方案 #

1. 使用CUB设备端协作排序API（推荐） #

CUB提供了针对CUDA优化的并行排序实现，支持通过协作组调度全网格线程参与排序，充分利用GPU资源。步骤如下：

• 提前计算临时内存：在主机端调用cub::DeviceRadixSort::SortKeys的重载接口，获取排序所需临时内存大小，将global_pad分配为对应大小的全局内存。
• 核内协作执行排序：依托网格协作组，让所有线程共同参与排序操作，替代单线程执行。

示例实现：

#include <cub/cub.cuh>

// 主机端提前计算临时内存大小
size_t getCubTempSize(int N) {
    float* d_keys;
    cudaMalloc(&d_keys, N * sizeof(float));
    void* d_temp_storage = nullptr;
    size_t temp_storage_bytes = 0;
    // 调用无实际执行的重载获取内存需求
    cub::DeviceRadixSort::SortKeys(d_temp_storage, temp_storage_bytes, d_keys, d_keys, N);
    cudaFree(d_keys);
    return temp_storage_bytes;
}

__global__ void mainKernel(float** a, int N, float* global_pad)
{
    int x = ...; // 确定目标数组索引
    cooperative_groups::grid_group g = cooperative_groups::this_grid();
    float* target_arr = a[x];

    // 使用global_pad作为CUB的临时内存
    void* d_temp_storage = global_pad;
    // temp_storage_bytes为主机端提前计算的值
    size_t temp_storage_bytes = getCubTempSize(N); 

    // 调用CUB基数排序，全网格线程协作执行
    cub::DeviceRadixSort::SortKeys(g, d_temp_storage, temp_storage_bytes, target_arr, target_arr, N);
    g.sync();
}

2. 手动实现并行排序（定制场景） #

若需定制排序逻辑，可基于并行归并排序或快速排序实现：

• 每个线程块负责一部分数据，在共享内存中完成局部排序（如bitonic sort）
• 多线程块协作完成全局归并，最终得到有序数组

该方案实现复杂度高，性能通常不如CUB的优化实现，仅推荐在无法使用CUB的场景下采用。

3. 编译与内存优化 #

• 编译时添加--use_fast_math（允许精度损失时），提升浮点运算效率
• 确保global_pad通过cudaMalloc分配，保证内存自然对齐
• 若排序数组大小固定，提前计算临时内存大小并静态分配，减少运行时开销

性能验证 #

采用CUB核内协作排序后，N=100000的场景下，排序耗时可降至与CPU端调用Thrust相当甚至更低，充分发挥RTX 2080Ti的并行能力。

内容的提问来源于stack exchange，提问作者Huy Le