如何在多GPU上执行cufftXt与CUDA核函数且避免重复分配内存？

当然有办法避免重复分配设备内存！核心思路是利用cufftXt的分布式内存管理机制，让你的CUDA核函数直接操作cufftXt管理的设备内存块，不需要为核函数单独分配内存。下面我会一步步拆解方案，并给出修改后的完整代码示例。

核心原理 #

cufftXt通过分布式数据描述符来管理跨GPU的内存，你可以通过官方API直接获取每个GPU上对应的内存指针，这样核函数就能直接访问这些内存，和cufftXt共享同一份设备内存空间。

具体实现步骤 #

1. 初始化cufftXt多GPU环境 #

首先创建cufft计划，并指定要使用的GPU列表。这一步会让cufftXt知道要在哪些设备上分配内存和执行运算。

2. 分配分布式内存 #

用cufftXtMalloc替代手动的cudaMalloc，为输入和输出数据分配跨GPU的分布式内存。这个函数会自动在各个指定GPU上分配对应大小的内存块。

3. 获取每个GPU上的内存指针 #

通过cufftXtGetSubarray函数，你可以获取到每个GPU上对应的设备内存指针和数据尺寸。这些指针可以直接传给你的CUDA核函数使用。

4. 执行核函数与FFT运算 #

切换到对应的GPU设备，用获取到的指针启动核函数处理数据；完成后直接调用cufftXtExec执行FFT，数据已经在cufftXt的分布式内存中，无需额外拷贝。

修改后的完整代码示例 #

#include <iostream>
#define _USE_MATH_DEFINES
#include <math.h>
#include <ctime>
#include <fstream>
#include <sstream>
#include <cstdlib>
#include <string>
#include <stdlib.h>
#include <stdio.h>
#include <cuda_runtime.h>
#include <cufft.h>
#include <cufftXt.h>
using namespace std;

#define gpuErrchk(ans) { gpuAssert((ans), __FILE__, __LINE__); }
inline void gpuAssert(cudaError_t code, const char *file, int line, bool abort=true) {
    if (code != cudaSuccess) {
        fprintf(stderr,"GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line);
        if (abort) exit(code);
    }
}

__global__ void Cube (cufftReal *data, cufftReal *data3, int N, int real_size) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < real_size){
        float x = (i % (N+2));
        if(x < N){
            data3[i] = pow(data[i], 3.0f);
        } else{
            data3[i] = 0.0f;
        }
    }
    __syncthreads();
}

int main (int argc, char **argv) {
    int x;
    int N = 8;
    int cplx_size = N * (N/2 + 1);
    int real_size = 2 * cplx_size;
    size_t mem_size = sizeof(cufftReal) * real_size;

    // 主机数据初始化
    cufftReal *h_data = (cufftReal*)malloc(mem_size);
    cufftReal *h_data3 = (cufftReal*)malloc(mem_size);
    for(int i=0; i<real_size; i++){
        x = (i % (N+2));
        if(x < N){h_data[i] = 2;} else{h_data[i] = 0;}
    }

    // --------------------------
    // 初始化cufftXt多GPU环境
    // --------------------------
    cufftHandle plan;
    gpuErrchk(cufftCreate(&plan));

    // 指定使用GPU 0和1
    int devices[] = {0, 1};
    gpuErrchk(cufftXtSetGPUs(plan, 2, devices));

    // 创建2D实数到复数的FFT计划（根据你的需求调整）
    int rank = 2;
    int n[] = {N, N};
    int istride = 1, idist = real_size;
    int ostride = 1, odist = cplx_size;
    size_t work_size = 0;
    gpuErrchk(cufftMakePlanMany(plan, rank, n,
                                NULL, istride, idist, CUFFT_R2C,
                                NULL, ostride, odist, CUFFT_C2R,
                                1, &work_size));

    // --------------------------
    // 分配分布式设备内存
    // --------------------------
    cufftReal *d_data = NULL;
    cufftReal *d_data3 = NULL;
    gpuErrchk(cufftXtMalloc(plan, (void**)&d_data, CUFFT_XT_FORMAT_INPLACE));
    gpuErrchk(cufftXtMalloc(plan, (void**)&d_data3, CUFFT_XT_FORMAT_INPLACE));

    // 将主机数据拷贝到分布式设备内存
    gpuErrchk(cufftXtMemcpy(plan, d_data, h_data, CUFFT_COPY_HOST_TO_DEVICE));

    cout << "Distributed device memory allocated and data copied" << endl;

    // --------------------------
    // 获取每个GPU上的内存指针并执行核函数
    // --------------------------
    int num_devices = 2;
    for (int dev_idx = 0; dev_idx < num_devices; dev_idx++) {
        int dev_id = devices[dev_idx];
        gpuErrchk(cudaSetDevice(dev_id));

        // 获取当前GPU上的数据子数组指针和尺寸
        cufftReal *sub_data, *sub_data3;
        size_t sub_size;
        gpuErrchk(cufftXtGetSubarray(plan, dev_idx, (void**)&sub_data, &sub_size));
        gpuErrchk(cufftXtGetSubarray(plan, dev_idx, (void**)&sub_data3, &sub_size));
        int half_real_size = sub_size / sizeof(cufftReal);

        // 启动核函数
        int maxThreads = (N > 1024) ? 1024 : N;
        int threadsPerBlock = maxThreads;
        int numBlocks = (half_real_size + threadsPerBlock - 1) / threadsPerBlock; // 确保整除

        Cube<<<numBlocks, threadsPerBlock>>>(sub_data, sub_data3, N, half_real_size);
        gpuErrchk(cudaPeekAtLastError());
        gpuErrchk(cudaDeviceSynchronize());
    }

    // --------------------------
    // 执行FFT运算
    // --------------------------
    // 这里以实数到复数FFT为例，根据你的需求调整方向
    gpuErrchk(cufftXtExecR2C(plan, d_data3, (cufftComplex*)d_data3));

    // 将结果拷贝回主机
    gpuErrchk(cufftXtMemcpy(plan, h_data3, d_data3, CUFFT_COPY_DEVICE_TO_HOST));

    // 打印结果（仅示例，大内存时建议注释）
    cout << endl << "Processed data output:" << endl;
    for(int i = 0; i<real_size; i++){
        cout << h_data3[i] << " ";
    }
    cout << endl;

    // --------------------------
    // 资源清理
    // --------------------------
    cufftXtFree(d_data);
    cufftXtFree(d_data3);
    cufftDestroy(plan);

    free(h_data);
    free(h_data3);

    return 0;
}

关键代码解释 #

• cufftXtSetGPUs: 告诉cufftXt要使用哪些GPU设备，这里指定了0和1。
• cufftXtMalloc: 自动在指定GPU上分配分布式内存，替代手动的cudaMalloc，内存会按最优策略拆分到各个GPU。
• cufftXtGetSubarray: 核心函数，获取指定GPU上的内存块指针和大小，让核函数可以直接操作cufftXt管理的内存。
• cufftXtMemcpy: 自动处理跨GPU的数据拷贝，无需手动拆分主机数据，简化了内存传输流程。

注意事项 #

• 确保你的CUDA和cuFFT版本支持cufftXt功能（通常CUDA 7.0及以上版本支持）。
• 核函数中的数据尺寸要和cufftXtGetSubarray返回的sub_size匹配，避免越界访问。
• 对于超大规模数据（数GB级别），这种方式能有效减少内存冗余，提升整体性能，因为数据不需要在多份内存之间拷贝。

内容的提问来源于stack exchange，提问作者kyle moats