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