xref: /petsc/src/mat/impls/fft/fftw/fftw.c (revision b2b77a04a644fb8949a04a8ba0dae8d616b7156d) 
1 #define PETSCMAT_DLL
2 
3 /*
4     Provides an interface to the FFTW package.
5     Testing examples can be found in ~src/mat/examples/tests
6 */
7 
8 #include "../src/mat/impls/fft/fft.h"   /*I "petscmat.h" I*/
9 EXTERN_C_BEGIN
10 #include "fftw3-mpi.h"
11 EXTERN_C_END
12 
13 typedef struct {
14   fftw_plan   p_forward,p_backward;
15   unsigned    p_flag; /* planner flags, FFTW_ESTIMATE,FFTW_MEASURE, FFTW_PATIENT, FFTW_EXHAUSTIVE */
16   PetscScalar *finarray,*foutarray,*binarray,*boutarray; /* keep track of arrays becaue fftw plan should be
17                                                             executed for the arrays with which the plan was created */
18 } Mat_FFTW;
19 
20 extern PetscErrorCode MatMult_SeqFFTW(Mat,Vec,Vec);
21 extern PetscErrorCode MatMultTranspose_SeqFFTW(Mat,Vec,Vec);
22 extern PetscErrorCode MatMult_MPIFFTW(Mat,Vec,Vec);
23 extern PetscErrorCode MatMultTranspose_MPIFFTW(Mat,Vec,Vec);
24 extern PetscErrorCode MatDestroy_FFTW(Mat);
25 extern PetscErrorCode VecDestroy_MPIFFTW(Vec);
26 extern PetscErrorCode MatGetVecs_FFTW(Mat,Vec*,Vec*);
27 
28 #undef __FUNCT__
29 #define __FUNCT__ "MatMult_SeqFFTW"
30 PetscErrorCode MatMult_SeqFFTW(Mat A,Vec x,Vec y)
31 {
32   PetscErrorCode ierr;
33   Mat_FFT        *fft  = (Mat_FFT*)A->data;
34   Mat_FFTW       *fftw = (Mat_FFTW*)fft->data;
35   PetscScalar    *x_array,*y_array;
36   PetscInt       ndim=fft->ndim,*dim=fft->dim;
37 
38   PetscFunctionBegin;
39 #if !defined(PETSC_USE_COMPLEX)
40   SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_SUP,"not support for real numbers");
41 #endif
42   ierr = VecGetArray(x,&x_array);CHKERRQ(ierr);
43   ierr = VecGetArray(y,&y_array);CHKERRQ(ierr);
44   if (!fftw->p_forward){ /* create a plan, then excute it */
45     switch (ndim){
46     case 1:
47       fftw->p_forward = fftw_plan_dft_1d(dim[0],(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_FORWARD,fftw->p_flag);
48       break;
49     case 2:
50       fftw->p_forward = fftw_plan_dft_2d(dim[0],dim[1],(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_FORWARD,fftw->p_flag);
51       break;
52     case 3:
53       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);
54       break;
55     default:
56       fftw->p_forward = fftw_plan_dft(ndim,dim,(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_FORWARD,fftw->p_flag);
57       break;
58     }
59     fftw->finarray  = x_array;
60     fftw->foutarray = y_array;
61     /* Warning: if (fftw->p_flag!==FFTW_ESTIMATE) The data in the in/out arrays is overwritten!
62                 planning should be done before x is initialized! See FFTW manual sec2.1 or sec4 */
63     fftw_execute(fftw->p_forward);
64   } else { /* use existing plan */
65     if (fftw->finarray != x_array || fftw->foutarray != y_array){ /* use existing plan on new arrays */
66       fftw_execute_dft(fftw->p_forward,(fftw_complex*)x_array,(fftw_complex*)y_array);
67     } else {
68       fftw_execute(fftw->p_forward);
69     }
70   }
71   ierr = VecRestoreArray(y,&y_array);CHKERRQ(ierr);
72   ierr = VecRestoreArray(x,&x_array);CHKERRQ(ierr);
73   PetscFunctionReturn(0);
74 }
75 
76 #undef __FUNCT__
77 #define __FUNCT__ "MatMultTranspose_SeqFFTW"
78 PetscErrorCode MatMultTranspose_SeqFFTW(Mat A,Vec x,Vec y)
79 {
80   PetscErrorCode ierr;
81   Mat_FFT        *fft = (Mat_FFT*)A->data;
82   Mat_FFTW       *fftw = (Mat_FFTW*)fft->data;
83   PetscScalar    *x_array,*y_array;
84   PetscInt       ndim=fft->ndim,*dim=fft->dim;
85 
86   PetscFunctionBegin;
87 #if !defined(PETSC_USE_COMPLEX)
88   SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_SUP,"not support for real numbers");
89 #endif
90   ierr = VecGetArray(x,&x_array);CHKERRQ(ierr);
91   ierr = VecGetArray(y,&y_array);CHKERRQ(ierr);
92   if (!fftw->p_backward){ /* create a plan, then excute it */
93     switch (ndim){
94     case 1:
95       fftw->p_backward = fftw_plan_dft_1d(dim[0],(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_BACKWARD,fftw->p_flag);
96       break;
97     case 2:
98       fftw->p_backward = fftw_plan_dft_2d(dim[0],dim[1],(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_BACKWARD,fftw->p_flag);
99       break;
100     case 3:
101       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);
102       break;
103     default:
104       fftw->p_backward = fftw_plan_dft(ndim,dim,(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_BACKWARD,fftw->p_flag);
105       break;
106     }
107     fftw->binarray  = x_array;
108     fftw->boutarray = y_array;
109     fftw_execute(fftw->p_backward);CHKERRQ(ierr);
110   } else { /* use existing plan */
111     if (fftw->binarray != x_array || fftw->boutarray != y_array){ /* use existing plan on new arrays */
112       fftw_execute_dft(fftw->p_backward,(fftw_complex*)x_array,(fftw_complex*)y_array);
113     } else {
114       fftw_execute(fftw->p_backward);CHKERRQ(ierr);
115     }
116   }
117   ierr = VecRestoreArray(y,&y_array);CHKERRQ(ierr);
118   ierr = VecRestoreArray(x,&x_array);CHKERRQ(ierr);
119   PetscFunctionReturn(0);
120 }
121 
122 #undef __FUNCT__
123 #define __FUNCT__ "MatMult_MPIFFTW"
124 PetscErrorCode MatMult_MPIFFTW(Mat A,Vec x,Vec y)
125 {
126   PetscErrorCode ierr;
127   Mat_FFT        *fft  = (Mat_FFT*)A->data;
128   Mat_FFTW       *fftw = (Mat_FFTW*)fft->data;
129   PetscScalar    *x_array,*y_array;
130   PetscInt       ndim=fft->ndim,*dim=fft->dim;
131   MPI_Comm       comm=((PetscObject)A)->comm;
132 
133   PetscFunctionBegin;
134 #if !defined(PETSC_USE_COMPLEX)
135   SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_SUP,"not support for real numbers");
136 #endif
137   ierr = VecGetArray(x,&x_array);CHKERRQ(ierr);
138   ierr = VecGetArray(y,&y_array);CHKERRQ(ierr);
139   if (!fftw->p_forward){ /* create a plan, then excute it */
140     switch (ndim){
141     case 1:
142       fftw->p_forward = fftw_mpi_plan_dft_1d(dim[0],(fftw_complex*)x_array,(fftw_complex*)y_array,comm,FFTW_FORWARD,fftw->p_flag);
143       break;
144     case 2:
145       fftw->p_forward = fftw_mpi_plan_dft_2d(dim[0],dim[1],(fftw_complex*)x_array,(fftw_complex*)y_array,comm,FFTW_FORWARD,fftw->p_flag);
146       break;
147     case 3:
148       fftw->p_forward = fftw_mpi_plan_dft_3d(dim[0],dim[1],dim[2],(fftw_complex*)x_array,(fftw_complex*)y_array,comm,FFTW_FORWARD,fftw->p_flag);
149       break;
150     default:
151       /*
152        fftw->p_forward = fftw_mpi_plan_dft(ndim,dim,(fftw_complex*)x_array,(fftw_complex*)y_array,comm,FFTW_FORWARD,fftw->p_flag);
153        */
154       SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_SUP,"Not supported yet");
155       break;
156     }
157     fftw->finarray  = x_array;
158     fftw->foutarray = y_array;
159     /* Warning: if (fftw->p_flag!==FFTW_ESTIMATE) The data in the in/out arrays is overwritten!
160                 planning should be done before x is initialized! See FFTW manual sec2.1 or sec4 */
161     fftw_execute(fftw->p_forward);
162   } else { /* use existing plan */
163     if (fftw->finarray != x_array || fftw->foutarray != y_array){ /* use existing plan on new arrays */
164       fftw_execute_dft(fftw->p_forward,(fftw_complex*)x_array,(fftw_complex*)y_array);
165     } else {
166       fftw_execute(fftw->p_forward);
167     }
168   }
169   ierr = VecRestoreArray(y,&y_array);CHKERRQ(ierr);
170   ierr = VecRestoreArray(x,&x_array);CHKERRQ(ierr);
171   PetscFunctionReturn(0);
172 }
173 
174 #undef __FUNCT__
175 #define __FUNCT__ "MatMultTranspose_MPIFFTW"
176 PetscErrorCode MatMultTranspose_MPIFFTW(Mat A,Vec x,Vec y)
177 {
178   PetscErrorCode ierr;
179   Mat_FFT        *fft  = (Mat_FFT*)A->data;
180   Mat_FFTW       *fftw = (Mat_FFTW*)fft->data;
181   PetscScalar    *x_array,*y_array;
182   PetscInt       ndim=fft->ndim,*dim=fft->dim;
183   MPI_Comm       comm=((PetscObject)A)->comm;
184 
185   PetscFunctionBegin;
186 #if !defined(PETSC_USE_COMPLEX)
187   SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_SUP,"not support for real numbers");
188 #endif
189   ierr = VecGetArray(x,&x_array);CHKERRQ(ierr);
190   ierr = VecGetArray(y,&y_array);CHKERRQ(ierr);
191   if (!fftw->p_backward){ /* create a plan, then excute it */
192     switch (ndim){
193     case 1:
194       fftw->p_backward = fftw_mpi_plan_dft_1d(dim[0],(fftw_complex*)x_array,(fftw_complex*)y_array,comm,FFTW_BACKWARD,fftw->p_flag);
195       break;
196     case 2:
197       fftw->p_backward = fftw_mpi_plan_dft_2d(dim[0],dim[1],(fftw_complex*)x_array,(fftw_complex*)y_array,comm,FFTW_BACKWARD,fftw->p_flag);
198       break;
199     case 3:
200       fftw->p_backward = fftw_mpi_plan_dft_3d(dim[0],dim[1],dim[2],(fftw_complex*)x_array,(fftw_complex*)y_array,comm,FFTW_BACKWARD,fftw->p_flag);
201       break;
202     default:
203       SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_SUP,"Not supported yet");
204       /* fftw->p_backward = fftw_mpi_plan_dft(ndim,dim,(fftw_complex*)x_array,(fftw_complex*)y_array,FFTW_BACKWARD,fftw->p_flag); */
205       break;
206     }
207     fftw->binarray  = x_array;
208     fftw->boutarray = y_array;
209     fftw_execute(fftw->p_backward);CHKERRQ(ierr);
210   } else { /* use existing plan */
211     if (fftw->binarray != x_array || fftw->boutarray != y_array){ /* use existing plan on new arrays */
212       fftw_execute_dft(fftw->p_backward,(fftw_complex*)x_array,(fftw_complex*)y_array);
213     } else {
214       fftw_execute(fftw->p_backward);CHKERRQ(ierr);
215     }
216   }
217   ierr = VecRestoreArray(y,&y_array);CHKERRQ(ierr);
218   ierr = VecRestoreArray(x,&x_array);CHKERRQ(ierr);
219   PetscFunctionReturn(0);
220 }
221 
222 #undef __FUNCT__
223 #define __FUNCT__ "MatDestroy_FFTW"
224 PetscErrorCode MatDestroy_FFTW(Mat A)
225 {
226   Mat_FFT        *fft = (Mat_FFT*)A->data;
227   Mat_FFTW       *fftw = (Mat_FFTW*)fft->data;
228   PetscErrorCode ierr;
229 
230   PetscFunctionBegin;
231 #if !defined(PETSC_USE_COMPLEX)
232   SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_SUP,"not support for real numbers");
233 #endif
234   fftw_destroy_plan(fftw->p_forward);
235   fftw_destroy_plan(fftw->p_backward);
236   ierr = PetscFree(fftw);CHKERRQ(ierr);
237   PetscFunctionReturn(0);
238 }
239 
240 #include "../src/vec/vec/impls/mpi/pvecimpl.h"   /*I  "petscvec.h"   I*/
241 #undef __FUNCT__
242 #define __FUNCT__ "VecDestroy_MPIFFTW"
243 PetscErrorCode VecDestroy_MPIFFTW(Vec v)
244 {
245   PetscErrorCode  ierr;
246   PetscScalar     *array;
247 
248   PetscFunctionBegin;
249 #if !defined(PETSC_USE_COMPLEX)
250   SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_SUP,"not support for real numbers");
251 #endif
252   ierr = VecGetArray(v,&array);CHKERRQ(ierr);
253   fftw_free((fftw_complex*)array);CHKERRQ(ierr);
254   ierr = VecRestoreArray(v,&array);CHKERRQ(ierr);
255   ierr = VecDestroy_MPI(v);CHKERRQ(ierr);
256   PetscFunctionReturn(0);
257 }
258 
259 #undef __FUNCT__
260 #define __FUNCT__ "MatGetVecs_FFTW"
261 /*
262    MatGetVecs_FFTW - Get vector(s) compatible with the matrix, i.e. with the
263      parallel layout determined by FFTW
264 
265    Collective on Mat
266 
267    Input Parameter:
268 .  mat - the matrix
269 
270    Output Parameter:
271 +   fin - (optional) input vector of forward FFTW
272 -   fout - (optional) output vector of forward FFTW
273 
274   Level: advanced
275 
276 .seealso: MatCreateFFTW()
277 */
278 PetscErrorCode  MatGetVecs_FFTW(Mat A,Vec *fin,Vec *fout)
279 {
280   PetscErrorCode ierr;
281   PetscMPIInt    size,rank;
282   MPI_Comm       comm=((PetscObject)A)->comm;
283   Mat_FFT        *fft = (Mat_FFT*)A->data;
284   PetscInt       N=fft->N;
285 
286   PetscFunctionBegin;
287 #if !defined(PETSC_USE_COMPLEX)
288   SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_SUP,"not support for real numbers");
289 #endif
290   PetscValidHeaderSpecific(A,MAT_CLASSID,1);
291   PetscValidType(A,1);
292 
293   ierr = MPI_Comm_size(comm, &size);CHKERRQ(ierr);
294   ierr = MPI_Comm_rank(comm, &rank);CHKERRQ(ierr);
295   if (size == 1){ /* sequential case */
296     if (fin) {ierr = VecCreateSeq(PETSC_COMM_SELF,N,fin);CHKERRQ(ierr);}
297     if (fout){ierr = VecCreateSeq(PETSC_COMM_SELF,N,fout);CHKERRQ(ierr);}
298   } else {        /* mpi case */
299     ptrdiff_t      alloc_local,local_n0,local_0_start;
300     PetscInt       ndim=fft->ndim,*dim=fft->dim,n=fft->n;
301     fftw_complex   *data_fin,*data_fout;
302 
303     switch (ndim){
304     case 1:
305       SETERRQ(((PetscObject)A)->comm,PETSC_ERR_SUP,"Not supported yet");
306       break;
307     case 2:
308       /* Get local size */
309       alloc_local = fftw_mpi_local_size_2d(dim[0],dim[1],comm,&local_n0,&local_0_start);
310       if (fin) {
311         data_fin  = (fftw_complex*)fftw_malloc(sizeof(fftw_complex)*alloc_local);
312         ierr = VecCreateMPIWithArray(comm,n,N,(const PetscScalar*)data_fin,fin);CHKERRQ(ierr);
313         (*fin)->ops->destroy   = VecDestroy_MPIFFTW;
314       }
315       if (fout) {
316         data_fout = (fftw_complex*)fftw_malloc(sizeof(fftw_complex)*alloc_local);
317         ierr = VecCreateMPIWithArray(comm,n,N,(const PetscScalar*)data_fout,fout);CHKERRQ(ierr);
318         (*fout)->ops->destroy   = VecDestroy_MPIFFTW;
319       }
320       break;
321     case 3:
322       SETERRQ(((PetscObject)A)->comm,PETSC_ERR_SUP,"Not supported yet");
323       break;
324     default:
325       SETERRQ(((PetscObject)A)->comm,PETSC_ERR_SUP,"Not supported yet");
326       break;
327     }
328   }
329   PetscFunctionReturn(0);
330 }
331 
332 EXTERN_C_BEGIN
333 #undef __FUNCT__
334 #define __FUNCT__ "MatCreate_FFTW"
335 /*
336       MatCreate_FFTW - Creates a matrix object that provides FFT
337   via the external package FFTW
338 
339   Options Database Keys:
340 + -mat_fftw_plannerflags - set FFTW planner flags
341 
342    Level: intermediate
343 
344 */
345 PetscErrorCode MatCreate_FFTW(Mat A)
346 {
347   PetscErrorCode ierr;
348   MPI_Comm       comm=((PetscObject)A)->comm;
349   Mat_FFT        *fft=(Mat_FFT*)A->data;
350   Mat_FFTW       *fftw;
351   PetscInt       n=fft->n,N=fft->N,ndim=fft->ndim,*dim = fft->dim;
352   const char     *p_flags[]={"FFTW_ESTIMATE","FFTW_MEASURE","FFTW_PATIENT","FFTW_EXHAUSTIVE"};
353   PetscBool      flg;
354   PetscInt       p_flag;
355   PetscMPIInt    size;
356 
357   PetscFunctionBegin;
358 #if !defined(PETSC_USE_COMPLEX)
359   SETERRQ(comm,PETSC_ERR_SUP,"not support for real numbers");
360 #endif
361 
362   ierr = MPI_Comm_size(comm, &size);CHKERRQ(ierr);
363   //ierr = MPI_Comm_rank(comm, &rank);CHKERRQ(ierr);
364 
365 
366   if (size == 1) {
367     ierr = MatSetSizes(A,N,N,N,N);CHKERRQ(ierr);
368     n = N;
369   } else {
370     ptrdiff_t alloc_local,local_n0,local_0_start;
371     switch (ndim){
372     case 1:
373       SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_SUP,"not implemented yet");
374       break;
375     case 2:
376       alloc_local = fftw_mpi_local_size_2d(dim[0],dim[1],comm,&local_n0,&local_0_start);
377       /*
378        PetscMPIInt    rank;
379        PetscSynchronizedPrintf(comm,"[%d] MatCreateSeqFFTW: local_n0, local_0_start %d %d, N %d,dim %d, %d\n",rank,(PetscInt)local_n0*dim[1],(PetscInt)local_0_start,m,dim[0],dim[1]);
380        PetscSynchronizedFlush(comm);
381        */
382       n = (PetscInt)local_n0*dim[1];
383       ierr = MatSetSizes(A,n,n,N,N);CHKERRQ(ierr);
384       break;
385     case 3:
386       SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_SUP,"not implemented yet");
387       break;
388     default:
389       SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_SUP,"not implemented yet");
390       break;
391     }
392   }
393   ierr = PetscObjectChangeTypeName((PetscObject)A,MATFFTW);CHKERRQ(ierr);
394 
395   ierr = PetscNewLog(A,Mat_FFTW,&fftw);CHKERRQ(ierr);
396   fft->data = (void*)fftw;
397 
398   fft->n           = n;
399   fftw->p_forward  = 0;
400   fftw->p_backward = 0;
401   fftw->p_flag     = FFTW_ESTIMATE;
402 
403   if (size == 1){
404     A->ops->mult          = MatMult_SeqFFTW;
405     A->ops->multtranspose = MatMultTranspose_SeqFFTW;
406   } else {
407     A->ops->mult          = MatMult_MPIFFTW;
408     A->ops->multtranspose = MatMultTranspose_MPIFFTW;
409   }
410   fft->matdestroy          = MatDestroy_FFTW;
411   A->ops->getvecs       = MatGetVecs_FFTW;
412   A->assembled          = PETSC_TRUE;
413 
414   /* get runtime options */
415   ierr = PetscOptionsBegin(((PetscObject)A)->comm,((PetscObject)A)->prefix,"FFTW Options","Mat");CHKERRQ(ierr);
416     ierr = PetscOptionsEList("-mat_fftw_plannerflags","Planner Flags","None",p_flags,4,p_flags[0],&p_flag,&flg);CHKERRQ(ierr);
417     if (flg) {fftw->p_flag = (unsigned)p_flag;}
418   PetscOptionsEnd();
419   PetscFunctionReturn(0);
420 }
421 EXTERN_C_END
422