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 :: 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 :: disend_batch => dstates_elec_all_to_all_communications_isend_batch
    procedure :: zisend_batch => zstates_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 :: 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 :: 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
 
    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.)
 
    this%nbatch_to_receive = states_elec_all_to_all_communications_eval_nreceive(st, task_from, this%nblock_to_receive)
    this%nbatch_to_send = states_elec_all_to_all_communications_eval_nsend(st, task_to, this%nblock_to_send)
 
    !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
 
  !------------------------------------------------------------
  integer function states_elec_all_to_all_communications_eval_nreceive(st, task_from, nblock_to_receive) result(nbatch_to_receive)
    type(states_elec_t),   intent(in)  :: st
    integer,               intent(in)  :: task_from
    integer,               intent(out) :: nblock_to_receive
 
    integer :: st_start, st_end, kpt_start, kpt_end, ib
 
    push_sub(states_elec_all_to_all_communications_eval_nreceive)
 
    nbatch_to_receive = 0
    nblock_to_receive = 0
 
    !What we receive for the wfn
    if (task_from > -1) then
      st_start   = st%st_kpt_task(task_from, 1)
      st_end     = st%st_kpt_task(task_from, 2)
      kpt_start  = st%st_kpt_task(task_from, 3)
      kpt_end    = st%st_kpt_task(task_from, 4)
 
      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
      nbatch_to_receive = 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 ', &
      nbatch_to_receive, ' batches from task ', task_from
    call messages_info(1, all_nodes=.true., debug_only=.true.)
 
    pop_sub(states_elec_all_to_all_communications_eval_nreceive)
  end function states_elec_all_to_all_communications_eval_nreceive
 
  !------------------------------------------------------------
  integer function states_elec_all_to_all_communications_eval_nsend(st, task_to, nblock_to_send) result(nbatch_to_send)
    type(states_elec_t),   intent(in)  :: st
    integer,               intent(in)  :: task_to
    integer,               intent(out) :: nblock_to_send
 
    push_sub(states_elec_all_to_all_communications_eval_nsend)
 
    nbatch_to_send = 0
    nblock_to_send = 0
 
    if (task_to > -1) then
      nblock_to_send = (st%group%block_end-st%group%block_start+1)
      nbatch_to_send = nblock_to_send*(st%d%kpt%end-st%d%kpt%start+1)
    end if
 
    write(message(1), '(a,i5,a,i5,a,i5)') 'Debug: Task ', st%st_kpt_mpi_grp%rank, ' will send ', nbatch_to_send, &
      ' batches to task ', task_to
    call messages_info(1, all_nodes=.true., debug_only=.true.)
 
    pop_sub(states_elec_all_to_all_communications_eval_nsend)
  end function states_elec_all_to_all_communications_eval_nsend
 
  !------------------------------------------------------------
  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
 
  !------------------------------------------------------------
  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"
 
end module states_elec_all_to_all_communications_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: