CUDA下将主机二维数组按索引分散至GPU的最优实现方案问询

需要将主机端的二维数组按指定索引分散至GPU数组（即arr_d[indices[i]] = arr_h[i]），arr_h规模最大可达10000×1000，所有行均需复制至GPU，寻求最快实现方式。

当前采用的方案是直接在GPU内核中访问主机端数组，内核代码如下：

__global__ void Scatter(
    int* dst, const int* src, const int* indices,
    const int len, const int lenRow) {
  constexpr int warpSize = 32;  
  const int numWarps = (blockDim.x * gridDim.x) >> 5;
  int warpId = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
  int laneId = threadIdx.x & 31;

  while (warpId < len) {  // Each warp copies one row
    int row = indices[warpId];
    // Each thread in the warp reads one column
    for (int col = laneId; col < lenRow; col += warpSize) {
      dst[row * lenRow + col] = src[warpId * lenRow + col];
    }
    warpId += numWarps;
  }
}

由于源数组位于主机端，推测该内核会产生大量主机与设备间的小数据传输，因此疑问：是否应先将整个arr_h复制至设备内存？若此方案更优，如何在设备端并行实现数组分散？

1. 先拷贝整个arr_h到设备内存是必然更优的选择 #

当前内核直接访问主机端的src指针，属于零拷贝内存访问，但这种方式会触发大量零散的PCIe小数据传输。PCIe的带宽优势体现在大块连续数据传输上，零散小数据的传输延迟极高，会严重拖慢整体性能。

将整个arr_h一次性拷贝到设备内存后，后续的分散操作完全在GPU内部进行，内存访问延迟仅为设备内存级别（远低于PCIe传输延迟），整体效率会有数量级的提升。

2. 设备端并行实现数组分散的两种方式 #

方式一：优化后的设备内核（自定义并行逻辑）

复用原内核的warp级行处理逻辑，但将src改为设备内存指针，所有内存访问都在GPU内部完成。可以进一步优化内存访问模式，比如使用连续的向量加载提升带宽利用率。

示例代码：

__global__ void ScatterDevice(
    int* dst, const int* src, const int* indices,
    const int len, const int lenRow) {
  constexpr int warpSize = 32;  
  const int numWarps = (blockDim.x * gridDim.x) >> 5;
  int warpId = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
  int laneId = threadIdx.x & 31;

  while (warpId < len) {  // 每个warp负责处理一行数据
    int src_row_idx = warpId;
    int dst_row_idx = indices[warpId];
    // 每个线程按步长遍历列，保证内存访问连续
    for (int col = laneId; col < lenRow; col += warpSize) {
      const int src_offset = src_row_idx * lenRow + col;
      const int dst_offset = dst_row_idx * lenRow + col;
      dst[dst_offset] = src[src_offset];
    }
    warpId += numWarps;
  }
}

// 主机端调用逻辑
int main() {
  // 假设arr_h是主机端二维数组，len是行数，lenRow是列数
  int len = 10000;
  int lenRow = 1000;
  int* arr_h = new int[len * lenRow];
  // ... 填充arr_h数据 ...

  // 分配设备内存并拷贝数据
  int* d_src;
  int* d_indices;
  int* arr_d; // 假设目标数组已按最大规模分配好内存
  cudaMalloc(&d_src, len * lenRow * sizeof(int));
  cudaMalloc(&d_indices, len * sizeof(int));
  cudaMemcpy(d_src, arr_h, len * lenRow * sizeof(int), cudaMemcpyHostToDevice);
  cudaMemcpy(d_indices, indices, len * sizeof(int), cudaMemcpyHostToDevice);

  // 启动内核：按warp数量计算网格大小，每个block设为256线程（8个warp）
  int block_size = 256;
  int warps_per_block = block_size / 32;
  int grid_size = (len + warps_per_block - 1) / warps_per_block;
  ScatterDevice<<<grid_size, block_size>>>(arr_d, d_src, d_indices, len, lenRow);

  // ... 后续操作 ...
  cudaFree(d_src);
  cudaFree(d_indices);
  delete[] arr_h;
  return 0;
}

方式二：使用Thrust库的scatter接口（简洁高效）

NVIDIA Thrust库提供了高度优化的并行算法，其中thrust::scatter可以直接实现批量数据的分散操作。只需将二维数组展平为一维，生成每个元素的目标索引即可。

示例代码：

#include <thrust/device_vector.h>
#include <thrust/scatter.h>
#include <thrust/sequence.h>

int main() {
  int len = 10000;
  int lenRow = 1000;
  int total_dst_rows = ...; // 目标数组的总行数
  int* arr_h = new int[len * lenRow];
  int* indices = new int[len];
  // ... 填充arr_h和indices数据 ...

  // 转换为设备向量
  thrust::device_vector<int> d_src(arr_h, arr_h + len * lenRow);
  thrust::device_vector<int> d_indices(indices, indices + len);
  thrust::device_vector<int> d_dst(total_dst_rows * lenRow, 0); // 初始化目标数组

  // 生成每个元素的目标索引：每行的目标索引从dst_row*lenRow开始连续递增
  thrust::device_vector<int> dst_indices(len * lenRow);
  for (int i = 0; i < len; ++i) {
    int dst_row = d_indices[i];
    thrust::sequence(dst_indices.begin() + i * lenRow, 
                     dst_indices.begin() + (i+1)*lenRow, 
                     dst_row * lenRow);
  }

  // 执行分散操作
  thrust::scatter(d_src.begin(), d_src.end(), dst_indices.begin(), d_dst.begin());

  // ... 后续操作 ...
  delete[] arr_h;
  delete[] indices;
  return 0;
}

3. 方案选择建议 #

• 如果需要高度自定义并行逻辑，或者无法使用第三方库，选择优化后的设备内核，还可以进一步加入向量加载（如int4）、共享内存缓存等优化手段提升带宽利用率。
• 如果追求代码简洁和最高性能，优先选择Thrust库的scatter接口，其底层实现经过NVIDIA深度优化，能充分利用GPU硬件特性。

内容的提问来源于stack exchange，提问作者Nicolás Tsu