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), ABI_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 46026 : 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 46026 : 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 46026 : call timab(tim_init,1,tsec)
193 :
194 46026 : xgTransposer%type = TRANS_TYPE_CONSTRUCTED
195 :
196 46026 : xgTransposer%xgBlock_linalg => xgBlock_linalg
197 46026 : xgTransposer%xgBlock_colsrows => xgBlock_colsrows
198 46026 : xgTransposer%state = state
199 46026 : xgTransposer%nspinor = nspinor
200 46026 : xgTransposer%me_g0_fft = me_g0_fft
201 46026 : xgTransposer%gpu_option = ABI_GPU_DISABLED
202 46026 : if(present(gpu_option)) xgTransposer%gpu_option = gpu_option
203 46026 : xgTransposer%gpu_thread_limit = 1
204 46026 : if(present(gpu_thread_limit)) xgTransposer%gpu_thread_limit = gpu_thread_limit
205 46026 : 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 46026 : else if (state==STATE_LINALG) then
211 46026 : commLinalg = comm(xgBlock_linalg)
212 : else
213 0 : ABI_ERROR("Invalid transposer state")
214 : end if
215 46026 : xgTransposer%mpiData(MPI_LINALG)%comm = commLinalg
216 46026 : xgTransposer%mpiData(MPI_LINALG)%rank = xmpi_comm_rank(commLinalg)
217 46026 : xgTransposer%mpiData(MPI_LINALG)%size = xmpi_comm_size(commLinalg)
218 :
219 : #if defined HAVE_MPI
220 46026 : 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 46016 : 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 46016 : 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 46016 : xgTransposer%mpiData(MPI_2DCART)%comm = xgTransposer%mpiData(MPI_LINALG)%comm
230 46016 : ncpuRows = xmpi_comm_size(comm_rows)
231 46016 : ncpuCols = xmpi_comm_size(comm_cols)
232 : end if
233 : #else
234 : ncpuRows = 1
235 : ncpuCols = 1
236 : #endif
237 46026 : 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 46026 : if ( ( algo == TRANS_ALL2ALL .or. algo == TRANS_GATHER ) ) then
244 46026 : 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 46026 : xgTransposer%nrowsLinalg => null()
258 46026 : xgTransposer%buffer => null()
259 46026 : 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 46026 : custom_ncolsColsRows_ = .false.
265 46026 : if (present(custom_ncolsColsRows)) custom_ncolsColsRows_ = custom_ncolsColsRows
266 46010 : 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 46026 : select case (state)
280 : case (STATE_LINALG)
281 :
282 : ! We are in the linalg representation.
283 : ! We need to construct the colsrows parallelization
284 46026 : call xgBlock_getSize(xgBlock_linalg,nrows,ncols)
285 :
286 : ! Get total number of rows in ColsRows representation
287 46026 : 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 46026 : 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 138078 : ABI_MALLOC(xgTransposer%lookup,(1:ncols))
304 : !ABI_MALLOC(xgTransposer%me_g0_lookup,())
305 476306 : do icol = 0, ncols-1
306 476306 : xgTransposer%lookup(icol+1) = MOD(icol,ncpuCols)
307 : end do
308 :
309 : #if defined HAVE_MPI
310 46026 : 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 46016 : comm_rows_=comm_rows
314 46016 : comm_cols_=comm_cols
315 : end if
316 46026 : xgTransposer%mpiData(MPI_ROWS)%comm = comm_rows_
317 46026 : 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 46026 : call xgTransposer_setComm(xgTransposer)
323 46026 : call xgTransposer_computeDistribution(xgTransposer)
324 92052 : 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 46026 : ABI_ERROR("State is undefined")
356 : end select
357 :
358 46026 : call timab(tim_init,2,tsec)
359 :
360 : !write(*,*) 'transposer construction done'
361 :
362 46026 : end subroutine xgTransposer_constructor
363 : !!***
364 :
365 : !!****f* m_xgTransposer/xgTransposer_copyConstructor
366 : !!
367 : !! NAME
368 : !! xgTransposer_copyConstructor
369 :
370 143852 : 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 143852 : call timab(tim_init,1,tsec)
386 :
387 143852 : xgTransposer%type = TRANS_TYPE_COPIED
388 :
389 143852 : xgTransposer%xgBlock_linalg => xgBlock_linalg
390 143852 : xgTransposer%xgBlock_colsrows => xgBlock_colsrows
391 143852 : xgTransposer%state = state
392 143852 : if (state==STATE_COLSROWS) then
393 0 : commLinalg = xgTransposerInitialized%mpiData(MPI_LINALG)%comm
394 143852 : else if (state==STATE_LINALG) then
395 143852 : commLinalg = comm(xgBlock_linalg)
396 : else
397 0 : ABI_ERROR("Invalid transposer state")
398 : end if
399 :
400 143852 : 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 719260 : xgTransposer%mpiData(:)%comm = xgTransposerInitialized%mpiData(:)%comm
405 143852 : xgTransposer%mpiData(MPI_LINALG)%rank = xmpi_comm_rank(commLinalg)
406 143852 : xgTransposer%mpiData(MPI_LINALG)%size = xmpi_comm_size(commLinalg)
407 143852 : xgTransposer%mpiData(MPI_COLS)%rank = xmpi_comm_rank(xgTransposer%mpiData(MPI_COLS)%comm)
408 143852 : xgTransposer%mpiData(MPI_COLS)%size = xmpi_comm_size(xgTransposer%mpiData(MPI_COLS)%comm)
409 143852 : xgTransposer%mpiData(MPI_ROWS)%rank = xmpi_comm_rank(xgTransposer%mpiData(MPI_ROWS)%comm)
410 143852 : xgTransposer%mpiData(MPI_ROWS)%size = xmpi_comm_size(xgTransposer%mpiData(MPI_ROWS)%comm)
411 :
412 143852 : xgTransposer%mpiAlgo = xgTransposerInitialized%mpiAlgo
413 143852 : xgTransposer%nspinor = xgTransposerInitialized%nspinor
414 143852 : xgTransposer%me_g0_fft = xgTransposerInitialized%me_g0_fft
415 143852 : xgTransposer%gpu_option = xgTransposerInitialized%gpu_option
416 :
417 : ! Nullify pointers in order to be able to query them
418 143852 : xgTransposer%nrowsLinalg => null()
419 143852 : xgTransposer%buffer => null()
420 143852 : xgTransposer%ncolsColsRows_sub => null()
421 : !write(std_out,*) '@nrowsLinalg', associated(xgTransposer%nrowsLinalg)
422 :
423 143852 : xgTransposer%custom_ncolsColsRows = xgTransposerInitialized%custom_ncolsColsRows
424 143852 : 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 143852 : ncpuCols = xgTransposer%mpiData(MPI_COLS)%size
430 143852 : ncpuRows = xgTransposer%mpiData(MPI_ROWS)%size
431 :
432 143852 : select case (state)
433 : case (STATE_LINALG)
434 :
435 143852 : call xgBlock_getSize(xgBlock_linalg,nrows,ncols)
436 :
437 : ! Get total number of rows in ColsRows representation
438 143852 : call xmpi_sum(nrows,commLinalg,ierr)
439 :
440 : ! Check that ncols can be uniformly distributed across ncolsCols MPI processes
441 143852 : 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 431556 : ABI_MALLOC(xgTransposer%lookup,(1:ncols))
454 143852 : 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 1169220 : xgTransposer%lookup(:) = xgTransposerInitialized%lookup(:)
460 : end if
461 :
462 143852 : call xgTransposer_computeDistribution(xgTransposer)
463 287704 : 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 143852 : ABI_ERROR("State is undefined")
485 : end select
486 :
487 143852 : call timab(tim_init,2,tsec)
488 :
489 143852 : end subroutine xgTransposer_copyConstructor
490 : !!***
491 :
492 : !!****f* m_xgTransposer/xgTransposer_setComm
493 : !!
494 : !! NAME
495 : !! xgTransposer_setComm
496 :
497 46026 : subroutine xgTransposer_setComm(xgTransposer)
498 :
499 : type(xgTransposer_t), intent(inout) :: xgTransposer
500 :
501 46026 : xgTransposer%mpiData(MPI_ROWS)%rank = xmpi_comm_rank(xgTransposer%mpiData(MPI_ROWS)%comm)
502 46026 : xgTransposer%mpiData(MPI_ROWS)%size = xmpi_comm_size(xgTransposer%mpiData(MPI_ROWS)%comm)
503 46026 : xgTransposer%mpiData(MPI_COLS)%rank = xmpi_comm_rank(xgTransposer%mpiData(MPI_COLS)%comm)
504 46026 : xgTransposer%mpiData(MPI_COLS)%size = xmpi_comm_size(xgTransposer%mpiData(MPI_COLS)%comm)
505 :
506 92052 : select case(xgTransposer%state)
507 : case (STATE_LINALG)
508 46026 : 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 46026 : ABI_ERROR("State is undefined")
513 : end select
514 :
515 46026 : 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 189878 : 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 569634 : ABI_MALLOC(xgTransposer%nrowsLinalg,(xgTransposer%mpiData(MPI_LINALG)%size))
586 189878 : nRealPairs = rows(xgTransposer%xgBlock_linalg)
587 189878 : if (MOD(nRealPairs,xgTransposer%nspinor)/=0) then
588 0 : ABI_ERROR('nspinor should divide nRealPairs!')
589 : end if
590 :
591 189878 : call xmpi_allgather(nRealPairs,xgTransposer%nrowsLinalg,xgTransposer%mpiData(MPI_LINALG)%comm,ierr)
592 189878 : if ( ierr /= xmpi_success ) then
593 0 : ABI_ERROR("Error while gathering number of rows in linalg")
594 : end if
595 :
596 189878 : ncpuCols = xgTransposer%mpiData(MPI_COLS)%size
597 189878 : ncpuRows = xgTransposer%mpiData(MPI_ROWS)%size
598 189878 : icpu_rows = xgTransposer%mpiData(MPI_ROWS)%rank
599 189878 : xgTransposer%nrowsColsRows = 0
600 686364 : do icpu_cols=0,ncpuCols-1
601 686364 : xgTransposer%nrowsColsRows = xgTransposer%nrowsColsRows + xgTransposer%nrowsLinalg(1+icpu_rows+icpu_cols*ncpuRows)
602 : end do
603 189878 : xgTransposer%ncolsColsRows = cols(xgTransposer%xgBlock_linalg)/ncpuCols
604 189878 : 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 189878 : 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 189878 : subroutine xgTransposer_makeXgBlock(xgTransposer)
680 :
681 : type(xgTransposer_t), intent(inout) :: xgTransposer
682 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD && !defined HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
683 : real(dp), ABI_CONTIGUOUS pointer :: xgTransposer__buffer(:,:)
684 : #endif
685 : integer :: ncolsColsRows, nrowsLinalgMe
686 :
687 379756 : select case (xgTransposer%state)
688 : case (STATE_LINALG)
689 :
690 189878 : ncolsColsRows = xgTransposer%ncolsColsRows
691 : !write(std_out,*) '@makeXgBlock dims', ncolsColsRows, xgTransposer%nrowsColsRows
692 : !call flush_unit(std_out)
693 :
694 : ! Assume xgBlock_colsrows is empty and not constructed because user cannot
695 : ! predict the size
696 189878 : if ( associated(xgTransposer%buffer) ) then
697 0 : if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
698 : #if defined HAVE_GPU && defined HAVE_YAKL
699 : ABI_FREE_MANAGED(xgTransposer%buffer)
700 : #endif
701 : else
702 : if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
703 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
704 : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
705 : !$OMP TARGET EXIT DATA MAP(delete:xgTransposer%buffer)
706 : #else
707 : xgTransposer__buffer => xgTransposer%buffer
708 : !$OMP TARGET EXIT DATA MAP(delete:xgTransposer__buffer)
709 : #endif
710 : #endif
711 : end if
712 0 : ABI_FREE(xgTransposer%buffer)
713 : end if
714 : end if
715 :
716 189878 : if ( xgTransposer%mpiData(MPI_COLS)%size == 1 ) then
717 37926 : xgTransposer%xgBlock_colsrows = xgTransposer%xgBlock_linalg
718 : else
719 151952 : if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
720 : #if defined HAVE_GPU && defined HAVE_YAKL
721 : ABI_MALLOC_MANAGED(xgTransposer%buffer,(/2,ncolsColsRows*xgTransposer%nrowsColsRows/))
722 : #endif
723 : else
724 455856 : ABI_MALLOC(xgTransposer%buffer,(2,ncolsColsRows*xgTransposer%nrowsColsRows))
725 : end if
726 : if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
727 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
728 : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
729 : !$OMP TARGET ENTER DATA MAP(alloc:xgTransposer%buffer)
730 : #else
731 : xgTransposer__buffer => xgTransposer%buffer
732 : !$OMP TARGET ENTER DATA MAP(alloc:xgTransposer__buffer)
733 : #endif
734 : #endif
735 : end if
736 : call xgBlock_map(xgTransposer%xgBlock_colsrows,xgTransposer%buffer,space(xgTransposer%xgBlock_linalg),&
737 : xgTransposer%nrowsColsRows,ncolsColsRows,xgTransposer%mpiData(MPI_ROWS)%comm,&
738 151952 : me_g0=xgTransposer%me_g0_fft,gpu_option=xgTransposer%gpu_option)
739 : end if
740 : case (STATE_COLSROWS)
741 :
742 : ! Assume xgBlock_linalg is empty and not constructed
743 :
744 : !write(std_out,*) 'db using @makeXgBlock not empty'
745 : !call flush_unit(std_out)
746 :
747 0 : nrowsLinalgMe = xgTransposer%nrowsLinalg(xgTransposer%mpiData(MPI_LINALG)%rank+1)
748 :
749 0 : if ( associated(xgTransposer%buffer) ) then
750 0 : if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
751 : #if defined HAVE_GPU && defined HAVE_YAKL
752 : ABI_FREE_MANAGED(xgTransposer%buffer)
753 : #endif
754 : else
755 : if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
756 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
757 : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
758 : !$OMP TARGET EXIT DATA MAP(delete:xgTransposer%buffer)
759 : #else
760 : xgTransposer__buffer => xgTransposer%buffer
761 : !$OMP TARGET EXIT DATA MAP(delete:xgTransposer__buffer)
762 : #endif
763 : #endif
764 : end if
765 0 : ABI_FREE(xgTransposer%buffer)
766 : end if
767 : end if
768 :
769 0 : if ( xgTransposer%mpiData(MPI_COLS)%size == 1 ) then
770 0 : xgTransposer%xgBlock_linalg = xgTransposer%xgBlock_colsrows
771 : else
772 0 : if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
773 : #if defined HAVE_GPU && defined HAVE_YAKL
774 : ABI_MALLOC_MANAGED(xgTransposer%buffer,(/2,xgTransposer%ncolsLinalg*nrowsLinalgMe/))
775 : #endif
776 : else
777 0 : ABI_MALLOC(xgTransposer%buffer,(2,xgTransposer%ncolsLinalg*nrowsLinalgMe))
778 : end if
779 : if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
780 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
781 : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
782 : !$OMP TARGET ENTER DATA MAP(alloc:xgTransposer%buffer)
783 : #else
784 : xgTransposer__buffer => xgTransposer%buffer
785 : !$OMP TARGET ENTER DATA MAP(alloc:xgTransposer__buffer)
786 : #endif
787 : #endif
788 : end if
789 : call xgBlock_map(xgTransposer%xgBlock_linalg,xgTransposer%buffer,space(xgTransposer%xgBlock_colsrows),&
790 : nrowsLinalgMe,xgTransposer%ncolsLinalg,xgTransposer%mpiData(MPI_COLS)%comm,&
791 0 : me_g0=xgTransposer%me_g0_fft,gpu_option=xgTransposer%gpu_option)
792 : end if
793 :
794 : case default
795 189878 : ABI_ERROR("State unknown")
796 : end select
797 :
798 : !if (associated(xgTransposer%buffer)) then
799 : ! write(std_out,*) '@makeXgBlock alloc buffer size', size(xgTransposer%buffer,1), size(xgTransposer%buffer,2)
800 : ! call flush_unit(std_out)
801 : !end if
802 :
803 189878 : end subroutine xgTransposer_makeXgBlock
804 : !!***
805 :
806 : !!****f* m_xgTransposer/xgTransposer_transpose
807 : !!
808 : !! NAME
809 : !! xgTransposer_transpose
810 :
811 2803692 : subroutine xgTransposer_transpose(xgTransposer,toState)
812 :
813 : type(xgTransposer_t), intent(inout) :: xgTransposer
814 : integer , intent(in ) :: toState
815 : double precision :: tsec(2)
816 :
817 2803692 : call timab(tim_transpose,1,tsec)
818 :
819 2803692 : if ( toState /= STATE_LINALG .and. toState /= STATE_COLSROWS ) then
820 0 : ABI_ERROR("Bad value for toState")
821 : end if
822 :
823 : !write(std_out,*) "linalg", rows(xgTransposer%xgBlock_linalg)*cols(xgTransposer%xgBlock_linalg)
824 : !write(std_out,*) "colsrows", rows(xgTransposer%xgBlock_colsrows)*cols(xgTransposer%xgBlock_colsrows)
825 2074480 : select case (toState)
826 : case (STATE_LINALG)
827 2074480 : if ( xgTransposer%state == STATE_LINALG ) then
828 0 : ABI_WARNING("Array linalg has already been transposed")
829 : end if
830 2074480 : if ( xgTransposer%mpiData(MPI_COLS)%size > 1 ) then
831 1212586 : call xgTransposer_toLinalg(xgTransposer)
832 : else
833 861894 : xgTransposer%state = STATE_LINALG
834 : end if
835 : case (STATE_COLSROWS)
836 729212 : if ( xgTransposer%state == STATE_COLSROWS ) then
837 0 : ABI_WARNING("Array colsrows has already been transposed")
838 : end if
839 3532904 : if ( xgTransposer%mpiData(MPI_COLS)%size > 1 ) then
840 440290 : call xgTransposer_toColsRows(xgTransposer)
841 : else
842 288922 : xgTransposer%state = STATE_COLSROWS
843 : end if
844 : end select
845 :
846 2803692 : call timab(tim_transpose,2,tsec)
847 :
848 2803692 : end subroutine xgTransposer_transpose
849 : !!***
850 :
851 : !!****f* m_xgTransposer/xgTransposer_toLinalg
852 : !!
853 : !! NAME
854 : !! xgTransposer_toLinalg
855 :
856 1212586 : subroutine xgTransposer_toLinalg(xgTransposer)
857 :
858 : type(xgTransposer_t), intent(inout) :: xgTransposer
859 1212586 : double precision, allocatable, target :: sendbuf(:,:)
860 1212586 : double precision, pointer :: recvbuf(:,:)
861 : !double precision, pointer :: buffer(:,:)
862 1212586 : integer, allocatable :: sendcounts(:), recvcounts(:)
863 1212586 : integer, allocatable :: sdispls(:), rdispls(:)
864 : integer :: ncpu_cols, ncpu_rows, comm, me_rows, me_cols
865 : integer :: nrowsColsRows
866 : integer :: ncolsColsRows
867 : integer :: nrowsLinalgMe
868 : integer :: icpu, ierr
869 : integer :: send_start,send_end
870 : !integer :: myrequest
871 : !integer, allocatable :: request(:), status(:)
872 1212586 : type(ptr_t), allocatable :: sendptrbuf(:)
873 1212586 : integer, pointer :: nrowsLinalg(:)
874 1212586 : integer, pointer :: ncolsColsRows_sub(:)
875 : double precision :: tsec(2)
876 :
877 : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
878 : double precision, allocatable :: recvbuf_mpi(:,:)
879 : integer(c_size_t) :: buffer_size
880 : #endif
881 :
882 1212586 : call timab(tim_toLinalg,1,tsec)
883 :
884 1212586 : ncpu_cols = xgTransposer%mpiData(MPI_COLS)%size
885 1212586 : ncpu_rows = xgTransposer%mpiData(MPI_ROWS)%size
886 1212586 : comm = xgTransposer%mpiData(MPI_COLS)%comm
887 : ! me = xgTransposer%mpiData(MPI_ROWS)%rank*ncpu
888 1212586 : me_rows = xgTransposer%mpiData(MPI_ROWS)%rank
889 :
890 : !write(std_out,'(a,i4,i4,i4)') '@xgTransposer_toLinalg for MPI proc in comm=', &
891 : ! xgTransposer%mpiData(MPI_LINALG)%rank, xgTransposer%mpiData(MPI_LINALG)%comm, comm
892 :
893 1212586 : nrowsColsRows = xgTransposer%nrowsColsRows
894 1212586 : ncolsColsRows = xgTransposer%ncolsColsRows
895 :
896 1212586 : nrowsLinalg => xgTransposer%nrowsLinalg
897 1212586 : if (xgTransposer%custom_ncolsColsRows) ncolsColsRows_sub => xgTransposer%ncolsColsRows_sub
898 1212586 : nrowsLinalgMe = nrowsLinalg(xgTransposer%mpiData(MPI_LINALG)%rank+1)
899 :
900 : !write(std_out,*) '@toLinalg nrowsColsRows, ncolsColsRows=', nrowsColsRows, ncolsColsRows
901 : !call flush_unit(std_out)
902 :
903 3637758 : ABI_MALLOC(sendbuf,(2,nrowsColsRows*ncolsColsRows))
904 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
905 : !$OMP TARGET ENTER DATA MAP(alloc:sendbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
906 : #endif
907 1212586 : call xgTransposer_reorganizeData(xgTransposer,sendbuf)
908 :
909 3637758 : ABI_MALLOC(recvcounts,(ncpu_cols))
910 2425172 : ABI_MALLOC(rdispls,(ncpu_cols))
911 1212586 : if (xgTransposer%custom_ncolsColsRows) then ! non uniform sendcounts
912 0 : recvcounts(:) = 2*nrowsLinalgMe*ncolsColsRows_sub(:)
913 : else
914 4164006 : recvcounts(:) = 2*nrowsLinalgMe*ncolsColsRows !! Thank you fortran for not starting at 0 !
915 : end if
916 1212586 : rdispls(1) = 0
917 2951420 : do icpu = 2, ncpu_cols
918 2951420 : rdispls(icpu) = rdispls(icpu-1)+recvcounts(icpu-1)
919 : end do
920 :
921 : call xgBlock_reverseMap(xgTransposer%xgBlock_linalg,recvbuf,&
922 1212586 : & rows=1,cols=cols(xgTransposer%xgBlock_linalg)*nrowsLinalgMe)
923 :
924 : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
925 : ! just for debug
926 : if( xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
927 : !call gpu_managed_ptr_status(C_LOC(recvbuf))
928 : buffer_size = size(recvbuf) * dp
929 : call gpu_data_prefetch_async(C_LOC(recvbuf), buffer_size, CPU_DEVICE_ID)
930 : call gpu_device_synchronize()
931 :
932 : ABI_MALLOC(recvbuf_mpi, (size(recvbuf,1), size(recvbuf,2)) )
933 :
934 : end if
935 : #endif
936 :
937 2425172 : ABI_MALLOC(sendcounts,(ncpu_cols))
938 2425172 : ABI_MALLOC(sdispls,(ncpu_cols))
939 : !sendcounts(:) = 2*nrowsLinalg(me+1:me+ncpu)*ncolsColsRows
940 4164006 : do icpu=0,ncpu_cols-1
941 4164006 : sendcounts(icpu+1) = 2*nrowsLinalg(me_rows+1+icpu*ncpu_rows)*ncolsColsRows
942 : !write(*,*) "nrowsLinalg ", icpu, ncpu_rows, nrowsLinalg(me_rows+1+icpu*ncpu_rows)
943 : end do
944 1212586 : sdispls(1) = 0
945 2951420 : do icpu = 2, ncpu_cols
946 2951420 : sdispls(icpu) = sdispls(icpu-1)+sendcounts(icpu-1)
947 : end do
948 :
949 2425172 : select case(xgTransposer%mpiAlgo)
950 : case (TRANS_ALL2ALL)
951 : !ABI_MALLOC(request,(1))
952 : !myrequest = 1
953 :
954 : ABI_NVTX_START_RANGE(NVTX_TRANSPOSER_MPI_ALL2ALL)
955 :
956 1212586 : if( xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
957 :
958 : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
959 : call timab(tim_all2allv,1,tsec)
960 : call xmpi_alltoallv(sendbuf, sendcounts, sdispls, &
961 : recvbuf_mpi, recvcounts, rdispls, &
962 : comm, ierr)
963 : call timab(tim_all2allv,2,tsec)
964 : !call xmpi_ialltoallv(sendbuf, sendcounts, sdispls, &
965 : ! recvbuf, recvcounts, rdispls, &
966 : ! comm, request(myrequest))
967 :
968 : ! copy back recvbuf_mpi into recvbuf
969 : recvbuf(:,:) = recvbuf_mpi(:,:)
970 :
971 : ABI_FREE(recvbuf_mpi)
972 : #endif
973 :
974 : else
975 :
976 : !write(*,*) "sendcounts for proc ", me_rows, xmpi_comm_rank(comm), ":", sendcounts(:)
977 : !write(*,*) "recvcounts for proc ", me_rows, xmpi_comm_rank(comm), ":", recvcounts(:)
978 :
979 1212586 : call timab(tim_all2allv,1,tsec)
980 : call xmpi_alltoallv(sendbuf, sendcounts, sdispls, &
981 : recvbuf, recvcounts, rdispls, &
982 1212586 : comm, ierr, use_omp_map=(xgTransposer%gpu_option==ABI_GPU_OPENMP))
983 1212586 : call timab(tim_all2allv,2,tsec)
984 : !call xmpi_ialltoallv(sendbuf, sendcounts, sdispls, &
985 : ! recvbuf, recvcounts, rdispls, &
986 : ! comm, request(myrequest))
987 :
988 : end if
989 :
990 : ABI_NVTX_END_RANGE()
991 :
992 : case (TRANS_GATHER)
993 :
994 : !ABI_MALLOC(request,(ncpu))
995 0 : me_cols = xgTransposer%mpiData(MPI_COLS)%rank
996 : !myrequest = me+1
997 :
998 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
999 : !$OMP TARGET UPDATE FROM(sendbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1000 : #endif
1001 0 : ABI_MALLOC(sendptrbuf,(1:ncpu_cols))
1002 0 : do icpu = 1, ncpu_cols
1003 0 : send_start = sdispls(icpu)/2+1
1004 0 : send_end = sdispls(icpu)/2+sendcounts(icpu)/2
1005 0 : sendptrbuf(icpu)%ptr => sendbuf(:,send_start:send_end)
1006 0 : call timab(tim_gatherv,1,tsec)
1007 0 : call xmpi_gatherv(sendptrbuf(icpu)%ptr,sendcounts(icpu),recvbuf,recvcounts,rdispls,icpu-1,comm,ierr)
1008 0 : call timab(tim_gatherv,2,tsec)
1009 : !call mpi_igatherv(sendptrbuf(me+1)%ptr,sendcounts(icpu),MPI_DOUBLE_PRECISION,&
1010 : ! recvbuf,recvcounts,rdispls,MPI_DOUBLE_PRECISION,icpu-1,comm,request(icpu),ierr)
1011 : end do
1012 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1013 : !$OMP TARGET UPDATE TO(recvbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1014 : #endif
1015 :
1016 : case default
1017 1212586 : ABI_BUG("This algo does not exist")
1018 : end select
1019 :
1020 1212586 : xgTransposer%state = STATE_LINALG
1021 :
1022 : !ABI_MALLOC(status,(MPI_STATUS_SIZE))
1023 : !call mpi_wait(request(myrequest),status,ierr)
1024 1212586 : if ( ierr /= xmpi_success ) then
1025 0 : ABI_ERROR("Error while waiting for mpi")
1026 : end if
1027 :
1028 1212586 : if ( allocated(sendcounts) ) then
1029 1212586 : ABI_FREE(sendcounts)
1030 : end if
1031 : if ( allocated(sdispls) ) then
1032 1212586 : ABI_FREE(sdispls)
1033 : end if
1034 :
1035 1212586 : ABI_FREE(recvcounts)
1036 1212586 : ABI_FREE(rdispls)
1037 :
1038 :
1039 1212586 : if ( allocated(sendptrbuf) ) then
1040 0 : ABI_FREE(sendptrbuf)
1041 : end if
1042 :
1043 : !do icpu = 1, size(request)
1044 : ! if ( icpu /= myrequest ) then
1045 : ! call mpi_wait(request(icpu),status,ierr)
1046 : ! if ( ierr /= MPI_SUCCESS ) then
1047 : ! ABI_ERROR("Error while waiting for other mpi")
1048 : ! end if
1049 : ! end if
1050 : !end do
1051 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1052 : !$OMP TARGET EXIT DATA MAP(delete:sendbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1053 : #endif
1054 1212586 : ABI_FREE(sendbuf)
1055 : !ABI_FREE(status)
1056 : !ABI_FREE(request)
1057 :
1058 1212586 : call timab(tim_toLinalg,2,tsec)
1059 :
1060 2425172 : end subroutine xgTransposer_toLinalg
1061 : !!***
1062 :
1063 : !!****f* m_xgTransposer/xgTransposer_toColsRows
1064 : !!
1065 : !! NAME
1066 : !! xgTransposer_toColsRows
1067 :
1068 440290 : subroutine xgTransposer_toColsRows(xgTransposer)
1069 :
1070 : type(xgTransposer_t), intent(inout) :: xgTransposer
1071 440290 : double precision, pointer :: sendbuf(:,:)
1072 440290 : double precision, allocatable :: recvbuf(:,:)
1073 : !double precision, allocatable :: buffer(:,:)
1074 440290 : integer, allocatable :: sendcounts(:), recvcounts(:)
1075 440290 : integer, allocatable :: sdispls(:), rdispls(:)
1076 : integer :: ncpu_rows, ncpu_cols, comm, me_rows, me_cols
1077 : integer :: nrowsColsRows
1078 : integer :: nrowsLinalgMe
1079 : integer :: ncolsColsRows
1080 : integer :: icpu,ierr
1081 : !integer :: myrequest
1082 : !integer, allocatable :: request(:), status(:)
1083 : type(xgBlock_t) :: xgBlock_toTransposed
1084 440290 : type(ptr_t), allocatable :: sendptrbuf(:)
1085 440290 : integer, pointer :: nrowsLinalg(:)
1086 440290 : integer, pointer :: ncolsColsRows_sub(:)
1087 : double precision :: tsec(2)
1088 :
1089 : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
1090 : double precision, allocatable :: sendbuf_mpi(:,:)
1091 : integer(c_size_t) :: buffer_size
1092 : #endif
1093 :
1094 440290 : call timab(tim_toColsRows,1,tsec)
1095 :
1096 440290 : ncpu_cols = xgTransposer%mpiData(MPI_COLS)%size
1097 440290 : ncpu_rows = xgTransposer%mpiData(MPI_ROWS)%size
1098 440290 : comm = xgTransposer%mpiData(MPI_COLS)%comm
1099 440290 : me_rows = xgTransposer%mpiData(MPI_ROWS)%rank
1100 :
1101 440290 : nrowsColsRows = xgTransposer%nrowsColsRows
1102 440290 : ncolsColsRows = xgTransposer%ncolsColsRows
1103 :
1104 440290 : nrowsLinalg => xgTransposer%nrowsLinalg
1105 440290 : if (xgTransposer%custom_ncolsColsRows) ncolsColsRows_sub => xgTransposer%ncolsColsRows_sub
1106 440290 : nrowsLinalgMe = nrowsLinalg(xgTransposer%mpiData(MPI_LINALG)%rank+1)
1107 :
1108 1320870 : ABI_MALLOC(recvbuf,(2,nrowsColsRows*ncolsColsRows))
1109 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1110 : !$OMP TARGET ENTER DATA MAP(alloc:recvbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1111 : #endif
1112 1320870 : ABI_MALLOC(recvcounts,(ncpu_cols))
1113 880580 : ABI_MALLOC(rdispls,(ncpu_cols))
1114 :
1115 : ! 2 * ncolsColsRows * nrowsLinalg_sub(:)
1116 1561694 : recvcounts(:) = 2*ncolsColsRows*nrowsLinalg(1+me_rows:1+me_rows+(ncpu_cols-1)*ncpu_rows:ncpu_rows)
1117 440290 : rdispls(1) = 0
1118 1121404 : do icpu = 2, ncpu_cols
1119 1121404 : rdispls(icpu) = rdispls(icpu-1)+recvcounts(icpu-1)
1120 : end do
1121 : !recvcounts(:) = 2*nrowsLinalg(me+1:me+ncpu)*ncolsColsRows
1122 : !rdispls(1) = 0
1123 : !do icpu = 2, ncpu
1124 : ! rdispls(icpu) = rdispls(icpu-1)+recvcounts(icpu-1)
1125 : !end do
1126 :
1127 880580 : select case(xgTransposer%mpiAlgo)
1128 : case (TRANS_ALL2ALL)
1129 880580 : ABI_MALLOC(sendcounts,(ncpu_cols))
1130 880580 : ABI_MALLOC(sdispls,(ncpu_cols))
1131 : !ABI_MALLOC(request,(1))
1132 : !myrequest = 1
1133 :
1134 : ! 2 * ncolsColsRows_sub(:) * nrowsLinalgMe
1135 440290 : if (xgTransposer%custom_ncolsColsRows) then ! non uniform sendcounts
1136 80 : sendcounts(:) = 2*nrowsLinalgMe*ncolsColsRows_sub(:)
1137 : else
1138 1561614 : sendcounts(:) = 2*nrowsLinalgMe*ncolsColsRows !! Thank you fortran for not starting at 0 !
1139 : end if
1140 440290 : sdispls(1) = 0
1141 1121404 : do icpu = 2, ncpu_cols
1142 1121404 : sdispls(icpu) = sdispls(icpu-1)+sendcounts(icpu-1)
1143 : end do
1144 :
1145 : call xgBlock_reverseMap(xgTransposer%xgBlock_linalg,sendbuf, &
1146 440290 : & rows=1,cols=cols(xgTransposer%xgBlock_linalg)*nrowsLinalgMe)
1147 : !write(*,*) "Before ialltoall"
1148 :
1149 : ABI_NVTX_START_RANGE(NVTX_TRANSPOSER_MPI_ALL2ALL)
1150 :
1151 : ! if gpu is enabled, data are located in GPU memory, so we copy them on a host buffer
1152 440290 : if( xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
1153 : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
1154 : ABI_MALLOC(sendbuf_mpi, (size(sendbuf,1), size(sendbuf,2)) )
1155 :
1156 : ! sync sendbuf on host and then copy to sendbuf_mpi
1157 : buffer_size = size(sendbuf) * dp
1158 : call gpu_data_prefetch_async(C_LOC(sendbuf), buffer_size, CPU_DEVICE_ID)
1159 : call gpu_device_synchronize()
1160 :
1161 : sendbuf_mpi(:,:) = sendbuf(:,:)
1162 :
1163 : call timab(tim_all2allv,1,tsec)
1164 : ! this is a cpu mpi comm
1165 : call xmpi_alltoallv(sendbuf_mpi, sendcounts, sdispls, &
1166 : recvbuf, recvcounts, rdispls, &
1167 : comm, ierr)
1168 : call timab(tim_all2allv,2,tsec)
1169 : !call xmpi_ialltoallv(sendbuf, sendcounts, sdispls, &
1170 : ! recvbuf, recvcounts, rdispls, &
1171 : ! comm, request(myrequest))
1172 : !write(*,*) "After ialltoall"
1173 :
1174 : ABI_FREE(sendbuf_mpi)
1175 : #endif
1176 : else
1177 :
1178 : !write(std_out,'(a,i4,i4,i4)') '@xgTransposer_toColsRows for MPI proc in comm=', &
1179 : ! xgTransposer%mpiData(MPI_LINALG)%rank, xgTransposer%mpiData(MPI_LINALG)%comm, comm
1180 : !if (xgTransposer%custom_ncolsColsRows) then
1181 : ! write(std_out,*) 'db ncolsColsRows_sub', ncolsColsRows_sub(:)
1182 : !end if
1183 : !write(std_out,*) 'db ncolsColsMe', ncolsColsRows
1184 : !write(std_out,*) 'db nrowsLinalg', nrowsLinalg(1+me_rows:1+me_rows+(ncpu_cols-1)*ncpu_rows:ncpu_rows)
1185 : !write(std_out,*) 'db nrowsLinalgMe', nrowsLinalgMe
1186 : !write(std_out,*) 'db sendbuf', cols(xgTransposer%xgBlock_linalg), nrowsLinalgMe
1187 : !write(std_out,*) 'db recvbuf', nrowsColsRows, ncolsColsRows
1188 : !write(std_out,*) 'db sendcounts', sendcounts(:)
1189 : !write(std_out,*) 'db recvcounts', recvcounts(:)
1190 : !write(std_out,*) 'db sdispls', sdispls(:)
1191 : !write(std_out,*) 'db rdispls', rdispls(:)
1192 : !call flush_unit(std_out)
1193 :
1194 440290 : call timab(tim_all2allv,1,tsec)
1195 : call xmpi_alltoallv(sendbuf, sendcounts, sdispls, &
1196 : recvbuf, recvcounts, rdispls, &
1197 440290 : comm, ierr, use_omp_map=(xgTransposer%gpu_option == ABI_GPU_OPENMP))
1198 440290 : call timab(tim_all2allv,2,tsec)
1199 : !call xmpi_ialltoallv(sendbuf, sendcounts, sdispls, &
1200 : ! recvbuf, recvcounts, rdispls, &
1201 : ! comm, request(myrequest))
1202 : !write(*,*) "After ialltoall"
1203 : end if
1204 :
1205 : ABI_NVTX_END_RANGE()
1206 :
1207 : case (TRANS_GATHER)
1208 :
1209 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1210 : !$OMP TARGET UPDATE FROM(sendbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1211 : #endif
1212 : !ABI_MALLOC(request,(ncpu))
1213 0 : me_cols = xgTransposer%mpiData(MPI_COLS)%rank
1214 : !myrequest = me+1
1215 :
1216 0 : ABI_MALLOC(sendptrbuf,(1:ncpu_cols))
1217 : !call flush(6)
1218 : !call xmpi_barrier(xgTransposer%mpiData(MPI_LINALG)%comm)
1219 0 : do icpu = 0, ncpu_cols-1
1220 : !write(*,*) me, "->", icpu, "from col ",icpu*ncolsColsRows+1, " number of rows:", nrowsLinalgMe
1221 0 : call xgBlock_setBlock(xgTransposer%xgBlock_linalg,xgBlock_toTransposed,nrowsLinalgMe,ncolsColsRows,fcol=icpu*ncolsColsRows+1)
1222 : call xgBlock_reverseMap(xgBlock_toTransposed,sendptrbuf(icpu+1)%ptr,&
1223 0 : rows=1,cols=ncolsColsRows*nrowsLinalgMe)
1224 0 : call timab(tim_gatherv,1,tsec)
1225 0 : call xmpi_gatherv(sendptrbuf(icpu+1)%ptr,2*ncolsColsRows*nrowsLinalgMe,recvbuf,recvcounts,rdispls,icpu,comm,ierr)
1226 0 : call timab(tim_gatherv,2,tsec)
1227 : !call mpi_igatherv(sendptrbuf(me+1)%ptr,2*ncolsColsRows*nrowsLinalgMe,MPI_DOUBLE_PRECISION,&
1228 : ! recvbuf,recvcounts,rdispls,MPI_DOUBLE_PRECISION,icpu,comm,request(icpu+1),ierr)
1229 : end do
1230 : !call xmpi_barrier(xgTransposer%mpiData(MPI_LINALG)%comm)
1231 : !call flush(6)
1232 : !write(*,*) me, request
1233 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1234 : !$OMP TARGET UPDATE TO(recvbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1235 : #endif
1236 :
1237 : case default
1238 440290 : ABI_BUG("This algo does not exist")
1239 : end select
1240 :
1241 : !ABI_MALLOC(status,(MPI_STATUS_SIZE))
1242 : !call mpi_wait(request(myrequest),status,ierr)
1243 : !write(*,*) "Request ended"
1244 440290 : if ( ierr /= xmpi_success ) then
1245 0 : ABI_ERROR("Error while waiting for mpi")
1246 : end if
1247 : !write(*,*) "with success"
1248 :
1249 440290 : call xgTransposer_reorganizeData(xgTransposer,recvbuf)
1250 :
1251 440290 : if ( allocated(sendcounts) ) then
1252 440290 : ABI_FREE(sendcounts)
1253 : end if
1254 440290 : if ( allocated(sdispls) ) then
1255 440290 : ABI_FREE(sdispls)
1256 : end if
1257 :
1258 440290 : ABI_FREE(recvcounts)
1259 440290 : ABI_FREE(rdispls)
1260 :
1261 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1262 : !$OMP TARGET EXIT DATA MAP(delete:recvbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
1263 : #endif
1264 440290 : ABI_FREE(recvbuf)
1265 :
1266 440290 : if ( allocated(sendptrbuf) ) then
1267 0 : ABI_FREE(sendptrbuf)
1268 : end if
1269 :
1270 440290 : xgTransposer%state = STATE_COLSROWS
1271 :
1272 : !do icpu = 1, size(request)
1273 : ! if ( icpu /= myrequest ) then
1274 : ! call mpi_wait(request(icpu),status,ierr)
1275 : ! if ( ierr /= MPI_SUCCESS ) then
1276 : ! ABI_ERROR("Error while waiting for other mpi")
1277 : ! end if
1278 : ! end if
1279 : !end do
1280 : !ABI_FREE(status)
1281 : !ABI_FREE(request)
1282 :
1283 440290 : call timab(tim_toColsRows,2,tsec)
1284 :
1285 880580 : end subroutine xgTransposer_toColsRows
1286 : !!***
1287 :
1288 : !!****f* m_xgTransposer/xgTransposer_reorganizeData
1289 : !!
1290 : !! NAME
1291 : !! xgTransposer_reorganizeData
1292 :
1293 1652876 : subroutine xgTransposer_reorganizeData(xgTransposer,bufferMess)
1294 :
1295 : type(xgTransposer_t), intent(inout) :: xgTransposer
1296 : double precision , intent(inout) :: bufferMess(:,:)
1297 : double precision, pointer :: bufferOrdered(:,:) => null()
1298 : integer :: nrowsColsRows,nthreads_bak
1299 : integer :: ncolsColsRows
1300 : integer :: tos,toe,froms,frome
1301 : integer :: col, icpu
1302 : integer :: me_rows,ncpu_cols,ncpu_rows
1303 : integer :: nPair,ispinor,nspinor
1304 : integer :: nrowsLinalgMe,nrowsLinalgMeSum
1305 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1306 : integer :: irow
1307 : integer, ABI_CONTIGUOUS pointer :: nrowsLinalg(:)
1308 : #else
1309 1652876 : integer, pointer :: nrowsLinalg(:)
1310 : #endif
1311 : double precision :: tsec(2)
1312 : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
1313 : integer(c_size_t) :: buffer_size
1314 : #endif
1315 :
1316 1652876 : call timab(tim_reorganize,1,tsec)
1317 :
1318 1652876 : me_rows = xgTransposer%mpiData(MPI_ROWS)%rank!*xgTransposer%mpiData(MPI_COLS)%size
1319 1652876 : ncpu_rows = xgTransposer%mpiData(MPI_ROWS)%size
1320 1652876 : ncpu_cols = xgTransposer%mpiData(MPI_COLS)%size
1321 :
1322 : ! These arrays are stored as
1323 : ! 1 2 3
1324 : ! 4 5 6
1325 : ! 7 8 9
1326 : !
1327 : ! stored as
1328 : !
1329 : ! column-major, one-base indexing : 1 4 7 2 5 8 3 6 9 | y + N_y*(x-1)
1330 : ! row-major , zero-based indexing: 1 2 3 4 5 6 7 8 9 | x + N_x*y
1331 1652876 : nspinor = xgTransposer%nspinor
1332 1652876 : nrowsLinalg => xgTransposer%nrowsLinalg
1333 :
1334 1652876 : nrowsColsRows = xgTransposer%nrowsColsRows
1335 1652876 : ncolsColsRows = xgTransposer%ncolsColsRows
1336 :
1337 : ! number of pairs depends on current MPI process
1338 1652876 : nPair = nrowsColsRows*ncolsColsRows
1339 1652876 : call xgBlock_reverseMap(xgTransposer%xgBlock_colsrows,bufferOrdered,rows=1,cols=nPair)
1340 :
1341 : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
1342 : ! if gpu is enabled, data are located in GPU memory, so we prefetch them on host
1343 : ! to do the following reorganization.
1344 : ! Alternatively, we should provide a GPU implementation of this data layout reorganization
1345 : if( xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
1346 : buffer_size = size(bufferOrdered) * dp
1347 : call gpu_data_prefetch_async(C_LOC(bufferOrdered), buffer_size, CPU_DEVICE_ID)
1348 : call gpu_device_synchronize()
1349 : end if
1350 :
1351 : ! if gpu enabled increase locally OpenMP num threads
1352 : if (xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
1353 : call xomp_set_num_threads(xgTransposer%gpu_kokkos_nthrd)
1354 : end if
1355 : #endif
1356 :
1357 1652876 : if (xgTransposer%gpu_option /= ABI_GPU_DISABLED .and. xgTransposer%gpu_thread_limit /= 0) then
1358 0 : nthreads_bak=xomp_get_max_threads()
1359 0 : call xomp_set_num_threads(min(xgTransposer%gpu_thread_limit,nthreads_bak))
1360 : end if
1361 :
1362 1652876 : select case (xgTransposer%state)
1363 : case (STATE_LINALG)
1364 : ! We are going to STATE_COLSROWS so we are after all2all
1365 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1366 : if (xgTransposer%gpu_option == ABI_GPU_OPENMP) then
1367 : ! bufferMess and bufferOrdered already live on the device: reorganize them there directly,
1368 : ! no CPU<->GPU round trip needed. nrowsLinalg is small and mapped just for this kernel.
1369 : ! Array-section assignment is not usable in device kernel bodies: loop explicitly instead.
1370 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1371 : !$OMP& PRIVATE(nrowsLinalgMe,nrowsLinalgMeSum,toe,tos,frome,froms,irow) &
1372 : !$OMP& MAP(to:bufferMess) MAP(to:bufferOrdered) MAP(to:nrowsLinalg)
1373 : do col = 1, ncolsColsRows
1374 : do icpu = 0, ncpu_cols-1
1375 : do ispinor = 1, nspinor
1376 : nrowsLinalgMe = nrowsLinalg(1+me_rows+icpu*ncpu_rows)
1377 : nrowsLinalgMeSum = sum(nrowsLinalg(1+me_rows:1+me_rows+(icpu-1)*ncpu_rows:ncpu_rows))
1378 : froms=1+(ispinor-1)*nrowsLinalgMe/nspinor+(col-1)*nrowsLinalgMe+nrowsLinalgMeSum*ncolsColsRows
1379 : frome=froms-1+nrowsLinalgMe/nspinor
1380 : tos=1+nrowsLinalgMeSum/nspinor+(ispinor-1)*nrowsColsRows/nspinor+(col-1)*nrowsColsRows
1381 : toe=tos-1+nrowsLinalgMe/nspinor
1382 : do irow = 0, toe-tos
1383 : bufferOrdered(1,tos+irow) = bufferMess(1,froms+irow)
1384 : bufferOrdered(2,tos+irow) = bufferMess(2,froms+irow)
1385 : end do
1386 : end do
1387 : end do
1388 : end do
1389 : else
1390 : #endif
1391 : !$omp parallel do private(nrowsLinalgMe,nrowsLinalgMeSum,toe,tos,frome,froms), collapse(3)
1392 1099232 : do col = 1, ncolsColsRows
1393 2795640 : do icpu = 0, ncpu_cols-1
1394 4264606 : do ispinor = 1, nspinor
1395 1909256 : nrowsLinalgMe = nrowsLinalg(1+me_rows+icpu*ncpu_rows)
1396 3833780 : nrowsLinalgMeSum = sum(nrowsLinalg(1+me_rows:1+me_rows+(icpu-1)*ncpu_rows:ncpu_rows))
1397 1909256 : froms=1+(ispinor-1)*nrowsLinalgMe/nspinor+(col-1)*nrowsLinalgMe+nrowsLinalgMeSum*ncolsColsRows
1398 1909256 : frome=froms-1+nrowsLinalgMe/nspinor
1399 1909256 : tos=1+nrowsLinalgMeSum/nspinor+(ispinor-1)*nrowsColsRows/nspinor+(col-1)*nrowsColsRows
1400 1909256 : toe=tos-1+nrowsLinalgMe/nspinor
1401 519291222 : bufferOrdered(:,tos:toe) = bufferMess(:,froms:frome)
1402 : end do
1403 : end do
1404 : end do
1405 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1406 : end if
1407 : #endif
1408 : case (STATE_COLSROWS)
1409 : ! We are going to STATE_LINALG so we are before all2all
1410 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1411 : if (xgTransposer%gpu_option == ABI_GPU_OPENMP) then
1412 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1413 : !$OMP& PRIVATE(nrowsLinalgMe,nrowsLinalgMeSum,toe,tos,frome,froms,irow) &
1414 : !$OMP& MAP(to:bufferMess) MAP(to:bufferOrdered) MAP(to:nrowsLinalg)
1415 : do col = 1, ncolsColsRows
1416 : do icpu = 0, ncpu_cols-1
1417 : do ispinor = 1, nspinor
1418 : nrowsLinalgMe = nrowsLinalg(1+me_rows+icpu*ncpu_rows)
1419 : nrowsLinalgMeSum = sum(nrowsLinalg(1+me_rows:1+me_rows+(icpu-1)*ncpu_rows:ncpu_rows))
1420 : froms=1+(ispinor-1)*nrowsLinalgMe/nspinor+(col-1)*nrowsLinalgMe+nrowsLinalgMeSum*ncolsColsRows
1421 : frome=froms-1+nrowsLinalgMe/nspinor
1422 : tos=1+nrowsLinalgMeSum/nspinor+(ispinor-1)*nrowsColsRows/nspinor+(col-1)*nrowsColsRows
1423 : toe=tos-1+nrowsLinalgMe/nspinor
1424 : do irow = 0, toe-tos
1425 : bufferMess(1,froms+irow) = bufferOrdered(1,tos+irow)
1426 : bufferMess(2,froms+irow) = bufferOrdered(2,tos+irow)
1427 : end do
1428 : end do
1429 : end do
1430 : end do
1431 : else
1432 : #endif
1433 : !$omp parallel do private(nrowsLinalgMe,nrowsLinalgMeSum,toe,tos,frome,froms), collapse(3)
1434 3490332 : do col = 1, ncolsColsRows
1435 7607026 : do icpu = 0, ncpu_cols-1
1436 11443888 : do ispinor = 1, nspinor
1437 5049448 : nrowsLinalgMe = nrowsLinalg(1+me_rows+icpu*ncpu_rows)
1438 9830780 : nrowsLinalgMeSum = sum(nrowsLinalg(1+me_rows:1+me_rows+(icpu-1)*ncpu_rows:ncpu_rows))
1439 5049448 : froms=1+(ispinor-1)*nrowsLinalgMe/nspinor+(col-1)*nrowsLinalgMe+nrowsLinalgMeSum*ncolsColsRows
1440 5049448 : frome=froms-1+nrowsLinalgMe/nspinor
1441 5049448 : tos=1+nrowsLinalgMeSum/nspinor+(ispinor-1)*nrowsColsRows/nspinor+(col-1)*nrowsColsRows
1442 5049448 : toe=tos-1+nrowsLinalgMe/nspinor
1443 1435687332 : bufferMess(:,froms:frome) = bufferOrdered(:,tos:toe)
1444 : end do
1445 : end do
1446 : end do
1447 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1448 : end if
1449 : #endif
1450 : end select
1451 :
1452 1652876 : if (xgTransposer%gpu_option /= ABI_GPU_DISABLED .and. xgTransposer%gpu_thread_limit /= 0) then
1453 0 : call xomp_set_num_threads(nthreads_bak)
1454 : end if
1455 :
1456 : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
1457 : ! if gpu enable restore OpenMP num threads to 1
1458 : if (xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
1459 : ! restore OMP_NUM_THREADS=1
1460 : call xomp_set_num_threads(1)
1461 : end if
1462 :
1463 : ! if gpu is enabled, transfer back data on GPU
1464 : if (xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
1465 : buffer_size = size(bufferOrdered) * dp
1466 : call gpu_data_prefetch_async(C_LOC(bufferOrdered), buffer_size)
1467 : call gpu_device_synchronize()
1468 : end if
1469 : #endif
1470 :
1471 1652876 : call timab(tim_reorganize,2,tsec)
1472 :
1473 1652876 : end subroutine xgTransposer_reorganizeData
1474 : !!***
1475 :
1476 : !!****f* m_xgTransposer/xgTransposer_getRank
1477 : !!
1478 : !! NAME
1479 : !! xgTransposer_getRank
1480 :
1481 0 : function xgTransposer_getRank(xgTransposer, comm) result(rank)
1482 : type(xgTransposer_t), intent(in ) :: xgTransposer
1483 : integer , intent(in ) :: comm
1484 : integer :: rank
1485 0 : if ( (comm > ubound(xgTransposer%mpiData,1)) .or. (comm < lbound(xgTransposer%mpiData,1)) ) then
1486 0 : ABI_ERROR("Value for communicator is wrong")
1487 : end if
1488 0 : rank = xgTransposer%mpiData(comm)%rank
1489 0 : end function xgTransposer_getRank
1490 : !!***
1491 :
1492 : !!****f* m_xgTransposer/xgTransposer_getComm
1493 : !!
1494 : !! NAME
1495 : !! xgTransposer_getComm
1496 :
1497 0 : function xgTransposer_getComm(xgTransposer, comm1) result(communicator)
1498 : type(xgTransposer_t), intent(in ) :: xgTransposer
1499 : integer , intent(in ) :: comm1
1500 : integer :: communicator
1501 0 : if ( (comm1 > ubound(xgTransposer%mpiData,1)) .or. (comm1 < lbound(xgTransposer%mpiData,1)) ) then
1502 0 : ABI_ERROR("Value for communicator is wrong")
1503 : end if
1504 0 : communicator = xgTransposer%mpiData(comm1)%comm
1505 0 : end function xgTransposer_getComm
1506 : !!***
1507 :
1508 : !!****f* m_xgTransposer/xgTransposer_free
1509 : !!
1510 : !! NAME
1511 : !! xgTransposer_free
1512 :
1513 189878 : subroutine xgTransposer_free(xgTransposer)
1514 :
1515 : type(xgTransposer_t), intent(inout) :: xgTransposer
1516 : double precision :: tsec(2)
1517 : integer :: i
1518 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD && !defined HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
1519 : real(dp), ABI_CONTIGUOUS pointer :: xgTransposer__buffer(:,:)
1520 : #endif
1521 :
1522 189878 : call timab(tim_free,1,tsec)
1523 : #ifdef HAVE_MPI
1524 189878 : if ( xgTransposer%type == TRANS_TYPE_CONSTRUCTED_NULL_COMM ) then
1525 10 : call mpi_comm_free(xgTransposer%mpiData(MPI_ROWS)%comm,i)
1526 10 : call mpi_comm_free(xgTransposer%mpiData(MPI_COLS)%comm,i)
1527 10 : call mpi_comm_free(xgTransposer%mpiData(MPI_2DCART)%comm,i)
1528 : end if
1529 : #else
1530 : ABI_UNUSED(i)
1531 : #endif
1532 :
1533 189878 : if ( allocated(xgTransposer%lookup) ) then
1534 189878 : ABI_FREE(xgTransposer%lookup)
1535 : end if
1536 :
1537 189878 : if ( associated(xgTransposer%nrowsLinalg) ) then
1538 189878 : ABI_FREE(xgTransposer%nrowsLinalg)
1539 : end if
1540 :
1541 189878 : if ( associated(xgTransposer%buffer) ) then
1542 151952 : if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
1543 : #if defined HAVE_GPU && defined HAVE_YAKL
1544 : ABI_FREE_MANAGED(xgTransposer%buffer)
1545 : #endif
1546 : else
1547 : if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
1548 : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
1549 : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
1550 : !$OMP TARGET EXIT DATA MAP(delete:xgTransposer%buffer)
1551 : #else
1552 : xgTransposer__buffer => xgTransposer%buffer
1553 : !$OMP TARGET EXIT DATA MAP(delete:xgTransposer__buffer)
1554 : #endif
1555 : #endif
1556 : end if
1557 151952 : ABI_FREE(xgTransposer%buffer)
1558 : end if
1559 : end if
1560 :
1561 189878 : call timab(tim_free,2,tsec)
1562 :
1563 189878 : end subroutine xgTransposer_free
1564 : !!***
1565 :
1566 0 : end module m_xgTransposer
1567 : !!***
|
