From f1efc03aed04e732ad99ff178842bf4fdb0bf5d6 Mon Sep 17 00:00:00 2001
From: "giovanni.lacopo" <giovanni.lacopo@inaf.it>
Date: Fri, 8 Sep 2023 15:21:32 +0200
Subject: [PATCH] cufftMP working with CUDA and OpenMP

---
 cuda_fft.cu | 188 ++++++++++++++++++++++++++++------------------------
 1 file changed, 101 insertions(+), 87 deletions(-)

diff --git a/cuda_fft.cu b/cuda_fft.cu
index 96530e0..0d88ea8 100644
--- a/cuda_fft.cu
+++ b/cuda_fft.cu
@@ -5,144 +5,158 @@
 #include <cuda_runtime.h>
 #include <complex.h>
 #include "cuComplex.h"
-#include "w-stacking.h"
+#include "proto.h"
+#include "errcodes.h"
 #include <time.h>
 
-#if defined(CUFFTMP) && !defined(USE_FFTW)
+#if defined(CUFFTMP) && defined(USE_FFTW)
 
 void cuda_fft(
-	int num_w_planes,
-	int grid_size_x,
-	int grid_size_y,
-	int xaxis,
-	int yaxis,
-	double * grid,
-	double * gridss,
-	MPI_Comm comm)
+	      int num_w_planes,
+	      int grid_size_x,
+	      int grid_size_y,
+	      int xaxis,
+	      int yaxis,
+	      double * grid,
+	      double * gridss,
+	      int rank,
+	      MPI_Comm comm)
 {
-#ifdef __CUDACC__
 
-        cudaError_t mmm;
-        cufftResult_t status;
+ #if !defined __CUDACC__
+  int ndevices;
+  cudaGetDeviceCount(&ndevices);
+  cudaSetDevice(rank % ndevices);
 
-	cufftDoubleComplex *fftwgrid;
-	fftwgrid = (cufftDoubleComplex*) malloc(sizeof(cufftDoubleComplex)*2*num_w_planes*yaxis*grid_size_x);
+  if ( rank == 0 ) {
+    if (0 == ndevices) {
+      
+      shutdown_wstacking(NO_ACCELERATORS_FOUND, "No accelerators found", __FILE__, __LINE__ );
+    }
+  }
+ #endif
+  
+  cudaError_t mmm;
+  cufftResult_t status;
 
+  cufftDoubleComplex *fftwgrid;
+  fftwgrid = (cufftDoubleComplex*) malloc(sizeof(cufftDoubleComplex)*2*num_w_planes*yaxis*grid_size_x);
 
 
-	// Plan creation
 
-	cufftHandle plan;
-        status = cufftCreate(&plan);
-        if (status != CUFFT_SUCCESS) {printf("!!! cufftCreate ERROR %d !!!\n", status);}
+  // Plan creation
 
-	cudaStream_t stream{};
-	cudaStreamCreate(&stream);
+  cufftHandle plan;
+  status = cufftCreate(&plan);
+  if (status != CUFFT_SUCCESS) {printf("!!! cufftCreate ERROR %d !!!\n", status);}
 
+  cudaStream_t stream{};
+  cudaStreamCreate(&stream);
 
-	status = cufftMpAttachComm(plan, CUFFT_COMM_MPI, &comm);
-	if (status != CUFFT_SUCCESS) {printf("!!! cufftMpAttachComm ERROR %d !!!\n", status);}
 
-	status = cufftSetStream(plan, stream);
-	if (status != CUFFT_SUCCESS) {printf("!!! cufftSetStream ERROR %d !!!\n", status);}
+  status = cufftMpAttachComm(plan, CUFFT_COMM_MPI, &comm);
+  if (status != CUFFT_SUCCESS) {printf("!!! cufftMpAttachComm ERROR %d !!!\n", status);}
 
-	size_t workspace;
-	status = cufftMakePlan2d(plan, grid_size_x, grid_size_y, CUFFT_Z2Z, &workspace);
-	if (status != CUFFT_SUCCESS) {printf("!!! cufftMakePlan2d ERROR %d !!!\n", status);}
-	cudaDeviceSynchronize();
+  status = cufftSetStream(plan, stream);
+  if (status != CUFFT_SUCCESS) {printf("!!! cufftSetStream ERROR %d !!!\n", status);}
 
+  size_t workspace;
+  status = cufftMakePlan2d(plan, grid_size_x, grid_size_y, CUFFT_Z2Z, &workspace);
+  if (status != CUFFT_SUCCESS) {printf("!!! cufftMakePlan2d ERROR %d !!!\n", status);}
+  cudaDeviceSynchronize();
 
 
-	long fftwindex = 0;
-	long fftwindex2D = 0;
-	double norm = 1.0/(double)(grid_size_x*grid_size_y);
 
+  uint fftwindex = 0;
+  uint fftwindex2D = 0;
+  double norm = 1.0/(double)(grid_size_x*grid_size_y);
 
 
 
-	// Grid composition
 
-	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];
-                        }
-                }
+  // Grid composition
+
+  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];
+	    }
+	}
 
 
-		cudaLibXtDesc *fftwgrid_g;
-        	cudaLibXtDesc *fftwgrid_g2;
+      cudaLibXtDesc *fftwgrid_g;
+      cudaLibXtDesc *fftwgrid_g2;
 
 
-		status = cufftXtMalloc(plan, &fftwgrid_g, CUFFT_XT_FORMAT_INPLACE);
-		if (status != CUFFT_SUCCESS) {printf("!!! cufftXtMalloc ERROR %d !!!\n", status);}
+      status = cufftXtMalloc(plan, &fftwgrid_g, CUFFT_XT_FORMAT_INPLACE);
+      if (status != CUFFT_SUCCESS) {printf("!!! cufftXtMalloc ERROR %d !!!\n", status);}
 
-	        status = cufftXtMalloc(plan, &fftwgrid_g2, CUFFT_XT_FORMAT_INPLACE);
-        	if (status != CUFFT_SUCCESS) {printf("!!! cufftXtMalloc 2 ERROR %d !!!\n", status);}
-	        cudaDeviceSynchronize();
+      status = cufftXtMalloc(plan, &fftwgrid_g2, CUFFT_XT_FORMAT_INPLACE);
+      if (status != CUFFT_SUCCESS) {printf("!!! cufftXtMalloc 2 ERROR %d !!!\n", status);}
+      cudaDeviceSynchronize();
 
-                mmm = cudaStreamSynchronize(stream);
-                if (mmm != cudaSuccess) {printf("!!! cudaStreamSynchronize ERROR %d !!!\n", mmm);}
+      mmm = cudaStreamSynchronize(stream);
+      if (mmm != cudaSuccess) {printf("!!! cudaStreamSynchronize ERROR %d !!!\n", mmm);}
 
-		status = cufftXtMemcpy(plan, fftwgrid_g, fftwgrid, CUFFT_COPY_HOST_TO_DEVICE);
-		if (status != CUFFT_SUCCESS) {printf("!!! cufftXtMemcpy htd fftwgrid_g ERROR %d !!!\n", status);}
+      status = cufftXtMemcpy(plan, fftwgrid_g, fftwgrid, CUFFT_COPY_HOST_TO_DEVICE);
+      if (status != CUFFT_SUCCESS) {printf("!!! cufftXtMemcpy htd fftwgrid_g ERROR %d !!!\n", status);}
 
-		cudaDeviceSynchronize();
+      cudaDeviceSynchronize();
 
-		status = cufftXtExecDescriptor(plan, fftwgrid_g, fftwgrid_g, CUFFT_INVERSE);
-		if (status != CUFFT_SUCCESS) {printf("!!! cufftXtExecDescriptor ERROR %d !!!\n", status);}
+      status = cufftXtExecDescriptor(plan, fftwgrid_g, fftwgrid_g, CUFFT_INVERSE);
+      if (status != CUFFT_SUCCESS) {printf("!!! cufftXtExecDescriptor ERROR %d !!!\n", status);}
 
 
-                mmm = cudaStreamSynchronize(stream);
-                if (mmm != cudaSuccess) {printf("!!! cudaStreamSynchronize 2 ERROR %d !!!\n", mmm);}
+      mmm = cudaStreamSynchronize(stream);
+      if (mmm != cudaSuccess) {printf("!!! cudaStreamSynchronize 2 ERROR %d !!!\n", mmm);}
 
-		cudaDeviceSynchronize();
+      cudaDeviceSynchronize();
 
 
-                status = cufftXtMemcpy(plan, fftwgrid_g2, fftwgrid_g, CUFFT_COPY_DEVICE_TO_DEVICE);
-                if (status != CUFFT_SUCCESS) {printf("!!! cufftXtMemcpy dtd fftwgrid ERROR %d !!!\n", status);}
+      status = cufftXtMemcpy(plan, fftwgrid_g2, fftwgrid_g, CUFFT_COPY_DEVICE_TO_DEVICE);
+      if (status != CUFFT_SUCCESS) {printf("!!! cufftXtMemcpy dtd fftwgrid ERROR %d !!!\n", status);}
 
 
-                mmm = cudaStreamSynchronize(stream);
-                if (mmm != cudaSuccess) {printf("!!! cudaStreamSynchronize 2 ERROR %d !!!\n", mmm);}
+      mmm = cudaStreamSynchronize(stream);
+      if (mmm != cudaSuccess) {printf("!!! cudaStreamSynchronize 2 ERROR %d !!!\n", mmm);}
 
-                cudaDeviceSynchronize();
+      cudaDeviceSynchronize();
 
-                status = cufftXtMemcpy(plan, fftwgrid, fftwgrid_g2, CUFFT_COPY_DEVICE_TO_HOST);
-                if (status != CUFFT_SUCCESS) {printf("!!! cufftXtMemcpy dth fftwgrid ERROR %d !!!\n", status);}
+      status = cufftXtMemcpy(plan, fftwgrid, fftwgrid_g2, CUFFT_COPY_DEVICE_TO_HOST);
+      if (status != CUFFT_SUCCESS) {printf("!!! cufftXtMemcpy dth fftwgrid ERROR %d !!!\n", status);}
 
 
-		cufftXtFree(fftwgrid_g);
-		cufftXtFree(fftwgrid_g2);
+      cufftXtFree(fftwgrid_g);
+      cufftXtFree(fftwgrid_g2);
 
-		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;
-			}
-        	}
+      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;
+	    }
 	}
+    }
 
 
 
-	status = cufftDestroy(plan);
-	if (status != CUFFT_SUCCESS) {printf("!!! cufftDestroy fftwgrid ERROR %d !!!\n", status);}
+  status = cufftDestroy(plan);
+  if (status != CUFFT_SUCCESS) {printf("!!! cufftDestroy fftwgrid ERROR %d !!!\n", status);}
 
 
-	cudaStreamDestroy(stream);
-	cudaDeviceSynchronize();
+  cudaStreamDestroy(stream);
+  cudaDeviceSynchronize();
 
-#endif // __CUDACC__
 }
 #endif 
+  
-- 
GitLab