CUDA——结合MPI多GPU实现

在最近的课题中需要计算颗粒速度矢量的最大值，并且这个程序是MPI+GPU并行的，因此初步想法是分别找出各个CPU进程对应的局部最大值，再通过MPI通讯找出全局最大值。
因此这个问题暂时转化成从一个数组中找出它的最大值的问题。

实现方案

Naive

由于规约算法一般会在输入数组的上面直接修改数据内容，因此对于只读类型的输入数据，需要创建一个同样大小的数组来将要处理的数据拷贝过来，然后再进行操作。
由于我这里的算法是首先计算出每一个block处理的数据的局部最大值，然后交给主机代码找出全局最大值。那么这个临时的拷贝数组就可以采用SMEM来完成。

__global__ void max_kernel(int *arr, int *max_arr, int nElem)
{
    __shared__ int arrT[BLOCKDIM_X];                        //- 申请静态共享内存，缓解高延时GMEM的访问
    int gid = blockIdx.x * blockDim.x + threadIdx.x;        //- 全局索引
    if(gid < nElem)
    {
        arrT[threadIdx.x] = arr[gid];                       //- 将GMEM中的数据转移到SMEM中
    }
    __syncthreads();                                        //- 但凡涉及到SMEM的存储，基本上需要块内同步
    for (int tid = blockDim.x / 2; tid > 0; tid = tid / 2)  //- 归并循环
    {
        if(threadIdx.x < tid)                               //- 前半个区域执行操作，避免warp发散
        {
            int lidLeft = threadIdx.x;
            int lidRight = lidLeft + tid;
            arrT[lidLeft] = arrT[lidLeft]>arrT[lidRight]?arrT[lidLeft]:arrT[lidRight];
        }
        __syncthreads();                                    //- 必须块内同步
    }
    if(threadIdx.x==0)                                      //- 当前块的局部最大值会汇总到首元素，并存到对应的数组中
        max_arr[blockIdx.x] = arrT[0];
}

可以看到，这里申请的是静态共享内存。根据计算能力限制，K80每个block限制的SMEM上限是48KB，本kernel每个block申请的是4KB，存在很大的空间。

MPI跨节点多GPU并行实现

单个GPU由于显存限制，能够处理的数组大小受限，并且大数组在主机和设备段的传输比较慢，使用多个GPU能够处理更大规模的数组，并且通过MPI能够实现多个节点之间合作。
基本思路：

1. 每个进程独立创建并初始化原始数组，并在对应的GPU上进行运行，这里假设每个进程处理的数据是1GB元素；
2. 每个进程遍历得出对应进程的局部最大值；
3. 由根进程汇总各个进程的局部最大值，并遍历得出全局最大值。

//- nvcc -arch sm_37 -I/opt/openmpi-4.1.2/include -L/opt/openmpi-4.1.2/lib --compiler-options -pthread -lmpi arrayMax2.cu -o test
#include "mpi.h"
#include "stdlib.h"
#include "stdio.h"
#define BLOCKDIM_X 1024

//- init_arr函数和max_kernel函数保持一致

int main(int argc, char **argv)
{
    const int GPU_ID[12] = {0,1,2,3,4,5,6,7,8,9,10,11};

    //- MPI初始化
    int cpuID, cpuNum;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &cpuID);
    MPI_Comm_size(MPI_COMM_WORLD, &cpuNum);
    
    //- 各进程独立任务
    cudaSetDevice(GPU_ID[cpuID]);
    int nElem_total = 1<<30;                       //- 数组总大小
    int nElem_per_cpu = nElem_total / cpuNum;      //- 每个进程需要处理的大小
    int *h_arr, *d_arr=NULL;
    h_arr = (int *)malloc(nElem_per_cpu * sizeof(int));
    cudaMalloc((void **)&d_arr, nElem_per_cpu * sizeof(int));
    init_arr(h_arr, nElem_per_cpu, cpuID);

    //- 创建线程组织参数
    dim3 block(BLOCKDIM_X);
    dim3 grid((nElem_per_cpu + block.x -1) / block.x);
    printf("CPU %d process <<<%d, %d>>>\n", cpuID, grid.x, block.x);
    cudaMemcpy(d_arr, h_arr, nElem_per_cpu * sizeof(int), cudaMemcpyHostToDevice);

    //- 创建block级最大值数组
    int *h_max_arr, *d_max_arr=NULL;
    h_max_arr = (int *)malloc(grid.x * sizeof(int));
    cudaMalloc((void**)&d_max_arr, grid.x * sizeof(int));

    //- 调用核函数
    max_kernel<<<grid, block>>>(d_arr, d_max_arr, nElem_per_cpu);
    cudaDeviceSynchronize();

    //- 拷贝结果到主机，计算出各进程的局部最大值
    cudaMemcpy(h_max_arr, d_max_arr, grid.x * sizeof(int), cudaMemcpyDeviceToHost);
    int max_cpu = 0;
    for(int i = 0; i < grid.x; i++)
    {
        max_cpu = max_cpu<h_max_arr[i]?h_max_arr[i]:max_cpu;
    }

    //- 通过MPI收集局部最大值
    int *max_cpu_arr=NULL;
    if(cpuID==0)
    {
        max_cpu_arr = new int[cpuNum * sizeof(int)];
    }
    MPI_Gather(&max_cpu, 1, MPI_INT, max_cpu_arr, 1, MPI_INT, 0, MPI_COMM_WORLD);

    //- 根进程计算全局最大值
    if(cpuID == 0)
    {
        int max_total = 0;
        for(int i = 0; i < cpuNum; i++)
        {
            max_total = max_total>max_cpu_arr[i]?max_total:max_cpu_arr[i];
        }
        printf("the max is %d\n", max_total);
    }

    MPI_Finalize();

    return 0;
}