Branch data Line data Source code
1 : : !!****m* ABINIT/m_xgTransposer
2 : : !! NAME
3 : : !! m_xgTransposer
4 : : !!
5 : : !! FUNCTION
6 : : !! This module is to be used to go to "KGB" representation and to "linear
7 : : !! algebra representation". It will replace most of prep_* subroutine
8 : : !! This should really help to do the transposition operation.
9 : : !!
10 : : !! NOTES
11 : : !! The transposer switches between two states. Let's assume for simplicity
12 : : !! that we have four MPI processes in the comm_cols communicator and
13 : : !! a single process in the comm_rows communicator:
14 : : !!
15 : : !! STATE_LINALG= STATE_COLSROWS=
16 : : !!
17 : : !! bands bands
18 : : !! |-------------------| |----|----|----|----|
19 : : !! | P0 | | | | | |
20 : : !! | | | | | | |
21 : : !! |-------------------| | | | | |
22 : : !! | P1 | | | | | |
23 : : !! | | | | | | |
24 : : !! pw |-------------------| pw | P0 | P1 | P2 | P3 |
25 : : !! | P2 | | | | | |
26 : : !! | | | | | | |
27 : : !! |-------------------| | | | | |
28 : : !! | P3 | | | | | |
29 : : !! | | | | | | |
30 : : !! |-------------------| |----|----|----|----|
31 : : !!
32 : : !! The user can define custom block sizes using ncolsColsRows_sub and nrowsLinalg_sub.
33 : : !!
34 : : !! TODO IML 04/04/2025
35 : : !! -clean versions to always pass from my array version
36 : : !! for ncolsColsRows or create wrappers so that if we
37 : : !! give an array then we turn it into scalar
38 : : !! -add unitary tests for my custom transposition
39 : : !!
40 : : !! COPYRIGHT
41 : : !! Copyright (C) 2017-2026 ABINIT group (J. Bieder, L. Baguet, IML)
42 : : !! This file is distributed under the terms of the
43 : : !! GNU General Public License, see ~abinit/COPYING
44 : : !! or http://www.gnu.org/copyleft/gpl.txt .
45 : : !!
46 : : !! NOTES
47 : : !!
48 : : !! SOURCE
49 : :
50 : : #if defined HAVE_CONFIG_H
51 : : #include "config.h"
52 : : #endif
53 : :
54 : : #include "abi_common.h"
55 : :
56 : : ! nvtx related macro definition
57 : : #include "nvtx_macros.h"
58 : :
59 : : module m_xgTransposer
60 : :
61 : : use, intrinsic :: iso_c_binding, only: c_double, c_size_t, c_loc
62 : :
63 : : use defs_basis, only : std_err, std_out, dp, ABI_GPU_KOKKOS, ABI_GPU_OPENMP, ABI_GPU_DISABLED
64 : : use m_xomp
65 : : use m_profiling_abi
66 : : USE_MPI
67 : : use m_xmpi
68 : : use m_errors
69 : : use m_xg
70 : : use m_time
71 : :
72 : : #if defined HAVE_YAKL
73 : : use gator_mod
74 : : #endif
75 : :
76 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_YAKL)
77 : : use m_gpu_toolbox, only : CPU_DEVICE_ID, gpu_device_synchronize, gpu_data_prefetch_async
78 : : #endif
79 : :
80 : : #if defined(HAVE_GPU_MARKERS)
81 : : use m_nvtx_data
82 : : #endif
83 : :
84 : : implicit none
85 : :
86 : : #ifdef HAVE_MPI1
87 : : include 'mpif.h'
88 : : #endif
89 : :
90 : : private
91 : :
92 : : integer, parameter, public :: STATE_LINALG = 1
93 : : integer, parameter, public :: STATE_COLSROWS = 2
94 : : integer, parameter, public :: STATE_UNKNOW = 3
95 : : integer, parameter, public :: MPI_LINALG = 1
96 : : integer, parameter, public :: MPI_ROWS = 2
97 : : integer, parameter, public :: MPI_COLS = 3
98 : : integer, parameter, public :: MPI_2DCART = 4
99 : : integer, parameter, public :: TRANS_ALL2ALL = 1
100 : : integer, parameter, public :: TRANS_GATHER = 2
101 : : integer, parameter :: TRANS_TYPE_CONSTRUCTED = 1
102 : : integer, parameter :: TRANS_TYPE_CONSTRUCTED_NULL_COMM = 2
103 : : integer, parameter :: TRANS_TYPE_COPIED = 3
104 : :
105 : : integer, parameter :: tim_toColsRows = 1662
106 : : integer, parameter :: tim_toLinalg = 1663
107 : : integer, parameter :: tim_all2allv = 1664
108 : : integer, parameter :: tim_gatherv = 1665
109 : : integer, parameter :: tim_reorganize = 1666
110 : : integer, parameter :: tim_init = 1667
111 : : integer, parameter :: tim_free = 1668
112 : : integer, parameter :: tim_transpose = 1669
113 : :
114 : : type, private :: mpiData_t
115 : : integer :: comm
116 : : integer :: rank
117 : : integer :: size
118 : : end type mpiData_t
119 : :
120 : : type, private :: ptr_t
121 : : double precision, pointer :: ptr(:,:) => null()
122 : : end type ptr_t
123 : :
124 : : type, public :: xgTransposer_t
125 : : type(xgBlock_t), pointer, private :: xgBlock_linalg => null()
126 : : type(xgBlock_t), pointer, private :: xgBlock_colsrows => null()
127 : : integer :: state
128 : : type(mpiData_t), private :: mpiData(4)
129 : : integer, allocatable, private :: lookup(:)
130 : : integer, pointer, private :: nrowsLinalg(:) => null()
131 : : integer, pointer, private :: ncolsColsRows_sub(:) => null() ! allows different bandpp per MPI process
132 : : integer :: nspinor
133 : : integer :: nrowsColsRows
134 : : integer :: ncolsColsRows
135 : : integer :: ncolsLinalg
136 : : integer :: mpiAlgo
137 : : integer :: type
138 : : integer :: me_g0_fft
139 : : integer :: gpu_option = ABI_GPU_DISABLED
140 : : integer :: gpu_kokkos_nthrd = 1
141 : : integer :: gpu_thread_limit = 1
142 : : logical :: custom_ncolsColsRows = .false. ! use ncolsColsRows_sub instead of ncolsColsRows
143 : : real(dp), contiguous, pointer:: buffer(:,:) => null()
144 : : end type xgTransposer_t
145 : :
146 : : public :: xgTransposer_constructor
147 : : public :: xgTransposer_copyConstructor
148 : : public :: xgTransposer_transpose
149 : : public :: xgTransposer_getRank
150 : : public :: xgTransposer_getComm
151 : : public :: xgTransposer_free
152 : :
153 : : contains
154 : : !!***
155 : :
156 : : !!****f* m_xgTransposer/xgTransposer_constructor
157 : : !!
158 : : !! NAME
159 : : !! xgTransposer_constructor
160 : :
161 : 46190 : subroutine xgTransposer_constructor(xgTransposer,xgBlock_linalg,xgBlock_colsrows,nspinor,&
162 : : state,algo,comm_rows,comm_cols,ncpu_cols,ncpu_rows,me_g0_fft,gpu_option,gpu_thread_limit,&
163 [ + + + - : 46190 : custom_ncolsColsRows,ncolsColsRows_sub,nrowsLinalg_sub)
- + - + -
- - - ]
164 : :
165 : : type(xgTransposer_t) , intent(inout) :: xgTransposer
166 : : type(xgBlock_t), target, intent(in ) :: xgBlock_linalg
167 : : type(xgBlock_t), target, intent(in ) :: xgBlock_colsrows
168 : : integer , intent(in ) :: comm_rows,comm_cols
169 : : integer , intent(in ) :: ncpu_rows,ncpu_cols
170 : : integer , intent(in ) :: nspinor
171 : : integer , intent(in ) :: state
172 : : integer , intent(in ) :: algo
173 : : integer , intent(in ) :: me_g0_fft
174 : : integer, optional , intent(in ) :: gpu_option,gpu_thread_limit
175 : : logical, optional , intent(in ) :: custom_ncolsColsRows ! query for pointer
176 : : integer, optional, target, intent(in) :: ncolsColsRows_sub(:)
177 : : integer, optional, target, intent(in) :: nrowsLinalg_sub(:)
178 : : integer :: commLinalg
179 : : integer :: ncols
180 : : integer :: nrows
181 : : integer :: ierr
182 : : integer :: icol
183 : : integer :: ncpuRows
184 : : integer :: ncpuCols
185 : : logical :: custom_ncolsColsRows_
186 : : #if defined HAVE_MPI
187 : : integer :: comm_rows_,comm_cols_
188 : : #endif
189 : : character(len=500) :: message
190 : : double precision :: tsec(2)
191 : :
192 : 46190 : call timab(tim_init,1,tsec)
193 : :
194 : 46190 : xgTransposer%type = TRANS_TYPE_CONSTRUCTED
195 : :
196 : 46190 : xgTransposer%xgBlock_linalg => xgBlock_linalg
197 : 46190 : xgTransposer%xgBlock_colsrows => xgBlock_colsrows
198 : 46190 : xgTransposer%state = state
199 : 46190 : xgTransposer%nspinor = nspinor
200 : 46190 : xgTransposer%me_g0_fft = me_g0_fft
201 : 46190 : xgTransposer%gpu_option = ABI_GPU_DISABLED
202 [ + + ]: 46190 : if(present(gpu_option)) xgTransposer%gpu_option = gpu_option
203 : 46190 : xgTransposer%gpu_thread_limit = 1
204 [ + + ]: 46190 : if(present(gpu_thread_limit)) xgTransposer%gpu_thread_limit = gpu_thread_limit
205 [ - + ]: 46190 : if (state==STATE_COLSROWS) then
206 [ # # # # ]: 0 : if (comm_rows==xmpi_comm_null.or.comm_cols==xmpi_comm_null) then
207 : 0 : ABI_ERROR('STATE_COLSROWS requires input comms')
208 : : end if
209 : 0 : commLinalg = comm_cols
210 [ + - ]: 46190 : else if (state==STATE_LINALG) then
211 : 46190 : commLinalg = comm(xgBlock_linalg)
212 : : else
213 : 0 : ABI_ERROR("Invalid transposer state")
214 : : end if
215 : 46190 : xgTransposer%mpiData(MPI_LINALG)%comm = commLinalg
216 : 46190 : xgTransposer%mpiData(MPI_LINALG)%rank = xmpi_comm_rank(commLinalg)
217 : 46190 : xgTransposer%mpiData(MPI_LINALG)%size = xmpi_comm_size(commLinalg)
218 : :
219 : : #if defined HAVE_MPI
220 [ + + + - ]: 46190 : if (comm_rows==xmpi_comm_null.and.comm_cols==xmpi_comm_null) then
221 : 10 : xgTransposer%type = TRANS_TYPE_CONSTRUCTED_NULL_COMM
222 : 10 : ncpuCols = ncpu_cols
223 : 10 : ncpuRows = ncpu_rows
224 [ - + ]: 46180 : else if (comm_rows==xmpi_comm_null) then
225 : 0 : ABI_ERROR("Comm_rows and comm_cols should have the same status : null or defined. Here only comm_rows is null.")
226 [ - + ]: 46180 : else if (comm_cols==xmpi_comm_null) then
227 : 0 : ABI_ERROR("Comm_rows and comm_cols should have the same status : null or defined. Here only comm_cols is null.")
228 : : else
229 : 46180 : xgTransposer%mpiData(MPI_2DCART)%comm = xgTransposer%mpiData(MPI_LINALG)%comm
230 : 46180 : ncpuRows = xmpi_comm_size(comm_rows)
231 : 46180 : ncpuCols = xmpi_comm_size(comm_cols)
232 : : end if
233 : : #else
234 : : ncpuRows = 1
235 : : ncpuCols = 1
236 : : #endif
237 [ - + ]: 46190 : if ( xgTransposer%mpiData(MPI_LINALG)%size < ncpuCols*ncpuRows ) then
238 : 0 : write(message,'(a,i6,a,i6,a)') "There is not enough MPI processes in the communication (", &
239 : 0 : xgTransposer%mpiData(MPI_LINALG)%size, "). Need at least ", ncpuCols*ncpuRows, " processes"
240 : 0 : ABI_ERROR(message)
241 : : end if
242 : :
243 [ + - ]: 46190 : if ( ( algo == TRANS_ALL2ALL .or. algo == TRANS_GATHER ) ) then
244 : 46190 : xgTransposer%mpiAlgo = algo
245 : : else
246 : 0 : xgTransposer%mpiAlgo = TRANS_ALL2ALL
247 : 0 : ABI_COMMENT("Bad value for transposition MPI_algo. Will use ALLTOALL")
248 : : end if
249 : :
250 : : !if ( xgTransposer%mpiAlgo == TRANS_ALL2ALL ) then
251 : : ! ABI_COMMENT("Using mpi_alltoall for transposition")
252 : : !else
253 : : ! ABI_COMMENT("Using mpi_gatherv for transposition")
254 : : !end if
255 : :
256 : : ! Nullify pointers in order to be able to query them
257 : 46190 : xgTransposer%nrowsLinalg => null()
258 : 46190 : xgTransposer%buffer => null()
259 : 46190 : xgTransposer%ncolsColsRows_sub => null()
260 : :
261 : : ! Attention: Do not check condition present() on the pointer because in Fortran
262 : : ! "An unallocated variable passed as an argument is not PRESENT"
263 : : ! in other words declared pointers with undefined status cannot be queried
264 : 46190 : custom_ncolsColsRows_ = .false.
265 [ + + ]: 46190 : if (present(custom_ncolsColsRows)) custom_ncolsColsRows_ = custom_ncolsColsRows
266 : 46174 : xgTransposer%custom_ncolsColsRows = custom_ncolsColsRows_
267 [ + - ]: 16 : if (custom_ncolsColsRows_) then
268 [ - + ]: 16 : if ( ncpuRows > 1 ) then
269 : 0 : write(message,'(a,i6)') "Custom ncolsColsRows not implemented for number of row MPI processes=", ncpuRows
270 : 0 : ABI_ERROR(message)
271 : : end if
272 : 16 : xgTransposer%custom_ncolsColsRows = .true.
273 [ + - ]: 16 : if (state==STATE_LINALG) xgTransposer%ncolsColsRows_sub => ncolsColsRows_sub
274 : : ! if state colsrows in will be constructed later in readDistribution
275 : : end if
276 : :
277 : : ! The following allocates memory buffer associated either to xgBlock_linalg or to xgBlock_colsrows
278 : : !! Assumes that only one of them is already allocated
279 : 46190 : select case (state)
280 : : case (STATE_LINALG)
281 : :
282 : : ! We are in the linalg representation.
283 : : ! We need to construct the colsrows parallelization
284 : 46190 : call xgBlock_getSize(xgBlock_linalg,nrows,ncols)
285 : :
286 : : ! Get total number of rows in ColsRows representation
287 : 46190 : call xmpi_sum(nrows,commLinalg,ierr)
288 : :
289 : : ! Check that ncols can be uniformly distributed across ncolsCols MPI processes
290 : : ! Load balance is uniform, each MPI process has the same ncolsColsRows
291 [ + + - + ]: 46190 : if ( (.not.xgTransposer%custom_ncolsColsRows) .and. MOD(ncols,ncpuCols) /=0 ) then
292 [ # # ]: 0 : if ( ncols > ncpuCols ) then
293 : 0 : write(message,'(a,i6,a,i6,a)') "Unbalanced parallelization : ", ncols, " columns for ", ncpuCols, " MPI"
294 : 0 : ABI_ERROR(message)
295 : : else
296 : 0 : write(message,'(i6,a)') (ncpuCols-ncols)*ncpuRows, " MPI will not be used because of the number of columns!!"
297 : 0 : ABI_ERROR(message)
298 : : !ncpuCols = ncols
299 : : end if
300 : : end if
301 : :
302 : : ! Build the lookup table that maps each column to an MPI process
303 [ - + + - : 138570 : ABI_MALLOC(xgTransposer%lookup,(1:ncols))
+ - + - -
+ - + -
+ ]
304 : : !ABI_MALLOC(xgTransposer%me_g0_lookup,())
305 [ + + ]: 477742 : do icol = 0, ncols-1
306 : 477742 : xgTransposer%lookup(icol+1) = MOD(icol,ncpuCols)
307 : : end do
308 : :
309 : : #if defined HAVE_MPI
310 [ + + + - ]: 46190 : if (comm_rows==xmpi_comm_null.and.comm_cols==xmpi_comm_null) then
311 : 10 : call xgTransposer_makeComm(xgTransposer,ncpuRows,ncpuCols,comm_rows_,comm_cols_)
312 : : else
313 : 46180 : comm_rows_=comm_rows
314 : 46180 : comm_cols_=comm_cols
315 : : end if
316 : 46190 : xgTransposer%mpiData(MPI_ROWS)%comm = comm_rows_
317 : 46190 : xgTransposer%mpiData(MPI_COLS)%comm = comm_cols_
318 : : #else
319 : : xgTransposer%mpiData(MPI_ROWS)%comm = xmpi_comm_null
320 : : xgTransposer%mpiData(MPI_COLS)%comm = xmpi_comm_null
321 : : #endif
322 : 46190 : call xgTransposer_setComm(xgTransposer)
323 : 46190 : call xgTransposer_computeDistribution(xgTransposer)
324 : 92380 : call xgTransposer_makeXgBlock(xgTransposer)
325 : :
326 : : case (STATE_COLSROWS)
327 : :
328 : : !write(std_out,'(a,i2,i2)') 'db @xgTransposer constructor enter case for # procs', ncpuRows,ncpuCols
329 : : !call flush_unit(std_out)
330 : :
331 : : ! We are in the colsrows representation.
332 : : ! We need to construct the linalg parallelization
333 : :
334 : : ! Load balance
335 : : ! nrows are distributed uniformly across ncpuCols MPI processes for 1,...,ncpuCols-1
336 : : ! The last MPI process uses remaining nrows. Read distribution from nrowsLinalg_sub
337 : 0 : call xgBlock_getSize(xgBlock_colsrows,nrows,ncols)
338 [ # # # # : 0 : ABI_MALLOC(xgTransposer%lookup,(1:ncols))
# # # # #
# # # #
# ]
339 : :
340 : : #if defined HAVE_MPI
341 : 0 : xgTransposer%mpiData(MPI_ROWS)%comm = comm_rows
342 : 0 : xgTransposer%mpiData(MPI_COLS)%comm = comm_cols
343 : : #else
344 : : xgTransposer%mpiData(MPI_ROWS)%comm = xmpi_comm_null
345 : : xgTransposer%mpiData(MPI_COLS)%comm = xmpi_comm_null
346 : : #endif
347 : :
348 : 0 : call xgTransposer_setComm(xgTransposer)
349 [ # # # # : 0 : ABI_MALLOC(xgTransposer%nrowsLinalg,(xgTransposer%mpiData(MPI_LINALG)%size))
# # # # #
# # # ]
350 [ # # ]: 0 : xgTransposer%nrowsLinalg(:) = nrowsLinalg_sub(:)
351 : 0 : call xgTransposer_readDistribution(xgTransposer)
352 : 0 : call xgTransposer_makeXgBlock(xgTransposer)
353 : :
354 : : case default
355 [ + - - ]: 46190 : ABI_ERROR("State is undefined")
356 : : end select
357 : :
358 : 46190 : call timab(tim_init,2,tsec)
359 : :
360 : : !write(*,*) 'transposer construction done'
361 : :
362 : 46190 : end subroutine xgTransposer_constructor
363 : : !!***
364 : :
365 : : !!****f* m_xgTransposer/xgTransposer_copyConstructor
366 : : !!
367 : : !! NAME
368 : : !! xgTransposer_copyConstructor
369 : :
370 : 144672 : subroutine xgTransposer_copyConstructor(xgTransposer,xgTransposerInitialized,xgBlock_linalg,xgBlock_colsrows,state)
371 : :
372 : : type(xgTransposer_t) , intent(inout) :: xgTransposer
373 : : type(xgTransposer_t) , intent(in ) :: xgTransposerInitialized
374 : : type(xgBlock_t), target, intent(in ) :: xgBlock_linalg
375 : : type(xgBlock_t), target, intent(in ) :: xgBlock_colsrows
376 : : integer , intent(in ) :: state
377 : : integer :: commLinalg
378 : : integer :: ncols, ncpuCols
379 : : integer :: nrows, ncpuRows
380 : : integer :: ierr
381 : : integer :: icol
382 : : character(len=500) :: message
383 : : double precision :: tsec(2)
384 : :
385 : 144672 : call timab(tim_init,1,tsec)
386 : :
387 : 144672 : xgTransposer%type = TRANS_TYPE_COPIED
388 : :
389 : 144672 : xgTransposer%xgBlock_linalg => xgBlock_linalg
390 : 144672 : xgTransposer%xgBlock_colsrows => xgBlock_colsrows
391 : 144672 : xgTransposer%state = state
392 [ - + ]: 144672 : if (state==STATE_COLSROWS) then
393 : 0 : commLinalg = xgTransposerInitialized%mpiData(MPI_LINALG)%comm
394 [ + - ]: 144672 : else if (state==STATE_LINALG) then
395 : 144672 : commLinalg = comm(xgBlock_linalg)
396 : : else
397 : 0 : ABI_ERROR("Invalid transposer state")
398 : : end if
399 : :
400 [ - + ]: 144672 : if ( commLinalg /= xgTransposerInitialized%mpiData(MPI_LINALG)%comm ) then
401 : 0 : ABI_ERROR("Linalg communicators are different for the two transposers, this is not allowed.")
402 : : end if
403 : :
404 [ + + ]: 723360 : xgTransposer%mpiData(:)%comm = xgTransposerInitialized%mpiData(:)%comm
405 : 144672 : xgTransposer%mpiData(MPI_LINALG)%rank = xmpi_comm_rank(commLinalg)
406 : 144672 : xgTransposer%mpiData(MPI_LINALG)%size = xmpi_comm_size(commLinalg)
407 : 144672 : xgTransposer%mpiData(MPI_COLS)%rank = xmpi_comm_rank(xgTransposer%mpiData(MPI_COLS)%comm)
408 : 144672 : xgTransposer%mpiData(MPI_COLS)%size = xmpi_comm_size(xgTransposer%mpiData(MPI_COLS)%comm)
409 : 144672 : xgTransposer%mpiData(MPI_ROWS)%rank = xmpi_comm_rank(xgTransposer%mpiData(MPI_ROWS)%comm)
410 : 144672 : xgTransposer%mpiData(MPI_ROWS)%size = xmpi_comm_size(xgTransposer%mpiData(MPI_ROWS)%comm)
411 : :
412 : 144672 : xgTransposer%mpiAlgo = xgTransposerInitialized%mpiAlgo
413 : 144672 : xgTransposer%nspinor = xgTransposerInitialized%nspinor
414 : 144672 : xgTransposer%me_g0_fft = xgTransposerInitialized%me_g0_fft
415 : 144672 : xgTransposer%gpu_option = xgTransposerInitialized%gpu_option
416 : :
417 : : ! Nullify pointers in order to be able to query them
418 : 144672 : xgTransposer%nrowsLinalg => null()
419 : 144672 : xgTransposer%buffer => null()
420 : 144672 : xgTransposer%ncolsColsRows_sub => null()
421 : : !write(std_out,*) '@nrowsLinalg', associated(xgTransposer%nrowsLinalg)
422 : :
423 : 144672 : xgTransposer%custom_ncolsColsRows = xgTransposerInitialized%custom_ncolsColsRows
424 [ - + - - ]: 144672 : if(xgTransposer%custom_ncolsColsRows .and. state==STATE_LINALG) then
425 : 0 : xgTransposer%ncolsColsRows_sub => xgTransposerInitialized%ncolsColsRows_sub
426 : : ! if state colsrows in will be constructed later in readDistribution
427 : : end if
428 : :
429 : 144672 : ncpuCols = xgTransposer%mpiData(MPI_COLS)%size
430 : 144672 : ncpuRows = xgTransposer%mpiData(MPI_ROWS)%size
431 : :
432 : 144672 : select case (state)
433 : : case (STATE_LINALG)
434 : :
435 : 144672 : call xgBlock_getSize(xgBlock_linalg,nrows,ncols)
436 : :
437 : : ! Get total number of rows in ColsRows representation
438 : 144672 : call xmpi_sum(nrows,commLinalg,ierr)
439 : :
440 : : ! Check that ncols can be uniformly distributed across ncolsCols MPI processes
441 [ + - - + ]: 144672 : if ( (.not.xgTransposer%custom_ncolsColsRows) .and. MOD(ncols,ncpuCols) /=0 ) then
442 [ # # ]: 0 : if ( ncols > ncpuCols ) then
443 : 0 : write(message,'(a,i6,a,i6,a)') "Unbalanced parallelization : ", ncols, " columns for ", ncpuCols, " MPI"
444 : 0 : ABI_ERROR(message)
445 : : else
446 : 0 : write(message,'(i6,a)') (ncpuCols-ncols)*ncpuRows, " MPI will not be used because of the number of columns!!"
447 : 0 : ABI_ERROR(message)
448 : : end if
449 : : end if
450 : :
451 : : ! Build the lookup table that maps each column to an MPI process
452 : : ! Assures load balance so that each MPI process has the same charge
453 [ - + + - : 434016 : ABI_MALLOC(xgTransposer%lookup,(1:ncols))
+ - + - -
+ - + -
+ ]
454 [ - + ]: 144672 : if ( cols(xgTransposerInitialized%xgBlock_linalg) /= ncols ) then
455 [ # # ]: 0 : do icol = 0, ncols-1
456 : 0 : xgTransposer%lookup(icol+1) = MOD(icol,ncpuCols)
457 : : end do
458 : : else
459 [ + + ]: 1176400 : xgTransposer%lookup(:) = xgTransposerInitialized%lookup(:)
460 : : end if
461 : :
462 : 144672 : call xgTransposer_computeDistribution(xgTransposer)
463 : 289344 : call xgTransposer_makeXgBlock(xgTransposer)
464 : :
465 : : case (STATE_COLSROWS)
466 : :
467 : : !write(std_out,*) 'db @xgTransposer copyConstructor enter case'
468 : : !call flush_unit(std_out)
469 : :
470 : 0 : nrows = rows(xgBlock_colsrows)
471 : 0 : ncols = cols(xgBlock_colsrows)
472 [ # # # # : 0 : if ( (.not.xgTransposer%custom_ncolsColsRows) .and. MOD(nrows,ncpuCols) /=0 .and. nrows > ncpuCols) then
# # ]
473 : 0 : write(message,'(a,i6,a,i6,a)') "Unbalanced parallelization : ", nrows, " rows for ", ncpuCols, " MPI"
474 : 0 : ABI_COMMENT(message)
475 : : end if
476 [ # # # # : 0 : ABI_MALLOC(xgTransposer%lookup,(1:ncols))
# # # # #
# # # #
# ]
477 : :
478 [ # # # # : 0 : ABI_MALLOC(xgTransposer%nrowsLinalg,(xgTransposer%mpiData(MPI_LINALG)%size))
# # # # #
# # # ]
479 [ # # ]: 0 : xgTransposer%nrowsLinalg(:) = xgTransposerInitialized%nrowsLinalg(:)
480 : 0 : call xgTransposer_readDistribution(xgTransposer)
481 : 0 : call xgTransposer_makeXgBlock(xgTransposer)
482 : :
483 : : case default
484 [ + - - ]: 144672 : ABI_ERROR("State is undefined")
485 : : end select
486 : :
487 : 144672 : call timab(tim_init,2,tsec)
488 : :
489 : 144672 : end subroutine xgTransposer_copyConstructor
490 : : !!***
491 : :
492 : : !!****f* m_xgTransposer/xgTransposer_setComm
493 : : !!
494 : : !! NAME
495 : : !! xgTransposer_setComm
496 : :
497 : 46190 : subroutine xgTransposer_setComm(xgTransposer)
498 : :
499 : : type(xgTransposer_t), intent(inout) :: xgTransposer
500 : :
501 : 46190 : xgTransposer%mpiData(MPI_ROWS)%rank = xmpi_comm_rank(xgTransposer%mpiData(MPI_ROWS)%comm)
502 : 46190 : xgTransposer%mpiData(MPI_ROWS)%size = xmpi_comm_size(xgTransposer%mpiData(MPI_ROWS)%comm)
503 : 46190 : xgTransposer%mpiData(MPI_COLS)%rank = xmpi_comm_rank(xgTransposer%mpiData(MPI_COLS)%comm)
504 : 46190 : xgTransposer%mpiData(MPI_COLS)%size = xmpi_comm_size(xgTransposer%mpiData(MPI_COLS)%comm)
505 : :
506 : 92380 : select case(xgTransposer%state)
507 : : case (STATE_LINALG)
508 : 46190 : call xgBlock_setComm(xgTransposer%xgBlock_colsrows,xgTransposer%mpiData(MPI_ROWS)%comm)
509 : : case (STATE_COLSROWS)
510 : 0 : call xgBlock_setComm(xgTransposer%xgBlock_linalg,xgTransposer%mpiData(MPI_ROWS)%comm)
511 : : case default
512 [ + - - ]: 46190 : ABI_ERROR("State is undefined")
513 : : end select
514 : :
515 : 46190 : end subroutine xgTransposer_setComm
516 : : !!***
517 : :
518 : : !!****f* m_xgTransposer/xgTransposer_makeComm
519 : : !!
520 : : !! NAME
521 : : !! xgTransposer_makeComm
522 : :
523 : 10 : subroutine xgTransposer_makeComm(xgTransposer,ncpuRows,ncpuCols,comm_rows,comm_cols)
524 : :
525 : : type(xgTransposer_t), intent(inout) :: xgTransposer
526 : : integer , intent(in ) :: ncpuRows
527 : : integer , intent(in ) :: ncpuCols
528 : : integer , intent(out ) :: comm_rows
529 : : integer , intent(out ) :: comm_cols
530 : : integer :: commColsRows
531 : : #if defined HAVE_MPI
532 : : integer :: sizeGrid(2) !coordInGrid(2),
533 : : logical :: periodic(2), selectDim(2), reorder
534 : : integer :: ierr
535 : :
536 : 10 : sizeGrid(1) = ncpuCols
537 : 10 : sizeGrid(2) = ncpuRows
538 : : !sizeGrid(1) = ncpuRows
539 : : !sizeGrid(2) = ncpuCols
540 : 10 : periodic = (/ .false., .false. /)
541 : 10 : reorder = .false.
542 : 10 : call mpi_cart_create(xgTransposer%mpiData(MPI_LINALG)%comm,2,sizeGrid,periodic,reorder,commColsRows,ierr)
543 [ - + ]: 10 : if ( ierr /= xmpi_success ) then
544 : 0 : ABI_ERROR("xgTransposer failed to creat cartesian grid")
545 : : end if
546 : 10 : xgTransposer%mpiData(MPI_2DCART)%comm = commColsRows
547 : :
548 : 10 : selectDim = (/ .false., .true. /)
549 : : !selectDim = (/ .true., .false. /)
550 : 10 : call mpi_cart_sub(commColsRows, selectDim, xgTransposer%mpiData(MPI_ROWS)%comm,ierr)
551 [ - + ]: 10 : if ( ierr /= xmpi_success ) then
552 : 0 : ABI_ERROR("xgTransposer failed to creat rows communicator")
553 : : end if
554 : 10 : selectDim = (/ .true., .false. /)
555 : : !selectDim = (/ .false., .true. /)
556 : 10 : call mpi_cart_sub(commColsRows, selectDim, xgTransposer%mpiData(MPI_COLS)%comm,ierr)
557 [ - + ]: 10 : if ( ierr /= xmpi_success ) then
558 : 0 : ABI_ERROR("xgTransposer failed to creat columns communicator")
559 : : end if
560 : : #else
561 : : commColsRows = xmpi_comm_null
562 : : xgTransposer%mpiData(MPI_ROWS)%comm = xmpi_comm_null
563 : : xgTransposer%mpiData(MPI_COLS)%comm = xmpi_comm_null
564 : : #endif
565 : 10 : comm_rows = xgTransposer%mpiData(MPI_ROWS)%comm
566 : 10 : comm_cols = xgTransposer%mpiData(MPI_COLS)%comm
567 : :
568 : 10 : end subroutine xgTransposer_makeComm
569 : : !!***
570 : :
571 : : !!****f* m_xgTransposer/xgTransposer_computeDistribution
572 : : !!
573 : : !! NAME
574 : : !! xgTransposer_computeDistribution
575 : :
576 : 190862 : subroutine xgTransposer_computeDistribution(xgTransposer)
577 : :
578 : : type(xgTransposer_t), intent(inout) :: xgTransposer
579 : : integer :: nRealPairs
580 : : integer :: ierr
581 : : integer :: icpu_cols,icpu_rows
582 : : integer :: ncpuCols,ncpuRows
583 : : integer :: ncolsColsRowsMe
584 : :
585 [ - + + - : 572586 : ABI_MALLOC(xgTransposer%nrowsLinalg,(xgTransposer%mpiData(MPI_LINALG)%size))
+ - + - -
+ - + ]
586 : 190862 : nRealPairs = rows(xgTransposer%xgBlock_linalg)
587 [ - + ]: 190862 : if (MOD(nRealPairs,xgTransposer%nspinor)/=0) then
588 : 0 : ABI_ERROR('nspinor should divide nRealPairs!')
589 : : end if
590 : :
591 : 190862 : call xmpi_allgather(nRealPairs,xgTransposer%nrowsLinalg,xgTransposer%mpiData(MPI_LINALG)%comm,ierr)
592 [ - + ]: 190862 : if ( ierr /= xmpi_success ) then
593 : 0 : ABI_ERROR("Error while gathering number of rows in linalg")
594 : : end if
595 : :
596 : 190862 : ncpuCols = xgTransposer%mpiData(MPI_COLS)%size
597 : 190862 : ncpuRows = xgTransposer%mpiData(MPI_ROWS)%size
598 : 190862 : icpu_rows = xgTransposer%mpiData(MPI_ROWS)%rank
599 : 190862 : xgTransposer%nrowsColsRows = 0
600 [ + + ]: 689316 : do icpu_cols=0,ncpuCols-1
601 : 689316 : xgTransposer%nrowsColsRows = xgTransposer%nrowsColsRows + xgTransposer%nrowsLinalg(1+icpu_rows+icpu_cols*ncpuRows)
602 : : end do
603 : 190862 : xgTransposer%ncolsColsRows = cols(xgTransposer%xgBlock_linalg)/ncpuCols
604 [ + + ]: 190862 : if (xgTransposer%custom_ncolsColsRows) then
605 : 16 : ncolsColsRowsMe = xgTransposer%ncolsColsRows_sub(xgTransposer%mpiData(MPI_LINALG)%rank+1)
606 : 16 : xgTransposer%ncolsColsRows = ncolsColsRowsMe ! overwrite
607 : : end if
608 : :
609 : : !write(*,*) "In linalg, # of real pairs:", xgTransposer%nrowsLinalg
610 : : !write(*,*) "In rows, # of real pairs for proc ", xgTransposer%mpiData(MPI_ROWS)%rank, ":", xgTransposer%nrowsColsRows
611 : : !write(*,*) "In cols, # of cols for proc ", xgTransposer%mpiData(MPI_COLS)%rank, ":", xgTransposer%ncolsColsRows
612 : :
613 : 190862 : end subroutine xgTransposer_computeDistribution
614 : : !!***
615 : :
616 : : !!****f* m_xgTransposer/xgTransposer_readDistribution
617 : : !!
618 : : !! NAME
619 : : !! xgTransposer_readDistribution
620 : : !!
621 : : !! FUNCTION
622 : : !! Read custom distribution from xgTransposer%nrowsLinalg(:)
623 : :
624 : 0 : subroutine xgTransposer_readDistribution(xgTransposer)
625 : :
626 : : ! Arguments
627 : : type(xgTransposer_t), intent(inout) :: xgTransposer
628 : : ! Variables
629 : : integer :: iproc,nRealPairs,tot_ncols,ncols,ierr
630 : : integer :: nprocs,commLinalg
631 : :
632 : : ! ***********************************************************
633 : :
634 : 0 : nprocs = xgTransposer%mpiData(MPI_LINALG)%size
635 : 0 : xgTransposer%nrowsColsRows = rows(xgTransposer%xgBlock_colsrows)
636 : 0 : ncols = cols(xgTransposer%xgBlock_colsrows)
637 : 0 : xgTransposer%ncolsColsRows = ncols ! for current MPI process
638 : :
639 [ # # ]: 0 : if(xgTransposer%custom_ncolsColsRows) then
640 [ # # # # : 0 : ABI_MALLOC(xgTransposer%ncolsColsRows_sub,(nprocs))
# # # # #
# # # ]
641 : :
642 : 0 : commLinalg = xgTransposer%mpiData(MPI_LINALG)%comm
643 : 0 : ncols = cols(xgTransposer%xgBlock_colsrows)
644 : 0 : call xmpi_allgather(ncols,xgTransposer%ncolsColsRows_sub,commLinalg,ierr)
645 [ # # ]: 0 : if ( ierr /= xmpi_success ) then
646 : 0 : ABI_ERROR("Error while gathering number of columns in colsrows")
647 : : end if
648 : 0 : call xmpi_sum(ncols,commLinalg,ierr)
649 [ # # ]: 0 : if ( ierr /= xmpi_success ) then
650 : 0 : ABI_ERROR("Error while summing number of columns in colsrows")
651 : : end if
652 : 0 : tot_ncols = ncols ! update sum
653 : : else
654 : 0 : tot_ncols = ncols*xgTransposer%mpiData(MPI_COLS)%size
655 : : end if
656 : :
657 [ # # ]: 0 : do iproc=1,size(xgTransposer%nrowsLinalg)
658 : 0 : nRealPairs = xgTransposer%nrowsLinalg(iproc)
659 [ # # ]: 0 : if (MOD(nRealPairs,xgTransposer%nspinor)/=0) then
660 : 0 : ABI_ERROR('nspinor should divide nRealPairs!')
661 : : end if
662 : : end do
663 : :
664 : 0 : xgTransposer%ncolsLinalg = tot_ncols
665 : :
666 : : !write(*,*) "In linalg, # of real pairs:", xgTransposer%nrowsLinalg
667 : : !write(*,*) "In linalg, # of cols:", xgTransposer%ncolsLinalg
668 : : !write(*,*) "In rows, # of real pairs for proc ", xgTransposer%mpiData(MPI_ROWS)%rank, ":", xgTransposer%nrowsColsRows
669 : : !write(*,*) "In cols, # of cols for proc ", xgTransposer%mpiData(MPI_COLS)%rank, ":", xgTransposer%ncolsColsRows
670 : :
671 : 0 : end subroutine xgTransposer_readDistribution
672 : : !!***
673 : :
674 : : !!****f* m_xgTransposer/xgTransposer_makeXgBlock
675 : : !!
676 : : !! NAME
677 : : !! xgTransposer_makeXgBlock
678 : :
679 : 190862 : subroutine xgTransposer_makeXgBlock(xgTransposer)
680 : :
681 : : type(xgTransposer_t), intent(inout) :: xgTransposer
682 : : integer :: ncolsColsRows, nrowsLinalgMe
683 : :
684 : 381724 : select case (xgTransposer%state)
685 : : case (STATE_LINALG)
686 : :
687 : 190862 : ncolsColsRows = xgTransposer%ncolsColsRows
688 : : !write(std_out,*) '@makeXgBlock dims', ncolsColsRows, xgTransposer%nrowsColsRows
689 : : !call flush_unit(std_out)
690 : :
691 : : ! Assume xgBlock_colsrows is empty and not constructed because user cannot
692 : : ! predict the size
693 [ - + ]: 190862 : if ( associated(xgTransposer%buffer) ) then
694 [ # # ]: 0 : if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
695 : : #if defined HAVE_GPU && defined HAVE_YAKL
696 : : ABI_FREE_MANAGED(xgTransposer%buffer)
697 : : #endif
698 : : else
699 : : if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
700 : : #ifdef HAVE_OPENMP_OFFLOAD
701 : : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
702 : : !$OMP TARGET EXIT DATA MAP(delete:xgTransposer%buffer)
703 : : #else
704 : : call xg_ompgpu_map_delete(xgTransposer%buffer)
705 : : #endif
706 : : #endif
707 : : end if
708 : 0 : ABI_FREE(xgTransposer%buffer)
709 : : end if
710 : : end if
711 : :
712 [ + + ]: 190862 : if ( xgTransposer%mpiData(MPI_COLS)%size == 1 ) then
713 : 37926 : xgTransposer%xgBlock_colsrows = xgTransposer%xgBlock_linalg
714 : : else
715 [ + - ]: 152936 : if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
716 : : #if defined HAVE_GPU && defined HAVE_YAKL
717 : : ABI_MALLOC_MANAGED(xgTransposer%buffer,(/2,ncolsColsRows*xgTransposer%nrowsColsRows/))
718 : : #endif
719 : : else
720 [ - + + - : 458808 : ABI_MALLOC(xgTransposer%buffer,(2,ncolsColsRows*xgTransposer%nrowsColsRows))
+ - + - -
+ - + ]
721 : : end if
722 : : if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
723 : : #ifdef HAVE_OPENMP_OFFLOAD
724 : : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
725 : : !$OMP TARGET ENTER DATA MAP(alloc:xgTransposer%buffer)
726 : : #else
727 : : call xg_ompgpu_map_alloc(xgTransposer%buffer)
728 : : #endif
729 : : #endif
730 : : end if
731 : : call xgBlock_map(xgTransposer%xgBlock_colsrows,xgTransposer%buffer,space(xgTransposer%xgBlock_linalg),&
732 : : xgTransposer%nrowsColsRows,ncolsColsRows,xgTransposer%mpiData(MPI_ROWS)%comm,&
733 : 152936 : me_g0=xgTransposer%me_g0_fft,gpu_option=xgTransposer%gpu_option)
734 : : end if
735 : : case (STATE_COLSROWS)
736 : :
737 : : ! Assume xgBlock_linalg is empty and not constructed
738 : :
739 : : !write(std_out,*) 'db using @makeXgBlock not empty'
740 : : !call flush_unit(std_out)
741 : :
742 : 0 : nrowsLinalgMe = xgTransposer%nrowsLinalg(xgTransposer%mpiData(MPI_LINALG)%rank+1)
743 : :
744 [ # # ]: 0 : if ( associated(xgTransposer%buffer) ) then
745 [ # # ]: 0 : if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
746 : : #if defined HAVE_GPU && defined HAVE_YAKL
747 : : ABI_FREE_MANAGED(xgTransposer%buffer)
748 : : #endif
749 : : else
750 : : if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
751 : : #ifdef HAVE_OPENMP_OFFLOAD
752 : : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
753 : : !$OMP TARGET EXIT DATA MAP(delete:xgTransposer%buffer)
754 : : #else
755 : : call xg_ompgpu_map_delete(xgTransposer%buffer)
756 : : #endif
757 : : #endif
758 : : end if
759 : 0 : ABI_FREE(xgTransposer%buffer)
760 : : end if
761 : : end if
762 : :
763 [ # # ]: 0 : if ( xgTransposer%mpiData(MPI_COLS)%size == 1 ) then
764 : 0 : xgTransposer%xgBlock_linalg = xgTransposer%xgBlock_colsrows
765 : : else
766 [ # # ]: 0 : if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
767 : : #if defined HAVE_GPU && defined HAVE_YAKL
768 : : ABI_MALLOC_MANAGED(xgTransposer%buffer,(/2,xgTransposer%ncolsLinalg*nrowsLinalgMe/))
769 : : #endif
770 : : else
771 [ # # # # : 0 : ABI_MALLOC(xgTransposer%buffer,(2,xgTransposer%ncolsLinalg*nrowsLinalgMe))
# # # # #
# # # ]
772 : : end if
773 : : if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
774 : : #ifdef HAVE_OPENMP_OFFLOAD
775 : : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
776 : : !$OMP TARGET ENTER DATA MAP(alloc:xgTransposer%buffer)
777 : : #else
778 : : call xg_ompgpu_map_alloc(xgTransposer%buffer)
779 : : #endif
780 : : #endif
781 : : end if
782 : : call xgBlock_map(xgTransposer%xgBlock_linalg,xgTransposer%buffer,space(xgTransposer%xgBlock_colsrows),&
783 : : nrowsLinalgMe,xgTransposer%ncolsLinalg,xgTransposer%mpiData(MPI_COLS)%comm,&
784 : 0 : me_g0=xgTransposer%me_g0_fft,gpu_option=xgTransposer%gpu_option)
785 : : end if
786 : :
787 : : case default
788 [ + - - ]: 190862 : ABI_ERROR("State unknown")
789 : : end select
790 : :
791 : : !if (associated(xgTransposer%buffer)) then
792 : : ! write(std_out,*) '@makeXgBlock alloc buffer size', size(xgTransposer%buffer,1), size(xgTransposer%buffer,2)
793 : : ! call flush_unit(std_out)
794 : : !end if
795 : :
796 : : #if defined HAVE_OPENMP_OFFLOAD && ! defined HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
797 : : contains
798 : : ! NOTE : routines to be used in order to avoid using derived-type syntax in OMP directives
799 : : subroutine xg_ompgpu_map_alloc(xgTransposerBuffer)
800 : : real(dp),intent(inout) :: xgTransposerBuffer(:,:)
801 : : !$OMP TARGET ENTER DATA MAP(alloc:xgTransposerBuffer)
802 : : end subroutine xg_ompgpu_map_alloc
803 : :
804 : : subroutine xg_ompgpu_map_delete(xgTransposerBuffer)
805 : : real(dp),intent(inout) :: xgTransposerBuffer(:,:)
806 : : !$OMP TARGET EXIT DATA MAP(delete:xgTransposerBuffer)
807 : : end subroutine xg_ompgpu_map_delete
808 : : #endif
809 : 190862 : end subroutine xgTransposer_makeXgBlock
810 : : !!***
811 : :
812 : : !!****f* m_xgTransposer/xgTransposer_transpose
813 : : !!
814 : : !! NAME
815 : : !! xgTransposer_transpose
816 : :
817 : 2806924 : subroutine xgTransposer_transpose(xgTransposer,toState)
818 : :
819 : : type(xgTransposer_t), intent(inout) :: xgTransposer
820 : : integer , intent(in ) :: toState
821 : : double precision :: tsec(2)
822 : :
823 : 2806924 : call timab(tim_transpose,1,tsec)
824 : :
825 [ - + ]: 2806924 : if ( toState /= STATE_LINALG .and. toState /= STATE_COLSROWS ) then
826 : 0 : ABI_ERROR("Bad value for toState")
827 : : end if
828 : :
829 : : !write(std_out,*) "linalg", rows(xgTransposer%xgBlock_linalg)*cols(xgTransposer%xgBlock_linalg)
830 : : !write(std_out,*) "colsrows", rows(xgTransposer%xgBlock_colsrows)*cols(xgTransposer%xgBlock_colsrows)
831 : 2076904 : select case (toState)
832 : : case (STATE_LINALG)
833 [ - + ]: 2076904 : if ( xgTransposer%state == STATE_LINALG ) then
834 : 0 : ABI_WARNING("Array linalg has already been transposed")
835 : : end if
836 [ + + ]: 2076904 : if ( xgTransposer%mpiData(MPI_COLS)%size > 1 ) then
837 : 1215010 : call xgTransposer_toLinalg(xgTransposer)
838 : : else
839 : 861894 : xgTransposer%state = STATE_LINALG
840 : : end if
841 : : case (STATE_COLSROWS)
842 [ - + ]: 730020 : if ( xgTransposer%state == STATE_COLSROWS ) then
843 : 0 : ABI_WARNING("Array colsrows has already been transposed")
844 : : end if
845 [ + + - + : 3536944 : if ( xgTransposer%mpiData(MPI_COLS)%size > 1 ) then
+ ]
846 : 441098 : call xgTransposer_toColsRows(xgTransposer)
847 : : else
848 : 288922 : xgTransposer%state = STATE_COLSROWS
849 : : end if
850 : : end select
851 : :
852 : 2806924 : call timab(tim_transpose,2,tsec)
853 : :
854 : 2806924 : end subroutine xgTransposer_transpose
855 : : !!***
856 : :
857 : : !!****f* m_xgTransposer/xgTransposer_toLinalg
858 : : !!
859 : : !! NAME
860 : : !! xgTransposer_toLinalg
861 : :
862 : 1215010 : subroutine xgTransposer_toLinalg(xgTransposer)
863 : :
864 : : type(xgTransposer_t), intent(inout) :: xgTransposer
865 : 1215010 : double precision, allocatable, target :: sendbuf(:,:)
866 : 1215010 : double precision, pointer :: recvbuf(:,:)
867 : : !double precision, pointer :: buffer(:,:)
868 : 1215010 : integer, allocatable :: sendcounts(:), recvcounts(:)
869 : 1215010 : integer, allocatable :: sdispls(:), rdispls(:)
870 : : integer :: ncpu_cols, ncpu_rows, comm, me_rows, me_cols
871 : : integer :: nrowsColsRows
872 : : integer :: ncolsColsRows
873 : : integer :: nrowsLinalgMe
874 : : integer :: icpu, ierr
875 : : integer :: send_start,send_end
876 : : !integer :: myrequest
877 : : !integer, allocatable :: request(:), status(:)
878 : 1215010 : type(ptr_t), allocatable :: sendptrbuf(:)
879 : 1215010 : integer, pointer :: nrowsLinalg(:)
880 : 1215010 : integer, pointer :: ncolsColsRows_sub(:)
881 : : double precision :: tsec(2)
882 : :
883 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
884 : : double precision, allocatable :: recvbuf_mpi(:,:)
885 : : integer(c_size_t) :: buffer_size
886 : : #endif
887 : :
888 : 1215010 : call timab(tim_toLinalg,1,tsec)
889 : :
890 : 1215010 : ncpu_cols = xgTransposer%mpiData(MPI_COLS)%size
891 : 1215010 : ncpu_rows = xgTransposer%mpiData(MPI_ROWS)%size
892 : 1215010 : comm = xgTransposer%mpiData(MPI_COLS)%comm
893 : : ! me = xgTransposer%mpiData(MPI_ROWS)%rank*ncpu
894 : 1215010 : me_rows = xgTransposer%mpiData(MPI_ROWS)%rank
895 : :
896 : : !write(std_out,'(a,i4,i4,i4)') '@xgTransposer_toLinalg for MPI proc in comm=', &
897 : : ! xgTransposer%mpiData(MPI_LINALG)%rank, xgTransposer%mpiData(MPI_LINALG)%comm, comm
898 : :
899 : 1215010 : nrowsColsRows = xgTransposer%nrowsColsRows
900 : 1215010 : ncolsColsRows = xgTransposer%ncolsColsRows
901 : :
902 : 1215010 : nrowsLinalg => xgTransposer%nrowsLinalg
903 [ - + ]: 1215010 : if (xgTransposer%custom_ncolsColsRows) ncolsColsRows_sub => xgTransposer%ncolsColsRows_sub
904 : 1215010 : nrowsLinalgMe = nrowsLinalg(xgTransposer%mpiData(MPI_LINALG)%rank+1)
905 : :
906 : : !write(std_out,*) '@toLinalg nrowsColsRows, ncolsColsRows=', nrowsColsRows, ncolsColsRows
907 : : !call flush_unit(std_out)
908 : :
909 [ - + + - : 3645030 : ABI_MALLOC(sendbuf,(2,nrowsColsRows*ncolsColsRows))
+ - + - -
+ - + ]
910 : : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
911 : : !$OMP TARGET ENTER DATA MAP(alloc:sendbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
912 : : #endif
913 : 1215010 : call xgTransposer_reorganizeData(xgTransposer,sendbuf)
914 : :
915 [ - + + - : 3645030 : ABI_MALLOC(recvcounts,(ncpu_cols))
+ - + - -
+ - + ]
916 [ - + + - : 2430020 : ABI_MALLOC(rdispls,(ncpu_cols))
- + - + ]
917 [ - + ]: 1215010 : if (xgTransposer%custom_ncolsColsRows) then ! non uniform sendcounts
918 [ # # ]: 0 : recvcounts(:) = 2*nrowsLinalgMe*ncolsColsRows_sub(:)
919 : : else
920 [ + + ]: 4171278 : recvcounts(:) = 2*nrowsLinalgMe*ncolsColsRows !! Thank you fortran for not starting at 0 !
921 : : end if
922 : 1215010 : rdispls(1) = 0
923 [ + + ]: 2956268 : do icpu = 2, ncpu_cols
924 : 2956268 : rdispls(icpu) = rdispls(icpu-1)+recvcounts(icpu-1)
925 : : end do
926 : :
927 : : call xgBlock_reverseMap(xgTransposer%xgBlock_linalg,recvbuf,&
928 : 1215010 : & rows=1,cols=cols(xgTransposer%xgBlock_linalg)*nrowsLinalgMe)
929 : :
930 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
931 : : ! just for debug
932 : : if( xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
933 : : !call gpu_managed_ptr_status(C_LOC(recvbuf))
934 : : buffer_size = size(recvbuf) * dp
935 : : call gpu_data_prefetch_async(C_LOC(recvbuf), buffer_size, CPU_DEVICE_ID)
936 : : call gpu_device_synchronize()
937 : :
938 : : ABI_MALLOC(recvbuf_mpi, (size(recvbuf,1), size(recvbuf,2)) )
939 : :
940 : : end if
941 : : #endif
942 : :
943 [ - + + - : 2430020 : ABI_MALLOC(sendcounts,(ncpu_cols))
- + - + ]
944 [ - + + - : 2430020 : ABI_MALLOC(sdispls,(ncpu_cols))
- + - + ]
945 : : !sendcounts(:) = 2*nrowsLinalg(me+1:me+ncpu)*ncolsColsRows
946 [ + + ]: 4171278 : do icpu=0,ncpu_cols-1
947 : 4171278 : sendcounts(icpu+1) = 2*nrowsLinalg(me_rows+1+icpu*ncpu_rows)*ncolsColsRows
948 : : !write(*,*) "nrowsLinalg ", icpu, ncpu_rows, nrowsLinalg(me_rows+1+icpu*ncpu_rows)
949 : : end do
950 : 1215010 : sdispls(1) = 0
951 [ + + ]: 2956268 : do icpu = 2, ncpu_cols
952 : 2956268 : sdispls(icpu) = sdispls(icpu-1)+sendcounts(icpu-1)
953 : : end do
954 : :
955 : 2430020 : select case(xgTransposer%mpiAlgo)
956 : : case (TRANS_ALL2ALL)
957 : : !ABI_MALLOC(request,(1))
958 : : !myrequest = 1
959 : :
960 : : ABI_NVTX_START_RANGE(NVTX_TRANSPOSER_MPI_ALL2ALL)
961 : :
962 [ + - ]: 1215010 : if( xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
963 : :
964 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
965 : : call timab(tim_all2allv,1,tsec)
966 : : call xmpi_alltoallv(sendbuf, sendcounts, sdispls, &
967 : : recvbuf_mpi, recvcounts, rdispls, &
968 : : comm, ierr)
969 : : call timab(tim_all2allv,2,tsec)
970 : : !call xmpi_ialltoallv(sendbuf, sendcounts, sdispls, &
971 : : ! recvbuf, recvcounts, rdispls, &
972 : : ! comm, request(myrequest))
973 : :
974 : : ! copy back recvbuf_mpi into recvbuf
975 : : recvbuf(:,:) = recvbuf_mpi(:,:)
976 : :
977 : : ABI_FREE(recvbuf_mpi)
978 : : #endif
979 : :
980 : : else
981 : :
982 : : !write(*,*) "sendcounts for proc ", me_rows, xmpi_comm_rank(comm), ":", sendcounts(:)
983 : : !write(*,*) "recvcounts for proc ", me_rows, xmpi_comm_rank(comm), ":", recvcounts(:)
984 : :
985 : 1215010 : call timab(tim_all2allv,1,tsec)
986 : : call xmpi_alltoallv(sendbuf, sendcounts, sdispls, &
987 : : recvbuf, recvcounts, rdispls, &
988 : 1215010 : comm, ierr, use_omp_map=(xgTransposer%gpu_option==ABI_GPU_OPENMP))
989 : 1215010 : call timab(tim_all2allv,2,tsec)
990 : : !call xmpi_ialltoallv(sendbuf, sendcounts, sdispls, &
991 : : ! recvbuf, recvcounts, rdispls, &
992 : : ! comm, request(myrequest))
993 : :
994 : : end if
995 : :
996 : : ABI_NVTX_END_RANGE()
997 : :
998 : : case (TRANS_GATHER)
999 : :
1000 : : !ABI_MALLOC(request,(ncpu))
1001 : 0 : me_cols = xgTransposer%mpiData(MPI_COLS)%rank
1002 : : !myrequest = me+1
1003 : :
1004 : : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1005 : : !$OMP TARGET UPDATE FROM(sendbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1006 : : #endif
1007 [ # # # # : 0 : ABI_MALLOC(sendptrbuf,(1:ncpu_cols))
# # # # #
# # # #
# ]
1008 [ # # ]: 0 : do icpu = 1, ncpu_cols
1009 : 0 : send_start = sdispls(icpu)/2+1
1010 : 0 : send_end = sdispls(icpu)/2+sendcounts(icpu)/2
1011 : 0 : sendptrbuf(icpu)%ptr => sendbuf(:,send_start:send_end)
1012 : 0 : call timab(tim_gatherv,1,tsec)
1013 : 0 : call xmpi_gatherv(sendptrbuf(icpu)%ptr,sendcounts(icpu),recvbuf,recvcounts,rdispls,icpu-1,comm,ierr)
1014 : 0 : call timab(tim_gatherv,2,tsec)
1015 : : !call mpi_igatherv(sendptrbuf(me+1)%ptr,sendcounts(icpu),MPI_DOUBLE_PRECISION,&
1016 : : ! recvbuf,recvcounts,rdispls,MPI_DOUBLE_PRECISION,icpu-1,comm,request(icpu),ierr)
1017 : : end do
1018 : : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1019 : : !$OMP TARGET UPDATE TO(recvbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1020 : : #endif
1021 : :
1022 : : case default
1023 [ + - - ]: 1215010 : ABI_BUG("This algo does not exist")
1024 : : end select
1025 : :
1026 : 1215010 : xgTransposer%state = STATE_LINALG
1027 : :
1028 : : !ABI_MALLOC(status,(MPI_STATUS_SIZE))
1029 : : !call mpi_wait(request(myrequest),status,ierr)
1030 [ - + ]: 1215010 : if ( ierr /= xmpi_success ) then
1031 : 0 : ABI_ERROR("Error while waiting for mpi")
1032 : : end if
1033 : :
1034 [ + - ]: 1215010 : if ( allocated(sendcounts) ) then
1035 : 1215010 : ABI_FREE(sendcounts)
1036 : : end if
1037 : : if ( allocated(sdispls) ) then
1038 : 1215010 : ABI_FREE(sdispls)
1039 : : end if
1040 : :
1041 [ - + ]: 1215010 : ABI_FREE(recvcounts)
1042 : 1215010 : ABI_FREE(rdispls)
1043 : :
1044 : :
1045 [ - + ]: 1215010 : if ( allocated(sendptrbuf) ) then
1046 : 0 : ABI_FREE(sendptrbuf)
1047 : : end if
1048 : :
1049 : : !do icpu = 1, size(request)
1050 : : ! if ( icpu /= myrequest ) then
1051 : : ! call mpi_wait(request(icpu),status,ierr)
1052 : : ! if ( ierr /= MPI_SUCCESS ) then
1053 : : ! ABI_ERROR("Error while waiting for other mpi")
1054 : : ! end if
1055 : : ! end if
1056 : : !end do
1057 : : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1058 : : !$OMP TARGET EXIT DATA MAP(delete:sendbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1059 : : #endif
1060 [ - + ]: 1215010 : ABI_FREE(sendbuf)
1061 : : !ABI_FREE(status)
1062 : : !ABI_FREE(request)
1063 : :
1064 : 1215010 : call timab(tim_toLinalg,2,tsec)
1065 : :
1066 [ - + - + ]: 2430020 : end subroutine xgTransposer_toLinalg
1067 : : !!***
1068 : :
1069 : : !!****f* m_xgTransposer/xgTransposer_toColsRows
1070 : : !!
1071 : : !! NAME
1072 : : !! xgTransposer_toColsRows
1073 : :
1074 : 441098 : subroutine xgTransposer_toColsRows(xgTransposer)
1075 : :
1076 : : type(xgTransposer_t), intent(inout) :: xgTransposer
1077 : 441098 : double precision, pointer :: sendbuf(:,:)
1078 : 441098 : double precision, allocatable :: recvbuf(:,:)
1079 : : !double precision, allocatable :: buffer(:,:)
1080 : 441098 : integer, allocatable :: sendcounts(:), recvcounts(:)
1081 : 441098 : integer, allocatable :: sdispls(:), rdispls(:)
1082 : : integer :: ncpu_rows, ncpu_cols, comm, me_rows, me_cols
1083 : : integer :: nrowsColsRows
1084 : : integer :: nrowsLinalgMe
1085 : : integer :: ncolsColsRows
1086 : : integer :: icpu,ierr
1087 : : !integer :: myrequest
1088 : : !integer, allocatable :: request(:), status(:)
1089 : : type(xgBlock_t) :: xgBlock_toTransposed
1090 : 441098 : type(ptr_t), allocatable :: sendptrbuf(:)
1091 : 441098 : integer, pointer :: nrowsLinalg(:)
1092 : 441098 : integer, pointer :: ncolsColsRows_sub(:)
1093 : : double precision :: tsec(2)
1094 : :
1095 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
1096 : : double precision, allocatable :: sendbuf_mpi(:,:)
1097 : : integer(c_size_t) :: buffer_size
1098 : : #endif
1099 : :
1100 : 441098 : call timab(tim_toColsRows,1,tsec)
1101 : :
1102 : 441098 : ncpu_cols = xgTransposer%mpiData(MPI_COLS)%size
1103 : 441098 : ncpu_rows = xgTransposer%mpiData(MPI_ROWS)%size
1104 : 441098 : comm = xgTransposer%mpiData(MPI_COLS)%comm
1105 : 441098 : me_rows = xgTransposer%mpiData(MPI_ROWS)%rank
1106 : :
1107 : 441098 : nrowsColsRows = xgTransposer%nrowsColsRows
1108 : 441098 : ncolsColsRows = xgTransposer%ncolsColsRows
1109 : :
1110 : 441098 : nrowsLinalg => xgTransposer%nrowsLinalg
1111 [ + + ]: 441098 : if (xgTransposer%custom_ncolsColsRows) ncolsColsRows_sub => xgTransposer%ncolsColsRows_sub
1112 : 441098 : nrowsLinalgMe = nrowsLinalg(xgTransposer%mpiData(MPI_LINALG)%rank+1)
1113 : :
1114 [ - + + - : 1323294 : ABI_MALLOC(recvbuf,(2,nrowsColsRows*ncolsColsRows))
+ - + - -
+ - + ]
1115 : : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1116 : : !$OMP TARGET ENTER DATA MAP(alloc:recvbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1117 : : #endif
1118 [ - + + - : 1323294 : ABI_MALLOC(recvcounts,(ncpu_cols))
+ - + - -
+ - + ]
1119 [ - + + - : 882196 : ABI_MALLOC(rdispls,(ncpu_cols))
- + - + ]
1120 : :
1121 : : ! 2 * ncolsColsRows * nrowsLinalg_sub(:)
1122 [ + + ]: 1564118 : recvcounts(:) = 2*ncolsColsRows*nrowsLinalg(1+me_rows:1+me_rows+(ncpu_cols-1)*ncpu_rows:ncpu_rows)
1123 : 441098 : rdispls(1) = 0
1124 [ + + ]: 1123020 : do icpu = 2, ncpu_cols
1125 : 1123020 : rdispls(icpu) = rdispls(icpu-1)+recvcounts(icpu-1)
1126 : : end do
1127 : : !recvcounts(:) = 2*nrowsLinalg(me+1:me+ncpu)*ncolsColsRows
1128 : : !rdispls(1) = 0
1129 : : !do icpu = 2, ncpu
1130 : : ! rdispls(icpu) = rdispls(icpu-1)+recvcounts(icpu-1)
1131 : : !end do
1132 : :
1133 : 882196 : select case(xgTransposer%mpiAlgo)
1134 : : case (TRANS_ALL2ALL)
1135 [ - + + - : 882196 : ABI_MALLOC(sendcounts,(ncpu_cols))
- + - + ]
1136 [ - + + - : 882196 : ABI_MALLOC(sdispls,(ncpu_cols))
- + - + ]
1137 : : !ABI_MALLOC(request,(1))
1138 : : !myrequest = 1
1139 : :
1140 : : ! 2 * ncolsColsRows_sub(:) * nrowsLinalgMe
1141 [ + + ]: 441098 : if (xgTransposer%custom_ncolsColsRows) then ! non uniform sendcounts
1142 [ + + ]: 80 : sendcounts(:) = 2*nrowsLinalgMe*ncolsColsRows_sub(:)
1143 : : else
1144 [ + + ]: 1564038 : sendcounts(:) = 2*nrowsLinalgMe*ncolsColsRows !! Thank you fortran for not starting at 0 !
1145 : : end if
1146 : 441098 : sdispls(1) = 0
1147 [ + + ]: 1123020 : do icpu = 2, ncpu_cols
1148 : 1123020 : sdispls(icpu) = sdispls(icpu-1)+sendcounts(icpu-1)
1149 : : end do
1150 : :
1151 : : call xgBlock_reverseMap(xgTransposer%xgBlock_linalg,sendbuf, &
1152 : 441098 : & rows=1,cols=cols(xgTransposer%xgBlock_linalg)*nrowsLinalgMe)
1153 : : !write(*,*) "Before ialltoall"
1154 : :
1155 : : ABI_NVTX_START_RANGE(NVTX_TRANSPOSER_MPI_ALL2ALL)
1156 : :
1157 : : ! if gpu is enabled, data are located in GPU memory, so we copy them on a host buffer
1158 [ + - ]: 441098 : if( xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
1159 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
1160 : : ABI_MALLOC(sendbuf_mpi, (size(sendbuf,1), size(sendbuf,2)) )
1161 : :
1162 : : ! sync sendbuf on host and then copy to sendbuf_mpi
1163 : : buffer_size = size(sendbuf) * dp
1164 : : call gpu_data_prefetch_async(C_LOC(sendbuf), buffer_size, CPU_DEVICE_ID)
1165 : : call gpu_device_synchronize()
1166 : :
1167 : : sendbuf_mpi(:,:) = sendbuf(:,:)
1168 : :
1169 : : call timab(tim_all2allv,1,tsec)
1170 : : ! this is a cpu mpi comm
1171 : : call xmpi_alltoallv(sendbuf_mpi, sendcounts, sdispls, &
1172 : : recvbuf, recvcounts, rdispls, &
1173 : : comm, ierr)
1174 : : call timab(tim_all2allv,2,tsec)
1175 : : !call xmpi_ialltoallv(sendbuf, sendcounts, sdispls, &
1176 : : ! recvbuf, recvcounts, rdispls, &
1177 : : ! comm, request(myrequest))
1178 : : !write(*,*) "After ialltoall"
1179 : :
1180 : : ABI_FREE(sendbuf_mpi)
1181 : : #endif
1182 : : else
1183 : :
1184 : : !write(std_out,'(a,i4,i4,i4)') '@xgTransposer_toColsRows for MPI proc in comm=', &
1185 : : ! xgTransposer%mpiData(MPI_LINALG)%rank, xgTransposer%mpiData(MPI_LINALG)%comm, comm
1186 : : !if (xgTransposer%custom_ncolsColsRows) then
1187 : : ! write(std_out,*) 'db ncolsColsRows_sub', ncolsColsRows_sub(:)
1188 : : !end if
1189 : : !write(std_out,*) 'db ncolsColsMe', ncolsColsRows
1190 : : !write(std_out,*) 'db nrowsLinalg', nrowsLinalg(1+me_rows:1+me_rows+(ncpu_cols-1)*ncpu_rows:ncpu_rows)
1191 : : !write(std_out,*) 'db nrowsLinalgMe', nrowsLinalgMe
1192 : : !write(std_out,*) 'db sendbuf', cols(xgTransposer%xgBlock_linalg), nrowsLinalgMe
1193 : : !write(std_out,*) 'db recvbuf', nrowsColsRows, ncolsColsRows
1194 : : !write(std_out,*) 'db sendcounts', sendcounts(:)
1195 : : !write(std_out,*) 'db recvcounts', recvcounts(:)
1196 : : !write(std_out,*) 'db sdispls', sdispls(:)
1197 : : !write(std_out,*) 'db rdispls', rdispls(:)
1198 : : !call flush_unit(std_out)
1199 : :
1200 : 441098 : call timab(tim_all2allv,1,tsec)
1201 : : call xmpi_alltoallv(sendbuf, sendcounts, sdispls, &
1202 : : recvbuf, recvcounts, rdispls, &
1203 : 441098 : comm, ierr, use_omp_map=(xgTransposer%gpu_option == ABI_GPU_OPENMP))
1204 : 441098 : call timab(tim_all2allv,2,tsec)
1205 : : !call xmpi_ialltoallv(sendbuf, sendcounts, sdispls, &
1206 : : ! recvbuf, recvcounts, rdispls, &
1207 : : ! comm, request(myrequest))
1208 : : !write(*,*) "After ialltoall"
1209 : : end if
1210 : :
1211 : : ABI_NVTX_END_RANGE()
1212 : :
1213 : : case (TRANS_GATHER)
1214 : :
1215 : : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1216 : : !$OMP TARGET UPDATE FROM(sendbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1217 : : #endif
1218 : : !ABI_MALLOC(request,(ncpu))
1219 : 0 : me_cols = xgTransposer%mpiData(MPI_COLS)%rank
1220 : : !myrequest = me+1
1221 : :
1222 [ # # # # : 0 : ABI_MALLOC(sendptrbuf,(1:ncpu_cols))
# # # # #
# # # #
# ]
1223 : : !call flush(6)
1224 : : !call xmpi_barrier(xgTransposer%mpiData(MPI_LINALG)%comm)
1225 [ # # ]: 0 : do icpu = 0, ncpu_cols-1
1226 : : !write(*,*) me, "->", icpu, "from col ",icpu*ncolsColsRows+1, " number of rows:", nrowsLinalgMe
1227 : 0 : call xgBlock_setBlock(xgTransposer%xgBlock_linalg,xgBlock_toTransposed,nrowsLinalgMe,ncolsColsRows,fcol=icpu*ncolsColsRows+1)
1228 : : call xgBlock_reverseMap(xgBlock_toTransposed,sendptrbuf(icpu+1)%ptr,&
1229 : 0 : rows=1,cols=ncolsColsRows*nrowsLinalgMe)
1230 : 0 : call timab(tim_gatherv,1,tsec)
1231 : 0 : call xmpi_gatherv(sendptrbuf(icpu+1)%ptr,2*ncolsColsRows*nrowsLinalgMe,recvbuf,recvcounts,rdispls,icpu,comm,ierr)
1232 : 0 : call timab(tim_gatherv,2,tsec)
1233 : : !call mpi_igatherv(sendptrbuf(me+1)%ptr,2*ncolsColsRows*nrowsLinalgMe,MPI_DOUBLE_PRECISION,&
1234 : : ! recvbuf,recvcounts,rdispls,MPI_DOUBLE_PRECISION,icpu,comm,request(icpu+1),ierr)
1235 : : end do
1236 : : !call xmpi_barrier(xgTransposer%mpiData(MPI_LINALG)%comm)
1237 : : !call flush(6)
1238 : : !write(*,*) me, request
1239 : : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1240 : : !$OMP TARGET UPDATE TO(recvbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1241 : : #endif
1242 : :
1243 : : case default
1244 [ + - - ]: 441098 : ABI_BUG("This algo does not exist")
1245 : : end select
1246 : :
1247 : : !ABI_MALLOC(status,(MPI_STATUS_SIZE))
1248 : : !call mpi_wait(request(myrequest),status,ierr)
1249 : : !write(*,*) "Request ended"
1250 [ - + ]: 441098 : if ( ierr /= xmpi_success ) then
1251 : 0 : ABI_ERROR("Error while waiting for mpi")
1252 : : end if
1253 : : !write(*,*) "with success"
1254 : :
1255 : 441098 : call xgTransposer_reorganizeData(xgTransposer,recvbuf)
1256 : :
1257 [ + - ]: 441098 : if ( allocated(sendcounts) ) then
1258 : 441098 : ABI_FREE(sendcounts)
1259 : : end if
1260 [ + - ]: 441098 : if ( allocated(sdispls) ) then
1261 : 441098 : ABI_FREE(sdispls)
1262 : : end if
1263 : :
1264 [ - + ]: 441098 : ABI_FREE(recvcounts)
1265 : 441098 : ABI_FREE(rdispls)
1266 : :
1267 : : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1268 : : !$OMP TARGET EXIT DATA MAP(delete:recvbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1269 : : #endif
1270 [ - + ]: 441098 : ABI_FREE(recvbuf)
1271 : :
1272 [ - + ]: 441098 : if ( allocated(sendptrbuf) ) then
1273 : 0 : ABI_FREE(sendptrbuf)
1274 : : end if
1275 : :
1276 : 441098 : xgTransposer%state = STATE_COLSROWS
1277 : :
1278 : : !do icpu = 1, size(request)
1279 : : ! if ( icpu /= myrequest ) then
1280 : : ! call mpi_wait(request(icpu),status,ierr)
1281 : : ! if ( ierr /= MPI_SUCCESS ) then
1282 : : ! ABI_ERROR("Error while waiting for other mpi")
1283 : : ! end if
1284 : : ! end if
1285 : : !end do
1286 : : !ABI_FREE(status)
1287 : : !ABI_FREE(request)
1288 : :
1289 : 441098 : call timab(tim_toColsRows,2,tsec)
1290 : :
1291 [ - + - + ]: 882196 : end subroutine xgTransposer_toColsRows
1292 : : !!***
1293 : :
1294 : : !!****f* m_xgTransposer/xgTransposer_reorganizeData
1295 : : !!
1296 : : !! NAME
1297 : : !! xgTransposer_reorganizeData
1298 : :
1299 [ - + ]: 1656108 : subroutine xgTransposer_reorganizeData(xgTransposer,bufferMess)
1300 : :
1301 : : type(xgTransposer_t), intent(inout) :: xgTransposer
1302 : : double precision , intent(inout) :: bufferMess(:,:)
1303 : : double precision, pointer :: bufferOrdered(:,:) => null()
1304 : : integer :: nrowsColsRows,nthreads_bak
1305 : : integer :: ncolsColsRows
1306 : : integer :: tos,toe,froms,frome
1307 : : integer :: col, icpu
1308 : : integer :: me_rows,ncpu_cols,ncpu_rows
1309 : : integer :: nPair,ispinor,nspinor
1310 : : integer :: nrowsLinalgMe,nrowsLinalgMeSum
1311 : : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1312 : : integer :: irow
1313 : : integer, contiguous, pointer :: nrowsLinalg(:)
1314 : : #else
1315 : 1656108 : integer, pointer :: nrowsLinalg(:)
1316 : : #endif
1317 : : double precision :: tsec(2)
1318 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
1319 : : integer(c_size_t) :: buffer_size
1320 : : #endif
1321 : :
1322 : 1656108 : call timab(tim_reorganize,1,tsec)
1323 : :
1324 : 1656108 : me_rows = xgTransposer%mpiData(MPI_ROWS)%rank!*xgTransposer%mpiData(MPI_COLS)%size
1325 : 1656108 : ncpu_rows = xgTransposer%mpiData(MPI_ROWS)%size
1326 : 1656108 : ncpu_cols = xgTransposer%mpiData(MPI_COLS)%size
1327 : :
1328 : : ! These arrays are stored as
1329 : : ! 1 2 3
1330 : : ! 4 5 6
1331 : : ! 7 8 9
1332 : : !
1333 : : ! stored as
1334 : : !
1335 : : ! column-major, one-base indexing : 1 4 7 2 5 8 3 6 9 | y + N_y*(x-1)
1336 : : ! row-major , zero-based indexing: 1 2 3 4 5 6 7 8 9 | x + N_x*y
1337 : 1656108 : nspinor = xgTransposer%nspinor
1338 : 1656108 : nrowsLinalg => xgTransposer%nrowsLinalg
1339 : :
1340 : 1656108 : nrowsColsRows = xgTransposer%nrowsColsRows
1341 : 1656108 : ncolsColsRows = xgTransposer%ncolsColsRows
1342 : :
1343 : : ! number of pairs depends on current MPI process
1344 : 1656108 : nPair = nrowsColsRows*ncolsColsRows
1345 : 1656108 : call xgBlock_reverseMap(xgTransposer%xgBlock_colsrows,bufferOrdered,rows=1,cols=nPair)
1346 : :
1347 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
1348 : : ! if gpu is enabled, data are located in GPU memory, so we prefetch them on host
1349 : : ! to do the following reorganization.
1350 : : ! Alternatively, we should provide a GPU implementation of this data layout reorganization
1351 : : if( xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
1352 : : buffer_size = size(bufferOrdered) * dp
1353 : : call gpu_data_prefetch_async(C_LOC(bufferOrdered), buffer_size, CPU_DEVICE_ID)
1354 : : call gpu_device_synchronize()
1355 : : end if
1356 : :
1357 : : ! if gpu enabled increase locally OpenMP num threads
1358 : : if (xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
1359 : : call xomp_set_num_threads(xgTransposer%gpu_kokkos_nthrd)
1360 : : end if
1361 : : #endif
1362 : :
1363 [ - + - - ]: 1656108 : if (xgTransposer%gpu_option /= ABI_GPU_DISABLED .and. xgTransposer%gpu_thread_limit /= 0) then
1364 : 0 : nthreads_bak=xomp_get_max_threads()
1365 : 0 : call xomp_set_num_threads(min(xgTransposer%gpu_thread_limit,nthreads_bak))
1366 : : end if
1367 : :
1368 : 1656108 : select case (xgTransposer%state)
1369 : : case (STATE_LINALG)
1370 : : ! We are going to STATE_COLSROWS so we are after all2all
1371 : : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1372 : : if (xgTransposer%gpu_option == ABI_GPU_OPENMP) then
1373 : : ! bufferMess and bufferOrdered already live on the device: reorganize them there directly,
1374 : : ! no CPU<->GPU round trip needed. nrowsLinalg is small and mapped just for this kernel.
1375 : : ! Array-section assignment is not usable in device kernel bodies: loop explicitly instead.
1376 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1377 : : !$OMP& PRIVATE(nrowsLinalgMe,nrowsLinalgMeSum,toe,tos,frome,froms,irow) &
1378 : : !$OMP& MAP(to:bufferMess) MAP(to:bufferOrdered) MAP(to:nrowsLinalg)
1379 : : do col = 1, ncolsColsRows
1380 : : do icpu = 0, ncpu_cols-1
1381 : : do ispinor = 1, nspinor
1382 : : nrowsLinalgMe = nrowsLinalg(1+me_rows+icpu*ncpu_rows)
1383 : : nrowsLinalgMeSum = sum(nrowsLinalg(1+me_rows:1+me_rows+(icpu-1)*ncpu_rows:ncpu_rows))
1384 : : froms=1+(ispinor-1)*nrowsLinalgMe/nspinor+(col-1)*nrowsLinalgMe+nrowsLinalgMeSum*ncolsColsRows
1385 : : frome=froms-1+nrowsLinalgMe/nspinor
1386 : : tos=1+nrowsLinalgMeSum/nspinor+(ispinor-1)*nrowsColsRows/nspinor+(col-1)*nrowsColsRows
1387 : : toe=tos-1+nrowsLinalgMe/nspinor
1388 : : do irow = 0, toe-tos
1389 : : bufferOrdered(1,tos+irow) = bufferMess(1,froms+irow)
1390 : : bufferOrdered(2,tos+irow) = bufferMess(2,froms+irow)
1391 : : end do
1392 : : end do
1393 : : end do
1394 : : end do
1395 : : else
1396 : : #endif
1397 : : !$omp parallel do private(nrowsLinalgMe,nrowsLinalgMeSum,toe,tos,frome,froms), collapse(3)
1398 [ + + ]: 1103196 : do col = 1, ncolsColsRows
1399 [ + + ]: 2805916 : do icpu = 0, ncpu_cols-1
1400 [ + + ]: 4283186 : do ispinor = 1, nspinor
1401 : 1918368 : nrowsLinalgMe = nrowsLinalg(1+me_rows+icpu*ncpu_rows)
1402 [ + + ]: 3847448 : nrowsLinalgMeSum = sum(nrowsLinalg(1+me_rows:1+me_rows+(icpu-1)*ncpu_rows:ncpu_rows))
1403 : 1918368 : froms=1+(ispinor-1)*nrowsLinalgMe/nspinor+(col-1)*nrowsLinalgMe+nrowsLinalgMeSum*ncolsColsRows
1404 : 1918368 : frome=froms-1+nrowsLinalgMe/nspinor
1405 : 1918368 : tos=1+nrowsLinalgMeSum/nspinor+(ispinor-1)*nrowsColsRows/nspinor+(col-1)*nrowsColsRows
1406 : 1918368 : toe=tos-1+nrowsLinalgMe/nspinor
1407 [ + + + + ]: 557592892 : bufferOrdered(:,tos:toe) = bufferMess(:,froms:frome)
1408 : : end do
1409 : : end do
1410 : : end do
1411 : : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1412 : : end if
1413 : : #endif
1414 : : case (STATE_COLSROWS)
1415 : : ! We are going to STATE_LINALG so we are before all2all
1416 : : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1417 : : if (xgTransposer%gpu_option == ABI_GPU_OPENMP) then
1418 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1419 : : !$OMP& PRIVATE(nrowsLinalgMe,nrowsLinalgMeSum,toe,tos,frome,froms,irow) &
1420 : : !$OMP& MAP(to:bufferMess) MAP(to:bufferOrdered) MAP(to:nrowsLinalg)
1421 : : do col = 1, ncolsColsRows
1422 : : do icpu = 0, ncpu_cols-1
1423 : : do ispinor = 1, nspinor
1424 : : nrowsLinalgMe = nrowsLinalg(1+me_rows+icpu*ncpu_rows)
1425 : : nrowsLinalgMeSum = sum(nrowsLinalg(1+me_rows:1+me_rows+(icpu-1)*ncpu_rows:ncpu_rows))
1426 : : froms=1+(ispinor-1)*nrowsLinalgMe/nspinor+(col-1)*nrowsLinalgMe+nrowsLinalgMeSum*ncolsColsRows
1427 : : frome=froms-1+nrowsLinalgMe/nspinor
1428 : : tos=1+nrowsLinalgMeSum/nspinor+(ispinor-1)*nrowsColsRows/nspinor+(col-1)*nrowsColsRows
1429 : : toe=tos-1+nrowsLinalgMe/nspinor
1430 : : do irow = 0, toe-tos
1431 : : bufferMess(1,froms+irow) = bufferOrdered(1,tos+irow)
1432 : : bufferMess(2,froms+irow) = bufferOrdered(2,tos+irow)
1433 : : end do
1434 : : end do
1435 : : end do
1436 : : end do
1437 : : else
1438 : : #endif
1439 : : !$omp parallel do private(nrowsLinalgMe,nrowsLinalgMeSum,toe,tos,frome,froms), collapse(3)
1440 [ + + - + : 3503032 : do col = 1, ncolsColsRows
+ ]
1441 [ + + ]: 7637854 : do icpu = 0, ncpu_cols-1
1442 [ + + ]: 11499628 : do ispinor = 1, nspinor
1443 : 5076784 : nrowsLinalgMe = nrowsLinalg(1+me_rows+icpu*ncpu_rows)
1444 [ + + ]: 9871784 : nrowsLinalgMeSum = sum(nrowsLinalg(1+me_rows:1+me_rows+(icpu-1)*ncpu_rows:ncpu_rows))
1445 : 5076784 : froms=1+(ispinor-1)*nrowsLinalgMe/nspinor+(col-1)*nrowsLinalgMe+nrowsLinalgMeSum*ncolsColsRows
1446 : 5076784 : frome=froms-1+nrowsLinalgMe/nspinor
1447 : 5076784 : tos=1+nrowsLinalgMeSum/nspinor+(ispinor-1)*nrowsColsRows/nspinor+(col-1)*nrowsColsRows
1448 : 5076784 : toe=tos-1+nrowsLinalgMe/nspinor
1449 [ + + + + ]: 1550592342 : bufferMess(:,froms:frome) = bufferOrdered(:,tos:toe)
1450 : : end do
1451 : : end do
1452 : : end do
1453 : : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1454 : : end if
1455 : : #endif
1456 : : end select
1457 : :
1458 [ - + - - ]: 1656108 : if (xgTransposer%gpu_option /= ABI_GPU_DISABLED .and. xgTransposer%gpu_thread_limit /= 0) then
1459 : 0 : call xomp_set_num_threads(nthreads_bak)
1460 : : end if
1461 : :
1462 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
1463 : : ! if gpu enable restore OpenMP num threads to 1
1464 : : if (xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
1465 : : ! restore OMP_NUM_THREADS=1
1466 : : call xomp_set_num_threads(1)
1467 : : end if
1468 : :
1469 : : ! if gpu is enabled, transfer back data on GPU
1470 : : if (xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
1471 : : buffer_size = size(bufferOrdered) * dp
1472 : : call gpu_data_prefetch_async(C_LOC(bufferOrdered), buffer_size)
1473 : : call gpu_device_synchronize()
1474 : : end if
1475 : : #endif
1476 : :
1477 : 1656108 : call timab(tim_reorganize,2,tsec)
1478 : :
1479 : 1656108 : end subroutine xgTransposer_reorganizeData
1480 : : !!***
1481 : :
1482 : : !!****f* m_xgTransposer/xgTransposer_getRank
1483 : : !!
1484 : : !! NAME
1485 : : !! xgTransposer_getRank
1486 : :
1487 : 0 : function xgTransposer_getRank(xgTransposer, comm) result(rank)
1488 : : type(xgTransposer_t), intent(in ) :: xgTransposer
1489 : : integer , intent(in ) :: comm
1490 : : integer :: rank
1491 [ # # ]: 0 : if ( (comm > ubound(xgTransposer%mpiData,1)) .or. (comm < lbound(xgTransposer%mpiData,1)) ) then
1492 : 0 : ABI_ERROR("Value for communicator is wrong")
1493 : : end if
1494 : 0 : rank = xgTransposer%mpiData(comm)%rank
1495 : 0 : end function xgTransposer_getRank
1496 : : !!***
1497 : :
1498 : : !!****f* m_xgTransposer/xgTransposer_getComm
1499 : : !!
1500 : : !! NAME
1501 : : !! xgTransposer_getComm
1502 : :
1503 : 0 : function xgTransposer_getComm(xgTransposer, comm1) result(communicator)
1504 : : type(xgTransposer_t), intent(in ) :: xgTransposer
1505 : : integer , intent(in ) :: comm1
1506 : : integer :: communicator
1507 [ # # ]: 0 : if ( (comm1 > ubound(xgTransposer%mpiData,1)) .or. (comm1 < lbound(xgTransposer%mpiData,1)) ) then
1508 : 0 : ABI_ERROR("Value for communicator is wrong")
1509 : : end if
1510 : 0 : communicator = xgTransposer%mpiData(comm1)%comm
1511 : 0 : end function xgTransposer_getComm
1512 : : !!***
1513 : :
1514 : : !!****f* m_xgTransposer/xgTransposer_free
1515 : : !!
1516 : : !! NAME
1517 : : !! xgTransposer_free
1518 : :
1519 : 190862 : subroutine xgTransposer_free(xgTransposer)
1520 : :
1521 : : type(xgTransposer_t), intent(inout) :: xgTransposer
1522 : : double precision :: tsec(2)
1523 : : integer :: i
1524 : :
1525 : 190862 : call timab(tim_free,1,tsec)
1526 : : #ifdef HAVE_MPI
1527 [ + + ]: 190862 : if ( xgTransposer%type == TRANS_TYPE_CONSTRUCTED_NULL_COMM ) then
1528 : 10 : call mpi_comm_free(xgTransposer%mpiData(MPI_ROWS)%comm,i)
1529 : 10 : call mpi_comm_free(xgTransposer%mpiData(MPI_COLS)%comm,i)
1530 : 10 : call mpi_comm_free(xgTransposer%mpiData(MPI_2DCART)%comm,i)
1531 : : end if
1532 : : #else
1533 : : ABI_UNUSED(i)
1534 : : #endif
1535 : :
1536 [ + - ]: 190862 : if ( allocated(xgTransposer%lookup) ) then
1537 : 190862 : ABI_FREE(xgTransposer%lookup)
1538 : : end if
1539 : :
1540 [ + - ]: 190862 : if ( associated(xgTransposer%nrowsLinalg) ) then
1541 : 190862 : ABI_FREE(xgTransposer%nrowsLinalg)
1542 : : end if
1543 : :
1544 [ + + ]: 190862 : if ( associated(xgTransposer%buffer) ) then
1545 [ + - ]: 152936 : if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
1546 : : #if defined HAVE_GPU && defined HAVE_YAKL
1547 : : ABI_FREE_MANAGED(xgTransposer%buffer)
1548 : : #endif
1549 : : else
1550 : : if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
1551 : : #ifdef HAVE_OPENMP_OFFLOAD
1552 : : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
1553 : : !$OMP TARGET EXIT DATA MAP(delete:xgTransposer%buffer)
1554 : : #else
1555 : : call xg_ompgpu_map_delete(xgTransposer%buffer)
1556 : : #endif
1557 : : #endif
1558 : : end if
1559 : 152936 : ABI_FREE(xgTransposer%buffer)
1560 : : end if
1561 : : end if
1562 : :
1563 : 190862 : call timab(tim_free,2,tsec)
1564 : :
1565 : : #if defined HAVE_OPENMP_OFFLOAD && ! defined HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
1566 : : contains
1567 : : ! NOTE : routine to be used in order to avoid using derived-type syntax in OMP directives
1568 : : subroutine xg_ompgpu_map_delete(xgTransposerBuffer)
1569 : : real(dp),intent(inout) :: xgTransposerBuffer(:,:)
1570 : : !$OMP TARGET EXIT DATA MAP(delete:xgTransposerBuffer)
1571 : : end subroutine xg_ompgpu_map_delete
1572 : : #endif
1573 : 190862 : end subroutine xgTransposer_free
1574 : : !!***
1575 : :
1576 [ # # # # : 0 : end module m_xgTransposer
# # # # #
# # # ]
1577 : : !!***
|