#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 TRLTOMAD_PACK_UNPACK USE BUFFERED_ALLOCATOR_MOD, ONLY: ALLOCATION_RESERVATION_HANDLE USE PARKIND_ECTRANS, ONLY: JPIM IMPLICIT NONE PRIVATE PUBLIC :: TRLTOMAD_PACK_HANDLE, PREPARE_TRLTOMAD_PACK, TRLTOMAD_PACK PUBLIC :: TRLTOMAD_UNPACK_HANDLE, PREPARE_TRLTOMAD_UNPACK, TRLTOMAD_UNPACK TYPE TRLTOMAD_PACK_HANDLE TYPE(ALLOCATION_RESERVATION_HANDLE) :: HREEL_COMPLEX END TYPE TYPE TRLTOMAD_UNPACK_HANDLE TYPE(ALLOCATION_RESERVATION_HANDLE) :: HPFBUF END TYPE INTEGER(KIND=JPIM) :: A = 8 !Alignment CONTAINS FUNCTION PREPARE_TRLTOMAD_PACK(ALLOCATOR, KF_FS) RESULT(HTRLTOM_PACK) USE PARKIND_ECTRANS, ONLY: JPIM, JPRBT, JPIB USE TPM_DISTR, ONLY: D USE ISO_C_BINDING, ONLY: C_SIZE_T, C_SIZEOF USE BUFFERED_ALLOCATOR_MOD, ONLY: BUFFERED_ALLOCATOR, RESERVE IMPLICIT NONE TYPE(BUFFERED_ALLOCATOR), INTENT(INOUT) :: ALLOCATOR INTEGER(KIND=JPIM), INTENT(IN) :: KF_FS TYPE(TRLTOMAD_PACK_HANDLE) :: HTRLTOM_PACK REAL(KIND=JPRBT) :: DUMMY HTRLTOM_PACK%HREEL_COMPLEX = RESERVE(ALLOCATOR, 1_JPIB*KF_FS*D%NLENGTF*C_SIZEOF(DUMMY), "HTRLTOMAD_PACK%HREEL_COMPLEX") END FUNCTION PREPARE_TRLTOMAD_PACK SUBROUTINE TRLTOMAD_PACK(ALLOCATOR,HTRLTOM_PACK,PREEL_COMPLEX,FOUBUF_IN,KF_FS) !**** *TRLTOMAD_PACK* - Copy fourier data from local array to buffer ! Purpose. ! -------- ! Routine for copying fourier data from local array to buffer !** Interface. ! ---------- ! CALL TRLTOM_PACK(...) ! Explicit arguments : PREEL - local fourier/GP array ! -------------------- KF_FS - number of fields ! ! Externals. None. ! ---------- ! Author. ! ------- ! Mats Hamrud *ECMWF* ! ------------------------------------------------------------------ USE BUFFERED_ALLOCATOR_MOD, ONLY: BUFFERED_ALLOCATOR, ASSIGN_PTR, GET_ALLOCATION USE PARKIND_ECTRANS, ONLY: JPIM, JPRBT, JPIB USE TPM_DISTR, ONLY: D, MYSETW USE TPM_GEOMETRY, ONLY: G USE TPM_DIM, ONLY: R USE ISO_C_BINDING, ONLY: C_SIZE_T, C_SIZEOF ! IMPLICIT NONE REAL(KIND=JPRBT), POINTER, INTENT(OUT) :: PREEL_COMPLEX(:) REAL(KIND=JPRBT), INTENT(IN) :: FOUBUF_IN(:) INTEGER(KIND=JPIM),INTENT(IN) :: KF_FS TYPE(BUFFERED_ALLOCATOR), INTENT(IN) :: ALLOCATOR TYPE(TRLTOMAD_PACK_HANDLE), INTENT(IN) :: HTRLTOM_PACK INTEGER(KIND=JPIM) :: JM,JF,IGLG,OFFSET_VAR,KGL,J,N INTEGER(KIND=JPIB) :: IOFF_LAT,ISTA REAL(KIND=JPRBT) :: SCAL ASSOCIATE(D_NSTAGTF=>D%NSTAGTF, D_NPNTGTB0=>D%NPNTGTB0, D_NPTRLS=>D%NPTRLS, & & D_NDGL_FS=>D%NDGL_FS, G_NMEN=>G%NMEN, G_NLOEN=>G%NLOEN, R_NSMAX=>R%NSMAX) CALL ASSIGN_PTR(PREEL_COMPLEX, GET_ALLOCATION(ALLOCATOR, HTRLTOM_PACK%HREEL_COMPLEX),& & 1_JPIB, 1_JPIB*KF_FS*D%NLENGTF*C_SIZEOF(PREEL_COMPLEX(1))) N = 1_JPIB*KF_FS*D%NLENGTF #ifdef ACCGPU !$ACC DATA PRESENT(PREEL_COMPLEX) !$ACC PARALLEL LOOP FIRSTPRIVATE(N) #endif #ifdef OMPGPU !$OMP TARGET DATA MAP(PRESENT,ALLOC:PREEL_COMPLEX) !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO DEFAULT(NONE) SHARED(N,PREEL_COMPLEX) #endif DO J=1,N PREEL_COMPLEX(J) = 0 ENDDO #ifdef ACCGPU !$ACC END DATA #endif #ifdef OMPGPU !$OMP END TARGET DATA #endif #ifdef OMPGPU !$OMP TARGET DATA MAP(PRESENT,ALLOC:G,G_NMEN,D,D_NPNTGTB0,FOUBUF_IN,PREEL_COMPLEX,D_NSTAGTF,& !$OMP& D_NDGL_FS,G_NLOEN,R,R_NSMAX) #endif #ifdef ACCGPU !$ACC DATA PRESENT(G,G_NMEN,D,D_NPNTGTB0,FOUBUF_IN,PREEL_COMPLEX,D_NSTAGTF,D_NDGL_FS,G_NLOEN, R,R_NSMAX) ASYNC(1) #endif ! scale results and move into next transformation buffer OFFSET_VAR=D_NPTRLS(MYSETW) #ifdef OMPGPU !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) PRIVATE(IGLG,IOFF_LAT,ISTA,SCAL) & !$OMP& SHARED(D,R,KF_FS,OFFSET_VAR,G,& !$OMP& PREEL_COMPLEX,FOUBUF_IN) MAP(TO:KF_FS,OFFSET_VAR) #endif #ifdef ACCGPU !$ACC PARALLEL LOOP PRIVATE(IGLG,IOFF_LAT,ISTA,SCAL) FIRSTPRIVATE(KF_FS,OFFSET_VAR) & !$ACC& TILE(32,16,1) DEFAULT(NONE) & #ifndef _CRAYFTN !$ACC& ASYNC(1) #else !$ACC& #endif #endif DO KGL=1,D_NDGL_FS DO JM=0,R_NSMAX !(note that R_NSMAX <= G_NMEN(IGLG) for all IGLG) DO JF=1,KF_FS IGLG = OFFSET_VAR+KGL-1 IF (JM <= G_NMEN(IGLG)) THEN IOFF_LAT = KF_FS*D_NSTAGTF(KGL)+(JF-1)*(D_NSTAGTF(KGL+1)-D_NSTAGTF(KGL)) SCAL = 1._JPRBT/REAL(G_NLOEN(IGLG),JPRBT) ISTA = 2_JPIB*D_NPNTGTB0(JM,KGL)*KF_FS PREEL_COMPLEX(IOFF_LAT+2*JM+1) = SCAL * FOUBUF_IN(ISTA+2*JF-1) PREEL_COMPLEX(IOFF_LAT+2*JM+2) = SCAL * FOUBUF_IN(ISTA+2*JF ) ENDIF ENDDO ENDDO ENDDO #ifdef ACCGPU !$ACC END DATA !$ACC WAIT(1) #endif #ifdef OMPGPU !$OMP END TARGET DATA #endif END ASSOCIATE END SUBROUTINE TRLTOMAD_PACK FUNCTION PREPARE_TRLTOMAD_UNPACK(ALLOCATOR, KF_FS) RESULT(HTRLTOM_UNPACK) USE PARKIND_ECTRANS, ONLY: JPIM, JPRBT, JPRD, JPIB USE BUFFERED_ALLOCATOR_MOD, ONLY: BUFFERED_ALLOCATOR, RESERVE USE LEDIR_MOD, ONLY: LEDIR_STRIDES USE ISO_C_BINDING, ONLY: C_SIZE_T, C_SIZEOF USE TPM_DISTR, ONLY: D IMPLICIT NONE TYPE(BUFFERED_ALLOCATOR), INTENT(INOUT) :: ALLOCATOR INTEGER(KIND=JPIM), INTENT(IN) :: KF_FS TYPE(TRLTOMAD_UNPACK_HANDLE) :: HTRLTOM_UNPACK INTEGER(KIND=JPIM) :: IIN_STRIDES0 INTEGER(KIND=JPIB) :: IIN_SIZE INTEGER(KIND=JPIM) :: IIN0_STRIDES0, IIN0_SIZE INTEGER(KIND=JPIB) :: ISIZE REAL(KIND=JPRBT) :: DUMMY CALL LEDIR_STRIDES(KF_FS,IIN_STRIDES0=IIN_STRIDES0,IIN_SIZE=IIN_SIZE,& IIN0_STRIDES0=IIN0_STRIDES0,IIN0_SIZE=IIN0_SIZE) HTRLTOM_UNPACK%HPFBUF = RESERVE(ALLOCATOR, 2_JPIB*D%NLENGT1B*KF_FS*C_SIZEOF(DUMMY), "HTRLTOM_UNPACK%HPFBUF") END FUNCTION PREPARE_TRLTOMAD_UNPACK SUBROUTINE TRLTOMAD_UNPACK(ALLOCATOR,HTRLTOM_UNPACK,FOUBUF,ZINPS,ZINPA,ZINPS0,ZINPA0,KF_FS,KF_UV) USE PARKIND_ECTRANS, ONLY: JPIM, JPRBT, JPRD, JPIB USE TPM_DIM, ONLY: R USE TPM_GEOMETRY, ONLY: G USE BUFFERED_ALLOCATOR_MOD, ONLY: BUFFERED_ALLOCATOR, ASSIGN_PTR, GET_ALLOCATION USE TPM_FIELDS, ONLY: F USE TPM_DISTR, ONLY: D USE LEDIR_MOD, ONLY: LEDIR_STRIDES USE, INTRINSIC :: ISO_C_BINDING, ONLY: C_SIZE_T, C_SIZEOF IMPLICIT NONE REAL(KIND=JPRBT), POINTER, INTENT(OUT) :: FOUBUF(:) REAL(KIND=JPRBT), INTENT(IN) :: ZINPS(:), ZINPA(:) REAL(KIND=JPRD), INTENT(IN) :: ZINPS0(:), ZINPA0(:) INTEGER(KIND=JPIM), INTENT(IN) :: KF_FS, KF_UV TYPE(BUFFERED_ALLOCATOR), INTENT(IN) :: ALLOCATOR TYPE(TRLTOMAD_UNPACK_HANDLE), INTENT(IN) :: HTRLTOM_UNPACK REAL(KIND=JPRBT), POINTER :: PREEL_COMPLEX(:) INTEGER(KIND=JPIM) :: IIN_STRIDES0 INTEGER(KIND=JPIB) :: IIN_SIZE INTEGER(KIND=JPIM) :: IIN0_STRIDES0, IIN0_SIZE INTEGER(KIND=JPIB) :: IALLOC_POS, IALLOC_SZ INTEGER(KIND=JPIB) :: JF, OFFSET1, OFFSET2 INTEGER(KIND=JPIM) :: KM, ISL, IGLS, JGL, KMLOC REAL(KIND=JPRBT) :: PAIA, PAIS ASSOCIATE(D_NUMP=>D%NUMP, R_NDGNH=>R%NDGNH, R_NDGL=>R%NDGL, F_RW=>F%RW, F_RACTHE=>F%RACTHE, & & D_MYMS=>D%MYMS, D_NPNTGTB1=>D%NPNTGTB1, D_OFFSETS_GEMM1=>D%OFFSETS_GEMM1, & & G_NDGLU=>G%NDGLU) CALL LEDIR_STRIDES(KF_FS,IIN_STRIDES0=IIN_STRIDES0,IIN_SIZE=IIN_SIZE,& IIN0_STRIDES0=IIN0_STRIDES0,IIN0_SIZE=IIN0_SIZE) CALL ASSIGN_PTR(FOUBUF, GET_ALLOCATION(ALLOCATOR, HTRLTOM_UNPACK%HPFBUF),& & 1_JPIB, 2_JPIB*D%NLENGT1B*KF_FS*C_SIZEOF(FOUBUF(1))) #ifdef OMPGPU !$OMP TARGET DATA MAP(PRESENT,ALLOC:ZINPS,ZINPA,ZINPS0,ZINPA0) & !$OMP& MAP(PRESENT,ALLOC:F,F_RW,F_RACTHE) & !$OMP& MAP(PRESENT,ALLOC:D,D_MYMS,D_NUMP,R,R_NDGNH,R_NDGL,G,G_NDGLU) & !$OMP& MAP(PRESENT,ALLOC:D_NPNTGTB1,D_OFFSETS_GEMM1,FOUBUF) !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) DEFAULT(NONE) & !$OMP& PRIVATE(KM,ISL,IGLS,OFFSET1,OFFSET2,PAIA,PAIS) & !$OMP& SHARED(D,R,KF_FS,G,FOUBUF,F,& !$OMP& IIN_STRIDES0,ZINPA,ZINPS,IIN0_STRIDES0,ZINPA0,ZINPS0,KF_UV) & !$OMP& MAP(TO:KF_FS,KF_UV,IIN_STRIDES0,IIN0_STRIDES0) #endif #ifdef ACCGPU !$ACC DATA & !$ACC& PRESENT(ZINPS,ZINPA,ZINPS0,ZINPA0) & !$ACC& PRESENT(F,F_RW,F_RACTHE) & !$ACC& PRESENT(D,D_MYMS,D_NUMP,R,R_NDGNH,R_NDGL,G,G_NDGLU) & !$ACC& PRESENT(D_NPNTGTB1,D_OFFSETS_GEMM1,FOUBUF) !$ACC PARALLEL LOOP DEFAULT(NONE) COLLAPSE(3) PRIVATE(KM,ISL,IGLS,OFFSET1,OFFSET2,JGL,PAIA,PAIS) & !$ACC& FIRSTPRIVATE(KF_FS,KF_UV,IIN_STRIDES0,IIN0_STRIDES0) & #ifndef _CRAYFTN !$ACC& ASYNC(1) #else !$ACC& #endif #endif DO KMLOC=1,D_NUMP DO JGL=1,R_NDGNH DO JF=1,KF_FS*2 KM = D_MYMS(KMLOC) ISL = R_NDGNH-G_NDGLU(KM)+1 IF (JGL >= ISL) THEN !(DO JGL=ISL,R_NDGNH) IGLS = R_NDGL+1-JGL OFFSET1 = 2_JPIB*D_NPNTGTB1(KMLOC,JGL )*KF_FS OFFSET2 = 2_JPIB*D_NPNTGTB1(KMLOC,IGLS)*KF_FS PAIA = 0 PAIS = 0 IF (KM /= 0) THEN PAIA=REAL(F_RW(JGL),JPRBT)*ZINPA(JF+(JGL-ISL)*IIN_STRIDES0+IIN_STRIDES0*D_OFFSETS_GEMM1(KMLOC)) PAIS=REAL(F_RW(JGL),JPRBT)*ZINPS(JF+(JGL-ISL)*IIN_STRIDES0+IIN_STRIDES0*D_OFFSETS_GEMM1(KMLOC)) ELSEIF (MOD(JF-1,2) == 0) THEN ! every other field is sufficient because Im(KM=0) == 0 PAIA=REAL(F_RW(JGL),JPRBT)*ZINPA0((JF-1)/2+1+(JGL-1)*IIN0_STRIDES0) PAIS=REAL(F_RW(JGL),JPRBT)*ZINPS0((JF-1)/2+1+(JGL-1)*IIN0_STRIDES0) ENDIF IF (JF <= 4*KF_UV) THEN ! Multiply in case of velocity PAIA = PAIA*REAL(F_RACTHE(JGL),JPRBT) PAIS = PAIS*REAL(F_RACTHE(JGL),JPRBT) ENDIF FOUBUF(OFFSET1+JF) = PAIA+PAIS FOUBUF(OFFSET2+JF) = -PAIA+PAIS ENDIF ENDDO ENDDO END DO #ifdef OMPGPU !$OMP END TARGET DATA #endif #ifdef ACCGPU !$ACC END DATA #endif END ASSOCIATE END SUBROUTINE TRLTOMAD_UNPACK END MODULE TRLTOMAD_PACK_UNPACK