Actual source code: fftw.c

  1: #define PETSCMAT_DLL

  3: /*
  4:     Provides an interface to the FFTW package.
  5:     Testing examples can be found in ~src/mat/examples/tests
  6: */

 8:  #include private/matimpl.h
 10: #if defined(PETSC_USE_COMPLEX)
 11: #include "fftw3.h"
 12: #endif

 15: typedef struct {
 16:   PetscInt       ndim;
 17:   PetscInt       *dim;
 18:   fftw_plan      p_forward,p_backward;
 19:   unsigned       p_flag; /* planner flags, FFTW_ESTIMATE,FFTW_MEASURE, FFTW_PATIENT, FFTW_EXHAUSTIVE */
 20: } Mat_FFTW;

 24: PetscErrorCode MatMult_SeqFFTW(Mat A,Vec x,Vec y)
 25: {
 27:   Mat_FFTW       *fftw = (Mat_FFTW*)A->data;
 28:   PetscScalar    *x_array,*y_array;
 29:   PetscInt       ndim=fftw->ndim,*dim=fftw->dim;

 32:   VecGetArray(x,&x_array);
 33:   VecGetArray(y,&y_array);
 34:   if (!fftw->p_forward){ /* create a plan, then excute it */
 35:     switch (ndim){
 36:     case 1:
 37:       fftw->p_forward = fftw_plan_dft_1d(dim[0],(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_FORWARD,fftw->p_flag);
 38:       break;
 39:     case 2:
 40:       fftw->p_forward = fftw_plan_dft_2d(dim[0],dim[1],(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_FORWARD,fftw->p_flag);
 41:       break;
 42:     case 3:
 43:       fftw->p_forward = fftw_plan_dft_3d(dim[0],dim[1],dim[2],(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_FORWARD,fftw->p_flag);
 44:       break;
 45:     default:
 46:       fftw->p_forward = fftw_plan_dft(ndim,dim,(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_FORWARD,fftw->p_flag);
 47:       break;
 48:     }
 49:     fftw_execute(fftw->p_forward);
 50:   } else {
 51:     /* use existing plan */
 52:     fftw_execute_dft(fftw->p_forward,(fftw_complex*)x_array,(fftw_complex*)y_array);
 53:   }
 54:   VecRestoreArray(y,&y_array);
 55:   VecRestoreArray(x,&x_array);
 56:   return(0);
 57: }

 61: PetscErrorCode MatMultTranspose_SeqFFTW(Mat A,Vec x,Vec y)
 62: {
 64:   Mat_FFTW       *fftw = (Mat_FFTW*)A->data;
 65:   PetscScalar    *x_array,*y_array;
 66:   PetscInt       ndim=fftw->ndim,*dim=fftw->dim;

 69:   VecGetArray(x,&x_array);
 70:   VecGetArray(y,&y_array);
 71:   if (!fftw->p_backward){ /* create a plan, then excute it */
 72:     switch (ndim){
 73:     case 1:
 74:       fftw->p_backward = fftw_plan_dft_1d(dim[0],(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_BACKWARD,fftw->p_flag);
 75:       break;
 76:     case 2:
 77:       fftw->p_backward = fftw_plan_dft_2d(dim[0],dim[1],(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_BACKWARD,fftw->p_flag);
 78:       break;
 79:     case 3:
 80:       fftw->p_backward = fftw_plan_dft_3d(dim[0],dim[1],dim[2],(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_BACKWARD,fftw->p_flag);
 81:       break;
 82:     default:
 83:       fftw->p_backward = fftw_plan_dft(ndim,dim,(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_BACKWARD,fftw->p_flag);
 84:       break;
 85:     }
 86:     fftw_execute(fftw->p_backward);
 87:   } else { /* use existing plan */
 88:     fftw_execute_dft(fftw->p_backward,(fftw_complex*)x_array,(fftw_complex*)y_array);
 89:   }
 90:   VecRestoreArray(y,&y_array);
 91:   VecRestoreArray(x,&x_array);
 92:   return(0);
 93: }

 97: PetscErrorCode MatDestroy_SeqFFTW(Mat A)
 98: {
 99:   Mat_FFTW       *fftw = (Mat_FFTW*)A->data;

103:   PetscFree(fftw->dim);
104:   fftw_destroy_plan(fftw->p_forward);
105:   fftw_destroy_plan(fftw->p_backward);
106:   PetscFree(fftw);
107:   PetscObjectChangeTypeName((PetscObject)A,0);
108:   return(0);
109: }

113: /*@
114:       MatCreateSeqFFTW - Creates a matrix object that provides sequential FFT
115:   via the external package FFTW

117:    Collective on MPI_Comm

119:    Input Parameter:
120: +   comm - MPI communicator, set to PETSC_COMM_SELF
121: .   ndim - the ndim-dimensional transform
122: -   dim - array of size ndim, dim[i] contains the vector length in the i-dimension

124:    Output Parameter:
125: .   A  - the matrix

127:   Options Database Keys:
128: + -mat_fftw_plannerflags - set FFTW planner flags

130:    Level: intermediate
131:    
132: @*/
133: PetscErrorCode  MatCreateSeqFFTW(MPI_Comm comm,PetscInt ndim,const PetscInt dim[],Mat* A)
134: {
136:   Mat_FFTW       *fftw;
137:   PetscInt       m,i;
138:   const char     *p_flags[]={"FFTW_ESTIMATE","FFTW_MEASURE","FFTW_PATIENT","FFTW_EXHAUSTIVE"};
139:   PetscTruth     flg;
140:   PetscInt       p_flag;

143:   if (ndim < 0) SETERRQ1(PETSC_ERR_USER,"ndim %d must be > 0",ndim);
144:   MatCreate(comm,A);
145:   m = 1;
146:   for (i=0; i<ndim; i++){
147:     if (dim[i] < 0) SETERRQ2(PETSC_ERR_USER,"dim[%d]=%d must be > 0",i,dim[i]);
148:     m *= dim[i];
149:   }
150:   MatSetSizes(*A,m,m,m,m);
151:   PetscObjectChangeTypeName((PetscObject)*A,MATSEQFFTW);

153:   PetscNewLog(*A,Mat_FFTW,&fftw);
154:   (*A)->data = (void*)fftw;
155:   PetscMalloc((ndim+1)*sizeof(PetscInt),&fftw->dim);
156:   PetscMemcpy(fftw->dim,dim,ndim*sizeof(PetscInt));
157:   fftw->ndim       = ndim;
158:   fftw->p_forward  = 0;
159:   fftw->p_backward = 0;
160:   fftw->p_flag     = FFTW_ESTIMATE;

162:   (*A)->ops->mult          = MatMult_SeqFFTW;
163:   (*A)->ops->multtranspose = MatMultTranspose_SeqFFTW;
164:   (*A)->assembled          = PETSC_TRUE;
165:   (*A)->ops->destroy       = MatDestroy_SeqFFTW;

167:   /* get runtime options */
168:   PetscOptionsBegin(((PetscObject)(*A))->comm,((PetscObject)(*A))->prefix,"FFTW Options","Mat");
169:   PetscOptionsEList("-mat_fftw_plannerflags","Planner Flags","None",p_flags,4,p_flags[0],&p_flag,&flg);
170:   if (flg) {fftw->p_flag = (unsigned)p_flag;}
171:   PetscOptionsEnd();
172:   return(0);
173: }