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 : !!***
|