Homec4science
Diffusion SPClibs (feature/fft-split) Browse

spclibs/pputils2/pptransp2.tpl
c723281d03b6feature/fft-split

Last Change
Raw File

pptransp2.tpl
View Options

!>
!> @file pptransp2.tpl
!>
!> @brief
!>
!> @copyright
!> Copyright (©) 2021 EPFL (Ecole Polytechnique Fédérale de Lausanne)
!> SPC (Swiss Plasma Center)
!>
!> SPClibs is free software: you can redistribute it and/or modify it under
!> the terms of the GNU Lesser General Public License as published by the Free
!> Software Foundation, either version 3 of the License, or (at your option)
!> any later version.
!>
!> SPClibs 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 Lesser General Public License
!> along with this program. If not, see <https://www.gnu.org/licenses/>.
!>
!> @author
!> (in alphabetical order)
!> @author Trach-Minh Tran <trach-minh.tran@epfl.ch>
!>
!
INTEGER :: me, npes, i, j, istr, iend, ierr
INTEGER, DIMENSION(:), ALLOCATABLE :: ids, idr
INTEGER, DIMENSION(:,:), ALLOCATABLE :: ndists
INTEGER, DIMENSION(:,:), ALLOCATABLE :: offsets
INTEGER :: dims(lastdim), np(2), npmx(2)
INTEGER :: n1p, nlp, n1pmx, nlpmx, bufsiz, scount
INTEGER :: status(MPI_STATUS_SIZE)
!----------------------------------------------------------------------
! 0. Prologue
!
CALL mpi_comm_rank(comm, me, ierr)
CALL mpi_comm_size(comm, npes, ierr)
!
! Determine send/receive proc. id
ALLOCATE(ids(npes), idr(npes))
CALL partners(comm, ids, idr)
!----------------------------------------------------------------------
! 1. Send/receive buffers
!
! Distribution of first and last partitionned dimensions
ALLOCATE(ndists(2,npes))
ALLOCATE(offsets(2,0:npes))
np(1) = SIZE(b,lastdim) ! Local first
np(2) = SIZE(a,lastdim) ! and last dimension
CALL mpi_allgather(np, 2, MPI_INTEGER, ndists, 2, MPI_INTEGER, comm, ierr)
offsets = 0
DO i=1,npes
offsets(:,i) = offsets(:,i-1) + ndists(:,i)
END DO
!
! Allocate send and receive 1d buffers
npmx = MAXVAL(ndists,2)
bufsiz = npmx(1)*npmx(2) ! Maximum size of send/receive buffers
DO i=2,lastdim-1
bufsiz = bufsiz * SIZE(a,i)
END DO
ALLOCATE(s_buf(bufsiz), r_buf(bufsiz) )
!----------------------------------------------------------------------
! 2. Exchange blocks
!
DO i=1,npes
istr = offsets(1,ids(i)) + 1 ! Partition a along first dim
iend = offsets(1,ids(i)+1)
dims = SHAPE(a)
dims(1) = iend-istr+1
scount = PRODUCT(dims)
s_buf(1:scount) = RESHAPE(a(istr:iend,:), (/scount/)) !*** dim dependant ***!
CALL MPI_SENDRECV(s_buf, scount, mpitype, ids(i), i,&
& r_buf, bufsiz, mpitype, idr(i), i,&
& comm, status, ierr)
istr = offsets(2,idr(i)) + 1 ! Partition b along first dim
iend = offsets(2,idr(i)+1)
dims = SHAPE(b)
dims(1) = iend-istr+1
b(istr:iend,:) = RESHAPE(r_buf, dims, order=(/lastdim, 1/)) !*** dim dependant ***!
END DO
!----------------------------------------------------------------------
! 9. Epilogue
!
DEALLOCATE(ids, idr)
DEALLOCATE(ndists, offsets)
DEALLOCATE(s_buf, r_buf)
!
c4science · Help