source: codes/icosagcm/trunk/src/parallel/transfert_mpi.f90

! Module for MPI communication of field halos
! This module uses Fortran 2003 features : move_alloc intrinsic, pointer bounds remapping, allocatable type fields
module transfert_mpi_mod
  use abort_mod, only : dynamico_abort, abort_acc
  use profiling_mod, only : enter_profile, exit_profile, register_id
  use domain_mod, only : ndomain, ndomain_glo, domain, domain_glo, domloc_glo_ind, domglo_rank, domglo_loc_ind
  use field_mod, only : t_field, field_T, field_U, field_Z
  use transfert_request_mod
  implicit none
  private

  ! Describes how to pack/unpack a message from a local domain to another
  type t_local_submessage
    integer :: src_ind_loc, dest_ind_loc ! index of local and remote domain
    integer :: npoints ! Number of cells to transfer (dim12)
    integer, allocatable :: displ_src(:) ! List of indexes to copy from domain src_ind_loc
    integer, allocatable :: displ_dest(:) ! List of indexes to copy to domain dest_ind_loc
    integer, allocatable :: sign(:) ! Sign change to be applied for vector requests
  end type

  ! Describes how to pack/unpack a message from a domain to another, and contains MPI buffer
  type t_submessage
    integer :: ind_loc, remote_ind_glo, remote_rank ! index of local and remote domain
    integer :: npoints ! Number of cells to transfer (dim12)
    integer, allocatable :: displs(:) ! List of indexes to copy from field to buffer for each level
    integer, allocatable :: sign(:) ! Sign change to be applied for vector requests
    integer :: mpi_buffer_displ = -1
  end type
  
  type t_compact_submessages
    integer :: npoints
    integer, allocatable, dimension(:) :: field_ind, field_displ, sign, remote_rank, mpi_displ, level_offset
  end type
  
  type t_compact_local_submessages
    integer :: npoints
    integer, allocatable, dimension(:) :: field_ind_src, field_displ_src, sign, field_ind_dest, field_displ_dest
  end type
  
  type mpi_buffer_t
    integer :: n
    real, allocatable :: buff(:)
  end type

  ! Describes how to exchange data for a field.
  type t_message
    type (t_field), pointer :: field(:) => null() ! Field to exchange
    type (t_request), pointer :: request(:) => null() ! Type of message to send
    type (t_local_submessage), pointer :: message_local(:) ! Local halo copies
    type (t_submessage), pointer :: message_in(:) ! Messages to recieve from remote ranks and to copy back to the field
    type (t_submessage), pointer :: message_out(:) ! Halos to copy to MPI buffer and to send to remote ranks
    type (mpi_buffer_t), pointer :: mpi_buffer_in(:)
    type (mpi_buffer_t), pointer :: mpi_buffer_out(:)
    type (t_compact_submessages), pointer :: message_in_compact
    type (t_compact_submessages), pointer :: message_out_compact
    type (t_compact_local_submessages), pointer  :: message_local_compact    
    integer, pointer :: mpi_requests_in(:) ! MPI requests used for message_in.
    integer, pointer :: mpi_requests_out(:) ! MPI requests used for message_out.
    ! NOTE : requests are persistant requests initialized in init_message. MPI_Start and MPI_Wait are then used to initiate and complete communications.
    ! ex : Give mpi_requests_in(i) to MPI_Start to send the buffer contained in message_in(i)
    integer :: send_seq ! Sequence number : send_seq is incremented each time send_message is called
    integer :: wait_seq ! Sequence number : wait_seq is incremented each time wait_message is called
    logical :: ondevice ! Ready to transfer ondevice field
  end type t_message

  public :: t_message, t_request, &
    req_i1, req_e1_scal, req_e1_vect, &
    req_i0, req_e0_scal, req_e0_vect, &
    req_z1_scal, &
    init_transfert, &
    init_message, &
    finalize_message, &
    send_message, &
    wait_message, &
    test_message

  ! ---- Private variables ----
  ! Profiling id for mpi
  integer :: profile_mpi, profile_mpi_copies, profile_mpi_waitall, profile_mpi_barrier
contains
  ! Initialize transfert : must be called before any other transfert_mpi routines
  subroutine init_transfert
    use mpi_mod, only : MPI_THREAD_SINGLE, MPI_THREAD_FUNNELED
    use mpipara, only : mpi_threading_mode
    use profiling_mod, only : register_id
    logical, parameter :: profile_mpi_detail = .true.

    !$omp master
    ! Check requested threads support
    if( mpi_threading_mode /= MPI_THREAD_SINGLE .and. mpi_threading_mode /= MPI_THREAD_FUNNELED ) call dynamico_abort("Only single and funneled threading mode are supported.")

    ! Register profiling ids
    call register_id("MPI", profile_mpi)
    if( profile_mpi_detail ) then
      call register_id("MPI_copies", profile_mpi_copies)
      call register_id("MPI_waitall", profile_mpi_waitall)
      call register_id("MPI_omp_barrier", profile_mpi_barrier)
    else
      profile_mpi_copies = profile_mpi
      profile_mpi_waitall = profile_mpi
      profile_mpi_barrier = profile_mpi
    endif

    ! Initialize requests
    call init_all_requests()
    !$omp end master
    !$omp barrier
  end subroutine

  subroutine init_message(field, request, message, name)
    use mpi_mod
    use mpipara
    type(t_field), pointer, intent(in) :: field(:)
    type(t_request),pointer, intent(in) :: request(:)
    type(t_message), target, intent(out) :: message ! Needs intent out for call to finalize_message
    character(len=*), intent(in),optional :: name
    integer, parameter :: INITIAL_ALLOC_SIZE = 10, GROW_FACTOR = 2

    type(t_submessage) :: submessage_in, submessage_out
    type(t_local_submessage) :: submessage_local
    integer :: dim3, dim4, npoints, last_point
    integer :: ind, ind_loc, remote_ind_glo, loc_ind_glo, i, k, remote_rank
    integer :: message_in_size, message_out_size, message_local_size, buffer_in_size, buffer_out_size
    type(t_local_submessage), allocatable :: message_local_tmp(:)
    type(t_submessage), allocatable :: message_in_tmp(:), message_out_tmp(:)
    integer :: field_type

    !$omp barrier
    !$omp master
    !init off-device
    message%ondevice=.false.
    message%send_seq = 0
    message%wait_seq = 0

    if( request(1)%field_type /= field(1)%field_type ) call dynamico_abort( "init_message : field_type/request mismatch" )
    field_type = request(1)%field_type
    
    dim3 = size(field(1)%rval4d,2)
    dim4 = size(field(1)%rval4d,3)
    message%field => field
    message%request => request
    ! Create list of inbound/outbound/local messages
    allocate(message_in_tmp(INITIAL_ALLOC_SIZE))
    message_in_size=0
    allocate(message_out_tmp(INITIAL_ALLOC_SIZE))
    message_out_size=0
    allocate(message_local_tmp(INITIAL_ALLOC_SIZE))
    message_local_size=0
    do loc_ind_glo = 1, ndomain_glo
      do remote_ind_glo = 1, ndomain_glo
        if(domglo_rank(loc_ind_glo) == mpi_rank) then 
          ind_loc = domglo_loc_ind(loc_ind_glo)
          if( domglo_rank(remote_ind_glo) == mpi_rank ) then ! If sending to local domain
            if(request(ind_loc)%points_HtoB(remote_ind_glo)%npoints > 0 ) then ! Add only non-empty messages
              ! Add local message ind_loc -> remote_ind_glo, aggregarting submessage_in and submessage_out into submessage_local
              submessage_out = make_submessage( field_type, request(ind_loc)%points_HtoB(remote_ind_glo), &
                                                ind_loc, remote_ind_glo, dim3, dim4, request(1)%vector )
              submessage_in = make_submessage( field_type, request(domglo_loc_ind(remote_ind_glo))%points_BtoH(domloc_glo_ind(ind_loc)), &
                                              domglo_loc_ind(remote_ind_glo), domloc_glo_ind(ind_loc), dim3, dim4, request(1)%vector)
              submessage_local%src_ind_loc = ind_loc
              submessage_local%dest_ind_loc = domglo_loc_ind(remote_ind_glo)
              submessage_local%npoints = submessage_out%npoints
              submessage_local%displ_src = submessage_out%displs
              submessage_local%displ_dest = submessage_in%displs
              submessage_local%sign = submessage_in%sign
              ! Add to local message list
              call array_append_local_submessage( message_local_tmp, message_local_size, submessage_local)
            endif
          else ! If remote domain
            ! When data to send to remote_domain, add submessage in message%message_out
            if( request(ind_loc)%points_HtoB(remote_ind_glo)%npoints > 0 ) then
              submessage_out = make_submessage( field_type, request(ind_loc)%points_HtoB(remote_ind_glo), &
                                                ind_loc, remote_ind_glo, dim3, dim4, request(1)%vector )
              call array_append_submessage( message_out_tmp, message_out_size, submessage_out )
            end if          
          end if
        end if
      end do
    end do
    ! Recv and Send submessages are transposed to recieve and send in same order
    ! We iterate over global domain index to match sends with recieves (local domains are not ordered like global domains)
    do remote_ind_glo = 1, ndomain_glo
      do loc_ind_glo = 1, ndomain_glo
        if( (domglo_rank(loc_ind_glo) == mpi_rank) .and. (domglo_rank(remote_ind_glo) /= mpi_rank) ) then
          ind_loc = domglo_loc_ind(loc_ind_glo)
          if( request(ind_loc)%points_BtoH(remote_ind_glo)%npoints > 0 ) then
            submessage_in = make_submessage( field_type, request(ind_loc)%points_BtoH(remote_ind_glo), &
                                             ind_loc, remote_ind_glo, dim3, dim4, request(1)%vector )
            call array_append_submessage( message_in_tmp, message_in_size, submessage_in )
          end if
        end if
      end do
    end do
    
    
    ! Trim message_xx_tmp and put it in message%message_xx
    allocate(message%message_in(message_in_size)); message%message_in(:) = message_in_tmp(:message_in_size)
    allocate(message%message_out(message_out_size)); message%message_out(:) = message_out_tmp(:message_out_size)
    allocate(message%message_local(message_local_size)); message%message_local(:) = message_local_tmp(:message_local_size)

    ! Allocate MPI buffers
    allocate( message%mpi_buffer_in(0:mpi_size-1) )
    allocate( message%mpi_buffer_out(0:mpi_size-1) )
    do i = 0, mpi_size-1
      buffer_in_size = dim3*dim4*sum( message%message_in(:)%npoints, message%message_in(:)%remote_rank == i )
      buffer_out_size = dim3*dim4*sum( message%message_out(:)%npoints, message%message_out(:)%remote_rank == i )
      !TODO : what if size == 0 ?
      allocate( message%mpi_buffer_in(i)%buff( buffer_in_size ) )
      allocate( message%mpi_buffer_out(i)%buff( buffer_out_size ) )
      message%mpi_buffer_in(i)%n=0
      message%mpi_buffer_out(i)%n=0
    end do
    ! Set offsets in submessages
    do i=1, size(message%message_out)
      remote_rank = message%message_out(i)%remote_rank
      message%message_out(i)%mpi_buffer_displ = message%mpi_buffer_out(remote_rank)%n
      message%mpi_buffer_out(remote_rank)%n = message%mpi_buffer_out(remote_rank)%n + message%message_out(i)%npoints*dim3*dim4
    end do
    do i=1, size(message%message_in)
      remote_rank = message%message_in(i)%remote_rank
      message%message_in(i)%mpi_buffer_displ = message%mpi_buffer_in(remote_rank)%n
      message%mpi_buffer_in(remote_rank)%n = message%mpi_buffer_in(remote_rank)%n + message%message_in(i)%npoints*dim3*dim4
    end do
    ! Create persistant MPI requests
    allocate( message%mpi_requests_in(0:mpi_size-1) )
    allocate( message%mpi_requests_out(0:mpi_size-1) )
    message%mpi_requests_in(0:mpi_size-1) = MPI_REQUEST_NULL
    message%mpi_requests_out(0:mpi_size-1) = MPI_REQUEST_NULL
    do i = 0, mpi_size-1
      if(  size(message%mpi_buffer_in(i)%buff) /= message%mpi_buffer_in(i)%n &
      .or. size(message%mpi_buffer_out(i)%buff) /= message%mpi_buffer_out(i)%n)&
        call dynamico_abort("Internal error in transfert_mpi : mpi buffer size different than expected")
      if( message%mpi_buffer_out(i)%n > 0) then 
        call MPI_Send_Init( message%mpi_buffer_out(i)%buff, message%mpi_buffer_out(i)%n, MPI_REAL8, i,&
                            100, comm_icosa, message%mpi_requests_out(i), ierr )
      endif
      if( message%mpi_buffer_in(i)%n > 0) then 
        call MPI_Recv_Init( message%mpi_buffer_in(i)%buff, message%mpi_buffer_in(i)%n, MPI_REAL8, i,&
                            100, comm_icosa, message%mpi_requests_in(i), ierr )
      endif
    end do
    
    allocate(message%message_in_compact)    
    message%message_in_compact%npoints = sum(message%message_in(:)%npoints)
    npoints = message%message_in_compact%npoints
    allocate(message%message_in_compact%field_ind(npoints))
    allocate(message%message_in_compact%field_displ(npoints))
    allocate(message%message_in_compact%sign(npoints))
    allocate(message%message_in_compact%remote_rank(npoints))
    allocate(message%message_in_compact%mpi_displ(npoints))
    allocate(message%message_in_compact%level_offset(npoints))

    last_point=0
    do i = 1, size( message%message_in )
      do k = 1, message%message_in(i)%npoints
        last_point = last_point+1
        message%message_in_compact%field_ind(last_point)   = message%message_in(i)%ind_loc
        message%message_in_compact%field_displ(last_point) = message%message_in(i)%displs(k)
        message%message_in_compact%sign(last_point)        = message%message_in(i)%sign(k)
        message%message_in_compact%remote_rank(last_point) = message%message_in(i)%remote_rank
        message%message_in_compact%mpi_displ(last_point)   = message%message_in(i)%mpi_buffer_displ + k
        message%message_in_compact%level_offset(last_point)= message%message_in(i)%npoints
      end do
    end do
    
    allocate(message%message_out_compact) 
    message%message_out_compact%npoints = sum(message%message_out(:)%npoints)
    npoints = message%message_out_compact%npoints
    allocate(message%message_out_compact%field_ind(npoints))
    allocate(message%message_out_compact%field_displ(npoints))
    allocate(message%message_out_compact%sign(npoints))
    allocate(message%message_out_compact%remote_rank(npoints))
    allocate(message%message_out_compact%mpi_displ(npoints))
    allocate(message%message_out_compact%level_offset(npoints))

    last_point=0
    do i = 1, size( message%message_out )
      do k = 1, message%message_out(i)%npoints
        last_point = last_point+1
        message%message_out_compact%field_ind(last_point)   = message%message_out(i)%ind_loc
        message%message_out_compact%field_displ(last_point) = message%message_out(i)%displs(k)
        message%message_out_compact%sign(last_point)        = message%message_out(i)%sign(k)
        message%message_out_compact%remote_rank(last_point) = message%message_out(i)%remote_rank
        message%message_out_compact%mpi_displ(last_point)   = message%message_out(i)%mpi_buffer_displ + k
        message%message_out_compact%level_offset(last_point)= message%message_out(i)%npoints
      end do
    end do 
    
    allocate(message%message_local_compact) 
    message%message_local_compact%npoints = sum(message%message_local(:)%npoints)
    npoints = message%message_local_compact%npoints
    allocate(message%message_local_compact%field_ind_src(npoints))
    allocate(message%message_local_compact%field_displ_src(npoints))
    allocate(message%message_local_compact%sign(npoints))
    allocate(message%message_local_compact%field_ind_dest(npoints))
    allocate(message%message_local_compact%field_displ_dest(npoints))

    last_point=0
    do i = 1, size( message%message_local )
      do k = 1, message%message_local(i)%npoints
        last_point = last_point+1
        message%message_local_compact%field_ind_src(last_point)   = message%message_local(i)%src_ind_loc
        message%message_local_compact%field_displ_src(last_point) = message%message_local(i)%displ_src(k)
        message%message_local_compact%sign(last_point)            = message%message_local(i)%sign(k)
        message%message_local_compact%field_ind_dest(last_point)  = message%message_local(i)%dest_ind_loc
        message%message_local_compact%field_displ_dest(last_point)= message%message_local(i)%displ_dest(k)
      end do
    end do  
    
    !$omp end master
    !$omp barrier
  contains
    ! Generate submessage from points
    function make_submessage(field_type, points, ind_loc, remote_ind_glo, dim3, dim4, vector) result(submessage)
      use dimensions, only : swap_dimensions, iim, u_pos, z_pos
      integer, intent(in) :: field_type
      type(t_points), intent(in) :: points
      integer, intent(in) :: ind_loc, remote_ind_glo, dim3, dim4
      logical, intent(in) :: vector
      integer :: k
      type(t_submessage) :: submessage

      call swap_dimensions(ind_loc)
      submessage%ind_loc = ind_loc
      submessage%remote_ind_glo = remote_ind_glo
      submessage%remote_rank = domglo_rank(remote_ind_glo)
      submessage%npoints = points%npoints
      submessage%mpi_buffer_displ = -1 ! Buffers not allocated yet
      allocate( submessage%displs( points%npoints ) )
      submessage%displs(:) = points%i + (points%j-1)*iim
      if(field_type == field_U) submessage%displs = submessage%displs + u_pos( points%elt )
      if(field_type == field_Z) submessage%displs = submessage%displs + z_pos( points%elt )
      allocate(submessage%sign( points%npoints ))
      if( vector ) then ! For U fields only
        submessage%sign(:) = (/( domain(ind_loc)%edge_assign_sign(points%elt(k)-1, points%i(k), points%j(k)) ,k=1,points%npoints)/)
      else
        submessage%sign(:) = 1
      endif
    end function

    ! Add element to array, and reallocate if necessary
    subroutine array_append_submessage( a, a_size, elt )
      type(t_submessage), allocatable, intent(inout) :: a(:)
      integer, intent(inout) :: a_size
      type(t_submessage), intent(in) :: elt
      type(t_submessage), allocatable :: a_tmp(:)
      integer, parameter :: GROW_FACTOR = 2

      if( size( a ) <= a_size ) then
        allocate( a_tmp ( a_size * GROW_FACTOR ) )
        a_tmp(1:a_size) = a(1:a_size)
        call move_alloc(a_tmp, a)
      end if
      a_size = a_size + 1
      a(a_size) = elt;
    end subroutine
    ! Add element to array, and reallocate if necessary
    subroutine array_append_local_submessage( a, a_size, elt )
      type(t_local_submessage), allocatable, intent(inout) :: a(:)
      integer, intent(inout) :: a_size
      type(t_local_submessage), intent(in) :: elt
      type(t_local_submessage), allocatable :: a_tmp(:)
      integer, parameter :: GROW_FACTOR = 2

      if( size( a ) <= a_size ) then
        allocate( a_tmp ( a_size * GROW_FACTOR ) )
        a_tmp(1:a_size) = a(1:a_size)
        call move_alloc(a_tmp, a)
      end if
      a_size = a_size + 1
      a(a_size) = elt;
    end subroutine
    ! Je demande pardon au dieu du copier-coller car j'ai péché
  end subroutine

  subroutine message_create_ondevice(message)
    use mpi_mod
    use mpipara, only : mpi_size, comm_icosa
    type(t_message), intent(inout) :: message
    integer :: i, ierr

    if( message%ondevice ) call dynamico_abort("Message already on device")

    !$acc enter data copyin(message) async
    !$acc enter data copyin(message%mpi_buffer_in(:)) async
    !$acc enter data copyin(message%mpi_buffer_out(:)) async
    do i = 0, mpi_size-1
      !$acc enter data copyin(message%mpi_buffer_in(i)%buff(:)) async
      !$acc enter data copyin(message%mpi_buffer_out(i)%buff(:)) async
    end do
    !!$acc enter data copyin(message%message_in(:)) async
    !do i = 1, size( message%message_in )
    !  !$acc enter data copyin(message%message_in(i)%displs(:)) async
    !  !$acc enter data copyin(message%message_in(i)%sign(:)) async
    !end do
    !!$acc enter data copyin(message%message_out(:)) async
    !do i = 1, size( message%message_out )
    !  !$acc enter data copyin(message%message_out(i)%displs(:)) async
      !!$acc enter data copyin(message%message_out(i)%sign(:)) async
    !end do
    !!$acc enter data copyin(message%message_local(:)) async
    !do i = 1, size( message%message_local )
    !  !$acc enter data copyin(message%message_local(i)%displ_src(:)) async
    !  !$acc enter data copyin(message%message_local(i)%displ_dest(:)) async
    !  !$acc enter data copyin(message%message_local(i)%sign(:)) async
    !end do
    !$acc enter data copyin(message%field(:)) async
    do i = 1, ndomain
      !$acc enter data copyin(message%field(i)%rval4d(:,:,:)) async
    end do
    
    !$acc enter data copyin(message%message_in_compact) async
    !$acc enter data copyin(message%message_in_compact%field_ind(:)) async
    !$acc enter data copyin(message%message_in_compact%field_displ(:)) async
    !$acc enter data copyin(message%message_in_compact%sign(:)) async
    !$acc enter data copyin(message%message_in_compact%remote_rank(:)) async
    !$acc enter data copyin(message%message_in_compact%mpi_displ(:)) async
    !$acc enter data copyin(message%message_in_compact%level_offset(:)) async
    
    !$acc enter data copyin(message%message_out_compact) async
    !$acc enter data copyin(message%message_out_compact%field_ind(:)) async
    !$acc enter data copyin(message%message_out_compact%field_displ(:)) async
    !$acc enter data copyin(message%message_out_compact%sign(:)) async
    !$acc enter data copyin(message%message_out_compact%remote_rank(:)) async
    !$acc enter data copyin(message%message_out_compact%mpi_displ(:)) async
    !$acc enter data copyin(message%message_out_compact%level_offset(:)) async 
    
    !$acc enter data copyin(message%message_local_compact) async
    !$acc enter data copyin(message%message_local_compact%field_ind_src(:)) async
    !$acc enter data copyin(message%message_local_compact%field_displ_src(:)) async
    !$acc enter data copyin(message%message_local_compact%sign(:)) async
    !$acc enter data copyin(message%message_local_compact%field_ind_dest(:)) async
    !$acc enter data copyin(message%message_local_compact%field_displ_dest(:)) async

    !$acc wait
    do i = 0, mpi_size-1
      if( message%mpi_requests_out(i) /= MPI_REQUEST_NULL ) then
        call MPI_Request_free(message%mpi_requests_out(i), ierr)
        !$acc host_data use_device(message%mpi_buffer_out(i)%buff)
          ! /!\ buff(1) is important for PGI to avoid temporary array copy
          call MPI_Send_Init( message%mpi_buffer_out(i)%buff(1), message%mpi_buffer_out(i)%n, MPI_REAL8, i,&
                              0, comm_icosa, message%mpi_requests_out(i), ierr )
        !$acc end host_data
      end if
      if( message%mpi_requests_in(i) /= MPI_REQUEST_NULL ) then 
        call MPI_Request_free(message%mpi_requests_in(i), ierr)
        !$acc host_data use_device(message%mpi_buffer_in(i)%buff)
          call MPI_Recv_Init( message%mpi_buffer_in(i)%buff(1), message%mpi_buffer_in(i)%n, MPI_REAL8, i,&
                              0, comm_icosa, message%mpi_requests_in(i), ierr )
        !$acc end host_data
      endif
    end do
    message%ondevice=.true.
    !!$acc update device(message%ondevice)
  end subroutine

  subroutine message_delete_ondevice(message)
    use mpipara, only : mpi_size
    type(t_message), intent(inout) :: message
    integer :: i

    if( .not. message%ondevice ) call dynamico_abort("Message not on device")

    !do i = 1, size( message%message_in )
    !  !$acc exit data delete(message%message_in(i)%displs(:)) async
    !  !$acc exit data delete(message%message_in(i)%sign(:)) async
    !end do
    !!$acc exit data delete(message%message_in(:)) async
    !do i = 1, size( message%message_out )
    !  !$acc exit data delete(message%message_out(i)%displs(:)) async
    !  !!$acc exit data delete(message%message_out(i)%sign(:)) async
    !end do
    !!$acc exit data delete(message%message_out(:)) async
    !do i = 1, size( message%message_local )
    !  !$acc exit data delete(message%message_local(i)%displ_src(:)) async
    !  !$acc exit data delete(message%message_local(i)%displ_dest(:)) async
    !  !$acc exit data delete(message%message_local(i)%sign(:)) async
    !end do
    !!$acc exit data delete(message%message_local(:)) async
    do i = 0, mpi_size-1
      !$acc exit data delete(message%mpi_buffer_in(i)%buff(:)) async
      !$acc exit data delete(message%mpi_buffer_out(i)%buff(:)) async
    end do
    !$acc exit data delete(message%mpi_buffer_in(:)) async
    !$acc exit data delete(message%mpi_buffer_out(:)) async
    do i = 1, ndomain
      !$acc exit data delete(message%field(i)%rval4d(:,:,:)) async
    end do
    !$acc exit data delete(message%field(:)) async
    !$acc exit data delete(message) async
    
    !$acc exit data delete(message%message_in_compact%field_ind(:)) async
    !$acc exit data delete(message%message_in_compact%field_displ(:)) async
    !$acc exit data delete(message%message_in_compact%sign(:)) async
    !$acc exit data delete(message%message_in_compact%remote_rank(:)) async
    !$acc exit data delete(message%message_in_compact%mpi_displ(:)) async
    !$acc exit data delete(message%message_in_compact%level_offset(:)) async
    !$acc exit data delete(message%message_in_compact) async
    
    !$acc exit data delete(message%message_out_compact%field_ind(:)) async
    !$acc exit data delete(message%message_out_compact%field_displ(:)) async
    !$acc exit data delete(message%message_out_compact%sign(:)) async
    !$acc exit data delete(message%message_out_compact%remote_rank(:)) async
    !$acc exit data delete(message%message_out_compact%mpi_displ(:)) async
    !$acc exit data delete(message%message_out_compact%level_offset(:)) async
    !$acc exit data delete(message%message_out_compact) async
    
    !$acc exit data delete(message%message_local_compact%field_ind_src(:)) async
    !$acc exit data delete(message%message_local_compact%field_displ_src(:)) async
    !$acc exit data delete(message%message_local_compact%sign(:)) async
    !$acc exit data delete(message%message_local_compact%field_ind_dest(:)) async
    !$acc exit data delete(message%message_local_compact%field_displ_dest(:)) async
    !$acc exit data delete(message%message_local_compact) async
    
    message%ondevice=.false.
  end subroutine

  subroutine finalize_message(message)
    use mpi_mod    
    use mpipara, only : mpi_size
    type(t_message), intent(inout) :: message
    integer :: i, ierr

    !$omp barrier
    !$omp master
    if(message%send_seq /= message%wait_seq) call dynamico_abort("No matching wait_message before finalization")

    if(message%ondevice) call message_delete_ondevice(message)
    deallocate(message%message_in)
    deallocate(message%message_out)
    deallocate(message%message_local)
    do i=0, mpi_size-1
      if(message%mpi_requests_in(i) /= MPI_REQUEST_NULL) call MPI_Request_free(message%mpi_requests_in(i), ierr)
      if(message%mpi_requests_out(i) /= MPI_REQUEST_NULL)call MPI_Request_free(message%mpi_requests_out(i), ierr)
      deallocate(message%mpi_buffer_in(i)%buff)
      deallocate(message%mpi_buffer_out(i)%buff)
    end do
    deallocate(message%mpi_buffer_in)
    deallocate(message%mpi_buffer_out)
    deallocate(message%mpi_requests_in)
    deallocate(message%mpi_requests_out)
    deallocate(message%message_in_compact)
    deallocate(message%message_out_compact)
    deallocate(message%message_local_compact)
    !$omp end master
    !$omp barrier
  end subroutine

  ! Halo to Buffer : copy outbound message to MPI buffers
  subroutine copy_HtoB(message)
    use domain_mod, only : assigned_domain
    use omp_para, only : distrib_level
    type(t_message), intent(inout) :: message
    integer :: dim3, dim4, d3_begin, d3_end
    integer :: k, d3, d4, i
    integer :: local_displ

    dim4 = size(message%field(1)%rval4d, 3)
    dim3 = size(message%field(1)%rval4d, 2)
    CALL distrib_level( 1, dim3, d3_begin, d3_end )
    
    !$acc parallel loop collapse(3) present(message) default(present) async if(message%ondevice)
    do d4 = 1, dim4
      do d3 = d3_begin, d3_end
        do i=1, message%message_out_compact%npoints
          message%mpi_buffer_out( message%message_out_compact%remote_rank(i) )%buff( message%message_out_compact%mpi_displ(i) + message%message_out_compact%level_offset(i)*( (d3-1) + dim3*(d4-1) ) ) &
            = message%field(message%message_out_compact%field_ind(i))%rval4d( message%message_out_compact%field_displ(i), d3, d4 )
        end do
      end do
    end do
    
  end subroutine

  ! Halo to Halo : copy local messages from source field to destination field
  subroutine copy_HtoH(message)
    use domain_mod, only : assigned_domain
    use omp_para, only : distrib_level
    type(t_message), intent(inout) :: message
    integer :: dim3, dim4, d3_begin, d3_end
    integer :: k, d3, d4, i

    dim4 = size(message%field(1)%rval4d, 3)
    dim3 = size(message%field(1)%rval4d, 2)
    CALL distrib_level( 1, dim3, d3_begin, d3_end )
    
    ! TODO : too many copies when tiles are distributed among threads
    !$acc parallel loop collapse(3) present(message) default(present) async if(message%ondevice)
    do d4 = 1, dim4
      do d3 = d3_begin, d3_end
        do i=1, message%message_local_compact%npoints
          message%field(message%message_local_compact%field_ind_dest(i))%rval4d( message%message_local_compact%field_displ_dest(i), d3, d4 ) &
            = message%message_local_compact%sign(i)*message%field(message%message_local_compact%field_ind_src(i))%rval4d( message%message_local_compact%field_displ_src(i), d3, d4 )
        end do
      end do
    end do
  end subroutine

  ! Buffer to Halo : copy inbound message to field
  subroutine copy_BtoH(message)
    use domain_mod, only : assigned_domain
    use omp_para, only : distrib_level
    type(t_message), intent(inout) :: message
    integer :: dim3, dim4, d3_begin, d3_end
    integer :: k, d3, d4, i
    integer :: last_point
    
    dim4 = size(message%field(1)%rval4d, 3)
    dim3 = size(message%field(1)%rval4d, 2)
    CALL distrib_level( 1, dim3, d3_begin, d3_end )
    
    !$acc parallel loop collapse(3) present(message) default(present) async if(message%ondevice)
    do d4 = 1, dim4
      do d3 = d3_begin, d3_end
        do i=1, message%message_in_compact%npoints
          message%field(message%message_in_compact%field_ind(i))%rval4d( message%message_in_compact%field_displ(i), d3, d4 ) &
            = message%message_in_compact%sign(i)*message%mpi_buffer_in( message%message_in_compact%remote_rank(i) )%buff( message%message_in_compact%mpi_displ(i) + message%message_in_compact%level_offset(i)*( (d3-1) + dim3*(d4-1) ) ) 
        end do
      end do
    end do

  end subroutine
    

  subroutine send_message(field, message)
    use mpi_mod
    use mpipara , only : mpi_size
    type(t_field),pointer :: field(:)
    type(t_message), target :: message
    integer :: ierr, i

    call enter_profile(profile_mpi)

    ! Needed because rval4d is set in init_message
    if( .not. associated( message%field, field ) ) &
      call dynamico_abort("send_message must be called with the same field used in init_message")

    !Prepare 'message' for on-device copies if field is on device
    !$omp master
    if( field(1)%ondevice .and. .not. message%ondevice ) call message_create_ondevice(message)
    if( field(1)%ondevice .neqv. message%ondevice ) call dynamico_abort("send_message : internal device/host memory synchronization error")
    ! Check if previous message has been waited
    if(message%send_seq /= message%wait_seq) &
      call dynamico_abort("No matching wait_message before new send_message")
    message%send_seq = message%send_seq + 1
    !$omp end master

    call enter_profile(profile_mpi_barrier)
    !$omp barrier
    call exit_profile(profile_mpi_barrier)

    call enter_profile(profile_mpi_copies)
    call copy_HtoB(message)
    call exit_profile(profile_mpi_copies)

    call enter_profile(profile_mpi_barrier)
    !$acc wait
    call exit_profile(profile_mpi_barrier)
    
    !$omp master
    do i=0, mpi_size-1
      if(message%mpi_requests_in(i) /= MPI_REQUEST_NULL) call MPI_Start( message%mpi_requests_in(i), ierr )
    end do
    !$omp end master

    call enter_profile(profile_mpi_barrier)
    !$omp barrier
    call exit_profile(profile_mpi_barrier)

    !$omp master
    do i=0, mpi_size-1
      if(message%mpi_requests_out(i) /= MPI_REQUEST_NULL) call MPI_Start( message%mpi_requests_out(i), ierr )
    end do
    !$omp end master

    call enter_profile(profile_mpi_copies)
    call copy_HtoH(message)
    call exit_profile(profile_mpi_copies)

    call exit_profile(profile_mpi)
  end subroutine

  subroutine test_message(message)
    use mpi_mod
    type(t_message) :: message
    integer :: ierr
    logical :: completed

    !!$omp master
    !call MPI_Testall( size(message%mpi_requests_out), message%mpi_requests_out, completed, MPI_STATUSES_IGNORE, ierr )
    !call MPI_Testall( size(message%mpi_requests_in), message%mpi_requests_in, completed, MPI_STATUSES_IGNORE, ierr )
    !!$omp end master
  end subroutine

  subroutine wait_message(message)
    use mpi_mod
    type(t_message), target :: message
    integer :: ierr

    ! Check if message has been sent and not recieved yet
    ! note : barrier needed between this and send_seq increment, and this and wait_seq increment
    ! note : watch out for integer overflow a = b+1 doesn't imply a>b
    if(message%send_seq /= message%wait_seq+1) then
      print*, "WARNING : wait_message called multiple times for one send_message, skipping"
      return ! Don't recieve message if already recieved
    end if

    call enter_profile(profile_mpi)

    call enter_profile(profile_mpi_waitall)
    !$omp master
    call MPI_Waitall( size(message%mpi_requests_out), message%mpi_requests_out, MPI_STATUSES_IGNORE, ierr )
    call MPI_Waitall( size(message%mpi_requests_in), message%mpi_requests_in, MPI_STATUSES_IGNORE, ierr )
    !$omp end master
    call exit_profile(profile_mpi_waitall)

    call enter_profile(profile_mpi_barrier)
    !$omp barrier
    call exit_profile(profile_mpi_barrier)

    call enter_profile(profile_mpi_copies) 
    call copy_BtoH(message)
    call exit_profile(profile_mpi_copies)

    !$omp master
    message%wait_seq = message%wait_seq + 1
    !$omp end master

    call enter_profile(profile_mpi_barrier)
    !$omp barrier
    call exit_profile(profile_mpi_barrier)

    call exit_profile(profile_mpi)
  end subroutine
end module