fft_work.fpp Source File


Contents

Source Code


Source Code

# include "define.inc"

#if FFT == _FFTW_
#error Support for FFTW2 has been removed, please switch to FFTW3
#endif

!> Provides some wrappers for managing FFTW plans.
!>
!> The initialisation routines ([[init_crfftw]], [[init_rcfftw]],
!> [[init_ccfftw]]) create FFTW plans, storing them in an instance of
!> [[fft_type]], along with some useful metadata, notably
!> [[fft_type:scale]].
!>
!> The preprocessor macro `FFTW_PREFIX` is provided to enable calling
!> the single or double precision FFTW routines.
module fft_work

! Modifications for using FFTW version 3:
! (c) The Numerical Algorithms Group (NAG) Ltd, 2009
!                                 on behalf of the HECToR project

  use constants, only: kind_id
  implicit none

  private

  public :: measure_plan
  public :: fft_type, delete_fft
  public :: init_ccfftw, init_crfftw, init_rcfftw, init_z
  public :: FFT_TO_SPECTRAL_SPACE, FFT_TO_REAL_SPACE
  public :: save_wisdom, load_wisdom
  public :: finish_fft_work, time_fft

  interface init_crfftw
     module procedure init_crfftw_1d
     module procedure init_crfftw_2d
  end interface

  interface init_rcfftw
     module procedure init_rcfftw_1d
     module procedure init_rcfftw_2d
  end interface

  !> Constant for forward transforms, from real space to spectral space
  integer, parameter :: FFT_TO_SPECTRAL_SPACE = -1
  !> Constant for inverse transforms, from spectral space to real space
  integer, parameter :: FFT_TO_REAL_SPACE = 1


  !> If true, then use FFTW_PATIENT to pick the optimal FFT algorithm,
  !> otherwise use FFTW_ESTIMATE (potentially slower but more
  !> reproducible results). FFTW_PATIENT takes longer to set up the
  !> plan initially, but may be faster overall.
  logical :: measure_plan=.true.

  !> Information about a planned FFT transform
  type :: fft_type
     !> True if the plan has been created
     logical :: created = .false.
     !> Number of points in the array to be transformed
     integer :: n
     !> Direction of the FFT, can be either [[FFT_FORWARD]] or [[FFTW_BACKWARD]]
     integer :: direction
     !> The FFTW plan object
     integer (kind_id) :: plan
     !> How many independent slices in the array to be transformed
     integer :: howmany
     !> True if the array is strided or 2D
     logical :: strided
#ifdef SHMEM
     !> start indices for rank partition in first 2 dims
     integer ts1, ts2
#endif
     !> Normalisation factor
     real :: scale
     !> If true, then the result of the FFT needs to be multiplied by
     !> [[fft_type:scale]]. If false, then the scaling can be skipped,
     !> saving some work
     logical :: must_scale
  end type fft_type

#if FFT == _FFTW3_
# ifdef ASLFFTW3
  include 'aslfftw3.f'
# else
  include 'fftw3.f'
# endif
#endif

! the token concatenation operator ## requires the preprocessor to
! support ANSI_CPP. Also, it cannot be being run in traditional mode
#ifdef ANSI_CPP

#ifdef SINGLE_PRECISION
#define FFTW_PREFIX(fn) sfftw##fn
#else
#define FFTW_PREFIX(fn) dfftw##fn
#endif

#else

#ifdef SINGLE_PRECISION
#define FFTW_PREFIX(fn) sfftw/**/fn
#else
#define FFTW_PREFIX(fn) dfftw/**/fn
#endif

#endif

  !> Total time taken setting up and performing FFTs
  real :: time_fft(2) = 0.

  !> True if this module has been initialised
  logical :: initialised = .false.
contains

  !> Perform any one off setup of the fft library.
  !> This is currently mostly only needed for builds with
  !> OpenMP for threaded transforms.
  subroutine init_fft_work
#if FFT == _FFTW3_
#ifdef OPENMP
    use omp_lib, only: omp_get_num_threads
    implicit none
    integer :: ierr, nthreads
    if (initialised) return
    initialised = .true.
    !$OMP PARALLEL
    nthreads = omp_get_num_threads()
    !$OMP END PARALLEL
    call FFTW_PREFIX(_init_threads)(ierr)
    call FFTW_PREFIX(_plan_with_nthreads)(nthreads)
#endif
#endif
  end subroutine init_fft_work

  !> Does some common initialisation for [[fft_type]]
  type(fft_type) function basic_init(n, direction, howmany, strided)
    use mp, only: mp_abort
    use print_utils, only: to_string
    !> Number of points
    integer, intent(in) :: n
    !> Direction of FFT, must be either [[FFT_TO_REAL_SPACE]] or [[FFT_TO_SPECTRAL_SPACE]]
    integer, intent(in) :: direction
    !> Number of independent FFTs
    integer, intent(in) :: howmany
    !> Is array strided/2D
    logical, intent(in) :: strided

    if (.not. initialised) call init_fft_work

    basic_init%n = n
    basic_init%direction = direction
    if (direction == FFT_TO_SPECTRAL_SPACE) then
      basic_init%scale = 1. / real(n)
      basic_init%must_scale = .true.
    else if (direction == FFT_TO_REAL_SPACE) then
      basic_init%scale = 1.
      basic_init%must_scale = .false.
    else
      call mp_abort(&
           "Unexpected FFT direction, expected one of FFT_TO_REAL_SPACE, FFT_TO_SPECTRAL_SPACE, got " &
           // to_string(direction))
    end if
    basic_init%howmany = howmany
    basic_init%strided = strided
    basic_init%created = .true.
  end function basic_init

# if FFT == _FFTW3_
  !> Set flags for FFTW3
  integer function fftw_flags()
    if (measure_plan) then
       fftw_flags = FFTW_PATIENT + FFTW_UNALIGNED
    else
       fftw_flags = FFTW_ESTIMATE + FFTW_UNALIGNED
    endif
  end function fftw_flags
#endif

  !> FIXME : Add documentation
  subroutine init_z (fft, direction, n, howmany)
    use job_manage, only: time_message
    implicit none
    type (fft_type), intent (out) :: fft
    integer, intent (in) :: direction, n
    integer, intent (in) :: howmany
# if FFT == _FFTW3_
    complex, dimension (:,:), allocatable :: dummy_in_data, dummy_out_data
    integer, dimension (1) :: array_n, embed
# endif
    call time_message(.false., time_fft, ' FFT')

    fft = basic_init(n, direction, howmany, .false.)

#if FFT == _FFTW3_
    array_n = n
    embed   = n+1

    allocate (dummy_in_data(n+1, howmany), dummy_out_data(n+1, howmany))

    call FFTW_PREFIX(_plan_many_dft)(fft%plan, 1, array_n, howmany, &
         dummy_in_data,  embed, 1, n+1, &
         dummy_out_data, embed, 1, n+1, &
         direction, fftw_flags())

    deallocate (dummy_in_data, dummy_out_data)
# endif
    call time_message(.false., time_fft, ' FFT')
  end subroutine init_z

  !> Create an FFTW plan for complex-to-complex transforms
  subroutine init_ccfftw (fft, direction, n, howmany, data_array, transpose, m)
    use mp, only: mp_abort
    use job_manage, only: time_message
    use optionals, only: get_option_with_default
    implicit none
    !> The created plan
    type (fft_type), intent (out) :: fft
    !> Direction of the transformation, either [[FFT_TO_SPECTRAL_SPACE]] or [[FFT_TO_REAL_SPACE]]
    integer, intent (in) :: direction
    !> Number of points in array to transform
    integer, intent (in) :: n
    !> Number of independent FFTs in array
    integer, intent (in) :: howmany
    !> Example array to transform
    complex, dimension(:,:), intent(inout) :: data_array
    !> If true, then take FFT of transpose of array
    logical, optional, intent(in) :: transpose
    !> Size of transpose direction. Required if `transpose` is true
    integer, optional, intent(in) :: m
# if FFT == _FFTW3_
    integer, dimension(1) :: array_n
# endif
    call time_message(.false., time_fft, ' FFT')
    fft = basic_init(n, direction, howmany, .false.)
#if FFT == _FFTW3_
    ! the planer expects this as an array of size 1
    array_n = n

    if (get_option_with_default(transpose, .false.)) then
       if (.not. present(m)) call mp_abort("Missing argument 'm' for transpose FFT")
       call FFTW_PREFIX(_plan_many_dft)(fft%plan, 1, array_n, howmany, &
            data_array, array_n, howmany*m, 1, &
            data_array, array_n, howmany*m, 1, &
            direction, fftw_flags())
    else
       call FFTW_PREFIX(_plan_many_dft)(fft%plan, 1, array_n, howmany, &
            data_array, array_n, 1, n, &
            data_array, array_n, 1, n, &
            direction, fftw_flags())
    endif
# endif
    call time_message(.false., time_fft, ' FFT')
  end subroutine init_ccfftw

  !> Create FFTW plan suitable for 1D real to complex transformation
  subroutine init_rcfftw_1d (fft, direction, n, howmany, transposed)
    use job_manage, only: time_message
    implicit none
    !> The created plan
    type (fft_type), intent (out) :: fft
    !> Direction of the transformation, either [[FFT_TO_SPECTRAL_SPACE]] or [[FFT_TO_REAL_SPACE]]
    integer, intent (in) :: direction
    !> Number of points in array to transform
    integer, intent (in) :: n
    !> Number of independent FFTs in array
    integer, intent (in) :: howmany
    !> If true, then take FFT of transpose of array
    logical, optional, intent(in) :: transposed

#if FFT == _FFTW3_
    integer, dimension(1) :: vector_sizes_real, vector_sizes_complex

    ! we need two dummy arrays for the planner.  Since at
    ! present we are not using SSE instructions, these are save
    ! to be local to this routine
    real, dimension(:,:), allocatable :: dummy_real_data
    complex, dimension (:,:), allocatable :: dummy_complex_data
#endif

    call time_message(.false., time_fft, ' FFT')

    fft = basic_init(n, direction, howmany, .false.)

#if FFT == _FFTW3_
    if (present(transposed)) then
       allocate (dummy_real_data(max (1, howmany), N))
       allocate (dummy_complex_data(max(1, howmany), N/2+1))
    else
       allocate (dummy_real_data(N, max (1, howmany)))
       allocate (dummy_complex_data(N/2+1, max(1, howmany)))
    endif

    vector_sizes_real = N
    vector_sizes_complex = N/2+1

    if ( present(transposed)) then
       call FFTW_PREFIX(_plan_many_dft_r2c) (fft%plan, 1, &
            vector_sizes_real, howmany, &
            dummy_real_data,    vector_sizes_real,   howmany,  1, &
            dummy_complex_data, vector_sizes_complex, howmany, 1, &
            fftw_flags())
    else
       call FFTW_PREFIX(_plan_many_dft_r2c) (fft%plan, 1, &
            vector_sizes_real, howmany, &
            dummy_real_data,    vector_sizes_real,    1, N,     &
            dummy_complex_data, vector_sizes_complex, 1, N/2+1, &
            fftw_flags())
       endif

    deallocate (dummy_real_data, dummy_complex_data)
#endif
    call time_message(.false., time_fft, ' FFT')
  end subroutine init_rcfftw_1d

  !> Create FFTW plan suitable for 1D complex to real transformation
  subroutine init_crfftw_1d (fft, direction, n, howmany, transposed)
    use job_manage, only: time_message
    implicit none
    !> The created plan
    type (fft_type), intent (out) :: fft
    !> Direction of the transformation, either [[FFT_TO_SPECTRAL_SPACE]] or [[FFT_TO_REAL_SPACE]]
    integer, intent (in) :: direction
    !> Number of points in array to transform
    integer, intent (in) :: n
    !> Number of independent FFTs in array
    integer, intent (in) :: howmany
    !> If true, then take FFT of transpose of array
    logical, optional, intent(in) :: transposed

#if FFT == _FFTW3_
    integer, dimension(1) :: vector_sizes_real, vector_sizes_complex

    ! we need two dummy arrays for the planner.  Since at
    ! present we are not using SSE instructions, these are save
    ! to be local to this routine
    real, dimension(:,:), allocatable :: dummy_real_data
    complex, dimension (:,:), allocatable :: dummy_complex_data
#endif

    call time_message(.false., time_fft, ' FFT')
    fft = basic_init(n, direction, howmany, .false.)

#if FFT == _FFTW3_
    if (present(transposed)) then
       allocate (dummy_real_data(max (1, howmany), N))
       allocate (dummy_complex_data(max(1, howmany), N/2+1))
    else
       allocate (dummy_real_data(N, max (1, howmany)))
       allocate (dummy_complex_data(N/2+1, max(1, howmany)))
    endif

    vector_sizes_real = N
    vector_sizes_complex = N/2+1

    if (present(transposed)) then
       call FFTW_PREFIX(_plan_many_dft_c2r) (fft%plan, 1, &
            vector_sizes_real, howmany, &
            dummy_complex_data, vector_sizes_complex, howmany,1, &
            dummy_real_data,    vector_sizes_real,    howmany,1, &
            fftw_flags())
    else
       call FFTW_PREFIX(_plan_many_dft_c2r) (fft%plan, 1, &
            vector_sizes_real, howmany, &
            dummy_complex_data, vector_sizes_complex, 1, N/2+1, &
            dummy_real_data,    vector_sizes_real,    1, N,     &
            fftw_flags())
    endif

    deallocate (dummy_real_data, dummy_complex_data)
#endif

    call time_message(.false., time_fft, ' FFT')
  end subroutine init_crfftw_1d

  !> Create FFTW plan suitable for 2D real to complex transformation
  subroutine init_rcfftw_2d (fft, direction, m, n, howmany, stride_)
    use job_manage, only: time_message
    implicit none
    !> The created plan
    type (fft_type), intent (out) :: fft
    !> Direction of the transformation, either [[FFT_TO_SPECTRAL_SPACE]] or [[FFT_TO_REAL_SPACE]]
    integer, intent (in) :: direction
    !> Number of points in first dimension of array to transform
    integer, intent (in) :: m
    !> Number of points in second dimension of array to transform
    integer, intent (in) :: n
    !> Number of independent FFTs in array
    integer, intent (in) :: howmany
    !> Stride of array if different from `howmany`
    integer, optional, intent(in) :: stride_

#if FFT == _FFTW3_
    integer, parameter :: fft_dimension = 2
    integer  :: stride

    ! to pass to FFTW3, we need vector for actual sizes of the real
    ! and complex data
    integer, dimension(fft_dimension) :: &
         vector_sizes_real, vector_sizes_complex


    ! we need two dummy arrays for the planner.  Since at
    ! present we are not using SSE instructions, these are save
    ! to be local to this routine
    real, dimension(:,:), allocatable :: dummy_real_data
    complex, dimension (:,:), allocatable :: dummy_complex_data
#endif
    call time_message(.false., time_fft, ' FFT')
    fft = basic_init(m * n, direction, howmany, .true.)

#if FFT == _FFTW3_
    stride  = howmany
    if(present(stride_)) stride=stride_

    allocate (dummy_real_data(stride, m*n))
    allocate (dummy_complex_data(stride, (m/2+1)*n))

    vector_sizes_real(1) = m
    vector_sizes_real(2) = n

    vector_sizes_complex(1) = m/2+1
    vector_sizes_complex(2) = n

    call FFTW_PREFIX(_plan_many_dft_r2c)(fft%plan, fft_dimension, &
         vector_sizes_real, howmany, &
         dummy_real_data, vector_sizes_real, stride, 1, &
         dummy_complex_data, vector_sizes_complex, stride, 1, &
         fftw_flags())

    deallocate (dummy_real_data, dummy_complex_data)
#endif
    call time_message(.false., time_fft, ' FFT')
  end subroutine init_rcfftw_2d

  !> Create FFTW plan suitable for 2D complex to real transformation
  subroutine init_crfftw_2d (fft, direction, m, n, howmany, stride_)
    use job_manage, only: time_message
    implicit none
    !> The created plan
    type (fft_type), intent (out) :: fft
    !> Direction of the transformation, either [[FFT_TO_SPECTRAL_SPACE]] or [[FFT_TO_REAL_SPACE]]
    integer, intent (in) :: direction
    !> Number of points in first dimension of array to transform
    integer, intent (in) :: m
    !> Number of points in second dimension of array to transform
    integer, intent (in) :: n
    !> Number of independent FFTs in array
    integer, intent (in) :: howmany
    !> Stride of array if different from `howmany`
    integer, optional, intent(in) :: stride_

#if FFT == _FFTW3_
    integer, parameter :: fft_dimension = 2
    integer  :: stride

    ! to pass to FFTW3, we need vector for actual sizes of the real
    ! and complex data
    integer, dimension(fft_dimension) :: &
         vector_sizes_real, vector_sizes_complex

    ! we need two dummy arrays for the planner.  Since at
    ! present we are not using SSE instructions, these are save
    ! to be local to this routine
    real, dimension(:,:), allocatable :: dummy_real_data
    complex, dimension (:,:), allocatable :: dummy_complex_data
#endif
    call time_message(.false., time_fft, ' FFT')
    fft = basic_init(m * n, direction, howmany, .true.)

#if FFT == _FFTW3_
    stride  = howmany
    if(present(stride_)) stride = stride_

    allocate (dummy_real_data(stride, m*n))
    allocate (dummy_complex_data(stride, (m/2+1)*n))

    vector_sizes_real(1) = m
    vector_sizes_real(2) = n

    vector_sizes_complex(1) = m/2+1
    vector_sizes_complex(2) = n

    call FFTW_PREFIX(_plan_many_dft_c2r)(fft%plan, fft_dimension, &
         vector_sizes_real, howmany, &
         dummy_complex_data, vector_sizes_complex, stride, 1, &
         dummy_real_data, vector_sizes_real, stride, 1, &
         fftw_flags())

    deallocate (dummy_real_data, dummy_complex_data)
#endif
    call time_message(.false., time_fft, ' FFT')
  end subroutine init_crfftw_2d

  !> Destroy FFT plan
  !>
  !> This is safe to call on previously destroyed or
  !> non-constructed/empty plans
  subroutine delete_fft(fft)
    use job_manage, only: time_message
    implicit none
    type (fft_type), intent (in out) :: fft

    !Don't try to delete ffts which have not been created
    if(.not.fft%created) return
#if FFT == _FFTW3_
       call FFTW_PREFIX(_destroy_plan)(fft%plan)
# endif
    fft%created=.false.
    call time_message(.false., time_fft, ' FFT')
  end subroutine delete_fft

  !> Restore FFTW to pristine state
  subroutine finish_fft_work
    use job_manage, only: time_message
    implicit none
# if FFT == _FFTW3_
    integer :: ierr
    call time_message(.false., time_fft, ' FFT')
    call FFTW_PREFIX(_cleanup)(ierr)
    call time_message(.false., time_fft, ' FFT')
# endif
  end subroutine finish_fft_work

  !> Save FFTW wisdom to file
  !>
  !> See FFTW documentation for details of wisdom
  subroutine save_wisdom(filename)
    use job_manage, only: time_message
# ifdef FFT
    use iso_c_binding, only: c_int, c_null_char
    integer(c_int) :: ret
# endif
    character(*), intent(in) :: filename

# ifdef FFT
    interface
      function fftw_export_wisdom_to_filename(c_name) &
          bind(c, name='fftw_export_wisdom_to_filename')
        use iso_c_binding
        integer(c_int) :: fftw_export_wisdom_to_filename
        character(c_char) :: c_name(*)
      end function
    end interface
    call time_message(.false., time_fft, ' FFT')
    ret = fftw_export_wisdom_to_filename(filename//c_null_char)
    call time_message(.false., time_fft, ' FFT')
# endif
  end subroutine save_wisdom

  !> Load FFTW wisdom from file
  !>
  !> See FFTW documentation for details of wisdom
  subroutine load_wisdom(filename)
    use job_manage, only: time_message
# ifdef FFT
    use iso_c_binding, only: c_int, c_null_char
    integer(c_int) :: ret
# endif
    character(*), intent(in) :: filename

# ifdef FFT
    interface
      function fftw_import_wisdom_from_filename(c_name) &
          bind(c, name='fftw_import_wisdom_from_filename')
        use iso_c_binding
        integer(c_int) :: fftw_import_wisdom_from_filename
        character(c_char) :: c_name(*)
      end function
    end interface
    call time_message(.false., time_fft, ' FFT')
    ret = fftw_import_wisdom_from_filename(filename//c_null_char)
    call time_message(.false., time_fft, ' FFT')
# endif
  end subroutine load_wisdom
end module fft_work