Variable LALL_FFTW to perform fftw transforms in one bunch of all fields

ecmwf-ifs · Jan 28, 2025 · 3e2cefc · 3e2cefc
1 parent 2411323
commit 3e2cefc
Show file tree

Hide file tree

Showing 7 changed files with 31 additions and 22 deletions.
diff --git a/src/trans/cpu/external/setup_trans.F90 b/src/trans/cpu/external/setup_trans.F90
@@ -10,7 +10,7 @@
 
 SUBROUTINE SETUP_TRANS(KSMAX,KDGL,KDLON,KLOEN,LDSPLIT,PSTRET,&
 &KTMAX,KRESOL,PWEIGHT,LDGRIDONLY,LDUSERPNM,LDKEEPRPNM,LDUSEFLT,&
-&LDSPSETUPONLY,LDPNMONLY,LDUSEFFTW,&
+&LDSPSETUPONLY,LDPNMONLY,LDUSEFFTW,LD_ALL_FFTW,&
 &LDLL,LDSHIFTLL,CDIO_LEGPOL,CDLEGPOLFNAME,KLEGPOLPTR,KLEGPOLPTR_LEN)
 
 !**** *SETUP_TRANS* - Setup transform package for specific resolution
@@ -54,6 +54,7 @@ SUBROUTINE SETUP_TRANS(KSMAX,KDGL,KDLON,KLOEN,LDSPLIT,PSTRET,&
 !                  FLT, otherwise always kept)
 !     LDPNMONLY  - Compute the Legendre polynomials only, not the FFTs.
 !     LDUSEFFTW - Use FFTW for FFTs (option deprecated - FFTW is now mandatory)
+!     LD_ALL_FFTW : T to transform all fields in one call, F to transforms fields one after another
 !     LDLL                 - Setup second set of input/output latitudes
 !                                 the number of input/output latitudes to transform is equal KDGL
 !                                 or KDGL+2 in the case that includes poles + equator
@@ -96,6 +97,7 @@ SUBROUTINE SETUP_TRANS(KSMAX,KDGL,KDLON,KLOEN,LDSPLIT,PSTRET,&
 !        G. Mozdzynski : Jun 2015 Support alternative FFTs to FFTW
 !        M.Hamrud/W.Deconinck : July 2015 IO options for Legenndre polynomials
 !        R. El Khatib 07-Mar-2016 Better flexibility for Legendre polynomials computation in stretched mode
+!        R. El Khatib  08-Jun-2023 LALL_FFTW for better flexibility
 !     ------------------------------------------------------------------
 
 USE PARKIND1,                    ONLY: JPIM, JPRD, JPRB
@@ -140,6 +142,7 @@ SUBROUTINE SETUP_TRANS(KSMAX,KDGL,KDLON,KLOEN,LDSPLIT,PSTRET,&
 REAL(KIND=JPRD)    ,OPTIONAL,INTENT(IN) :: PSTRET
 LOGICAL   ,OPTIONAL,INTENT(IN):: LDGRIDONLY
 LOGICAL   ,OPTIONAL,INTENT(IN):: LDUSEFLT
+LOGICAL   ,OPTIONAL,INTENT(IN):: LD_ALL_FFTW
 LOGICAL   ,OPTIONAL,INTENT(IN):: LDUSERPNM
 LOGICAL   ,OPTIONAL,INTENT(IN):: LDKEEPRPNM
 LOGICAL   ,OPTIONAL,INTENT(IN):: LDSPSETUPONLY
@@ -228,6 +231,7 @@ SUBROUTINE SETUP_TRANS(KSMAX,KDGL,KDLON,KLOEN,LDSPLIT,PSTRET,&
 S%LUSE_BELUSOV=.TRUE. ! use Belusov algorithm to compute RPNM array instead of per m
 S%LKEEPRPNM=.FALSE. ! Keep Legendre polonomials (RPNM)
 S%LUSEFLT=.FALSE. ! Use fast legendre transforms
+TW%LALL_FFTW=.FALSE. ! transform fields one at a time
 LLSPSETUPONLY = .FALSE. ! Only create distributed spectral setup
 S%LDLL = .FALSE. ! use mapping to/from second set of latitudes
 S%LSHIFTLL = .FALSE. ! shift output lat-lon by 0.5dx, 0.5dy
@@ -340,6 +344,10 @@ SUBROUTINE SETUP_TRANS(KSMAX,KDGL,KDLON,KLOEN,LDSPLIT,PSTRET,&
 ! Setup distribution independent dimensions
 CALL SETUP_DIMS
 
+IF(PRESENT(LD_ALL_FFTW)) THEN
+  TW%LALL_FFTW=LD_ALL_FFTW
+ENDIF
+
 S%LSOUTHPNM=.FALSE.
 IF(PRESENT(PSTRET)) THEN
   IF (ABS(PSTRET-1.0_JPRD)>100._JPRD*EPSILON(1._JPRD)) THEN

diff --git a/src/trans/cpu/internal/ftdir_mod.F90 b/src/trans/cpu/internal/ftdir_mod.F90
@@ -43,14 +43,14 @@ SUBROUTINE FTDIR(PREEL,KFIELDS,KGL)
 !        D. Degrauwe  (Feb 2012): Alternative extension zone (E')
 !        G. Mozdzynski (Oct 2014): support for FFTW transforms
 !        G. Mozdzynski (Jun 2015): Support alternative FFTs to FFTW
-
+!        R. El Khatib  08-Jun-2023 LALL_FFTW for better flexibility
 !     ------------------------------------------------------------------
 
 USE PARKIND1  ,ONLY : JPIM, JPIB, JPRB
 
 USE TPM_DISTR       ,ONLY : D, MYSETW
 USE TPM_GEOMETRY    ,ONLY : G
-USE TPM_FFTW        ,ONLY : EXEC_FFTW
+USE TPM_FFTW        ,ONLY : TW, EXEC_FFTW
 USE TPM_DIM         ,ONLY : R
 !
 
@@ -61,7 +61,6 @@ SUBROUTINE FTDIR(PREEL,KFIELDS,KGL)
 
 INTEGER(KIND=JPIM) :: IGLG,IST,ILEN,JJ,IST1
 INTEGER(KIND=JPIM) :: IOFF,IRLEN,ICLEN, ITYPE
-LOGICAL :: LL_ALL=.FALSE. ! T=do kfields ffts in one batch, F=do kfields ffts one at a time
 
 !     ------------------------------------------------------------------
 
@@ -75,7 +74,7 @@ SUBROUTINE FTDIR(PREEL,KFIELDS,KGL)
   IRLEN=G%NLOEN(IGLG)+R%NNOEXTZL
   ICLEN=(IRLEN/2+1)*2
 
-  CALL EXEC_FFTW(ITYPE,IRLEN,ICLEN,IOFF,KFIELDS,LL_ALL,PREEL)
+  CALL EXEC_FFTW(ITYPE,IRLEN,ICLEN,IOFF,KFIELDS,TW%LALL_FFTW,PREEL)
 ENDIF
 
 IST1=1

diff --git a/src/trans/cpu/internal/ftdirad_mod.F90 b/src/trans/cpu/internal/ftdirad_mod.F90
@@ -42,14 +42,14 @@ SUBROUTINE FTDIRAD(PREEL,KFIELDS,KGL)
 !        D. Degrauwe  (Feb 2012): Alternative extension zone (E')
 !        G. Mozdzynski (Oct 2014): support for FFTW transforms
 !        G. Mozdzynski (Jun 2015): Support alternative FFTs to FFTW 
-
+!        R. El Khatib  08-Jun-2023 LALL_FFTW for better flexibility
 !     ------------------------------------------------------------------
 
 USE PARKIND1  ,ONLY : JPIM, JPRB
 
 USE TPM_DISTR       ,ONLY : D, MYSETW
 USE TPM_GEOMETRY    ,ONLY : G
-USE TPM_FFTW        ,ONLY : EXEC_FFTW
+USE TPM_FFTW        ,ONLY : TW, EXEC_FFTW
 USE TPM_DIM         ,ONLY : R
 
 IMPLICIT NONE
@@ -60,7 +60,6 @@ SUBROUTINE FTDIRAD(PREEL,KFIELDS,KGL)
 INTEGER(KIND=JPIM) :: IGLG,IST,ILEN,JJ,JF,ILOEN
 INTEGER(KIND=JPIM) :: IOFF,IRLEN,ICLEN,ITYPE
 REAL(KIND=JPRB) :: ZMUL
-LOGICAL :: LL_ALL=.FALSE. ! T=do kfields ffts in one batch, F=do kfields ffts one at a time
 !     ------------------------------------------------------------------
 
 ITYPE = 1
@@ -78,7 +77,7 @@ SUBROUTINE FTDIRAD(PREEL,KFIELDS,KGL)
   ENDDO
 ENDDO
 
-CALL EXEC_FFTW(ITYPE,IRLEN,ICLEN,IOFF,KFIELDS,LL_ALL,PREEL)
+CALL EXEC_FFTW(ITYPE,IRLEN,ICLEN,IOFF,KFIELDS,TW%LALL_FFTW,PREEL)
 
   ! Change of metric (not in forward routine)
 

diff --git a/src/trans/cpu/internal/ftinv_mod.F90 b/src/trans/cpu/internal/ftinv_mod.F90
@@ -42,13 +42,14 @@ SUBROUTINE FTINV(PREEL,KFIELDS,KGL)
 !        D. Degrauwe  (Feb 2012): Alternative extension zone (E')
 !        G. Mozdzynski (Oct 2014): support for FFTW transforms
 !        G. Mozdzynski (Jun 2015): Support alternative FFTs to FFTW
+!        R. El Khatib  08-Jun-2023 LALL_FFTW for better flexibility
 !     ------------------------------------------------------------------
 
 USE PARKIND1  ,ONLY : JPIM, JPRB
 
 USE TPM_DISTR       ,ONLY : D, MYSETW
 USE TPM_GEOMETRY    ,ONLY : G
-USE TPM_FFTW        ,ONLY : EXEC_FFTW
+USE TPM_FFTW        ,ONLY : TW, EXEC_FFTW
 USE TPM_DIM         ,ONLY : R
 
 IMPLICIT NONE
@@ -58,7 +59,6 @@ SUBROUTINE FTINV(PREEL,KFIELDS,KGL)
 
 INTEGER(KIND=JPIM) :: IGLG,IST,ILEN,JJ,JF,IST1
 INTEGER(KIND=JPIM) :: IOFF,IRLEN,ICLEN, ITYPE
-LOGICAL :: LL_ALL=.FALSE. ! T=do kfields ffts in one batch, F=do kfields ffts one at a time
 
 !     ------------------------------------------------------------------
 
@@ -80,7 +80,7 @@ SUBROUTINE FTINV(PREEL,KFIELDS,KGL)
   IRLEN=G%NLOEN(IGLG)+R%NNOEXTZL
   ICLEN=(IRLEN/2+1)*2
 
-  CALL EXEC_FFTW(ITYPE,IRLEN,ICLEN,IOFF,KFIELDS,LL_ALL,PREEL)
+  CALL EXEC_FFTW(ITYPE,IRLEN,ICLEN,IOFF,KFIELDS,TW%LALL_FFTW,PREEL)
 ENDIF
 
 !     ------------------------------------------------------------------

diff --git a/src/trans/cpu/internal/ftinvad_mod.F90 b/src/trans/cpu/internal/ftinvad_mod.F90
@@ -42,15 +42,15 @@ SUBROUTINE FTINVAD(PREEL,KFIELDS,KGL)
 !        D. Degrauwe  (Feb 2012): Alternative extension zone (E')
 !        G. Mozdzynski (Oct 2014): support for FFTW transforms
 !        G. Mozdzynski (Jun 2015): Support alternative FFTs to FFTW
-
+!        R. El Khatib  08-Jun-2023 LALL_FFTW for better flexibility
 !     ------------------------------------------------------------------
 
 USE PARKIND1  ,ONLY : JPIM, JPIB, JPRB
 
 USE TPM_DISTR       ,ONLY : D, MYSETW
 USE TPM_DIM         ,ONLY : R
 USE TPM_GEOMETRY    ,ONLY : G
-USE TPM_FFTW        ,ONLY : EXEC_FFTW
+USE TPM_FFTW        ,ONLY : TW, EXEC_FFTW
 USE TPM_DIM         ,ONLY : R
 !
 
@@ -61,7 +61,6 @@ SUBROUTINE FTINVAD(PREEL,KFIELDS,KGL)
 
 INTEGER(KIND=JPIM) :: IGLG,IST,ILEN,JJ,JF,ILOEN
 INTEGER(KIND=JPIM) :: IOFF,IRLEN,ICLEN,ITYPE
-LOGICAL :: LL_ALL=.FALSE. ! T=do kfields ffts in one batch, F=do kfields ffts one at a time
 
 !     ------------------------------------------------------------------
 
@@ -81,7 +80,7 @@ SUBROUTINE FTINVAD(PREEL,KFIELDS,KGL)
   ENDDO
 ENDDO
 
-CALL EXEC_FFTW(ITYPE,IRLEN,ICLEN,IOFF,KFIELDS,LL_ALL,PREEL)
+CALL EXEC_FFTW(ITYPE,IRLEN,ICLEN,IOFF,KFIELDS,TW%LALL_FFTW,PREEL)
 
 DO JJ=1,ILEN
   DO JF=1,KFIELDS

diff --git a/src/trans/cpu/internal/tpm_fftw.F90 b/src/trans/cpu/internal/tpm_fftw.F90
@@ -17,6 +17,7 @@ MODULE TPM_FFTW
 !   -------------- 
 !     Original      October 2014
 !     R. El Khatib 01-Sep-2015 More subroutines for better modularity
+!        R. El Khatib  08-Jun-2023 LALL_FFTW for better flexibility
 
 USE, INTRINSIC :: ISO_C_BINDING
 
@@ -44,7 +45,8 @@ MODULE TPM_FFTW
   INTEGER(KIND=JPIM),ALLOCATABLE :: N_PLANS(:)
   TYPE(FFTW_PLAN),POINTER :: FFTW_PLANS(:) => NULL()
   INTEGER(KIND=JPIM) :: N_MAX=0         ! maximum number of latitudes
-  INTEGER(KIND=JPIM) :: N_MAX_PLANS=4   ! maximum number of plans for each active latitudes
+  INTEGER(KIND=JPIM) :: N_MAX_PLANS=4   ! maximum number of plans for each active latitude
+  LOGICAL            :: LALL_FFTW=.FALSE. ! T=do kfields ffts in one batch, F=do kfields ffts one at a time
 END TYPE FFTW_TYPE
 
 
@@ -422,8 +424,8 @@ SUBROUTINE EXEC_EFFTW(KTYPE,KRLEN,KCLEN,KOFF,KFIELDS,LD_ALL,PREEL)
        CALL SFFTW_EXECUTE_DFT_C2R(IPLAN_C2R,ZFFT,ZFFT)
     END IF
     IF (LHOOK) CALL DR_HOOK('FFTW_EXECUTE_DFT_C2R',1,ZHOOK_HANDLE2)
-    DO JJ=1,KRLEN
-      DO JF=1,KFIELDS
+    DO JF=1,KFIELDS
+      DO JJ=1,KRLEN
         PREEL(KOFF+JJ-1,JF)=ZFFT(JJ,JF)
       ENDDO
     ENDDO
@@ -440,8 +442,8 @@ SUBROUTINE EXEC_EFFTW(KTYPE,KRLEN,KCLEN,KOFF,KFIELDS,LD_ALL,PREEL)
        CALL SFFTW_EXECUTE_DFT_R2C(IPLAN_C2R,ZFFT,ZFFT)
     END IF
     IF (LHOOK) CALL DR_HOOK('FFTW_EXECUTE_DFT_R2C',1,ZHOOK_HANDLE2)
-    DO JJ=1,KCLEN
-      DO JF=1,KFIELDS
+    DO JF=1,KFIELDS
+      DO JJ=1,KCLEN
         PREEL(KOFF+JJ-1,JF)=ZFFT(JJ,JF)/REAL(KRLEN,JPRB)
       ENDDO
     ENDDO

diff --git a/src/trans/include/ectrans/setup_trans.h b/src/trans/include/ectrans/setup_trans.h
@@ -11,7 +11,7 @@
 INTERFACE
 SUBROUTINE SETUP_TRANS(KSMAX,KDGL,KDLON,KLOEN,LDSPLIT,PSTRET,&
 &KTMAX,KRESOL,PWEIGHT,LDGRIDONLY,LDUSERPNM,LDKEEPRPNM,LDUSEFLT,&
-&LDSPSETUPONLY,LDPNMONLY,LDUSEFFTW,&
+&LDSPSETUPONLY,LDPNMONLY,LDUSEFFTW,LD_ALL_FFTW,&
 &LDLL,LDSHIFTLL,CDIO_LEGPOL,CDLEGPOLFNAME,KLEGPOLPTR,KLEGPOLPTR_LEN)
 
 !**** *SETUP_TRANS* - Setup transform package for specific resolution
@@ -52,6 +52,7 @@ SUBROUTINE SETUP_TRANS(KSMAX,KDGL,KDLON,KLOEN,LDSPLIT,PSTRET,&
 !                  FLT, otherwise always kept)
 !     LDPNMONLY  - Compute the Legendre polynomialsonly, not the FFTs.
 !     LDUSEFFTW - Use FFTW for FFTs (option deprecated - FFTW is now mandatory)
+!     LD_ALL_FFTW : T to transform all fields in one call, F to transforms fields one after another
 !     LDLL                 - Setup second set of input/output latitudes
 !                                 the number of input/output latitudes to transform is equal KDGL 
 !                                 or KDGL+2 in the case that includes poles + equator
@@ -97,6 +98,7 @@ REAL(KIND=JPRD)    ,OPTIONAL,INTENT(IN) :: PWEIGHT(:)
 REAL(KIND=JPRD)    ,OPTIONAL,INTENT(IN) :: PSTRET
 LOGICAL   ,OPTIONAL,INTENT(IN):: LDGRIDONLY
 LOGICAL   ,OPTIONAL,INTENT(IN):: LDUSEFLT
+LOGICAL   ,OPTIONAL,INTENT(IN):: LD_ALL_FFTW
 LOGICAL   ,OPTIONAL,INTENT(IN):: LDUSERPNM
 LOGICAL   ,OPTIONAL,INTENT(IN):: LDKEEPRPNM
 LOGICAL   ,OPTIONAL,INTENT(IN):: LDPNMONLY