Single node multi-GPU cuFFT

#include <math.h>

#include <stdlib.h>

#include <stdio.h>

#ifndef MULTI_GPU

#include <cufft.h>

#include <cufftw.h>

#else

#include <cufftXt.h>

#endif

#include "w-stacking.h"

#include <cuda_runtime.h>

#include <complex.h>

#include "cuComplex.h"

/*************************************************************************/

//	CUDA KERNELS FOR fftwgrid

/*

#ifdef __CUDACC__

__global__ void write_grid(

	int num_w_planes,

	int xaxis,

	int yaxis,

	cufftDoubleComplex * fftwgrid,

	double * grid)

{

	long fftwindex2D = blockIdx.x;

	long fftwindex = 2*(blockIdx.x+threadIdx.x*xaxis*yaxis);

	fftwgrid[fftwindex2D].x = grid[fftwindex];

	fftwgrid[fftwindex2D].y = grid[fftwindex+1];

}

#endif

*/

/*

#ifdef __CUDACC__

__global__ void write_gridss(

	int num_w_planes,

	int xaxis,

	int yaxis,

	cufftDoubleComplex * fftwgrid,

	double * gridss,

	int grid_size_x,

	int grid_size_y)

{

	long fftwindex2D = blockIdx.x;

	long fftwindex = 2*(blockIdx.x+threadIdx.x*xaxis*yaxis);

	double norm = 1.0/(double)(grid_size_x*grid_size_y);

	gridss[fftwindex] = norm*fftwgrid[fftwindex2D].x;

	gridss[fftwindex+1] = norm*fftwgrid[fftwindex2D].y;

}

#endif

*/

/*************************************************************************/

void cuda_fft(

int num_w_planes,

int grid_size_x,

int grid_size_y,

int xaxis,

int yaxis,

double * grid,

double * gridss)

{

#ifdef __CUDACC__

	cudaError_t mmm;

	cufftResult_t status;

	clock_t startcu, endcu;

	double cufft_exec_time;

	struct timespec begincu, finishcu;

	// Define CUDA setup

	long Nbl = (long)(xaxis+yaxis*xaxis);

	int Nth = num_w_planes;

	if(NWORKERS == 1) {Nbl = 1; Nth = 1;}

#ifndef MULTI_GPU

	printf("Performing single GPU cuFFT\n");

	cufftDoubleComplex *fftwgrid;

	cufftDoubleComplex *fftwgrid_g;

	//double * grid_g;

	//double * gridss_g;

	fftwgrid = (cufftDoubleComplex*) malloc(sizeof(cufftDoubleComplex)*2*num_w_planes*xaxis*yaxis);

	mmm = cudaMalloc(&fftwgrid_g, sizeof(cufftDoubleComplex)*2*num_w_planes*xaxis*yaxis);

	//mmm = cudaMalloc(&grid_g, sizeof(double)*2*num_w_planes*xaxis*yaxis);

	//mmm = cudaMalloc(&gridss_g, sizeof(double)*2*num_w_planes*xaxis*yaxis);

	//mmm = cudaMemcpy(grid_g, grid, sizeof(double)*2*num_w_planes*xaxis*yaxis, cudaMemcpyHostToDevice);

	if (mmm != cudaSuccess) {printf("!!! cudaMemcpy grid ERROR %d !!!\n", mmm);}

	cudaDeviceSynchronize();

	cufftHandle plan;

	cufftCreate(&plan);

	cufftPlan1d(&plan, num_w_planes*xaxis*yaxis, CUFFT_Z2Z, 1);

	double norm = 1.0/(double)(grid_size_x*grid_size_y);

	long fftwindex = 0;

	long fftwindex2D = 0;

	for (int iw=0; iw<num_w_planes; iw++)

	{

		printf("select the %d w-plane to transform\n", iw);

		for (int iv=0; iv<yaxis; iv++)

		{

			for (int iu=0; iu<xaxis; iu++)

			{

			fftwindex2D = iu + iv*xaxis;

			fftwindex = 2*(fftwindex2D + iw*xaxis*yaxis);

			fftwgrid[fftwindex2D].x = grid[fftwindex];

			fftwgrid[fftwindex2D].y = grid[fftwindex+1];

			}

		}

	}

	mmm=cudaMemcpy(fftwgrid_g, fftwgrid, sizeof(cufftDoubleComplex)*2*num_w_planes*xaxis*yaxis, cudaMemcpyHostToDevice);

	if (mmm != cudaSuccess) {printf("cudaMemcpy ERROR %d\n",mmm);}

	cudaDeviceSynchronize();

	/*

	// CUDA kernel grid composition

	write_grid<<<Nbl, Nth>>>(num_w_planes, xaxis, yaxis, fftwgrid_g, grid_g);

	cudaDeviceSynchronize();

	*/

	// cuFFT plan execution

	clock_gettime(CLOCK_MONOTONIC, &begincu);

	startcu = clock();

	status = cufftExecZ2Z(plan, fftwgrid_g, fftwgrid_g, CUFFT_INVERSE);

	cudaDeviceSynchronize();

	endcu = clock();

	clock_gettime(CLOCK_MONOTONIC, &finishcu);

	cufft_exec_time = ((double) (endcu - startcu)) / CLOCKS_PER_SEC;

	printf("Time for cuFFT plan execution is %f sec\n", cufft_exec_time);

	if (status != CUFFT_SUCCESS) {printf("!!! cufftExec ERROR %d !!!\n", status);}

	// CUDA kernel grid decomposition

	/*

	write_gridss<<<Nbl, Nth>>>(num_w_planes, xaxis, yaxis, fftwgrid_g, gridss_g, grid_size_x, grid_size_y);

	cudaDeviceSynchronize();

	mmm = cudaMemcpy(gridss, gridss_g, sizeof(double)*2*num_w_planes*xaxis*yaxis, cudaMemcpyDeviceToHost);

	if (mmm != cudaSuccess) {printf("!!! cudaMemcpy gridss ERROR %d !!!\n", mmm);}

	cudaDeviceSynchronize();

	*/

	mmm = cudaMemcpy(fftwgrid, fftwgrid_g, sizeof(cufftDoubleComplex)*2*num_w_planes*xaxis*yaxis, cudaMemcpyDeviceToHost);

	if (mmm != cudaSuccess) {printf("!!! cudaMemcpy fftwgrid ERROR %d !!!\n", mmm);}

	cudaDeviceSynchronize();

	for (int iw=0; iw<num_w_planes; iw++)

	{

		for (int iv=0; iv<yaxis; iv++)

		{

			for (int iu=0; iu<xaxis; iu++)

			{

			fftwindex2D = iu + iv*xaxis;

			fftwindex = 2*(fftwindex2D + iw*xaxis*yaxis);

			gridss[fftwindex] = norm*fftwgrid[fftwindex2D].x;

			gridss[fftwindex+1] = norm*fftwgrid[fftwindex2D].y;

			}

		}

	}

	cufftDestroy(plan);

	mmm = cudaFree(fftwgrid_g);

	//mmm = cudaFree(grid_g);

	//mmm = cudaFree(gridss_g);

	cudaDeviceSynchronize();

#else

	printf("Performing multiple GPU cuFFT\n");

	cufftDoubleComplex *fftwgrid;

	fftwgrid = (cufftDoubleComplex*) malloc(sizeof(cufftDoubleComplex)*2*num_w_planes*xaxis*yaxis);

	double norm = 1.0/(double)(grid_size_x*grid_size_y);

	long fftwindex2D = 0;

	long fftwindex = 0;

	for (int iw=0; iw<num_w_planes; iw++)

	{

		for (int iv=0; iv<yaxis; iv++)

		{

			for (int iu=0; iu<xaxis; iu++)

			{

			fftwindex2D = iu + iv*xaxis;

			fftwindex = 2*(fftwindex2D + iw*xaxis*yaxis);

			fftwgrid[fftwindex2D].x = grid[fftwindex];

			fftwgrid[fftwindex2D].y = grid[fftwindex+1];

			}

		}

	}

	// Specify how many GPUs to be used and their Ids

	const int nDevices = 4;

	int deviceIds[nDevices] = {0, 1, 2, 3};

	size_t workspace_sizes[nDevices];

	cudaLibXtDesc *fftwgrid_multig;

	cufftHandle plan;

	cufftCreate(&plan);

	#if CUFFT_VERSION >= 10400

	cudaStream_t stream{};

	cudaStreamCreate(&stream);

	cufftSetStream(plan, stream);

	#endif

	status = cufftXtSetGPUs(plan, nDevices, deviceIds);

	if (status != CUFFT_SUCCESS) {printf("!!! cufftXtSetGPUs ERROR %d !!!\n", status);}

	status = cufftMakePlan1d(plan, num_w_planes*xaxis*yaxis, CUFFT_Z2Z, 1, workspace_sizes);

	if (status != CUFFT_SUCCESS) {printf("!!! cufftMakePaln1d ERROR %d !!!\n", status);}

	cudaDeviceSynchronize();

	status = cufftXtMalloc(plan, &fftwgrid_multig, CUFFT_XT_FORMAT_INPLACE);

	if (status != CUFFT_SUCCESS) {printf("!!! cuufftXtMalloc ERROR %d !!!\n", status);}

	cudaDeviceSynchronize();

	status = cufftXtMemcpy(plan, fftwgrid_multig, fftwgrid, CUFFT_COPY_HOST_TO_DEVICE);

	if (status != CUFFT_SUCCESS) {printf("!!! cufftXtMemcpy fftwgrid_multig ERROR %d !!!\n", status);}

	cudaDeviceSynchronize();

	clock_gettime(CLOCK_MONOTONIC, &begincu);

	startcu = clock();

	status = cufftXtExecDescriptorZ2Z(plan, fftwgrid_multig, fftwgrid_multig, CUFFT_INVERSE);

	if (status != CUFFT_SUCCESS) {printf("!!! cufftXtExec ERROR %d !!!\n", status);}

	cudaDeviceSynchronize();

	endcu = clock();

	clock_gettime(CLOCK_MONOTONIC, &finishcu);

	cufft_exec_time = ((double) (endcu - startcu)) / CLOCKS_PER_SEC;

	printf("Time for multiGPU cuFFT plan execution is %f sec\n", cufft_exec_time);

	status = cufftXtMemcpy(plan, fftwgrid, fftwgrid_multig, CUFFT_COPY_DEVICE_TO_HOST);

	if (status != CUFFT_SUCCESS) {printf("!!! cufftXtMemcpy fftwgrid ERROR %d !!!\n", status);}

	cudaDeviceSynchronize();

	for (int iw=0; iw<num_w_planes; iw++)

	{

		for (int iv=0; iv<yaxis; iv++)

		{

			for (int iu=0; iu<xaxis; iu++)

			{

			fftwindex2D = iu + iv*xaxis;

			fftwindex = 2*(fftwindex2D + iw*xaxis*yaxis);

			gridss[fftwindex] = norm*fftwgrid[fftwindex2D].x;

			gridss[fftwindex+1] = norm*fftwgrid[fftwindex2D].y;

			}

		}

	}

	#if CUFFT_VERSION >= 10400

	CUDA_RT_CALL(cudaStreamDestroy(stream));

	#endif

	cufftXtFree(fftwgrid_multig);

	cufftDestroy(plan);

#endif // MULTI_GPU

#endif // __CUDACC__

}