states_elec_all_to_all_communications.F90

Go to the documentation of this file.

!! Copyright (C) 2023 N. Tancogne-Dejean
!!
!! This program is free software; you can redistribute it and/or modify
!! it under the terms of the GNU General Public License as published by
!! the Free Software Foundation; either version 2, or (at your option)
!! any later version.
!!
!! This program is distributed in the hope that it will be useful,
!! but WITHOUT ANY WARRANTY; without even the implied warranty of
!! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!! GNU General Public License for more details.
!!
!! You should have received a copy of the GNU General Public License
!! along with this program; if not, write to the Free Software
!! Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
!! 02110-1301, USA.
!!
 
#include "global.h"
 
 
module states_elec_all_to_all_communications_oct_m
  use accel_oct_m
  use batch_oct_m
  use debug_oct_m
  use global_oct_m
  use math_oct_m
  use mesh_oct_m
  use messages_oct_m
  use mpi_oct_m
  use namespace_oct_m
  use profiling_oct_m
  use states_elec_oct_m
  use states_elec_parallel_oct_m
  use wfs_elec_oct_m
 
  implicit none
 
  private
 
  public ::                                            &
    states_elec_all_to_all_communications_t
 
  type states_elec_all_to_all_communications_t
    private
 
    integer :: task_from
    integer :: task_to
 
    integer :: nbatch_to_receive
    integer :: nbatch_to_send
    integer :: n_comms
 
    integer :: nblock_to_receive
    integer :: nblock_to_send
 
    logical :: gpu_aware = .false. 
 
    type(MPI_Request), allocatable :: send_req(:)
  contains
    procedure :: start => states_elec_all_to_all_communications_start
    procedure :: get_nreceiv => states_elec_all_to_all_communications_get_nreceive
    procedure :: get_nsend => states_elec_all_to_all_communications_get_nsend
    procedure :: get_ncom => states_elec_all_to_all_communications_get_ncom
    procedure :: process_receives => states_elec_all_to_all_communications_process_receives
    procedure :: process_sends => states_elec_all_to_all_communications_process_sends
    procedure :: alloc_receive_batch => states_elec_all_to_all_communications_alloc_receive_batch
    procedure :: get_send_indices => states_elec_all_to_all_communications_get_send_indices
    procedure :: get_receive_indices => states_elec_all_to_all_communications_get_receive_indices
    procedure :: dpost_all_mpi_isend => dstates_elec_all_to_all_communications_post_all_mpi_isend
    procedure :: zpost_all_mpi_isend => zstates_elec_all_to_all_communications_post_all_mpi_isend
    procedure :: spost_all_mpi_isend => sstates_elec_all_to_all_communications_post_all_mpi_isend
    procedure :: cpost_all_mpi_isend => cstates_elec_all_to_all_communications_post_all_mpi_isend
    procedure :: disend_batch => dstates_elec_all_to_all_communications_isend_batch
    procedure :: zisend_batch => zstates_elec_all_to_all_communications_isend_batch
    procedure :: sisend_batch => sstates_elec_all_to_all_communications_isend_batch
    procedure :: cisend_batch => cstates_elec_all_to_all_communications_isend_batch
    procedure :: disend_batch_to => dstates_elec_all_to_all_communications_isend_batch_to
    procedure :: zisend_batch_to => zstates_elec_all_to_all_communications_isend_batch_to
    procedure :: dmpi_irecv_batch => dstates_elec_all_to_all_communications_mpi_irecv_batch
    procedure :: zmpi_irecv_batch => zstates_elec_all_to_all_communications_mpi_irecv_batch
    procedure :: sisend_batch_to => sstates_elec_all_to_all_communications_isend_batch_to
    procedure :: cisend_batch_to => cstates_elec_all_to_all_communications_isend_batch_to
    procedure :: smpi_irecv_batch => sstates_elec_all_to_all_communications_mpi_irecv_batch
    procedure :: cmpi_irecv_batch => cstates_elec_all_to_all_communications_mpi_irecv_batch
    procedure :: drecv_batch_from => dstates_elec_all_to_all_communications_recv_batch_from
    procedure :: zrecv_batch_from => zstates_elec_all_to_all_communications_recv_batch_from
    procedure :: srecv_batch_from => sstates_elec_all_to_all_communications_recv_batch_from
    procedure :: crecv_batch_from => cstates_elec_all_to_all_communications_recv_batch_from
    procedure :: wait_all_isend => states_elec_all_to_all_communications_wait_all_isend
  end type states_elec_all_to_all_communications_t
 
contains
 
  subroutine states_elec_all_to_all_communications_start(this, st, task_from, task_to, gpu_aware)
    class(states_elec_all_to_all_communications_t),  intent(inout) :: this
    type(states_elec_t),   intent(in) :: st
    integer,               intent(in) :: task_from, task_to
    logical, optional,     intent(in) :: gpu_aware
 
    integer :: kpt_start, kpt_end
 
    push_sub(states_elec_all_to_all_communications_start)
 
    this%task_from = task_from
    this%task_to = task_to
 
    this%gpu_aware = optional_default(gpu_aware, .false.)
 
    !What we receive for the wfn
    this%nblock_to_receive = set_nblock_to_receive(st, task_from)
    this%nbatch_to_receive = 0
    if (task_from > -1) then
      kpt_start  = st%st_kpt_task(task_from, 3)
      kpt_end    = st%st_kpt_task(task_from, 4)
      this%nbatch_to_receive = this%nblock_to_receive * (kpt_end - kpt_start + 1)
    end if
 
    write(message(1), '(a,i5,a,i5,a,i5)') 'Debug: Task ', st%st_kpt_mpi_grp%rank, ' will receive ', &
      this%nbatch_to_receive, ' batches from task ', task_from
    call messages_info(1, all_nodes=.true., debug_only=.true.)
 
    !What we send for the wfn
    this%nblock_to_send = merge(st%group%block_end - st%group%block_start + 1, 0, task_to > -1)
    this%nbatch_to_send = this%nblock_to_send * (st%d%kpt%end - st%d%kpt%start + 1)
 
    write(message(1), '(a,i5,a,i5,a,i5)') 'Debug: Task ', st%st_kpt_mpi_grp%rank, ' will send ', &
      this%nbatch_to_send, ' batches to task ', task_to
    call messages_info(1, all_nodes=.true., debug_only=.true.)
 
    !Number of communications
    this%n_comms = max(this%nbatch_to_send, this%nbatch_to_receive)
 
    pop_sub(states_elec_all_to_all_communications_start)
  end subroutine states_elec_all_to_all_communications_start
 
  ! Private free function
 
  function set_nblock_to_receive(st, task_from) result(nblock_to_receive)
    type(states_elec_t),   intent(in)  :: st
    integer,               intent(in)  :: task_from
    integer                            :: nblock_to_receive
 
    integer :: st_start, st_end, ib
 
    push_sub(set_nblock_to_receive)
 
    nblock_to_receive = 0
 
    if (task_from > -1) then
      st_start   = st%st_kpt_task(task_from, 1)
      st_end     = st%st_kpt_task(task_from, 2)
      nblock_to_receive = 0
      do ib = 1, st%group%nblocks
        if (st%group%block_range(ib, 1) >= st_start .and. st%group%block_range(ib, 2) <= st_end) then
          nblock_to_receive = nblock_to_receive + 1
        end if
      end do
    end if
 
    pop_sub(set_nblock_to_receive)
 
  end function set_nblock_to_receive
 
  ! Type-bound member functions
 
  !------------------------------------------------------------
  integer pure function states_elec_all_to_all_communications_get_ncom(this) result(n_comms)
    class(states_elec_all_to_all_communications_t),  intent(in) :: this
 
    n_comms = this%n_comms
  end function states_elec_all_to_all_communications_get_ncom
 
  !------------------------------------------------------------
  integer pure function states_elec_all_to_all_communications_get_nreceive(this) result(nbatch_to_receive)
    class(states_elec_all_to_all_communications_t),  intent(in) :: this
 
    nbatch_to_receive = this%nbatch_to_receive
  end function states_elec_all_to_all_communications_get_nreceive
 
  !------------------------------------------------------------
  integer pure function states_elec_all_to_all_communications_get_nsend(this) result(nbatch_to_send)
    class(states_elec_all_to_all_communications_t),  intent(in) :: this
 
    nbatch_to_send = this%nbatch_to_send
  end function states_elec_all_to_all_communications_get_nsend
 
  !------------------------------------------------------------
  logical pure function states_elec_all_to_all_communications_process_receives(this, icom) result(receives)
    class(states_elec_all_to_all_communications_t),  intent(in) :: this
    integer,                                         intent(in) :: icom
 
    receives = this%task_from > -1 .and. icom <= this%nbatch_to_receive
  end function states_elec_all_to_all_communications_process_receives
 
  !------------------------------------------------------------
  logical pure function states_elec_all_to_all_communications_process_sends(this, icom) result(sends)
    class(states_elec_all_to_all_communications_t),  intent(in) :: this
    integer,                                         intent(in) :: icom
 
    sends = this%task_to > -1 .and. icom <= this%nbatch_to_send
  end function states_elec_all_to_all_communications_process_sends
 
  !------------------------------------------------------------
  subroutine states_elec_all_to_all_communications_alloc_receive_batch(this, st, icom, np, psib)
    class(states_elec_all_to_all_communications_t),  intent(in)  :: this
    type(states_elec_t),                             intent(in)  :: st
    integer,                                         intent(in)  :: icom
    integer,                                         intent(in)  :: np
    type(wfs_elec_t),                                intent(out) :: psib
 
    integer :: block_id, ib, ik
 
    push_sub(states_elec_all_to_all_communications_alloc_receive_batch)
 
    ! Given the icom, returns the id of the block of state communicated
    block_id = mod(icom-1, this%nblock_to_receive)+1
    ik = int((icom-block_id)/this%nblock_to_receive) + st%st_kpt_task(this%task_from, 3)
    ib = block_id - 1 + st%group%iblock(st%st_kpt_task(this%task_from, 1))
 
    write(message(1), '(a,i5,a,i5,a,i5)') 'Debug: Task ', st%st_kpt_mpi_grp%rank, ' allocates memory for block ', &
      ib, ' and k-point ', ik
    call messages_info(1, all_nodes=.true., debug_only=.true.)
 
    call states_elec_parallel_allocate_batch(st, psib, np, ib, ik, packed=.true.)
 
    ! For CUDA-aware MPI, the batch must be on GPU to receive directly in the buffer.
    if (this%gpu_aware) then
      call psib%do_pack(batch_device_packed, copy = .false.)
    end if
 
    pop_sub(states_elec_all_to_all_communications_alloc_receive_batch)
  end subroutine states_elec_all_to_all_communications_alloc_receive_batch
 
  !------------------------------------------------------------
  subroutine states_elec_all_to_all_communications_get_send_indices(this, st, icom, ib, ik)
    class(states_elec_all_to_all_communications_t),  intent(in)  :: this
    type(states_elec_t),                             intent(in)  :: st
    integer,                                         intent(in)  :: icom
    integer,                                         intent(out) :: ib
    integer,                                         intent(out) :: ik
 
    push_sub(states_elec_all_to_all_communications_get_send_indices)
 
    ! Given the icom, returns the id of the block of state communicated
    ib = mod(icom-1, this%nblock_to_send) + 1
    ik = int((icom-ib)/this%nblock_to_send) + st%d%kpt%start
    ib = ib - 1 + st%group%iblock(st%st_start)
 
    write(message(1), '(a,i5,a,i5,a,i5)') 'Debug: Task ', st%st_kpt_mpi_grp%rank, ' will send the block ', &
      ib, ' with k-point ', ik
    call messages_info(1, all_nodes=.true., debug_only=.true.)
 
    pop_sub(states_elec_all_to_all_communications_get_send_indices)
  end subroutine states_elec_all_to_all_communications_get_send_indices
 
  !------------------------------------------------------------
  subroutine states_elec_all_to_all_communications_get_receive_indices(this, st, icom, ib, ik)
    class(states_elec_all_to_all_communications_t),  intent(in)  :: this
    type(states_elec_t),                             intent(in)  :: st
    integer,                                         intent(in)  :: icom
    integer,                                         intent(out) :: ib
    integer,                                         intent(out) :: ik
 
    push_sub(states_elec_all_to_all_communications_get_receive_indices)
 
    ! Given the icom, returns the id of the block of state communicated
    ib = mod(icom-1, this%nblock_to_receive)+1
    ik = int((icom-ib)/this%nblock_to_receive) +  st%st_kpt_task(this%task_from, 3)
    ib = ib - 1 + st%group%iblock(st%st_kpt_task(this%task_from, 1))
 
    if (debug%info) then
      write(message(1), '(a,i5,a,i5,a,i5)') 'Task ', st%st_kpt_mpi_grp%rank, ' will receive the block ', &
        ib, ' with k-point ', ik
      call messages_info(1, all_nodes=.true.)
    end if
 
    pop_sub(states_elec_all_to_all_communications_get_receive_indices)
  end subroutine states_elec_all_to_all_communications_get_receive_indices
 
  !------------------------------------------------------------
  subroutine states_elec_all_to_all_communications_wait_all_isend(this, st)
    class(states_elec_all_to_all_communications_t),  intent(inout) :: this
    type(states_elec_t),                             intent(in)    :: st
 
    push_sub(states_elec_all_to_all_communications_wait_all_isend)
    call profiling_in("ALL_TO_ALL_COMM")
 
    if (allocated(this%send_req)) then
      call st%st_kpt_mpi_grp%wait(this%nbatch_to_send, this%send_req)
      safe_deallocate_a(this%send_req)
    end if
 
    call profiling_out("ALL_TO_ALL_COMM")
    pop_sub(states_elec_all_to_all_communications_wait_all_isend)
  end subroutine states_elec_all_to_all_communications_wait_all_isend
 
 
#include "undef.F90"
#include "real.F90"
#include "states_elec_all_to_all_communications_inc.F90"
 
#include "undef.F90"
#include "complex.F90"
#include "states_elec_all_to_all_communications_inc.F90"
 
#include "undef.F90"
#include "real32.F90"
#include "states_elec_all_to_all_communications_inc.F90"
 
#include "undef.F90"
#include "complex32.F90"
#include "states_elec_all_to_all_communications_inc.F90"
#include "undef.F90"
 
end module states_elec_all_to_all_communications_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: