ecearth3/sources/tm5mp/base/arr_decomp.F90

!
#define TRACEBACK write (gol,'("in ",a," (",a,", line",i5,")")') rname, __FILE__, __LINE__; call goErr
#define IF_NOTOK_RETURN(action) if (status/=0) then; TRACEBACK; action; return; end if
#define IF_ERROR_RETURN(action) if (status> 0) then; TRACEBACK; action; return; end if
!
#define IF_NOTOK_MPI(action) if (ierr/=MPI_SUCCESS) then; TRACEBACK; action; return; end if
!
#include "tm5.inc"
!
!----------------------------------------------------------------------------
!                    TM5                                                    !
!----------------------------------------------------------------------------
!BOP
!
! !MODULE:      ARR_DECOMP
!
! !DESCRIPTION: Define a distributed ARRAY object and its methods.
!               This is regardless of the lat/lon grid, i.e. can be applied
!                to any 1- or 2D array.
!                
!               ** ARRAYS ARE EVENLY DISRIBUTED ALONG THE 1ST DIMENSION ONLY **
!               
!               See subroutine TESTDA for basic examples.
!\\
!\\
! !INTERFACE: 
!
MODULE ARR_DECOMP
  !
  ! !USES:
  !
  use Go,    only : goErr, goPr, gol ! go = general objects
  use dims,  only : okdebug
  use partools  ! to include mpif.h, ierr, localComm,...
  
  IMPLICIT NONE
  PRIVATE
  !
  ! !PUBLIC MEMBER FUNCTIONS:
  !
  public :: Set_Darr, Done_Darr  ! life cycle routines
  public :: testda               ! unit test
  public :: Get_Darr
  
  public :: GATHER, SCATTER  ! communication routines  
  !
  ! !PUBLIC TYPES:
  !
  TYPE, PUBLIC :: DIST_ARR
     PRIVATE
     integer :: im     ! global size, 1st dim
     integer :: jm     ! global size, 2nd dim if any
     integer :: i_strt ! begin local arr index
     integer :: i_stop ! end   local arr index
     logical :: lactiv ! may be inactive if there are more processors than data
     
     ! i_start, i_stop of all PEs: (2,npes)
     integer, pointer :: bounds(:,:)

   END TYPE DIST_ARR
  !
  ! !PRIVATE DATA MEMBERS:
  !    
  character(len=*), parameter  ::  mname='Arr_Decomp_MOD_'
  !
  ! !INTERFACE:
  !  
  INTERFACE Gather
    MODULE PROCEDURE gather_1d_i
    MODULE PROCEDURE gather_1d_r
    MODULE PROCEDURE gather_2d_i
    MODULE PROCEDURE gather_2d_r
  END INTERFACE
  
  INTERFACE Scatter
    MODULE PROCEDURE scatter_1d_i
    MODULE PROCEDURE scatter_1d_r
    MODULE PROCEDURE scatter_2d_i
    MODULE PROCEDURE scatter_2d_r
  END INTERFACE
  !
  ! !REVISION HISTORY: 
  !   Monday, May 28, 2018 - P. Le Sager - v0
  !
  ! !REMARKS:
  !     
  ! (1) GATHER & SCATTER :
  !     - global arrays have to really be global on root only, and can be 
  !        (1,1,..) on other processes.
  !     - if not using MPI, gather and scatter just copy arrays, so it may be
  !        better to not call them to save memory in that case.
  !
  ! (2) Be careful when passing a slice (or a pointer to a slice) as argument:
  !     
  !       Passing a subarray will *NOT* work most of the time, unless it is a
  !       slice on the last dimension(s). This will give erroneous results:
  !       
  !           allocate( gbl_arr(-3:imr, 1:jmr ))
  !           call gather( darr, local_arr, gbl_arr(1:imr,:), status)
  !       
  !       but passing full size array will work:
  !       
  !           allocate( gbl_arr(-3:imr, 1:jmr ))
  !           allocate( temp(1:imr,1:jmr) )
  !           call gather( darr, local_arr, temp, status)
  !           gbl_arr(1:imr,:) = temp
  !EOP
  !------------------------------------------------------------------------

CONTAINS
  
  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    DARR_RANGE
  !
  ! !DESCRIPTION: Give range of indices covered by rank when using nprocs.
  !               This is used for one dimension.
  !               Distribution is done evenly. Eg: it will distribute 11 boxes
  !               across 3 processes as 4,4,3 on each pe.
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE DARR_RANGE(ij, rank, nprocs, istart, iend)
    !
    ! !INPUT PARAMETERS:
    !
    integer, intent(in) :: ij      ! range (1,..,ij) to be distributed
    integer, intent(in) :: rank    ! current process, in (0,.., nprocs-1)
    integer, intent(in) :: nprocs  ! total # of processes
    !
    ! !OUTPUT PARAMETERS:
    !
    integer, intent(out):: istart, iend    ! index range covered by rank
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    integer :: iwork1, iwork2

    iwork1 = ij/nprocs
    iwork2 = mod(ij,nprocs)

    istart = rank * iwork1 + 1 + min(rank, iwork2)
    iend = istart + iwork1 - 1
    if (iwork2 > rank) iend = iend + 1

  END SUBROUTINE DARR_RANGE
  !EOC

  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    SET_DARR
  !
  ! !DESCRIPTION: initialize a distributed array object
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE SET_DARR( darr, im, istart, iend, status, jm)
    !
    ! !USES:
    !    
    !
    ! !RETURN VALUE:
    !
    type(dist_arr), intent(inout) :: darr
    !
    ! !INPUT PARAMETERS:
    !
    integer, intent(in)           :: im       ! number of points to distribute - 1st dimension of global array
    integer, optional, intent(in) :: jm       ! 2nd dimension of global array - default is 0
    !
    ! !OUTPUT PARAMETERS:
    !
    integer, intent(out) :: istart, iend, status
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    character(len=*), parameter  ::  rname = mname//'set_darr'

    integer :: lshape(2)
    
    !---------------------------------------------
    ! initialize distributed array info
    !---------------------------------------------
    call done_darr(darr, status)
    call darr_range(im, myid, npes, istart, iend) ! index range covered by current process

    ! Need to acount for array size smaller than the number of available processor. Just inactivate
    !  the processors.
    darr%lactiv = .true.
    if ( (iend-istart+1) < 1 ) then
      write(gol,*)" Inactivate ", myid
      darr%lactiv = .false.
    end if
    
    darr%i_strt = istart
    darr%i_stop = iend
    darr%im = im
    darr%jm = 0
    if (present(jm)) darr%jm = jm
    
    !---------------------------------------------
    ! store local shapes info of all PE on all PE
    !---------------------------------------------
#ifdef MPI    
    allocate(darr%bounds(2,0:npes-1))
    lshape = (/ darr%i_strt, darr%i_stop /)
    call MPI_ALLGATHER(lshape, 2, MPI_INTEGER, darr%bounds, 2, MPI_INTEGER, localComm, ierr)
#endif

    status = 0
        
  END SUBROUTINE SET_DARR
  !EOC
  
  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    DONE_DARR
  !
  ! !DESCRIPTION: deallocate distributed object elements
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE DONE_DARR( darr, status )
    !
    ! !INPUT PARAMETERS:
    !
    type(dist_arr), intent(inout) :: darr
    !
    ! !OUTPUT PARAMETERS:
    !
    integer, intent(out) :: status
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    character(len=*), parameter  ::  rname = mname//'Done_Darr'
    
    if (associated(darr%bounds)) deallocate(darr%bounds)

    status=0 

  END SUBROUTINE DONE_DARR
  !EOC
  
  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    GET_DARR
  !
  ! !DESCRIPTION: provide quick access to object elements (bounds and grids),
  !               while keeping them private.
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE GET_DARR(Darr, istart, istop, im)
    !
    ! !INPUT PARAMETERS:
    !
    type(dist_arr),  intent(in) :: Darr
    integer, optional :: istart, istop, im
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    if (present(istart)) istart = darr%i_strt
    if (present(istop))  istop = darr%i_stop
    if (present(im))     im = darr%im
    
  END SUBROUTINE GET_DARR
  !EOC
  
#ifdef MPI  /* MPI TYPES */

  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    GET_INTERIOR_TYPE
  !
  ! !DESCRIPTION: Returns derived MPI types that describe the interior arrs
  !               needed for collective communications.
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE GET_INTERIOR_TYPE( Darr, shp, datatype, linterior, ginterior, status )
    !
    ! !INPUT PARAMETERS:
    !
    type(dist_arr), intent(in) :: Darr
    integer,        intent(in) :: shp(:)     ! shape of local array
    integer,        intent(in) :: datatype   ! basic MPI datatype
    !
    ! !OUTPUT PARAMETERS:
    !
    ! derived MPI datatypes describing distributed interior arrs:
    integer, intent(out) :: ginterior(npes-1) ! within global array (used by root, as many as NPES-1)
    integer, intent(out) :: linterior         ! within local array (used by non-root)
    integer, intent(out) :: status
    !
    ! !REVISION HISTORY: 
    !   Monday, May 28, 2018 - P. Le Sager - v0
    !
    ! !REMARKS:
    !  (1) input must be checked beforehand by calling CHECK_DIST_ARR first
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    character(len=*), parameter  ::  rname = mname//'get_interior_type'

    integer :: gslice, lslice     ! intermediate datatypes
    integer :: n, m               ! sizes to build datatypes
    integer :: hlstride, hgstride ! strides to build datatypes
    integer :: stack, sz, klm
    integer(kind=MPI_ADDRESS_KIND) :: sizeoftype, lb, vlstride, vgstride

    ! init : number of dimensions, default value
    sz = size(shp)
    ginterior = MPI_DATATYPE_NULL
    linterior = MPI_DATATYPE_NULL

    ! collapse third and above dimensions
    stack = 1
    if (sz > 2) stack = product(shp(3:))
    
    ! size of data slice to carry 
    n  = darr%I_STOP - darr%I_STRT + 1
    m  = darr%jm

    CALL MPI_TYPE_GET_EXTENT( datatype, lb, sizeoftype, ierr) 
    IF_NOTOK_MPI(status=1)
    
    ! horizontal global stride (in data)
    hgstride = darr%im

    ! vertical global stride (in byte)
    vgstride = darr%im * darr%jm * sizeoftype

    ! local strides (may be different than n and n*m because of halo)
    hlstride = shp(1)                    ! in data
    vlstride = shp(1)*shp(2)*sizeoftype  ! in byte

    if ( isRoot ) then

       do klm=1,npes-1

          ! horizontal chunk is M lines of lenght N, separated by a global
          ! stride 
          n = darr%bounds(2,klm) - darr%bounds(1,klm) + 1
          m = darr%jm

          call MPI_TYPE_VECTOR (m, n, hgstride, datatype, gslice, ierr)
          IF_NOTOK_MPI(status=1)
          
          ! stack 3rd and above dimensions
          if (stack == 1) then
             ginterior(klm) = gslice
          else             
             ! note : also works with stack=1
             call MPI_TYPE_CREATE_HVECTOR(stack, 1, vgstride, gslice, ginterior(klm), ierr)
             IF_NOTOK_MPI(status=1)

             call MPI_TYPE_FREE(gslice, ierr)
             IF_NOTOK_MPI(status=1)
          end if
          
          call MPI_TYPE_COMMIT (ginterior(klm), ierr)
          IF_NOTOK_MPI(status=1)
          
       end do

    else
    
       ! local interior is basically M lines of lenght N, separated by a local
       ! stride 
       call MPI_TYPE_VECTOR (m, n, hlstride, datatype, lslice, ierr)
       IF_NOTOK_MPI(status=1)
       
       ! stack 3rd and above dimensions
       if (stack == 1) then
          linterior = lslice
       else
          ! note : also works with stack=1
          call MPI_TYPE_CREATE_HVECTOR (stack, 1, vlstride, lslice, linterior, ierr)
          IF_NOTOK_MPI(status=1)

          call MPI_TYPE_FREE(lslice, ierr)
          IF_NOTOK_MPI(status=1)
       end if
       
       call MPI_TYPE_COMMIT (linterior, ierr)
       IF_NOTOK_MPI(status=1)
       
    end if

    status=0
    
  END SUBROUTINE GET_INTERIOR_TYPE
  !EOC

  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    FREE_DERIVED_TYPE
  !
  ! !DESCRIPTION: free all MPI derived datatypes in a vector
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE FREE_DERIVED_TYPE( datatype )
    !
    ! !INPUT/OUTPUT PARAMETERS:
    !
    integer, intent(inout) :: datatype(:)      ! set of derived MPI datatypes
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC

    integer :: i, j
    integer :: res(size(datatype))  ! hold unique elements 
    integer :: k                    ! number of unique elements
    
    ! Get list of unique handle(s)
    ! ----------------------------
    k = 1
    res(1) = 1
    
    outer: do i=2,size(datatype)

       ! look for a match
       do j=1,k
          if (datatype(res(j)) == datatype(i)) cycle outer  ! match
       end do
       
       ! no match : add it to the list
       k = k + 1
       res(k) = i
       
    end do outer

    ! Free handle(s)
    ! ---------------------------
    do i=1,k
       if (datatype(res(i)) /= MPI_DATATYPE_NULL) &
            call MPI_TYPE_FREE(datatype(res(i)), ierr)
    end do

  END SUBROUTINE FREE_DERIVED_TYPE
  !EOC

#endif /* MPI TYPES */
  
  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    CHECK_DIST_ARR
  !
  ! !DESCRIPTION: Check that the shape of a local array correspond to an array
  !               distributed on current process. This check is done on the
  !               first 2 dimensions only, along which the arr
  !               decomposition is done. 
  !               
  !               Optionally: check shape of a global array. If arrays are 3D
  !               or more, the 3rd and above dimensions of local and global
  !               arrays are also compared.
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE CHECK_DIST_ARR( darr, shp, status, glbl_shp, has_global )
    !
    ! !INPUT PARAMETERS:
    !
    type(dist_arr),    intent(in) :: darr
    integer,           intent(in) :: shp(:)       ! shape of local array
    !
    integer, intent(in), optional :: glbl_shp(:)  ! shape of global array
    logical, intent(in), optional :: has_global   ! current proc hold global array (default is root)
    !
    ! !OUTPUT PARAMETERS:
    !
    integer, intent(out) :: status
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    character(len=*), parameter  ::  rname = mname//'check_dist_arr '
    integer :: n, m, sz, sg
    integer, allocatable :: glbsz(:)
    logical :: hold_global

    status = 0

    ! by default global array is expected on root
    hold_global = isRoot
    if (present(has_global)) hold_global=has_global
        
    ! required size
    n = darr%i_stop - darr%i_strt + 1

    ! check the 1st dimension, which is distributed
    if ((shp(1) /= n) ) then
      write (gol,*) "CHECK_DIST_ARR: local array shape does not conform" ; call goErr
      write (gol,'("  local array : ",2i4)') shp(1)                      ; call goErr    
      write (gol,'("  should be   : ",2i4)') n                           ; call goErr
      write (gol,'("  w/ start & stop : ", i4)') darr%i_strt, darr%i_stop ; call goErr
      TRACEBACK; status=1; return
    end if

    ! check shape of global array on root
    sz = size(shp)
    if ( present(glbl_shp) ) then
      sg = size(glbl_shp)

      if (sz /= sg ) then
        write (gol,'("CHECK_DIST_ARR : global and local arrays have different rank:")') ; call goErr
        write (gol,'("  local rank  : ", i4)') sz                                       ; call goErr
        write (gol,'("  global rank : ", i4)') sg                                       ; call goErr
        TRACEBACK; status=1; return          
      end if
    
      if ((sz == 2) .and. hold_global) then
        if (shp(2) /= glbl_shp(2)) then
          write (gol,'("CHECK_DIST_ARR : global and local arrays have different J-size:")') ; call goErr
          write (gol,'("  local size(2)  : ", i4)') shp(2)                                  ; call goErr
          write (gol,'("  global size(2) : ", i4)') glbl_shp(2)                             ; call goErr
          TRACEBACK; status=1; return          
        end if
      end if
      
    end if

  END SUBROUTINE CHECK_DIST_ARR
  !EOC
  
  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    GATHER_1D_R
  !
  ! !DESCRIPTION: gather local 1D arrays into a global 1D array  
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE GATHER_1D_R( darr, lcl_array, glbl_array, status )
    !
    ! !INPUT PARAMETERS:
    !
    type(dist_arr), intent(in) :: darr
    real,           intent(in) :: lcl_array(darr%i_strt:)
    !
    ! !OUTPUT PARAMETERS:
    !
    real,    intent(out) :: glbl_array(:)
    integer, intent(out) :: status
    !
    ! !REMARKS:
    ! (1) I have not use mpi_gatherv because of varying *receiving* size
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    character(len=*), parameter  ::  rname = mname//'gather_1d_r'

#ifndef MPI
    glbl_array = lcl_array( darr%i_strt:darr%i_stop)
    status = 0
#else    

    integer :: stat(MPI_STATUS_SIZE)
    integer :: i, j, klm, i0, j0, i1, j1, sz(1)

    status=0

    ! check input, get derived types
    !-------------------------------
    sz = shape(lcl_array)

!    if(okdebug)then
      CALL CHECK_DIST_ARR( Darr, sz, status, shape(glbl_array))
      IF_NOTOK_RETURN(status=1)
!    end if

    i0 = darr%i_strt
    i1 = darr%i_stop

    ! ------- GATHER -------------
    if ( isRoot ) then

      ! get first chunk locally
      glbl_array(i0:i1) = lcl_array(i0:i1)

      ! get remaining chunks from other pes
      do klm=1,npes-1
        
        i = darr%bounds(1,klm)
        j = darr%bounds(2,klm)
        
        call MPI_RECV( glbl_array(i), j-i+1, my_real, klm, 1, localComm, stat, ierr)
      end do
    else
      call MPI_SEND( lcl_array(i0), i1-i0+1, my_real, root, 1, localComm, ierr) 
    end if
#endif

  END SUBROUTINE GATHER_1D_R

  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    GATHER_2D_R
  !
  ! !DESCRIPTION: gather local 2D arrays into a global 2D array  
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE GATHER_2D_R( darr, lcl_array, glbl_array, status )
    !
    ! !INPUT PARAMETERS:
    !
    type(dist_arr), intent(in) :: darr
    real,           intent(in) :: lcl_array(darr%i_strt:,:)
    !
    ! !OUTPUT PARAMETERS:
    !
    real,    intent(out) :: glbl_array(:,:)
    integer, intent(out) :: status
    !
    ! !REMARKS:
    ! (1) I have not use mpi_gatherv because of varying *receiving* size
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    character(len=*), parameter  ::  rname = mname//'gather_2d_r'

#ifndef MPI
    glbl_array = lcl_array( darr%i_strt:darr%i_stop,:)
    status = 0
#else    
    integer :: stat(MPI_STATUS_SIZE), linterior, ginterior(npes-1)
    integer :: i, j, klm, i0, j0, i1, j1, sz(2)

    status=0
    
    ! check input, get derived types
    !-------------------------------
    sz = shape(lcl_array)
    
!    if(okdebug)then
       call check_dist_arr( darr, sz, status, shape(glbl_array))
       IF_NOTOK_RETURN(status=1)
!    end if

    call get_interior_type( Darr, sz, my_real, linterior, ginterior, status )
    IF_NOTOK_RETURN(status=1)
    
    i0  = darr%i_strt

    ! ------- GATHER -------------    
    if ( isRoot ) then

       ! get first chunk locally
       i1 = darr%i_stop

       glbl_array(i0:i1,:) = lcl_array(i0:i1,:)

       ! get remaining chunks from other pes
       do klm=1,npes-1

          i = darr%bounds(1,klm)

          call MPI_RECV( glbl_array(i,1), 1, ginterior(klm), klm, 1, &
                         localComm, stat, ierr)
       end do

       call FREE_DERIVED_TYPE( ginterior )
    else
       call MPI_SEND( lcl_array(i0,1), 1, linterior, root, 1, localComm, ierr) 
       call MPI_TYPE_FREE(linterior, ierr)
    end if
#endif
    
  END SUBROUTINE GATHER_2D_R
  !EOC
      
  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    GATHER_1D_I
  !
  ! !DESCRIPTION: gather local 1D arrays into a global 1D array  
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE GATHER_1D_I( darr, lcl_array, glbl_array, status )
    !
    ! !INPUT PARAMETERS:
    !
    type(dist_arr), intent(in) :: darr
    integer,        intent(in) :: lcl_array(darr%i_strt:)
    !
    ! !OUTPUT PARAMETERS:
    !
    integer, intent(out) :: glbl_array(:)
    integer, intent(out) :: status
    !
    ! !REMARKS:
    ! (1) I have not use mpi_gatherv because of varying *receiving* size
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    character(len=*), parameter  ::  rname = mname//'gather_1d_r'

#ifndef MPI
    glbl_array = lcl_array( darr%i_strt:darr%i_stop)
    status = 0
#else    

    integer :: stat(MPI_STATUS_SIZE)
    integer :: i, j, klm, i0, j0, i1, j1, sz(1)

    status=0

    ! check input, get derived types
    !-------------------------------
    sz = shape(lcl_array)

    !    if(okdebug)then
    CALL CHECK_DIST_ARR( Darr, sz, status, shape(glbl_array))
    IF_NOTOK_RETURN(status=1)
    !    end if

    i0 = darr%i_strt
    i1 = darr%i_stop

    ! ------- GATHER -------------
    if ( isRoot ) then

      ! get first chunk locally
      glbl_array(i0:i1) = lcl_array(i0:i1)

      ! get remaining chunks from other pes
      do klm=1,npes-1

        i = darr%bounds(1,klm)
        j = darr%bounds(2,klm)

        call MPI_RECV( glbl_array(i), j-i+1, MPI_INTEGER, klm, 1, localComm, stat, ierr)

      end do
    else
      call MPI_SEND( lcl_array(i0), i1-i0+1, MPI_INTEGER, root, 1, localComm, ierr) 
    end if
#endif

  END SUBROUTINE GATHER_1D_I
  
  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    GATHER_2D_I
  !
  ! !DESCRIPTION: gather local 2D arrays into a global 2D array (integer version)  
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE GATHER_2D_I( Darr, lcl_array, glbl_array, status )
    !
    ! !INPUT PARAMETERS:
    !
    type(dist_arr), intent(in) :: Darr
    integer,         intent(in) :: lcl_array(darr%i_strt:,:)
    !
    ! !OUTPUT PARAMETERS:
    !
    integer, intent(out) :: glbl_array(:,:)
    integer, intent(out) :: status
    !
    ! !REVISION HISTORY: 
    !   Monday, May 28, 2018 - P. Le Sager - v0
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    character(len=*), parameter  ::  rname = mname//'gather_2d_i'

#ifndef MPI
    glbl_array = lcl_array( darr%i_strt:darr%i_stop, : )
    status = 0
#else    

    integer :: stat(MPI_STATUS_SIZE), linterior, ginterior(npes-1)
    integer :: i, j, klm, i0, j0, i1, j1, sz(2)

    status=0

    ! check input, get derived types
    !-------------------------------
    sz = shape(lcl_array)
    
!    if(okdebug)then
       call check_dist_arr( darr, sz, status, shape(glbl_array))
       IF_NOTOK_RETURN(status=1)
!    end if

    call get_interior_type( darr, sz, MPI_INTEGER, linterior, ginterior, status )
    IF_NOTOK_RETURN(status=1)
    
    i0  = darr%i_strt

    ! ------- GATHER -------------
    if ( isRoot ) then

       ! get first chunk locally
       i1 = darr%i_stop

       glbl_array(i0:i1,:) = lcl_array(i0:i1,:)

       ! get remaining chunks from other pes
       do klm=1,npes-1

          i = darr%bounds(1,klm)

          call MPI_RECV( glbl_array(i,1), 1, ginterior(klm), klm, 1, &
                         localComm, stat, ierr)
       end do
        
       call FREE_DERIVED_TYPE( ginterior )

    else
       call MPI_SEND( lcl_array(i0,1), 1, linterior, root, 1, localComm, ierr) 
       call MPI_TYPE_FREE(linterior, ierr)
    end if
#endif

  END SUBROUTINE GATHER_2D_I
  !EOC
  
  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    SCATTER_1D_R
  !
  ! !DESCRIPTION: scatter a 2D global real array  
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE SCATTER_1D_R( darr, lcl_array, glbl_array, status )
    !
    ! !INPUT PARAMETERS:
    !
    type(dist_arr), intent(in) :: darr
    real,           intent(in) :: glbl_array(:)
    !
    ! !OUTPUT PARAMETERS:
    !
    real,    intent(out) :: lcl_array(darr%i_strt:)
    integer, intent(out) :: status
    !
    ! !REVISION HISTORY: 
    !   Monday, May 28, 2018 - P. Le Sager - v0
    !   21 Jun 2013 - P. Le Sager - MPI_SEND -> MPI_SSEND to avoid buffering
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    character(len=*), parameter  ::  rname = mname//'scatter_1d_r'

#ifndef MPI
    lcl_array( darr%i_strt:darr%i_stop) = glbl_array
    status = 0
#else

    integer :: stat(MPI_STATUS_SIZE)
    integer :: i, j, klm, i0, j0, i1, j1, sz(1)

    status=0

    ! ------- Check input & get derived types 
    sz = shape(lcl_array)

!    if(okdebug)then
      CALL CHECK_DIST_ARR( darr, sz, status, shape(glbl_array))
      IF_NOTOK_RETURN(status=1)
!    end if

    i0 = darr%i_strt
    i1 = darr%i_stop

    ! ------- SCATTER -------------
    if ( isRoot ) then  ! send

      ! get first chunk locally
      lcl_array(i0:i1) = glbl_array(i0:i1)

      ! send remaining chunks to other pes
      do klm=1,npes-1

        i = darr%bounds(1,klm)
        j = darr%bounds(2,klm)

        call MPI_SSEND( glbl_array(i), j-i+1, my_real, klm, klm, localComm, ierr)
        IF_NOTOK_MPI(status=1)

      end do

    else

      call MPI_COMM_RANK(localComm, klm, ierr)
      IF_NOTOK_MPI(status=1)

      call MPI_RECV( lcl_array(i0), i1-i0+1, my_real, root, klm, localComm, stat, ierr)
      IF_NOTOK_MPI(status=1)

    end if
#endif
  END SUBROUTINE SCATTER_1D_R
  !EOC
  
  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    SCATTER_2D_R
  !
  ! !DESCRIPTION: scatter a 2D global real array  
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE SCATTER_2D_R( darr, lcl_array, glbl_array, status )
    !
    ! !INPUT PARAMETERS:
    !
    type(dist_arr), intent(in) :: darr
    real,           intent(in) :: glbl_array(:,:)
    !
    ! !OUTPUT PARAMETERS:
    !
    real,    intent(out) :: lcl_array(darr%i_strt:,:)
    integer, intent(out) :: status
    !
    ! !REVISION HISTORY: 
    !   Monday, May 28, 2018 - P. Le Sager - v0
    !   21 Jun 2013 - P. Le Sager - MPI_SEND -> MPI_SSEND to avoid buffering
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    character(len=*), parameter  ::  rname = mname//'scatter_2d_r'

#ifndef MPI
    lcl_array( darr%i_strt:darr%i_stop, : ) = glbl_array
    status = 0
#else
    
    integer :: stat(MPI_STATUS_SIZE), linterior, ginterior(npes-1)
    integer :: i, j, klm, i0, j0, i1, j1, sz(2)

    status=0
 
    ! ------- Check input & get derived types 
    sz = shape(lcl_array)

!    if(okdebug)then
       call check_dist_arr( darr, sz, status, shape(glbl_array))
       IF_NOTOK_RETURN(status=1)
!    end if

    call get_interior_type( darr, sz, my_real, linterior, ginterior, status )
    IF_NOTOK_RETURN(status=1)

    i0  = darr%i_strt

    ! ------- SCATTER -------------
    if ( isRoot ) then

       ! get first chunk locally
       i1 = darr%i_stop

       lcl_array(i0:i1,:) = glbl_array(i0:i1,:)

       ! send remaining chunks to other pes
       do klm=1,npes-1

          i = darr%bounds(1,klm)

          call MPI_SSEND( glbl_array(i,1), 1, ginterior(klm), klm, klm, localComm, ierr)
          IF_NOTOK_MPI(status=1)

       end do

       call FREE_DERIVED_TYPE( ginterior )

    else

       call MPI_COMM_RANK(localComm, klm, ierr)
       IF_NOTOK_MPI(status=1)

       call MPI_RECV( lcl_array(i0,1), 1, linterior, root, klm, localComm, stat, ierr)
       IF_NOTOK_MPI(status=1)
       
       call MPI_TYPE_FREE(linterior, ierr)
       IF_NOTOK_MPI(status=1)

    end if
#endif
  END SUBROUTINE SCATTER_2D_R
  !EOC

  
  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    SCATTER_1D_I
  !
  ! !DESCRIPTION: scatter a 2D global integer array  
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE SCATTER_1D_I( darr, lcl_array, glbl_array, status )
    !
    ! !INPUT PARAMETERS:
    !
    type(dist_arr), intent(in) :: darr
    integer,        intent(in) :: glbl_array(:)
    !
    ! !OUTPUT PARAMETERS:
    !
    integer, intent(out) :: lcl_array(darr%i_strt:)
    integer, intent(out) :: status
    !
    ! !REVISION HISTORY: 
    !   Monday, May 28, 2018 - P. Le Sager - v0
    !   21 Jun 2013 - P. Le Sager - MPI_SEND -> MPI_SSEND to avoid buffering
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    character(len=*), parameter  ::  rname = mname//'scatter_1d_i'

#ifndef MPI
    lcl_array( darr%i_strt:darr%i_stop) = glbl_array
    status = 0
#else

    integer :: stat(MPI_STATUS_SIZE)
    integer :: i, j, klm, i0, j0, i1, j1, sz(1)

    status=0

    ! ------- Check input & get derived types 
    sz = shape(lcl_array)

    !    if(okdebug)then
    CALL CHECK_DIST_ARR( darr, sz, status, shape(glbl_array))
    IF_NOTOK_RETURN(status=1)
    !    end if

    i0 = darr%i_strt
    i1 = darr%i_stop

    ! ------- SCATTER -------------
    if ( isRoot ) then  ! send

      ! get first chunk locally
      lcl_array(i0:i1) = glbl_array(i0:i1)

      ! send remaining chunks to other pes
      do klm=1,npes-1

        i = darr%bounds(1,klm)
        j = darr%bounds(2,klm)

        call MPI_SSEND( glbl_array(i), j-i+1, MPI_INTEGER, klm, klm, localComm, ierr)
        IF_NOTOK_MPI(status=1)

      end do

    else

      call MPI_COMM_RANK(localComm, klm, ierr)
      IF_NOTOK_MPI(status=1)

      call MPI_RECV( lcl_array(i0), i1-i0+1, MPI_INTEGER, root, klm, localComm, stat, ierr)
      IF_NOTOK_MPI(status=1)

    end if
#endif
  END SUBROUTINE SCATTER_1D_I
  !EOC

  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    SCATTER_2D_I
  !
  ! !DESCRIPTION: scatter a 2D global integer array  
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE SCATTER_2D_I( darr, lcl_array, glbl_array, status )
    !
    ! !INPUT PARAMETERS:
    !
    type(dist_arr), intent(in) :: darr
    integer,           intent(in) :: glbl_array(:,:)
    !
    ! !OUTPUT PARAMETERS:
    !
    integer,    intent(out) :: lcl_array(darr%i_strt:,:)
    integer, intent(out) :: status
    !
    ! !REVISION HISTORY: 
    !   Monday, May 28, 2018 - P. Le Sager - v0
    !   21 Jun 2013 - P. Le Sager - MPI_SEND -> MPI_SSEND to avoid buffering
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    character(len=*), parameter  ::  rname = mname//'scatter_2d_i'

#ifndef MPI
    lcl_array( darr%i_strt:darr%i_stop, : ) = glbl_array
    status = 0
#else

    integer :: stat(MPI_STATUS_SIZE), linterior, ginterior(npes-1)
    integer :: i, j, klm, i0, j0, i1, j1, sz(2)

    status=0

    ! ------- Check input & get derived types 
    sz = shape(lcl_array)

    !    if(okdebug)then
    call check_dist_arr( darr, sz, status, shape(glbl_array))
    IF_NOTOK_RETURN(status=1)
    !    end if

    call get_interior_type( darr, sz, MPI_INTEGER, linterior, ginterior, status )
    IF_NOTOK_RETURN(status=1)

    i0  = darr%i_strt

    ! ------- SCATTER -------------
    if ( isRoot ) then

      ! get first chunk locally
      i1 = darr%i_stop

      lcl_array(i0:i1,:) = glbl_array(i0:i1,:)

      ! send remaining chunks to other pes
      do klm=1,npes-1

        i = darr%bounds(1,klm)

        call MPI_SSEND( glbl_array(i,1), 1, ginterior(klm), klm, klm, localComm, ierr)
        IF_NOTOK_MPI(status=1)

      end do

      call FREE_DERIVED_TYPE( ginterior )

    else

      call MPI_COMM_RANK(localComm, klm, ierr)
      IF_NOTOK_MPI(status=1)

      call MPI_RECV( lcl_array(i0,1), 1, linterior, root, klm, localComm, stat, ierr)
      IF_NOTOK_MPI(status=1)

      call MPI_TYPE_FREE(linterior, ierr)
      IF_NOTOK_MPI(status=1)

    end if
#endif
  END SUBROUTINE SCATTER_2D_I
  !EOC
  
  !--------------------------------------------------------------------------
  !                    TM5                                                  !
  !--------------------------------------------------------------------------
  !BOP
  !
  ! !IROUTINE:    TESTDA
  !
  ! !DESCRIPTION: check if the communications are working as expected.
  !               Currently checked:
  !               - GATHER, SCATTER
  !\\
  !\\
  ! !INTERFACE:
  !
  SUBROUTINE TESTDA( status )
    !
    ! !OUTPUT PARAMETERS:
    !
    integer,        intent(out) :: status
    !
    ! !REMARKS:
    !
    !EOP
    !------------------------------------------------------------------------
    !BOC
    character(len=*), parameter  ::  rname = mname//'testcomm'

    real, allocatable :: glb(:), lcl(:), ttglb(:)
    real, allocatable :: glb2(:,:), lcl2(:,:), ttglb2(:,:)

    integer :: i, j, istrt, istop, gsize, jm
    type(dist_arr) :: darr
    
    ! General
    gsize = 20
    jm = 359
    
    ! ----- 1D -------------------------------------------------------------
    ! test array
    if (isRoot) then
      allocate(TTglb(gsize))
      TTglb = (/(real(i)+0.123, i=1,gsize)/)
    endif

    call Set_Darr( darr, gsize, istrt, istop, status)
    
    ! Global array (only on root it needs to be filled with data and be the full size)
    if (isRoot) then
      allocate(glb(gsize))
      glb = (/(i, i=1,gsize)/)
    else
      allocate(glb(1))
    endif

    ! fill local array by scattering
    allocate(lcl(istrt:istop))
    call scatter(darr, lcl, glb, status)

    ! do something on the local array
    lcl = lcl + 0.123
    
    call gather(darr, lcl, glb, status)
    
    call done_darr(darr, status)
    
    if (isRoot) then
      if ( any(TTglb /= glb) ) then
        write(gol,*) "1D: We have a problem!!!" ; call goPr
        status=1
      else
        write(gol,*) "1D: SUCCESS !!!" ; call goPr
        status=0
      endif
    endif
    
    call Par_Broadcast_Status( status, root )
    IF_NOTOK_RETURN(status=1)

    deallocate(glb,lcl)
    if (isRoot) deallocate(TTglb)
    
    ! ----- 2D -------------------------------------------------------------
    ! test array
    if (isRoot) then
      allocate(TTglb2(gsize,jm))
      TTglb2 = reshape((/(real(i)+0.123, i=1,gsize*jm)/), shape(TTglb2))
    endif

    call Set_Darr( darr, gsize, istrt, istop, status, jm)

    ! Global array (only on root it needs to be filled with data and be the full size)
    if (isRoot) then
      allocate(glb2(gsize,jm))
      glb2 = reshape((/(i, i=1,gsize*jm)/), shape(glb2))
    else
      allocate(glb2(1,1))
    endif

    ! fill local array by scattering
    allocate(lcl2(istrt:istop,jm))
    call scatter(darr, lcl2, glb2, status)

    ! do something on the local array
    lcl2 = lcl2 + 0.123

    call gather(darr, lcl2, glb2, status)
    
    ! - test
    if (isRoot) then
      if ( any(TTglb2 /= glb2) ) then
        write(gol,*) "2D: We have a problem!!!" ; call goPr
        status=1
        do i=1,gsize
          do j=1,jm
            write(gol,*) i,j, glb2(i,j), TTglb2(i,j); call goPr
          enddo
        enddo
        write(gol,*)'--- FAIL TEST ----'; call goPr
      else
        write(gol,*) "2D: SUCCESS !!!" ; call goPr
        status=0
      endif
    endif

    call Par_Broadcast_Status( status, root )
    IF_NOTOK_RETURN(status=1)

    deallocate(glb2,lcl2)
    if (isRoot) deallocate(TTglb2)
    
    ! clean DA
    call done_darr(darr, status)

    
  END SUBROUTINE TESTDA
  !EOC

END MODULE ARR_DECOMP