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

Question

我正在尝试在 C++ 中使用 CUBLAS 重写一个 python/tensorflow 脚本，该脚本对批量输入样本（形状为 BxD，B：BatchSize，D：扁平二维矩阵的深度）进行操作

第一步，我决定使用 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* 都指向给定输入矩阵中每个批次的起始元素。

任何意见和建议将不胜感激。

Answer 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);

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