partition.F90

Go to the documentation of this file.

!! Copyright (C) 2013 M. Oliveira
!! Copyright (C) 2021 S. Ohlmann
!!
!! 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 partition_oct_m
  use debug_oct_m
  use global_oct_m
  use io_oct_m
  use io_binary_oct_m
  use messages_oct_m
  use mpi_oct_m
  use mpi_debug_oct_m
  use profiling_oct_m
 
  implicit none
 
  private
 
  public ::                         &
    partition_t,                    &
    partition_init,                 &
    partition_end,                  &
    partition_set,                  &
    partition_dump,                 &
    partition_load,                 &
    partition_get_local_size,       &
    partition_get_global,           &
    partition_get_partition_number, &
    partition_get_np_local,         &
    partition_get_npart,            &
    partition_get_part,             &
    partition_get_local
 
 
  !
  type partition_t
    private
 
    ! The following components are the same for all processes:
    type(mpi_grp_t) :: mpi_grp
    integer(int64) ::     np_global
    integer        ::     npart = 1
    integer, allocatable        :: np_local_vec(:)
    integer(int64), allocatable :: istart_vec(:)
 
    ! The following components are process-dependent:
    integer :: partno
    integer :: np_local
    integer(int64) :: istart
    integer, allocatable :: part(:)
  end type partition_t
 
  interface partition_get_partition_number
    module procedure partition_get_partition_number_4, partition_get_partition_number_8
  end interface partition_get_partition_number
 
 
contains
 
  !---------------------------------------------------------
  !
  subroutine partition_init(partition, np_global, mpi_grp)
    type(partition_t), intent(out) :: partition
    integer(int64),       intent(in)  :: np_global
    type(mpi_grp_t),   intent(in)  :: mpi_grp
 
    integer :: ipart
    integer(int64) :: iend
 
    push_sub(partition_init)
 
    !Global variables
    partition%mpi_grp = mpi_grp
    partition%np_global = np_global
    partition%npart = mpi_grp%size
    partition%partno = mpi_grp%rank + 1
 
    safe_allocate(partition%np_local_vec(1:partition%npart))
    safe_allocate(partition%istart_vec(1:partition%npart))
 
    ! use block data decomposition
    do ipart = 1, partition%npart
      ! according to Fortran 2008 standard, Sect. 7.1.9.3, point 4, the kind parameter of an
      ! operation is the one with a larger decimal range, i.e., everything is promoted to i8
      partition%istart_vec(ipart) = (ipart-1) * np_global/partition%npart + 1
      iend  = ipart * np_global/partition%npart
      partition%np_local_vec(ipart) = i8_to_i4(iend - partition%istart_vec(ipart) + 1)
    end do
    partition%istart = partition%istart_vec(partition%partno)
    partition%np_local = partition%np_local_vec(partition%partno)
 
    !Allocate memory for the partition
    safe_allocate(partition%part(1:partition%np_local))
 
    pop_sub(partition_init)
  end subroutine partition_init
 
  ! ---------------------------------------------------------
  subroutine partition_end(partition)
    type(partition_t), intent(inout) :: partition
 
    push_sub(partition_end)
 
    safe_deallocate_a(partition%part)
    safe_deallocate_a(partition%np_local_vec)
    safe_deallocate_a(partition%istart_vec)
 
    pop_sub(partition_end)
  end subroutine partition_end
 
  ! ---------------------------------------------------------
  subroutine partition_set(partition, part)
    type(partition_t), intent(inout) :: partition
    integer,           intent(in)    :: part(:)
 
    push_sub(partition_set)
 
    partition%part(1:partition%np_local) = part(1:partition%np_local)
 
    pop_sub(partition_set)
  end subroutine partition_set
 
  ! ---------------------------------------------------------
  !
  subroutine partition_dump(partition, dir, filename, ierr)
    type(partition_t), intent(in)  :: partition
    character(len=*),  intent(in)  :: dir
    character(len=*),  intent(in)  :: filename
    integer,           intent(out) :: ierr
 
    integer :: err
    character(len=MAX_PATH_LEN) :: full_filename
 
    push_sub(partition_dump)
 
    ierr = 0
    full_filename = trim(dir)//'/'//trim(filename)
 
    ! Write the header (root only) and wait
    if (partition%mpi_grp%is_root()) then
      call iwrite_header(full_filename, partition%np_global, err)
      if (err /= 0) ierr = ierr + 1
    end if
    call partition%mpi_grp%bcast(ierr, 1, mpi_integer, 0)
 
    assert(all(partition%part(:) > 0))
 
    ! Each process writes a portion of the partition
    if (ierr == 0) then
      call mpi_debug_in(partition%mpi_grp%comm, c_mpi_file_write)
      ! Only one rank per partition group should write the partition restart information
      ! Otherwise, more than once is trying to be written data
      ! TODO(Martin): this "one writer per replicated partition copy" selection uses mpi_world
      ! coordinates; under per-system sub-communicators it must use the grp the partition was carved
      ! from (needs a parent grp stored on partition_t).
      ! This need to be done when parallising over ensembles (see #1494).
      if (mod(mpi_world%rank, mpi_world%size/partition%mpi_grp%size) == 0) then
        call io_binary_write_parallel(full_filename, partition%mpi_grp%comm, partition%istart, &
          partition%np_local, partition%part, err)
        call mpi_debug_out(partition%mpi_grp%comm, c_mpi_file_write)
        if (err /= 0) ierr = ierr + 2
      end if
    end if
 
    pop_sub(partition_dump)
  end subroutine partition_dump
 
  ! ---------------------------------------------------------
  !
  subroutine partition_load(partition, dir, filename, ierr)
    use iso_c_binding, only: c_sizeof
    type(partition_t), intent(inout) :: partition
    character(len=*),  intent(in)    :: dir
    character(len=*),  intent(in)    :: filename
    integer,           intent(out)   :: ierr
 
    integer :: err
    integer(int64) :: np, file_size
    integer, allocatable :: scounts(:)
    integer(int64), allocatable :: sdispls(:)
    character(len=MAX_PATH_LEN) :: full_filename
 
    push_sub(partition_load)
 
    ierr = 0
    full_filename = trim(dir)//'/'//trim(filename)
 
    ! This is a writing to avoid an optimization of gfortran with -O3
    write(message(1),'(a,i8)') "Info: number of points in the partition (in root process) =", size(partition%part)
    call messages_info(1)
 
    ! Check if the file exists and has the proper size (only world root)
    ! TODO(Martin): root-read + bcast is over mpi_world; under per-system sub-communicators this
    ! should use partition%mpi_grp so each system reads its own restart. This should be changed
    ! as part of issue #1494 (parallelisation over ensembles)
    if (mpi_world%is_root()) then
      call io_binary_get_info(full_filename, np, file_size, err)
      assert(np == partition%np_global)
    end if
 
    ! All nodes need to know the result
    call mpi_world%bcast(err, 1, mpi_integer, 0)
    call mpi_world%bcast(file_size, 1, mpi_integer8, 0)
 
    if (err /= 0) then
      ierr = ierr + 1
      pop_sub(partition_load)
      return
    end if
    ! The size of the file is not as big as np_global
    if (file_size - 64 /= partition%np_global * c_sizeof(int(0))) then
      ierr = ierr + 2
      pop_sub(partition_load)
      return
    end if
 
    ! Calculate displacements for reading
    safe_allocate(scounts(1:partition%npart))
    safe_allocate(sdispls(1:partition%npart))
 
    scounts(:) = partition%np_local_vec(:)
    sdispls(:) = partition%istart_vec(:) - 1
 
    assert(sum(scounts(:)) == partition%np_global)
 
#ifdef HAVE_MPI
    call mpi_debug_in(partition%mpi_grp%comm, c_mpi_file_read)
    call io_binary_read_parallel(full_filename, partition%mpi_grp%comm, partition%istart, &
      partition%np_local, partition%part, err)
    call mpi_debug_out(partition%mpi_grp%comm, c_mpi_file_read)
    if (err /= 0) ierr = ierr + 4
#endif
 
    if (any(partition%part(:) <= 0)) then
      write(message(1),'(a)') 'Internal error: some elements of partition are <= 0.'
      write(message(2),*) 'filename = ', full_filename
      write(message(3),*) 'scounts = ', scounts(:)
      write(message(4),*) 'sdispls = ', sdispls(:)
      call messages_fatal(6)
    end if
 
    safe_deallocate_a(scounts)
    safe_deallocate_a(sdispls)
 
    pop_sub(partition_load)
  end subroutine partition_load
 
  ! ---------------------------------------------------------
  subroutine partition_get_local_size(partition, istart, np_local)
    type(partition_t), intent(in)  :: partition
    integer(int64),       intent(out) :: istart
    integer,           intent(out) :: np_local
 
    push_sub(partition_get_local_size)
 
    istart = partition%istart
    np_local = partition%np_local
 
    pop_sub(partition_get_local_size)
  end subroutine partition_get_local_size
 
  ! ---------------------------------------------------------
  subroutine partition_get_global(partition, part_global, root)
    type(partition_t), intent(in)  :: partition
    integer,           intent(out) :: part_global(:)
    integer, optional, intent(in)  :: root
 
    integer(int64), allocatable :: rdispls(:)
    integer, allocatable :: rcounts(:)
 
    push_sub(partition_get_global)
 
    safe_allocate(rdispls(1:partition%npart))
    safe_allocate(rcounts(1:partition%npart))
 
    rcounts(:) = partition%np_local_vec(:)
    rdispls(:) = partition%istart_vec(:) - 1
 
    assert(all(partition%part(1:partition%np_local) > 0))
 
    if (present(root)) then
      call partition%mpi_grp%gatherv(partition%part, partition%np_local, mpi_integer, &
        part_global, rcounts, rdispls, mpi_integer, root)
    else
      call partition%mpi_grp%allgatherv(partition%part, partition%np_local, mpi_integer, &
        part_global, rcounts, rdispls, mpi_integer)
    end if
 
    if (.not. present(root) .or. partition%mpi_grp%is_root()) then
      assert(all(part_global(:) > 0))
    end if
 
#ifndef HAVE_MPI
    part_global = partition%part
#endif
 
    safe_deallocate_a(rdispls)
    safe_deallocate_a(rcounts)
 
    pop_sub(partition_get_global)
  end subroutine partition_get_global
 
  ! ---------------------------------------------------------
  subroutine partition_get_partition_number_4(partition, np, points, partno)
    type(partition_t), intent(in)  :: partition
    integer,           intent(in)  :: np
    integer(int64),       intent(in)  :: points(:)
    integer,           intent(out) :: partno(:)
 
    integer :: ip, nproc, rnp
    integer(int64), allocatable :: sbuffer(:), rbuffer(:)
    integer, allocatable :: scounts(:), rcounts(:)
    integer, allocatable :: sdispls(:), rdispls(:)
    integer, allocatable :: ipos(:), order(:)
 
    push_sub(partition_get_partition_number_4)
 
    safe_allocate(scounts(1:partition%npart))
    safe_allocate(rcounts(1:partition%npart))
    safe_allocate(sdispls(1:partition%npart))
    safe_allocate(rdispls(1:partition%npart))
 
    ! How many points will we have to send/receive from each process?
    scounts = 0
    do ip = 1, np
      ! Who knows where points(ip) is stored?
      nproc = partition_get_number(partition, points(ip))
      ! We increase the respective counter
      scounts(nproc) = scounts(nproc) + 1
    end do
 
 
    !Tell each process how many points we will need from it
    call partition%mpi_grp%alltoall(scounts, 1, mpi_integer, &
      rcounts, 1, mpi_integer)
 
    !Build displacement arrays
    sdispls(1) = 0
    rdispls(1) = 0
    do ip = 2, partition%npart
      sdispls(ip) = sdispls(ip-1) + scounts(ip-1)
      rdispls(ip) = rdispls(ip-1) + rcounts(ip-1)
    end do
 
 
    rnp = sum(rcounts)
    safe_allocate(sbuffer(1:np))
    safe_allocate(rbuffer(1:rnp))
 
    !Put points in correct order for sending
    safe_allocate(ipos(1:partition%npart))
    safe_allocate(order(1:np))
    ipos = 0
    do ip = 1, np
      ! Who knows where points(ip) is stored?
      nproc = partition_get_number(partition, points(ip))
 
      !We increase the respective counter
      ipos(nproc) = ipos(nproc) + 1
 
      !Put the point in the correct place
      order(ip) = sdispls(nproc) + ipos(nproc)
      sbuffer(order(ip)) = points(ip) ! global index of the point
    end do
    safe_deallocate_a(ipos)
 
    !Send the global indices of the points to the process that knows what is the corresponding partition
    call partition%mpi_grp%alltoallv(sbuffer, scounts, sdispls, mpi_integer8, &
      rbuffer, rcounts, rdispls, mpi_integer8)
 
    !We get the partition number from the global index. This will be send back.
    do ip = 1, rnp
      if (rbuffer(ip) == 0) cycle
      rbuffer(ip) = partition%part(rbuffer(ip) - partition%istart + 1)
    end do
 
    !Now we send the information backwards
    call partition%mpi_grp%alltoallv(rbuffer, rcounts, rdispls, mpi_integer8, &
      sbuffer, scounts, sdispls, mpi_integer8)
 
    !Reorder the points
    do ip = 1, np
      partno(ip) = i8_to_i4(sbuffer(order(ip)))
    end do
 
    !Deallocate memory
    safe_deallocate_a(order)
    safe_deallocate_a(sbuffer)
    safe_deallocate_a(scounts)
    safe_deallocate_a(sdispls)
    safe_deallocate_a(rbuffer)
    safe_deallocate_a(rcounts)
    safe_deallocate_a(rdispls)
 
    pop_sub(partition_get_partition_number_4)
  end subroutine partition_get_partition_number_4
 
  !---------------------------------------------------------
  subroutine partition_get_partition_number_8(partition, np, points, partno)
    type(partition_t), intent(in)  :: partition
    integer(int64),       intent(in)  :: np
    integer(int64),       intent(in)  :: points(:)
    integer,           intent(out) :: partno(:)
 
    integer(int64) :: rounds, offset, iround
 
    push_sub(partition_get_partition_number_8)
 
    rounds = np/huge(0_int32)
    do iround = 1, rounds
      offset = (iround - 1)*huge(0_int32) + 1
      call partition_get_partition_number(partition, huge(0_int32), points(offset:offset+huge(0_int32)), &
        partno(offset:offset+huge(0_int32)))
    end do
    offset = rounds*huge(0_int32) + 1
    call partition_get_partition_number(partition, i8_to_i4(np-offset+1), points(offset:np), &
      partno(offset:np))
 
    pop_sub(partition_get_partition_number_8)
  end subroutine partition_get_partition_number_8
 
  !---------------------------------------------------------
  subroutine partition_get_np_local(partition, np_local_vec)
    type(partition_t),    intent(in)  :: partition
    integer, contiguous,  intent(inout) :: np_local_vec(:)
 
    integer, allocatable :: np_local_vec_tmp(:)
    integer :: ip
 
    push_sub(partition_get_np_local)
 
    assert(ubound(np_local_vec, 1) >= partition%npart)
    assert(partition%npart > 0)
    assert(all(partition%part(:) > 0))
    safe_allocate(np_local_vec_tmp(1:partition%npart))
    np_local_vec_tmp = 0
 
    ! Calculate locally the local points of each partition
    do ip = 1, partition%np_local
      np_local_vec_tmp(partition%part(ip)) = np_local_vec_tmp(partition%part(ip)) + 1
    end do
 
    ! Collect all the local points
    call partition%mpi_grp%allreduce(np_local_vec_tmp, np_local_vec, partition%npart, mpi_integer, mpi_sum)
    safe_deallocate_a(np_local_vec_tmp)
 
    pop_sub(partition_get_np_local)
  end subroutine partition_get_np_local
 
  !---------------------------------------------------------
  pure integer function partition_get_npart(partition) result(npart)
    type(partition_t), intent(in) :: partition
    npart = partition%npart
  end function partition_get_npart
 
  !---------------------------------------------------------
  pure integer function partition_get_part(partition, local_point) result(part)
    type(partition_t), intent(in) :: partition
    integer,           intent(in) :: local_point
    part = partition%part(local_point)
  end function partition_get_part
 
  !---------------------------------------------------------
  pure integer function partition_get_number(partition, global_point) result(part)
    type(partition_t), intent(in) :: partition
    integer(int64),       intent(in) :: global_point
 
    if (global_point == 0) then
      part = partition%partno
    else
      part = i8_to_i4((partition%npart*global_point - 1)/partition%np_global + 1)
    end if
  end function partition_get_number
 
  !---------------------------------------------------------
  subroutine partition_get_local(partition, rbuffer, np_local)
    type(partition_t),        intent(in)    :: partition
    integer(int64), contiguous,  intent(inout) :: rbuffer(:)
    integer,                  intent(out)   :: np_local
 
    integer :: ip, ipart
    integer(int64) :: istart
    integer(int64), allocatable :: sdispls(:), rdispls(:), sbuffer(:)
    integer, allocatable :: scounts(:), rcounts(:)
 
    push_sub(partition_get_local)
 
    safe_allocate(sdispls(1:partition%npart))
    safe_allocate(scounts(1:partition%npart))
    safe_allocate(rcounts(1:partition%npart))
 
    ! Calculate the starting point of the running process
    istart = partition%istart - 1
 
    scounts = 0
    ! Count and store the local points for each partition
    do ip = 1, partition%np_local
      ipart = partition%part(ip)
      scounts(ipart) = scounts(ipart) + 1
    end do
 
    ! Create displacements
    sdispls(1) = 0
    do ipart = 2, partition%npart
      sdispls(ipart) = sdispls(ipart-1) + scounts(ipart-1)
    end do
 
    ! Allocate and fill the send buffer
    np_local = sum(scounts)
    scounts = 0
    safe_allocate(sbuffer(1:np_local))
    do ip = 1, np_local
      ipart = partition%part(ip)
      scounts(ipart) = scounts(ipart) + 1
      sbuffer(sdispls(ipart) + scounts(ipart)) = ip + istart
    end do
 
    ! Tell each process how many points we will need from it
    call partition%mpi_grp%alltoall(scounts, 1, mpi_integer, &
      rcounts, 1, mpi_integer)
 
    ! Create the displacement vector from the counts vector
    np_local = sum(rcounts)
    safe_allocate(rdispls(1:partition%npart))
    assert(ubound(rbuffer, 1) >= np_local)
 
    rdispls(1) = 0
    do ipart = 2, partition%npart
      rdispls(ipart) = rcounts(ipart-1) + rdispls(ipart-1)
    end do
 
    ! Collect the corresponding points to each process
    call partition%mpi_grp%alltoallv(sbuffer, scounts, sdispls, mpi_integer8, &
      rbuffer, rcounts, rdispls, mpi_integer8)
 
    safe_deallocate_a(sdispls)
    safe_deallocate_a(scounts)
    safe_deallocate_a(sbuffer)
    safe_deallocate_a(rcounts)
    safe_deallocate_a(rdispls)
 
    pop_sub(partition_get_local)
  end subroutine partition_get_local
end module partition_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: