NERSCPowering Scientific Discovery Since 1974

FFTW2 MPI 3D Example Fortran Source

implicit none

      include "mpif.h"

      integer n, nx, ny, nz, nit
      parameter (n=8, nx=8, ny=8, nz=8, nit=2)
      integer lnx, lxs, lnyt, lyst, lsize

      integer fftw_fwd, fftw_bkwd, fftw_est, fftw_norm_order
      parameter (fftw_fwd=-1, fftw_bkwd=1, fftw_est=0,
     .           fftw_norm_order=0)
      integer*8 plan, iplan
      double complex data, work
      dimension data(nx,ny,nz), work(nx,ny,nz)

      integer ierr, rank
      integer i,j,k,it
c
      call mpi_init(ierr)
      call mpi_comm_rank(mpi_comm_world,rank,ierr)

c
      print *,"PE, nx, ny, nz = ",rank, nx, ny, nz

      do i=1,n
         do j=1,n
            do k=1,n
               data(i,j,k) = (1.0, 0.0)
               if ((i.lt.2).and.(j.lt.2).and.(k.lt.2)) print *,
     .          rank," input: ",i,j,k,data(i,j,k)
            enddo
         enddo
      enddo

      print *,"fftw_fwd =         ",fftw_fwd
      print *,"fftw_bkwd =        ",fftw_bkwd
      print *,"fftw_est =         ",fftw_est
      print *,"fftw_norm_order =  ",fftw_norm_order

      call fftw3d_f77_mpi_create_plan(
     .      plan, mpi_comm_world, nx, ny, nz,
     .      fftw_fwd, fftw_est)

      call fftw3d_f77_mpi_create_plan(
     .      iplan, mpi_comm_world, nx, ny, nz,
     .      fftw_bkwd, fftw_est)

      print *,"address of fwd plan =   ",plan
      print *,"address of bckwd plan = ",iplan

c   slab decomposition

      call fftwnd_f77_mpi_local_sizes(
     .      plan, lnx, lxs, lnyt, lyst, lsize)

      print *,"data_decomp: ",rank, "xslab ",lxs, " -> ",(lxs+lnx),
     . lnx, "(", lyst, ") "

      call mpi_barrier(mpi_comm_world, ierr)
      if (rank.eq.0) print *,"number of iterations: ",nit

      call fftwnd_f77_mpi(plan, 1, data, work, 0, fftw_norm_order)

      do i=1,2
         do j=1,2
            do k=1,2
               print *,rank," forward: ",i,j,k,data(i,j,k)
            enddo
         enddo
      enddo

      call mpi_barrier (mpi_comm_world, ierr)
      call fftwnd_f77_mpi(iplan, 1, data, work, 0, fftw_norm_order)

      do i=1,2
         do j=1,2
            do k=1,2
               print *,rank," reverse: ",i,j,k,data(i,j,k)
            enddo
         enddo
      enddo

      call mpi_barrier(mpi_comm_world, ierr)
      print *,"Wrap up PE ",rank
      call fftwnd_f77_mpi_destroy_plan(plan)
      call fftwnd_f77_mpi_destroy_plan(iplan)

      call mpi_finalize(ierr)

      end