如何将thrust::device_vector数组传递给CUDA核函数？

这个问题我之前也碰到过，其实核心是要搞清楚thrust::device_vector数组的内存分布：你定义的thrust::device_vector<int> A[N]是主机端的数组，每个元素是device_vector对象（这些对象本身在主机内存），但它们管理的实际数据是在设备内存里的。核函数跑在设备上，没法直接访问主机端的A数组，所以得换个思路——把每个device_vector的底层设备指针提取出来，拼成一个设备端的指针数组，再传给核函数。

下面是具体的实现步骤和代码示例：

1. 准备主机端的device_vector数组 #

先初始化你的device_vector数组，比如给每个vector填充一些数据：

#define N 5 // 假设数组大小为5

int main() {
    // 主机端的device_vector数组
    thrust::device_vector<int> A[N];
    
    // 初始化每个device_vector
    for (int i = 0; i < N; ++i) {
        A[i].resize(10); // 每个vector存10个int
        thrust::fill(A[i].begin(), A[i].end(), i); // 填充值为i
    }

2. 提取底层设备指针并拷贝到设备端 #

我们需要把每个device_vector的底层设备指针收集到一个主机端数组，再把这个数组拷贝到设备内存中：

// 主机端指针数组，存储每个device_vector的底层设备指针
    int* h_ptr_array[N];
    for (int i = 0; i < N; ++i) {
        h_ptr_array[i] = thrust::raw_pointer_cast(A[i].data());
    }

    // 分配设备内存存储指针数组
    int** d_ptr_array;
    cudaMalloc(&d_ptr_array, N * sizeof(int*));
    // 把主机端的指针数组拷贝到设备端
    cudaMemcpy(d_ptr_array, h_ptr_array, N * sizeof(int*), cudaMemcpyHostToDevice);

3. 修改核函数并调用 #

核函数需要接收设备端的指针数组，以及每个vector的大小（如果需要访问元素的话）：

// 调用核函数
    kernel<<<1, N>>>(d_ptr_array, 10);
    cudaDeviceSynchronize(); // 等待核函数执行完成

    // 释放设备内存
    cudaFree(d_ptr_array);
    return 0;
}

// 核函数定义
__global__ void kernel(int** ptr_array, int vec_size) {
    int thread_idx = threadIdx.x;
    if (thread_idx >= N) return; // 防止越界

    // 获取对应device_vector的底层设备指针
    int* vec_data = ptr_array[thread_idx];
    // 现在可以像操作普通设备指针一样访问元素了
    printf("Thread %d: 对应vector的第一个元素是 %d\n", thread_idx, vec_data[0]);
}

关键注意点 #

• 为什么不能直接传递主机端的A数组？因为A是主机内存中的数组，核函数运行在设备上无法直接访问主机内存的对象，而且device_vector的内部结构是主机端的封装，设备端无法解析。
• 如果你用的是动态大小的device_vector数组（比如std::vector<thrust::device_vector<int>>），方法是完全一样的：遍历每个元素提取raw_pointer_cast后的指针，再拷贝到设备端的指针数组即可。

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