LCOV - code coverage report
Current view: top level - src/dbt/tas - dbt_tas_global.F (source / functions) Hit Total Coverage
Test: CP2K Regtests (git:e7e05ae) Lines: 137 137 100.0 %
Date: 2024-04-18 06:59:28 Functions: 28 28 100.0 %

          Line data    Source code
       1             : !--------------------------------------------------------------------------------------------------!
       2             : !   CP2K: A general program to perform molecular dynamics simulations                              !
       3             : !   Copyright 2000-2024 CP2K developers group <https://cp2k.org>                                   !
       4             : !                                                                                                  !
       5             : !   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
       6             : !--------------------------------------------------------------------------------------------------!
       7             : 
       8             : ! **************************************************************************************************
       9             : !> \brief Global data (distribution and block sizes) for tall-and-skinny matrices
      10             : !>        For very sparse matrices with one very large dimension, storing array data of the same
      11             : !>        size as the matrix dimensions may require too much memory and we need to compute them on
      12             : !>        the fly for a given row or column. Hence global array data such as distribution and block
      13             : !>        sizes are specified as function objects, leaving up to the caller how to efficiently store
      14             : !>        global data.
      15             : !> \author Patrick Seewald
      16             : ! **************************************************************************************************
      17             : MODULE dbt_tas_global
      18             :    USE kinds,                           ONLY: dp,&
      19             :                                               int_8
      20             :    USE util,                            ONLY: sort
      21             : #include "../../base/base_uses.f90"
      22             : 
      23             :    IMPLICIT NONE
      24             :    PRIVATE
      25             : 
      26             :    CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'dbt_tas_global'
      27             : 
      28             :    PUBLIC :: &
      29             :       dbt_tas_blk_size_arb, &
      30             :       dbt_tas_blk_size_repl, &
      31             :       dbt_tas_blk_size_one, &
      32             :       dbt_tas_dist_arb, &
      33             :       dbt_tas_dist_arb_default, &
      34             :       dbt_tas_dist_cyclic, &
      35             :       dbt_tas_dist_repl, &
      36             :       dbt_tas_distribution, &
      37             :       dbt_tas_rowcol_data, &
      38             :       dbt_tas_default_distvec
      39             : 
      40             : ! **************************************************************************************************
      41             : !> \brief abstract type for distribution vectors along one dimension
      42             : !> \var nprowcol number of process rows / columns
      43             : !> \var nmrowcol number of matrix rows / columns
      44             : !> \var dist     map matrix rows/cols to distribution rows/cols
      45             : !> \var rowcols  map distribution rows/cols to matrix rows/cols
      46             : ! **************************************************************************************************
      47             :    TYPE, ABSTRACT :: dbt_tas_distribution
      48             :       INTEGER                           :: nprowcol = -1
      49             :       INTEGER(KIND=int_8)               :: nmrowcol = -1
      50             :    CONTAINS
      51             :       PROCEDURE(rowcol_dist), deferred  :: dist
      52             :       PROCEDURE(dist_rowcols), deferred :: rowcols
      53             :    END TYPE
      54             : 
      55             : ! **************************************************************************************************
      56             : !> \brief type for cyclic (round robin) distribution:
      57             : !>        - may not be load balanced for arbitrary block sizes
      58             : !>        - memory efficient for large dimensions
      59             : ! **************************************************************************************************
      60             :    TYPE, EXTENDS(dbt_tas_distribution) :: dbt_tas_dist_cyclic
      61             :       INTEGER   :: split_size = -1
      62             :    CONTAINS
      63             :       PROCEDURE :: dist => cyclic_dist
      64             :       PROCEDURE :: rowcols => cyclic_rowcols
      65             :    END TYPE
      66             : 
      67             : ! **************************************************************************************************
      68             : !> \brief type for arbitrary distributions
      69             : !>        - stored as an array
      70             : !>        - not memory efficient for large dimensions
      71             : ! **************************************************************************************************
      72             :    TYPE, EXTENDS(dbt_tas_distribution) :: dbt_tas_dist_arb
      73             :       INTEGER, DIMENSION(:), ALLOCATABLE :: dist_vec
      74             :    CONTAINS
      75             :       PROCEDURE                          :: dist => arb_dist
      76             :       PROCEDURE                          :: rowcols => arb_rowcols
      77             :    END TYPE
      78             : 
      79             : ! **************************************************************************************************
      80             : !> \brief type for replicated distribution
      81             : !>        - a submatrix distribution replicated on all process groups
      82             : !>        - memory efficient for large dimensions
      83             : ! **************************************************************************************************
      84             :    TYPE, EXTENDS(dbt_tas_distribution) :: dbt_tas_dist_repl
      85             :       INTEGER, DIMENSION(:), ALLOCATABLE :: dist_vec
      86             :       INTEGER                           :: nmrowcol_local = -1
      87             :       INTEGER                           :: n_repl = -1
      88             :       INTEGER                           :: dist_size = -1
      89             :    CONTAINS
      90             :       PROCEDURE                         :: dist => repl_dist
      91             :       PROCEDURE                         :: rowcols => repl_rowcols
      92             :    END TYPE
      93             : 
      94             : ! **************************************************************************************************
      95             : !> \brief abstract type for integer data (e.g. block sizes) along one dimension
      96             : !> \var nmrowcol     number of matrix rows / columns (blocks)
      97             : !> \var nfullrowcol  number of matrix rows / columns (elements)
      98             : !> \var data         integer data for each block row / col
      99             : ! **************************************************************************************************
     100             :    TYPE, ABSTRACT :: dbt_tas_rowcol_data
     101             :       INTEGER(KIND=int_8)              :: nmrowcol = -1
     102             :       INTEGER(KIND=int_8)              :: nfullrowcol = -1
     103             :    CONTAINS
     104             :       PROCEDURE(rowcol_data), deferred :: DATA
     105             :    END TYPE
     106             : 
     107             : ! **************************************************************************************************
     108             : !> \brief type for arbitrary block sizes
     109             : !>        - stored as an array
     110             : !>        - not memory efficient for large dimensions
     111             : ! **************************************************************************************************
     112             :    TYPE, EXTENDS(dbt_tas_rowcol_data) :: dbt_tas_blk_size_arb
     113             :       INTEGER, DIMENSION(:), ALLOCATABLE :: blk_size_vec
     114             :    CONTAINS
     115             :       PROCEDURE                          :: DATA => blk_size_arb
     116             :    END TYPE
     117             : 
     118             : ! **************************************************************************************************
     119             : !> \brief type for replicated block sizes
     120             : !>        - submatrix block sizes replicated on all process groups
     121             : !>        - memory efficient for large dimensions
     122             : ! **************************************************************************************************
     123             :    TYPE, EXTENDS(dbt_tas_rowcol_data) :: dbt_tas_blk_size_repl
     124             :       INTEGER, DIMENSION(:), ALLOCATABLE :: blk_size_vec
     125             :       INTEGER                            :: nmrowcol_local = -1
     126             :    CONTAINS
     127             :       PROCEDURE                          :: DATA => blk_size_repl
     128             :    END TYPE
     129             : 
     130             : ! **************************************************************************************************
     131             : !> \brief type for blocks of size one
     132             : !>        - memory efficient for large dimensions
     133             : ! **************************************************************************************************
     134             :    TYPE, EXTENDS(dbt_tas_rowcol_data) :: dbt_tas_blk_size_one
     135             :    CONTAINS
     136             :       PROCEDURE :: DATA => blk_size_one
     137             :    END TYPE
     138             : 
     139             :    ABSTRACT INTERFACE
     140             : 
     141             : ! **************************************************************************************************
     142             : !> \brief map matrix rows/cols to distribution rows/cols
     143             : !> \param t ...
     144             : !> \param rowcol ...
     145             : !> \return ...
     146             : ! **************************************************************************************************
     147             :       FUNCTION rowcol_dist(t, rowcol)
     148             :          IMPORT :: dbt_tas_distribution, int_8
     149             :          CLASS(dbt_tas_distribution), INTENT(IN) :: t
     150             :          INTEGER(KIND=int_8), INTENT(IN) :: rowcol
     151             :          INTEGER :: rowcol_dist
     152             :       END FUNCTION
     153             : 
     154             : ! **************************************************************************************************
     155             : !> \brief map distribution rows/cols to matrix rows/cols
     156             : !> \param t ...
     157             : !> \param dist ...
     158             : !> \return ...
     159             : ! **************************************************************************************************
     160             :       FUNCTION dist_rowcols(t, dist)
     161             :          IMPORT :: dbt_tas_distribution, int_8
     162             :          CLASS(dbt_tas_distribution), INTENT(IN) :: t
     163             :          INTEGER, INTENT(IN) :: dist
     164             :          INTEGER(KIND=int_8), DIMENSION(:), ALLOCATABLE :: dist_rowcols
     165             :       END FUNCTION
     166             : 
     167             : ! **************************************************************************************************
     168             : !> \brief integer data for each block row / col
     169             : !> \param t ...
     170             : !> \param rowcol ...
     171             : !> \return ...
     172             : ! **************************************************************************************************
     173             :       FUNCTION rowcol_data(t, rowcol)
     174             :          IMPORT :: dbt_tas_rowcol_data, int_8
     175             :          CLASS(dbt_tas_rowcol_data), INTENT(IN) :: t
     176             :          INTEGER(KIND=int_8), INTENT(IN) :: rowcol
     177             :          INTEGER :: rowcol_data
     178             :       END FUNCTION
     179             : 
     180             :    END INTERFACE
     181             : 
     182             :    INTERFACE dbt_tas_dist_cyclic
     183             :       MODULE PROCEDURE new_block_tas_dist_cyclic
     184             :    END INTERFACE
     185             : 
     186             :    INTERFACE dbt_tas_dist_arb
     187             :       MODULE PROCEDURE new_block_tas_dist_arb
     188             :    END INTERFACE
     189             : 
     190             :    INTERFACE dbt_tas_dist_repl
     191             :       MODULE PROCEDURE new_block_tas_dist_repl
     192             :    END INTERFACE
     193             : 
     194             :    INTERFACE dbt_tas_blk_size_arb
     195             :       MODULE PROCEDURE new_block_tas_blk_size_arb
     196             :    END INTERFACE
     197             : 
     198             :    INTERFACE dbt_tas_blk_size_repl
     199             :       MODULE PROCEDURE new_block_tas_blk_size_repl
     200             :    END INTERFACE
     201             : 
     202             :    INTERFACE dbt_tas_blk_size_one
     203             :       MODULE PROCEDURE new_block_tas_blk_size_one
     204             :    END INTERFACE
     205             : 
     206             : CONTAINS
     207             : 
     208             : ! **************************************************************************************************
     209             : !> \brief ...
     210             : !> \param t ...
     211             : !> \param rowcol ...
     212             : !> \return ...
     213             : !> \author Patrick Seewald
     214             : ! **************************************************************************************************
     215     2694791 :    FUNCTION blk_size_arb(t, rowcol)
     216             :       CLASS(dbt_tas_blk_size_arb), INTENT(IN) :: t
     217             :       INTEGER(KIND=int_8), INTENT(IN) :: rowcol
     218             :       INTEGER :: blk_size_arb
     219     2694791 :       blk_size_arb = t%blk_size_vec(rowcol)
     220     2694791 :    END FUNCTION
     221             : 
     222             : ! **************************************************************************************************
     223             : !> \brief ...
     224             : !> \param t ...
     225             : !> \param rowcol ...
     226             : !> \return ...
     227             : !> \author Patrick Seewald
     228             : ! **************************************************************************************************
     229     2651679 :    FUNCTION blk_size_repl(t, rowcol)
     230             :       CLASS(dbt_tas_blk_size_repl), INTENT(IN) :: t
     231             :       INTEGER(KIND=int_8), INTENT(IN) :: rowcol
     232             :       INTEGER :: blk_size_repl
     233             :       INTEGER :: igroup
     234             :       INTEGER :: rowcol_local
     235             : 
     236     2651679 :       igroup = INT((rowcol - 1_int_8)/t%nmrowcol_local)
     237     2651679 :       rowcol_local = INT(MOD(rowcol - 1_int_8, INT(t%nmrowcol_local, KIND=int_8))) + 1
     238     2651679 :       blk_size_repl = t%blk_size_vec(rowcol_local)
     239             : 
     240     2651679 :    END FUNCTION
     241             : 
     242             : ! **************************************************************************************************
     243             : !> \brief ...
     244             : !> \param t ...
     245             : !> \param rowcol ...
     246             : !> \return ...
     247             : !> \author Patrick Seewald
     248             : ! **************************************************************************************************
     249     5381221 :    FUNCTION blk_size_one(t, rowcol)
     250             :       CLASS(dbt_tas_blk_size_one), INTENT(IN) :: t
     251             :       INTEGER(KIND=int_8), INTENT(IN) :: rowcol
     252             :       INTEGER :: blk_size_one
     253             : 
     254             :       MARK_USED(t)
     255             :       MARK_USED(rowcol)
     256     5381221 :       blk_size_one = 1
     257     5381221 :    END FUNCTION
     258             : 
     259             : ! **************************************************************************************************
     260             : !> \brief ...
     261             : !> \param blk_size_vec ...
     262             : !> \return ...
     263             : !> \author Patrick Seewald
     264             : ! **************************************************************************************************
     265      205884 :    FUNCTION new_block_tas_blk_size_arb(blk_size_vec)
     266             :       INTEGER, DIMENSION(:), INTENT(IN)                  :: blk_size_vec
     267             :       TYPE(dbt_tas_blk_size_arb)                         :: new_block_tas_blk_size_arb
     268             : 
     269      308826 :       ALLOCATE (new_block_tas_blk_size_arb%blk_size_vec(SIZE(blk_size_vec)))
     270      922033 :       new_block_tas_blk_size_arb%blk_size_vec(:) = blk_size_vec(:)
     271      102942 :       new_block_tas_blk_size_arb%nmrowcol = SIZE(blk_size_vec)
     272      922033 :       new_block_tas_blk_size_arb%nfullrowcol = SUM(blk_size_vec)
     273      102942 :    END FUNCTION
     274             : 
     275             : ! **************************************************************************************************
     276             : !> \brief ...
     277             : !> \param blk_size_vec ...
     278             : !> \param n_repl ...
     279             : !> \return ...
     280             : !> \author Patrick Seewald
     281             : ! **************************************************************************************************
     282      205836 :    FUNCTION new_block_tas_blk_size_repl(blk_size_vec, n_repl)
     283             :       INTEGER, DIMENSION(:), INTENT(IN)                  :: blk_size_vec
     284             :       INTEGER, INTENT(IN)                                :: n_repl
     285             :       TYPE(dbt_tas_blk_size_repl)                        :: new_block_tas_blk_size_repl
     286             : 
     287      102918 :       new_block_tas_blk_size_repl%nmrowcol_local = SIZE(blk_size_vec)
     288      308754 :       ALLOCATE (new_block_tas_blk_size_repl%blk_size_vec(new_block_tas_blk_size_repl%nmrowcol_local))
     289      926083 :       new_block_tas_blk_size_repl%blk_size_vec(:) = blk_size_vec(:)
     290      102918 :       new_block_tas_blk_size_repl%nmrowcol = new_block_tas_blk_size_repl%nmrowcol_local*n_repl
     291      926083 :       new_block_tas_blk_size_repl%nfullrowcol = SUM(blk_size_vec)*n_repl
     292      102918 :    END FUNCTION
     293             : 
     294             : ! **************************************************************************************************
     295             : !> \brief ...
     296             : !> \param nrowcol ...
     297             : !> \return ...
     298             : !> \author Patrick Seewald
     299             : ! **************************************************************************************************
     300      130692 :    FUNCTION new_block_tas_blk_size_one(nrowcol)
     301             :       INTEGER(KIND=int_8), INTENT(IN)                    :: nrowcol
     302             :       TYPE(dbt_tas_blk_size_one)                         :: new_block_tas_blk_size_one
     303             : 
     304      130692 :       new_block_tas_blk_size_one%nmrowcol = nrowcol
     305      130692 :       new_block_tas_blk_size_one%nfullrowcol = nrowcol
     306      130692 :    END FUNCTION
     307             : 
     308             : ! **************************************************************************************************
     309             : !> \brief ...
     310             : !> \param t ...
     311             : !> \param rowcol ...
     312             : !> \return ...
     313             : !> \author Patrick Seewald
     314             : ! **************************************************************************************************
     315     9976275 :    FUNCTION arb_dist(t, rowcol)
     316             :       CLASS(dbt_tas_dist_arb), INTENT(IN) :: t
     317             :       INTEGER(KIND=int_8), INTENT(IN) :: rowcol
     318             :       INTEGER :: arb_dist
     319             : 
     320     9976275 :       arb_dist = t%dist_vec(rowcol)
     321     9976275 :    END FUNCTION
     322             : 
     323             : ! **************************************************************************************************
     324             : !> \brief ...
     325             : !> \param t ...
     326             : !> \param rowcol ...
     327             : !> \return ...
     328             : !> \author Patrick Seewald
     329             : ! **************************************************************************************************
     330     6827417 :    FUNCTION repl_dist(t, rowcol)
     331             :       CLASS(dbt_tas_dist_repl), INTENT(IN) :: t
     332             :       INTEGER(KIND=int_8), INTENT(IN) :: rowcol
     333             :       INTEGER :: repl_dist
     334             :       INTEGER :: rowcol_local
     335             :       INTEGER :: igroup
     336             : 
     337     6827417 :       igroup = INT((rowcol - 1_int_8)/t%nmrowcol_local)
     338     6827417 :       rowcol_local = INT(MOD(rowcol - 1_int_8, INT(t%nmrowcol_local, KIND=int_8))) + 1
     339             : 
     340     6827417 :       repl_dist = t%dist_vec(rowcol_local) + igroup*t%dist_size
     341             : 
     342     6827417 :    END FUNCTION
     343             : 
     344             : ! **************************************************************************************************
     345             : !> \brief ...
     346             : !> \param t ...
     347             : !> \param dist ...
     348             : !> \return ...
     349             : !> \author Patrick Seewald
     350             : ! **************************************************************************************************
     351      205836 :    FUNCTION repl_rowcols(t, dist)
     352             :       CLASS(dbt_tas_dist_repl), INTENT(IN) :: t
     353             :       INTEGER, INTENT(IN) :: dist
     354             :       INTEGER :: nrowcols
     355      205836 :       INTEGER(KIND=int_8), DIMENSION(:), ALLOCATABLE :: repl_rowcols, rowcols_tmp
     356             :       INTEGER :: igroup
     357             :       INTEGER :: rowcol, count
     358             :       LOGICAL :: cond
     359             : 
     360      205836 :       igroup = dist/t%dist_size
     361             : 
     362      205836 :       nrowcols = t%nmrowcol_local
     363      205836 :       count = 0
     364      617508 :       ALLOCATE (rowcols_tmp(nrowcols))
     365     1852166 :       rowcols_tmp(:) = 0
     366     1852166 :       DO rowcol = 1, nrowcols
     367     1646330 :          cond = t%dist_vec(rowcol) + igroup*t%dist_size == dist
     368             : 
     369     1852166 :          IF (cond) THEN
     370     1646330 :             count = count + 1
     371     1646330 :             rowcols_tmp(count) = rowcol
     372             :          END IF
     373             :       END DO
     374             : 
     375      617508 :       ALLOCATE (repl_rowcols(count))
     376     1852166 :       repl_rowcols(:) = rowcols_tmp(1:count) + igroup*t%nmrowcol_local
     377             : 
     378             :    END FUNCTION
     379             : 
     380             : ! **************************************************************************************************
     381             : !> \brief ...
     382             : !> \param t ...
     383             : !> \param dist ...
     384             : !> \return ...
     385             : !> \author Patrick Seewald
     386             : ! **************************************************************************************************
     387      333728 :    FUNCTION arb_rowcols(t, dist)
     388             :       CLASS(dbt_tas_dist_arb), INTENT(IN) :: t
     389             :       INTEGER, INTENT(IN) :: dist
     390             :       INTEGER(KIND=int_8) :: rowcol, nrowcols
     391      333728 :       INTEGER(KIND=int_8), DIMENSION(:), ALLOCATABLE :: arb_rowcols, rowcols_tmp
     392             :       INTEGER :: count
     393             : 
     394      333728 :       nrowcols = t%nmrowcol
     395      333728 :       count = 0
     396     1001184 :       ALLOCATE (rowcols_tmp(nrowcols))
     397     2732110 :       rowcols_tmp(:) = 0
     398     2732110 :       DO rowcol = 1, nrowcols
     399     2732110 :          IF (t%dist_vec(rowcol) == dist) THEN
     400     1674198 :             count = count + 1
     401     1674198 :             rowcols_tmp(count) = rowcol
     402             :          END IF
     403             :       END DO
     404             : 
     405     1000504 :       ALLOCATE (arb_rowcols(count))
     406     2007926 :       arb_rowcols(:) = rowcols_tmp(1:count)
     407             :    END FUNCTION
     408             : 
     409             : ! **************************************************************************************************
     410             : !> \brief ...
     411             : !> \param split_size ...
     412             : !> \param nprowcol ...
     413             : !> \param nmrowcol ...
     414             : !> \return ...
     415             : !> \author Patrick Seewald
     416             : ! **************************************************************************************************
     417         216 :    FUNCTION new_block_tas_dist_cyclic(split_size, nprowcol, nmrowcol)
     418             :       INTEGER, INTENT(IN)                                :: split_size, nprowcol
     419             :       INTEGER(KIND=int_8), INTENT(IN)                    :: nmrowcol
     420             :       TYPE(dbt_tas_dist_cyclic)                          :: new_block_tas_dist_cyclic
     421             : 
     422         216 :       new_block_tas_dist_cyclic%split_size = split_size
     423         216 :       new_block_tas_dist_cyclic%nprowcol = nprowcol
     424         216 :       new_block_tas_dist_cyclic%nmrowcol = nmrowcol
     425         216 :    END FUNCTION
     426             : 
     427             : ! **************************************************************************************************
     428             : !> \brief ...
     429             : !> \param dist_vec ...
     430             : !> \param nprowcol ...
     431             : !> \param nmrowcol ...
     432             : !> \return ...
     433             : !> \author Patrick Seewald
     434             : ! **************************************************************************************************
     435      631728 :    FUNCTION new_block_tas_dist_arb(dist_vec, nprowcol, nmrowcol)
     436             :       INTEGER, DIMENSION(:), INTENT(IN)                  :: dist_vec
     437             :       INTEGER, INTENT(IN)                                :: nprowcol
     438             :       INTEGER(KIND=int_8), INTENT(IN)                    :: nmrowcol
     439             :       TYPE(dbt_tas_dist_arb)                             :: new_block_tas_dist_arb
     440             : 
     441      947592 :       ALLOCATE (new_block_tas_dist_arb%dist_vec(nmrowcol))
     442     2890003 :       new_block_tas_dist_arb%dist_vec(:) = dist_vec(:)
     443      315864 :       new_block_tas_dist_arb%nprowcol = nprowcol
     444      315864 :       new_block_tas_dist_arb%nmrowcol = nmrowcol
     445      315864 :    END FUNCTION
     446             : 
     447             : ! **************************************************************************************************
     448             : !> \brief Distribution that is more or less cyclic (round robin) and load balanced with different
     449             : !>        weights for each element.
     450             : !>        This is used for creating adhoc distributions whenever matrices are mapped to new grids.
     451             : !>        Only for small dimensions since distribution is created as an array
     452             : !> \param nprowcol ...
     453             : !> \param nmrowcol ...
     454             : !> \param dbt_sizes ...
     455             : !> \return ...
     456             : !> \author Patrick Seewald
     457             : ! **************************************************************************************************
     458      212946 :    FUNCTION dbt_tas_dist_arb_default(nprowcol, nmrowcol, dbt_sizes)
     459             :       INTEGER                                            :: nprowcol
     460             :       INTEGER(KIND=int_8), INTENT(IN)                    :: nmrowcol
     461             : 
     462             :       CLASS(dbt_tas_rowcol_data), INTENT(IN) :: dbt_sizes
     463             :       TYPE(dbt_tas_dist_arb)            :: dbt_tas_dist_arb_default
     464      425892 :       INTEGER, DIMENSION(nmrowcol) :: dist_vec, bsize_vec
     465             :       INTEGER(KIND=int_8) :: ind
     466             : 
     467     1969034 :       DO ind = 1, nmrowcol
     468     1969034 :          bsize_vec(ind) = dbt_sizes%data(ind)
     469             :       END DO
     470             : 
     471      212946 :       CALL dbt_tas_default_distvec(INT(nmrowcol), nprowcol, bsize_vec, dist_vec)
     472      212946 :       dbt_tas_dist_arb_default = dbt_tas_dist_arb(dist_vec, nprowcol, nmrowcol)
     473             : 
     474      212946 :    END FUNCTION
     475             : 
     476             : ! **************************************************************************************************
     477             : !> \brief get a load-balanced and randomized distribution along one tensor dimension
     478             : !> \param nblk number of blocks (along one tensor dimension)
     479             : !> \param nproc number of processes (along one process grid dimension)
     480             : !> \param blk_size block sizes
     481             : !> \param dist distribution
     482             : !> \author Patrick Seewald
     483             : ! **************************************************************************************************
     484      249456 :    SUBROUTINE dbt_tas_default_distvec(nblk, nproc, blk_size, dist)
     485             :       INTEGER, INTENT(IN)                                :: nblk, nproc
     486             :       INTEGER, DIMENSION(nblk), INTENT(IN)               :: blk_size
     487             :       INTEGER, DIMENSION(nblk), INTENT(OUT)              :: dist
     488             : 
     489      249456 :       CALL distribute_lpt_random(nblk, nproc, blk_size, dist)
     490             : 
     491      249456 :    END SUBROUTINE
     492             : 
     493             : ! **************************************************************************************************
     494             : !> \brief distribute `nel` elements with weights `weights` over `nbin` bins.
     495             : !>        load balanced distribution is obtained by using LPT algorithm together with randomization
     496             : !>        over equivalent bins
     497             : !>        (i.e. randomization over all bins with the smallest accumulated weight)
     498             : !> \param nel ...
     499             : !> \param nbin ...
     500             : !> \param weights ...
     501             : !> \param dist ...
     502             : !> \author Patrick Seewald
     503             : ! **************************************************************************************************
     504      249456 :    SUBROUTINE distribute_lpt_random(nel, nbin, weights, dist)
     505             :       !!
     506             :       INTEGER, INTENT(IN)                                :: nel, nbin
     507             :       INTEGER, DIMENSION(nel), INTENT(IN)                :: weights
     508             :       INTEGER, DIMENSION(nel), INTENT(OUT)               :: dist
     509             : 
     510             :       INTEGER, PARAMETER                                 :: n_idle = 1000
     511             : 
     512             :       INTEGER                                            :: i, i_select, ibin, iel, min_occup, &
     513             :                                                             n_avail
     514      249456 :       INTEGER, ALLOCATABLE, DIMENSION(:)                 :: bins_avail
     515             :       INTEGER, DIMENSION(4)                              :: iseed
     516      498912 :       INTEGER, DIMENSION(nel)                            :: sort_index, weights_s
     517      498912 :       INTEGER, DIMENSION(nbin)                           :: occup
     518      498912 :       LOGICAL, DIMENSION(nbin)                           :: bin_mask
     519             :       REAL(dp)                                           :: rand
     520             : 
     521             :       ! initialize seed based on input arguments such that random numbers are deterministic across all processes
     522     4249364 :       iseed(1) = nel; iseed(2) = nbin; iseed(3) = MAXVAL(weights); iseed(4) = MINVAL(weights)
     523             : 
     524      249456 :       iseed(4) = iseed(4)*2 + 1 ! odd
     525             : 
     526     1247280 :       iseed(:) = MODULO(iseed(:), 2**12)
     527             : 
     528   249705456 :       DO i = 1, n_idle
     529   249705456 :          CALL dlarnv(1, iseed, 1, rand)
     530             :       END DO
     531             : 
     532      593052 :       occup(:) = 0
     533     2249410 :       weights_s = weights
     534      249456 :       CALL sort(weights_s, nel, sort_index)
     535             : 
     536      593052 :       occup(:) = 0
     537     2249410 :       DO iel = nel, 1, -1
     538     4586884 :          min_occup = MINVAL(occup, 1)
     539             : 
     540             :          ! available bins with min. occupancy
     541     4586884 :          bin_mask = occup == min_occup
     542     4586884 :          n_avail = COUNT(bin_mask)
     543     5999862 :          ALLOCATE (bins_avail(n_avail))
     544     9173768 :          bins_avail(:) = PACK((/(i, i=1, nbin)/), MASK=bin_mask)
     545             : 
     546     1999954 :          CALL dlarnv(1, iseed, 1, rand)
     547     1999954 :          i_select = FLOOR(rand*n_avail) + 1
     548     1999954 :          ibin = bins_avail(i_select)
     549     1999954 :          DEALLOCATE (bins_avail)
     550             : 
     551     1999954 :          dist(sort_index(iel)) = ibin - 1
     552     2249410 :          occup(ibin) = occup(ibin) + weights_s(iel)
     553             :       END DO
     554             : 
     555      249456 :    END SUBROUTINE
     556             : 
     557             : ! **************************************************************************************************
     558             : !> \brief ...
     559             : !> \param dist_vec ...
     560             : !> \param nprowcol ...
     561             : !> \param nmrowcol ...
     562             : !> \param n_repl ...
     563             : !> \param dist_size ...
     564             : !> \return ...
     565             : !> \author Patrick Seewald
     566             : ! **************************************************************************************************
     567      205836 :    FUNCTION new_block_tas_dist_repl(dist_vec, nprowcol, nmrowcol, n_repl, dist_size)
     568             :       INTEGER, DIMENSION(:), INTENT(IN)                  :: dist_vec
     569             :       INTEGER, INTENT(IN)                                :: nprowcol, nmrowcol, n_repl, dist_size
     570             :       TYPE(dbt_tas_dist_repl)                            :: new_block_tas_dist_repl
     571             : 
     572      102918 :       new_block_tas_dist_repl%n_repl = n_repl
     573      102918 :       new_block_tas_dist_repl%dist_size = dist_size
     574      308754 :       ALLOCATE (new_block_tas_dist_repl%dist_vec(nmrowcol))
     575      926083 :       new_block_tas_dist_repl%dist_vec(:) = MOD(dist_vec(:), dist_size)
     576      102918 :       new_block_tas_dist_repl%nprowcol = nprowcol
     577      102918 :       new_block_tas_dist_repl%nmrowcol_local = nmrowcol
     578      102918 :       new_block_tas_dist_repl%nmrowcol = nmrowcol*n_repl
     579      102918 :    END FUNCTION
     580             : 
     581             : ! **************************************************************************************************
     582             : !> \brief ...
     583             : !> \param t ...
     584             : !> \param rowcol ...
     585             : !> \return ...
     586             : !> \author Patrick Seewald
     587             : ! **************************************************************************************************
     588       74356 :    FUNCTION cyclic_dist(t, rowcol)
     589             :       CLASS(dbt_tas_dist_cyclic), INTENT(IN) :: t
     590             :       INTEGER(KIND=int_8), INTENT(IN) :: rowcol
     591             :       INTEGER :: cyclic_dist
     592             : 
     593       74356 :       cyclic_dist = INT(MOD((rowcol - 1)/INT(t%split_size, KIND=int_8), INT(t%nprowcol, KIND=int_8)))
     594             : 
     595       74356 :    END FUNCTION
     596             : 
     597             : ! **************************************************************************************************
     598             : !> \brief ...
     599             : !> \param t ...
     600             : !> \param dist ...
     601             : !> \return ...
     602             : !> \author Patrick Seewald
     603             : ! **************************************************************************************************
     604         312 :    FUNCTION cyclic_rowcols(t, dist)
     605             :       CLASS(dbt_tas_dist_cyclic), INTENT(IN) :: t
     606             :       INTEGER, INTENT(IN) :: dist
     607             :       INTEGER(KIND=int_8), DIMENSION(:), ALLOCATABLE :: cyclic_rowcols
     608             :       INTEGER :: count, nsplit, isplit, irowcol, max_size
     609             :       INTEGER(KIND=int_8) :: rowcol
     610         312 :       INTEGER(KIND=int_8), DIMENSION(:), ALLOCATABLE :: rowcols_tmp
     611             : 
     612         312 :       nsplit = INT((t%nmrowcol - 1)/INT(t%split_size, KIND=int_8) + 1_int_8)
     613         312 :       max_size = nsplit*t%split_size
     614         936 :       ALLOCATE (rowcols_tmp(max_size))
     615       31016 :       rowcols_tmp(:) = 0
     616             :       count = 0
     617       12036 :       loop: DO isplit = 1, nsplit
     618       27628 :          DO irowcol = 1, t%split_size
     619             :             rowcol = INT((dist + (isplit - 1)*t%nprowcol), KIND=int_8)*INT(t%split_size, KIND=int_8) + &
     620       15592 :                      INT(irowcol, KIND=int_8)
     621       27316 :             IF (rowcol > t%nmrowcol) THEN
     622             :                EXIT loop
     623             :             ELSE
     624       15280 :                count = count + 1
     625       15280 :                rowcols_tmp(count) = rowcol
     626             :             END IF
     627             :          END DO
     628             :       END DO loop
     629             : 
     630         936 :       ALLOCATE (cyclic_rowcols(count))
     631       15592 :       cyclic_rowcols(:) = rowcols_tmp(1:count)
     632             :    END FUNCTION
     633             : 
     634     4822962 : END MODULE

Generated by: LCOV version 1.15