mp_diag.f90 9.93 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11
!
! Copyright (C) 2013 Quantum ESPRESSO group
! This file is distributed under the terms of the
! GNU General Public License. See the file `License'
! in the root directory of the present distribution,
! or http://www.gnu.org/copyleft/gpl.txt .
!
!----------------------------------------------------------------------------
MODULE mp_diag
  !----------------------------------------------------------------------------
  !
12
  USE mp, ONLY : mp_size, mp_rank, mp_sum, mp_comm_free, mp_comm_split
13 14 15 16
  !
  ! The following variables are needed in order to set up the communicator
  ! for scalapack
  !
17
  USE mp_world, ONLY : world_comm
18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39
  USE mp_pools, ONLY : npool, nproc_pool, my_pool_id
  USE mp_bands, ONLY : nbgrp, my_bgrp_id
  !
  USE parallel_include
  !
  IMPLICIT NONE 
  SAVE
  !
  ! ... linear-algebra group (also known as "ortho" or "diag" group). 
  ! ... Used for parallelization of dense-matrix diagonalization used in
  ! ... iterative diagonalization/orthonormalization, matrix-matrix products
  !
  INTEGER :: np_ortho(2) = 1  ! size of the processor grid used in ortho
  INTEGER :: me_ortho(2) = 0  ! coordinates of the processors
  INTEGER :: me_ortho1   = 0  ! task id for the ortho group
  INTEGER :: nproc_ortho = 1  ! size of the ortho group:
  INTEGER :: leg_ortho   = 1  ! the distance in the father communicator
                              ! of two neighbour processors in ortho_comm
  INTEGER :: ortho_comm  = 0  ! communicator for the ortho group
  INTEGER :: ortho_row_comm  = 0  ! communicator for the ortho row group
  INTEGER :: ortho_col_comm  = 0  ! communicator for the ortho col group
  INTEGER :: ortho_comm_id= 0 ! id of the ortho_comm
40
  INTEGER :: ortho_parent_comm  = 0  ! parent communicator from which ortho group has been created
41 42 43 44
  !
#if defined __SCALAPACK
  INTEGER :: me_blacs   =  0  ! BLACS processor index starting from 0
  INTEGER :: np_blacs   =  1  ! BLACS number of processor
45 46
#endif
  !
47 48 49 50 51 52
  INTEGER :: world_cntx = -1  ! BLACS context of all processor 
  INTEGER :: ortho_cntx = -1  ! BLACS context for ortho_comm
  !
CONTAINS
  !
  !----------------------------------------------------------------------------
53
  SUBROUTINE mp_start_diag( ndiag_, parent_comm )
54 55 56 57 58 59
    !---------------------------------------------------------------------------
    !
    ! ... Ortho/diag/linear algebra group initialization
    !
    IMPLICIT NONE
    !
60
    INTEGER, INTENT(IN) :: ndiag_, parent_comm
61 62 63 64 65 66 67 68
    !
    INTEGER :: nproc_ortho_try
    INTEGER :: parent_nproc  ! nproc of the parent group
    INTEGER :: ierr = 0
    !
    parent_nproc = mp_size( parent_comm )
    !
#if defined __SCALAPACK
69 70
    np_blacs     = mp_size( world_comm )
    me_blacs     = mp_rank( world_comm )
71
    !
72 73 74 75 76 77
    ! define a 1D grid containing all MPI tasks of the global communicator
    ! NOTE: world_cntx has the MPI communicator on entry and the BLACS context on exit
    !       BLACS_GRID_INIT() will create a copy of the communicator, which can be
    !       later retrieved using CALL BLACS_GET(world_cntx, 10, comm_copy)
    !
    world_cntx = world_comm
78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137
    CALL BLACS_GRIDINIT( world_cntx, 'Row', 1, np_blacs )
    !
#endif
    !
    IF( ndiag_ > 0 ) THEN
       ! command-line argument -ndiag N or -northo N set to a value N
       ! use the command line value ensuring that it falls in the proper range
       nproc_ortho_try = MIN( ndiag_ , parent_nproc )
    ELSE 
       ! no command-line argument -ndiag N or -northo N is present
       ! insert here custom architecture specific default definitions
#if defined __SCALAPACK
       nproc_ortho_try = MAX( parent_nproc/2, 1 )
#else
       nproc_ortho_try = 1
#endif
    END IF
    !
    ! the ortho group for parallel linear algebra is a sub-group of the pool,
    ! then there are as many ortho groups as pools.
    !
    CALL init_ortho_group( nproc_ortho_try, parent_comm )
    !  
    RETURN
    !
  END SUBROUTINE mp_start_diag
  !
  !
  SUBROUTINE init_ortho_group( nproc_try_in, comm_all )
    !
    IMPLICIT NONE

    INTEGER, INTENT(IN) :: nproc_try_in, comm_all

    LOGICAL, SAVE :: first = .true.
    INTEGER :: ierr, color, key, me_all, nproc_all, nproc_try

#if defined __SCALAPACK
    INTEGER, ALLOCATABLE :: blacsmap(:,:)
    INTEGER, ALLOCATABLE :: ortho_cntx_pe(:,:,:)
    INTEGER :: nprow, npcol, myrow, mycol, i, j, k
    INTEGER, EXTERNAL :: BLACS_PNUM
    !
    INTEGER :: nparent=1
    INTEGER :: total_nproc=1
    INTEGER :: total_mype=0
    INTEGER :: nproc_parent=1
    INTEGER :: my_parent_id=0
#endif


#if defined __MPI

    me_all    = mp_rank( comm_all )
    !
    nproc_all = mp_size( comm_all )
    !
    nproc_try = MIN( nproc_try_in, nproc_all )
    nproc_try = MAX( nproc_try, 1 )

138
    IF( .NOT. first ) CALL clean_ortho_group ( ) 
139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184

    !  find the square closer (but lower) to nproc_try
    !
    CALL grid2d_dims( 'S', nproc_try, np_ortho(1), np_ortho(2) )
    !
    !  now, and only now, it is possible to define the number of tasks
    !  in the ortho group for parallel linear algebra
    !
    nproc_ortho = np_ortho(1) * np_ortho(2)
    !
    IF( nproc_all >= 4*nproc_ortho ) THEN
       !
       !  here we choose a processor every 4, in order not to stress memory BW
       !  on multi core procs, for which further performance enhancements are
       !  possible using OpenMP BLAS inside regter/cegter/rdiaghg/cdiaghg
       !  (to be implemented)
       !
       color = 0
       IF( me_all < 4*nproc_ortho .AND. MOD( me_all, 4 ) == 0 ) color = 1
       !
       leg_ortho = 4
       !
    ELSE IF( nproc_all >= 2*nproc_ortho ) THEN
       !
       !  here we choose a processor every 2, in order not to stress memory BW
       !
       color = 0
       IF( me_all < 2*nproc_ortho .AND. MOD( me_all, 2 ) == 0 ) color = 1
       !
       leg_ortho = 2
       !
    ELSE
       !
       !  here we choose the first processors
       !
       color = 0
       IF( me_all < nproc_ortho ) color = 1
       !
       leg_ortho = 1
       !
    END IF
    !
    key   = me_all
    !
    !  initialize the communicator for the new group by splitting the input communicator
    !
185
    CALL mp_comm_split ( comm_all, color, key, ortho_comm )
186
    !
187 188 189 190
    ! and remember where it comes from
    !
    ortho_parent_comm = comm_all
    !
191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206
    !  Computes coordinates of the processors, in row maior order
    !
    me_ortho1   = mp_rank( ortho_comm )
    !
    IF( me_all == 0 .AND. me_ortho1 /= 0 ) &
         CALL errore( " init_ortho_group ", " wrong root task in ortho group ", ierr )
    !
    if( color == 1 ) then
       ortho_comm_id = 1
       CALL GRID2D_COORDS( 'R', me_ortho1, np_ortho(1), np_ortho(2), me_ortho(1), me_ortho(2) )
       CALL GRID2D_RANK( 'R', np_ortho(1), np_ortho(2), me_ortho(1), me_ortho(2), ierr )
       IF( ierr /= me_ortho1 ) &
            CALL errore( " init_ortho_group ", " wrong task coordinates in ortho group ", ierr )
       IF( me_ortho1*leg_ortho /= me_all ) &
            CALL errore( " init_ortho_group ", " wrong rank assignment in ortho group ", ierr )

207 208
       CALL mp_comm_split( ortho_comm, me_ortho(2), me_ortho(1), ortho_col_comm)
       CALL mp_comm_split( ortho_comm, me_ortho(1), me_ortho(2), ortho_row_comm)
209 210 211 212 213 214 215 216 217 218 219 220

    else
       ortho_comm_id = 0
       me_ortho(1) = me_ortho1
       me_ortho(2) = me_ortho1
    endif
#if defined __SCALAPACK
    !
    !  This part is used to eliminate the image dependency from ortho groups
    !  SCALAPACK is now independent of whatever level of parallelization
    !  is present on top of pool parallelization
    !
221 222
    total_nproc = mp_size( world_comm )
    total_mype = mp_rank( world_comm )
223 224 225 226 227 228 229 230 231 232 233 234 235 236
    nparent = total_nproc/npool/nproc_pool
    nproc_parent = total_nproc/nparent
    my_parent_id = total_mype/nproc_parent
    !
    !
    ALLOCATE( ortho_cntx_pe( npool, nbgrp, nparent ) )
    ALLOCATE( blacsmap( np_ortho(1), np_ortho(2) ) )

    DO j = 1, nparent

     DO k = 1, nbgrp

       DO i = 1, npool

237
         CALL BLACS_GET(world_cntx, 10, ortho_cntx_pe( i, k, j)) ! retrieve communicator of world context
238 239 240 241 242 243 244 245 246 247 248
         blacsmap = 0
         nprow = np_ortho(1)
         npcol = np_ortho(2)

         IF( ( j == ( my_parent_id + 1 ) ) .and. ( k == ( my_bgrp_id + 1 ) ) .and.  &
             ( i == ( my_pool_id  + 1 ) ) .and. ( ortho_comm_id > 0 ) ) THEN

           blacsmap( me_ortho(1) + 1, me_ortho(2) + 1 ) = BLACS_PNUM( world_cntx, 0, me_blacs )

         END IF

249
         ! All MPI tasks defined in the global communicator take part in the definition of the BLACS grid
250

251
         CALL mp_sum( blacsmap, world_comm ) 
252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295

         CALL BLACS_GRIDMAP( ortho_cntx_pe(i,k,j), blacsmap, nprow, nprow, npcol )

         CALL BLACS_GRIDINFO( ortho_cntx_pe(i,k,j), nprow, npcol, myrow, mycol )

         IF( ( j == ( my_parent_id + 1 ) ) .and. ( k == ( my_bgrp_id + 1 ) ) .and. &
             ( i == ( my_pool_id  + 1 ) ) .and. ( ortho_comm_id > 0 ) ) THEN

            IF(  np_ortho(1) /= nprow ) &
               CALL errore( ' init_ortho_group ', ' problem with SCALAPACK, wrong no. of task rows ', 1 )
            IF(  np_ortho(2) /= npcol ) &
               CALL errore( ' init_ortho_group ', ' problem with SCALAPACK, wrong no. of task columns ', 1 )
            IF(  me_ortho(1) /= myrow ) &
               CALL errore( ' init_ortho_group ', ' problem with SCALAPACK, wrong task row ID ', 1 )
            IF(  me_ortho(2) /= mycol ) &
               CALL errore( ' init_ortho_group ', ' problem with SCALAPACK, wrong task columns ID ', 1 )

            ortho_cntx = ortho_cntx_pe(i,k,j)

         END IF

       END DO

     END DO

    END DO 

    DEALLOCATE( blacsmap )
    DEALLOCATE( ortho_cntx_pe )


#endif

#else

    ortho_comm_id = 1

#endif

    first = .false.

    RETURN
  END SUBROUTINE init_ortho_group
  !
296 297 298 299 300 301 302 303 304 305
  SUBROUTINE clean_ortho_group ( )
    !  
    !  free resources associated to the communicator
    !
    CALL mp_comm_free( ortho_comm )
    IF(  ortho_comm_id > 0  ) THEN
       CALL mp_comm_free( ortho_col_comm )
       CALL mp_comm_free( ortho_row_comm )
    ENDIF
#if defined __SCALAPACK
306
    IF(  ortho_cntx /= -1 ) CALL BLACS_GRIDEXIT( ortho_cntx )
307 308 309 310 311
    ortho_cntx = -1
#endif
    !
  END SUBROUTINE clean_ortho_group
  !
312
END MODULE mp_diag