Use the threadpool to parallelize the Green function calculation in Fourier space.

src/mesh_gravity.c

@@ -345,6 +345,137 @@ void mesh_to_gparts_CIC(struct gpart* gp, const double* pot, const int N,
 #endif
 }
 
+/**
+ * @brief Shared information about the Green function to be used by all the
+ * threads in the pool.
+ */
+struct Green_function_data {
+
+  int N;
+  fftw_complex* frho;
+  double green_fac;
+  double a_smooth2;
+  double k_fac;
+};
+
+/**
+ * @brief Mapper function for the application of the Green function.
+ *
+ * @param map_data The array of the density field Fourier transform.
+ * @param num The number of elements to iterate on (along the x-axis).
+ * @param extra The properties of the Green function.
+ */
+void mesh_apply_Green_function_mapper(void* map_data, const int num,
+                                      void* extra) {
+
+  struct Green_function_data* data = (struct Green_function_data*)extra;
+
+  /* Unpack the array */
+  fftw_complex* const frho = data->frho;
+  const int N = data->N;
+  const int N_half = N / 2;
+
+  /* Unpack the Green function properties */
+  const double green_fac = data->green_fac;
+  const double a_smooth2 = data->a_smooth2;
+  const double k_fac = data->k_fac;
+
+  /* Range handled by this call */
+  const int i_start = (fftw_complex*)map_data - frho;
+  const int i_end = i_start + num;
+
+  /* Loop over the x range corresponding to this thread */
+  for (int i = i_start; i < i_end; ++i) {
+
+    /* kx component of vector in Fourier space and 1/sinc(kx) */
+    const int kx = (i > N_half ? i - N : i);
+    const double kx_d = (double)kx;
+    const double fx = k_fac * kx_d;
+    const double sinc_kx_inv = (kx != 0) ? fx / sin(fx) : 1.;
+
+    for (int j = 0; j < N; ++j) {
+
+      /* ky component of vector in Fourier space and 1/sinc(ky) */
+      const int ky = (j > N_half ? j - N : j);
+      const double ky_d = (double)ky;
+      const double fy = k_fac * ky_d;
+      const double sinc_ky_inv = (ky != 0) ? fy / sin(fy) : 1.;
+
+      for (int k = 0; k < N_half + 1; ++k) {
+
+        /* kz component of vector in Fourier space and 1/sinc(kz) */
+        const int kz = (k > N_half ? k - N : k);
+        const double kz_d = (double)kz;
+        const double fz = k_fac * kz_d;
+        const double sinc_kz_inv = (kz != 0) ? fz / (sin(fz) + FLT_MIN) : 1.;
+
+        /* Norm of vector in Fourier space */
+        const double k2 = (kx_d * kx_d + ky_d * ky_d + kz_d * kz_d);
+
+        /* Avoid FPEs... */
+        if (k2 == 0.) continue;
+
+        /* Green function */
+        double W = 1.;
+        fourier_kernel_long_grav_eval(k2 * a_smooth2, &W);
+        const double green_cor = green_fac * W / (k2 + FLT_MIN);
+
+        /* Deconvolution of CIC */
+        const double CIC_cor = sinc_kx_inv * sinc_ky_inv * sinc_kz_inv;
+        const double CIC_cor2 = CIC_cor * CIC_cor;
+        const double CIC_cor4 = CIC_cor2 * CIC_cor2;
+
+        /* Combined correction */
+        const double total_cor = green_cor * CIC_cor4;
+
+        /* Apply to the mesh */
+        const int index = N * (N_half + 1) * i + (N_half + 1) * j + k;
+        frho[index][0] *= total_cor;
+        frho[index][1] *= total_cor;
+      }
+    }
+  }
+}
+
+/**
+ * @brief Apply the Green function in Fourier space to the density
+ * array to get the potential.
+ *
+ * Also deconvolves the CIC kernel.
+ *
+ * @param tp The threadpool.
+ * @param frho The NxNx(N/2) complex array of the Fourier transform of the
+ * density field.
+ * @param N The dimension of the array.
+ * @param r_s The Green function smoothing scale.
+ * @param box_size The physical size of the simulation box.
+ */
+void mesh_apply_Green_function(struct threadpool* tp, fftw_complex* frho,
+                               const int N, const double r_s,
+                               const double box_size) {
+
+  /* Some common factors */
+  struct Green_function_data data;
+  data.frho = frho;
+  data.N = N;
+  data.green_fac = -1. / (M_PI * box_size);
+  data.a_smooth2 = 4. * M_PI * M_PI * r_s * r_s / (box_size * box_size);
+  data.k_fac = M_PI / (double)N;
+
+  /* Parallelize the Green function application using the threadpool
+     to split the x-axis loop over the threads. */
+  if (N < 64) {
+    mesh_apply_Green_function_mapper(frho, N, &data);
+  } else {
+    threadpool_map(tp, mesh_apply_Green_function_mapper, frho, N,
+                   sizeof(fftw_complex), 0, &data);
+  }
+
+  /* Correct singularity at (0,0,0) */
+  frho[0][0] = 0.;
+  frho[0][1] = 0.;
+}
+
 #endif
 
 /**
@@ -362,7 +493,7 @@ void mesh_to_gparts_CIC(struct gpart* gp, const double* pot, const int N,
  * @param verbose Are we talkative?
  */
 void pm_mesh_compute_potential(struct pm_mesh* mesh, const struct space* s,
-                               struct threadpool* tp, int verbose) {
+                               struct threadpool* tp, const int verbose) {
 
 #ifdef HAVE_FFTW
 
@@ -456,66 +587,8 @@ void pm_mesh_compute_potential(struct pm_mesh* mesh, const struct space* s,
 
   tic = getticks();
 
-  /* Some common factors */
-  const double green_fac = -1. / (M_PI * box_size);
-  const double a_smooth2 = 4. * M_PI * M_PI * r_s * r_s / (box_size * box_size);
-  const double k_fac = M_PI / (double)N;
-
   /* Now de-convolve the CIC kernel and apply the Green function */
-  for (int i = 0; i < N; ++i) {
-
-    /* kx component of vector in Fourier space and 1/sinc(kx) */
-    const int kx = (i > N_half ? i - N : i);
-    const double kx_d = (double)kx;
-    const double fx = k_fac * kx_d;
-    const double sinc_kx_inv = (kx != 0) ? fx / sin(fx) : 1.;
-
-    for (int j = 0; j < N; ++j) {
-
-      /* ky component of vector in Fourier space and 1/sinc(ky) */
-      const int ky = (j > N_half ? j - N : j);
-      const double ky_d = (double)ky;
-      const double fy = k_fac * ky_d;
-      const double sinc_ky_inv = (ky != 0) ? fy / sin(fy) : 1.;
-
-      for (int k = 0; k < N_half + 1; ++k) {
-
-        /* kz component of vector in Fourier space and 1/sinc(kz) */
-        const int kz = (k > N_half ? k - N : k);
-        const double kz_d = (double)kz;
-        const double fz = k_fac * kz_d;
-        const double sinc_kz_inv = (kz != 0) ? fz / (sin(fz) + FLT_MIN) : 1.;
-
-        /* Norm of vector in Fourier space */
-        const double k2 = (kx_d * kx_d + ky_d * ky_d + kz_d * kz_d);
-
-        /* Avoid FPEs... */
-        if (k2 == 0.) continue;
-
-        /* Green function */
-        double W = 1.;
-        fourier_kernel_long_grav_eval(k2 * a_smooth2, &W);
-        const double green_cor = green_fac * W / (k2 + FLT_MIN);
-
-        /* Deconvolution of CIC */
-        const double CIC_cor = sinc_kx_inv * sinc_ky_inv * sinc_kz_inv;
-        const double CIC_cor2 = CIC_cor * CIC_cor;
-        const double CIC_cor4 = CIC_cor2 * CIC_cor2;
-
-        /* Combined correction */
-        const double total_cor = green_cor * CIC_cor4;
-
-        /* Apply to the mesh */
-        const int index = N * (N_half + 1) * i + (N_half + 1) * j + k;
-        frho[index][0] *= total_cor;
-        frho[index][1] *= total_cor;
-      }
-    }
-  }
-
-  /* Correct singularity at (0,0,0) */
-  frho[0][0] = 0.;
-  frho[0][1] = 0.;
+  mesh_apply_Green_function(tp, frho, N, r_s, box_size);
 
   if (verbose)
     message("Applying Green function took %.3f %s.",
@@ -537,7 +610,7 @@ void pm_mesh_compute_potential(struct pm_mesh* mesh, const struct space* s,
   mesh->potential = rho;
 
   /* message("\n\n\n POTENTIAL"); */
-  /* print_array(potential, N); */
+  /* print_array(mesh->potential, N); */
 
   /* Clean-up the mess */
   fftw_destroy_plan(forward_plan);