fftwnd.c

FFTW, a collection of fast C routines to compute the Discrete Fourier Transform in one or more dime

	       }
	  } else {		/* use buffered transform */
	       destroy_plan_array(rank, plans_nobuf);
	  }
     }

     return p;
}

fftwnd_plan fftw2d_create_plan_specific(int nx, int ny,
					fftw_direction dir, int flags,
					fftw_complex *in, int istride,
					fftw_complex *out, int ostride)
{
     int n[2];

     n[0] = nx;
     n[1] = ny;

     return fftwnd_create_plan_specific(2, n, dir, flags,
					in, istride, out, ostride);
}

fftwnd_plan fftw3d_create_plan_specific(int nx, int ny, int nz,
					fftw_direction dir, int flags,
					fftw_complex *in, int istride,
					fftw_complex *out, int ostride)
{
     int n[3];

     n[0] = nx;
     n[1] = ny;
     n[2] = nz;

     return fftwnd_create_plan_specific(3, n, dir, flags,
					in, istride, out, ostride);
}

/* Create a generic fftwnd plan: */

fftwnd_plan fftwnd_create_plan(int rank, const int *n,
			       fftw_direction dir, int flags)
{
     return fftwnd_create_plan_specific(rank, n, dir, flags, 0, 1, 0, 1);
}

fftwnd_plan fftw2d_create_plan(int nx, int ny,
			       fftw_direction dir, int flags)
{
     return fftw2d_create_plan_specific(nx, ny, dir, flags, 0, 1, 0, 1);
}

fftwnd_plan fftw3d_create_plan(int nx, int ny, int nz,
			       fftw_direction dir, int flags)
{
     return fftw3d_create_plan_specific(nx, ny, nz, dir, flags, 0, 1, 0, 1);
}

/************************ Freeing the FFTWND Plan ************************/

static void destroy_plan_array(int rank, fftw_plan *plans)
{
     if (plans) {
	  int i, j;

	  for (i = 0; i < rank; ++i) {
	       for (j = i - 1;
		    j >= 0 && plans[i] != plans[j];
		    --j);
	       if (j < 0 && plans[i])
		    fftw_destroy_plan(plans[i]);
	  }
	  fftw_free(plans);
     }
}

void fftwnd_destroy_plan(fftwnd_plan plan)
{
     if (plan) {
	  destroy_plan_array(plan->rank, plan->plans);

	  if (plan->n)
	       fftw_free(plan->n);

	  if (plan->n_before)
	       fftw_free(plan->n_before);

	  if (plan->n_after)
	       fftw_free(plan->n_after);

	  if (plan->work)
	       fftw_free(plan->work);

	  fftw_free(plan);
     }
}

/************************ Printing the FFTWND Plan ************************/

void fftwnd_fprint_plan(FILE *f, fftwnd_plan plan)
{
     if (plan) {
	  int i, j;

	  if (plan->rank == 0) {
	       fprintf(f, "plan for rank 0 (null) transform.\n");
	       return;
	  }
	  fprintf(f, "plan for ");
	  for (i = 0; i < plan->rank; ++i)
	       fprintf(f, "%s%d", i ? "x" : "", plan->n[i]);
	  fprintf(f, " transform:\n");

	  if (plan->nbuffers > 0)
	       fprintf(f, "  -- using buffered transforms (%d buffers)\n",
		       plan->nbuffers);
	  else
	       fprintf(f, "  -- using unbuffered transform\n");

	  for (i = 0; i < plan->rank; ++i) {
	       fprintf(f, "* dimension %d (size %d) ", i, plan->n[i]);

	       for (j = i - 1; j >= 0; --j)
		    if (plan->plans[j] == plan->plans[i])
			 break;

	       if (j < 0)
		    fftw_fprint_plan(f, plan->plans[i]);
	       else
		    fprintf(f, "plan is same as dimension %d plan.\n", j);
	  }
     }
}

void fftwnd_print_plan(fftwnd_plan plan)
{
     fftwnd_fprint_plan(stdout, plan);
}

/********************* Buffered FFTW (in-place) *********************/

void fftw_buffered(fftw_plan p, int howmany,
		   fftw_complex *in, int istride, int idist,
		   fftw_complex *work,
		   int nbuffers, fftw_complex *buffers)
{
     int i = 0, n, nb;

     n = p->n;
     nb = n + FFTWND_BUFFER_PADDING;

     do {
	  for (; i <= howmany - nbuffers; i += nbuffers) {
	       fftw_complex *cur_in = in + i * idist;
	       int j, buf;

	       /* 
	        * First, copy nbuffers strided arrays to the
	        * contiguous buffer arrays (reading consecutive
	        * locations, assuming that idist is 1):
	        */
	       for (j = 0; j < n; ++j) {
		    fftw_complex *cur_in2 = cur_in + j * istride;
		    fftw_complex *cur_buffers = buffers + j;

		    for (buf = 0; buf <= nbuffers - 4; buf += 4) {
			 fftw_real r0, i0, r1, i1, r2, i2, r3, i3;
			 r0 = c_re(cur_in2[0]);
			 i0 = c_im(cur_in2[0]);
			 r1 = c_re(cur_in2[idist]);
			 i1 = c_im(cur_in2[idist]);
			 r2 = c_re(cur_in2[2 * idist]);
			 i2 = c_im(cur_in2[2 * idist]);
			 r3 = c_re(cur_in2[3 * idist]);
			 i3 = c_im(cur_in2[3 * idist]);
			 c_re(cur_buffers[0]) = r0;
			 c_im(cur_buffers[0]) = i0;
			 c_re(cur_buffers[nb]) = r1;
			 c_im(cur_buffers[nb]) = i1;
			 c_re(cur_buffers[2 * nb]) = r2;
			 c_im(cur_buffers[2 * nb]) = i2;
			 c_re(cur_buffers[3 * nb]) = r3;
			 c_im(cur_buffers[3 * nb]) = i3;
			 cur_buffers += 4 * nb;
			 cur_in2 += 4 * idist;
		    }
		    for (; buf < nbuffers; ++buf) {
			 *cur_buffers = *cur_in2;
			 cur_buffers += nb;
			 cur_in2 += idist;
		    }
	       }

	       /* 
	        * Now, compute the FFTs in the buffers (in-place
	        * using work): 
	        */
	       fftw(p, nbuffers, buffers, 1, nb, work, 1, 0);

	       /* 
	        * Finally, copy the results back from the contiguous
	        * buffers to the strided arrays (writing consecutive
	        * locations):
	        */
	       for (j = 0; j < n; ++j) {
		    fftw_complex *cur_in2 = cur_in + j * istride;
		    fftw_complex *cur_buffers = buffers + j;

		    for (buf = 0; buf <= nbuffers - 4; buf += 4) {
			 fftw_real r0, i0, r1, i1, r2, i2, r3, i3;
			 r0 = c_re(cur_buffers[0]);
			 i0 = c_im(cur_buffers[0]);
			 r1 = c_re(cur_buffers[nb]);
			 i1 = c_im(cur_buffers[nb]);
			 r2 = c_re(cur_buffers[2 * nb]);
			 i2 = c_im(cur_buffers[2 * nb]);
			 r3 = c_re(cur_buffers[3 * nb]);
			 i3 = c_im(cur_buffers[3 * nb]);
			 c_re(cur_in2[0]) = r0;
			 c_im(cur_in2[0]) = i0;
			 c_re(cur_in2[idist]) = r1;
			 c_im(cur_in2[idist]) = i1;
			 c_re(cur_in2[2 * idist]) = r2;
			 c_im(cur_in2[2 * idist]) = i2;
			 c_re(cur_in2[3 * idist]) = r3;
			 c_im(cur_in2[3 * idist]) = i3;
			 cur_buffers += 4 * nb;
			 cur_in2 += 4 * idist;
		    }
		    for (; buf < nbuffers; ++buf) {
			 *cur_in2 = *cur_buffers;
			 cur_buffers += nb;
			 cur_in2 += idist;
		    }
	       }
	  }

	  /* 
	   * we skip howmany % nbuffers ffts at the end of the loop,
	   * so we have to go back and do them: 
	   */
	  nbuffers = howmany - i;
     } while (i < howmany);
}

/********************* Computing the N-Dimensional FFT *********************/

void fftwnd_aux(fftwnd_plan p, int cur_dim,
		fftw_complex *in, int istride,
		fftw_complex *out, int ostride,
		fftw_complex *work)
{
     int n_after = p->n_after[cur_dim], n = p->n[cur_dim];

     if (cur_dim == p->rank - 2) {
	  /* just do the last dimension directly: */
	  if (p->is_in_place)
	       fftw(p->plans[p->rank - 1], n,
		    in, istride, n_after * istride,
		    work, 1, 0);
	  else
	       fftw(p->plans[p->rank - 1], n,
		    in, istride, n_after * istride,
		    out, ostride, n_after * ostride);
     } else {			/* we have at least two dimensions to go */
	  int i;

	  /* 
	   * process the subsequent dimensions recursively, in hyperslabs,
	   * to get maximum locality: 
	   */
	  for (i = 0; i < n; ++i)
	       fftwnd_aux(p, cur_dim + 1,
			  in + i * n_after * istride, istride,
			  out + i * n_after * ostride, ostride, work);
     }

     /* do the current dimension (in-place): */
     if (p->nbuffers == 0) {
	  fftw(p->plans[cur_dim], n_after,
	       out, n_after * ostride, ostride,
	       work, 1, 0);
     } else			/* using contiguous copy buffers: */
	  fftw_buffered(p->plans[cur_dim], n_after,
			out, n_after * ostride, ostride,
			work, p->nbuffers, work + n);
}

/*
 * alternate version of fftwnd_aux -- this version pushes the howmany
 * loop down to the leaves of the computation, for greater locality in
 * cases where dist < stride
 */
void fftwnd_aux_howmany(fftwnd_plan p, int cur_dim,
			int howmany,
			fftw_complex *in, int istride, int idist,
			fftw_complex *out, int ostride, int odist,
			fftw_complex *work)
{
     int n_after = p->n_after[cur_dim], n = p->n[cur_dim];
     int k;

     if (cur_dim == p->rank - 2) {
	  /* just do the last dimension directly: */
	  if (p->is_in_place)
	       for (k = 0; k < n; ++k)
		    fftw(p->plans[p->rank - 1], howmany,
			 in + k * n_after * istride, istride, idist,
			 work, 1, 0);
	  else
	       for (k = 0; k < n; ++k)
		    fftw(p->plans[p->rank - 1], howmany,
			 in + k * n_after * istride, istride, idist,
			 out + k * n_after * ostride, ostride, odist);
     } else {			/* we have at least two dimensions to go */
	  int i;

	  /* 
	   * process the subsequent dimensions recursively, in
	   * hyperslabs, to get maximum locality:
	   */
	  for (i = 0; i < n; ++i)
	       fftwnd_aux_howmany(p, cur_dim + 1, howmany,
			      in + i * n_after * istride, istride, idist,
				  out + i * n_after * ostride, ostride, odist,
				  work);
     }

     /* do the current dimension (in-place): */
     if (p->nbuffers == 0)
	  for (k = 0; k < n_after; ++k)
	       fftw(p->plans[cur_dim], howmany,
		    out + k * ostride, n_after * ostride, odist,
		    work, 1, 0);
     else			/* using contiguous copy buffers: */
	  for (k = 0; k < n_after; ++k)
	       fftw_buffered(p->plans[cur_dim], howmany,
			     out + k * ostride, n_after * ostride, odist,
			     work, p->nbuffers, work + n);
}

void fftwnd(fftwnd_plan p, int howmany,
	    fftw_complex *in, int istride, int idist,
	    fftw_complex *out, int ostride, int odist)
{
     fftw_complex *work;

#ifdef FFTW_DEBUG
     if (p->rank > 0 && (p->plans[0]->flags & FFTW_THREADSAFE)
	 && p->nwork && p->work)
	  fftw_die("bug with FFTW_THREADSAFE flag\n");
#endif

     if (p->nwork && !p->work)
	  work = (fftw_complex *) fftw_malloc(p->nwork * sizeof(fftw_complex));
     else
	  work = p->work;

     switch (p->rank) {
	 case 0:
	      break;
	 case 1:
	      if (p->is_in_place)	/* fft is in-place */
		   fftw(p->plans[0], howmany, in, istride, idist,
			work, 1, 0);
	      else
		   fftw(p->plans[0], howmany, in, istride, idist,
			out, ostride, odist);
	      break;
	 default:		/* rank >= 2 */
	      {
		   if (p->is_in_place) {
			out = in;
			ostride = istride;
			odist = idist;
		   }
		   if (howmany > 1 && odist < ostride)
			fftwnd_aux_howmany(p, 0, howmany,
					   in, istride, idist,
					   out, ostride, odist,
					   work);
		   else {
			int i;

			for (i = 0; i < howmany; ++i)
			     fftwnd_aux(p, 0,
					in + i * idist, istride,
					out + i * odist, ostride,
					work);
		   }
	      }
     }

     if (p->nwork && !p->work)
	  fftw_free(work);

}

void fftwnd_one(fftwnd_plan p, fftw_complex *in, fftw_complex *out)
{
     fftwnd(p, 1, in, 1, 1, out, 1, 1);
}