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 DBT tensor Input / Output
      10             : !> \author Patrick Seewald
      11             : ! **************************************************************************************************
      12             : MODULE dbt_io
      13             : 
      14             :    #:include "dbt_macros.fypp"
      15             :    #:set maxdim = maxrank
      16             :    #:set ndims = range(2,maxdim+1)
      17             : 
      18             :    USE dbt_types, ONLY: &
      19             :       dbt_get_info, dbt_type, ndims_tensor, dbt_get_num_blocks, dbt_get_num_blocks_total, &
      20             :       blk_dims_tensor, dbt_get_stored_coordinates, dbt_get_nze, dbt_get_nze_total, &
      21             :       dbt_pgrid_type, dbt_nblks_total
      22             :    USE kinds, ONLY: default_string_length, int_8, dp
      23             :    USE message_passing, ONLY: mp_comm_type
      24             :    USE dbt_block, ONLY: &
      25             :       dbt_iterator_type, dbt_iterator_next_block, dbt_iterator_start, &
      26             :       dbt_iterator_blocks_left, dbt_iterator_stop, dbt_get_block
      27             :    USE dbt_tas_io, ONLY: dbt_tas_write_split_info
      28             : 
      29             : #include "../base/base_uses.f90"
      30             : 
      31             :    IMPLICIT NONE
      32             :    PRIVATE
      33             :    CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'dbt_types'
      34             : 
      35             :    PUBLIC :: &
      36             :       dbt_write_tensor_info, &
      37             :       dbt_write_tensor_dist, &
      38             :       dbt_write_blocks, &
      39             :       dbt_write_block, &
      40             :       dbt_write_block_indices, &
      41             :       dbt_write_split_info, &
      42             :       prep_output_unit
      43             : 
      44             : CONTAINS
      45             : 
      46             : ! **************************************************************************************************
      47             : !> \brief Write tensor global info: block dimensions, full dimensions and process grid dimensions
      48             : !> \param full_info Whether to print distribution and block size vectors
      49             : !> \author Patrick Seewald
      50             : ! **************************************************************************************************
      51      135836 :    SUBROUTINE dbt_write_tensor_info(tensor, unit_nr, full_info)
      52             :       TYPE(dbt_type), INTENT(IN) :: tensor
      53             :       INTEGER, INTENT(IN)            :: unit_nr
      54             :       LOGICAL, OPTIONAL, INTENT(IN)  :: full_info
      55      271672 :       INTEGER, DIMENSION(ndims_tensor(tensor)) :: nblks_total, nfull_total, pdims, my_ploc, nblks_local, nfull_local
      56             : 
      57             :       #:for idim in range(1, maxdim+1)
      58      271672 :          INTEGER, DIMENSION(dbt_nblks_total(tensor, ${idim}$)) :: proc_dist_${idim}$
      59      271672 :          INTEGER, DIMENSION(dbt_nblks_total(tensor, ${idim}$)) :: blk_size_${idim}$
      60      135836 :          INTEGER, DIMENSION(dbt_nblks_total(tensor, ${idim}$)) :: blks_local_${idim}$
      61             :       #:endfor
      62             :       CHARACTER(len=default_string_length)                     :: name
      63             :       INTEGER                                                  :: idim
      64             :       INTEGER                                                  :: iblk
      65             :       INTEGER                                                  :: unit_nr_prv
      66             : 
      67      135836 :       unit_nr_prv = prep_output_unit(unit_nr)
      68      135836 :       IF (unit_nr_prv == 0) RETURN
      69             : 
      70             :       CALL dbt_get_info(tensor, nblks_total, nfull_total, nblks_local, nfull_local, pdims, my_ploc, &
      71             :                         ${varlist("blks_local")}$, ${varlist("proc_dist")}$, ${varlist("blk_size")}$, &
      72      135836 :                         name=name)
      73             : 
      74      135836 :       IF (unit_nr_prv > 0) THEN
      75             :          WRITE (unit_nr_prv, "(T2,A)") &
      76          45 :             "GLOBAL INFO OF "//TRIM(name)
      77          45 :          WRITE (unit_nr_prv, "(T4,A,1X)", advance="no") "block dimensions:"
      78         184 :          DO idim = 1, ndims_tensor(tensor)
      79         184 :             WRITE (unit_nr_prv, "(I6)", advance="no") nblks_total(idim)
      80             :          END DO
      81          45 :          WRITE (unit_nr_prv, "(/T4,A,1X)", advance="no") "full dimensions:"
      82         184 :          DO idim = 1, ndims_tensor(tensor)
      83         184 :             WRITE (unit_nr_prv, "(I8)", advance="no") nfull_total(idim)
      84             :          END DO
      85          45 :          WRITE (unit_nr_prv, "(/T4,A,1X)", advance="no") "process grid dimensions:"
      86         184 :          DO idim = 1, ndims_tensor(tensor)
      87         184 :             WRITE (unit_nr_prv, "(I6)", advance="no") pdims(idim)
      88             :          END DO
      89          45 :          WRITE (unit_nr_prv, *)
      90             : 
      91          45 :          IF (PRESENT(full_info)) THEN
      92          45 :             IF (full_info) THEN
      93           0 :                WRITE (unit_nr_prv, '(T4,A)', advance='no') "Block sizes:"
      94             :                #:for dim in range(1, maxdim+1)
      95           0 :                   IF (ndims_tensor(tensor) >= ${dim}$) THEN
      96           0 :                      WRITE (unit_nr_prv, '(/T8,A,1X,I1,A,1X)', advance='no') 'Dim', ${dim}$, ':'
      97           0 :                      DO iblk = 1, SIZE(blk_size_${dim}$)
      98           0 :                         WRITE (unit_nr_prv, '(I2,1X)', advance='no') blk_size_${dim}$ (iblk)
      99             :                      END DO
     100             :                   END IF
     101             :                #:endfor
     102           0 :                WRITE (unit_nr_prv, '(/T4,A)', advance='no') "Block distribution:"
     103             :                #:for dim in range(1, maxdim+1)
     104           0 :                   IF (ndims_tensor(tensor) >= ${dim}$) THEN
     105           0 :                      WRITE (unit_nr_prv, '(/T8,A,1X,I1,A,1X)', advance='no') 'Dim', ${dim}$, ':'
     106           0 :                      DO iblk = 1, SIZE(proc_dist_${dim}$)
     107           0 :                         WRITE (unit_nr_prv, '(I3,1X)', advance='no') proc_dist_${dim}$ (iblk)
     108             :                      END DO
     109             :                   END IF
     110             :                #:endfor
     111             :             END IF
     112          45 :             WRITE (unit_nr_prv, *)
     113             :          END IF
     114             :       END IF
     115             : 
     116             :    END SUBROUTINE
     117             : 
     118             : ! **************************************************************************************************
     119             : !> \brief Write info on tensor distribution & load balance
     120             : !> \author Patrick Seewald
     121             : ! **************************************************************************************************
     122      135836 :    SUBROUTINE dbt_write_tensor_dist(tensor, unit_nr)
     123             :       TYPE(dbt_type), INTENT(IN) :: tensor
     124             :       INTEGER, INTENT(IN)            :: unit_nr
     125             :       INTEGER                        :: nproc, nblock_max, nelement_max
     126             :       INTEGER(KIND=int_8)            :: nblock_sum, nelement_sum, nblock_tot
     127             :       INTEGER                        :: nblock, nelement, unit_nr_prv
     128             :       INTEGER, DIMENSION(2)          :: tmp
     129      135836 :       INTEGER, DIMENSION(ndims_tensor(tensor)) :: bdims
     130             :       REAL(KIND=dp)              :: occupation
     131             : 
     132      135836 :       unit_nr_prv = prep_output_unit(unit_nr)
     133      135836 :       IF (unit_nr_prv == 0) RETURN
     134             : 
     135      135836 :       nproc = tensor%pgrid%mp_comm_2d%num_pe
     136             : 
     137      135836 :       nblock = dbt_get_num_blocks(tensor)
     138      135836 :       nelement = dbt_get_nze(tensor)
     139             : 
     140      135836 :       nblock_sum = dbt_get_num_blocks_total(tensor)
     141      135836 :       nelement_sum = dbt_get_nze_total(tensor)
     142             : 
     143      407508 :       tmp = (/nblock, nelement/)
     144      135836 :       CALL tensor%pgrid%mp_comm_2d%max(tmp)
     145      135836 :       nblock_max = tmp(1); nelement_max = tmp(2)
     146             : 
     147      135836 :       CALL blk_dims_tensor(tensor, bdims)
     148      500514 :       nblock_tot = PRODUCT(INT(bdims, KIND=int_8))
     149             : 
     150      135836 :       occupation = -1.0_dp
     151      135836 :       IF (nblock_tot .NE. 0) occupation = 100.0_dp*REAL(nblock_sum, dp)/REAL(nblock_tot, dp)
     152             : 
     153      135836 :       IF (unit_nr_prv > 0) THEN
     154             :          WRITE (unit_nr_prv, "(T2,A)") &
     155          45 :             "DISTRIBUTION OF "//TRIM(tensor%name)
     156          45 :          WRITE (unit_nr_prv, "(T15,A,T68,I13)") "Number of non-zero blocks:", nblock_sum
     157          45 :          WRITE (unit_nr_prv, "(T15,A,T75,F6.2)") "Percentage of non-zero blocks:", occupation
     158          45 :          WRITE (unit_nr_prv, "(T15,A,T68,I13)") "Average number of blocks per CPU:", (nblock_sum + nproc - 1)/nproc
     159          45 :          WRITE (unit_nr_prv, "(T15,A,T68,I13)") "Maximum number of blocks per CPU:", nblock_max
     160          45 :          WRITE (unit_nr_prv, "(T15,A,T68,I13)") "Average number of matrix elements per CPU:", (nelement_sum + nproc - 1)/nproc
     161          45 :          WRITE (unit_nr_prv, "(T15,A,T68,I13)") "Maximum number of matrix elements per CPU:", nelement_max
     162             :       END IF
     163             : 
     164             :    END SUBROUTINE
     165             : 
     166             : ! **************************************************************************************************
     167             : !> \brief Write all tensor blocks
     168             : !> \param io_unit_master for global output
     169             : !> \param io_unit_all for local output
     170             : !> \param write_int convert to integers (useful for testing with integer tensors)
     171             : !> \author Patrick Seewald
     172             : ! **************************************************************************************************
     173           0 :    SUBROUTINE dbt_write_blocks(tensor, io_unit_master, io_unit_all, write_int)
     174             :       TYPE(dbt_type), INTENT(INOUT)                  :: tensor
     175             :       INTEGER, INTENT(IN)                                :: io_unit_master, io_unit_all
     176             :       LOGICAL, INTENT(IN), OPTIONAL                      :: write_int
     177           0 :       INTEGER, DIMENSION(ndims_tensor(tensor))          :: blk_index, blk_size
     178             :       #:for ndim in ndims
     179             :          REAL(KIND=dp), ALLOCATABLE, &
     180           0 :             DIMENSION(${shape_colon(ndim)}$)                :: blk_values_${ndim}$
     181             :       #:endfor
     182             :       TYPE(dbt_iterator_type)                        :: iterator
     183             :       INTEGER                                            :: proc, mynode
     184             :       LOGICAL                                            :: found
     185             : 
     186           0 :       IF (io_unit_master > 0) THEN
     187           0 :          WRITE (io_unit_master, '(T7,A)') "(block index) @ process: (array index) value"
     188             :       END IF
     189           0 :       CALL dbt_iterator_start(iterator, tensor)
     190           0 :       DO WHILE (dbt_iterator_blocks_left(iterator))
     191           0 :          CALL dbt_iterator_next_block(iterator, blk_index, blk_size=blk_size)
     192           0 :          CALL dbt_get_stored_coordinates(tensor, blk_index, proc)
     193           0 :          mynode = tensor%pgrid%mp_comm_2d%mepos
     194           0 :          CPASSERT(proc .EQ. mynode)
     195             :          #:for ndim in ndims
     196           0 :             IF (ndims_tensor(tensor) == ${ndim}$) THEN
     197           0 :                CALL dbt_get_block(tensor, blk_index, blk_values_${ndim}$, found)
     198           0 :                CPASSERT(found)
     199             :                CALL dbt_write_block(tensor%name, blk_size, blk_index, proc, io_unit_all, &
     200           0 :                                     blk_values_${ndim}$=blk_values_${ndim}$, write_int=write_int)
     201           0 :                DEALLOCATE (blk_values_${ndim}$)
     202             :             END IF
     203             :          #:endfor
     204             :       END DO
     205           0 :       CALL dbt_iterator_stop(iterator)
     206           0 :    END SUBROUTINE
     207             : 
     208             : ! **************************************************************************************************
     209             : !> \brief Write a tensor block
     210             : !> \param name tensor name
     211             : !> \param blk_size block size
     212             : !> \param blk_index block index
     213             : !> \param blk_values_i block values for 2 dimensions
     214             : !> \param write_int write_int convert values to integers
     215             : !> \param unit_nr unit number
     216             : !> \param proc which process am I
     217             : !> \author Patrick Seewald
     218             : ! **************************************************************************************************
     219           0 :    SUBROUTINE dbt_write_block(name, blk_size, blk_index, proc, unit_nr, &
     220           0 :                               ${varlist("blk_values",nmin=2)}$, write_int)
     221             :       CHARACTER(LEN=*), INTENT(IN)                       :: name
     222             :       INTEGER, DIMENSION(:), INTENT(IN)                  :: blk_size
     223             :       INTEGER, DIMENSION(:), INTENT(IN)                  :: blk_index
     224             :       #:for ndim in ndims
     225             :          REAL(KIND=dp), &
     226             :             DIMENSION(${arrlist("blk_size", nmax=ndim)}$), &
     227             :             INTENT(IN), OPTIONAL                            :: blk_values_${ndim}$
     228             :       #:endfor
     229             :       LOGICAL, INTENT(IN), OPTIONAL                      :: write_int
     230             :       LOGICAL                                            :: write_int_prv
     231             :       INTEGER, INTENT(IN)                                :: unit_nr
     232             :       INTEGER, INTENT(IN)                                :: proc
     233             :       INTEGER                                            :: ${varlist("i")}$
     234             :       INTEGER                                            :: ndim
     235             : 
     236           0 :       IF (PRESENT(write_int)) THEN
     237           0 :          write_int_prv = write_int
     238             :       ELSE
     239             :          write_int_prv = .FALSE.
     240             :       END IF
     241             : 
     242           0 :       ndim = SIZE(blk_size)
     243             : 
     244           0 :       IF (unit_nr > 0) THEN
     245             :          #:for ndim in ndims
     246           0 :             IF (ndim == ${ndim}$) THEN
     247             :                #:for idim in range(ndim,0,-1)
     248           0 :                   DO i_${idim}$ = 1, blk_size(${idim}$)
     249             :                      #:endfor
     250           0 :                      IF (write_int_prv) THEN
     251             :                         WRITE (unit_nr, '(T7,A,T16,A,${ndim}$I3,1X,A,1X,I3,A,1X,A,${ndim}$I3,1X,A,1X,I20)') &
     252           0 :                            TRIM(name), "(", blk_index, ") @", proc, ':', &
     253           0 :                            "(", ${varlist("i", nmax=ndim)}$, ")", &
     254           0 :                            INT(blk_values_${ndim}$ (${varlist("i", nmax=ndim)}$), KIND=int_8)
     255             :                      ELSE
     256             :                         WRITE (unit_nr, '(T7,A,T16,A,${ndim}$I3,1X,A,1X,I3,A,1X,A,${ndim}$I3,1X,A,1X,F10.5)') &
     257           0 :                            TRIM(name), "(", blk_index, ") @", proc, ':', &
     258           0 :                            "(", ${varlist("i", nmax=ndim)}$, ")", &
     259           0 :                            blk_values_${ndim}$ (${varlist("i", nmax=ndim)}$)
     260             :                      END IF
     261             :                      #:for idim in range(ndim,0,-1)
     262             :                         END DO
     263             :                      #:endfor
     264             :                   END IF
     265             :                #:endfor
     266             :             END IF
     267           0 :          END SUBROUTINE
     268             : 
     269             : ! **************************************************************************************************
     270             : !> \author Patrick Seewald
     271             : ! **************************************************************************************************
     272           0 :          SUBROUTINE dbt_write_block_indices(tensor, io_unit_master, io_unit_all)
     273             :             TYPE(dbt_type), INTENT(INOUT)                  :: tensor
     274             :             INTEGER, INTENT(IN)                                :: io_unit_master, io_unit_all
     275             :             TYPE(dbt_iterator_type)                        :: iterator
     276           0 :             INTEGER, DIMENSION(ndims_tensor(tensor))          :: blk_index, blk_size
     277             :             INTEGER                                            :: mynode, proc
     278             : 
     279           0 :             IF (io_unit_master > 0) THEN
     280           0 :                WRITE (io_unit_master, '(T7,A)') "(block index) @ process: size"
     281             :             END IF
     282             : 
     283           0 :             CALL dbt_iterator_start(iterator, tensor)
     284           0 :             DO WHILE (dbt_iterator_blocks_left(iterator))
     285           0 :                CALL dbt_iterator_next_block(iterator, blk_index, blk_size=blk_size)
     286           0 :                CALL dbt_get_stored_coordinates(tensor, blk_index, proc)
     287           0 :                mynode = tensor%pgrid%mp_comm_2d%mepos
     288           0 :                CPASSERT(proc .EQ. mynode)
     289             :                #:for ndim in ndims
     290           0 :                   IF (ndims_tensor(tensor) == ${ndim}$) THEN
     291             :                      WRITE (io_unit_all, '(T7,A,T16,A,${ndim}$I3,1X,A,1X,I3,A2,${ndim}$I3)') &
     292           0 :                         TRIM(tensor%name), "blk index (", blk_index, ") @", proc, ":", blk_size
     293             :                   END IF
     294             :                #:endfor
     295             :             END DO
     296           0 :             CALL dbt_iterator_stop(iterator)
     297           0 :          END SUBROUTINE
     298             : 
     299             : ! **************************************************************************************************
     300             : !> \author Patrick Seewald
     301             : ! **************************************************************************************************
     302           0 :          SUBROUTINE dbt_write_split_info(pgrid, unit_nr)
     303             :             TYPE(dbt_pgrid_type), INTENT(IN) :: pgrid
     304             :             INTEGER, INTENT(IN) :: unit_nr
     305             : 
     306           0 :             IF (ALLOCATED(pgrid%tas_split_info)) THEN
     307           0 :                CALL dbt_tas_write_split_info(pgrid%tas_split_info, unit_nr)
     308             :             END IF
     309           0 :          END SUBROUTINE
     310             : 
     311             : ! **************************************************************************************************
     312             : !> \author Patrick Seewald
     313             : ! **************************************************************************************************
     314     1390944 :          FUNCTION prep_output_unit(unit_nr) RESULT(unit_nr_out)
     315             :             INTEGER, INTENT(IN), OPTIONAL :: unit_nr
     316             :             INTEGER                       :: unit_nr_out
     317             : 
     318     1390944 :             IF (PRESENT(unit_nr)) THEN
     319      804196 :                unit_nr_out = unit_nr
     320             :             ELSE
     321             :                unit_nr_out = 0
     322             :             END IF
     323             : 
     324     1390944 :          END FUNCTION
     325             : 
     326             :       END MODULE