Homec4science
Diffusion SPClibs (feature/fft-split) Browse

spclibs/pputils2/pptransp3.tpl
03540012fedefeature/fft-split

Last Change
Raw File

pptransp3.tpl
View Options

!>
!> @file pptransp3.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/>.
!>
!> @authors
!> (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 dim1 and dim2 partitionned dimensions
ALLOCATE(ndists(2,npes))
ALLOCATE(offsets(2,0:npes))
np(1) = SIZE(b, dim2) ! Local first
np(2) = SIZE(a, dim2) ! and second 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=1,lastdim
IF ( (i .NE. dim1) .AND. (i .NE. dim2) ) bufsiz = bufsiz * SIZE(a,i)
END DO
ALLOCATE(s_buf(bufsiz), r_buf(bufsiz) )
!----------------------------------------------------------------------
! 2. Exchange blocks
!
IF ( (dim1 .EQ. 1) .AND. ( dim2 .EQ. 2 ) ) THEN !*** dim dependant ***!
recv_order = (/2,1,3/) !*** dim dependant ***!
ELSE IF ( (dim1 .EQ. 1) .AND. ( dim2 .eq. 3 ) ) THEN !*** dim dependant ***!
recv_order = (/3,2,1/) !*** dim dependant ***!
ELSE IF ( (dim1 .EQ. 2) .AND. ( dim2 .eq. 3 ) ) THEN !*** dim dependant ***!
recv_order = (/1,3,2/) !*** dim dependant ***!
ELSE
IF ( me .EQ. 0 ) THEN
WRITE(*, '(a,i4,a,i4,a)') 'pptransp3: Cannot handle case dim1 = ', dim1, ', dim2 = ', dim2, '!'
CALL mpi_abort(MPI_COMM_WORLD, -1, ierr)
END IF
END IF
!
DO i=1,npes
istr = offsets(1,ids(i)) + 1 ! Partition a along dimension dim1
iend = offsets(1,ids(i)+1)
dims = SHAPE(a)
dims(dim1) = iend-istr+1
scount = PRODUCT(dims)
IF (dim1 .EQ. 1) THEN !*** dim dependant ***!
s_buf(1:scount) = RESHAPE(a(istr:iend,:,:), (/scount/)) !*** dim dependant ***!
ELSE IF (dim1 .EQ. 2) THEN !*** dim dependant ***!
s_buf(1:scount) = RESHAPE(a(:,istr:iend,:), (/scount/)) !*** dim dependant ***!
END IF !*** 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 dimension dim1
iend = offsets(2,idr(i)+1)
dims = SHAPE(b)
dims(dim1) = iend-istr+1
IF (dim1 .EQ. 1) THEN !*** dim dependant ***!
b(istr:iend,:,:) = RESHAPE(r_buf, dims, order=recv_order) !*** dim dependant ***!
ELSE IF (dim1 .EQ. 2) THEN !*** dim dependant ***!
b(:,istr:iend,:) = RESHAPE(r_buf, dims, order=recv_order) !*** dim dependant ***!
END IF !*** dim dependant ***!
END DO
!----------------------------------------------------------------------
! 9. Epilogue
!
DEALLOCATE(ids, idr)
DEALLOCATE(ndists, offsets)
DEALLOCATE(s_buf, r_buf)
!
c4science · Help