LCOV - code coverage report
Current view: top level - src/45_xgTools - m_xgTransposer.F90 (source / functions) Coverage Total Hit
Test: coverage.info Lines: 71.9 % 452 325
Test Date: 2026-11-06 14:42:53 Functions: 68.8 % 16 11
Branches: 40.8 % 542 221

             Branch data     Line data    Source code
       1                 :             : !!****m* ABINIT/m_xgTransposer
       2                 :             : !! NAME
       3                 :             : !!  m_xgTransposer
       4                 :             : !!
       5                 :             : !! FUNCTION
       6                 :             : !! This module is to be used to go to "KGB" representation and to "linear
       7                 :             : !! algebra representation". It will replace most of prep_* subroutine
       8                 :             : !! This should really help to do the transposition operation.
       9                 :             : !!
      10                 :             : !! NOTES
      11                 :             : !! The transposer switches between two states. Let's assume for simplicity
      12                 :             : !! that we have four MPI processes in the comm_cols communicator and
      13                 :             : !! a single process in the comm_rows communicator:
      14                 :             : !!
      15                 :             : !!  STATE_LINALG=                       STATE_COLSROWS=
      16                 :             : !!
      17                 :             : !!              bands                            bands
      18                 :             : !!      |-------------------|            |----|----|----|----|
      19                 :             : !!      |        P0         |            |    |    |    |    |
      20                 :             : !!      |                   |            |    |    |    |    |
      21                 :             : !!      |-------------------|            |    |    |    |    |
      22                 :             : !!      |        P1         |            |    |    |    |    |
      23                 :             : !!      |                   |            |    |    |    |    |
      24                 :             : !!  pw  |-------------------|        pw  | P0 | P1 | P2 | P3 |
      25                 :             : !!      |        P2         |            |    |    |    |    |
      26                 :             : !!      |                   |            |    |    |    |    |
      27                 :             : !!      |-------------------|            |    |    |    |    |
      28                 :             : !!      |        P3         |            |    |    |    |    |
      29                 :             : !!      |                   |            |    |    |    |    |
      30                 :             : !!      |-------------------|            |----|----|----|----|
      31                 :             : !!
      32                 :             : !! The user can define custom block sizes using ncolsColsRows_sub and nrowsLinalg_sub.
      33                 :             : !!
      34                 :             : !! TODO                                  IML 04/04/2025
      35                 :             : !! -clean versions to always pass from my array version
      36                 :             : !!  for ncolsColsRows or create wrappers so that if we
      37                 :             : !!  give an array then we turn it into scalar
      38                 :             : !! -add unitary tests for my custom transposition
      39                 :             : !!
      40                 :             : !! COPYRIGHT
      41                 :             : !!  Copyright (C) 2017-2026 ABINIT group (J. Bieder, L. Baguet, IML)
      42                 :             : !!  This file is distributed under the terms of the
      43                 :             : !!  GNU General Public License, see ~abinit/COPYING
      44                 :             : !!  or http://www.gnu.org/copyleft/gpl.txt .
      45                 :             : !!
      46                 :             : !! NOTES
      47                 :             : !!
      48                 :             : !! SOURCE
      49                 :             : 
      50                 :             : #if defined HAVE_CONFIG_H
      51                 :             : #include "config.h"
      52                 :             : #endif
      53                 :             : 
      54                 :             : #include "abi_common.h"
      55                 :             : 
      56                 :             : ! nvtx related macro definition
      57                 :             : #include "nvtx_macros.h"
      58                 :             : 
      59                 :             : module m_xgTransposer
      60                 :             : 
      61                 :             :   use, intrinsic :: iso_c_binding, only: c_double, c_size_t, c_loc
      62                 :             : 
      63                 :             :   use defs_basis, only : std_err, std_out, dp, ABI_GPU_KOKKOS, ABI_GPU_OPENMP, ABI_GPU_DISABLED
      64                 :             :   use m_xomp
      65                 :             :   use m_profiling_abi
      66                 :             :   USE_MPI
      67                 :             :   use m_xmpi
      68                 :             :   use m_errors
      69                 :             :   use m_xg
      70                 :             :   use m_time
      71                 :             : 
      72                 :             : #if defined HAVE_YAKL
      73                 :             :   use gator_mod
      74                 :             : #endif
      75                 :             : 
      76                 :             : #if defined(HAVE_GPU_CUDA) && defined(HAVE_YAKL)
      77                 :             :  use m_gpu_toolbox, only : CPU_DEVICE_ID, gpu_device_synchronize, gpu_data_prefetch_async
      78                 :             : #endif
      79                 :             : 
      80                 :             : #if defined(HAVE_GPU_MARKERS)
      81                 :             :   use m_nvtx_data
      82                 :             : #endif
      83                 :             : 
      84                 :             :   implicit none
      85                 :             : 
      86                 :             : #ifdef HAVE_MPI1
      87                 :             :  include 'mpif.h'
      88                 :             : #endif
      89                 :             : 
      90                 :             :   private
      91                 :             : 
      92                 :             :   integer, parameter, public :: STATE_LINALG = 1
      93                 :             :   integer, parameter, public :: STATE_COLSROWS  = 2
      94                 :             :   integer, parameter, public :: STATE_UNKNOW    = 3
      95                 :             :   integer, parameter, public :: MPI_LINALG = 1
      96                 :             :   integer, parameter, public :: MPI_ROWS   = 2
      97                 :             :   integer, parameter, public :: MPI_COLS   = 3
      98                 :             :   integer, parameter, public :: MPI_2DCART = 4
      99                 :             :   integer, parameter, public :: TRANS_ALL2ALL = 1
     100                 :             :   integer, parameter, public :: TRANS_GATHER = 2
     101                 :             :   integer, parameter         :: TRANS_TYPE_CONSTRUCTED = 1
     102                 :             :   integer, parameter         :: TRANS_TYPE_CONSTRUCTED_NULL_COMM = 2
     103                 :             :   integer, parameter         :: TRANS_TYPE_COPIED = 3
     104                 :             : 
     105                 :             :   integer, parameter :: tim_toColsRows  = 1662
     106                 :             :   integer, parameter :: tim_toLinalg    = 1663
     107                 :             :   integer, parameter :: tim_all2allv    = 1664
     108                 :             :   integer, parameter :: tim_gatherv     = 1665
     109                 :             :   integer, parameter :: tim_reorganize  = 1666
     110                 :             :   integer, parameter :: tim_init        = 1667
     111                 :             :   integer, parameter :: tim_free        = 1668
     112                 :             :   integer, parameter :: tim_transpose   = 1669
     113                 :             : 
     114                 :             :   type, private :: mpiData_t
     115                 :             :     integer :: comm
     116                 :             :     integer :: rank
     117                 :             :     integer :: size
     118                 :             :   end type mpiData_t
     119                 :             : 
     120                 :             :   type, private :: ptr_t
     121                 :             :     double precision, pointer :: ptr(:,:) => null()
     122                 :             :   end type ptr_t
     123                 :             : 
     124                 :             :   type, public :: xgTransposer_t
     125                 :             :     type(xgBlock_t), pointer, private :: xgBlock_linalg => null()
     126                 :             :     type(xgBlock_t), pointer, private :: xgBlock_colsrows => null()
     127                 :             :     integer :: state
     128                 :             :     type(mpiData_t), private :: mpiData(4)
     129                 :             :     integer, allocatable, private :: lookup(:)
     130                 :             :     integer, pointer, private :: nrowsLinalg(:) => null()
     131                 :             :     integer, pointer, private :: ncolsColsRows_sub(:) => null() ! allows different bandpp per MPI process
     132                 :             :     integer :: nspinor
     133                 :             :     integer :: nrowsColsRows
     134                 :             :     integer :: ncolsColsRows
     135                 :             :     integer :: ncolsLinalg
     136                 :             :     integer :: mpiAlgo
     137                 :             :     integer :: type
     138                 :             :     integer :: me_g0_fft
     139                 :             :     integer :: gpu_option = ABI_GPU_DISABLED
     140                 :             :     integer :: gpu_kokkos_nthrd = 1
     141                 :             :     integer :: gpu_thread_limit = 1
     142                 :             :     logical :: custom_ncolsColsRows = .false. ! use ncolsColsRows_sub instead of ncolsColsRows
     143                 :             :     real(dp), contiguous, pointer:: buffer(:,:) => null()
     144                 :             :   end type xgTransposer_t
     145                 :             : 
     146                 :             :   public :: xgTransposer_constructor
     147                 :             :   public :: xgTransposer_copyConstructor
     148                 :             :   public :: xgTransposer_transpose
     149                 :             :   public :: xgTransposer_getRank
     150                 :             :   public :: xgTransposer_getComm
     151                 :             :   public :: xgTransposer_free
     152                 :             : 
     153                 :             :   contains
     154                 :             : !!***
     155                 :             : 
     156                 :             : !!****f* m_xgTransposer/xgTransposer_constructor
     157                 :             : !!
     158                 :             : !! NAME
     159                 :             : !! xgTransposer_constructor
     160                 :             : 
     161                 :       46182 :   subroutine xgTransposer_constructor(xgTransposer,xgBlock_linalg,xgBlock_colsrows,nspinor,&
     162                 :             :       state,algo,comm_rows,comm_cols,ncpu_cols,ncpu_rows,me_g0_fft,gpu_option,gpu_thread_limit,&
     163   [ +  +  +  -  :       46182 :       custom_ncolsColsRows,ncolsColsRows_sub,nrowsLinalg_sub)
          -  +  -  +  -  
                -  -  - ]
     164                 :             : 
     165                 :             :     type(xgTransposer_t)   , intent(inout) :: xgTransposer
     166                 :             :     type(xgBlock_t), target, intent(in   ) :: xgBlock_linalg
     167                 :             :     type(xgBlock_t), target, intent(in   ) :: xgBlock_colsrows
     168                 :             :     integer                , intent(in   ) :: comm_rows,comm_cols
     169                 :             :     integer                , intent(in   ) :: ncpu_rows,ncpu_cols
     170                 :             :     integer                , intent(in   ) :: nspinor
     171                 :             :     integer                , intent(in   ) :: state
     172                 :             :     integer                , intent(in   ) :: algo
     173                 :             :     integer                , intent(in   ) :: me_g0_fft
     174                 :             :     integer, optional      , intent(in   ) :: gpu_option,gpu_thread_limit
     175                 :             :     logical, optional      , intent(in   ) :: custom_ncolsColsRows ! query for pointer
     176                 :             :     integer, optional, target, intent(in) :: ncolsColsRows_sub(:)
     177                 :             :     integer, optional, target, intent(in) :: nrowsLinalg_sub(:)
     178                 :             :     integer :: commLinalg
     179                 :             :     integer :: ncols
     180                 :             :     integer :: nrows
     181                 :             :     integer :: ierr
     182                 :             :     integer :: icol
     183                 :             :     integer :: ncpuRows
     184                 :             :     integer :: ncpuCols
     185                 :             :     logical :: custom_ncolsColsRows_
     186                 :             : #if defined HAVE_MPI
     187                 :             :     integer :: comm_rows_,comm_cols_
     188                 :             : #endif
     189                 :             :     character(len=500) :: message
     190                 :             :     double precision :: tsec(2)
     191                 :             : 
     192                 :       46182 :     call timab(tim_init,1,tsec)
     193                 :             : 
     194                 :       46182 :     xgTransposer%type = TRANS_TYPE_CONSTRUCTED
     195                 :             : 
     196                 :       46182 :     xgTransposer%xgBlock_linalg => xgBlock_linalg
     197                 :       46182 :     xgTransposer%xgBlock_colsrows => xgBlock_colsrows
     198                 :       46182 :     xgTransposer%state = state
     199                 :       46182 :     xgTransposer%nspinor = nspinor
     200                 :       46182 :     xgTransposer%me_g0_fft = me_g0_fft
     201                 :       46182 :     xgTransposer%gpu_option = ABI_GPU_DISABLED
     202         [ +  + ]:       46182 :     if(present(gpu_option)) xgTransposer%gpu_option = gpu_option
     203                 :       46182 :     xgTransposer%gpu_thread_limit = 1
     204         [ +  + ]:       46182 :     if(present(gpu_thread_limit)) xgTransposer%gpu_thread_limit = gpu_thread_limit
     205         [ -  + ]:       46182 :     if (state==STATE_COLSROWS) then
     206   [ #  #  #  # ]:           0 :         if (comm_rows==xmpi_comm_null.or.comm_cols==xmpi_comm_null) then
     207                 :           0 :             ABI_ERROR('STATE_COLSROWS requires input comms')
     208                 :             :         end if
     209                 :           0 :         commLinalg = comm_cols
     210         [ +  - ]:       46182 :     else if (state==STATE_LINALG) then
     211                 :       46182 :         commLinalg = comm(xgBlock_linalg)
     212                 :             :     else
     213                 :           0 :         ABI_ERROR("Invalid transposer state")
     214                 :             :     end if
     215                 :       46182 :     xgTransposer%mpiData(MPI_LINALG)%comm = commLinalg
     216                 :       46182 :     xgTransposer%mpiData(MPI_LINALG)%rank = xmpi_comm_rank(commLinalg)
     217                 :       46182 :     xgTransposer%mpiData(MPI_LINALG)%size = xmpi_comm_size(commLinalg)
     218                 :             : 
     219                 :             : #if defined HAVE_MPI
     220   [ +  +  +  - ]:       46182 :     if (comm_rows==xmpi_comm_null.and.comm_cols==xmpi_comm_null) then
     221                 :          10 :       xgTransposer%type = TRANS_TYPE_CONSTRUCTED_NULL_COMM
     222                 :          10 :       ncpuCols = ncpu_cols
     223                 :          10 :       ncpuRows = ncpu_rows
     224         [ -  + ]:       46172 :     else if (comm_rows==xmpi_comm_null) then
     225                 :           0 :       ABI_ERROR("Comm_rows and comm_cols should have the same status : null or defined. Here only comm_rows is null.")
     226         [ -  + ]:       46172 :     else if (comm_cols==xmpi_comm_null) then
     227                 :           0 :       ABI_ERROR("Comm_rows and comm_cols should have the same status : null or defined. Here only comm_cols is null.")
     228                 :             :     else
     229                 :       46172 :       xgTransposer%mpiData(MPI_2DCART)%comm = xgTransposer%mpiData(MPI_LINALG)%comm
     230                 :       46172 :       ncpuRows = xmpi_comm_size(comm_rows)
     231                 :       46172 :       ncpuCols = xmpi_comm_size(comm_cols)
     232                 :             :     end if
     233                 :             : #else
     234                 :             :     ncpuRows = 1
     235                 :             :     ncpuCols = 1
     236                 :             : #endif
     237         [ -  + ]:       46182 :     if ( xgTransposer%mpiData(MPI_LINALG)%size < ncpuCols*ncpuRows ) then
     238                 :           0 :       write(message,'(a,i6,a,i6,a)') "There is not enough MPI processes in the communication (", &
     239                 :           0 :         xgTransposer%mpiData(MPI_LINALG)%size, "). Need at least ", ncpuCols*ncpuRows, " processes"
     240                 :           0 :       ABI_ERROR(message)
     241                 :             :     end if
     242                 :             : 
     243         [ +  - ]:       46182 :     if ( ( algo == TRANS_ALL2ALL .or. algo == TRANS_GATHER ) ) then
     244                 :       46182 :       xgTransposer%mpiAlgo = algo
     245                 :             :     else
     246                 :           0 :       xgTransposer%mpiAlgo = TRANS_ALL2ALL
     247                 :           0 :       ABI_COMMENT("Bad value for transposition MPI_algo. Will use ALLTOALL")
     248                 :             :     end if
     249                 :             : 
     250                 :             :     !if ( xgTransposer%mpiAlgo == TRANS_ALL2ALL ) then
     251                 :             :     !  ABI_COMMENT("Using mpi_alltoall for transposition")
     252                 :             :     !else
     253                 :             :     !  ABI_COMMENT("Using mpi_gatherv for transposition")
     254                 :             :     !end if
     255                 :             : 
     256                 :             :     ! Nullify pointers in order to be able to query them
     257                 :       46182 :     xgTransposer%nrowsLinalg => null()
     258                 :       46182 :     xgTransposer%buffer => null()
     259                 :       46182 :     xgTransposer%ncolsColsRows_sub => null()
     260                 :             : 
     261                 :             :     ! Attention: Do not check condition present() on the pointer because in Fortran
     262                 :             :     ! "An unallocated variable passed as an argument is not PRESENT"
     263                 :             :     ! in other words declared pointers with undefined status cannot be queried
     264                 :       46182 :     custom_ncolsColsRows_ = .false.
     265         [ +  + ]:       46182 :     if (present(custom_ncolsColsRows)) custom_ncolsColsRows_ = custom_ncolsColsRows
     266                 :       46166 :     xgTransposer%custom_ncolsColsRows = custom_ncolsColsRows_
     267         [ +  - ]:          16 :     if (custom_ncolsColsRows_) then
     268         [ -  + ]:          16 :       if ( ncpuRows > 1 ) then
     269                 :           0 :         write(message,'(a,i6)') "Custom ncolsColsRows not implemented for number of row MPI processes=", ncpuRows
     270                 :           0 :         ABI_ERROR(message)
     271                 :             :       end if
     272                 :          16 :       xgTransposer%custom_ncolsColsRows = .true.
     273         [ +  - ]:          16 :       if (state==STATE_LINALG) xgTransposer%ncolsColsRows_sub => ncolsColsRows_sub
     274                 :             :       ! if state colsrows in will be constructed later in readDistribution
     275                 :             :     end if
     276                 :             : 
     277                 :             :     ! The following allocates memory buffer associated either to xgBlock_linalg or to xgBlock_colsrows
     278                 :             :     !! Assumes that only one of them is already allocated
     279                 :       46182 :     select case (state)
     280                 :             :     case (STATE_LINALG)
     281                 :             : 
     282                 :             :       ! We are in the linalg representation.
     283                 :             :       ! We need to construct the colsrows parallelization
     284                 :       46182 :       call xgBlock_getSize(xgBlock_linalg,nrows,ncols)
     285                 :             : 
     286                 :             :       ! Get total number of rows in ColsRows representation
     287                 :       46182 :       call xmpi_sum(nrows,commLinalg,ierr)
     288                 :             : 
     289                 :             :       ! Check that ncols can be uniformly distributed across ncolsCols MPI processes
     290                 :             :       ! Load balance is uniform, each MPI process has the same ncolsColsRows
     291   [ +  +  -  + ]:       46182 :       if ( (.not.xgTransposer%custom_ncolsColsRows) .and. MOD(ncols,ncpuCols) /=0 ) then
     292         [ #  # ]:           0 :         if ( ncols > ncpuCols ) then
     293                 :           0 :           write(message,'(a,i6,a,i6,a)') "Unbalanced parallelization : ", ncols, " columns for ", ncpuCols, " MPI"
     294                 :           0 :           ABI_ERROR(message)
     295                 :             :         else
     296                 :           0 :           write(message,'(i6,a)') (ncpuCols-ncols)*ncpuRows, " MPI will not be used because of the number of columns!!"
     297                 :           0 :           ABI_ERROR(message)
     298                 :             :           !ncpuCols = ncols
     299                 :             :         end if
     300                 :             :       end if
     301                 :             : 
     302                 :             :       ! Build the lookup table that maps each column to an MPI process
     303   [ -  +  +  -  :      138546 :       ABI_MALLOC(xgTransposer%lookup,(1:ncols))
          +  -  +  -  -  
             +  -  +  -  
                      + ]
     304                 :             :       !ABI_MALLOC(xgTransposer%me_g0_lookup,())
     305         [ +  + ]:      476942 :       do icol = 0, ncols-1
     306                 :      476942 :         xgTransposer%lookup(icol+1) = MOD(icol,ncpuCols)
     307                 :             :       end do
     308                 :             : 
     309                 :             : #if defined HAVE_MPI
     310   [ +  +  +  - ]:       46182 :       if (comm_rows==xmpi_comm_null.and.comm_cols==xmpi_comm_null) then
     311                 :          10 :         call xgTransposer_makeComm(xgTransposer,ncpuRows,ncpuCols,comm_rows_,comm_cols_)
     312                 :             :       else
     313                 :       46172 :         comm_rows_=comm_rows
     314                 :       46172 :         comm_cols_=comm_cols
     315                 :             :       end if
     316                 :       46182 :       xgTransposer%mpiData(MPI_ROWS)%comm = comm_rows_
     317                 :       46182 :       xgTransposer%mpiData(MPI_COLS)%comm = comm_cols_
     318                 :             : #else
     319                 :             :       xgTransposer%mpiData(MPI_ROWS)%comm = xmpi_comm_null
     320                 :             :       xgTransposer%mpiData(MPI_COLS)%comm = xmpi_comm_null
     321                 :             : #endif
     322                 :       46182 :       call xgTransposer_setComm(xgTransposer)
     323                 :       46182 :       call xgTransposer_computeDistribution(xgTransposer)
     324                 :       92364 :       call xgTransposer_makeXgBlock(xgTransposer)
     325                 :             : 
     326                 :             :     case (STATE_COLSROWS)
     327                 :             : 
     328                 :             :       !write(std_out,'(a,i2,i2)') 'db @xgTransposer constructor enter case for # procs', ncpuRows,ncpuCols
     329                 :             :       !call flush_unit(std_out)
     330                 :             : 
     331                 :             :       ! We are in the colsrows representation.
     332                 :             :       ! We need to construct the linalg parallelization
     333                 :             : 
     334                 :             :       ! Load balance
     335                 :             :       ! nrows are distributed uniformly across ncpuCols MPI processes for 1,...,ncpuCols-1
     336                 :             :       ! The last MPI process uses remaining nrows. Read distribution from nrowsLinalg_sub
     337                 :           0 :       call xgBlock_getSize(xgBlock_colsrows,nrows,ncols)
     338   [ #  #  #  #  :           0 :       ABI_MALLOC(xgTransposer%lookup,(1:ncols))
          #  #  #  #  #  
             #  #  #  #  
                      # ]
     339                 :             : 
     340                 :             : #if defined HAVE_MPI
     341                 :           0 :       xgTransposer%mpiData(MPI_ROWS)%comm = comm_rows
     342                 :           0 :       xgTransposer%mpiData(MPI_COLS)%comm = comm_cols
     343                 :             : #else
     344                 :             :       xgTransposer%mpiData(MPI_ROWS)%comm = xmpi_comm_null
     345                 :             :       xgTransposer%mpiData(MPI_COLS)%comm = xmpi_comm_null
     346                 :             : #endif
     347                 :             : 
     348                 :           0 :       call xgTransposer_setComm(xgTransposer)
     349   [ #  #  #  #  :           0 :       ABI_MALLOC(xgTransposer%nrowsLinalg,(xgTransposer%mpiData(MPI_LINALG)%size))
          #  #  #  #  #  
                #  #  # ]
     350         [ #  # ]:           0 :       xgTransposer%nrowsLinalg(:) = nrowsLinalg_sub(:)
     351                 :           0 :       call xgTransposer_readDistribution(xgTransposer)
     352                 :           0 :       call xgTransposer_makeXgBlock(xgTransposer)
     353                 :             : 
     354                 :             :     case default
     355      [ +  -  - ]:       46182 :       ABI_ERROR("State is undefined")
     356                 :             :     end select
     357                 :             : 
     358                 :       46182 :     call timab(tim_init,2,tsec)
     359                 :             : 
     360                 :             :     !write(*,*) 'transposer construction done'
     361                 :             : 
     362                 :       46182 :   end subroutine xgTransposer_constructor
     363                 :             : !!***
     364                 :             : 
     365                 :             : !!****f* m_xgTransposer/xgTransposer_copyConstructor
     366                 :             : !!
     367                 :             : !! NAME
     368                 :             : !! xgTransposer_copyConstructor
     369                 :             : 
     370                 :      144632 :   subroutine xgTransposer_copyConstructor(xgTransposer,xgTransposerInitialized,xgBlock_linalg,xgBlock_colsrows,state)
     371                 :             : 
     372                 :             :     type(xgTransposer_t)   , intent(inout) :: xgTransposer
     373                 :             :     type(xgTransposer_t)   , intent(in   ) :: xgTransposerInitialized
     374                 :             :     type(xgBlock_t), target, intent(in   ) :: xgBlock_linalg
     375                 :             :     type(xgBlock_t), target, intent(in   ) :: xgBlock_colsrows
     376                 :             :     integer                , intent(in   ) :: state
     377                 :             :     integer :: commLinalg
     378                 :             :     integer :: ncols, ncpuCols
     379                 :             :     integer :: nrows, ncpuRows
     380                 :             :     integer :: ierr
     381                 :             :     integer :: icol
     382                 :             :     character(len=500) :: message
     383                 :             :     double precision :: tsec(2)
     384                 :             : 
     385                 :      144632 :     call timab(tim_init,1,tsec)
     386                 :             : 
     387                 :      144632 :     xgTransposer%type = TRANS_TYPE_COPIED
     388                 :             : 
     389                 :      144632 :     xgTransposer%xgBlock_linalg => xgBlock_linalg
     390                 :      144632 :     xgTransposer%xgBlock_colsrows => xgBlock_colsrows
     391                 :      144632 :     xgTransposer%state = state
     392         [ -  + ]:      144632 :     if (state==STATE_COLSROWS) then
     393                 :           0 :         commLinalg = xgTransposerInitialized%mpiData(MPI_LINALG)%comm
     394         [ +  - ]:      144632 :     else if (state==STATE_LINALG) then
     395                 :      144632 :         commLinalg = comm(xgBlock_linalg)
     396                 :             :     else
     397                 :           0 :         ABI_ERROR("Invalid transposer state")
     398                 :             :     end if
     399                 :             : 
     400         [ -  + ]:      144632 :     if ( commLinalg /= xgTransposerInitialized%mpiData(MPI_LINALG)%comm ) then
     401                 :           0 :       ABI_ERROR("Linalg communicators are different for the two transposers, this is not allowed.")
     402                 :             :     end if
     403                 :             : 
     404         [ +  + ]:      723160 :     xgTransposer%mpiData(:)%comm = xgTransposerInitialized%mpiData(:)%comm
     405                 :      144632 :     xgTransposer%mpiData(MPI_LINALG)%rank = xmpi_comm_rank(commLinalg)
     406                 :      144632 :     xgTransposer%mpiData(MPI_LINALG)%size = xmpi_comm_size(commLinalg)
     407                 :      144632 :     xgTransposer%mpiData(MPI_COLS)%rank = xmpi_comm_rank(xgTransposer%mpiData(MPI_COLS)%comm)
     408                 :      144632 :     xgTransposer%mpiData(MPI_COLS)%size = xmpi_comm_size(xgTransposer%mpiData(MPI_COLS)%comm)
     409                 :      144632 :     xgTransposer%mpiData(MPI_ROWS)%rank = xmpi_comm_rank(xgTransposer%mpiData(MPI_ROWS)%comm)
     410                 :      144632 :     xgTransposer%mpiData(MPI_ROWS)%size = xmpi_comm_size(xgTransposer%mpiData(MPI_ROWS)%comm)
     411                 :             : 
     412                 :      144632 :     xgTransposer%mpiAlgo = xgTransposerInitialized%mpiAlgo
     413                 :      144632 :     xgTransposer%nspinor = xgTransposerInitialized%nspinor
     414                 :      144632 :     xgTransposer%me_g0_fft = xgTransposerInitialized%me_g0_fft
     415                 :      144632 :     xgTransposer%gpu_option = xgTransposerInitialized%gpu_option
     416                 :             : 
     417                 :             :     ! Nullify pointers in order to be able to query them
     418                 :      144632 :     xgTransposer%nrowsLinalg => null()
     419                 :      144632 :     xgTransposer%buffer => null()
     420                 :      144632 :     xgTransposer%ncolsColsRows_sub => null()
     421                 :             :     !write(std_out,*) '@nrowsLinalg', associated(xgTransposer%nrowsLinalg)
     422                 :             : 
     423                 :      144632 :     xgTransposer%custom_ncolsColsRows = xgTransposerInitialized%custom_ncolsColsRows
     424   [ -  +  -  - ]:      144632 :     if(xgTransposer%custom_ncolsColsRows .and. state==STATE_LINALG) then
     425                 :           0 :         xgTransposer%ncolsColsRows_sub => xgTransposerInitialized%ncolsColsRows_sub
     426                 :             :         ! if state colsrows in will be constructed later in readDistribution
     427                 :             :     end if
     428                 :             : 
     429                 :      144632 :     ncpuCols = xgTransposer%mpiData(MPI_COLS)%size
     430                 :      144632 :     ncpuRows = xgTransposer%mpiData(MPI_ROWS)%size
     431                 :             : 
     432                 :      144632 :     select case (state)
     433                 :             :     case (STATE_LINALG)
     434                 :             : 
     435                 :      144632 :       call xgBlock_getSize(xgBlock_linalg,nrows,ncols)
     436                 :             : 
     437                 :             :       ! Get total number of rows in ColsRows representation
     438                 :      144632 :       call xmpi_sum(nrows,commLinalg,ierr)
     439                 :             : 
     440                 :             :       ! Check that ncols can be uniformly distributed across ncolsCols MPI processes
     441   [ +  -  -  + ]:      144632 :       if ( (.not.xgTransposer%custom_ncolsColsRows) .and. MOD(ncols,ncpuCols) /=0 ) then
     442         [ #  # ]:           0 :         if ( ncols > ncpuCols ) then
     443                 :           0 :           write(message,'(a,i6,a,i6,a)') "Unbalanced parallelization : ", ncols, " columns for ", ncpuCols, " MPI"
     444                 :           0 :           ABI_ERROR(message)
     445                 :             :         else
     446                 :           0 :           write(message,'(i6,a)') (ncpuCols-ncols)*ncpuRows, " MPI will not be used  because of the number of columns!!"
     447                 :           0 :           ABI_ERROR(message)
     448                 :             :         end if
     449                 :             :       end if
     450                 :             : 
     451                 :             :       ! Build the lookup table that maps each column to an MPI process
     452                 :             :       ! Assures load balance so that each MPI process has the same charge
     453   [ -  +  +  -  :      433896 :       ABI_MALLOC(xgTransposer%lookup,(1:ncols))
          +  -  +  -  -  
             +  -  +  -  
                      + ]
     454         [ -  + ]:      144632 :       if ( cols(xgTransposerInitialized%xgBlock_linalg) /= ncols ) then
     455         [ #  # ]:           0 :         do icol = 0, ncols-1
     456                 :           0 :             xgTransposer%lookup(icol+1) = MOD(icol,ncpuCols)
     457                 :             :         end do
     458                 :             :       else
     459         [ +  + ]:     1172400 :         xgTransposer%lookup(:) = xgTransposerInitialized%lookup(:)
     460                 :             :       end if
     461                 :             : 
     462                 :      144632 :       call xgTransposer_computeDistribution(xgTransposer)
     463                 :      289264 :       call xgTransposer_makeXgBlock(xgTransposer)
     464                 :             : 
     465                 :             :     case (STATE_COLSROWS)
     466                 :             : 
     467                 :             :       !write(std_out,*) 'db @xgTransposer copyConstructor enter case'
     468                 :             :       !call flush_unit(std_out)
     469                 :             : 
     470                 :           0 :       nrows = rows(xgBlock_colsrows)
     471                 :           0 :       ncols = cols(xgBlock_colsrows)
     472   [ #  #  #  #  :           0 :       if ( (.not.xgTransposer%custom_ncolsColsRows) .and. MOD(nrows,ncpuCols) /=0 .and. nrows > ncpuCols) then
                   #  # ]
     473                 :           0 :         write(message,'(a,i6,a,i6,a)') "Unbalanced parallelization : ", nrows, " rows for ", ncpuCols, " MPI"
     474                 :           0 :         ABI_COMMENT(message)
     475                 :             :       end if
     476   [ #  #  #  #  :           0 :       ABI_MALLOC(xgTransposer%lookup,(1:ncols))
          #  #  #  #  #  
             #  #  #  #  
                      # ]
     477                 :             : 
     478   [ #  #  #  #  :           0 :       ABI_MALLOC(xgTransposer%nrowsLinalg,(xgTransposer%mpiData(MPI_LINALG)%size))
          #  #  #  #  #  
                #  #  # ]
     479         [ #  # ]:           0 :       xgTransposer%nrowsLinalg(:) = xgTransposerInitialized%nrowsLinalg(:)
     480                 :           0 :       call xgTransposer_readDistribution(xgTransposer)
     481                 :           0 :       call xgTransposer_makeXgBlock(xgTransposer)
     482                 :             : 
     483                 :             :     case default
     484      [ +  -  - ]:      144632 :       ABI_ERROR("State is undefined")
     485                 :             :     end select
     486                 :             : 
     487                 :      144632 :     call timab(tim_init,2,tsec)
     488                 :             : 
     489                 :      144632 :   end subroutine xgTransposer_copyConstructor
     490                 :             : !!***
     491                 :             : 
     492                 :             : !!****f* m_xgTransposer/xgTransposer_setComm
     493                 :             : !!
     494                 :             : !! NAME
     495                 :             : !! xgTransposer_setComm
     496                 :             : 
     497                 :       46182 :   subroutine xgTransposer_setComm(xgTransposer)
     498                 :             : 
     499                 :             :     type(xgTransposer_t), intent(inout) :: xgTransposer
     500                 :             : 
     501                 :       46182 :     xgTransposer%mpiData(MPI_ROWS)%rank = xmpi_comm_rank(xgTransposer%mpiData(MPI_ROWS)%comm)
     502                 :       46182 :     xgTransposer%mpiData(MPI_ROWS)%size = xmpi_comm_size(xgTransposer%mpiData(MPI_ROWS)%comm)
     503                 :       46182 :     xgTransposer%mpiData(MPI_COLS)%rank = xmpi_comm_rank(xgTransposer%mpiData(MPI_COLS)%comm)
     504                 :       46182 :     xgTransposer%mpiData(MPI_COLS)%size = xmpi_comm_size(xgTransposer%mpiData(MPI_COLS)%comm)
     505                 :             : 
     506                 :       92364 :     select case(xgTransposer%state)
     507                 :             :     case (STATE_LINALG)
     508                 :       46182 :         call xgBlock_setComm(xgTransposer%xgBlock_colsrows,xgTransposer%mpiData(MPI_ROWS)%comm)
     509                 :             :     case (STATE_COLSROWS)
     510                 :           0 :         call xgBlock_setComm(xgTransposer%xgBlock_linalg,xgTransposer%mpiData(MPI_ROWS)%comm)
     511                 :             :     case default
     512      [ +  -  - ]:       46182 :       ABI_ERROR("State is undefined")
     513                 :             :     end select
     514                 :             : 
     515                 :       46182 :   end subroutine xgTransposer_setComm
     516                 :             : !!***
     517                 :             : 
     518                 :             : !!****f* m_xgTransposer/xgTransposer_makeComm
     519                 :             : !!
     520                 :             : !! NAME
     521                 :             : !! xgTransposer_makeComm
     522                 :             : 
     523                 :          10 :   subroutine xgTransposer_makeComm(xgTransposer,ncpuRows,ncpuCols,comm_rows,comm_cols)
     524                 :             : 
     525                 :             :     type(xgTransposer_t), intent(inout) :: xgTransposer
     526                 :             :     integer             , intent(in   ) :: ncpuRows
     527                 :             :     integer             , intent(in   ) :: ncpuCols
     528                 :             :     integer             , intent(out  ) :: comm_rows
     529                 :             :     integer             , intent(out  ) :: comm_cols
     530                 :             :     integer :: commColsRows
     531                 :             : #if defined HAVE_MPI
     532                 :             :     integer :: sizeGrid(2) !coordInGrid(2),
     533                 :             :     logical :: periodic(2), selectDim(2), reorder
     534                 :             :     integer :: ierr
     535                 :             : 
     536                 :          10 :     sizeGrid(1) = ncpuCols
     537                 :          10 :     sizeGrid(2) = ncpuRows
     538                 :             :     !sizeGrid(1) = ncpuRows
     539                 :             :     !sizeGrid(2) = ncpuCols
     540                 :          10 :     periodic = (/ .false., .false. /)
     541                 :          10 :     reorder  = .false.
     542                 :          10 :     call mpi_cart_create(xgTransposer%mpiData(MPI_LINALG)%comm,2,sizeGrid,periodic,reorder,commColsRows,ierr)
     543         [ -  + ]:          10 :     if ( ierr /= xmpi_success ) then
     544                 :           0 :       ABI_ERROR("xgTransposer failed to creat cartesian grid")
     545                 :             :     end if
     546                 :          10 :     xgTransposer%mpiData(MPI_2DCART)%comm = commColsRows
     547                 :             : 
     548                 :          10 :     selectDim = (/ .false., .true. /)
     549                 :             :     !selectDim = (/ .true., .false. /)
     550                 :          10 :     call mpi_cart_sub(commColsRows, selectDim, xgTransposer%mpiData(MPI_ROWS)%comm,ierr)
     551         [ -  + ]:          10 :     if ( ierr /= xmpi_success ) then
     552                 :           0 :       ABI_ERROR("xgTransposer failed to creat rows communicator")
     553                 :             :     end if
     554                 :          10 :     selectDim = (/ .true., .false. /)
     555                 :             :     !selectDim = (/ .false., .true. /)
     556                 :          10 :     call mpi_cart_sub(commColsRows, selectDim, xgTransposer%mpiData(MPI_COLS)%comm,ierr)
     557         [ -  + ]:          10 :     if ( ierr /= xmpi_success ) then
     558                 :           0 :       ABI_ERROR("xgTransposer failed to creat columns communicator")
     559                 :             :     end if
     560                 :             : #else
     561                 :             :     commColsRows = xmpi_comm_null
     562                 :             :     xgTransposer%mpiData(MPI_ROWS)%comm = xmpi_comm_null
     563                 :             :     xgTransposer%mpiData(MPI_COLS)%comm = xmpi_comm_null
     564                 :             : #endif
     565                 :          10 :     comm_rows = xgTransposer%mpiData(MPI_ROWS)%comm
     566                 :          10 :     comm_cols = xgTransposer%mpiData(MPI_COLS)%comm
     567                 :             : 
     568                 :          10 :   end subroutine xgTransposer_makeComm
     569                 :             : !!***
     570                 :             : 
     571                 :             : !!****f* m_xgTransposer/xgTransposer_computeDistribution
     572                 :             : !!
     573                 :             : !! NAME
     574                 :             : !! xgTransposer_computeDistribution
     575                 :             : 
     576                 :      190814 :   subroutine xgTransposer_computeDistribution(xgTransposer)
     577                 :             : 
     578                 :             :     type(xgTransposer_t), intent(inout) :: xgTransposer
     579                 :             :     integer :: nRealPairs
     580                 :             :     integer :: ierr
     581                 :             :     integer :: icpu_cols,icpu_rows
     582                 :             :     integer :: ncpuCols,ncpuRows
     583                 :             :     integer :: ncolsColsRowsMe
     584                 :             : 
     585   [ -  +  +  -  :      572442 :     ABI_MALLOC(xgTransposer%nrowsLinalg,(xgTransposer%mpiData(MPI_LINALG)%size))
          +  -  +  -  -  
                +  -  + ]
     586                 :      190814 :     nRealPairs = rows(xgTransposer%xgBlock_linalg)
     587         [ -  + ]:      190814 :     if (MOD(nRealPairs,xgTransposer%nspinor)/=0) then
     588                 :           0 :         ABI_ERROR('nspinor should divide nRealPairs!')
     589                 :             :     end if
     590                 :             : 
     591                 :      190814 :     call xmpi_allgather(nRealPairs,xgTransposer%nrowsLinalg,xgTransposer%mpiData(MPI_LINALG)%comm,ierr)
     592         [ -  + ]:      190814 :     if ( ierr /= xmpi_success ) then
     593                 :           0 :         ABI_ERROR("Error while gathering number of rows in linalg")
     594                 :             :     end if
     595                 :             : 
     596                 :      190814 :     ncpuCols = xgTransposer%mpiData(MPI_COLS)%size
     597                 :      190814 :     ncpuRows = xgTransposer%mpiData(MPI_ROWS)%size
     598                 :      190814 :     icpu_rows = xgTransposer%mpiData(MPI_ROWS)%rank
     599                 :      190814 :     xgTransposer%nrowsColsRows = 0
     600         [ +  + ]:      688884 :     do icpu_cols=0,ncpuCols-1
     601                 :      688884 :         xgTransposer%nrowsColsRows = xgTransposer%nrowsColsRows + xgTransposer%nrowsLinalg(1+icpu_rows+icpu_cols*ncpuRows)
     602                 :             :     end do
     603                 :      190814 :     xgTransposer%ncolsColsRows = cols(xgTransposer%xgBlock_linalg)/ncpuCols
     604         [ +  + ]:      190814 :     if (xgTransposer%custom_ncolsColsRows) then
     605                 :          16 :       ncolsColsRowsMe = xgTransposer%ncolsColsRows_sub(xgTransposer%mpiData(MPI_LINALG)%rank+1)
     606                 :          16 :       xgTransposer%ncolsColsRows = ncolsColsRowsMe ! overwrite
     607                 :             :     end if
     608                 :             : 
     609                 :             :     !write(*,*) "In linalg, # of real pairs:", xgTransposer%nrowsLinalg
     610                 :             :     !write(*,*) "In rows, # of real pairs for proc ", xgTransposer%mpiData(MPI_ROWS)%rank, ":", xgTransposer%nrowsColsRows
     611                 :             :     !write(*,*) "In cols, # of cols for proc ", xgTransposer%mpiData(MPI_COLS)%rank, ":", xgTransposer%ncolsColsRows
     612                 :             : 
     613                 :      190814 :   end subroutine xgTransposer_computeDistribution
     614                 :             : !!***
     615                 :             : 
     616                 :             : !!****f* m_xgTransposer/xgTransposer_readDistribution
     617                 :             : !!
     618                 :             : !! NAME
     619                 :             : !! xgTransposer_readDistribution
     620                 :             : !!
     621                 :             : !! FUNCTION
     622                 :             : !! Read custom distribution from xgTransposer%nrowsLinalg(:)
     623                 :             : 
     624                 :           0 :   subroutine xgTransposer_readDistribution(xgTransposer)
     625                 :             : 
     626                 :             :       ! Arguments
     627                 :             :       type(xgTransposer_t), intent(inout) :: xgTransposer
     628                 :             :       ! Variables
     629                 :             :       integer :: iproc,nRealPairs,tot_ncols,ncols,ierr
     630                 :             :       integer :: nprocs,commLinalg
     631                 :             : 
     632                 :             :       ! ***********************************************************
     633                 :             : 
     634                 :           0 :       nprocs = xgTransposer%mpiData(MPI_LINALG)%size
     635                 :           0 :       xgTransposer%nrowsColsRows = rows(xgTransposer%xgBlock_colsrows)
     636                 :           0 :       ncols = cols(xgTransposer%xgBlock_colsrows)
     637                 :           0 :       xgTransposer%ncolsColsRows = ncols ! for current MPI process
     638                 :             : 
     639         [ #  # ]:           0 :       if(xgTransposer%custom_ncolsColsRows) then
     640   [ #  #  #  #  :           0 :           ABI_MALLOC(xgTransposer%ncolsColsRows_sub,(nprocs))
          #  #  #  #  #  
                #  #  # ]
     641                 :             : 
     642                 :           0 :           commLinalg = xgTransposer%mpiData(MPI_LINALG)%comm
     643                 :           0 :           ncols = cols(xgTransposer%xgBlock_colsrows)
     644                 :           0 :           call xmpi_allgather(ncols,xgTransposer%ncolsColsRows_sub,commLinalg,ierr)
     645         [ #  # ]:           0 :           if ( ierr /= xmpi_success ) then
     646                 :           0 :               ABI_ERROR("Error while gathering number of columns in colsrows")
     647                 :             :           end if
     648                 :           0 :           call xmpi_sum(ncols,commLinalg,ierr)
     649         [ #  # ]:           0 :           if ( ierr /= xmpi_success ) then
     650                 :           0 :               ABI_ERROR("Error while summing number of columns in colsrows")
     651                 :             :           end if
     652                 :           0 :           tot_ncols = ncols ! update sum
     653                 :             :       else
     654                 :           0 :           tot_ncols = ncols*xgTransposer%mpiData(MPI_COLS)%size
     655                 :             :       end if
     656                 :             : 
     657         [ #  # ]:           0 :       do iproc=1,size(xgTransposer%nrowsLinalg)
     658                 :           0 :         nRealPairs = xgTransposer%nrowsLinalg(iproc)
     659         [ #  # ]:           0 :         if (MOD(nRealPairs,xgTransposer%nspinor)/=0) then
     660                 :           0 :             ABI_ERROR('nspinor should divide nRealPairs!')
     661                 :             :         end if
     662                 :             :       end do
     663                 :             : 
     664                 :           0 :       xgTransposer%ncolsLinalg = tot_ncols
     665                 :             : 
     666                 :             :       !write(*,*) "In linalg, # of real pairs:", xgTransposer%nrowsLinalg
     667                 :             :       !write(*,*) "In linalg, # of cols:", xgTransposer%ncolsLinalg
     668                 :             :       !write(*,*) "In rows, # of real pairs for proc ", xgTransposer%mpiData(MPI_ROWS)%rank, ":", xgTransposer%nrowsColsRows
     669                 :             :       !write(*,*) "In cols, # of cols for proc ", xgTransposer%mpiData(MPI_COLS)%rank, ":", xgTransposer%ncolsColsRows
     670                 :             : 
     671                 :           0 :   end subroutine xgTransposer_readDistribution
     672                 :             : !!***
     673                 :             : 
     674                 :             : !!****f* m_xgTransposer/xgTransposer_makeXgBlock
     675                 :             : !!
     676                 :             : !! NAME
     677                 :             : !! xgTransposer_makeXgBlock
     678                 :             : 
     679                 :      190814 :   subroutine xgTransposer_makeXgBlock(xgTransposer)
     680                 :             : 
     681                 :             :     type(xgTransposer_t), intent(inout) :: xgTransposer
     682                 :             :     integer :: ncolsColsRows, nrowsLinalgMe
     683                 :             : 
     684                 :      381628 :     select case (xgTransposer%state)
     685                 :             :     case (STATE_LINALG)
     686                 :             : 
     687                 :      190814 :       ncolsColsRows = xgTransposer%ncolsColsRows
     688                 :             :       !write(std_out,*) '@makeXgBlock dims', ncolsColsRows, xgTransposer%nrowsColsRows
     689                 :             :       !call flush_unit(std_out)
     690                 :             : 
     691                 :             :       ! Assume xgBlock_colsrows is empty and not constructed because user cannot
     692                 :             :       ! predict the size
     693         [ -  + ]:      190814 :       if ( associated(xgTransposer%buffer) ) then
     694         [ #  # ]:           0 :         if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
     695                 :             : #if defined HAVE_GPU && defined HAVE_YAKL
     696                 :             :           ABI_FREE_MANAGED(xgTransposer%buffer)
     697                 :             : #endif
     698                 :             :         else
     699                 :             :           if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
     700                 :             : #ifdef HAVE_OPENMP_OFFLOAD
     701                 :             : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
     702                 :             :             !$OMP TARGET EXIT DATA MAP(delete:xgTransposer%buffer)
     703                 :             : #else
     704                 :             :             call xg_ompgpu_map_delete(xgTransposer%buffer)
     705                 :             : #endif
     706                 :             : #endif
     707                 :             :           end if
     708                 :           0 :           ABI_FREE(xgTransposer%buffer)
     709                 :             :         end if
     710                 :             :       end if
     711                 :             : 
     712         [ +  + ]:      190814 :       if ( xgTransposer%mpiData(MPI_COLS)%size == 1 ) then
     713                 :       38214 :         xgTransposer%xgBlock_colsrows = xgTransposer%xgBlock_linalg
     714                 :             :       else
     715         [ +  - ]:      152600 :         if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
     716                 :             : #if defined HAVE_GPU && defined HAVE_YAKL
     717                 :             :           ABI_MALLOC_MANAGED(xgTransposer%buffer,(/2,ncolsColsRows*xgTransposer%nrowsColsRows/))
     718                 :             : #endif
     719                 :             :         else
     720   [ -  +  +  -  :      457800 :           ABI_MALLOC(xgTransposer%buffer,(2,ncolsColsRows*xgTransposer%nrowsColsRows))
          +  -  +  -  -  
                +  -  + ]
     721                 :             :         end if
     722                 :             :         if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
     723                 :             : #ifdef HAVE_OPENMP_OFFLOAD
     724                 :             : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
     725                 :             :           !$OMP TARGET ENTER DATA MAP(alloc:xgTransposer%buffer)
     726                 :             : #else
     727                 :             :           call xg_ompgpu_map_alloc(xgTransposer%buffer)
     728                 :             : #endif
     729                 :             : #endif
     730                 :             :         end if
     731                 :             :         call xgBlock_map(xgTransposer%xgBlock_colsrows,xgTransposer%buffer,space(xgTransposer%xgBlock_linalg),&
     732                 :             :           xgTransposer%nrowsColsRows,ncolsColsRows,xgTransposer%mpiData(MPI_ROWS)%comm,&
     733                 :      152600 :           me_g0=xgTransposer%me_g0_fft,gpu_option=xgTransposer%gpu_option)
     734                 :             :       end if
     735                 :             :     case (STATE_COLSROWS)
     736                 :             : 
     737                 :             :       ! Assume xgBlock_linalg is empty and not constructed
     738                 :             : 
     739                 :             :       !write(std_out,*) 'db using @makeXgBlock not empty'
     740                 :             :       !call flush_unit(std_out)
     741                 :             : 
     742                 :           0 :       nrowsLinalgMe = xgTransposer%nrowsLinalg(xgTransposer%mpiData(MPI_LINALG)%rank+1)
     743                 :             : 
     744         [ #  # ]:           0 :       if ( associated(xgTransposer%buffer) ) then
     745         [ #  # ]:           0 :         if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
     746                 :             : #if defined HAVE_GPU && defined HAVE_YAKL
     747                 :             :           ABI_FREE_MANAGED(xgTransposer%buffer)
     748                 :             : #endif
     749                 :             :         else
     750                 :             :           if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
     751                 :             : #ifdef HAVE_OPENMP_OFFLOAD
     752                 :             : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
     753                 :             :             !$OMP TARGET EXIT DATA MAP(delete:xgTransposer%buffer)
     754                 :             : #else
     755                 :             :             call xg_ompgpu_map_delete(xgTransposer%buffer)
     756                 :             : #endif
     757                 :             : #endif
     758                 :             :           end if
     759                 :           0 :           ABI_FREE(xgTransposer%buffer)
     760                 :             :         end if
     761                 :             :       end if
     762                 :             : 
     763         [ #  # ]:           0 :       if ( xgTransposer%mpiData(MPI_COLS)%size == 1 ) then
     764                 :           0 :         xgTransposer%xgBlock_linalg = xgTransposer%xgBlock_colsrows
     765                 :             :       else
     766         [ #  # ]:           0 :         if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
     767                 :             : #if defined HAVE_GPU && defined HAVE_YAKL
     768                 :             :           ABI_MALLOC_MANAGED(xgTransposer%buffer,(/2,xgTransposer%ncolsLinalg*nrowsLinalgMe/))
     769                 :             : #endif
     770                 :             :         else
     771   [ #  #  #  #  :           0 :           ABI_MALLOC(xgTransposer%buffer,(2,xgTransposer%ncolsLinalg*nrowsLinalgMe))
          #  #  #  #  #  
                #  #  # ]
     772                 :             :         end if
     773                 :             :         if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
     774                 :             : #ifdef HAVE_OPENMP_OFFLOAD
     775                 :             : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
     776                 :             :           !$OMP TARGET ENTER DATA MAP(alloc:xgTransposer%buffer)
     777                 :             : #else
     778                 :             :           call xg_ompgpu_map_alloc(xgTransposer%buffer)
     779                 :             : #endif
     780                 :             : #endif
     781                 :             :         end if
     782                 :             :         call xgBlock_map(xgTransposer%xgBlock_linalg,xgTransposer%buffer,space(xgTransposer%xgBlock_colsrows),&
     783                 :             :           nrowsLinalgMe,xgTransposer%ncolsLinalg,xgTransposer%mpiData(MPI_COLS)%comm,&
     784                 :           0 :           me_g0=xgTransposer%me_g0_fft,gpu_option=xgTransposer%gpu_option)
     785                 :             :       end if
     786                 :             : 
     787                 :             :     case default
     788      [ +  -  - ]:      190814 :       ABI_ERROR("State unknown")
     789                 :             :     end select
     790                 :             : 
     791                 :             :     !if (associated(xgTransposer%buffer)) then
     792                 :             :     !    write(std_out,*) '@makeXgBlock alloc buffer size', size(xgTransposer%buffer,1), size(xgTransposer%buffer,2)
     793                 :             :     !    call flush_unit(std_out)
     794                 :             :     !end if
     795                 :             : 
     796                 :             : #if defined HAVE_OPENMP_OFFLOAD && ! defined HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
     797                 :             :  contains
     798                 :             :    ! NOTE : routines to be used in order to avoid using derived-type syntax in OMP directives
     799                 :             :    subroutine xg_ompgpu_map_alloc(xgTransposerBuffer)
     800                 :             :     real(dp),intent(inout) :: xgTransposerBuffer(:,:)
     801                 :             :     !$OMP TARGET ENTER DATA MAP(alloc:xgTransposerBuffer)
     802                 :             :    end subroutine xg_ompgpu_map_alloc
     803                 :             : 
     804                 :             :    subroutine xg_ompgpu_map_delete(xgTransposerBuffer)
     805                 :             :     real(dp),intent(inout) :: xgTransposerBuffer(:,:)
     806                 :             :     !$OMP TARGET EXIT DATA MAP(delete:xgTransposerBuffer)
     807                 :             :    end subroutine xg_ompgpu_map_delete
     808                 :             : #endif
     809                 :      190814 :   end subroutine xgTransposer_makeXgBlock
     810                 :             : !!***
     811                 :             : 
     812                 :             : !!****f* m_xgTransposer/xgTransposer_transpose
     813                 :             : !!
     814                 :             : !! NAME
     815                 :             : !! xgTransposer_transpose
     816                 :             : 
     817                 :     2853804 :   subroutine xgTransposer_transpose(xgTransposer,toState)
     818                 :             : 
     819                 :             :     type(xgTransposer_t), intent(inout) :: xgTransposer
     820                 :             :     integer             , intent(in   ) :: toState
     821                 :             :     double precision :: tsec(2)
     822                 :             : 
     823                 :     2853804 :     call timab(tim_transpose,1,tsec)
     824                 :             : 
     825         [ -  + ]:     2853804 :     if ( toState /= STATE_LINALG .and. toState /= STATE_COLSROWS ) then
     826                 :           0 :       ABI_ERROR("Bad value for toState")
     827                 :             :     end if
     828                 :             : 
     829                 :             :     !write(std_out,*) "linalg", rows(xgTransposer%xgBlock_linalg)*cols(xgTransposer%xgBlock_linalg)
     830                 :             :     !write(std_out,*) "colsrows", rows(xgTransposer%xgBlock_colsrows)*cols(xgTransposer%xgBlock_colsrows)
     831                 :     2112064 :     select case (toState)
     832                 :             :     case (STATE_LINALG)
     833         [ -  + ]:     2112064 :       if ( xgTransposer%state == STATE_LINALG ) then
     834                 :           0 :         ABI_WARNING("Array linalg has already been transposed")
     835                 :             :       end if
     836         [ +  + ]:     2112064 :       if ( xgTransposer%mpiData(MPI_COLS)%size > 1 ) then
     837                 :     1214170 :         call xgTransposer_toLinalg(xgTransposer)
     838                 :             :       else
     839                 :      897894 :         xgTransposer%state = STATE_LINALG
     840                 :             :       end if
     841                 :             :     case (STATE_COLSROWS)
     842         [ -  + ]:      741740 :       if ( xgTransposer%state == STATE_COLSROWS ) then
     843                 :           0 :         ABI_WARNING("Array colsrows has already been transposed")
     844                 :             :       end if
     845   [ +  +  -  +  :     3595544 :       if ( xgTransposer%mpiData(MPI_COLS)%size > 1 ) then
                      + ]
     846                 :      440818 :         call xgTransposer_toColsRows(xgTransposer)
     847                 :             :       else
     848                 :      300922 :         xgTransposer%state = STATE_COLSROWS
     849                 :             :       end if
     850                 :             :     end select
     851                 :             : 
     852                 :     2853804 :     call timab(tim_transpose,2,tsec)
     853                 :             : 
     854                 :     2853804 :   end subroutine xgTransposer_transpose
     855                 :             : !!***
     856                 :             : 
     857                 :             : !!****f* m_xgTransposer/xgTransposer_toLinalg
     858                 :             : !!
     859                 :             : !! NAME
     860                 :             : !! xgTransposer_toLinalg
     861                 :             : 
     862                 :     1214170 :   subroutine xgTransposer_toLinalg(xgTransposer)
     863                 :             : 
     864                 :             :    type(xgTransposer_t), intent(inout) :: xgTransposer
     865                 :     1214170 :    double precision, allocatable, target :: sendbuf(:,:)
     866                 :     1214170 :    double precision, pointer :: recvbuf(:,:)
     867                 :             :    !double precision, pointer :: buffer(:,:)
     868                 :     1214170 :    integer, allocatable :: sendcounts(:), recvcounts(:)
     869                 :     1214170 :    integer, allocatable :: sdispls(:), rdispls(:)
     870                 :             :    integer :: ncpu_cols, ncpu_rows, comm, me_rows, me_cols
     871                 :             :    integer :: nrowsColsRows
     872                 :             :    integer :: ncolsColsRows
     873                 :             :    integer :: nrowsLinalgMe
     874                 :             :    integer :: icpu, ierr
     875                 :             :    integer :: send_start,send_end
     876                 :             :    !integer :: myrequest
     877                 :             :    !integer, allocatable :: request(:), status(:)
     878                 :     1214170 :    type(ptr_t), allocatable :: sendptrbuf(:)
     879                 :     1214170 :    integer, pointer :: nrowsLinalg(:)
     880                 :     1214170 :    integer, pointer :: ncolsColsRows_sub(:)
     881                 :             :    double precision :: tsec(2)
     882                 :             : 
     883                 :             : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
     884                 :             :    double precision, allocatable :: recvbuf_mpi(:,:)
     885                 :             :    integer(c_size_t) :: buffer_size
     886                 :             : #endif
     887                 :             : 
     888                 :     1214170 :    call timab(tim_toLinalg,1,tsec)
     889                 :             : 
     890                 :     1214170 :    ncpu_cols = xgTransposer%mpiData(MPI_COLS)%size
     891                 :     1214170 :    ncpu_rows = xgTransposer%mpiData(MPI_ROWS)%size
     892                 :     1214170 :    comm = xgTransposer%mpiData(MPI_COLS)%comm
     893                 :             : !   me = xgTransposer%mpiData(MPI_ROWS)%rank*ncpu
     894                 :     1214170 :    me_rows = xgTransposer%mpiData(MPI_ROWS)%rank
     895                 :             : 
     896                 :             :    !write(std_out,'(a,i4,i4,i4)') '@xgTransposer_toLinalg for MPI proc in comm=', &
     897                 :             :    !    xgTransposer%mpiData(MPI_LINALG)%rank, xgTransposer%mpiData(MPI_LINALG)%comm, comm
     898                 :             : 
     899                 :     1214170 :    nrowsColsRows = xgTransposer%nrowsColsRows
     900                 :     1214170 :    ncolsColsRows = xgTransposer%ncolsColsRows
     901                 :             : 
     902                 :     1214170 :    nrowsLinalg => xgTransposer%nrowsLinalg
     903         [ -  + ]:     1214170 :    if (xgTransposer%custom_ncolsColsRows) ncolsColsRows_sub => xgTransposer%ncolsColsRows_sub
     904                 :     1214170 :    nrowsLinalgMe = nrowsLinalg(xgTransposer%mpiData(MPI_LINALG)%rank+1)
     905                 :             : 
     906                 :             :    !write(std_out,*) '@toLinalg nrowsColsRows, ncolsColsRows=', nrowsColsRows, ncolsColsRows
     907                 :             :    !call flush_unit(std_out)
     908                 :             : 
     909   [ -  +  +  -  :     3642510 :    ABI_MALLOC(sendbuf,(2,nrowsColsRows*ncolsColsRows))
          +  -  +  -  -  
                +  -  + ]
     910                 :             : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
     911                 :             :    !$OMP TARGET ENTER DATA MAP(alloc:sendbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
     912                 :             : #endif
     913                 :     1214170 :    call xgTransposer_reorganizeData(xgTransposer,sendbuf)
     914                 :             : 
     915   [ -  +  +  -  :     3642510 :    ABI_MALLOC(recvcounts,(ncpu_cols))
          +  -  +  -  -  
                +  -  + ]
     916   [ -  +  +  -  :     2428340 :    ABI_MALLOC(rdispls,(ncpu_cols))
             -  +  -  + ]
     917         [ -  + ]:     1214170 :    if (xgTransposer%custom_ncolsColsRows) then ! non uniform sendcounts
     918         [ #  # ]:           0 :      recvcounts(:) = 2*nrowsLinalgMe*ncolsColsRows_sub(:)
     919                 :             :    else
     920         [ +  + ]:     4168758 :      recvcounts(:) = 2*nrowsLinalgMe*ncolsColsRows !! Thank you fortran for not starting at 0 !
     921                 :             :    end if
     922                 :     1214170 :    rdispls(1) = 0
     923         [ +  + ]:     2954588 :    do icpu = 2, ncpu_cols
     924                 :     2954588 :      rdispls(icpu) = rdispls(icpu-1)+recvcounts(icpu-1)
     925                 :             :    end do
     926                 :             : 
     927                 :             :    call xgBlock_reverseMap(xgTransposer%xgBlock_linalg,recvbuf,&
     928                 :     1214170 :      & rows=1,cols=cols(xgTransposer%xgBlock_linalg)*nrowsLinalgMe)
     929                 :             : 
     930                 :             : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
     931                 :             :    ! just for debug
     932                 :             :    if( xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
     933                 :             :       !call gpu_managed_ptr_status(C_LOC(recvbuf))
     934                 :             :       buffer_size = size(recvbuf) * dp
     935                 :             :       call gpu_data_prefetch_async(C_LOC(recvbuf), buffer_size, CPU_DEVICE_ID)
     936                 :             :       call gpu_device_synchronize()
     937                 :             : 
     938                 :             :       ABI_MALLOC(recvbuf_mpi, (size(recvbuf,1), size(recvbuf,2)) )
     939                 :             : 
     940                 :             :    end if
     941                 :             : #endif
     942                 :             : 
     943   [ -  +  +  -  :     2428340 :    ABI_MALLOC(sendcounts,(ncpu_cols))
             -  +  -  + ]
     944   [ -  +  +  -  :     2428340 :    ABI_MALLOC(sdispls,(ncpu_cols))
             -  +  -  + ]
     945                 :             :    !sendcounts(:) = 2*nrowsLinalg(me+1:me+ncpu)*ncolsColsRows
     946         [ +  + ]:     4168758 :    do icpu=0,ncpu_cols-1
     947                 :     4168758 :      sendcounts(icpu+1) = 2*nrowsLinalg(me_rows+1+icpu*ncpu_rows)*ncolsColsRows
     948                 :             :      !write(*,*) "nrowsLinalg ", icpu, ncpu_rows, nrowsLinalg(me_rows+1+icpu*ncpu_rows)
     949                 :             :    end do
     950                 :     1214170 :    sdispls(1) = 0
     951         [ +  + ]:     2954588 :    do icpu = 2, ncpu_cols
     952                 :     2954588 :      sdispls(icpu) = sdispls(icpu-1)+sendcounts(icpu-1)
     953                 :             :    end do
     954                 :             : 
     955                 :     2428340 :    select case(xgTransposer%mpiAlgo)
     956                 :             :    case (TRANS_ALL2ALL)
     957                 :             :      !ABI_MALLOC(request,(1))
     958                 :             :      !myrequest = 1
     959                 :             : 
     960                 :             :      ABI_NVTX_START_RANGE(NVTX_TRANSPOSER_MPI_ALL2ALL)
     961                 :             : 
     962         [ +  - ]:     1214170 :      if( xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
     963                 :             : 
     964                 :             : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
     965                 :             :        call timab(tim_all2allv,1,tsec)
     966                 :             :        call xmpi_alltoallv(sendbuf,     sendcounts, sdispls, &
     967                 :             :                            recvbuf_mpi, recvcounts, rdispls, &
     968                 :             :                            comm, ierr)
     969                 :             :        call timab(tim_all2allv,2,tsec)
     970                 :             :        !call xmpi_ialltoallv(sendbuf, sendcounts, sdispls, &
     971                 :             :        !                    recvbuf, recvcounts, rdispls, &
     972                 :             :        !                    comm, request(myrequest))
     973                 :             : 
     974                 :             :        ! copy back recvbuf_mpi into recvbuf
     975                 :             :        recvbuf(:,:) = recvbuf_mpi(:,:)
     976                 :             : 
     977                 :             :        ABI_FREE(recvbuf_mpi)
     978                 :             : #endif
     979                 :             : 
     980                 :             :      else
     981                 :             : 
     982                 :             :        !write(*,*) "sendcounts for proc ", me_rows, xmpi_comm_rank(comm), ":", sendcounts(:)
     983                 :             :        !write(*,*) "recvcounts for proc ", me_rows, xmpi_comm_rank(comm), ":", recvcounts(:)
     984                 :             : 
     985                 :     1214170 :        call timab(tim_all2allv,1,tsec)
     986                 :             :        call xmpi_alltoallv(sendbuf, sendcounts, sdispls, &
     987                 :             :                            recvbuf, recvcounts, rdispls, &
     988                 :     1214170 :                            comm, ierr, use_omp_map=(xgTransposer%gpu_option==ABI_GPU_OPENMP))
     989                 :     1214170 :        call timab(tim_all2allv,2,tsec)
     990                 :             :        !call xmpi_ialltoallv(sendbuf, sendcounts, sdispls, &
     991                 :             :        !                    recvbuf, recvcounts, rdispls, &
     992                 :             :        !                    comm, request(myrequest))
     993                 :             : 
     994                 :             :      end if
     995                 :             : 
     996                 :             :      ABI_NVTX_END_RANGE()
     997                 :             : 
     998                 :             :    case (TRANS_GATHER)
     999                 :             : 
    1000                 :             :      !ABI_MALLOC(request,(ncpu))
    1001                 :           0 :      me_cols = xgTransposer%mpiData(MPI_COLS)%rank
    1002                 :             :      !myrequest = me+1
    1003                 :             : 
    1004                 :             : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
    1005                 :             :      !$OMP TARGET UPDATE FROM(sendbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
    1006                 :             : #endif
    1007   [ #  #  #  #  :           0 :      ABI_MALLOC(sendptrbuf,(1:ncpu_cols))
          #  #  #  #  #  
             #  #  #  #  
                      # ]
    1008         [ #  # ]:           0 :      do icpu = 1, ncpu_cols
    1009                 :           0 :        send_start = sdispls(icpu)/2+1
    1010                 :           0 :        send_end   = sdispls(icpu)/2+sendcounts(icpu)/2
    1011                 :           0 :        sendptrbuf(icpu)%ptr => sendbuf(:,send_start:send_end)
    1012                 :           0 :        call timab(tim_gatherv,1,tsec)
    1013                 :           0 :        call xmpi_gatherv(sendptrbuf(icpu)%ptr,sendcounts(icpu),recvbuf,recvcounts,rdispls,icpu-1,comm,ierr)
    1014                 :           0 :        call timab(tim_gatherv,2,tsec)
    1015                 :             :        !call mpi_igatherv(sendptrbuf(me+1)%ptr,sendcounts(icpu),MPI_DOUBLE_PRECISION,&
    1016                 :             :        !  recvbuf,recvcounts,rdispls,MPI_DOUBLE_PRECISION,icpu-1,comm,request(icpu),ierr)
    1017                 :             :      end do
    1018                 :             : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
    1019                 :             :      !$OMP TARGET UPDATE TO(recvbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
    1020                 :             : #endif
    1021                 :             : 
    1022                 :             :    case default
    1023      [ +  -  - ]:     1214170 :      ABI_BUG("This algo does not exist")
    1024                 :             :    end select
    1025                 :             : 
    1026                 :     1214170 :    xgTransposer%state = STATE_LINALG
    1027                 :             : 
    1028                 :             :    !ABI_MALLOC(status,(MPI_STATUS_SIZE))
    1029                 :             :    !call mpi_wait(request(myrequest),status,ierr)
    1030         [ -  + ]:     1214170 :    if ( ierr /= xmpi_success ) then
    1031                 :           0 :      ABI_ERROR("Error while waiting for mpi")
    1032                 :             :    end if
    1033                 :             : 
    1034         [ +  - ]:     1214170 :    if ( allocated(sendcounts) ) then
    1035                 :     1214170 :      ABI_FREE(sendcounts)
    1036                 :             :    end if
    1037                 :             :    if ( allocated(sdispls) ) then
    1038                 :     1214170 :      ABI_FREE(sdispls)
    1039                 :             :    end if
    1040                 :             : 
    1041         [ -  + ]:     1214170 :    ABI_FREE(recvcounts)
    1042                 :     1214170 :    ABI_FREE(rdispls)
    1043                 :             : 
    1044                 :             : 
    1045         [ -  + ]:     1214170 :    if ( allocated(sendptrbuf) ) then
    1046                 :           0 :      ABI_FREE(sendptrbuf)
    1047                 :             :    end if
    1048                 :             : 
    1049                 :             :    !do icpu = 1, size(request)
    1050                 :             :    !  if ( icpu /= myrequest ) then
    1051                 :             :    !    call mpi_wait(request(icpu),status,ierr)
    1052                 :             :    !    if ( ierr /= MPI_SUCCESS ) then
    1053                 :             :    !      ABI_ERROR("Error while waiting for other mpi")
    1054                 :             :    !    end if
    1055                 :             :    !  end if
    1056                 :             :    !end do
    1057                 :             : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
    1058                 :             :    !$OMP TARGET EXIT DATA MAP(delete:sendbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
    1059                 :             : #endif
    1060         [ -  + ]:     1214170 :    ABI_FREE(sendbuf)
    1061                 :             :    !ABI_FREE(status)
    1062                 :             :    !ABI_FREE(request)
    1063                 :             : 
    1064                 :     1214170 :     call timab(tim_toLinalg,2,tsec)
    1065                 :             : 
    1066   [ -  +  -  + ]:     2428340 :   end subroutine xgTransposer_toLinalg
    1067                 :             : !!***
    1068                 :             : 
    1069                 :             : !!****f* m_xgTransposer/xgTransposer_toColsRows
    1070                 :             : !!
    1071                 :             : !! NAME
    1072                 :             : !! xgTransposer_toColsRows
    1073                 :             : 
    1074                 :      440818 :   subroutine xgTransposer_toColsRows(xgTransposer)
    1075                 :             : 
    1076                 :             :    type(xgTransposer_t), intent(inout) :: xgTransposer
    1077                 :      440818 :    double precision, pointer :: sendbuf(:,:)
    1078                 :      440818 :    double precision, allocatable :: recvbuf(:,:)
    1079                 :             :    !double precision, allocatable :: buffer(:,:)
    1080                 :      440818 :    integer, allocatable :: sendcounts(:), recvcounts(:)
    1081                 :      440818 :    integer, allocatable :: sdispls(:), rdispls(:)
    1082                 :             :    integer :: ncpu_rows, ncpu_cols, comm, me_rows, me_cols
    1083                 :             :    integer :: nrowsColsRows
    1084                 :             :    integer :: nrowsLinalgMe
    1085                 :             :    integer :: ncolsColsRows
    1086                 :             :    integer :: icpu,ierr
    1087                 :             :    !integer :: myrequest
    1088                 :             :    !integer, allocatable :: request(:), status(:)
    1089                 :             :    type(xgBlock_t) :: xgBlock_toTransposed
    1090                 :      440818 :    type(ptr_t), allocatable :: sendptrbuf(:)
    1091                 :      440818 :    integer, pointer :: nrowsLinalg(:)
    1092                 :      440818 :    integer, pointer :: ncolsColsRows_sub(:)
    1093                 :             :    double precision :: tsec(2)
    1094                 :             : 
    1095                 :             : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
    1096                 :             :    double precision, allocatable :: sendbuf_mpi(:,:)
    1097                 :             :    integer(c_size_t) :: buffer_size
    1098                 :             : #endif
    1099                 :             : 
    1100                 :      440818 :     call timab(tim_toColsRows,1,tsec)
    1101                 :             : 
    1102                 :      440818 :    ncpu_cols = xgTransposer%mpiData(MPI_COLS)%size
    1103                 :      440818 :    ncpu_rows = xgTransposer%mpiData(MPI_ROWS)%size
    1104                 :      440818 :    comm = xgTransposer%mpiData(MPI_COLS)%comm
    1105                 :      440818 :    me_rows = xgTransposer%mpiData(MPI_ROWS)%rank
    1106                 :             : 
    1107                 :      440818 :    nrowsColsRows = xgTransposer%nrowsColsRows
    1108                 :      440818 :    ncolsColsRows = xgTransposer%ncolsColsRows
    1109                 :             : 
    1110                 :      440818 :    nrowsLinalg => xgTransposer%nrowsLinalg
    1111         [ +  + ]:      440818 :    if (xgTransposer%custom_ncolsColsRows) ncolsColsRows_sub => xgTransposer%ncolsColsRows_sub
    1112                 :      440818 :    nrowsLinalgMe = nrowsLinalg(xgTransposer%mpiData(MPI_LINALG)%rank+1)
    1113                 :             : 
    1114   [ -  +  +  -  :     1322454 :    ABI_MALLOC(recvbuf,(2,nrowsColsRows*ncolsColsRows))
          +  -  +  -  -  
                +  -  + ]
    1115                 :             : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
    1116                 :             :    !$OMP TARGET ENTER DATA MAP(alloc:recvbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
    1117                 :             : #endif
    1118   [ -  +  +  -  :     1322454 :    ABI_MALLOC(recvcounts,(ncpu_cols))
          +  -  +  -  -  
                +  -  + ]
    1119   [ -  +  +  -  :      881636 :    ABI_MALLOC(rdispls,(ncpu_cols))
             -  +  -  + ]
    1120                 :             : 
    1121                 :             :    ! 2 * ncolsColsRows * nrowsLinalg_sub(:)
    1122         [ +  + ]:     1563278 :    recvcounts(:) = 2*ncolsColsRows*nrowsLinalg(1+me_rows:1+me_rows+(ncpu_cols-1)*ncpu_rows:ncpu_rows)
    1123                 :      440818 :    rdispls(1) = 0
    1124         [ +  + ]:     1122460 :    do icpu = 2, ncpu_cols
    1125                 :     1122460 :      rdispls(icpu) = rdispls(icpu-1)+recvcounts(icpu-1)
    1126                 :             :    end do
    1127                 :             :    !recvcounts(:) = 2*nrowsLinalg(me+1:me+ncpu)*ncolsColsRows
    1128                 :             :    !rdispls(1) = 0
    1129                 :             :    !do icpu = 2, ncpu
    1130                 :             :    !  rdispls(icpu) = rdispls(icpu-1)+recvcounts(icpu-1)
    1131                 :             :    !end do
    1132                 :             : 
    1133                 :      881636 :    select case(xgTransposer%mpiAlgo)
    1134                 :             :    case (TRANS_ALL2ALL)
    1135   [ -  +  +  -  :      881636 :      ABI_MALLOC(sendcounts,(ncpu_cols))
             -  +  -  + ]
    1136   [ -  +  +  -  :      881636 :      ABI_MALLOC(sdispls,(ncpu_cols))
             -  +  -  + ]
    1137                 :             :      !ABI_MALLOC(request,(1))
    1138                 :             :      !myrequest = 1
    1139                 :             : 
    1140                 :             :      ! 2 * ncolsColsRows_sub(:) * nrowsLinalgMe
    1141         [ +  + ]:      440818 :      if (xgTransposer%custom_ncolsColsRows) then ! non uniform sendcounts
    1142         [ +  + ]:          80 :         sendcounts(:) = 2*nrowsLinalgMe*ncolsColsRows_sub(:)
    1143                 :             :      else
    1144         [ +  + ]:     1563198 :         sendcounts(:) = 2*nrowsLinalgMe*ncolsColsRows !! Thank you fortran for not starting at 0 !
    1145                 :             :      end if
    1146                 :      440818 :      sdispls(1) = 0
    1147         [ +  + ]:     1122460 :      do icpu = 2, ncpu_cols
    1148                 :     1122460 :        sdispls(icpu) = sdispls(icpu-1)+sendcounts(icpu-1)
    1149                 :             :      end do
    1150                 :             : 
    1151                 :             :      call xgBlock_reverseMap(xgTransposer%xgBlock_linalg,sendbuf, &
    1152                 :      440818 : &      rows=1,cols=cols(xgTransposer%xgBlock_linalg)*nrowsLinalgMe)
    1153                 :             :      !write(*,*) "Before ialltoall"
    1154                 :             : 
    1155                 :             :      ABI_NVTX_START_RANGE(NVTX_TRANSPOSER_MPI_ALL2ALL)
    1156                 :             : 
    1157                 :             :     ! if gpu is enabled, data are located in GPU memory, so we copy them on a host buffer
    1158         [ +  - ]:      440818 :     if( xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
    1159                 :             : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
    1160                 :             :       ABI_MALLOC(sendbuf_mpi, (size(sendbuf,1), size(sendbuf,2)) )
    1161                 :             : 
    1162                 :             :       ! sync sendbuf on host and then copy to sendbuf_mpi
    1163                 :             :       buffer_size = size(sendbuf) * dp
    1164                 :             :       call gpu_data_prefetch_async(C_LOC(sendbuf), buffer_size, CPU_DEVICE_ID)
    1165                 :             :       call gpu_device_synchronize()
    1166                 :             : 
    1167                 :             :       sendbuf_mpi(:,:) = sendbuf(:,:)
    1168                 :             : 
    1169                 :             :       call timab(tim_all2allv,1,tsec)
    1170                 :             :       ! this is a cpu mpi comm
    1171                 :             :       call xmpi_alltoallv(sendbuf_mpi, sendcounts, sdispls, &
    1172                 :             :                           recvbuf,     recvcounts, rdispls, &
    1173                 :             :                           comm, ierr)
    1174                 :             :       call timab(tim_all2allv,2,tsec)
    1175                 :             :       !call xmpi_ialltoallv(sendbuf, sendcounts, sdispls, &
    1176                 :             :       !                    recvbuf, recvcounts, rdispls, &
    1177                 :             :       !                    comm, request(myrequest))
    1178                 :             :       !write(*,*) "After ialltoall"
    1179                 :             : 
    1180                 :             :       ABI_FREE(sendbuf_mpi)
    1181                 :             : #endif
    1182                 :             :     else
    1183                 :             : 
    1184                 :             :       !write(std_out,'(a,i4,i4,i4)') '@xgTransposer_toColsRows for MPI proc in comm=', &
    1185                 :             :       !    xgTransposer%mpiData(MPI_LINALG)%rank, xgTransposer%mpiData(MPI_LINALG)%comm, comm
    1186                 :             :       !if (xgTransposer%custom_ncolsColsRows) then
    1187                 :             :       !    write(std_out,*) 'db ncolsColsRows_sub', ncolsColsRows_sub(:)
    1188                 :             :       !end if
    1189                 :             :       !write(std_out,*) 'db ncolsColsMe', ncolsColsRows
    1190                 :             :       !write(std_out,*) 'db nrowsLinalg', nrowsLinalg(1+me_rows:1+me_rows+(ncpu_cols-1)*ncpu_rows:ncpu_rows)
    1191                 :             :       !write(std_out,*) 'db nrowsLinalgMe', nrowsLinalgMe
    1192                 :             :       !write(std_out,*) 'db sendbuf', cols(xgTransposer%xgBlock_linalg), nrowsLinalgMe
    1193                 :             :       !write(std_out,*) 'db recvbuf', nrowsColsRows, ncolsColsRows
    1194                 :             :       !write(std_out,*) 'db sendcounts', sendcounts(:)
    1195                 :             :       !write(std_out,*) 'db recvcounts', recvcounts(:)
    1196                 :             :       !write(std_out,*) 'db sdispls', sdispls(:)
    1197                 :             :       !write(std_out,*) 'db rdispls', rdispls(:)
    1198                 :             :       !call flush_unit(std_out)
    1199                 :             : 
    1200                 :      440818 :       call timab(tim_all2allv,1,tsec)
    1201                 :             :       call xmpi_alltoallv(sendbuf, sendcounts, sdispls, &
    1202                 :             :                           recvbuf, recvcounts, rdispls, &
    1203                 :      440818 :                           comm, ierr, use_omp_map=(xgTransposer%gpu_option == ABI_GPU_OPENMP))
    1204                 :      440818 :       call timab(tim_all2allv,2,tsec)
    1205                 :             :       !call xmpi_ialltoallv(sendbuf, sendcounts, sdispls, &
    1206                 :             :       !                    recvbuf, recvcounts, rdispls, &
    1207                 :             :       !                    comm, request(myrequest))
    1208                 :             :       !write(*,*) "After ialltoall"
    1209                 :             :     end if
    1210                 :             : 
    1211                 :             :     ABI_NVTX_END_RANGE()
    1212                 :             : 
    1213                 :             :    case (TRANS_GATHER)
    1214                 :             : 
    1215                 :             : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
    1216                 :             :      !$OMP TARGET UPDATE FROM(sendbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
    1217                 :             : #endif
    1218                 :             :      !ABI_MALLOC(request,(ncpu))
    1219                 :           0 :      me_cols = xgTransposer%mpiData(MPI_COLS)%rank
    1220                 :             :      !myrequest = me+1
    1221                 :             : 
    1222   [ #  #  #  #  :           0 :      ABI_MALLOC(sendptrbuf,(1:ncpu_cols))
          #  #  #  #  #  
             #  #  #  #  
                      # ]
    1223                 :             :      !call flush(6)
    1224                 :             :      !call xmpi_barrier(xgTransposer%mpiData(MPI_LINALG)%comm)
    1225         [ #  # ]:           0 :      do icpu = 0, ncpu_cols-1
    1226                 :             :        !write(*,*) me, "->", icpu, "from col ",icpu*ncolsColsRows+1, " number of rows:", nrowsLinalgMe
    1227                 :           0 :        call xgBlock_setBlock(xgTransposer%xgBlock_linalg,xgBlock_toTransposed,nrowsLinalgMe,ncolsColsRows,fcol=icpu*ncolsColsRows+1)
    1228                 :             :        call xgBlock_reverseMap(xgBlock_toTransposed,sendptrbuf(icpu+1)%ptr,&
    1229                 :           0 :          rows=1,cols=ncolsColsRows*nrowsLinalgMe)
    1230                 :           0 :        call timab(tim_gatherv,1,tsec)
    1231                 :           0 :        call xmpi_gatherv(sendptrbuf(icpu+1)%ptr,2*ncolsColsRows*nrowsLinalgMe,recvbuf,recvcounts,rdispls,icpu,comm,ierr)
    1232                 :           0 :        call timab(tim_gatherv,2,tsec)
    1233                 :             :        !call mpi_igatherv(sendptrbuf(me+1)%ptr,2*ncolsColsRows*nrowsLinalgMe,MPI_DOUBLE_PRECISION,&
    1234                 :             :        !  recvbuf,recvcounts,rdispls,MPI_DOUBLE_PRECISION,icpu,comm,request(icpu+1),ierr)
    1235                 :             :      end do
    1236                 :             :      !call xmpi_barrier(xgTransposer%mpiData(MPI_LINALG)%comm)
    1237                 :             :      !call flush(6)
    1238                 :             :      !write(*,*) me, request
    1239                 :             : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
    1240                 :             :      !$OMP TARGET UPDATE TO(recvbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
    1241                 :             : #endif
    1242                 :             : 
    1243                 :             :    case default
    1244      [ +  -  - ]:      440818 :      ABI_BUG("This algo does not exist")
    1245                 :             :    end select
    1246                 :             : 
    1247                 :             :    !ABI_MALLOC(status,(MPI_STATUS_SIZE))
    1248                 :             :    !call mpi_wait(request(myrequest),status,ierr)
    1249                 :             :    !write(*,*) "Request ended"
    1250         [ -  + ]:      440818 :    if ( ierr /= xmpi_success ) then
    1251                 :           0 :      ABI_ERROR("Error while waiting for mpi")
    1252                 :             :    end if
    1253                 :             :    !write(*,*) "with success"
    1254                 :             : 
    1255                 :      440818 :    call xgTransposer_reorganizeData(xgTransposer,recvbuf)
    1256                 :             : 
    1257         [ +  - ]:      440818 :    if ( allocated(sendcounts) ) then
    1258                 :      440818 :      ABI_FREE(sendcounts)
    1259                 :             :    end if
    1260         [ +  - ]:      440818 :    if ( allocated(sdispls) ) then
    1261                 :      440818 :      ABI_FREE(sdispls)
    1262                 :             :    end if
    1263                 :             : 
    1264         [ -  + ]:      440818 :    ABI_FREE(recvcounts)
    1265                 :      440818 :    ABI_FREE(rdispls)
    1266                 :             : 
    1267                 :             : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
    1268                 :             :    !$OMP TARGET EXIT DATA MAP(delete:recvbuf) if(xgTransposer%gpu_option == ABI_GPU_OPENMP)
    1269                 :             : #endif
    1270         [ -  + ]:      440818 :    ABI_FREE(recvbuf)
    1271                 :             : 
    1272         [ -  + ]:      440818 :    if ( allocated(sendptrbuf) ) then
    1273                 :           0 :      ABI_FREE(sendptrbuf)
    1274                 :             :    end if
    1275                 :             : 
    1276                 :      440818 :    xgTransposer%state = STATE_COLSROWS
    1277                 :             : 
    1278                 :             :    !do icpu = 1, size(request)
    1279                 :             :    !  if ( icpu /= myrequest ) then
    1280                 :             :    !    call mpi_wait(request(icpu),status,ierr)
    1281                 :             :    !    if ( ierr /= MPI_SUCCESS ) then
    1282                 :             :    !      ABI_ERROR("Error while waiting for other mpi")
    1283                 :             :    !    end if
    1284                 :             :    !  end if
    1285                 :             :    !end do
    1286                 :             :    !ABI_FREE(status)
    1287                 :             :    !ABI_FREE(request)
    1288                 :             : 
    1289                 :      440818 :     call timab(tim_toColsRows,2,tsec)
    1290                 :             : 
    1291   [ -  +  -  + ]:      881636 :   end subroutine xgTransposer_toColsRows
    1292                 :             : !!***
    1293                 :             : 
    1294                 :             : !!****f* m_xgTransposer/xgTransposer_reorganizeData
    1295                 :             : !!
    1296                 :             : !! NAME
    1297                 :             : !! xgTransposer_reorganizeData
    1298                 :             : 
    1299         [ -  + ]:     1654988 :   subroutine xgTransposer_reorganizeData(xgTransposer,bufferMess)
    1300                 :             : 
    1301                 :             :     type(xgTransposer_t), intent(inout) :: xgTransposer
    1302                 :             :     double precision    , intent(inout) :: bufferMess(:,:)
    1303                 :             :     double precision, pointer :: bufferOrdered(:,:) => null()
    1304                 :             :     integer :: nrowsColsRows,nthreads_bak
    1305                 :             :     integer :: ncolsColsRows
    1306                 :             :     integer :: tos,toe,froms,frome
    1307                 :             :     integer :: col, icpu
    1308                 :             :     integer :: me_rows,ncpu_cols,ncpu_rows
    1309                 :             :     integer :: nPair,ispinor,nspinor
    1310                 :             :     integer :: nrowsLinalgMe,nrowsLinalgMeSum
    1311                 :             : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
    1312                 :             :     integer :: irow
    1313                 :             :     integer, contiguous, pointer :: nrowsLinalg(:)
    1314                 :             : #else
    1315                 :     1654988 :     integer, pointer :: nrowsLinalg(:)
    1316                 :             : #endif
    1317                 :             :     double precision :: tsec(2)
    1318                 :             : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
    1319                 :             :     integer(c_size_t) :: buffer_size
    1320                 :             : #endif
    1321                 :             : 
    1322                 :     1654988 :     call timab(tim_reorganize,1,tsec)
    1323                 :             : 
    1324                 :     1654988 :     me_rows = xgTransposer%mpiData(MPI_ROWS)%rank!*xgTransposer%mpiData(MPI_COLS)%size
    1325                 :     1654988 :     ncpu_rows = xgTransposer%mpiData(MPI_ROWS)%size
    1326                 :     1654988 :     ncpu_cols = xgTransposer%mpiData(MPI_COLS)%size
    1327                 :             : 
    1328                 :             :     ! These arrays are stored as
    1329                 :             :     ! 1 2 3
    1330                 :             :     ! 4 5 6
    1331                 :             :     ! 7 8 9
    1332                 :             :     !
    1333                 :             :     ! stored as
    1334                 :             :     !
    1335                 :             :     ! column-major, one-base indexing  : 1 4 7 2 5 8 3 6 9    | y + N_y*(x-1)
    1336                 :             :     ! row-major   , zero-based indexing: 1 2 3 4 5 6 7 8 9    | x + N_x*y
    1337                 :     1654988 :     nspinor = xgTransposer%nspinor
    1338                 :     1654988 :     nrowsLinalg => xgTransposer%nrowsLinalg
    1339                 :             : 
    1340                 :     1654988 :     nrowsColsRows = xgTransposer%nrowsColsRows
    1341                 :     1654988 :     ncolsColsRows = xgTransposer%ncolsColsRows
    1342                 :             : 
    1343                 :             :     ! number of pairs depends on current MPI process
    1344                 :     1654988 :     nPair = nrowsColsRows*ncolsColsRows
    1345                 :     1654988 :     call xgBlock_reverseMap(xgTransposer%xgBlock_colsrows,bufferOrdered,rows=1,cols=nPair)
    1346                 :             : 
    1347                 :             : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
    1348                 :             :     ! if gpu is enabled, data are located in GPU memory, so we prefetch them on host
    1349                 :             :     ! to do the following reorganization.
    1350                 :             :     ! Alternatively, we should provide a GPU implementation of this data layout reorganization
    1351                 :             :     if( xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
    1352                 :             :        buffer_size = size(bufferOrdered) * dp
    1353                 :             :        call gpu_data_prefetch_async(C_LOC(bufferOrdered), buffer_size, CPU_DEVICE_ID)
    1354                 :             :        call gpu_device_synchronize()
    1355                 :             :     end if
    1356                 :             : 
    1357                 :             :     ! if gpu enabled increase locally OpenMP num threads
    1358                 :             :     if (xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
    1359                 :             :        call xomp_set_num_threads(xgTransposer%gpu_kokkos_nthrd)
    1360                 :             :     end if
    1361                 :             : #endif
    1362                 :             : 
    1363   [ -  +  -  - ]:     1654988 :     if (xgTransposer%gpu_option /= ABI_GPU_DISABLED .and. xgTransposer%gpu_thread_limit /= 0) then
    1364                 :           0 :       nthreads_bak=xomp_get_max_threads()
    1365                 :           0 :       call xomp_set_num_threads(min(xgTransposer%gpu_thread_limit,nthreads_bak))
    1366                 :             :     end if
    1367                 :             : 
    1368                 :     1654988 :     select case (xgTransposer%state)
    1369                 :             :     case (STATE_LINALG)
    1370                 :             :       ! We are going to STATE_COLSROWS so we are after all2all
    1371                 :             : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
    1372                 :             :       if (xgTransposer%gpu_option == ABI_GPU_OPENMP) then
    1373                 :             :         ! bufferMess and bufferOrdered already live on the device: reorganize them there directly,
    1374                 :             :         ! no CPU<->GPU round trip needed. nrowsLinalg is small and mapped just for this kernel.
    1375                 :             :         ! Array-section assignment is not usable in device kernel bodies: loop explicitly instead.
    1376                 :             :         !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
    1377                 :             :         !$OMP& PRIVATE(nrowsLinalgMe,nrowsLinalgMeSum,toe,tos,frome,froms,irow) &
    1378                 :             :         !$OMP& MAP(to:bufferMess) MAP(to:bufferOrdered) MAP(to:nrowsLinalg)
    1379                 :             :         do col = 1, ncolsColsRows
    1380                 :             :           do icpu = 0, ncpu_cols-1
    1381                 :             :             do ispinor = 1, nspinor
    1382                 :             :               nrowsLinalgMe = nrowsLinalg(1+me_rows+icpu*ncpu_rows)
    1383                 :             :               nrowsLinalgMeSum = sum(nrowsLinalg(1+me_rows:1+me_rows+(icpu-1)*ncpu_rows:ncpu_rows))
    1384                 :             :               froms=1+(ispinor-1)*nrowsLinalgMe/nspinor+(col-1)*nrowsLinalgMe+nrowsLinalgMeSum*ncolsColsRows
    1385                 :             :               frome=froms-1+nrowsLinalgMe/nspinor
    1386                 :             :               tos=1+nrowsLinalgMeSum/nspinor+(ispinor-1)*nrowsColsRows/nspinor+(col-1)*nrowsColsRows
    1387                 :             :               toe=tos-1+nrowsLinalgMe/nspinor
    1388                 :             :               do irow = 0, toe-tos
    1389                 :             :                 bufferOrdered(1,tos+irow) = bufferMess(1,froms+irow)
    1390                 :             :                 bufferOrdered(2,tos+irow) = bufferMess(2,froms+irow)
    1391                 :             :               end do
    1392                 :             :             end do
    1393                 :             :           end do
    1394                 :             :         end do
    1395                 :             :       else
    1396                 :             : #endif
    1397                 :             :         !$omp parallel do private(nrowsLinalgMe,nrowsLinalgMeSum,toe,tos,frome,froms), collapse(3)
    1398         [ +  + ]:     1100816 :         do col = 1, ncolsColsRows
    1399         [ +  + ]:     2799336 :           do icpu = 0, ncpu_cols-1
    1400         [ +  + ]:     4269886 :             do ispinor = 1, nspinor
    1401                 :     1911368 :               nrowsLinalgMe = nrowsLinalg(1+me_rows+icpu*ncpu_rows)
    1402         [ +  + ]:     3836948 :               nrowsLinalgMeSum = sum(nrowsLinalg(1+me_rows:1+me_rows+(icpu-1)*ncpu_rows:ncpu_rows))
    1403                 :     1911368 :               froms=1+(ispinor-1)*nrowsLinalgMe/nspinor+(col-1)*nrowsLinalgMe+nrowsLinalgMeSum*ncolsColsRows
    1404                 :     1911368 :               frome=froms-1+nrowsLinalgMe/nspinor
    1405                 :     1911368 :               tos=1+nrowsLinalgMeSum/nspinor+(ispinor-1)*nrowsColsRows/nspinor+(col-1)*nrowsColsRows
    1406                 :     1911368 :               toe=tos-1+nrowsLinalgMe/nspinor
    1407   [ +  +  +  + ]:   554331942 :               bufferOrdered(:,tos:toe) = bufferMess(:,froms:frome)
    1408                 :             :             end do
    1409                 :             :           end do
    1410                 :             :         end do
    1411                 :             : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
    1412                 :             :       end if
    1413                 :             : #endif
    1414                 :             :     case (STATE_COLSROWS)
    1415                 :             :       ! We are going to STATE_LINALG so we are before all2all
    1416                 :             : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
    1417                 :             :       if (xgTransposer%gpu_option == ABI_GPU_OPENMP) then
    1418                 :             :         !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
    1419                 :             :         !$OMP& PRIVATE(nrowsLinalgMe,nrowsLinalgMeSum,toe,tos,frome,froms,irow) &
    1420                 :             :         !$OMP& MAP(to:bufferMess) MAP(to:bufferOrdered) MAP(to:nrowsLinalg)
    1421                 :             :         do col = 1, ncolsColsRows
    1422                 :             :           do icpu = 0, ncpu_cols-1
    1423                 :             :             do ispinor = 1, nspinor
    1424                 :             :               nrowsLinalgMe = nrowsLinalg(1+me_rows+icpu*ncpu_rows)
    1425                 :             :               nrowsLinalgMeSum = sum(nrowsLinalg(1+me_rows:1+me_rows+(icpu-1)*ncpu_rows:ncpu_rows))
    1426                 :             :               froms=1+(ispinor-1)*nrowsLinalgMe/nspinor+(col-1)*nrowsLinalgMe+nrowsLinalgMeSum*ncolsColsRows
    1427                 :             :               frome=froms-1+nrowsLinalgMe/nspinor
    1428                 :             :               tos=1+nrowsLinalgMeSum/nspinor+(ispinor-1)*nrowsColsRows/nspinor+(col-1)*nrowsColsRows
    1429                 :             :               toe=tos-1+nrowsLinalgMe/nspinor
    1430                 :             :               do irow = 0, toe-tos
    1431                 :             :                 bufferMess(1,froms+irow) = bufferOrdered(1,tos+irow)
    1432                 :             :                 bufferMess(2,froms+irow) = bufferOrdered(2,tos+irow)
    1433                 :             :               end do
    1434                 :             :             end do
    1435                 :             :           end do
    1436                 :             :         end do
    1437                 :             :       else
    1438                 :             : #endif
    1439                 :             :         !$omp parallel do private(nrowsLinalgMe,nrowsLinalgMeSum,toe,tos,frome,froms), collapse(3)
    1440   [ +  +  -  +  :     3495612 :         do col = 1, ncolsColsRows
                      + ]
    1441         [ +  + ]:     7618114 :           do icpu = 0, ncpu_cols-1
    1442         [ +  + ]:    11459728 :             do ispinor = 1, nspinor
    1443                 :     5055784 :               nrowsLinalgMe = nrowsLinalg(1+me_rows+icpu*ncpu_rows)
    1444         [ +  + ]:     9840284 :               nrowsLinalgMeSum = sum(nrowsLinalg(1+me_rows:1+me_rows+(icpu-1)*ncpu_rows:ncpu_rows))
    1445                 :     5055784 :               froms=1+(ispinor-1)*nrowsLinalgMe/nspinor+(col-1)*nrowsLinalgMe+nrowsLinalgMeSum*ncolsColsRows
    1446                 :     5055784 :               frome=froms-1+nrowsLinalgMe/nspinor
    1447                 :     5055784 :               tos=1+nrowsLinalgMeSum/nspinor+(ispinor-1)*nrowsColsRows/nspinor+(col-1)*nrowsColsRows
    1448                 :     5055784 :               toe=tos-1+nrowsLinalgMe/nspinor
    1449   [ +  +  +  + ]:  1540809492 :               bufferMess(:,froms:frome) = bufferOrdered(:,tos:toe)
    1450                 :             :             end do
    1451                 :             :           end do
    1452                 :             :         end do
    1453                 :             : #if defined HAVE_GPU && defined HAVE_OPENMP_OFFLOAD
    1454                 :             :       end if
    1455                 :             : #endif
    1456                 :             :     end select
    1457                 :             : 
    1458   [ -  +  -  - ]:     1654988 :     if (xgTransposer%gpu_option /= ABI_GPU_DISABLED .and. xgTransposer%gpu_thread_limit /= 0) then
    1459                 :           0 :       call xomp_set_num_threads(nthreads_bak)
    1460                 :             :     end if
    1461                 :             : 
    1462                 :             : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
    1463                 :             :     ! if gpu enable restore OpenMP num threads to 1
    1464                 :             :     if (xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
    1465                 :             :        ! restore OMP_NUM_THREADS=1
    1466                 :             :        call xomp_set_num_threads(1)
    1467                 :             :     end if
    1468                 :             : 
    1469                 :             :     ! if gpu is enabled, transfer back data on GPU
    1470                 :             :     if (xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
    1471                 :             :        buffer_size = size(bufferOrdered) * dp
    1472                 :             :        call gpu_data_prefetch_async(C_LOC(bufferOrdered), buffer_size)
    1473                 :             :        call gpu_device_synchronize()
    1474                 :             :     end if
    1475                 :             : #endif
    1476                 :             : 
    1477                 :     1654988 :     call timab(tim_reorganize,2,tsec)
    1478                 :             : 
    1479                 :     1654988 :   end subroutine xgTransposer_reorganizeData
    1480                 :             : !!***
    1481                 :             : 
    1482                 :             : !!****f* m_xgTransposer/xgTransposer_getRank
    1483                 :             : !!
    1484                 :             : !! NAME
    1485                 :             : !! xgTransposer_getRank
    1486                 :             : 
    1487                 :           0 :   function xgTransposer_getRank(xgTransposer, comm) result(rank)
    1488                 :             :     type(xgTransposer_t), intent(in   ) :: xgTransposer
    1489                 :             :     integer             , intent(in   ) :: comm
    1490                 :             :     integer :: rank
    1491         [ #  # ]:           0 :     if ( (comm > ubound(xgTransposer%mpiData,1)) .or.  (comm < lbound(xgTransposer%mpiData,1)) ) then
    1492                 :           0 :       ABI_ERROR("Value for communicator is wrong")
    1493                 :             :     end if
    1494                 :           0 :     rank = xgTransposer%mpiData(comm)%rank
    1495                 :           0 :   end function xgTransposer_getRank
    1496                 :             : !!***
    1497                 :             : 
    1498                 :             : !!****f* m_xgTransposer/xgTransposer_getComm
    1499                 :             : !!
    1500                 :             : !! NAME
    1501                 :             : !! xgTransposer_getComm
    1502                 :             : 
    1503                 :           0 :   function xgTransposer_getComm(xgTransposer, comm1) result(communicator)
    1504                 :             :     type(xgTransposer_t), intent(in   ) :: xgTransposer
    1505                 :             :     integer             , intent(in   ) :: comm1
    1506                 :             :     integer :: communicator
    1507         [ #  # ]:           0 :     if ( (comm1 > ubound(xgTransposer%mpiData,1)) .or.  (comm1 < lbound(xgTransposer%mpiData,1)) ) then
    1508                 :           0 :       ABI_ERROR("Value for communicator is wrong")
    1509                 :             :     end if
    1510                 :           0 :     communicator = xgTransposer%mpiData(comm1)%comm
    1511                 :           0 :   end function xgTransposer_getComm
    1512                 :             : !!***
    1513                 :             : 
    1514                 :             : !!****f* m_xgTransposer/xgTransposer_free
    1515                 :             : !!
    1516                 :             : !! NAME
    1517                 :             : !! xgTransposer_free
    1518                 :             : 
    1519                 :      190814 :   subroutine xgTransposer_free(xgTransposer)
    1520                 :             : 
    1521                 :             :     type(xgTransposer_t), intent(inout) :: xgTransposer
    1522                 :             :     double precision :: tsec(2)
    1523                 :             :     integer :: i
    1524                 :             : 
    1525                 :      190814 :     call timab(tim_free,1,tsec)
    1526                 :             : #ifdef HAVE_MPI
    1527         [ +  + ]:      190814 :     if ( xgTransposer%type == TRANS_TYPE_CONSTRUCTED_NULL_COMM ) then
    1528                 :          10 :       call mpi_comm_free(xgTransposer%mpiData(MPI_ROWS)%comm,i)
    1529                 :          10 :       call mpi_comm_free(xgTransposer%mpiData(MPI_COLS)%comm,i)
    1530                 :          10 :       call mpi_comm_free(xgTransposer%mpiData(MPI_2DCART)%comm,i)
    1531                 :             :     end if
    1532                 :             : #else
    1533                 :             :     ABI_UNUSED(i)
    1534                 :             : #endif
    1535                 :             : 
    1536         [ +  - ]:      190814 :     if ( allocated(xgTransposer%lookup) ) then
    1537                 :      190814 :       ABI_FREE(xgTransposer%lookup)
    1538                 :             :     end if
    1539                 :             : 
    1540         [ +  - ]:      190814 :     if ( associated(xgTransposer%nrowsLinalg) ) then
    1541                 :      190814 :       ABI_FREE(xgTransposer%nrowsLinalg)
    1542                 :             :     end if
    1543                 :             : 
    1544         [ +  + ]:      190814 :     if ( associated(xgTransposer%buffer) ) then
    1545         [ +  - ]:      152600 :       if(xgTransposer%gpu_option == ABI_GPU_KOKKOS) then
    1546                 :             : #if defined HAVE_GPU && defined HAVE_YAKL
    1547                 :             :         ABI_FREE_MANAGED(xgTransposer%buffer)
    1548                 :             : #endif
    1549                 :             :       else
    1550                 :             :         if(xgTransposer%gpu_option == ABI_GPU_OPENMP) then
    1551                 :             : #ifdef HAVE_OPENMP_OFFLOAD
    1552                 :             : #ifdef HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
    1553                 :             :           !$OMP TARGET EXIT DATA MAP(delete:xgTransposer%buffer)
    1554                 :             : #else
    1555                 :             :           call xg_ompgpu_map_delete(xgTransposer%buffer)
    1556                 :             : #endif
    1557                 :             : #endif
    1558                 :             :         end if
    1559                 :      152600 :         ABI_FREE(xgTransposer%buffer)
    1560                 :             :       end if
    1561                 :             :     end if
    1562                 :             : 
    1563                 :      190814 :     call timab(tim_free,2,tsec)
    1564                 :             : 
    1565                 :             : #if defined HAVE_OPENMP_OFFLOAD && ! defined HAVE_OPENMP_OFFLOAD_DATASTRUCTURE
    1566                 :             :  contains
    1567                 :             :    ! NOTE : routine to be used in order to avoid using derived-type syntax in OMP directives
    1568                 :             :    subroutine xg_ompgpu_map_delete(xgTransposerBuffer)
    1569                 :             :     real(dp),intent(inout) :: xgTransposerBuffer(:,:)
    1570                 :             :     !$OMP TARGET EXIT DATA MAP(delete:xgTransposerBuffer)
    1571                 :             :    end subroutine xg_ompgpu_map_delete
    1572                 :             : #endif
    1573                 :      190814 :   end subroutine xgTransposer_free
    1574                 :             : !!***
    1575                 :             : 
    1576   [ #  #  #  #  :           0 : end module m_xgTransposer
          #  #  #  #  #  
                #  #  # ]
    1577                 :             : !!***
        

Generated by: LCOV version 2.3-1