-
Notifications
You must be signed in to change notification settings - Fork 67
/
mpi_transpose_mod.f90
119 lines (102 loc) · 3.66 KB
/
mpi_transpose_mod.f90
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
!!
!! Copyright (C) 2009-2017 Johns Hopkins University
!!
!! This file is part of lesgo.
!!
!! lesgo 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 3 of the License, or
!! (at your option) any later version.
!!
!! lesgo 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 lesgo. If not, see <http://www.gnu.org/licenses/>.
!!
!*******************************************************************************
module mpi_transpose_mod
!*******************************************************************************
use types, only : rprec
use param, only : nproc, comm, coord, ierr, MPI_CPREC
use mpi
implicit none
save
private
public :: mpi_transpose
contains
!*******************************************************************************
subroutine mpi_transpose (mx, my, mz, a, b)
!*******************************************************************************
!
!--the sizes of a, b are assumed shape here to allow us to use arrays
! dimensioned for the Nyquist frequency in x-direction
!--only a(1:mx, 1:my, 1:mz) & b(1:mz*nproc, 1:my, 1:mx/nproc) are used
!
implicit none
! declare as arguments so untransposing does not require additional code
integer, intent (in) :: mx, my, mz
complex(rprec), intent (in) :: a(:, :, :)
complex(rprec), intent (out) :: b(:, :, :)
integer :: bs
integer :: ip
integer :: up, down
integer :: status(MPI_STATUS_SIZE)
integer :: i, k, jx, jz
logical, save :: init = .false.
complex(rprec) :: tmpout(mx/nproc, my, mz), tmpin(mx/nproc, my, mz)
integer, save, allocatable, dimension(:) :: src, dest
logical, save :: arrays_allocated = .false.
if( .not. arrays_allocated ) then
allocate(src(nproc-1))
allocate(dest(nproc-1))
arrays_allocated = .true.
endif
if (.not. init) then
do ip = 1, nproc-1
! this is a bit awkward: really want periodic topology so can
! use cart_shift here but this is not good for the finite
! differences, so perhaps create a new cartisian topology from
! existing one, but we no not allow reordering and we do allow
! periodicity
up = modulo (coord + ip, nproc) ! corresponds to dest(ip)
down = modulo (coord - ip, nproc) ! corresponds to src(ip)
call MPI_cart_rank (comm, (/ up /), dest(ip), ierr)
call MPI_cart_rank (comm, (/ down /), src(ip), ierr)
end do
init = .true.
end if
!--block size
bs = mx*my*mz/nproc
do ip = 1, nproc-1
up = modulo (coord + ip, nproc) ! corresponds to dest(ip)
down = modulo (coord - ip, nproc) ! corresponds to src(ip)
! copy chunk "up" into buffer (no local transpose)
do jz = 1, mz
do jx = 1, mx/nproc
tmpout(jx, :, jz) = a(up*mx/nproc+jx, :, jz)
end do
end do
call MPI_sendrecv (tmpout(1,1,1), bs, MPI_CPREC, dest(ip), ip, &
tmpin(1,1,1), bs, MPI_CPREC, src(ip), ip, comm, status, ierr)
! copy chunk "down" from tmpin to b, in transposed order
do i = 1, mx/nproc
do k = 1, mz
jz = down*mz + k
b(jz, :, i) = tmpin(i, :, k)
end do
end do
end do
! local transpose on non-transferred data
! chunk 'rank' should not have have been sent/received
do i = 1, mx/nproc
do k = 1, mz
jx = coord*mx/nproc + i
jz = coord*mz + k
b(jz, :, i) = a(jx, :, k)
end do
end do
end subroutine mpi_transpose
end module mpi_transpose_mod