在设备上获取指向一维连续数组特定元素的指针

时间:2018-06-21 15:30:29

标签: cuda gpgpu cublas

我正在尝试在C ++中使用CUBLAS重写一个python / tensorflow脚本,该脚本正在处理一批输入样本(形状为BxD,B:BatchSize,D:扁平2D矩阵的深度)

第一步,我决定使用CUBLAS cublasSgemmBatched计算一批矩阵的MatMul。

我在link to the question中找到了几个工作示例代码, 但是我想要分配一个大的连续设备阵列来存储一批扁平的相同形状的矩阵。我不想将彼此分开的批次存储在设备内存中(因为它们在给定StackOverflow问题的链接中提供的示例代码中)

从我的想象中,我必须以某种方式获得指向设备内存中每个批处理的起始元素的指针列表。像这样的东西:

float **device_batch_ptr;
cudaMalloc((void**)&device_batch_ptr, batch_size*sizeof(float *));
for(int i = 0 ; i < batch_size; i++ ) {
    // set device_batch_ptr[i] to starting point of i'th batch on device memory array.
}

请注意,cublasSgemmBatched需要一个float **,其中的每个float *都指向给定输入矩阵中每个批次的起始元素。

任何建议将不胜感激。

1 个答案:

答案 0 :(得分:1)

如果阵列位于连续的线性内存(device_array)中,那么您要做的就是使用标准指针算法计算偏移量并将设备地址存储在主机阵列中,然后将其复制到设备中。像这样:

float** device_batch_ptr;
float** h_device_batch_ptr = new float*[batch_size];

cudaMalloc((void**)&device_batch_ptr, batch_size*sizeof(float *));
size_t nelementsperrarray = N * N;
for(int i = 0 ; i < batch_size; i++ ) {
    // set h_device_batch_ptr[i] to starting point of i'th batch on device memory array.
    h_device_batch_ptr[i] = device_array + i * nelementsperarray;
}
cudaMemcpy(device_batch_ptr, h_device_batch_ptr, batch_size*sizeof(float *)),
            cudaMemcpyHostToDevice);

[显然从未编译或测试,使用风险自负]