Actual source code: fftw.c

  1: #define PETSCMAT_DLL

  3: /*
  4:     Provides an interface to the FFTW package
  5: */

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

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

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

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

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

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

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

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

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

116:    Collective on MPI_Comm

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

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

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

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

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

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

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

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