快速入门

以下代码是一个独立示例，展示了如何使用 cuFFTMp。

#include <cufftMp.h>
#include <mpi.h>
#include <vector>
#include <complex>

int main() {

    // Initialize MPI, pick a device
    MPI_Init(NULL, NULL);
    MPI_Comm comm = MPI_COMM_WORLD;
    
    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    
    int ndevices;
    cudaGetDeviceCount(&ndevices);
    cudaSetDevice(rank % ndevices);
    
    // Allocate CPU memory
    size_t N = 32;
    std::vector<std::complex<float>> cpu_data((N / size) * N * N, {314, 0});
    
    // Create plan, attach to communicator, make plan
    cufftHandle plan = 0;
    size_t workspace;
    cufftCreate(&plan);
    cufftMpAttachComm(plan, CUFFT_COMM_MPI, &comm);
    cufftMakePlan3d(plan, N, N, N, CUFFT_C2C, &workspace);
    
    // Allocate memory, copy CPU data to GPU
    cudaLibXtDesc *desc;
    cufftXtMalloc(plan, &desc, CUFFT_XT_FORMAT_INPLACE);
    cufftXtMemcpy(plan, desc, cpu_data.data(), CUFFT_COPY_HOST_TO_DEVICE);
    
    // Run C2C FFT Forward
    cufftXtExecDescriptor(plan, desc, desc, CUFFT_FORWARD);
    
    // Copy back to CPU
    cufftXtMemcpy(plan, cpu_data.data(), desc, CUFFT_COPY_DEVICE_TO_HOST);
    
    // Data in cpu_data is now distributed along the Y dimension, of size N * (N / size) * N
    // Test by comparing the very first entry on rank 0
    if(rank == 0) {
        if(cpu_data[0].real() == 314 * N * N * N) {
            printf("PASSED\n");
        } else {
            printf("FAILED\n");
        }
    }
    
    // Cleanup
    cufftXtFree(desc);
    cufftDestroy(plan);
    MPI_Finalize();
    
}

注意

更多示例可以在 Github 上找到。

此示例计算跨 size 个 GPU 的 32 x 32 x 32 复数到复数 FFT。它执行以下操作

1. 它初始化 MPI，并为当前进程选择一个设备。

2. 它在每个进程上分配一个 (32 / size) x 32 x 32 CPU 数组。每个进程拥有 (32 / size) x 32 x 32 个元素（即，沿 X 分布的全局数据的 Slab）。

3. 它创建一个 cuFFT plan，使用 cufftMpAttach 启用多进程 API 并创建 plan。

4. 它分配 GPU 内存并将本地 CPU 缓冲区复制到本地 GPU 缓冲区。

5. 它执行 C2C plan。

6. 它将数据从 GPU 复制回 CPU。此时，每个进程拥有 32 x (32 / size) x 32 个元素，即沿 Y 分布的 Slab。

该示例可以像这样轻松地编译和运行

$ nvcc app.cu -I/path/to/cufftMp/include -L/path/to/cufftMp/lib -lcufftMp  -L/path/to/nvshmem/lib -lnvshmem_host -I/path/to/mpi/include -L/path/to/mpi/lib -lmpi -o app
$ mpirun -n 2 ./app
PASSED