首页 > 解决方案 > 获取指向设备上一维连续数组的特定元素的指针

问题描述

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

第一步,我决定使用 CUBLAScublasSgemmBatched为批量矩阵计算 MatMul。

我找到了几个工作示例代码作为问题的链接,但我想要的是分配一个大的连续设备阵列来存储成批的扁平相同形状的矩阵。我不想在设备内存上存储彼此分开的批次(因为它们在给定的 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* 都指向给定输入矩阵中每个批次的起始元素。

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

标签: cudagpgpucublas

解决方案


如果您的数组位于连续的线性内存 ( 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);

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


推荐阅读