Move more TPM modules and setup routines to common

src/trans/common/CMakeLists.txt

@@ -18,19 +18,24 @@ list( APPEND ectrans_common_src
   internal/cpledn_mod.F90
   internal/field_split_mod.F90
   internal/gawl_mod.F90
+  internal/interpol_decomp_mod.F90
   internal/sugaw_mod.F90
   internal/supol_mod.F90
   internal/supolf_mod.F90
   internal/tpm_constants.F90
   internal/tpm_ctl.F90
   internal/tpm_dim.F90
+  internal/tpm_fields.F90
   internal/tpm_gen.F90
   internal/tpm_geometry.F90
   internal/tpm_pol.F90
   internal/tpm_distr.F90
   internal/pe2set_mod.F90
   internal/set2pe_mod.F90
   internal/eq_regions_mod.F90
+  internal/pre_suleg_mod.F90
+  internal/setup_dims_mod.F90
+  internal/setup_geom_mod.F90
   internal/shuffle_mod.F90
   internal/sump_trans0_mod.F90
   internal/sustaonl_mod.F90
@@ -42,6 +47,7 @@ list( APPEND ectrans_common_src
   internal/myrecvset_mod.F90
   internal/suwavedi_mod.F90
   internal/sump_trans_preleg_mod.F90
+  internal/wts500_mod.F90
   external/get_current.F90
   external/setup_trans0.F90
   external/ini_spec_dist.F90

src/trans/gpu/internal/pre_suleg_mod.F90 → src/trans/common/internal/pre_suleg_mod.F90

@@ -12,7 +12,7 @@ MODULE PRE_SULEG_MOD
 IMPLICIT NONE
 CONTAINS
 SUBROUTINE PRE_SULEG
-USE PARKIND1,      ONLY: JPRD, JPIM
+USE EC_PARKIND,    ONLY: JPRD, JPIM
 USE TPM_GEN,       ONLY: NPRINTLEV, NOUT
 USE TPM_DIM,       ONLY: R
 USE TPM_CONSTANTS, ONLY: RA

src/trans/gpu/internal/setup_dims_mod.F90 → src/trans/common/internal/setup_dims_mod.F90

@@ -11,11 +11,13 @@
 MODULE SETUP_DIMS_MOD
 CONTAINS
 SUBROUTINE SETUP_DIMS
+! EXPECTED TO BE SET ALREADY:
+!   - NSMAX
+!   - NTMAX
+!   - NDGL
 
-USE PARKIND1, ONLY: JPIM
-USE TPM_DIM,  ONLY: R
-USE TPM_FLT,  ONLY: S
-!
+USE EC_PARKIND ,ONLY : JPIM
+USE TPM_DIM    ,ONLY : R
 
 IMPLICIT NONE
 
@@ -38,7 +40,6 @@ SUBROUTINE SETUP_DIMS
 
 R%NLEI1 = R%NSMAX+4+MOD(R%NSMAX+4+1,2)
 R%NLEI3 = R%NDGNH+MOD(R%NDGNH+2,2)
-IF (S%LSOUTHPNM) R%NLEI3=2*R%NLEI3 
 
 R%NLED3 = R%NTMAX+2+MOD(R%NTMAX+3,2)
 R%NLED4 = R%NTMAX+3+MOD(R%NTMAX+4,2)

src/trans/cpu/internal/setup_geom_mod.F90 → src/trans/common/internal/setup_geom_mod.F90

@@ -12,7 +12,7 @@ MODULE SETUP_GEOM_MOD
 CONTAINS
 SUBROUTINE SETUP_GEOM
 
-USE PARKIND1  ,ONLY : JPRD, JPIM
+USE EC_PARKIND      ,ONLY : JPRD, JPIM
 
 USE TPM_GEN         ,ONLY : NOUT, NPRINTLEV
 USE TPM_DIM         ,ONLY : R

src/trans/common/internal/tpm_fields.F90

@@ -0,0 +1,39 @@
+! (C) Copyright 2000- ECMWF.
+! (C) Copyright 2000- Meteo-France.
+! (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_FIELDS
+
+USE EC_PARKIND, ONLY: JPIM, JPRD
+
+IMPLICIT NONE
+
+SAVE
+
+TYPE FIELDS_TYPE
+REAL(KIND=JPRD)    ,ALLOCATABLE :: RPNM(:,:) ! Legendre polynomials
+REAL(KIND=JPRD)    ,ALLOCATABLE :: RMU(:)    ! sin(theta) for Gaussian latitudes
+REAL(KIND=JPRD)    ,ALLOCATABLE :: RW(:)     ! Weights of the Gaussian quadrature
+REAL(KIND=JPRD)    ,ALLOCATABLE :: R1MU2(:)  ! 1.-MU*MU, cos(theta)**2
+REAL(KIND=JPRD)    ,ALLOCATABLE :: RACTHE(:) ! 1./SQRT(R1MU2), 1/(cos(theta))
+
+REAL(KIND=JPRD)    ,ALLOCATABLE :: REPSNM(:) ! eps(n,m) used in the Legendre transforms
+REAL(KIND=JPRD)    ,ALLOCATABLE :: RN(:)     ! n (to avoid integer to real conversion)
+REAL(KIND=JPRD)    ,ALLOCATABLE :: RLAPIN(:) ! eigen-values of the inverse Laplace operator
+INTEGER(KIND=JPIM) ,ALLOCATABLE :: NLTN(:)   ! R%NTMAX+2-JN
+
+REAL(KIND=JPRD)    ,ALLOCATABLE :: RMU2(:)   ! sin(theta) for dual input/output latitudes
+REAL(KIND=JPRD)    ,ALLOCATABLE :: RACTHE2(:)! 1./SQRT(R1MU2), 1/(cos(theta)) dual input/output latitudes
+END TYPE FIELDS_TYPE
+
+TYPE(FIELDS_TYPE),ALLOCATABLE,TARGET :: FIELDS_RESOL(:)
+TYPE(FIELDS_TYPE),POINTER     :: F
+
+END MODULE TPM_FIELDS

src/trans/cpu/algor/wts500_mod.F90 → src/trans/common/internal/wts500_mod.F90

@@ -12,11 +12,11 @@ MODULE WTS500_MOD
 CONTAINS
 SUBROUTINE WTS500(PX,PW,KN)
 
-USE PARKIND1  ,ONLY : JPIM     ,JPRB
+USE EC_PARKIND  ,ONLY : JPIM, JPRD
 IMPLICIT NONE
 
 INTEGER(KIND=JPIM), INTENT(IN)  :: KN
-REAL(KIND=JPRB),    INTENT(OUT) :: PX(:),PW(:)
+REAL(KIND=JPRD),    INTENT(OUT) :: PX(:),PW(:)
 
 ! This routine returns a set of Gaussian nodes and weights for
 ! integrating the functions exp(lambda*x)dx over the range x=0 to x=infinity.

src/trans/cpu/algor/seefmm_mix.F90

@@ -39,7 +39,7 @@ module seefmm_mix
 !     ------------------------------------------------------------------
 
 
-use parkind1,only : jpim     ,jprb, jprd
+use parkind1,   only : jpim, jprb, jprd
 use ecsort_mix, only : keysort
 use wts500_mod, only: wts500
 
@@ -86,12 +86,13 @@ recursive subroutine setup_seefmm(kx,px,ky,py,ydfmm,pdiff)
 integer(kind=jpim),intent(in)  :: kx    
 real(kind=jprd)   ,intent(in)  :: px(:)
 integer(kind=jpim),intent(in)  :: ky
-real(kind=jprb)   ,intent(in)  :: py(:)
+real(kind=jprd)   ,intent(in)  :: py(:)
 type(fmm_type)    ,intent(out) :: ydfmm
-real(kind=jprb),optional,intent(in)  :: pdiff(:,:)
+real(kind=jprd),optional,intent(in)  :: pdiff(:,:)
 
-real(kind=jprb) :: zxy(kx+ky),zrt(56),zcik((kx+ky)*(kx+ky))
-real(kind=jprb) :: zr
+real(kind=jprd) :: zxy(kx+ky),zcik((kx+ky)*(kx+ky))
+real(kind=jprd) :: zr, zrt(56), zrw(56)
+real(kind=jprd), allocatable :: zrdexp(:,:)
 integer(kind=jpim) :: ixy
 !---------------------------------------------------------------------------
 ydfmm%nx=kx
@@ -103,16 +104,19 @@ recursive subroutine setup_seefmm(kx,px,ky,py,ydfmm,pdiff)
 ! Combine px and py to form xxy, compute ascending index for xxy
 call comb_xy(kx,px,ky,py,ixy,zxy,ydfmm%index)
 ! Setup quadrature, scale (see 3.1.1 in [1])
-call suquad(ixy,zxy(ydfmm%index(1))-zxy(ydfmm%index(ixy)),&
- & ydfmm%nquad,ydfmm%rw,zrt,zr)
-allocate(ydfmm%rdexp(ydfmm%nquad,ixy))
+call suquad(ixy,zxy(ydfmm%index(1))-zxy(ydfmm%index(ixy)),ydfmm%nquad,&
+  & zrw,zrt,zr)
+allocate(zrdexp(ydfmm%nquad,ixy))
 allocate(ydfmm%nclose(ixy))
 ! Main pre-computation
-call prepotf(kx,ixy,ydfmm%nquad,ydfmm%rw,zrt,zr,zxy,ydfmm%index,&
- & ydfmm%rdexp,ydfmm%nclose,zcik,ydfmm%ncik,pdiff)
-! Needed as size of cik unknown beforehand
+call prepotf(kx,ixy,ydfmm%nquad,zrw,zrt,zr,zxy,ydfmm%index,&
+ & zrdexp,ydfmm%nclose,zcik,ydfmm%ncik,pdiff)
+
+allocate(ydfmm%rdexp(ydfmm%nquad,ixy))
 allocate(ydfmm%cik(ydfmm%ncik))
-ydfmm%cik(:)=zcik(1:ydfmm%ncik)
+ydfmm%rw(:)      = real(zrw(:),jprb)
+ydfmm%rdexp(:,:) = real(zrdexp(:,:),jprb)
+ydfmm%cik(:)     = real(zcik(1:ydfmm%ncik),jprb)
 
 end subroutine setup_seefmm
 !==========================================================================
@@ -432,19 +436,19 @@ end subroutine potfm
 recursive subroutine suquad(kn,prange,kquad,prw,prt,pr)
 implicit none
 
-integer(kind=jpim)        ,intent(in)  :: kn
-real(kind=jprb),intent(in)  :: prange
-integer(kind=jpim)        ,intent(in) :: kquad
-real(kind=jprb),intent(out) :: prw(:)
-real(kind=jprb),intent(out) :: prt(:)
-real(kind=jprb),intent(out) :: pr
+integer(kind=jpim)  ,intent(in)  :: kn
+real(kind=jprd)     ,intent(in)  :: prange
+integer(kind=jpim)  ,intent(in)  :: kquad
+real(kind=jprd)     ,intent(out) :: prw(:)
+real(kind=jprd)     ,intent(out) :: prt(:)
+real(kind=jprd)     ,intent(out) :: pr
 
-real(kind=jprb) :: za,zb,zs
+real(kind=jprd) :: za,zb,zs
 integer(kind=jpim) :: jm
 !-------------------------------------------------------------------------
 
-za=1.0
-zb=500.0
+za=1.0_jprd
+zb=500.0_jprd
 zs=zb/prange
 pr=za/zs
 call wts500(prt,prw,kquad)
@@ -461,23 +465,17 @@ recursive subroutine comb_xy(kx,px,ky,py,kxy,pxy,kindex)
 
 integer(kind=jpim), intent(in)  :: kx,ky
 real(kind=jprd),    intent(in)  :: px(:)
-real(kind=jprb),    intent(in)  :: py(:)
+real(kind=jprd),    intent(in)  :: py(:)
 integer(kind=jpim), intent(in)  :: kxy
-real(kind=jprb),    intent(out) :: pxy(:)
+real(kind=jprd),    intent(out) :: pxy(:)
 integer(kind=jpim), intent(out) :: kindex(:)
 integer(kind=jpim) :: iret
-!integer(kind=jpim) :: jxy
 
 !-------------------------------------------------------------------------
 
 pxy(1:kx)=px(1:kx)
 pxy(kx+1:kx+ky)=py(1:ky)
-!call m01daf(pxy,1,kxy,'D',irank,ifail)
 call keysort(iret,pxy,kxy,descending=.true.,index=kindex,init=.true.)
-!!$do jxy=1,kxy
-!!$  kindex(irank(jxy))=jxy
-!!$enddo
-
 end subroutine comb_xy
 !==========================================================================
 recursive subroutine prepotf(kx,kxy,kquad,prw,prt,pr,pxy,kindex,prdexp,&
@@ -488,20 +486,20 @@ recursive subroutine prepotf(kx,kxy,kquad,prw,prt,pr,pxy,kindex,prdexp,&
 integer(kind=jpim), intent(in)  :: kx
 integer(kind=jpim), intent(in)  :: kxy
 integer(kind=jpim), intent(in)  :: kquad
-real(kind=jprb),    intent(in)  :: pxy(:)
-real(kind=jprb),    intent(in)  :: prw(:)
-real(kind=jprb),    intent(in)  :: pr
-real(kind=jprb),    intent(in)  :: prt(:)
+real(kind=jprd),    intent(in)  :: pxy(:)
+real(kind=jprd),    intent(in)  :: prw(:)
+real(kind=jprd),    intent(in)  :: pr
+real(kind=jprd),    intent(in)  :: prt(:)
 integer(kind=jpim), intent(in)  :: kindex(:)
-real(kind=jprb),    intent(out) :: prdexp(:,:)
+real(kind=jprd),    intent(out) :: prdexp(:,:)
 integer(kind=jpim), intent(out) :: kclosel(:)
-real(kind=jprb),    intent(out) :: pcik(:)
+real(kind=jprd),    intent(out) :: pcik(:)
 integer(kind=jpim), intent(out) :: knocik
-real(kind=jprb),optional, intent(in)  :: pdiff(:,:)
+real(kind=jprd),optional, intent(in)  :: pdiff(:,:)
 
-real(kind=jprb) :: zdx
-real(kind=jprb) :: zsum
-real(kind=jprb) :: zdiff(kxy,kxy)
+real(kind=jprd) :: zdx
+real(kind=jprd) :: zsum
+real(kind=jprd) :: zdiff(kxy,kxy)
 integer(kind=jpim)  :: jxy,jq,isize,jdist,ixy,ixym1,i1,i1pd,j1,j2
 logical :: llexit
 !-------------------------------------------------------------------------
@@ -535,11 +533,11 @@ recursive subroutine prepotf(kx,kxy,kquad,prw,prt,pr,pxy,kindex,prdexp,&
       kclosel(jxy)=kclosel(jxy)+1
       if((i1 > kx .and. i1pd <= kx) .or. (i1pd > kx .and.  i1 <= kx)) then
         knocik=knocik+1
-        zsum=0.0_jprb
+        zsum=0.0_jprd
         do jq=1,kquad
           zsum=zsum+prw(jq)*exp(-zdx*prt(jq))
         enddo
-        pcik(knocik)=1.0_jprb/zdx-zsum
+        pcik(knocik)=1.0_jprd/zdx-zsum
       endif
     else
       exit

src/trans/cpu/external/setup_trans.F90

@@ -337,11 +337,15 @@ SUBROUTINE SETUP_TRANS(KSMAX,KDGL,KDLON,KLOEN,LDSPLIT,PSTRET,&
   WRITE(NOUT,*) 'FFTW is now mandatory so this option is deprecated'
 ENDIF
 
+! Setup distribution independent dimensions
+CALL SETUP_DIMS
+
 S%LSOUTHPNM=.FALSE.
 IF(PRESENT(PSTRET)) THEN
   IF (ABS(PSTRET-1.0_JPRD)>100._JPRD*EPSILON(1._JPRD)) THEN
     G%RSTRET=PSTRET
     S%LSOUTHPNM=.TRUE.
+    R%NLEI3=2*R%NLEI3 ! double
   ENDIF
 ENDIF
 
@@ -388,9 +392,6 @@ SUBROUTINE SETUP_TRANS(KSMAX,KDGL,KDLON,KLOEN,LDSPLIT,PSTRET,&
 !     Setup resolution dependent structures
 !     -------------------------------------
 
-! Setup distribution independent dimensions
-CALL SETUP_DIMS
-
 ! First part of setup of distributed environment
 CALL SUMP_TRANS_PRELEG
 

src/trans/cpu/external/trans_pnm.F90

@@ -167,7 +167,7 @@ SUBROUTINE TRANS_PNM(KRESOL,KM,PRPNM,LDTRANSPOSE,LDCHEAP)
 
 !$OMP PARALLEL DO SCHEDULE(DYNAMIC,1) PRIVATE(JGL,ZLPOL,JI,JN)
 DO JGL=1,IDGLU
-  CALL SUPOLF(KM,INMAX,REAL (F%RMU(ISL+JGL-1), JPRD),ZLPOL(0:INMAX),KCHEAP=ICHEAP_ANTISYM)
+  CALL SUPOLF(KM,INMAX,F%RMU(ISL+JGL-1),ZLPOL(0:INMAX),KCHEAP=ICHEAP_ANTISYM)
   IF (LLTRANSPOSE) THEN
     DO JI=1,ILA
       PRPNM(IA+(JI-1)*2,ISL+JGL-1) = ZLPOL(KM+2*(ILA-JI)+1)
@@ -177,7 +177,7 @@ SUBROUTINE TRANS_PNM(KRESOL,KM,PRPNM,LDTRANSPOSE,LDCHEAP)
       PRPNM(ISL+JGL-1,IA+(JI-1)*2) = ZLPOL(KM+2*(ILA-JI)+1)
     ENDDO
   ENDIF
-  CALL SUPOLF(KM,INMAX,REAL (F%RMU(ISL+JGL-1), JPRD),ZLPOL(0:INMAX),KCHEAP=ICHEAP_SYM)
+  CALL SUPOLF(KM,INMAX,F%RMU(ISL+JGL-1),ZLPOL(0:INMAX),KCHEAP=ICHEAP_SYM)
   IF (LLTRANSPOSE) THEN
     DO JI=1,ILS
       PRPNM(IS+(JI-1)*2,ISL+JGL-1) = ZLPOL(KM+2*(ILS-JI))

src/trans/cpu/internal/fsc_mod.F90

@@ -73,15 +73,15 @@ SUBROUTINE FSC(KGL,KF_UV,KF_SCALARS,KF_SCDERS,&
 !     ------------------------------------------------------------------
 
 IGLG    = D%NPTRLS(MYSETW)+KGL-1
-ZACHTE  = F%RACTHE(IGLG)
+ZACHTE  = REAL(F%RACTHE(IGLG),JPRB)
 IMEN    = G%NMEN(IGLG)
 ISTAGTF = D%NSTAGTF(KGL)
-ZACHTE2  = F%RACTHE(IGLG)
+ZACHTE2 = REAL(F%RACTHE(IGLG),JPRB)
 
 IF( LATLON.AND.S%LDLL ) THEN
   ZPI = 2.0_JPRB*ASIN(1.0_JPRB)
   ZACHTE2 = 1._JPRB
-  ZACHTE  = F%RACTHE2(IGLG)
+  ZACHTE  = REAL(F%RACTHE2(IGLG),JPRB)
 
   ! apply shift for (even) lat-lon output grid
   IF( S%LSHIFTLL ) THEN

src/trans/cpu/internal/fscad_mod.F90

@@ -72,7 +72,7 @@ SUBROUTINE FSCAD(KGL,KF_UV,KF_SCALARS,KF_SCDERS,&
 !     ------------------------------------------------------------------
 
 IGLG = D%NPTRLS(MYSETW)+KGL-1
-ZACHTE  = F%RACTHE(IGLG)
+ZACHTE  = REAL(F%RACTHE(IGLG),JPRB)
 IMEN    = G%NMEN(IGLG)
 ISTAGTF = D%NSTAGTF(KGL)
 

src/trans/cpu/internal/gpnorm_trans_ctl_mod.F90

@@ -80,7 +80,7 @@ SUBROUTINE GPNORM_TRANS_CTL(PGP,KFIELDS,KPROMA,PAVE,PMIN,PMAX,LDAVE_ONLY,PW)
 INTEGER(KIND=JPIM),INTENT(IN)    :: KFIELDS
 INTEGER(KIND=JPIM),INTENT(IN)    :: KPROMA
 LOGICAL           ,INTENT(IN)    :: LDAVE_ONLY
-REAL(KIND=JPRB)   ,INTENT(IN)    :: PW(R%NDGL)
+REAL(KIND=JPRD)   ,INTENT(IN)    :: PW(R%NDGL)
 
 !ifndef INTERFACE
 
@@ -216,7 +216,7 @@ SUBROUTINE GPNORM_TRANS_CTL(PGP,KFIELDS,KPROMA,PAVE,PMIN,PMAX,LDAVE_ONLY,PW)
   DO JGL=IBEG,IEND
     IGL = D%NPTRLS(MYSETW) + JGL - 1
     DO JF=1,IF_FS
-      ZAVE(JF,JGL)=ZAVE(JF,JGL)*PW(IGL)/G%NLOEN(IGL)
+      ZAVE(JF,JGL)=ZAVE(JF,JGL)*REAL(PW(IGL),JPRB)/G%NLOEN(IGL)
     ENDDO
   ENDDO
 

src/trans/cpu/internal/ldfou2_mod.F90

@@ -70,6 +70,9 @@ SUBROUTINE LDFOU2(KM,KF_UV,PAIA,PSIA)
 
 REAL(KIND=JPRB) ,INTENT(INOUT) :: PSIA(:,:),   PAIA(:,:)
 
+!     LOCAL REAL SCALARS
+REAL(KIND=JPRB)   :: ZACTHE
+
 !     LOCAL INTEGER SCALARS
 INTEGER(KIND=JPIM) :: J, JGL ,IFLD ,ISL
 
@@ -85,9 +88,10 @@ SUBROUTINE LDFOU2(KM,KF_UV,PAIA,PSIA)
 !*       1.1      U AND V
 
 DO JGL=ISL,R%NDGNH
+  ZACTHE = REAL(F%RACTHE(JGL),JPRB)
   DO J=1,IFLD
-    PAIA(J,JGL) = PAIA(J,JGL)*F%RACTHE(JGL)
-    PSIA(J,JGL) = PSIA(J,JGL)*F%RACTHE(JGL)
+    PAIA(J,JGL) = PAIA(J,JGL)*ZACTHE
+    PSIA(J,JGL) = PSIA(J,JGL)*ZACTHE
   ENDDO
 ENDDO