! (C) Copyright 2014- ECMWF.
! (C) Copyright 2022- NVIDIA.
!
! This software is licensed under the terms of the Apache Licence Version 2.0
! which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
! In applying this licence, ECMWF does not waive the privileges and immunities
! granted to it by virtue of its status as an intergovernmental organisation
! nor does it submit to any jurisdiction.
!

MODULE TPM_HICFFT

  !   Author.
  !   -------
  !     George Mozdzynski
  !
  !   Modifications.
  !   --------------
  !     Original      October 2014
  !     HICFFT abstraction for CUDA and HIP     August 2023     B. Reuter

  USE, INTRINSIC :: ISO_C_BINDING, ONLY: C_FLOAT, C_DOUBLE, C_LOC
  USE GROWING_ALLOCATOR_MOD,       ONLY: GROWING_ALLOCATION_TYPE

  IMPLICIT NONE

  SAVE

  PRIVATE
  PUBLIC EXECUTE_DIR_FFT, EXECUTE_INV_FFT
  PUBLIC CLEAN_FFT

INTERFACE EXECUTE_DIR_FFT
  MODULE PROCEDURE EXECUTE_DIR_FFT_FLOAT,EXECUTE_DIR_FFT_DOUBLE
END INTERFACE

INTERFACE EXECUTE_INV_FFT
  MODULE PROCEDURE EXECUTE_INV_FFT_FLOAT,EXECUTE_INV_FFT_DOUBLE
END INTERFACE

INTERFACE
  SUBROUTINE CLEAN_FFT(RESOL_ID) BIND(C, NAME="clean_fft")
    USE ISO_C_BINDING
    INTEGER(KIND=C_INT), INTENT(IN), VALUE :: RESOL_ID
  END SUBROUTINE
END INTERFACE



  ! ------------------------------------------------------------------
  CONTAINS
  ! ------------------------------------------------------------------


SUBROUTINE EXECUTE_DIR_FFT_FLOAT(PREEL_REAL,PREEL_COMPLEX,RESOL_ID,KFIELD,LOENS,OFFSETS,ALLOC)
  USE EC_PARKIND ,ONLY : JPIM, JPIB

  IMPLICIT NONE

  REAL(KIND=C_FLOAT), POINTER, INTENT(IN) :: PREEL_REAL(:)
  REAL(KIND=C_FLOAT), POINTER, INTENT(OUT) :: PREEL_COMPLEX(:)
  INTEGER(KIND=JPIM),INTENT(IN)  :: RESOL_ID
  INTEGER(KIND=JPIM),INTENT(IN)  :: KFIELD
  INTEGER(KIND=JPIM),INTENT(IN)  :: LOENS(:)
  INTEGER(KIND=JPIB),INTENT(IN)  :: OFFSETS(:)
  TYPE(GROWING_ALLOCATION_TYPE), INTENT(IN), POINTER :: ALLOC
  INTERFACE
    SUBROUTINE EXECUTE_DIR_FFT_FLOAT_C(PREEL_REAL,PREEL_COMPLEX,RESOL_ID,KFIELD,LOENS,OFFSETS,NFFT,ALLOC) &
 &                   BIND(C, NAME="execute_dir_fft_float")
      USE ISO_C_BINDING, ONLY: C_FLOAT, C_INT, C_PTR, C_INT64_T
      TYPE(C_PTR), VALUE :: PREEL_REAL
      TYPE(C_PTR), VALUE :: PREEL_COMPLEX
      INTEGER(KIND=C_INT),INTENT(IN),VALUE  :: RESOL_ID
      INTEGER(KIND=C_INT),INTENT(IN),VALUE  :: KFIELD
      INTEGER(KIND=C_INT),INTENT(IN)  :: LOENS(*)
      INTEGER(KIND=C_INT64_T),INTENT(IN)  :: OFFSETS(*)
      INTEGER(KIND=C_INT),INTENT(IN),VALUE :: NFFT
      TYPE(C_PTR), INTENT(IN), VALUE :: ALLOC
    END SUBROUTINE
  END INTERFACE

#ifdef OMPGPU
  !$OMP TARGET DATA USE_DEVICE_ADDR(PREEL_REAL,PREEL_COMPLEX)
#endif
#ifdef ACCGPU
  !$ACC HOST_DATA USE_DEVICE(PREEL_REAL,PREEL_COMPLEX)
#endif
  CALL EXECUTE_DIR_FFT_FLOAT_C(C_LOC(PREEL_REAL),C_LOC(PREEL_COMPLEX),RESOL_ID,KFIELD,LOENS,OFFSETS,SIZE(LOENS),C_LOC(ALLOC))
#ifdef ACCGPU
  !$ACC END HOST_DATA
#endif
#ifdef OMPGPU
  !$OMP END TARGET DATA
#endif

END SUBROUTINE EXECUTE_DIR_FFT_FLOAT
SUBROUTINE EXECUTE_DIR_FFT_DOUBLE(PREEL_REAL,PREEL_COMPLEX,RESOL_ID,KFIELD,LOENS,OFFSETS,ALLOC)
  USE EC_PARKIND ,ONLY : JPIM, JPIB

  IMPLICIT NONE

  REAL(KIND=C_DOUBLE), POINTER, INTENT(IN) :: PREEL_REAL(:)
  REAL(KIND=C_DOUBLE), POINTER, INTENT(OUT) :: PREEL_COMPLEX(:)
  INTEGER(KIND=JPIM),INTENT(IN)  :: RESOL_ID
  INTEGER(KIND=JPIM),INTENT(IN)  :: KFIELD
  INTEGER(KIND=JPIM),INTENT(IN)  :: LOENS(:)
  INTEGER(KIND=JPIB),INTENT(IN)  :: OFFSETS(:)
  TYPE(GROWING_ALLOCATION_TYPE), INTENT(IN), POINTER :: ALLOC
  INTERFACE
    SUBROUTINE EXECUTE_DIR_FFT_DOUBLE_C(PREEL_REAL,PREEL_COMPLEX,RESOL_ID,KFIELD,LOENS,OFFSETS,NFFT,ALLOC) &
 &                   BIND(C, NAME="execute_dir_fft_double")
      USE ISO_C_BINDING, ONLY: C_DOUBLE, C_INT, C_PTR, C_INT64_T
      TYPE(C_PTR), VALUE :: PREEL_REAL
      TYPE(C_PTR), VALUE :: PREEL_COMPLEX
      INTEGER(KIND=C_INT),INTENT(IN),VALUE  :: RESOL_ID
      INTEGER(KIND=C_INT),INTENT(IN),VALUE  :: KFIELD
      INTEGER(KIND=C_INT),INTENT(IN)  :: LOENS(*)
      INTEGER(KIND=C_INT64_T),INTENT(IN)  :: OFFSETS(*)
      INTEGER(KIND=C_INT),INTENT(IN),VALUE :: NFFT
      TYPE(C_PTR), INTENT(IN), VALUE :: ALLOC
    END SUBROUTINE
  END INTERFACE

#ifdef OMPGPU
  !$OMP TARGET DATA USE_DEVICE_ADDR(PREEL_REAL,PREEL_COMPLEX)
#endif
#ifdef ACCGPU
  !$ACC HOST_DATA USE_DEVICE(PREEL_REAL,PREEL_COMPLEX)
#endif
  CALL EXECUTE_DIR_FFT_DOUBLE_C(C_LOC(PREEL_REAL),C_LOC(PREEL_COMPLEX),RESOL_ID,KFIELD,LOENS,OFFSETS,SIZE(LOENS),C_LOC(ALLOC))
#ifdef ACCGPU
  !$ACC END HOST_DATA
#endif
#ifdef OMPGPU
  !$OMP END TARGET DATA
#endif

END SUBROUTINE EXECUTE_DIR_FFT_DOUBLE

SUBROUTINE EXECUTE_INV_FFT_FLOAT(PREEL_COMPLEX,PREEL_REAL,RESOL_ID,KFIELD,LOENS,OFFSETS,ALLOC)
  USE EC_PARKIND ,ONLY : JPIM, JPIB

  IMPLICIT NONE

  REAL(KIND=C_FLOAT), POINTER, INTENT(IN) :: PREEL_COMPLEX(:)
  REAL(KIND=C_FLOAT), POINTER, INTENT(OUT) :: PREEL_REAL(:)
  INTEGER(KIND=JPIM),INTENT(IN)  :: RESOL_ID
  INTEGER(KIND=JPIM),INTENT(IN)  :: KFIELD
  INTEGER(KIND=JPIM),INTENT(IN)  :: LOENS(:)
  INTEGER(KIND=JPIB),INTENT(IN)  :: OFFSETS(:)
  TYPE(GROWING_ALLOCATION_TYPE), INTENT(IN), POINTER :: ALLOC
  INTERFACE
    SUBROUTINE EXECUTE_INV_FFT_FLOAT_C(PREEL_COMPLEX,PREEL_REAL,RESOL_ID,KFIELD,LOENS,OFFSETS,NFFT,ALLOC) &
 &                   BIND(C, NAME="execute_inv_fft_float")
      USE ISO_C_BINDING, ONLY: C_FLOAT, C_INT, C_PTR, C_INT64_T
      TYPE(C_PTR), VALUE :: PREEL_COMPLEX
      TYPE(C_PTR), VALUE :: PREEL_REAL
      INTEGER(KIND=C_INT),INTENT(IN),VALUE  :: RESOL_ID
      INTEGER(KIND=C_INT),INTENT(IN),VALUE  :: KFIELD
      INTEGER(KIND=C_INT),INTENT(IN)  :: LOENS(*)
      INTEGER(KIND=C_INT64_T),INTENT(IN)  :: OFFSETS(*)
      INTEGER(KIND=C_INT),INTENT(IN),VALUE :: NFFT
      TYPE(C_PTR), INTENT(IN), VALUE :: ALLOC
    END SUBROUTINE
  END INTERFACE

#ifdef OMPGPU
  !$OMP TARGET DATA USE_DEVICE_ADDR(PREEL_COMPLEX,PREEL_REAL)
#endif
#ifdef ACCGPU
  !$ACC HOST_DATA USE_DEVICE(PREEL_COMPLEX,PREEL_REAL)
#endif
  CALL EXECUTE_INV_FFT_FLOAT_C(C_LOC(PREEL_COMPLEX),C_LOC(PREEL_REAL),RESOL_ID,KFIELD,LOENS,OFFSETS,SIZE(LOENS),C_LOC(ALLOC))
#ifdef ACCGPU
  !$ACC END HOST_DATA
#endif
#ifdef OMPGPU
  !$OMP END TARGET DATA
#endif

END SUBROUTINE

SUBROUTINE EXECUTE_INV_FFT_DOUBLE(PREEL_COMPLEX,PREEL_REAL,RESOL_ID,KFIELD,LOENS,OFFSETS,ALLOC)
  USE EC_PARKIND ,ONLY : JPIM, JPIB

  IMPLICIT NONE

  REAL(KIND=C_DOUBLE), POINTER, INTENT(IN) :: PREEL_COMPLEX(:)
  REAL(KIND=C_DOUBLE), POINTER, INTENT(OUT) :: PREEL_REAL(:)
  INTEGER(KIND=JPIM),INTENT(IN)  :: RESOL_ID
  INTEGER(KIND=JPIM),INTENT(IN)  :: KFIELD
  INTEGER(KIND=JPIM),INTENT(IN)  :: LOENS(:)
  INTEGER(KIND=JPIB),INTENT(IN)  :: OFFSETS(:)
  TYPE(GROWING_ALLOCATION_TYPE), INTENT(IN), POINTER :: ALLOC
  INTERFACE
    SUBROUTINE EXECUTE_INV_FFT_DOUBLE_C(PREEL_COMPLEX,PREEL_REAL,RESOL_ID,KFIELD,LOENS,OFFSETS,NFFT,ALLOC) &
 &                   BIND(C, NAME="execute_inv_fft_double")
      USE ISO_C_BINDING, ONLY: C_DOUBLE, C_INT, C_PTR, C_INT64_T
      TYPE(C_PTR), VALUE :: PREEL_COMPLEX
      TYPE(C_PTR), VALUE :: PREEL_REAL
      INTEGER(KIND=C_INT),INTENT(IN),VALUE  :: RESOL_ID
      INTEGER(KIND=C_INT),INTENT(IN),VALUE  :: KFIELD
      INTEGER(KIND=C_INT),INTENT(IN)  :: LOENS(*)
      INTEGER(KIND=C_INT64_T),INTENT(IN)  :: OFFSETS(*)
      INTEGER(KIND=C_INT),INTENT(IN),VALUE :: NFFT
      TYPE(C_PTR), INTENT(IN), VALUE :: ALLOC
    END SUBROUTINE
  END INTERFACE

#ifdef OMPGPU
  !$OMP TARGET DATA USE_DEVICE_ADDR(PREEL_COMPLEX,PREEL_REAL)
#endif
#ifdef ACCGPU
  !$ACC HOST_DATA USE_DEVICE(PREEL_COMPLEX,PREEL_REAL)
#endif
  CALL EXECUTE_INV_FFT_DOUBLE_C(C_LOC(PREEL_COMPLEX),C_LOC(PREEL_REAL),RESOL_ID,KFIELD,LOENS,OFFSETS,SIZE(LOENS),C_LOC(ALLOC))
#ifdef ACCGPU
  !$ACC END HOST_DATA
#endif
#ifdef OMPGPU
  !$OMP END TARGET DATA
#endif

END SUBROUTINE


END MODULE TPM_HICFFT
