Multithread fft

3 files changed, 42 insertions, 20 deletions

diff --git a/src/main.c b/src/main.c
index 904b5bf..5c7baf7 100644
--- a/src/main.c
+++ b/src/main.c
@@ -13,7 +13,7 @@
 int print_usage();
 
 // read command line arguments
-int read_args(int argc, const char* argv[], char** params, unsigned int* driving_force, unsigned int* command);
+int read_args(int argc, const char* argv[], char** params, unsigned int* driving_force, unsigned int* command, unsigned int* nthreads);
 int read_params(char* params, int* K1, int* K2, int* N1, int* N2, unsigned int* nsteps, double* nu, double* delta, unsigned int* print_freq);
 int set_parameter(char* lhs, char* rhs, int* K1, int* K2, int* N1, int* N2, unsigned int* nsteps, double* nu, double* delta, unsigned int* print_freq, bool* setN1, bool* setN2);
 
@@ -38,12 +38,13 @@ int main (
   int ret;
   unsigned int driving,command;
   unsigned int print_freq;
+  unsigned int nthreads=1;
 
   command=0;
   driving=0;
 
   // read command line arguments
-  ret=read_args(argc, argv, &params, &driving, &command);
+  ret=read_args(argc, argv, &params, &driving, &command, &nthreads);
   if(ret<0){
     return(-1);
   }
@@ -65,13 +66,13 @@ int main (
 
   // run command
   if (command==COMMAND_UK){
-    uk(K1, K2, N1, N2, nsteps, nu, delta, g, print_freq);
+    uk(K1, K2, N1, N2, nsteps, nu, delta, g, print_freq, nthreads);
   }
   else if(command==COMMAND_ENSTROPHY){
-    enstrophy(K1, K2, N1, N2, nsteps, nu, delta, g, print_freq);
+    enstrophy(K1, K2, N1, N2, nsteps, nu, delta, g, print_freq, nthreads);
   }
   else if(command==COMMAND_QUIET){
-    quiet(K1, K2, N1, N2, nsteps, nu, delta, g);
+    quiet(K1, K2, N1, N2, nsteps, nu, delta, g, nthreads);
   }
   else if(command==0){
     fprintf(stderr, "error: no command specified\n");
@@ -90,14 +91,17 @@ int print_usage(){
 // read command line arguments
 #define CP_FLAG_PARAMS 1
 #define CP_FLAG_DRIVING 2
+#define CP_FLAG_NTHREADS 3
 int read_args(
   int argc,
   const char* argv[],
   char** params,
   unsigned int* driving_force,
-  unsigned int* command
+  unsigned int* command,
+  unsigned int* nthreads
 ){
   int i;
+  int ret;
   // pointers
   char* ptr;
   // flag that indicates what argument is being read
@@ -109,14 +113,15 @@ int read_args(
     if(argv[i][0]=='-'){
       for(ptr=((char*)argv[i])+1;*ptr!='\0';ptr++){
 	switch(*ptr){
-	// timestep
 	case 'p':
 	  flag=CP_FLAG_PARAMS;
 	  break;
-	// nsteps
 	case 'g':
 	  flag=CP_FLAG_DRIVING;
 	  break;
+	case 't':
+	  flag=CP_FLAG_NTHREADS;
+	  break;
 	default:
 	  fprintf(stderr, "unrecognized option '-%c'\n", *ptr);
 	  print_usage();
@@ -141,6 +146,15 @@ int read_args(
       }
       flag=0;
     }
+    // nthreads
+    else if(flag==CP_FLAG_NTHREADS){
+      ret=sscanf(argv[i],"%u",nthreads);
+      if(ret!=1){
+	fprintf(stderr, "error: '-t' should be followed by an unsigned integer\n       got '%s'\n",argv[i]);
+	return(-1);
+      }
+      flag=0;
+    }
     // computation to run
     else{
       if(strcmp(argv[i], "uk")==0){
diff --git a/src/navier-stokes.c b/src/navier-stokes.c
index 87f5aba..7f22c31 100644
--- a/src/navier-stokes.c
+++ b/src/navier-stokes.c
@@ -14,7 +14,8 @@ int uk(
   double nu,
   double delta,
   _Complex double (*g)(int,int),
-  unsigned int print_freq
+  unsigned int print_freq,
+  unsigned int nthreads
 ){
   _Complex double* u;
   _Complex double* tmp1;
@@ -26,7 +27,7 @@ int uk(
   fft_vect ifft;
   int kx,ky;
 
-  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &fft1, &fft2, &ifft, K1, K2, N1, N2);
+  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &fft1, &fft2, &ifft, K1, K2, N1, N2, nthreads);
   ns_init_u(u, K1, K2);
 
   // iterate
@@ -65,7 +66,8 @@ int enstrophy(
   double nu,
   double delta,
   _Complex double (*g)(int,int),
-  unsigned int print_freq
+  unsigned int print_freq,
+  unsigned int nthreads
 ){
   _Complex double* u;
   _Complex double* tmp1;
@@ -78,7 +80,7 @@ int enstrophy(
   fft_vect fft2;
   fft_vect ifft;
 
-  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &fft1, &fft2, &ifft, K1, K2, N1, N2);
+  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &fft1, &fft2, &ifft, K1, K2, N1, N2, nthreads);
   ns_init_u(u, K1, K2);
 
 
@@ -114,7 +116,8 @@ int quiet(
   unsigned int nsteps,
   double nu,
   double delta,
-  _Complex double (*g)(int,int)
+  _Complex double (*g)(int,int),
+  unsigned int nthreads
 ){
   _Complex double* u;
   _Complex double* tmp1;
@@ -125,7 +128,7 @@ int quiet(
   fft_vect fft2;
   fft_vect ifft;
 
-  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &fft1, &fft2, &ifft, K1, K2, N1, N2);
+  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &fft1, &fft2, &ifft, K1, K2, N1, N2, nthreads);
   ns_init_u(u, K1, K2);
 
   // iterate
@@ -150,7 +153,8 @@ int  ns_init_tmps(
   int K1,
   int K2,
   int N1,
-  int N2
+  int N2,
+  unsigned int nthreads
 ){
   // velocity field
   *u=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));
@@ -160,6 +164,10 @@ int  ns_init_tmps(
   *tmp2=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));
   *tmp3=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));
 
+  // init threads
+  fftw_init_threads();
+  fftw_plan_with_nthreads(nthreads);
+
   // prepare vectors for fft
   fft1->fft=fftw_malloc(sizeof(fftw_complex)*N1*N2);
   fft1->fft_plan=fftw_plan_dft_2d(N1,N2, fft1->fft, fft1->fft, FFTW_FORWARD, FFTW_MEASURE);
@@ -189,7 +197,7 @@ int ns_free_tmps(
   fftw_free(fft2.fft);
   fftw_free(ifft.fft);
 
-  fftw_cleanup();
+  fftw_cleanup_threads();
 
   free(tmp3);
   free(tmp2);
diff --git a/src/navier-stokes.h b/src/navier-stokes.h
index d144fce..707dc52 100644
--- a/src/navier-stokes.h
+++ b/src/navier-stokes.h
@@ -12,17 +12,17 @@ typedef struct fft_vects {
 } fft_vect;
 
 // compute u_k
-int uk( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, _Complex double (*g)(int,int), unsigned int print_freq);
+int uk( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, _Complex double (*g)(int,int), unsigned int print_freq, unsigned int nthreads);
 
 // compute enstrophy
-int enstrophy( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, _Complex double (*g)(int,int), unsigned int print_freq);
+int enstrophy( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, _Complex double (*g)(int,int), unsigned int print_freq, unsigned int nthreads);
 
 // compute solution as a function of time, but do not print anything (useful for debugging)
-int quiet( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, _Complex double (*g)(int,int));
+int quiet( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, _Complex double (*g)(int,int), unsigned int nthreads);
 
 
 // initialize vectors for computation
-int  ns_init_tmps( _Complex double **u, _Complex double ** tmp1, _Complex double **tmp2, _Complex double **tmp3, fft_vect* fft1, fft_vect *fft2, fft_vect *ifft, int K1, int K2, int N1, int N2);
+int  ns_init_tmps( _Complex double **u, _Complex double ** tmp1, _Complex double **tmp2, _Complex double **tmp3, fft_vect* fft1, fft_vect *fft2, fft_vect *ifft, int K1, int K2, int N1, int N2, unsigned int nthreads);
 // release vectors
 int ns_free_tmps( _Complex double* u, _Complex double* tmp1, _Complex double *tmp2,_Complex double *tmp3, fft_vect fft1, fft_vect fft2, fft_vect ifft);