CUDA指针数组Kernel函数

DechinPhy

发布于 2024-03-08 08:18:45

1140

发布于 2024-03-08 08:18:45

技术背景

在前面的一篇文章中，我们介绍了在C++中使用指针数组的方式实现的一个不规则的二维数组。那么如果我们希望可以在CUDA中也能够使用到这种类似形式的不规则的数组，有没有办法可以直接实现呢？可能过程会稍微有一点麻烦，因为我们需要在Host和Device之间来回的转换，需要使用到很多CUDA内置的cudaMalloc和cudaMemcpy函数，以下做一个完整的介绍。

原始代码及修改思路

在上一篇文章中我们使用到的案例代码是这样的：

// g++ main.cpp -o main && ./main
#include <iostream>
 
struct bucket{
    int num;
    int *ptr;
};
 
void print_bucket(bucket *bc, int shape[]){
    for (int i=0; i<4; i++){
        bucket bc_i = bc[i];
        printf("%d: ", bc_i.num);
        for (int j=0; j<shape[i]; j++){
            printf("%d,", bc_i.ptr[j]);
        }
        printf("\n");
    }
}
 
int main(){
    // 定长数组
    int arr[4][3] = {{0,1,2},{1,2,3},{2,3,4},{3,4,5}};
    // 有效长度
    int shape[4] = {2,3,2,1};
    // 先构建结构体数组
	bucket _bc[4];
    for (int i=0; i<4; i++){
        _bc[i].num = shape[i];
        _bc[i].ptr = arr[i];
        _bc[i].ptr += 3-shape[i];
    }
    // 再把结构体数组赋值给结构体指针
    bucket *bc = _bc;
    // 打印结构体的所有内容
    print_bucket(bc, shape);
    return 0;
}

通过定义一个bucket结构体，用双重的指针数组实现了一个不规则数组的存储。第一重的指针对应于不规则数组的第一个维度，这里长度一般是固定的。第二重的指针指向不规则数组的第二个维度，这个维度的长度大小是不一致的，因为我们在结构体中存储的只是一个指针和该维度的数组长度，因此可以实现不规则数组的存储。那么上述代码的运行结果为：

$ g++ main.cpp -o main && ./main
2: 1,2,
3: 1,2,3,
2: 3,4,
1: 5,

打印的第一列是当前数组的长度，也就是不规则数组的第二个维度。后面的数字是对应的数组内容，当然，这里需要注意的点是，我们在初始化的时候，尤其是跟Python等语言进行交互的时候，初始化阶段使用的还是一个固定长度的Tensor，而不需要使用的那些位置需要填充或者叫padding一些数字，常见的就是-1和0。

那么如果我们希望可以在CUDA上实现一个类似的功能，首先需要考虑到以下几个方面：

首先我们需要把数据拷贝到CUDA的Device Memory里面才能用来计算；
Host侧和Device侧指针不能共享，也需要使用Memcpy来进行拷贝；
Kernel函数需要分配一定的计算资源，关于GPU计算资源分配的内容，可以参考之前写的这一篇博客。

CUDA实现

根据以上提到的几个修改点，我们可以这样逐个解决：分别在Host侧定义好相关的数组、指针和结构体之后，使用CUDA的内置函数将相应的内容拷贝到Device侧，两侧同时保留数据，所有的数据更新也都在CUDA上实现。如果有回传数据的需要，我们再把最终的Device侧数据拷贝到Host侧进行同步。完成CUDA的计算之后，同步所有CUDA的线程，并且释放不必要的内存。以下是具体代码实现：

// 文件名：main.cu
// 编译运行指令：nvcc -Xcompiler -fPIC -o main main.cu && ./main
#include <iostream>
#include "cuda_runtime.h"

struct bucket{
    int num;
    int *ptr;
};
// CUDA Kernel函数，该函数主要用于打印bucket结构体的内部数据
__global__ void print_bucket_cuda(bucket *bc, int *shape){
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    if (i < 4){
        bucket bc_i = bc[i];
        for (int j=0; j<shape[i]; j++){
            printf("%d %d\n", i, bc_i.ptr[j]);
        }
    }
}

int main(){
    // 定义Host侧数据
    int arr[4][3] = {{0,1,2},{1,2,3},{2,3,4},{3,4,5}};
    int shape[4] = {2,3,2,1};
    // 先定义Host侧结构体，但是第二重指针在Device侧分配和定义
	bucket _bc[4];
    for (int i=0; i<4; i++){
        _bc[i].num = shape[i];
        cudaMalloc((void**)&(_bc[i].ptr), shape[i]*4);
        cudaMemcpy(_bc[i].ptr, arr[i]+3-shape[i], shape[i]*4, cudaMemcpyHostToDevice);
    }
    // 定义Device侧的结构体
    bucket *d_bc;
    cudaMalloc((void**)&d_bc, sizeof(bucket)*4);  
    int *d_shape; 
    cudaMalloc((void**)&d_shape, sizeof(int)*4);  
    // 将Host侧结构体拷贝到Device侧
    cudaMemcpy(d_bc, _bc, sizeof(bucket)*4, cudaMemcpyHostToDevice);  
    cudaMemcpy(d_shape, shape, sizeof(int)*4, cudaMemcpyHostToDevice); 
    // 运行Kernel打印函数
    print_bucket_cuda<<<4, 1>>>(d_bc, d_shape);
    // CUDA线程同步
    cudaDeviceSynchronize();
    // 释放CUDA显存
    cudaFree(d_bc);  
    cudaFree(d_shape);  
    return 0;
}

在这个实现中，比较重要的一个难点是，我们从Host侧拷贝一个双重指针去Device侧，如果直接拷贝第一重的指针，会出现一个问题是在Device侧无法读取在Host上存储的第二重指针的数据。因此我们在Host侧拷贝数据给Device侧时，我们应该先定义一个Host侧的结构体，但该结构体的第二重指针应该指向Device侧的内存。然后再将第一重的指针拷贝到Device侧，这样才完成了整个结构体的内容拷贝，在Device上才可以识别。该代码的运行结果如下所示：

$ nvcc -Xcompiler -fPIC -o main main.cu && ./main
2 3
3 5
1 1
0 1
2 4
0 2
1 2
1 3

这里是乱序的打印，因为CUDA在计算时几乎是同一时间完成的，因此打印任务也是同时执行的，至于哪一个结果先被输出出来，其实是有一定的随机性的。但是通过对比，我们发现这里输出的数据内容跟前面C++的代码输出内容是一致的。第一列的数据表示第一个维度的索引ID，如果输出是0也就对应上面C++输出的第一行内容。例如这里首位是0的数据，第二列对应元素有1和2，这里就跟C++第一行输出的数组内容对应上了。

总结概要

继上一篇文章学习使用C++存储一个不规则二维数组之后，这里介绍如何在C语言版的CUDA中实现一个不规则的二维数组。总体的实现思路跟前面一篇文章一样，使用了一个二维的指针数组来存储。其中主要的不同点大概就是在Host和Device之间的内存交互上，需要不断的分配、拷贝和释放内存，最终我们还是用一个CUDA的Kernel函数实现了一个不规则数组的输出。