#define ALIGN(I, A) (((I)+(A)-1)/(A)*(A)) ! (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 LTINVAD_MOD USE BUFFERED_ALLOCATOR_MOD, ONLY: ALLOCATION_RESERVATION_HANDLE IMPLICIT NONE PRIVATE PUBLIC :: LTINVAD, LTINVAD_HANDLE, PREPARE_LTINVAD TYPE LTINVAD_HANDLE TYPE(ALLOCATION_RESERVATION_HANDLE) :: HPIA_AND_IN END TYPE CONTAINS FUNCTION PREPARE_LTINVAD(ALLOCATOR,KF_UV,KF_SCALARS,LVORGP,LDIVGP,LSCDERS) RESULT(HLTINVAD) USE PARKIND_ECTRANS, ONLY: JPIM, JPRBT, JPRD, JPIB USE TPM_DISTR, ONLY: D USE TPM_DIM, ONLY: R USE ISO_C_BINDING, ONLY: C_SIZE_T, C_SIZEOF USE LEINV_MOD, ONLY: LEINV_STRIDES USE BUFFERED_ALLOCATOR_MOD, ONLY: BUFFERED_ALLOCATOR, RESERVE IMPLICIT NONE TYPE(BUFFERED_ALLOCATOR), INTENT(INOUT) :: ALLOCATOR INTEGER(KIND=JPIM), INTENT(IN) :: KF_UV,KF_SCALARS LOGICAL, INTENT(IN) :: LVORGP,LDIVGP,LSCDERS TYPE(LTINVAD_HANDLE) :: HLTINVAD INTEGER(KIND=JPIB) :: IALLOC_SZ, IPIA_SZ INTEGER(KIND=JPIM) :: IOUT_STRIDES0 INTEGER(KIND=JPIB) :: IOUT_SIZE INTEGER(KIND=JPIM) :: IOUT0_STRIDES0, IOUT0_SIZE INTEGER(KIND=JPIM) :: IIN_STRIDES0 INTEGER(KIND=JPIB) :: IIN_SIZE INTEGER(KIND=JPIM) :: IIN0_STRIDES0, IIN0_SIZE REAL(KIND=JPRBT) :: ZPRBT_DUMMY REAL(KIND=JPRD) :: ZPRD_DUMMY INTEGER(KIND=JPIM) :: IF_READIN, IF_LEG ! # fields that are initially read. We always read vorticity ! and divergence! Also keep in mind that we actually have 2X ! this number of levels because real+complex IF_READIN = 0 IF_READIN = IF_READIN + KF_UV ! Vorticity or divergence IF_READIN = IF_READIN + KF_UV ! Vorticity or divergence IF_READIN = IF_READIN + KF_UV ! U IF_READIN = IF_READIN + KF_UV ! V IF_READIN = IF_READIN + KF_SCALARS ! Scalars IF (LSCDERS) & IF_READIN = IF_READIN + KF_SCALARS ! Scalars NS Derivatives IPIA_SZ = ALIGN(2_JPIB*IF_READIN*(R%NSMAX+3)*D%NUMP*C_SIZEOF(ZPRBT_DUMMY),128) ! In Legendre space, we then ignore vorticity/divergence, if ! they don't need to be transformed. IF_LEG = IF_READIN IF(.NOT. LVORGP) IF_LEG = IF_LEG - KF_UV ! No vorticity needed IF(.NOT. LDIVGP) IF_LEG = IF_LEG - KF_UV ! No divergence needed CALL LEINV_STRIDES(IF_LEG,IOUT_STRIDES0,IOUT_SIZE,IIN_STRIDES0,IIN_SIZE,& IOUT0_STRIDES0,IOUT0_SIZE,IIN0_STRIDES0,IIN0_SIZE) ! PIA IALLOC_SZ = IPIA_SZ ! ZINP IALLOC_SZ = IALLOC_SZ + ALIGN(IIN_SIZE*C_SIZEOF(ZPRBT_DUMMY),128) ! ZINP0 IALLOC_SZ = IALLOC_SZ + ALIGN(IIN0_SIZE*C_SIZEOF(ZPRD_DUMMY),128) HLTINVAD%HPIA_AND_IN = RESERVE(ALLOCATOR, IALLOC_SZ, "HLTINVAD_HPIA_AND_IN") END FUNCTION PREPARE_LTINVAD SUBROUTINE LTINVAD(ALLOCATOR,HLTINVAD,KF_UV,KF_SCALARS,& & PSPVOR,PSPDIV,PSPSCALAR,PSPSC3A,PSPSC3B,PSPSC2, & & ZOUTS,ZOUTA,ZOUTS0,ZOUTA0) USE PARKIND_ECTRANS, ONLY: JPIM, JPRB, JPRBT, JPRD, JPIB USE YOMHOOK, ONLY: LHOOK, DR_HOOK, JPHOOK USE TPM_DIM, ONLY: R USE TPM_TRANS, ONLY: LDIVGP, LVORGP, NF_SC2, NF_SC3A, NF_SC3B, LSCDERS USE TPM_GEOMETRY, ONLY: G USE BUFFERED_ALLOCATOR_MOD, ONLY: BUFFERED_ALLOCATOR, ASSIGN_PTR, GET_ALLOCATION USE TPM_DISTR, ONLY: D USE PRFI1BAD_MOD, ONLY: PRFI1BAD USE VDTUVAD_MOD, ONLY: VDTUVAD USE SPNSDEAD_MOD, ONLY: SPNSDEAD USE LEINV_MOD, ONLY: LEINV_STRIDES USE LEDIR_MOD, ONLY: LEDIR USE ABORT_TRANS_MOD, ONLY: ABORT_TRANS USE TPM_FIELDS_GPU, ONLY: FG USE MPL_MODULE, ONLY: MPL_BARRIER,MPL_ALL_MS_COMM USE TPM_GEN, ONLY: LSYNC_TRANS USE TPM_STATS, ONLY: GSTATS => GSTATS_NVTX USE ISO_C_BINDING, ONLY: C_SIZE_T, C_LOC, C_SIZEOF !**** *LTINVAD* - adjoint of inverse Legendre transform ! ! Purpose. ! -------- ! Adjoint of the "tranform from Laplace space to Fourier space, compute U and V ! and north/south derivatives of state variables". !** Interface. ! ---------- ! *CALL* *LTINVAD(...) ! Explicit arguments : ! -------------------- ! KM - zonal wavenumber ! KMLOC - local zonal wavenumber ! PSPVOR - spectral vorticity ! PSPDIV - spectral divergence ! PSPSCALAR - spectral scalar variables ! Implicit arguments : The Laplace arrays of the model. ! -------------------- The values of the Legendre polynomials ! The grid point arrays of the model ! Method. ! ------- ! Externals. ! ---------- ! PREPSNM - prepare REPSNM for wavenumber KM ! PRFI1B - prepares the spectral fields ! VDTUV - compute u and v from vorticity and divergence ! SPNSDE - compute north-south derivatives ! LEINV - Inverse Legendre transform ! Reference. ! ---------- ! ECMWF Research Department documentation of the IFS ! Temperton, 1991, MWR 119 p1303 ! Author. ! ------- ! Mats Hamrud *ECMWF* ! Modifications. ! -------------- ! Original : 00-02-01 From LTINV in IFS CY22R1 ! ------------------------------------------------------------------ IMPLICIT NONE INTEGER(KIND=JPIM), INTENT(IN) :: KF_UV INTEGER(KIND=JPIM), INTENT(IN) :: KF_SCALARS REAL(KIND=JPRB) ,OPTIONAL,INTENT(INOUT) :: PSPVOR(:,:) REAL(KIND=JPRB) ,OPTIONAL,INTENT(INOUT) :: PSPDIV(:,:) REAL(KIND=JPRB) ,OPTIONAL,INTENT(INOUT) :: PSPSCALAR(:,:) REAL(KIND=JPRB) ,OPTIONAL,INTENT(INOUT) :: PSPSC2(:,:) REAL(KIND=JPRB) ,OPTIONAL,INTENT(INOUT) :: PSPSC3A(:,:,:) REAL(KIND=JPRB) ,OPTIONAL,INTENT(INOUT) :: PSPSC3B(:,:,:) REAL(KIND=JPRBT) , INTENT(IN) :: ZOUTS(:), ZOUTA(:) REAL(KIND=JPRD) , INTENT(IN) :: ZOUTS0(:), ZOUTA0(:) INTEGER(KIND=JPIM) :: IFIRST, J3 REAL(KIND=JPRB), POINTER :: PIA_L(:), PIA(:,:,:) REAL(KIND=JPRB), POINTER :: PU(:,:,:), PV(:,:,:), PVOR(:,:,:), PDIV(:,:,:) REAL(KIND=JPRB), POINTER :: PSCALARS(:,:,:), PSCALARS_NSDER(:,:,:) REAL(KIND=JPHOOK) :: ZHOOK_HANDLE TYPE(BUFFERED_ALLOCATOR), INTENT(IN) :: ALLOCATOR TYPE(LTINVAD_HANDLE), INTENT(IN) :: HLTINVAD INTEGER(KIND=JPIM) :: IOUT_STRIDES0 INTEGER(KIND=JPIB) :: IOUT_SIZE INTEGER(KIND=JPIM) :: IOUT0_STRIDES0, IOUT0_SIZE INTEGER(KIND=JPIM) :: IIN_STRIDES0 INTEGER(KIND=JPIB) :: IIN_SIZE INTEGER(KIND=JPIM) :: IIN0_STRIDES0, IIN0_SIZE INTEGER(KIND=JPIM) :: JK, J, KMLOC, KM INTEGER(KIND=JPIM) :: IF_READIN, IF_LEG INTEGER(KIND=JPIB) :: IALLOC_POS, IALLOC_SZ REAL(KIND=JPRBT), POINTER :: ZINP(:) REAL(KIND=JPRD), POINTER :: ZINP0(:) REAL(KIND=JPRB), POINTER :: PIA_LEDIR(:,:,:) ASSOCIATE(ZEPSNM=>FG%ZEPSNM, D_NUMP=>D%NUMP, D_MYMS=>D%MYMS) ! ------------------------------------------------------------------ !* 1. PERFORM LEGENDRE TRANFORM. ! -------------------------- IF (LHOOK) CALL DR_HOOK('LTINVAD_MOD',0,ZHOOK_HANDLE) ! Get all pointers IF_READIN = 0 IF_READIN = IF_READIN + KF_UV ! Vorticity or divergence IF_READIN = IF_READIN + KF_UV ! Vorticity or divergence IF_READIN = IF_READIN + KF_UV ! U IF_READIN = IF_READIN + KF_UV ! V IF_READIN = IF_READIN + KF_SCALARS ! Scalars IF (LSCDERS) & IF_READIN = IF_READIN + KF_SCALARS ! Scalars NS Derivatives ! In Legendre space, we then ignore vorticity/divergence, if ! they don't need to be transformed. IF_LEG = IF_READIN IF(.NOT. LVORGP) IF_LEG = IF_LEG - KF_UV ! No vorticity needed IF(.NOT. LDIVGP) IF_LEG = IF_LEG - KF_UV ! No divergence needed CALL LEINV_STRIDES(IF_LEG,IOUT_STRIDES0,IOUT_SIZE,IIN_STRIDES0,IIN_SIZE,& IOUT0_STRIDES0,IOUT0_SIZE,IIN0_STRIDES0,IIN0_SIZE) IALLOC_POS = 1 ! PIA IALLOC_SZ = ALIGN(2_JPIB*IF_READIN*(R%NTMAX+3)*D%NUMP*C_SIZEOF(PIA_L(1)),128) CALL ASSIGN_PTR(PIA_L, GET_ALLOCATION(ALLOCATOR, HLTINVAD%HPIA_AND_IN),& & IALLOC_POS, IALLOC_SZ) CALL C_F_POINTER(C_LOC(PIA_L), PIA, (/ 2*IF_READIN, R%NTMAX+3, D%NUMP /)) IALLOC_POS = IALLOC_POS + IALLOC_SZ ! ZINP IALLOC_SZ = ALIGN(IIN_SIZE*C_SIZEOF(ZINP(1)),128) CALL ASSIGN_PTR(ZINP, GET_ALLOCATION(ALLOCATOR, HLTINVAD%HPIA_AND_IN),& & IALLOC_POS, IALLOC_SZ) IALLOC_POS = IALLOC_POS + IALLOC_SZ ! ZINP0 IALLOC_SZ = ALIGN(IIN0_SIZE*C_SIZEOF(ZINP0(1)),128) CALL ASSIGN_PTR(ZINP0, GET_ALLOCATION(ALLOCATOR, HLTINVAD%HPIA_AND_IN),& & IALLOC_POS, IALLOC_SZ) IALLOC_POS = IALLOC_POS + IALLOC_SZ ! Assign pointers do the different components of PIA IFIRST = 0 IF (.NOT. LVORGP .OR. LDIVGP) THEN ! Usually we want to store vorticity first PVOR => PIA(IFIRST+1:IFIRST+2*KF_UV,:,:) IFIRST = IFIRST + 2*KF_UV ! Vorticity PDIV => PIA(IFIRST+1:IFIRST+2*KF_UV,:,:) IFIRST = IFIRST + 2*KF_UV ! Divergence ELSE ! Except if we want to translate Vorticity but not Divergence, we should have Divergence first ! Then we have all buffers that move on in a contiguous buffer PDIV => PIA(IFIRST+1:IFIRST+2*KF_UV,:,:) IFIRST = IFIRST + 2*KF_UV ! Divergence PVOR => PIA(IFIRST+1:IFIRST+2*KF_UV,:,:) IFIRST = IFIRST + 2*KF_UV ! Vorticity ENDIF PU => PIA(IFIRST+1:IFIRST+2*KF_UV,:,:) IFIRST = IFIRST + 2*KF_UV ! U PV => PIA(IFIRST+1:IFIRST+2*KF_UV,:,:) IFIRST = IFIRST + 2*KF_UV ! V PSCALARS => PIA(IFIRST+1:IFIRST+2*KF_SCALARS,:,:) IFIRST = IFIRST + 2*KF_SCALARS ! Scalars IF (LSCDERS) THEN PSCALARS_NSDER => PIA(IFIRST+1:IFIRST+2*KF_SCALARS,:,:) IFIRST = IFIRST + 2*KF_SCALARS ! Scalars NS Derivatives ENDIF ! ------------------------------------------------------------------ !* 4. Adjoint of INVERSE LEGENDRE TRANSFORM. ! --------------------------- ! Legendre transforms. When converting PIA into ZOUT, we ignore the first entries of LEINV. ! This is because vorticity and divergence is not necessarily converted to GP space. PIA_LEDIR => PIA(2*(IF_READIN-IF_LEG)+1:IF_READIN,:,:) CALL LEDIR(ALLOCATOR,ZOUTS,ZOUTA,ZOUTS0,ZOUTA0,ZINP,ZINP0,PIA_LEDIR,IF_LEG) #ifdef OMPGPU !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) DEFAULT(NONE) SHARED(D,PIA_LEDIR,IF_LEG) PRIVATE(KM) MAP(TO:IF_LEG) #endif #ifdef ACCGPU !$ACC DATA PRESENT(D,D_MYMS,D_NUMP,PIA_LEDIR) !$ACC PARALLEL LOOP COLLAPSE(2) PRIVATE(KMLOC,KM,J) & !$ACC& FIRSTPRIVATE(IF_LEG) DEFAULT(NONE) #endif DO KMLOC=1,D_NUMP DO JK=1,IF_LEG KM = D_MYMS(KMLOC) IF(KM .eq. 0)THEN #ifdef ACCGPU !$ACC LOOP SEQ #endif DO J=1,SIZE(PIA_LEDIR,2) PIA_LEDIR(2*JK,J,KMLOC) = 0.0_JPRB ENDDO ENDIF ENDDO ENDDO #ifdef ACCGPU !$ACC END DATA #endif ! ------------------------------------------------------------------ !* 3. Adjoint of SPECTRAL COMPUTATIONS FOR U,V AND DERIVATIVES. ! ---------------------------------------------- #ifdef OMPGPU !$OMP TARGET DATA MAP(FROM:PSPVOR,PSPDIV) IF(KF_UV > 0) !$OMP TARGET DATA MAP(FROM:PSPSCALAR) IF(PRESENT(PSPSCALAR) .AND. KF_SCALARS > 0) !$OMP TARGET DATA MAP(FROM:PSPSC2) IF(NF_SC2 > 0) !$OMP TARGET DATA MAP(FROM:PSPSC3A) IF(NF_SC3A > 0) !$OMP TARGET DATA MAP(FROM:PSPSC3B) IF(NF_SC3B > 0) #endif #ifdef ACCGPU !$ACC DATA COPYOUT(PSPVOR,PSPDIV) IF(KF_UV > 0) !$ACC DATA COPYOUT(PSPSCALAR) IF(PRESENT(PSPSCALAR) .AND. KF_SCALARS > 0) !$ACC DATA COPYOUT(PSPSC2) IF(NF_SC2 > 0) !$ACC DATA COPYOUT(PSPSC3A) IF(NF_SC3A > 0) !$ACC DATA COPYOUT(PSPSC3B) IF(NF_SC3B > 0) #endif ! Compute NS derivatives if needed IF (LSCDERS) THEN CALL SPNSDEAD(KF_SCALARS,ZEPSNM,PSCALARS,PSCALARS_NSDER) ENDIF IF (KF_SCALARS > 0) THEN IF(PRESENT(PSPSCALAR)) THEN CALL PRFI1BAD(PSCALARS,PSPSCALAR,KF_SCALARS,UBOUND(PSPSCALAR,2)) ELSE IFIRST = 1 IF(PRESENT(PSPSC2) .AND. NF_SC2 > 0) THEN CALL PRFI1BAD(PSCALARS(IFIRST:IFIRST+2*NF_SC2-1,:,:),PSPSC2(:,:),NF_SC2,UBOUND(PSPSC2,2)) IFIRST = IFIRST+2*NF_SC2 ENDIF IF(PRESENT(PSPSC3A) .AND. NF_SC3A > 0) THEN DO J3=1,UBOUND(PSPSC3A,3) CALL PRFI1BAD(PSCALARS(IFIRST:IFIRST+2*NF_SC3A-1,:,:),PSPSC3A(:,:,J3),NF_SC3A,UBOUND(PSPSC3A,2)) IFIRST = IFIRST+2*NF_SC3A ENDDO ENDIF IF(PRESENT(PSPSC3B) .AND. NF_SC3B > 0) THEN DO J3=1,UBOUND(PSPSC3B,3) CALL PRFI1BAD(PSCALARS(IFIRST:IFIRST+2*NF_SC3B-1,:,:),PSPSC3B(:,:,J3),NF_SC3B,UBOUND(PSPSC3B,2)) IFIRST = IFIRST+2*NF_SC3B ENDDO ENDIF IF(IFIRST-1 /= 2*KF_SCALARS) THEN WRITE(0,*) 'LTINV:KF_SCALARS,IFIRST',KF_SCALARS,IFIRST CALL ABORT_TRANS('LTINV_MOD:IFIRST /= 2*KF_SCALARS') ENDIF ENDIF ENDIF IF (KF_UV > 0) THEN ! Compute U and V for VOR and DIV CALL VDTUVAD(KF_UV,ZEPSNM,PVOR,PDIV,PU,PV) CALL PRFI1BAD(PVOR,PSPVOR,KF_UV,UBOUND(PSPVOR,2)) CALL PRFI1BAD(PDIV,PSPDIV,KF_UV,UBOUND(PSPDIV,2)) ENDIF IF (LSYNC_TRANS) THEN CALL GSTATS(440,0) CALL MPL_BARRIER(MPL_ALL_MS_COMM,CDSTRING='') CALL GSTATS(440,1) ENDIF CALL GSTATS(422,0) #ifdef OMPGPU !$OMP END TARGET DATA !$OMP END TARGET DATA !$OMP END TARGET DATA !$OMP END TARGET DATA !$OMP END TARGET DATA #endif #ifdef ACCGPU !$ACC WAIT(1) !$ACC END DATA !$ACC END DATA !$ACC END DATA !$ACC END DATA !$ACC END DATA #endif IF (LSYNC_TRANS) THEN CALL GSTATS(442,0) CALL MPL_BARRIER(MPL_ALL_MS_COMM,CDSTRING='') CALL GSTATS(442,1) ENDIF CALL GSTATS(422,1) IF (LHOOK) CALL DR_HOOK('LTINVAD_MOD',1,ZHOOK_HANDLE) END ASSOCIATE ! ------------------------------------------------------------------ END SUBROUTINE LTINVAD END MODULE LTINVAD_MOD