! (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_DISTR ! Module for distributed memory environment. USE PARKIND_ECTRANS ,ONLY : JPIM ,JPRBT IMPLICIT NONE SAVE !* Variables describing distributed memory parallelization INTEGER(KIND=JPIM) :: NPROC ! Number of processors (NPRGPNS*NPRGPEW) INTEGER(KIND=JPIM) :: NPRGPNS ! No. of sets in N-S direction (grid-point space) INTEGER(KIND=JPIM) :: NPRGPEW ! No. of sets in E-W direction (grid-point space) INTEGER(KIND=JPIM) :: NPRTRW ! No. of sets in wave direction (spectral space) INTEGER(KIND=JPIM) :: NPRTRV ! NPROC/NPRTRW INTEGER(KIND=JPIM) :: NPRTRNS ! No. of sets in N-S direction (Fourier space) ! (always equal to NPRTRW) LOGICAL :: LEQ_REGIONS ! TRUE - Use new eq_regions partitioning ! FALSE- Use old NPRGPNS x NPRGPEW partitioning INTEGER(KIND=JPIM) :: MYPROC ! My processor number INTEGER(KIND=JPIM) :: MYSETW ! My set number in wave direction (spectral space) INTEGER(KIND=JPIM) :: MYSETV ! My set number in field direction(S.S and F.S) INTEGER(KIND=JPIM) :: NCOMBFLEN ! Size of communication buffer INTEGER(KIND=JPIM) :: MTAGLETR ! Tag INTEGER(KIND=JPIM) :: MTAGML ! Tag INTEGER(KIND=JPIM) :: MTAGLG ! Tag INTEGER(KIND=JPIM) :: MTAGGL ! Tag INTEGER(KIND=JPIM) :: MTAGPART ! Tag INTEGER(KIND=JPIM) :: MTAGDISTSP ! Tag INTEGER(KIND=JPIM) :: MTAGLM ! Tag INTEGER(KIND=JPIM) :: MTAGDISTGP ! Tag INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPRCIDS(:) ! Array containing the process ids TYPE DISTR_TYPE LOGICAL :: LGRIDONLY ! TRUE - only grid space structures are available LOGICAL :: LWEIGHTED_DISTR ! TRUE - weighted distribution LOGICAL :: LSPLIT ! TRUE - latitudes are shared between a-sets LOGICAL :: LCPNMONLY ! TRUE - Compute Legendre polynomials only, not FFTs ! SPECTRAL SPACE INTEGER(KIND=JPIM) :: NUMP ! No. of spectral waves handled by this processor INTEGER(KIND=JPIM) :: NSPEC ! No. of complex spectral coefficients (on this PE) INTEGER(KIND=JPIM) :: NSPEC2 ! 2*NSPEC INTEGER(KIND=JPIM) :: NSPEC2MX ! maximun NSPEC2 among all PEs INTEGER(KIND=JPIM) :: NTPEC2 ! cf. NSPEC2 but for truncation NTMAX INTEGER(KIND=JPIM) :: NUMTP ! cf. NUMP but for truncation NTMAX INTEGER(KIND=JPIM) :: NSPOLEGL ! No. of legendre polynomials on this PE INTEGER(KIND=JPIM) :: NLEI3D ! (NLEI3-1)/NPRTRW+1 INTEGER(KIND=JPIM) ,ALLOCATABLE :: MYMS(:) ! Wave numbers handled by this PE INTEGER(KIND=JPIM) ,ALLOCATABLE :: NUMPP(:) ! No. of wave numbers each wave set is ! responsible for INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPOSSP(:) ! Not needed in transform? INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPROCM(:) ! Process that does the calc. for certain ! wavenumber M INTEGER(KIND=JPIM) ,ALLOCATABLE :: NDIM0G(:) ! Defines partitioning of global spectral ! fields among PEs INTEGER(KIND=JPIM) ,ALLOCATABLE :: NASM0(:) ! Address in a spectral array of (m, n=m) INTEGER(KIND=JPIM) ,ALLOCATABLE :: NATM0(:) ! Same as NASM0 but for NTMAX INTEGER(KIND=JPIM) ,ALLOCATABLE :: NALLMS(:) ! Wave numbers for all a-set concatenated ! together to give all wave numbers in a-set ! order. Used when global spectral norms ! have to be gathered. INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPTRMS(:) ! Pointer to the first wave number of a given ! a-set in nallms array. ! Legendre polynomials INTEGER(KIND=JPIM) ,ALLOCATABLE :: NLATLS(:,:) ! First latitude for which each a-set,bset calcul. INTEGER(KIND=JPIM) ,ALLOCATABLE :: NLATLE(:,:) ! Last latitude for which each a-set,bset calcul. INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPMT(:) ! Adress for legendre polynomial for ! given M (NTMAX) INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPMS(:) ! Adress for legendre polynomial for ! given M (NSMAX) INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPMG(:) ! Global version of NPMS ! FOURIER SPACE INTEGER(KIND=JPIM) :: NDGL_FS ! Number of rows of latitudes for which this process is ! performing Fourier Space calculations INTEGER(KIND=JPIM) ,ALLOCATABLE :: NSTAGTF(:) ! Offset for specific latitude in ! Fourier/gridpoint buffer INTEGER(KIND=JPIM) :: NLENGTF ! Second dimension of Fourier/gridpoint buffer ! (sum of (NLOEN+3) over local latitudes) INTEGER(KIND=JPIM) ,ALLOCATABLE :: NULTPP(:) ! No of lats. for each wave_set (F.S) INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPROCL(:) ! Process responsible for each lat. (F.S) INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPTRLS(:) ! Pointer to first lat. (F.S) ! NSTAGT0B to NLENGT1B: help arrays for spectral to fourier space transposition INTEGER(KIND=JPIM) ,ALLOCATABLE :: NSTAGT0B(:) ! Start adresses for segments within buffer ! (according to processors to whom data ! is going to be sent) INTEGER(KIND=JPIM) ,ALLOCATABLE :: NSTAGT1B(:) INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPNTGTB0(:,:) INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPNTGTB1(:,:) INTEGER(KIND=JPIM) ,ALLOCATABLE :: NLTSFTB(:) INTEGER(KIND=JPIM) ,ALLOCATABLE :: NLTSGTB(:) INTEGER(KIND=JPIM) ,ALLOCATABLE :: MSTABF(:) INTEGER(KIND=JPIM) :: NLENGT0B ! dimension INTEGER(KIND=JPIM) :: NLENGT1B ! dimension ! GRIDPOINT SPACE INTEGER(KIND=JPIM) :: NDGL_GP ! D%NLSTLAT(MY_REGION_NS)-D%NFRSTLOFF INTEGER(KIND=JPIM) ,ALLOCATABLE :: NFRSTLAT(:) ! First lat of each a-set INTEGER(KIND=JPIM) ,ALLOCATABLE :: NLSTLAT(:) ! Last lat of each a-set INTEGER(KIND=JPIM) :: NFRSTLOFF ! Offset for first lat of own a-set ! i.e. NFRSTLOFF=NFRSTLAT(MYSETA)-1 INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPTRLAT(:) ! Pointer to start of latitude INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPTRFRSTLAT(:) ! Pointer to the first latitude of each ! a-set in NSTA and NONL arrays INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPTRLSTLAT(:) ! Pointer to the last latitude of each ! a-set in NSTA and NONL arrays INTEGER(KIND=JPIM) :: NPTRFLOFF ! Offset for pointer to the first latitude of own a-set ! NSTA and NONL arrays, i.e. NPTRFRSTLAT(MYSETA)-1 LOGICAL ,ALLOCATABLE :: LSPLITLAT(:) ! True if latitude is split over 2 a-sets ! NSTA(R%NDGL+NPRGPNS-1,NPRGPEW) : Position of first grid column ! for the latitudes on a processor. The information is ! available for all processors. The b-sets are distinguished ! by the last dimension of NSTA(). The latitude band for ! each a-set is addressed by NPTRFRSTLAT(JASET), ! NPTRLSTLAT(JASET), and NPTRFLOFF=NPTRFRSTLAT(MYSETA) on ! this processors a-set. Each split latitude has two entries ! in NSTA(,:) which necessitates the rather complex ! addressing of NSTA(,:) and the overdimensioning of NSTA by ! NPRGPNS. ! NONL(R%NDGL+NPRGPNS-1,NPRGPEW) : Number of grid columns for ! the latitudes on a processor. Similar to NSTA() in data ! structure. INTEGER(KIND=JPIM) ,ALLOCATABLE :: NSTA(:,:) INTEGER(KIND=JPIM) ,ALLOCATABLE :: NONL(:,:) INTEGER(KIND=JPIM) :: NGPTOT ! Total number of grid columns on this PE INTEGER(KIND=JPIM) :: NGPTOTG ! Total number of grid columns on the Globe INTEGER(KIND=JPIM) :: NGPTOTMX ! Maximum number of grid columns on any of the PEs INTEGER(KIND=JPIM) ,ALLOCATABLE :: NGPTOTL(:,:) ! Number of grid columns on each PE. REAL(KIND=JPRBT) ,ALLOCATABLE :: RWEIGHT(:) ! Weight per grid-point (if weighted distribution) INTEGER(KIND=JPIM) ,ALLOCATABLE :: NPROCA_GP(:) ! Number of grid-points per a-set INTEGER(KIND=JPIM), ALLOCATABLE :: OFFSETS_GEMM1(:), OFFSETS_GEMM2(:) END TYPE DISTR_TYPE !flat versions of the above INTEGER(KIND=JPIM) :: D_NUMP ! No. of spectral waves handled by this processor INTEGER(KIND=JPIM) ,ALLOCATABLE :: D_MYMS(:) ! Wave numbers handled by this PE INTEGER(KIND=JPIM) ,ALLOCATABLE :: D_NSTAGT0B(:) ! Start adresses for segments within buffer ! (according to processors to whom data ! is going to be sent) INTEGER(KIND=JPIM) ,ALLOCATABLE :: D_NSTAGT1B(:) INTEGER(KIND=JPIM) ,ALLOCATABLE :: D_NPROCL(:) ! Process responsible for each lat. (F.S) INTEGER(KIND=JPIM) ,ALLOCATABLE :: D_NPNTGTB1(:,:) INTEGER(KIND=JPIM) ,ALLOCATABLE :: D_NASM0(:) ! Address in a spectral array of (m, n=m) INTEGER(KIND=JPIM) ,ALLOCATABLE :: D_NSTAGTF(:) ! Offset for specific latitude in INTEGER(KIND=JPIM) :: D_NDGL_FS ! Number of rows of latitudes for which this process is ! performing Fourier Space calculations INTEGER(KIND=JPIM) ,ALLOCATABLE :: D_MSTABF(:) INTEGER(KIND=JPIM) ,ALLOCATABLE :: D_NPNTGTB0(:,:) INTEGER(KIND=JPIM) ,ALLOCATABLE :: D_NPROCM(:) ! Process that does the calc. for certain INTEGER(KIND=JPIM) ,ALLOCATABLE :: D_NPTRLS(:) ! Pointer to first lat. (F.S) ! The offsets in the input and output arrays to the gemms. ! (1) are the offsets in the "inputs" of dirtrans ("outputs" invtrans) ! (2) are the offsets in the "outputs" of invtrans ("inputs" dirtrans) INTEGER(KIND=JPIM), POINTER :: D_OFFSETS_GEMM1(:), D_OFFSETS_GEMM2(:) TYPE(DISTR_TYPE),ALLOCATABLE,TARGET :: DISTR_RESOL(:) TYPE(DISTR_TYPE),POINTER :: D END MODULE TPM_DISTR