Branch data Line data Source code
1 : : !!****f* ABINIT/m_slice_task
2 : : !! NAME
3 : : !! m_slice_task
4 : : !!
5 : : !! FUNCTION
6 : : !! This module contains types and routines to implement scheduler and resource allocator
7 : : !! for slice tasks. Implements scheduler and resource allocator for slice tasks.
8 : : !! Implements logic for asynchronous slice memory avoiding race condition in read/write.
9 : : !!
10 : : !! NOTES
11 : : !! The logic for parallel slice treatment (a task) is the following.
12 : : !! Memory and workload are distributed using a 2D cartesian grid. Let's assume
13 : : !! for simplicity that we have four MPI processes in the spacecom communicator.
14 : : !! Matrix X is distributed along plane-waves at the beginning:
15 : : !!
16 : : !! bands
17 : : !! |-------------------|
18 : : !! | P0 |
19 : : !! | |
20 : : !! |-------------------|
21 : : !! | P1 |
22 : : !! | |
23 : : !! pw |-------------------|
24 : : !! | P2 |
25 : : !! | |
26 : : !! |-------------------|
27 : : !! | P3 |
28 : : !! | |
29 : : !! |-------------------|
30 : : !!
31 : : !! At the start, we use xgTransposer to MPI transpose the matrix X
32 : : !! achieving a custom layout for bandpp, and we end up with:
33 : : !!
34 : : !! bands
35 : : !! |-------|---|---|---|
36 : : !! | | | | |
37 : : !! | | | | |
38 : : !! | | | | |
39 : : !! | | | | |
40 : : !! | | | | |
41 : : !! pw | P0 |P1 |P2 |P3 |
42 : : !! | | | | |
43 : : !! | | | | |
44 : : !! | | | | |
45 : : !! | | | | |
46 : : !! | | | | |
47 : : !! |-------|---|---|---|
48 : : !!
49 : : !! From there, we can define slices acting on subgroup of processes.
50 : : !! For example, slice one can have process 0 and slice two the remaining
51 : : !! 1,2,3 processes. MPI transposing to Linalg representation using
52 : : !! the slice sub-communicators yields:
53 : : !!
54 : : !! bands
55 : : !! |-------|-----------|
56 : : !! | | |
57 : : !! | | P1 |
58 : : !! | | |
59 : : !! | |-----------|
60 : : !! | | |
61 : : !! pw | P0 | P2 |
62 : : !! | | |
63 : : !! | |-----------|
64 : : !! | | |
65 : : !! | | P3 |
66 : : !! | | |
67 : : !! |-------|-----------|
68 : : !!
69 : : !! At this point, parallel Rayleigh-Ritz is possible.
70 : : !
71 : : !!
72 : : !! COPYRIGHT
73 : : !! Copyright (C) 2018-2026 ABINIT group (IML)
74 : : !! This file is distributed under the terms of the
75 : : !! gnu general public license, see ~abinit/COPYING
76 : : !! or http://www.gnu.org/copyleft/gpl.txt .
77 : : !! for the initials of contributors, see ~abinit/doc/developers/contributors.txt .
78 : : !!
79 : : !! SOURCE
80 : :
81 : : #if defined HAVE_CONFIG_H
82 : : #include "config.h"
83 : : #endif
84 : :
85 : : #include "abi_common.h"
86 : :
87 : : ! nvtx related macro definition
88 : : #include "nvtx_macros.h"
89 : :
90 : : module m_slice_task
91 : :
92 : : use defs_basis
93 : : use defs_abitypes
94 : : use m_abicore
95 : : use m_errors
96 : : use m_time, only : timab
97 : : use m_sort, only: sort_dp
98 : :
99 : : use m_cgtools
100 : : use m_xg
101 : : use m_xgTransposer
102 : :
103 : : use m_xmpi
104 : : use m_xomp
105 : : #ifdef HAVE_OPENMP
106 : : use omp_lib
107 : : #endif
108 : :
109 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_YAKL)
110 : : use m_gpu_toolbox, only : CPU_DEVICE_ID, gpu_device_synchronize
111 : : #endif
112 : :
113 : : #if defined(HAVE_GPU_MARKERS)
114 : : use m_nvtx_data
115 : : #endif
116 : :
117 : : implicit none
118 : :
119 : : private
120 : :
121 : : ! Timers
122 : : !---------------------------------------------------
123 : : integer, parameter :: tim_swap = 1761
124 : : integer, parameter :: tim_RR_q = 1759
125 : : integer, parameter :: tim_barrier = 1764
126 : : integer, parameter :: tim_copy = 1765
127 : : integer, parameter :: tim_getAX_BX = 1754
128 : : integer, parameter :: tim_invovl = 1755
129 : :
130 : : ! Public 'matrixInfo' datatype
131 : : !-------------------------------------------------
132 : : type, public :: matrixInfo_t
133 : :
134 : : integer :: comm_rows ! xmpi_comm_self ...
135 : : integer :: comm_cols ! same as spacecom
136 : : integer :: spacecom ! same as comm_cols
137 : : integer :: neigenpairs ! total number of bands (=number of eigenpairs)
138 : : integer :: total_spacedim ! total number of plane-waves
139 : : integer :: spacedim ! nb of plane-waves per process in linalg representation
140 : : integer :: space ! real or complex eigenvectors
141 : : integer :: gpu_kokkos_nthrd
142 : : integer :: gpu_thread_limit
143 : : integer :: gpu_option ! enable GPU
144 : : integer :: paral_kgb ! enable parallel (k-points, G basis, bands)
145 : : integer :: me_g0
146 : : integer :: me_g0_fft
147 : : integer :: nspinor
148 : :
149 : : end type matrixInfo_t
150 : :
151 : : ! Public 'activeTask' datatype for active slice in use (me=current MPI process)
152 : : !-------------------------------------------------
153 : : type, public :: activeTask_t
154 : :
155 : : ! MPI-related information for active task
156 : :
157 : : integer :: me_g0 ! process contains G(0,0,0)
158 : : integer :: me_g0_fft ! process contains G(0,0,0) for fft
159 : : integer :: me_nproc ! number of processes reserved to slice in use
160 : : integer :: me_comm ! (sub-)communicator reserved to slice in use
161 : : integer :: me_comm_rows ! (sub-)communicator for rows in Transposer
162 : : integer :: me_comm_cols ! (sub-)communicator for cols in Transposer
163 : : integer :: me_id_slice ! identifier from 1 to nslice of slice task in use
164 : : integer :: me_neigenpairs ! total number of eigenpairs of slice in use
165 : : integer :: me_bandpp ! number of distributed bands per process for slice in use
166 : : integer :: me_lowb ! lower slice interval bound
167 : : integer :: me_uppb ! upper slice interval bound
168 : : integer :: me_ndeg ! polynomial filter degree for slice in use
169 : : integer :: me_fcol_async ! first column of task load in async memory
170 : :
171 : : ! flags for slice location in spectrum
172 : : logical :: is_lowpass
173 : : logical :: is_last
174 : :
175 : : !integer, allocatable :: me_cols_X(:) ! columns of shared memory used in task
176 : : !integer, allocatable :: me_cols_Xext(:) ! columns of asynchronous memory used in task
177 : : integer, allocatable :: me_mask_Xext(:) ! converged columns
178 : :
179 : : ! MPI column and row distribution for active task
180 : : integer, allocatable :: me_ncolsColsRows(:) ! ncol of colsrows representation of asyncMem
181 : : integer, allocatable :: me_nrowsLinalg(:) ! nrow of linalg representation of asyncMem
182 : :
183 : : ! Memory address used for reads/writes
184 : : type(xgBlock_t) :: me_Xext ! eigenvector memory in use by active slice
185 : : type(xgBlock_t) :: me_eigen ! eigenvalue memory in use by active slice
186 : : type(xgBlock_t) :: me_resid ! residual memory in use by active slice
187 : :
188 : : ! Memory space used for computation
189 : : !type(chebfi_t) :: chebfi ! should not use chebfi objects
190 : :
191 : : end type activeTask_t
192 : :
193 : : ! Public 'taskScheduler' datatype for asynchronous slice treatment
194 : : ! it is basically for asynchronous treatment of the 'async' memory
195 : : !-------------------------------------------------
196 : : type, public :: taskScheduler_t
197 : :
198 : : integer :: ntasks
199 : : integer :: nprocs
200 : : integer :: tot_load ! size of async memory buffer (columns)
201 : :
202 : : ! Fixed quantities (global)
203 : : integer, allocatable :: load_per_task(:) ! number of eigenpairs per slice
204 : : integer, allocatable :: nproc_per_task(:) ! number of used processes per slice
205 : :
206 : : ! next task info (updated on every task)
207 : : integer :: next_task_id ! number from 1 to ntasks (1-base)
208 : : integer :: next_load ! number of eigenpairs
209 : : integer, allocatable :: next_lookup_proc(:) ! which slice each MPI rank serves (0-base)
210 : :
211 : : end type taskScheduler_t
212 : :
213 : : ! Public 'asyncMemory' datatype for slice input/output without race condition
214 : : ! is basically created by blocks of number of asynchronous tasks
215 : : !-------------------------------------------------
216 : : type, public :: asyncMemory_t
217 : :
218 : : integer :: nslice
219 : : integer :: neigenpairs_ext ! total number of extended columns
220 : : integer :: paral_kgb
221 : :
222 : : ! Memory buffers
223 : : type(xg_t) :: X_ext ! eigenvector memory used by all slices for I/O
224 : : type(xg_t) :: eigen_ext ! eigenvalue memory used by all slices for I/O
225 : : type(xg_t) :: resid_ext ! residual memory used by all slices for I/O
226 : : type(xgTransposer_t) :: xgTransposerXext ! transposer datastructure for eigenvectors
227 : :
228 : : ! Pointers
229 : : type(xgBlock_t) :: XextLinalg
230 : :
231 : : ! Arrays related to data distribution (all-slice phase)
232 : : integer, allocatable :: lookup_cols_Xext(:) ! which slice each column serves in extented memory
233 : :
234 : : ! Flags for MPI distribution state
235 : : logical :: use_linalg = .false. ! use linalg representation
236 : : logical :: use_colsrows = .false. ! use colsrows representation
237 : : logical :: has_transposer
238 : :
239 : : end type asyncMemory_t
240 : :
241 : : ! Public methods
242 : : !-------------------------------------------------
243 : : public :: init_matrixInfo ! wrapper for various xgBlock parameters
244 : : public :: slice_task_allocAsyncMemory ! allocates async memory buffer
245 : : public :: slice_task_copyToAsyncMemory ! fills async memory colwise (deep copy)
246 : : public :: slice_task_freeAsyncMemory ! deallocates async buffer
247 : : public :: slice_task_initSchedule ! compute 'process-to-slices' distribution
248 : : public :: slice_task_freeSchedule ! deallocate all MPI distribution info
249 : : public :: slice_task_printSchedule ! print a table of schedule
250 : : public :: slice_task_initNextTask ! compute MPI distribution for next task
251 : : public :: slice_task_freeActiveTask ! deallocate memory for active task
252 : : public :: slice_task_enableAsync ! apply MPI distribution for next task
253 : : public :: slice_task_runActiveTask ! execute Subspace Iteration for active vectors
254 : : !public :: mark_active_task
255 : : !public :: allocate_active_task
256 : : !public :: execute_active_task
257 : : !public :: mask_active_task
258 : : !public :: compress_extended_memory
259 : :
260 : : CONTAINS
261 : : !=====================================================================
262 : : !!***
263 : :
264 : : !!****f* m_slice_task/init_matrixInfo
265 : : !! NAME
266 : : !! init_matrixInfo
267 : : !!
268 : : !! SOURCE
269 : :
270 : 18 : subroutine init_matrixInfo(matrixInfo, comm_rows, comm_cols, spacecom, neigenpairs,&
271 : : total_spacedim, spacedim, space, gpu_kokkos_nthrd, gpu_thread_limit, gpu_option,&
272 : : paral_kgb, me_g0, me_g0_fft, nspinor)
273 : :
274 : : implicit none
275 : :
276 : : type(matrixInfo_t), intent(inout) :: matrixInfo
277 : : integer, intent(in) :: comm_rows, comm_cols, spacecom, neigenpairs, total_spacedim
278 : : integer, intent(in) :: spacedim, space, gpu_kokkos_nthrd, gpu_thread_limit, gpu_option
279 : : integer, intent(in) :: paral_kgb, me_g0, me_g0_fft, nspinor
280 : :
281 : 18 : matrixInfo%comm_rows = comm_rows
282 : 18 : matrixInfo%comm_cols = comm_cols
283 : 18 : matrixInfo%spacecom = spacecom
284 : 18 : matrixInfo%neigenpairs = neigenpairs
285 : 18 : matrixInfo%total_spacedim = total_spacedim
286 : 18 : matrixInfo%spacedim = spacedim
287 : 18 : matrixInfo%space = space
288 : 18 : matrixInfo%gpu_kokkos_nthrd = gpu_kokkos_nthrd
289 : 18 : matrixInfo%gpu_thread_limit = gpu_thread_limit
290 : 18 : matrixInfo%gpu_option = gpu_option
291 : 18 : matrixInfo%paral_kgb = paral_kgb
292 : 18 : matrixInfo%me_g0 = me_g0
293 : 18 : matrixInfo%me_g0_fft = me_g0_fft
294 : 18 : matrixInfo%nspinor = nspinor
295 : :
296 : 18 : end subroutine init_matrixInfo
297 : : !!***
298 : :
299 : : !----------------------------------------------------------------------
300 : :
301 : : !!****f* m_slice_task/slice_task_allocAsyncMemory
302 : : !! NAME
303 : : !! slice_task_allocAsyncMemory
304 : : !!
305 : : !! FUNCTION
306 : : !! Allocate async memory buffers and distribute them according
307 : : !! to slice logic. Essentially allocates memory work%Xext
308 : : !! favoring data overlap over communication overlap.
309 : : !!
310 : : !! SOURCE
311 : :
312 : 18 : subroutine slice_task_allocAsyncMemory(work, minfo, ncol_ext)
313 : :
314 : : implicit none
315 : : type(asyncMemory_t), intent(inout) :: work
316 : : type(matrixInfo_t), intent(in) :: minfo
317 : : integer, intent(in) :: ncol_ext
318 : : integer :: gpu_option
319 : :
320 : 18 : work%paral_kgb = minfo%paral_kgb
321 : 18 : work%neigenpairs_ext = ncol_ext
322 : 18 : work%use_linalg = .true. ! for sanity
323 : 18 : work%use_colsrows = .false. ! for sanity
324 : 18 : work%has_transposer = .false.
325 : 18 : gpu_option = minfo%gpu_option
326 : :
327 : 18 : call slice_task_freeAsyncMemory(work)
328 : :
329 [ + - - + : 54 : ABI_MALLOC_IFNOT(work%lookup_cols_Xext, (ncol_ext))
+ - + - +
- - + - +
- + ]
330 : :
331 : : ! Allocate extended space in linalg representation
332 : : call xg_init(work%X_ext, minfo%space, minfo%spacedim, work%neigenpairs_ext, &
333 : 18 : minfo%spacecom, me_g0=minfo%me_g0, gpu_option=gpu_option)
334 : :
335 : : ! row-array so that columns of tasks can be pointed more easily
336 : 18 : call xg_init(work%resid_ext, SPACE_R, 1, work%neigenpairs_ext, gpu_option=gpu_option)
337 : 18 : call xg_init(work%eigen_ext, SPACE_R, 1, work%neigenpairs_ext, gpu_option=gpu_option)
338 : :
339 : 18 : end subroutine slice_task_allocAsyncMemory
340 : : !!***
341 : :
342 : : !----------------------------------------------------------------------
343 : :
344 : : !!****f* m_slice_task/slice_task_freeAsyncMemory
345 : : !! NAME
346 : : !! slice_task_freeAsyncMemory
347 : : !!
348 : : !! SOURCE
349 : :
350 : 36 : subroutine slice_task_freeAsyncMemory(work)
351 : :
352 : : implicit none
353 : : type(asyncMemory_t), intent(inout) :: work
354 : :
355 : 36 : call xg_free(work%X_ext)
356 : 36 : call xg_free(work%eigen_ext)
357 : 36 : call xg_free(work%resid_ext)
358 [ + + ]: 36 : ABI_SFREE(work%lookup_cols_Xext)
359 [ + + ]: 36 : if (work%has_transposer) then
360 : 16 : call xgTransposer_free(work%xgTransposerXext)
361 : : end if
362 : :
363 : 36 : end subroutine slice_task_freeAsyncMemory
364 : : !!***
365 : :
366 : : !----------------------------------------------------------------------
367 : :
368 : : !!****f* m_slice_task/slice_task_copyToAsyncMemory
369 : : !! NAME
370 : : !! slice_task_copyToAsyncMemory
371 : : !!
372 : : !! FUNCTION
373 : : !! Initialize async memory content (for all tasks)
374 : : !! IML dev note:
375 : : !! current version initializes using sketching of wanted size. Might also need
376 : : !! to test if choosing directly random vectors is better.
377 : : !!
378 : : !! SOURCE
379 : :
380 [ - + ]: 18 : subroutine slice_task_copyToAsyncMemory(work, X0, minfo, mapper, ncols_per_task)
381 : :
382 : : implicit none
383 : :
384 : : !Arguments ------------------------------------
385 : : type(asyncMemory_t), intent(inout) :: work
386 : : type(matrixInfo_t), intent(in) :: minfo
387 : : type(xgBlock_t), intent(in) :: X0
388 : : logical, pointer, intent(in) :: mapper(:,:)
389 : : integer, intent(in) :: ncols_per_task(:)
390 : :
391 : : !Local variables-------------------------------
392 : : integer :: nslice, islice, fcol, fcol_ext_prev, fcol_ext
393 : : integer :: nrows, ncols, k_sketch, me_ncompl, nselect, fcol_compl
394 : : integer :: space, gpu_option, spacecom, fcol_ext_sketch
395 : 18 : integer, allocatable :: nrand_per_task(:)
396 : 18 : integer, allocatable :: ncompl_per_task(:)
397 : : ! typed
398 : : type(xg_t) :: X0_compl, X_sketch
399 : : type(xgBlock_t) :: col_in, col_out, Xext_last
400 : :
401 : : ! *********************************************************************
402 : :
403 : : ! Sanity check
404 [ - + ]: 18 : if ((.not. work%use_linalg) .or. work%use_colsrows) then
405 : 0 : ABI_ERROR("not in linalg representation")
406 : : end if
407 : :
408 : 18 : nrows = rows(X0)
409 : 18 : ncols = cols(X0)
410 : 18 : space = minfo%space
411 : 18 : gpu_option = minfo%gpu_option
412 : 18 : spacecom = minfo%spacecom
413 : 18 : nslice = size(mapper, dim=2)
414 : :
415 : 18 : work%XextLinalg = work%X_ext%self
416 : :
417 [ - + + - : 54 : ABI_MALLOC(nrand_per_task, (nslice))
+ - + - -
+ - + ]
418 [ - + + - : 36 : ABI_MALLOC(ncompl_per_task, (nslice))
- + - + ]
419 : :
420 : : ! Copy X to XextLinalg to achieve contiguous column blocks
421 : 18 : fcol_ext_prev = 1
422 : 18 : fcol_ext = 0
423 [ + + ]: 54 : do islice=1, nslice
424 : 36 : nselect = 0
425 [ + - ]: 36 : do fcol=1, ncols
426 [ + - ]: 36 : if (.not. mapper(fcol, islice)) then
427 : : exit
428 : : end if
429 : 0 : nselect = nselect + 1
430 : 0 : fcol_ext = fcol_ext + nselect
431 : 0 : call xgBlock_setBlock(X0, col_in, nrows, 1, fcol=fcol)
432 : 0 : call xgBlock_setBlock(work%XextLinalg, col_out, nrows, 1, fcol=fcol_ext)
433 : :
434 : : ! Reminder: xgBlock_copy is always on CPU except if both blocks are on GPU
435 : 36 : call xgBlock_copy(col_in, col_out)
436 : : end do
437 : 36 : nrand_per_task(islice) = ncols_per_task(islice) - nselect
438 : 36 : ncompl_per_task(islice) = ncols - nselect
439 : 36 : write(std_out,*) 'selected', nselect, 'out of', ncols_per_task(islice), 'then rand is', &
440 : 72 : nrand_per_task(islice)
441 : 36 : fcol_ext = ncols_per_task(islice) ! jump to end of slice
442 [ + + ]: 2154 : work%lookup_cols_Xext(fcol_ext_prev:fcol_ext) = islice
443 : 54 : fcol_ext_prev = fcol_ext
444 : : end do
445 : :
446 : : ! recover the remaining columns not mapped to the slice and sketch them
447 [ + + ]: 54 : do islice=1, nslice
448 : 36 : k_sketch = nrand_per_task(islice) ! size of sketch
449 : 36 : me_ncompl = ncompl_per_task(islice) ! size of complement in X0
450 [ - + ]: 36 : if (k_sketch==0 .or. me_ncompl==0) then
451 : : exit
452 : : end if
453 : 36 : call xg_init(X0_compl, space, nrows, me_ncompl, spacecom, gpu_option=gpu_option)
454 : 36 : fcol_compl = 0
455 : 36 : fcol_ext_sketch = 1
456 [ + + ]: 6948 : do fcol=1, ncols
457 [ + - ]: 6948 : if (.not. mapper(fcol, islice)) then ! for remaining dimensions not in slice
458 : 6912 : fcol_compl = fcol_compl + 1
459 : 6912 : call xgBlock_setBlock(X0, col_in, nrows, 1, fcol=fcol)
460 : 6912 : call xgBlock_setBlock(X0_compl%self, col_out, nrows, 1, fcol=fcol_compl)
461 : 6912 : call xgBlock_copy(col_in, col_out)
462 : : end if
463 : : end do
464 : : ! Y = X * Omega where Omega sketch matrix to capture all directions at once (linalg distribution)
465 : 36 : call xg_init(X_sketch, space, nrows, k_sketch, spacecom, gpu_option=gpu_option)
466 : 36 : call xgBlock_randomSketching(X0_compl%self, X_sketch%self, k_sketch)
467 : 36 : call xgBlock_setBlock(work%XextLinalg, Xext_last, nrows, k_sketch, fcol=fcol_ext_sketch)
468 : 36 : call xgBlock_copy(X_sketch%self, Xext_last)
469 : 36 : call xg_free(X0_compl)
470 : 36 : call xg_free(X_sketch)
471 : 54 : fcol_ext_sketch = ncols_per_task(islice) ! jump to end of slice
472 : : end do
473 : :
474 [ - + ]: 18 : ABI_FREE(nrand_per_task)
475 [ - + ]: 18 : ABI_FREE(ncompl_per_task)
476 : :
477 [ + - ]: 18 : end subroutine slice_task_copyToAsyncMemory
478 : : !!***
479 : :
480 : : !----------------------------------------------------------------------
481 : :
482 : : !!****f* m_slice_task/slice_task_initSchedule
483 : : !! NAME
484 : : !! slice_task_initSchedule
485 : : !!
486 : : !! FUNCTION
487 : : !! Initialization of scheduler object from 'nproc' available resources
488 : : !! using execution options given by 'paral_kgb' and 'paral_task'.
489 : : !! Essentially allows to compute and apply (via wrapper to xgTransposer)
490 : : !! an intermediate level of MPI distribution along tasks, on top of paral_kgb level.
491 : : !! - If enable_paral then we should first divide processes to slices
492 : : !! - If disable_paral then we should use all available processes for every slice (no division)
493 : : !!
494 : : !! OUTPUT
495 : : !! scheduler%load_per_task
496 : : !! scheduler%nproc_per_task
497 : : !! scheduler%next_loookup_proc
498 : : !!
499 : : !! SOURCE
500 : :
501 [ - + ]: 18 : subroutine slice_task_initSchedule(scheduler, ntasks, nprocs, load, paral_kgb, paral_task)
502 : :
503 : : implicit none
504 : :
505 : : type(taskScheduler_t), intent(inout) :: scheduler
506 : : integer, intent(in) :: ntasks, nprocs, paral_kgb, paral_task
507 : : integer, intent(in) :: load(:)
508 : 18 : integer, allocatable :: weights(:)
509 : :
510 [ + + ]: 54 : scheduler%tot_load = sum(load)
511 : 18 : scheduler%ntasks = ntasks
512 : 18 : scheduler%nprocs = nprocs
513 : 18 : scheduler%next_task_id = 1 ! fixme this is not true if paral slices
514 : 18 : scheduler%next_load = load(1) ! fixme this is not true if paral slices
515 : :
516 : 18 : call slice_task_allocSchedule(scheduler)
517 [ + + ]: 54 : scheduler%load_per_task(:) = load(:)
518 : :
519 [ + + + + ]: 18 : if (paral_kgb==0 .or. paral_task==0) then
520 : : ! do not divide available resources to slices at all
521 [ + + ]: 30 : scheduler%nproc_per_task = nprocs !! use all MPI, can be 1
522 [ + + ]: 44 : scheduler%next_lookup_proc = 0
523 : : else
524 : :
525 : : ! computation to divide resources
526 [ - + + - : 24 : ABI_MALLOC_IFNOT(weights, (ntasks))
+ - + - -
+ - + ]
527 [ + + ]: 24 : weights = 1
528 : : ! weights = slice%poly_degrees ! IML works less well
529 [ - + - - : 8 : call fair_allocation(ntasks, load, weights, nprocs, scheduler%nproc_per_task)
- - - + ]
530 : 8 : call assign_tasks_to_processes(scheduler%nproc_per_task, scheduler%next_lookup_proc)
531 [ + - ]: 8 : ABI_SFREE(weights)
532 : :
533 : : end if
534 : :
535 : 18 : end subroutine slice_task_initSchedule
536 : : !!***
537 : :
538 : : !----------------------------------------------------------------------
539 : :
540 : : !!****f* m_slice_task/slice_task_allocSchedule
541 : : !! NAME
542 : : !! slice_task_allocSchedule
543 : : !!
544 : : !! FUNCTION
545 : : !! Constructor for scheduler object
546 : :
547 : 18 : subroutine slice_task_allocSchedule(scheduler)
548 : :
549 : : implicit none
550 : :
551 : : type(taskScheduler_t), intent(inout) :: scheduler
552 : :
553 : 18 : call slice_task_freeSchedule(scheduler)
554 [ + - - + : 54 : ABI_MALLOC_IFNOT(scheduler%load_per_task, (scheduler%ntasks))
+ - + - +
- - + - +
- + ]
555 [ + - - + : 54 : ABI_MALLOC_IFNOT(scheduler%nproc_per_task, (scheduler%ntasks))
+ - + - +
- - + - +
- + ]
556 [ + - - + : 54 : ABI_MALLOC_IFNOT(scheduler%next_lookup_proc, (scheduler%nprocs))
+ - + - +
- - + - +
- + ]
557 : :
558 : 18 : end subroutine slice_task_allocSchedule
559 : : !!***
560 : :
561 : : !----------------------------------------------------------------------
562 : :
563 : : !!****f* m_slice_task/slice_task_freeSchedule
564 : : !! NAME
565 : : !! slice_task_freeSchedule
566 : : !!
567 : : !! FUNCTION
568 : : !! Destructor for scheduler object
569 : :
570 : 36 : subroutine slice_task_freeSchedule(scheduler)
571 : :
572 : : implicit none
573 : :
574 : : type(taskScheduler_t), intent(inout) :: scheduler
575 : :
576 [ + + ]: 36 : ABI_SFREE(scheduler%load_per_task)
577 [ + + ]: 36 : ABI_SFREE(scheduler%nproc_per_task)
578 [ + + ]: 36 : ABI_SFREE(scheduler%next_lookup_proc)
579 : :
580 : 36 : end subroutine slice_task_freeSchedule
581 : : !!***
582 : :
583 : : !----------------------------------------------------------------------
584 : :
585 : : !!****f* m_slice_task/slice_task_printSchedule
586 : : !! NAME
587 : : !! slice_task_printSchedule
588 : : !!
589 : :
590 : 18 : subroutine slice_task_printSchedule(scheduler, wout)
591 : :
592 : : implicit none
593 : :
594 : : type(taskScheduler_t), intent(inout) :: scheduler
595 : : integer, intent(in) :: wout
596 : :
597 : 18 : write(wout,*) '###### Schedule info ######'
598 : 18 : write(wout,*) 'Process per task =', scheduler%nproc_per_task
599 : 18 : write(wout,*) 'Task by process =', scheduler%next_lookup_proc
600 : 18 : flush(wout)
601 : :
602 : 18 : end subroutine slice_task_printSchedule
603 : : !!***
604 : :
605 : : !----------------------------------------------------------------------
606 : :
607 : : !!****f* m_slice_task/slice_task_initNextTask
608 : : !! NAME
609 : : !! slice_task_initNextTask
610 : : !!
611 : : !! FUNCTION
612 : : !! Apply logic for task execution and prepare MPI distribution
613 : : !! either a slice has some MPI processes or a slice has ALL MPI processes.
614 : : !! Mark allocated portion as actively used by setting me_* variables.
615 : : !! My process only marks resources its assigned slice has reserved.
616 : : !!
617 : : !! OUTPUT
618 : : !! Active 'task' object with a valid MPI distribution
619 : : !!
620 : : !! SOURCE
621 : :
622 : 108 : subroutine slice_task_initNextTask(scheduler, task, minfo)
623 : :
624 : : implicit none
625 : :
626 : : type(taskScheduler_t), intent(inout) :: scheduler
627 : : type(activeTask_t), intent(inout) :: task
628 : : type(matrixInfo_t), intent(in) :: minfo
629 : :
630 : : integer :: global_comm, comm_rows, comm_cols, ierr
631 : : integer :: my_rank, my_rank_sub, my_task, sanity_check
632 : :
633 : : ! *********************************************************************
634 : :
635 : 18 : global_comm = minfo%spacecom
636 : 18 : my_rank = xmpi_comm_rank(global_comm)
637 : 18 : my_task = scheduler%next_lookup_proc(my_rank + 1) + 1
638 : 18 : scheduler%next_task_id = my_task
639 : :
640 : 18 : task%me_id_slice = my_task
641 : 18 : task%me_neigenpairs = scheduler%load_per_task(my_task)
642 : 18 : task%me_fcol_async = 1
643 [ + + ]: 18 : if (my_task>1) then
644 [ + + ]: 10 : task%me_fcol_async = sum(scheduler%load_per_task(1:my_task-1))+1
645 : : end if
646 : 18 : task%me_nproc = scheduler%nproc_per_task(my_task)
647 : :
648 : : ! Split global comm into disjoint sub-comms, only procs with the same color (my_task) communicate
649 : 18 : comm_rows = minfo%comm_rows
650 : 18 : comm_cols = minfo%comm_cols
651 : 18 : call xmpi_comm_split(global_comm, my_task, my_rank, task%me_comm, ierr)
652 [ - + ]: 18 : if ( ierr /= xmpi_success ) then
653 : 0 : ABI_ERROR("Error while creating slice spacecom subcommunicator")
654 : : end if
655 : 18 : call xmpi_comm_split(comm_rows, my_task, my_rank, task%me_comm_rows, ierr)
656 [ - + ]: 18 : if ( ierr /= xmpi_success ) then
657 : 0 : ABI_ERROR("Error while creating slice row subcommunicator")
658 : : end if
659 : 18 : call xmpi_comm_split(comm_cols, my_task, my_rank, task%me_comm_cols, ierr)
660 [ - + ]: 18 : if ( ierr /= xmpi_success ) then
661 : 0 : ABI_ERROR("Error while creating slice col subcommunicator")
662 : : end if
663 : :
664 : : ! process waits for others to create their subcommunicators before using its own
665 : 18 : call xmpi_barrier(global_comm)
666 : :
667 : : ! If using more than MPI processes, compute column and row distributions across processes
668 : 18 : call slice_task_freeActiveTask(task)
669 [ + - - + : 54 : ABI_MALLOC_IFNOT(task%me_ncolsColsRows, (scheduler%nprocs))
+ - + - +
- - + - +
- + ]
670 [ + - - + : 54 : ABI_MALLOC_IFNOT(task%me_nrowsLinalg, (scheduler%nprocs))
+ - + - +
- - + - +
- + ]
671 [ + + ]: 18 : if (task%me_nproc>1) then
672 : 15 : call distribute_vectors(task%me_neigenpairs, task%me_nproc, task%me_ncolsColsRows)
673 : 15 : call distribute_vectors(minfo%total_spacedim, task%me_nproc, task%me_nrowsLinalg)
674 : 15 : my_rank_sub = xmpi_comm_rank(task%me_comm)
675 : 15 : task%me_bandpp = task%me_ncolsColsRows(my_rank_sub + 1)
676 : : else
677 : 3 : task%me_bandpp = task%me_neigenpairs
678 [ + + ]: 9 : task%me_nrowsLinalg = minfo%total_spacedim
679 : : end if
680 : :
681 : : ! Concatenate task%me_bandpp into collective task%me_ncolsColsRows (global)
682 : 18 : call xmpi_allgather(task%me_bandpp, task%me_ncolsColsRows, global_comm, ierr)
683 [ - + ]: 18 : if ( ierr /= xmpi_success ) then
684 : 0 : ABI_ERROR("Error while gathering number of columns in colsrows for all slices")
685 : : end if
686 : :
687 : 18 : sanity_check = task%me_bandpp
688 : 18 : call xmpi_sum(sanity_check, global_comm, ierr)
689 [ + + ]: 18 : if (sanity_check == scheduler%tot_load) then
690 : 8 : write(std_out,*) 'parallel execution of tasks, sum of bands across procs=', sanity_check
691 : : else
692 : 10 : write(std_out,*) 'sequential execution of tasks, sum of bands across procs=', sanity_check
693 : : end if
694 : :
695 : : ! at some point fixme
696 : : ! ABI_SFREE(task%me_ncolsColsRows)
697 : :
698 : 18 : write(std_out,*) '@task current rank has slice=', task%me_id_slice
699 : 18 : write(std_out,*) '@task active process=', task%me_nproc
700 : 18 : write(std_out,*) '@task bands per active process=', task%me_bandpp
701 : 18 : flush(std_out)
702 : :
703 : 18 : end subroutine slice_task_initNextTask
704 : : !!***
705 : :
706 : : !----------------------------------------------------------------------
707 : :
708 : : !!****f* m_slice_task/slice_task_freeActiveTask
709 : : !! NAME
710 : : !! slice_task_freeActiveTask
711 : :
712 : 36 : subroutine slice_task_freeActiveTask(task)
713 : :
714 : : implicit none
715 : : type(activeTask_t), intent(inout) :: task
716 : :
717 : : !ABI_SFREE(task%me_cols_X)
718 : : !ABI_SFREE(task%me_cols_Xext)
719 [ + + ]: 36 : ABI_SFREE(task%me_ncolsColsRows)
720 [ + + ]: 36 : ABI_SFREE(task%me_nrowsLinalg)
721 : :
722 : 36 : end subroutine slice_task_freeActiveTask
723 : : !!***
724 : :
725 : : !----------------------------------------------------------------------
726 : :
727 : : !!****f* m_slice_task/slice_task_enableAsync
728 : : !! NAME
729 : : !! slice_task_enableAsync
730 : : !!
731 : : !! FUNCTION
732 : : !! Enable asynchronous memory treatment by MPI transposing global data
733 : : !! across all available MPI processes. This transposition allows for a
734 : : !! slice to not see others. It serves as a transition from global
735 : : !! communicator to slice communicator. After the transposition each
736 : : !! process contains the correct bandpp corresponding to the slice so
737 : : !! that no additional communication has to be performed in order to
738 : : !! bring band slices to processes.
739 : : !!
740 : : !! INPUT
741 : : !! asyncMemory buffer in linalg distribution
742 : : !! target colsrows distribution in task%me_ncolsColsRows of global data
743 : : !!
744 : : !! OUTPUT
745 : : !! asyncMemory buffer in colsrows distribution, stored in task%me_Xext
746 : : !! allocated if multiple MPI ranks or just pointer if MPI disabled.
747 : : !! Same for task%me_eigen, task%me_resid.
748 : : !!
749 : : !! SOURCE
750 : :
751 : 18 : subroutine slice_task_enableAsync(work, task, minfo)
752 : :
753 : : implicit none
754 : :
755 : : type(asyncMemory_t), intent(inout) :: work
756 : : type(activeTask_t), intent(inout) :: task
757 : : type(matrixInfo_t), intent(in) :: minfo
758 : :
759 : : integer :: nrows, ncols, neigen, fcol
760 : :
761 : : ! *********************************************************************
762 : :
763 : : ! todo add timers and make nvtx markers consistent
764 [ + + ]: 18 : if (minfo%paral_kgb==1) then
765 : :
766 : 16 : write(std_out,*) 'using distro as target='
767 : 16 : write(std_out,*) task%me_ncolsColsRows
768 : 16 : flush(std_out)
769 : :
770 : : call xgTransposer_constructor(work%xgTransposerXext, work%XextLinalg, task%me_Xext,&
771 : : minfo%nspinor, STATE_LINALG, TRANS_ALL2ALL, minfo%comm_rows, minfo%comm_cols, &
772 : : 0, 0, minfo%me_g0_fft, gpu_option=minfo%gpu_option, &
773 : : gpu_thread_limit=minfo%gpu_thread_limit, custom_ncolsColsRows=.true.,&
774 : 16 : ncolsColsRows_sub=task%me_ncolsColsRows)
775 : :
776 : 16 : work%xgTransposerXext%gpu_kokkos_nthrd = minfo%gpu_kokkos_nthrd
777 : :
778 : : ABI_NVTX_START_RANGE(NVTX_SLICE_TRANSPOSE)
779 : 16 : call xgTransposer_transpose(work%xgTransposerXext, STATE_COLSROWS)
780 : : ABI_NVTX_END_RANGE()
781 : :
782 : 16 : work%use_colsrows = .true.
783 : 16 : work%use_linalg = .false.
784 : 16 : work%has_transposer = .true.
785 : : else
786 : 2 : nrows = rows(work%XextLinalg)
787 : 2 : ncols = cols(work%XextLinalg)
788 : 2 : call xgBlock_setBlock(work%XextLinalg, task%me_Xext, nrows, ncols)
789 : : end if
790 : :
791 : : ! Every process has all eigen and resid of slice (not distributed)
792 : 18 : neigen = task%me_neigenpairs
793 : 18 : fcol = task%me_fcol_async
794 : 18 : call xgBlock_setBlock(work%eigen_ext%self, task%me_eigen, 1, neigen, fcol=fcol)
795 : 18 : call xgBlock_setBlock(work%resid_ext%self, task%me_resid, 1, neigen, fcol=fcol)
796 : 18 : call xgBlock_reshape(task%me_eigen, neigen, 1)
797 : 18 : call xgBlock_reshape(task%me_resid, neigen, 1)
798 : :
799 : 18 : end subroutine slice_task_enableAsync
800 : : !!***
801 : :
802 : : !----------------------------------------------------------------------
803 : :
804 : : !!****f* m_slice_task/slice_task_runActiveTask
805 : : !! NAME
806 : : !! slice_task_runActiveTask
807 : : !!
808 : : !! FUNCTION
809 : : !! Execute Subspace iteration using active task memory.
810 : : !! Runs entirely independently of other tasks thanks to asynchronous memory.
811 : : !! Allocate internal memory of slice based on intermediate 'chebfi' structure.
812 : : !!
813 : : !! INPUT
814 : : !! task%me_Xext: initial guess of size (total_spacedim, task%me_bandpp)
815 : : !! task%me_eigen: empty array of size (task%me_neigenpairs, 1)
816 : : !! task%me_resid: empty array of size (task%me_neigenpairs, 1)
817 : : !!
818 : : !! OUTPUT
819 : : !! eigen=converged eigenvalue array of size (neigenpairs,1)
820 : : !! residu=slice residual array of size (neigenpairs,1)
821 : : !! /IML\ true output work%XextLinalg= guess/converged eigenvectors for all slices
822 : : !!
823 : : !! SOURCE
824 : :
825 : 0 : subroutine slice_task_runActiveTask(work, task)
826 : :
827 : : implicit none
828 : : type(asyncMemory_t), intent(inout) :: work
829 : : type(activeTask_t), intent(inout) :: task
830 : : !type(slice_t), intent(inout) :: slice ! should not use slice objects at all!!!
831 : : !type(chebfi_t), intent(inout) :: chebfi ! should not use chebfi objects at all !!!
832 : :
833 : : integer :: nbdbuf, oracle, num_proc
834 : : real(dp) :: oracle_factor, oracle_min_occ
835 : :
836 : : ! fix dummy
837 : 0 : nbdbuf = 0
838 : 0 : oracle = 0
839 : 0 : num_proc = 0
840 : 0 : oracle_factor = 1.0
841 : 0 : oracle_min_occ = 1.0
842 : 0 : write(std_out,*) 'dummy arg=', work%nslice
843 : 0 : write(std_out,*) 'dummu arg=', task%me_g0
844 : 0 : flush(std_out)
845 : :
846 : : ! can also set slice params?? needs slice object. Maybe do a different called initActiveTask
847 : : !
848 : : ! ! Get parameters of active task
849 : : ! neigenpairs = task%me_neigenpairs
850 : : ! bandpp = task%me_bandpp
851 : : ! ndeg_filter = task%me_ndeg
852 : : ! comm = task%me_comm
853 : : ! comm_rows = task%me_comm_rows
854 : : ! comm_cols = task%me_comm_cols
855 : :
856 : :
857 : : ! if (slice%me_id_slice==1) then
858 : : ! is_lowpass = .true. ! [lambda_minus,lambda_plus) to be diminished
859 : : ! lambda_minus = slice%upp_bounds(task%me_id_slice)
860 : : ! lambda_plus = slice%maxeig_global
861 : : ! else
862 : : ! is_lowpass = .false. ! [lambda_minus,lambda_plus) to be amplified
863 : : ! lambda_minus = slice%low_bounds(task%me_id_slice)
864 : : ! lambda_plus = slice%upp_bounds(task%me_id_slice)
865 : : ! end if
866 : : ! ! deactivate chebfi oracle
867 : : ! oracle = 0
868 : : ! nbdbuf = 0
869 : : ! oracle_factor = 1.d0
870 : : ! oracle_min_occ = 0.d0
871 : : !
872 : : ! num_proc = xmpi_comm_size(comm)
873 : : ! ABI_MALLOC_IFNOT(nrowsLinalg,(num_proc))
874 : : ! nrowsLinalg_ptr => nrowsLinalg
875 : : ! nrowsLinalg = task%me_nrowsLinalg_slice
876 : : !
877 : : ! write(std_out,*) "Allocating slice space..", slice%total_spacedim, neigenpairs
878 : : ! write(std_out,*) "bands per process=", bandpp
879 : : ! flush(std_out)
880 : : !
881 : : ! ! Initialize chebfi object in MPI Colsrows distribution
882 : : ! call chebfi_init(chebfi,neigenpairs,slice%total_spacedim,slice%tolerance,slice%ecut,slice%paral_kgb,bandpp,&
883 : : ! ndeg_filter,nbdbuf,slice%space,1,comm,task%me_g0,task%me_g0_fft,slice%paw,comm_rows,comm_cols,&
884 : : ! oracle,oracle_factor,oracle_min_occ,slice%gpu_option,gpu_kokkos_nthrd=slice%gpu_kokkos_nthrd,&
885 : : ! gpu_thread_limit=slice%gpu_thread_limit,from_linalg=.false.)
886 : : !
887 : : ! ! Recover actively used array
888 : : ! X0_active = task%me_Xext
889 : : !
890 : : ! ! fixme move this inside chebfi_runSI?
891 : : ! task%chebfi%xXColsRows = X0_active
892 : : !
893 : : ! !call chebfi_runSlice(chebfi, X0_active, getAX_BX, getBm1X, eigen_active, residu_active, nspinor,&
894 : : ! ! slice%mineig_global, slice%maxeig_global, lambda_minus, lambda_plus, is_lowpass, slice%neigenpairs,&
895 : : ! ! nrowsLinalg_ptr)
896 : : !
897 : : ! ! todo give k=m+p where p is oversample
898 : : ! ! residual will be converged for m values. Give m as input
899 : : ! k_conv = slice%neigenpairs - 20 ! hardcoded assuming offset 20 fixme
900 : : !
901 : : ! !call chebfi_runSubspaceIteration(chebfi, X0_active, getAX_BX, getBm1X, eigen_active, residu_active, &
902 : : ! ! nspinor, slice%mineig_global, slice%maxeig_global, lambda_minus, lambda_plus, is_lowpass, &
903 : : ! ! k_conv, nrowsLinalg_ptr)
904 : : !
905 : : ! call chebfi_runSubspaceIterationDummy(task%chebfi, X0_active, getAX_BX, getBm1X, eigen_active, residu_active, &
906 : : ! nspinor, slice%mineig_global, slice%maxeig_global, lambda_minus, lambda_plus, is_lowpass, &
907 : : ! k_conv, nrowsLinalg_ptr)
908 : : !
909 : : ! ! why?
910 : : ! if (slice%gpu_option == ABI_GPU_OPENMP) then
911 : : ! call xgBlock_copy_to_gpu(eigen_active)
912 : : ! call xgBlock_copy_to_gpu(residu_active)
913 : : ! end if
914 : :
915 : 0 : end subroutine slice_task_runActiveTask
916 : : !!***
917 : :
918 : : !----------------------------------------------------------------------
919 : :
920 : : !!****f* m_slice_task/mask_active_task
921 : : !! NAME
922 : : !! mask_active_task
923 : : !!
924 : : !! FUNCTION
925 : : !! Mask converged solutions in active task (asynchronous).
926 : : !! Create mask for converged solutions in active task
927 : : !! that will be used to combine all active tasks to extended memory
928 : : !!
929 : : !! SOURCE
930 : :
931 : : subroutine mask_active_task(task, tol)
932 : :
933 : : implicit none
934 : :
935 : : ! Arguments ------------------------------------
936 : : type(activeTask_t), intent(inout) :: task
937 : : real(dp), intent(inout) :: tol
938 : :
939 : : ! Local variables-------------------------------
940 : : integer :: n_active, iband
941 : : logical :: selected
942 : : real(dp) :: theta, res
943 : : real(dp), pointer :: thetas_conv(:,:) => null()
944 : : real(dp), pointer :: residu_conv(:,:) => null()
945 : :
946 : : ! *********************************************************************
947 : :
948 : : n_active = rows(task%me_eigen)
949 : : call xgBlock_reverseMap(task%me_eigen, thetas_conv, rows=n_active, cols=1)
950 : : call xgBlock_reverseMap(task%me_resid, residu_conv, rows=n_active, cols=1)
951 : :
952 : : ! Hard acceptance criterion so that slices do not overlap
953 : : do iband=1, n_active
954 : : res = residu_conv(iband, 1)
955 : : theta = thetas_conv(iband, 1)
956 : : selected = .false.
957 : : if (task%is_lowpass) then
958 : : selected = (res < tol .and. theta < task%me_uppb)
959 : : else if (task%is_last) then
960 : : selected = (res < tol .and. theta > task%me_lowb)
961 : : else
962 : : selected = (res < tol .and. theta < task%me_uppb .and. theta > task%me_lowb)
963 : : end if
964 : : if (selected) then
965 : : task%me_mask_Xext(iband) = 1
966 : : end if
967 : : end do
968 : :
969 : : end subroutine mask_active_task
970 : : !!***
971 : :
972 : : !----------------------------------------------------------------------
973 : :
974 : : !!****f* m_slice_task/compress_extended_memory
975 : : !! NAME
976 : : !! compress_extended_memory
977 : : !!
978 : : !! FUNCTION
979 : : !! Copy data from asyncMemory to spectrum I/O memory (shared)
980 : : !! Copy masked async memory to spectrum memory
981 : : !!
982 : : !! SOURCE
983 : :
984 : : !subroutine compress_extended_memory(task, work, X0, eigen, resid)
985 : : !
986 : : ! implicit none
987 : : !
988 : : ! ! Arguments ------------------------------------
989 : : ! type(activeTask_t), intent(inout) :: task
990 : : ! type(asyncMemory_t), intent(inout) :: work
991 : : ! type(xgBlock_t), intent(inout) :: X0
992 : : ! type(xgBlock_t), intent(inout) :: eigen
993 : : ! type(xgBlock_t), intent(inout) :: resid
994 : : !
995 : : ! ! Local variables-------------------------------
996 : : ! integer :: nrows, ncols, fcol
997 : : !
998 : : ! ! *********************************************************************
999 : : !
1000 : : ! if (work%paral_kgb==1) then
1001 : : !
1002 : : ! ! Sanity check
1003 : : ! if ((.not. work%use_linalg) .or. work%use_colsrows) ) then
1004 : : ! ABI_ERROR("not in linalg")
1005 : : ! end if
1006 : : !
1007 : : ! ! Recover pointer task%me_Xext into memory work%XextLinalg
1008 : : ! call xmpi_barrier(slice%spacecom)
1009 : : ! ABI_NVTX_START_RANGE(NVTX_SLICE_TRANSPOSE)
1010 : : ! call xgTransposer_transpose(work%xgTransposerXext, STATE_LINALG)
1011 : : ! ABI_NVTX_END_RANGE()
1012 : : ! else
1013 : : !
1014 : : ! nrows =
1015 : : ! ncols =
1016 : : ! fcol =
1017 : : ! xgBlock_setBlock(work%XextLinalg, task%me_Xext, rows(work%XextLinalg), cols(
1018 : : !
1019 : : ! end if
1020 : : !
1021 : : ! ! Detect missing or extra eigenvalues
1022 : : ! if (tot_ncols_kept < slice%neigenpairs) then
1023 : : ! ABI_WARNING("Not enough converged eigenvalues in slice")
1024 : : ! else if (tot_ncols_kept > slice%neigenpairs) then
1025 : : ! ABI_WARNING("Too many converged eigenvalues kept. Decrease tolfilter or nstep_mixed.")
1026 : : ! end if
1027 : : !
1028 : : ! ! Copy from async memory to regular memory
1029 : : ! do islice=1,slice%nslice
1030 : : ! fcol = slice%fcol_in_X(islice)
1031 : : ! fcol_ext = slice%fcol_in_Xext(islice)
1032 : : ! neigenpairs_slice = slice%neigenpairs_per_slice(islice)
1033 : : ! write(std_out,*) 'block copy from fcol, ncols=', fcol_ext, neigenpairs_slice
1034 : : ! write(std_out,*) 'block copy to fcol, ncols=', fcol, neigenpairs_slice
1035 : : ! ! Blocks to copy from
1036 : : ! call xgBlock_setBlock(slice%XextLinalg, X_kept, rows=slice%spacedim, cols=neigenpairs_slice, fcol=fcol_ext)
1037 : : ! call xgBlock_setBlock(eigen_ext%self, eigen_kept, rows=1, cols=neigenpairs_slice, fcol=fcol_ext)
1038 : : ! call xgBlock_setBlock(resid_ext%self, resid_kept, rows=1, cols=neigenpairs_slice, fcol=fcol_ext)
1039 : : ! ! Blocks to copy to
1040 : : ! call xgBlock_setBlock(X0, X0_out, rows=slice%spacedim, cols=neigenpairs_slice, fcol=fcol)
1041 : : ! call xgBlock_setBlock(eigen, eigen_out, rows=1, cols=neigenpairs_slice, fcol=fcol)
1042 : : ! call xgBlock_setBlock(resid, resid_out, rows=1, cols=neigenpairs_slice, fcol=fcol)
1043 : : ! ! copy
1044 : : ! call xgBlock_copy(X_kept, X0_out)
1045 : : ! call xgBlock_copy(eigen_kept, eigen_out)
1046 : : ! call xgBlock_copy(resid_kept, resid_out)
1047 : : ! end do
1048 : : !
1049 : : ! ! Recover dimensions
1050 : : ! call xgBlock_reshape(eigen, slice%neigenpairs, 1)
1051 : : ! call xgBlock_reshape(resid, slice%neigenpairs, 1)
1052 : : !
1053 : : ! ! Free memory
1054 : : ! call xg_free(eigen_ext)
1055 : : ! call xg_free(resid_ext)
1056 : : !
1057 : : !end subroutine compress_extended_memory
1058 : : !!***
1059 : :
1060 : : !----------------------------------------------------------------------
1061 : :
1062 : : !!****f* m_slice_task/assign_tasks_to_processes
1063 : : !! NAME
1064 : : !! assign_tasks_to_processes
1065 : : !!
1066 : : !! FUNCTION
1067 : : !! Perform the inverse of the allocation operation, assigning
1068 : : !! processes to slices based on the allocation array.
1069 : : !!
1070 : : !! SOURCE
1071 : :
1072 [ - + - + ]: 8 : subroutine assign_tasks_to_processes(allocations, processes)
1073 : :
1074 : : implicit none
1075 : : integer, intent(in) :: allocations(:)
1076 : : integer, intent(out) :: processes(:)
1077 : : integer :: i, j
1078 : :
1079 : : ! *********************************************************************
1080 : :
1081 : 8 : j = 1
1082 [ + + ]: 24 : do i = 1, size(allocations)
1083 [ + + ]: 48 : processes(j:j + allocations(i) - 1) = i - 1
1084 : 24 : j = j + allocations(i)
1085 : : end do
1086 : :
1087 : 8 : end subroutine assign_tasks_to_processes
1088 : : !!***
1089 : :
1090 : : !----------------------------------------------------------------------
1091 : :
1092 : : !!****f* m_slice_task/distribute_vectors
1093 : : !! NAME
1094 : : !! distribute_vectors
1095 : : !!
1096 : : !! FUNCTION
1097 : : !! Distribute m vectors as uniformly as possible across n processes.
1098 : : !! Assumptions:
1099 : : !! -The remainder should be distributed evenly to the first n-1 processes.
1100 : : !! -The last process should always get fewer vectors.
1101 : : !! -The sum of the allocations should be exactly m.
1102 : : !!
1103 : : !! SOURCE
1104 : :
1105 : 30 : subroutine distribute_vectors(m, n, allocation)
1106 : :
1107 : : implicit none
1108 : : integer, intent(in) :: m, n
1109 : : integer, intent(out) :: allocation(n)
1110 : : integer :: i, base, remainder
1111 : :
1112 : : ! *********************************************************************
1113 : :
1114 : 30 : base = m / n
1115 : 30 : remainder = m - base * n
1116 [ + + ]: 128 : allocation = base
1117 [ + + ]: 30 : if (remainder /= 0) then
1118 [ + + ]: 94 : do i = 1, n-1
1119 [ + + ]: 94 : if (remainder > 0) then
1120 : 58 : allocation(i) = allocation(i) + 1
1121 : 58 : remainder = remainder - 1
1122 : : end if
1123 : : end do
1124 : : end if
1125 : :
1126 : 30 : end subroutine distribute_vectors
1127 : : !!***
1128 : :
1129 : : !----------------------------------------------------------------------
1130 : :
1131 : : !!****f* m_slice_task/fair_allocation
1132 : : !! NAME
1133 : : !! fair_allocation
1134 : : !!
1135 : : !! FUNCTION
1136 : : !! Solve integer optimization problem under constraint:
1137 : : !!
1138 : : !! min_{x_1,..,x_s} max_{1,..,s} f_i(x_i)
1139 : : !! subject to: x_1 + .. + x_s = p
1140 : : !! x_i integers
1141 : : !!
1142 : : !! with objective cost function f_i(x)=m_i*w_i/x.
1143 : : !! The solution x_i is the amount of resource allocated to the i-th task.
1144 : : !! The algorithm uses binary search for integer rounding.
1145 : : !! Note that this is better than greedy but not optimal.
1146 : : !! Exhaustive search is too expensive (=(p+1)^n combinations).
1147 : : !! Feature: ensures the total allocation is exactly equal to p while
1148 : : !! minimizing the allocation imbalance.
1149 : : !!
1150 : : !! INPUTS
1151 : : !! arrays m and w (length n), and integer p
1152 : : !! m can be the group size, w can be another measure or need (weight)
1153 : : !! p in the number of total resources
1154 : : !!
1155 : : !! OUTPUT
1156 : : !! integer array x of size s such that sum(x) = p and max(m_i*n_i/x_i) is minimized
1157 : : !!
1158 : : !! SOURCE
1159 : :
1160 : 8 : subroutine fair_allocation(n, m, w, p, x)
1161 : :
1162 : : implicit none
1163 : :
1164 : : ! Arguments
1165 : : integer, intent(in) :: n ! Number of groups (slices)
1166 : : integer, intent(in) :: m(n) ! Array: size of each group
1167 : : integer, intent(in) :: p ! Total resources to allocate
1168 : : integer, intent(in) :: w(n) ! Weight of each group
1169 : : integer, intent(out) :: x(n) ! Array: allocated resources per group
1170 : :
1171 : : ! Local variables
1172 : : integer :: i, total_allocated
1173 : : real(dp) :: total_work, lower, upper, mid, multiplier
1174 : 8 : integer :: allocation(n)
1175 : :
1176 : : ! *********************************************************************
1177 : :
1178 [ + + ]: 24 : total_work = dot_product(m, w)
1179 : :
1180 : : ! Binary search for the optimal multiplier: divide search interval in half
1181 : 8 : lower = 0.d0
1182 : 8 : upper = real(p)
1183 [ + + ]: 24 : do while (upper - lower > 1.d0)
1184 : 16 : mid = (lower + upper) / 2.d0
1185 [ + + ]: 48 : do i = 1, n
1186 : 48 : allocation(i) = int((real(w(i)) * real(m(i)) / total_work) * mid + 0.d5)
1187 : : end do
1188 [ + + ]: 48 : total_allocated = sum(allocation(1:n))
1189 : :
1190 : : ! Adjust binary search bounds
1191 [ - + ]: 24 : if (total_allocated > p) then
1192 : : upper = mid
1193 : : else
1194 : 16 : lower = mid
1195 : : end if
1196 : : end do
1197 : :
1198 : : ! Fair rounding: tasks with heavier workload get extra resource units
1199 : :
1200 : : ! Final allocation after binary search converges
1201 : 8 : multiplier = (lower + upper) / 2.d0
1202 [ + + ]: 24 : do i = 1, n
1203 : 24 : allocation(i) = int((real(w(i)) * real(m(i)) / total_work) * multiplier + 0.d5)
1204 : : end do
1205 : :
1206 : : ! Adjust total allocation to exactly match p
1207 [ + + ]: 24 : total_allocated = sum(allocation(1:n))
1208 : :
1209 [ + - ]: 8 : if (total_allocated < p) then
1210 [ + + ]: 20 : do while (total_allocated < p)
1211 : : ! Add one resource to the group closest to its ideal allocation
1212 : 12 : call adjust_allocation(n, m, w, allocation, total_work, p, total_allocated)
1213 [ + + ]: 44 : total_allocated = sum(allocation(1:n))
1214 : : end do
1215 [ # # ]: 0 : else if (total_allocated > p) then ! FIXME error infinite loop
1216 [ # # ]: 0 : do while (total_allocated > p)
1217 : : ! Remove one resource from the over-allocated group
1218 : 0 : call reduce_allocation(n, m, w, allocation, total_work, p, total_allocated)
1219 [ # # ]: 0 : total_allocated = sum(allocation(1:n))
1220 : : end do
1221 : : end if
1222 : :
1223 : : ! Assign the final allocation to the output variable
1224 [ + + ]: 24 : x = allocation
1225 : :
1226 : : !write(std_out,'(a)') 'Memory allocation info:'
1227 : : !do i=1,n
1228 : : ! if (x(i) .ne. 0) then
1229 : : ! write(std_out,'(a,i4,i6,i5)') '#task #workload #allocated resources', i, w(i)*m(i)/x(i), x(i)
1230 : : ! else
1231 : : ! write(std_out,'(a,i4,i4,i4,i4)') 'Allocation error: x(i)= w(i)= m(i)= for i=', x(i), w(i), m(i), i
1232 : : ! end if
1233 : : !end do
1234 : :
1235 : 8 : end subroutine fair_allocation
1236 : : !!***
1237 : :
1238 : : !----------------------------------------------------------------------
1239 : :
1240 : : !!****f* m_slice_task/adjust_allocation
1241 : : !! NAME
1242 : : !! adjust_allocation
1243 : : !!
1244 : : !! FUNCTION
1245 : : !! Adjust allocation by adding resources to the group closest to its ideal allocation
1246 : :
1247 : 12 : subroutine adjust_allocation(n, m, w, allocation, total_weight, p, total_allocated)
1248 : :
1249 : : implicit none
1250 : :
1251 : : integer, intent(in) :: n, m(n)
1252 : : integer, intent(in) :: w(n), p
1253 : : real(dp), intent(in) :: total_weight
1254 : : integer, intent(inout) :: allocation(n)
1255 : : integer, intent(inout) :: total_allocated
1256 : :
1257 : : integer :: i, closest_group
1258 : : real(dp) :: max_diff, diff, ideal
1259 : :
1260 : : ! *********************************************************************
1261 : :
1262 : : ! Find the group with the largest difference between current and ideal allocation
1263 : 12 : max_diff = -1.0_dp
1264 : 12 : closest_group = -1
1265 : :
1266 [ + + ]: 36 : do i = 1, n
1267 : 24 : ideal = (real(w(i), dp) * real(m(i), dp) * real(p, dp)) / total_weight
1268 : 24 : diff = abs(real(allocation(i), dp) - ideal)
1269 [ + + ]: 36 : if (diff > max_diff) then
1270 : 20 : max_diff = diff
1271 : 20 : closest_group = i
1272 : : end if
1273 : : end do
1274 : :
1275 : : ! Only allocate if total does not exceed limit
1276 [ + - + - ]: 12 : if (total_allocated < p .and. closest_group > 0) then
1277 : 12 : allocation(closest_group) = allocation(closest_group) + 1
1278 : 12 : total_allocated = total_allocated + 1
1279 : : end if
1280 : :
1281 : 12 : end subroutine adjust_allocation
1282 : : !!***
1283 : :
1284 : : !----------------------------------------------------------------------
1285 : :
1286 : : !!****f* m_slice_task/reduce_allocation
1287 : : !! NAME
1288 : : !! reduce_allocation
1289 : : !!
1290 : : !! FUNCTION
1291 : : !! Reduce allocation by removing resources from the over-allocated group
1292 : :
1293 : 0 : subroutine reduce_allocation(n, m, w, allocation, total_weight, p, total_allocated)
1294 : :
1295 : : implicit none
1296 : :
1297 : : integer, intent(in) :: n, m(n)
1298 : : integer, intent(in) :: w(n), p
1299 : : real(dp), intent(in) :: total_weight
1300 : : integer, intent(inout) :: allocation(n)
1301 : : integer, intent(inout) :: total_allocated
1302 : :
1303 : : integer :: i, target_group
1304 : : real(dp) :: max_diff, expected, diff
1305 : : integer :: reduce_by
1306 : :
1307 : : ! *********************************************************************
1308 : :
1309 : : ! Find the group with the largest *positive* over-allocation
1310 : 0 : max_diff = -1.0_dp
1311 : 0 : target_group = -1
1312 : :
1313 [ # # ]: 0 : do i = 1, n
1314 : 0 : expected = (real(w(i), dp) * real(m(i), dp) * real(p, dp)) / total_weight
1315 : 0 : diff = real(allocation(i), dp) - expected
1316 [ # # # # : 0 : if (diff > max_diff .and. diff > 0.0_dp .and. allocation(i) > 0) then
# # ]
1317 : 0 : max_diff = diff
1318 : 0 : target_group = i
1319 : : end if
1320 : : end do
1321 : :
1322 : : ! If we found an over-allocated group, reduce its allocation
1323 [ # # # # ]: 0 : if (target_group > 0 .and. max_diff > 0.0_dp) then
1324 : 0 : reduce_by = min(1, allocation(target_group)) ! Only reduce if it's > 0
1325 : 0 : reduce_by = min(1, allocation(target_group))
1326 : 0 : allocation(target_group) = allocation(target_group) - reduce_by
1327 : 0 : total_allocated = total_allocated - reduce_by
1328 : : end if
1329 : :
1330 : 0 : end subroutine reduce_allocation
1331 : : !!***
1332 : :
1333 [ # # # # : 0 : end module m_slice_task
# # # # #
# # # # #
# # # # ]
1334 : : !!***
|
