LCOV - code coverage report
Current view: top level - src/48_diago - m_slice_task.F90 (source / functions) Coverage Total Hit
Test: coverage.info Lines: 86.0 % 279 240
Test Date: 2026-09-21 19:39:32 Functions: 60.0 % 25 15

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

Generated by: LCOV version 2.3-1