Branch data Line data Source code
1 : : !!****m* ABINIT/m_prep_kgb
2 : : !! NAME
3 : : !! m_prep_kgb
4 : : !!
5 : : !! FUNCTION
6 : : !! This module provides wrappers that used to apply the full Hamiltonian or just the Vnl part
7 : : !! or to perform the FFT of the wavefunctions when the orbitals are distributed in linalg mode (paral_kgb = 1).
8 : : !!
9 : : !! COPYRIGHT
10 : : !! Copyright (C) 1998-2026 ABINIT group (FBottin,MT,GZ,MD,FDahm)
11 : : !! This file is distributed under the terms of the
12 : : !! GNU General Public License, see ~abinit/COPYING
13 : : !! or http://www.gnu.org/copyleft/gpl.txt .
14 : : !!
15 : : !! SOURCE
16 : :
17 : : #if defined HAVE_CONFIG_H
18 : : #include "config.h"
19 : : #endif
20 : :
21 : : #include "abi_common.h"
22 : :
23 : : module m_prep_kgb
24 : :
25 : : use defs_basis
26 : : use m_abicore
27 : : use m_errors
28 : : use m_xmpi
29 : : use m_xomp
30 : : use m_gputk
31 : : use m_abi_linalg
32 : :
33 : : use, intrinsic :: iso_c_binding, only: c_loc, c_size_t
34 : : use defs_abitypes, only : MPI_type
35 : : use m_time, only : timab
36 : : use m_bandfft_kpt, only : bandfft_kpt, bandfft_kpt_get_ikpt, bandfft_kpt_type
37 : : use m_pawcprj, only : pawcprj_type
38 : : use m_hamiltonian, only : gs_hamiltonian_type
39 : : use m_nonlop, only : nonlop
40 : : use m_getghc, only : multithreaded_getghc
41 : : use m_fft, only : fourwf, fourwf_optmem
42 : :
43 : : #if defined HAVE_GPU_CUDA
44 : : use m_manage_cuda
45 : : #endif
46 : :
47 : : #if defined HAVE_YAKL
48 : : use gator_mod
49 : : #endif
50 : :
51 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_YAKL)
52 : : use m_gpu_toolbox, only : CPU_DEVICE_ID, gpu_device_synchronize, gpu_data_prefetch_async
53 : : #endif
54 : :
55 : : implicit none
56 : :
57 : : private
58 : : !!***
59 : :
60 : : public :: prep_getghc
61 : : public :: prep_nonlop
62 : : public :: prep_fourwf
63 : : public :: prep_wavef_sym_do
64 : : public :: prep_wavef_sym_undo
65 : : public :: prep_index_wavef_bandpp
66 : : public :: prep_sort_wavef_spin
67 : : !!***
68 : :
69 : : contains
70 : : !!***
71 : :
72 : : !!****f* ABINIT/prep_getghc
73 : : !! NAME
74 : : !! prep_getghc
75 : : !!
76 : : !! FUNCTION
77 : : !! this routine prepares the data to the call of getghc.
78 : : !!
79 : : !! INPUTS
80 : : !! blocksize= size of block for FFT
81 : : !! cpopt=flag defining the status of cprjin%cp(:)=<Proj_i|Cnk> scalars (see below, side effects)
82 : : !! cwavef(2,npw*my_nspinor*blocksize)=planewave coefficients of wavefunction.
83 : : !! gs_hamk <type(gs_hamiltonian_type)>=all data for the hamiltonian at k
84 : : !! gvnlxc=matrix elements <G|Vnonlocal+VFockACE|C>
85 : : !! lambda=factor to be used when computing <G|H-lambda.S|C> - only for sij_opt=-1
86 : : !! Typically lambda is the eigenvalue (or its guess)
87 : : !! mpi_enreg=information about mpi parallelization
88 : : !! prtvol=control print volume and debugging output
89 : : !! sij_opt= -PAW ONLY- if 0, only matrix elements <G|H|C> have to be computed
90 : : !! (S=overlap) if 1, matrix elements <G|S|C> have to be computed in gsc in addition to ghc
91 : : !! if -1, matrix elements <G|H-lambda.S|C> have to be computed in ghc (gsc not used)
92 : : !!
93 : : !! OUTPUT
94 : : !! gwavef=(2,npw*my_nspinor*blocksize)=matrix elements <G|H|C> (if sij_opt>=0)
95 : : !! or <G|H-lambda.S|C> (if sij_opt=-1).
96 : : !! swavef=(2,npw*my_nspinor*blocksize)=matrix elements <G|S|C>.
97 : : !!
98 : : !! SIDE EFFECTS
99 : : !! ====== if gs_hamk%usepaw==1
100 : : !! cwaveprj(natom,my_nspinor*bandpp)= wave functions at k projected with nl projectors
101 : : !!
102 : : !! SOURCE
103 : :
104 [ - + - + : 59132 : subroutine prep_getghc(cwavef, gs_hamk, gvnlxc, gwavef, swavef, lambda, blocksize, &
- + - + ]
105 [ - + ]: 59132 : mpi_enreg, prtvol, sij_opt, cpopt, cwaveprj, &
106 : : already_transposed) ! optional argument
107 : :
108 : : !Arguments ------------------------------------
109 : : !scalars
110 : : integer,intent(in) :: blocksize,cpopt,prtvol,sij_opt
111 : : logical, intent(in),optional :: already_transposed
112 : : real(dp),intent(in) :: lambda
113 : : type(gs_hamiltonian_type),intent(inout) :: gs_hamk
114 : : type(mpi_type),intent(in) :: mpi_enreg
115 : : !arrays
116 : : real(dp),intent(in) :: cwavef(:,:)
117 : : real(dp),intent(inout) :: gvnlxc (:,:),gwavef(:,:),swavef(:,:)
118 : : type(pawcprj_type), intent(inout) :: cwaveprj(:,:)
119 : :
120 : : !Local variables-------------------------------
121 : : !scalars
122 : : integer,parameter :: tim_getghc=6
123 : : integer :: bandpp,bandpp_sym,idatarecv0,ier,ikpt_this_proc,iscalc,mcg,my_nspinor
124 : : integer :: nbval,ndatarecv,ndatarecv_tot,ndatasend_sym,nproc_band,nproc_fft
125 : : integer :: spaceComm
126 : : logical :: flag_inv_sym, do_transpose, local_gvnlxc
127 : : !character(len=500) :: msg
128 : : !arrays
129 : 59132 : integer,allocatable :: index_wavef_band(:),index_wavef_send(:),index_wavef_spband(:)
130 : 59132 : integer,allocatable :: rdisplsloc(:),recvcountsloc(:),sdisplsloc(:),sendcountsloc(:)
131 : 59132 : integer,contiguous, pointer :: kg_k_gather_sym(:,:)
132 : 59132 : integer,contiguous, pointer :: rdispls(:),rdispls_sym(:)
133 : 59132 : integer,contiguous, pointer :: recvcounts(:),recvcounts_sym(:),recvcounts_sym_tot(:)
134 : 59132 : integer,contiguous, pointer :: sdispls(:),sdispls_sym(:)
135 : 59132 : integer,contiguous, pointer :: sendcounts(:),sendcounts_sym(:),sendcounts_sym_all(:)
136 : 59132 : integer,contiguous, pointer :: tab_proc(:)
137 : : real(dp) :: tsec(2)
138 : 59132 : real(dp),allocatable,target :: cwavef_alltoall1(:,:), gvnlxc_alltoall1(:,:)
139 : 59132 : real(dp),allocatable,target :: gwavef_alltoall1(:,:), swavef_alltoall1(:,:)
140 : :
141 : : #if defined HAVE_GPU && defined HAVE_YAKL
142 : : real(c_double), contiguous, pointer :: cwavef_alltoall2(:,:) => null()
143 : : real(c_double), contiguous, pointer :: gvnlxc_alltoall2(:,:) => null()
144 : : real(c_double), contiguous, pointer :: gwavef_alltoall2(:,:) => null()
145 : : real(c_double), contiguous, pointer :: swavef_alltoall2(:,:) => null()
146 : : #else
147 : 59132 : real(dp),allocatable,target :: cwavef_alltoall2(:,:)
148 : 59132 : real(dp),allocatable,target :: gvnlxc_alltoall2(:,:)
149 : 59132 : real(dp),allocatable,target :: gwavef_alltoall2(:,:)
150 : 59132 : real(dp),allocatable,target :: swavef_alltoall2(:,:)
151 : : #endif
152 : :
153 : 59132 : real(dp),pointer :: ewavef_alltoall_sym(:,:)
154 : 59132 : real(dp),pointer :: gvnlxc_alltoall_sym(:,:)
155 : 59132 : real(dp),pointer :: gwavef_alltoall_sym(:,:)
156 : 59132 : real(dp),pointer :: swavef_alltoall_sym(:,:)
157 : :
158 : : ! *************************************************************************
159 : :
160 : 59132 : call timab(630,1,tsec)
161 : 59132 : call timab(631,3,tsec)
162 : :
163 : : !Some inits
164 : 59132 : nproc_band = mpi_enreg%nproc_band
165 : 59132 : nproc_fft = mpi_enreg%nproc_fft
166 : 59132 : bandpp = mpi_enreg%bandpp
167 : 59132 : my_nspinor = max(1,gs_hamk%nspinor/mpi_enreg%nproc_spinor)
168 : :
169 : 59132 : do_transpose = .true.
170 [ - + ]: 59132 : if(present(already_transposed)) then
171 [ + + ]: 59132 : if(already_transposed) do_transpose = .false.
172 : : end if
173 : :
174 [ + + + + : 59132 : flag_inv_sym = (gs_hamk%istwf_k==2 .and. any(gs_hamk%ngfft(7) == [401,402,312,512]))
+ - ]
175 : : if (flag_inv_sym) then
176 : 3274 : gs_hamk%istwf_k = 1
177 [ + + ]: 3274 : if (modulo(bandpp,2)==0) bandpp_sym = bandpp/2
178 [ + + ]: 3274 : if (modulo(bandpp,2)/=0) bandpp_sym = bandpp
179 : : end if
180 : :
181 : : !Check sizes
182 : 59132 : mcg=2*gs_hamk%npw_fft_k*my_nspinor*bandpp
183 [ + + ]: 59132 : if (do_transpose) mcg=2*gs_hamk%npw_k*my_nspinor*blocksize
184 [ + + - + ]: 177396 : if (size(cwavef)<mcg) then
185 : 0 : ABI_BUG('wrong size for cwavef!')
186 : : end if
187 [ + + - + ]: 177396 : if (size(gwavef)<mcg) then
188 : 0 : ABI_BUG('wrong size for gwavef!')
189 : : end if
190 : 59132 : local_gvnlxc = .false.
191 [ + + + + ]: 177396 : if (size(gvnlxc)<=1) then
192 : 27896 : local_gvnlxc = .true.
193 : : end if
194 [ + + + + : 177396 : if ((.not.local_gvnlxc).and.(size(gvnlxc)<mcg)) then
- + ]
195 : 0 : ABI_BUG('wrong size for gvnlxc!')
196 : : end if
197 [ + + ]: 59132 : if (sij_opt==1) then
198 [ + + - + ]: 75864 : if (size(swavef)<mcg) then
199 : 0 : ABI_BUG('wrong size for swavef!')
200 : : end if
201 : : end if
202 [ + + + + ]: 59132 : if (gs_hamk%usepaw==1.and.cpopt>=0) then
203 [ + + - + ]: 77808 : if (size(cwaveprj)<gs_hamk%natom*my_nspinor*bandpp) then
204 : 0 : ABI_BUG('wrong size for cwaveprj!')
205 : : end if
206 : : end if
207 : :
208 : : !====================================================================================
209 : :
210 : 59132 : spaceComm=mpi_enreg%comm_fft
211 [ + - ]: 59132 : if(mpi_enreg%paral_kgb==1) spaceComm=mpi_enreg%comm_band
212 : :
213 : 59132 : ikpt_this_proc=bandfft_kpt_get_ikpt()
214 : :
215 [ - + + - : 177396 : ABI_MALLOC(sendcountsloc,(nproc_band))
+ - + - -
+ - + ]
216 [ - + + - : 118264 : ABI_MALLOC(sdisplsloc ,(nproc_band))
- + - + ]
217 [ - + + - : 118264 : ABI_MALLOC(recvcountsloc,(nproc_band))
- + - + ]
218 [ - + + - : 118264 : ABI_MALLOC(rdisplsloc ,(nproc_band))
- + - + ]
219 : :
220 : 59132 : recvcounts =>bandfft_kpt(ikpt_this_proc)%recvcounts(:)
221 : 59132 : sendcounts =>bandfft_kpt(ikpt_this_proc)%sendcounts(:)
222 : 59132 : rdispls =>bandfft_kpt(ikpt_this_proc)%rdispls (:)
223 : 59132 : sdispls =>bandfft_kpt(ikpt_this_proc)%sdispls (:)
224 : 59132 : ndatarecv = bandfft_kpt(ikpt_this_proc)%ndatarecv
225 : :
226 [ + + ]: 59132 : if (flag_inv_sym ) then
227 : 3274 : idatarecv0 = bandfft_kpt(ikpt_this_proc)%idatarecv0
228 : 3274 : ndatarecv_tot = bandfft_kpt(ikpt_this_proc)%ndatarecv_tot
229 : 3274 : ndatasend_sym = bandfft_kpt(ikpt_this_proc)%ndatasend_sym
230 : 3274 : kg_k_gather_sym =>bandfft_kpt(ikpt_this_proc)%kg_k_gather_sym(:,:)
231 : 3274 : rdispls_sym =>bandfft_kpt(ikpt_this_proc)%rdispls_sym(:)
232 : 3274 : recvcounts_sym =>bandfft_kpt(ikpt_this_proc)%recvcounts_sym(:)
233 : 3274 : recvcounts_sym_tot =>bandfft_kpt(ikpt_this_proc)%recvcounts_sym_tot(:)
234 : 3274 : sdispls_sym =>bandfft_kpt(ikpt_this_proc)%sdispls_sym(:)
235 : 3274 : sendcounts_sym =>bandfft_kpt(ikpt_this_proc)%sendcounts_sym(:)
236 : 3274 : sendcounts_sym_all =>bandfft_kpt(ikpt_this_proc)%sendcounts_sym_all(:)
237 : 3274 : tab_proc =>bandfft_kpt(ikpt_this_proc)%tab_proc(:)
238 : : end if
239 : 59132 : iscalc=(sij_opt+1)/2 ! 0 if S not calculated, 1 otherwise
240 : 59132 : nbval=(ndatarecv*my_nspinor*bandpp)*iscalc
241 : :
242 : : if ( ((.not.flag_inv_sym) .and. bandpp==1 .and. mpi_enreg%paral_spinor==0 .and. my_nspinor==2 ).or. &
243 [ + + + - : 59132 : & ((.not.flag_inv_sym) .and. bandpp>1) .or. flag_inv_sym ) then
+ - + + +
+ ]
244 [ - + + - : 130320 : ABI_MALLOC(cwavef_alltoall1,(2,ndatarecv*my_nspinor*bandpp))
+ - + - -
+ - + ]
245 [ - + + - : 86880 : ABI_MALLOC(gwavef_alltoall1,(2,ndatarecv*my_nspinor*bandpp))
- + - + ]
246 [ - + + - : 86880 : ABI_MALLOC(swavef_alltoall1,(2,ndatarecv*my_nspinor*bandpp))
- + - + ]
247 [ + + ]: 43440 : if (local_gvnlxc) then
248 [ - + ]: 27896 : ABI_MALLOC(gvnlxc_alltoall1,(0,0))
249 : : else
250 [ - + + - : 31088 : ABI_MALLOC(gvnlxc_alltoall1,(2,ndatarecv*my_nspinor*bandpp))
- + - + ]
251 : : end if
252 [ + + + + ]: 228702600 : swavef_alltoall1(:,:)=zero
253 [ + + + + ]: 228702600 : cwavef_alltoall1(:,:)=zero
254 [ + + + + ]: 228702600 : gwavef_alltoall1(:,:)=zero
255 [ + + + + : 152784600 : if (.not.local_gvnlxc) gvnlxc_alltoall1(:,:)=zero
+ + ]
256 : : end if
257 : :
258 [ + - ]: 59132 : if(gs_hamk%gpu_option==ABI_GPU_KOKKOS) then
259 : : #if defined HAVE_GPU && defined HAVE_YAKL
260 : : ABI_MALLOC_MANAGED(cwavef_alltoall2,(/2,ndatarecv*my_nspinor*bandpp/))
261 : : ABI_MALLOC_MANAGED(gwavef_alltoall2,(/2,ndatarecv*my_nspinor*bandpp/))
262 : : ABI_MALLOC_MANAGED(swavef_alltoall2,(/2,ndatarecv*my_nspinor*bandpp/))
263 : : if (local_gvnlxc) then
264 : : ABI_MALLOC_MANAGED(gvnlxc_alltoall2,(/0,0/))
265 : : else
266 : : ABI_MALLOC_MANAGED(gvnlxc_alltoall2,(/2,ndatarecv*my_nspinor*bandpp/))
267 : : end if
268 : : #endif
269 : : else
270 [ - + + - : 177396 : ABI_MALLOC(cwavef_alltoall2,(2,ndatarecv*my_nspinor*bandpp))
+ - + - -
+ - + ]
271 [ - + + - : 118264 : ABI_MALLOC(gwavef_alltoall2,(2,ndatarecv*my_nspinor*bandpp))
- + - + ]
272 [ - + + - : 118264 : ABI_MALLOC(swavef_alltoall2,(2,ndatarecv*my_nspinor*bandpp))
- + - + -
+ ]
273 [ + + ]: 59132 : if (local_gvnlxc) then
274 [ - + ]: 27896 : ABI_MALLOC(gvnlxc_alltoall2,(0,0))
275 : : else
276 [ - + + - : 62472 : ABI_MALLOC(gvnlxc_alltoall2,(2,ndatarecv*my_nspinor*bandpp))
- + - + ]
277 : : end if
278 : : end if
279 : :
280 [ + + + + ]: 251503076 : swavef_alltoall2(:,:)=zero
281 [ + + + + ]: 251503076 : cwavef_alltoall2(:,:)=zero
282 [ + + + + ]: 251503076 : gwavef_alltoall2(:,:)=zero
283 [ + + + + : 175585076 : if (.not.local_gvnlxc) gvnlxc_alltoall2(:,:)=zero
+ + ]
284 : :
285 [ + + ]: 204868 : recvcountsloc(:)=recvcounts(:)*2*my_nspinor*bandpp
286 [ + + ]: 204868 : rdisplsloc(:)=rdispls(:)*2*my_nspinor*bandpp
287 [ + + ]: 204868 : sendcountsloc(:)=sendcounts(:)*2*my_nspinor
288 [ + + ]: 204868 : sdisplsloc(:)=sdispls(:)*2*my_nspinor
289 : 59132 : call timab(631,2,tsec)
290 : :
291 [ + + ]: 59132 : if(do_transpose) then
292 : 30900 : call timab(545,3,tsec)
293 : : if ( ((.not.flag_inv_sym) .and. bandpp==1 .and. mpi_enreg%paral_spinor==0 .and. my_nspinor==2 ).or. &
294 [ + + + - : 30900 : & ((.not.flag_inv_sym) .and. bandpp>1) .or. flag_inv_sym ) then
+ - + + +
+ ]
295 : : call xmpi_alltoallv(cwavef,sendcountsloc,sdisplsloc,cwavef_alltoall1,&
296 : 15208 : & recvcountsloc,rdisplsloc,spaceComm,ier)
297 : : else
298 : : call xmpi_alltoallv(cwavef,sendcountsloc,sdisplsloc,cwavef_alltoall2,&
299 : 15692 : & recvcountsloc,rdisplsloc,spaceComm,ier)
300 : : end if
301 : 30900 : call timab(545,2,tsec)
302 : : else
303 : : ! Here, we cheat, and use DCOPY to bypass some compiler's overzealous bound-checking
304 : : ! (ndatarecv*my_nspinor*bandpp might be greater than the declared size of cwavef)
305 [ - + - - : 28232 : call DCOPY(2*ndatarecv*my_nspinor*bandpp, cwavef, 1, cwavef_alltoall2, 1)
- - - - -
+ - - -
- ]
306 : : end if
307 : :
308 : : !====================================================================
309 [ + + ]: 59132 : if ((.not.(flag_inv_sym)) .and. (bandpp==1)) then
310 [ + - + - : 15692 : if (do_transpose .and. mpi_enreg%paral_spinor==0.and.my_nspinor==2)then
- + ]
311 : 0 : call timab(632,3,tsec)
312 : : ! Sort to have all ispinor=1 first, then all ispinor=2
313 : 0 : call prep_sort_wavef_spin(nproc_band,my_nspinor,ndatarecv,recvcounts,rdispls,index_wavef_spband)
314 [ # # # # ]: 0 : cwavef_alltoall2(:,:)=cwavef_alltoall1(:,index_wavef_spband)
315 : 0 : call timab(632,2,tsec)
316 : : end if
317 : :
318 : 15692 : call timab(635,3,tsec)
319 : : call multithreaded_getghc(cpopt,cwavef_alltoall2,cwaveprj,gwavef_alltoall2,swavef_alltoall2(:,1:nbval),&
320 : 15692 : & gs_hamk,gvnlxc_alltoall2,lambda,mpi_enreg,1,prtvol,sij_opt,tim_getghc,0)
321 : 15692 : call timab(635,2,tsec)
322 : :
323 [ + - + - : 15692 : if (do_transpose .and. mpi_enreg%paral_spinor==0.and.my_nspinor==2)then
- + ]
324 : 0 : call timab(634,3,tsec)
325 [ # # # # ]: 0 : gwavef_alltoall1(:,index_wavef_spband)=gwavef_alltoall2(:,:)
326 [ # # # # : 0 : if (sij_opt==1) swavef_alltoall1(:,index_wavef_spband)=swavef_alltoall2(:,:)
# # ]
327 [ # # # # : 0 : if (.not.local_gvnlxc) gvnlxc_alltoall1(:,index_wavef_spband)=gvnlxc_alltoall2(:,:)
# # ]
328 [ # # ]: 0 : ABI_FREE(index_wavef_spband)
329 : 0 : call timab(634,2,tsec)
330 : : end if
331 : :
332 [ + + ]: 43440 : else if ((.not.(flag_inv_sym)) .and. (bandpp>1)) then
333 : : ! -------------------------------------------------------------
334 : : ! Computation of the index to class the waves functions below bandpp
335 : : ! -------------------------------------------------------------
336 : :
337 [ + + ]: 40166 : if(do_transpose) then
338 : 13590 : call timab(632,3,tsec)
339 : : call prep_index_wavef_bandpp(nproc_band,bandpp,&
340 : 13590 : & my_nspinor,ndatarecv, recvcounts,rdispls, index_wavef_band)
341 : : ! -------------------------------------------------------
342 : : ! Sorting of the waves functions below bandpp
343 : : ! -------------------------------------------------------
344 [ + + + + ]: 125187702 : cwavef_alltoall2(:,:) = cwavef_alltoall1(:,index_wavef_band)
345 : 13590 : call timab(632,2,tsec)
346 : : end if
347 : :
348 : : ! ----------------------
349 : : ! Fourier transformation
350 : : ! ----------------------
351 : 40166 : call timab(636,3,tsec)
352 : : call multithreaded_getghc(cpopt,cwavef_alltoall2,cwaveprj,gwavef_alltoall2,swavef_alltoall2,gs_hamk,&
353 : 40166 : & gvnlxc_alltoall2,lambda,mpi_enreg,bandpp,prtvol,sij_opt,tim_getghc,0)
354 : 40166 : call timab(636,2,tsec)
355 : :
356 : : ! -----------------------------------------------------
357 : : ! Sorting of waves functions below the processors
358 : : ! -----------------------------------------------------
359 [ + + ]: 40166 : if(do_transpose) then
360 : 13590 : call timab(634,3,tsec)
361 [ + + + + ]: 125187702 : gwavef_alltoall1(:,index_wavef_band) = gwavef_alltoall2(:,:)
362 [ + + + + : 15165318 : if (sij_opt==1) swavef_alltoall1(:,index_wavef_band) = swavef_alltoall2(:,:)
+ + ]
363 [ + - + + : 125187702 : if (.not.local_gvnlxc) gvnlxc_alltoall1(:,index_wavef_band) = gvnlxc_alltoall2(:,:)
+ + ]
364 [ - + ]: 13590 : ABI_FREE(index_wavef_band)
365 : 13590 : call timab(634,2,tsec)
366 : : end if
367 : :
368 : :
369 [ + - ]: 3274 : else if (flag_inv_sym) then
370 : :
371 : : ! -------------------------------------------------------------
372 : : ! Computation of the index to class the waves functions below bandpp
373 : : ! -------------------------------------------------------------
374 [ + + ]: 3274 : if(do_transpose) then
375 : 1618 : call timab(632,3,tsec)
376 : : call prep_index_wavef_bandpp(nproc_band,bandpp,&
377 : : & my_nspinor,ndatarecv,&
378 : : & recvcounts,rdispls,&
379 : 1618 : & index_wavef_band)
380 : :
381 : : ! -------------------------------------------------------
382 : : ! Sorting the wave functions below bandpp
383 : : ! -------------------------------------------------------
384 [ + + + + ]: 26785450 : cwavef_alltoall2(:,:) = cwavef_alltoall1(:,index_wavef_band)
385 : : end if
386 : :
387 : : ! ------------------------------------------------------------
388 : : ! We associate the waves functions by two
389 : : ! ------------------------------------------------------------
390 : : call prep_wavef_sym_do(mpi_enreg,bandpp,my_nspinor,&
391 : : & ndatarecv,&
392 : : & ndatarecv_tot,ndatasend_sym,tab_proc,&
393 : : & cwavef_alltoall2,&
394 : : & sendcounts_sym,sdispls_sym,&
395 : : & recvcounts_sym,rdispls_sym,&
396 : : & ewavef_alltoall_sym,&
397 : 3274 : & index_wavef_send)
398 : :
399 : : ! ------------------------------------------------------------
400 : : ! Allocation
401 : : ! ------------------------------------------------------------
402 [ - + + - : 9822 : ABI_MALLOC(gwavef_alltoall_sym,(2,ndatarecv_tot*bandpp_sym))
+ - + - -
+ - + ]
403 [ + + + - : 9822 : ABI_MALLOC(swavef_alltoall_sym,(2,(ndatarecv_tot*bandpp_sym)*iscalc))
+ - + + -
+ - + ]
404 [ + + ]: 3274 : if (local_gvnlxc) then
405 [ - + ]: 1656 : ABI_MALLOC(gvnlxc_alltoall_sym ,(0,0))
406 : : else
407 [ - + + - : 3236 : ABI_MALLOC(gvnlxc_alltoall_sym ,(2,ndatarecv_tot*bandpp_sym))
- + - + ]
408 : : end if
409 [ + + + + ]: 27990616 : gwavef_alltoall_sym(:,:)=zero
410 [ + + + + ]: 665260 : swavef_alltoall_sym(:,:)=zero
411 [ + + + + : 26778514 : if (.not.local_gvnlxc) gvnlxc_alltoall_sym(:,:)=zero
+ + ]
412 : :
413 : 3274 : call timab(632,2,tsec)
414 : :
415 : : ! ------------------------------------------------------------
416 : : ! Fourier calculation
417 : : ! ------------------------------------------------------------
418 : 3274 : call timab(637,3,tsec)
419 : : call multithreaded_getghc(cpopt,ewavef_alltoall_sym,cwaveprj,gwavef_alltoall_sym,swavef_alltoall_sym,gs_hamk,&
420 : : & gvnlxc_alltoall_sym,lambda,mpi_enreg,bandpp_sym,prtvol,sij_opt,tim_getghc,1,&
421 : 3274 : & kg_fft_k=kg_k_gather_sym,filter_dilatmx_loc=.false.)
422 : 3274 : call timab(637,2,tsec)
423 : :
424 : 3274 : call timab(633,3,tsec)
425 : :
426 : : ! ------------------------------------------------------------
427 : : ! We dissociate each wave function in two waves functions
428 : : ! gwavef is classed below of bandpp
429 : : ! ------------------------------------------------------------
430 : : call prep_wavef_sym_undo(mpi_enreg,bandpp,my_nspinor,&
431 : : & ndatarecv,&
432 : : & ndatarecv_tot,ndatasend_sym,idatarecv0,&
433 : : & gwavef_alltoall2,&
434 : : & sendcounts_sym,sdispls_sym,&
435 : : & recvcounts_sym,rdispls_sym,&
436 : : & gwavef_alltoall_sym,&
437 : 3274 : & index_wavef_send)
438 [ + + ]: 3274 : if (sij_opt==1)then
439 : : call prep_wavef_sym_undo(mpi_enreg,bandpp,my_nspinor,&
440 : : & ndatarecv,&
441 : : & ndatarecv_tot,ndatasend_sym,idatarecv0,&
442 : : & swavef_alltoall2,&
443 : : & sendcounts_sym,sdispls_sym,&
444 : : & recvcounts_sym,rdispls_sym,&
445 : : & swavef_alltoall_sym,&
446 : 1062 : & index_wavef_send)
447 : : end if
448 [ + + ]: 3274 : if (.not.local_gvnlxc) call prep_wavef_sym_undo(mpi_enreg,bandpp,my_nspinor,&
449 : : & ndatarecv,&
450 : : & ndatarecv_tot,ndatasend_sym,idatarecv0,&
451 : : & gvnlxc_alltoall2,&
452 : : & sendcounts_sym,sdispls_sym,&
453 : : & recvcounts_sym,rdispls_sym,&
454 : : & gvnlxc_alltoall_sym,&
455 : 1618 : & index_wavef_send)
456 : :
457 [ - + ]: 3274 : ABI_FREE(ewavef_alltoall_sym)
458 [ - + ]: 3274 : ABI_FREE(index_wavef_send)
459 [ - + ]: 3274 : ABI_FREE(gwavef_alltoall_sym)
460 [ - + ]: 3274 : ABI_FREE(swavef_alltoall_sym)
461 [ - + ]: 3274 : ABI_FREE(gvnlxc_alltoall_sym)
462 : :
463 : : ! -------------------------------------------
464 : : ! We call getghc to calculate the nl matrix elements.
465 : : ! --------------------------------------------
466 : 3274 : gs_hamk%istwf_k=2
467 : : !!write(std_out,*)"Setting iswfk_k to 2"
468 : :
469 : 3274 : call timab(633,2,tsec)
470 : :
471 : 3274 : call timab(638,3,tsec)
472 : : call multithreaded_getghc(cpopt,cwavef_alltoall2,cwaveprj,gwavef_alltoall2,swavef_alltoall2,gs_hamk,&
473 : 3274 : & gvnlxc_alltoall2,lambda,mpi_enreg,bandpp,prtvol,sij_opt,tim_getghc,2)
474 : 3274 : call timab(638,2,tsec)
475 : :
476 : 3274 : call timab(634,3,tsec)
477 : :
478 : 3274 : gs_hamk%istwf_k=1
479 : :
480 : : ! -------------------------------------------------------
481 : : ! Sorting the wave functions below the processors
482 : : ! -------------------------------------------------------
483 [ + + ]: 3274 : if(do_transpose) then
484 : : ! cwavef_alltoall(:,index_wavef_band) = cwavef_alltoall(:,:) ! NOT NEEDED
485 [ + + + + ]: 26785450 : gwavef_alltoall1(:,index_wavef_band) = gwavef_alltoall2(:,:)
486 [ + + + + : 374842 : if (sij_opt==1) swavef_alltoall1(:,index_wavef_band) = swavef_alltoall2(:,:)
+ + ]
487 [ + - + + : 26785450 : if (.not.local_gvnlxc) gvnlxc_alltoall1(:,index_wavef_band) = gvnlxc_alltoall2(:,:)
+ + ]
488 [ - + ]: 1618 : ABI_FREE(index_wavef_band)
489 : 1618 : call timab(634,2,tsec)
490 : : end if
491 : :
492 : : end if
493 : : !====================================================================
494 : :
495 [ + - ]: 30900 : if(do_transpose) then
496 : :
497 : 30900 : call timab(545,3,tsec)
498 : : if ( ((.not.flag_inv_sym) .and. bandpp==1 .and. mpi_enreg%paral_spinor==0 .and. my_nspinor==2 ).or. &
499 [ + + + + : 30900 : & ((.not.flag_inv_sym) .and. bandpp>1) .or. flag_inv_sym ) then
+ - + - +
+ ]
500 [ + + ]: 15208 : if (sij_opt==1) then
501 : : call xmpi_alltoallv(swavef_alltoall1,recvcountsloc,rdisplsloc,swavef,&
502 : 9088 : & sendcountsloc,sdisplsloc,spaceComm,ier)
503 : : end if
504 [ + - ]: 15208 : if (.not.local_gvnlxc) call xmpi_alltoallv(gvnlxc_alltoall1,recvcountsloc,rdisplsloc,gvnlxc,&
505 : 15208 : & sendcountsloc,sdisplsloc,spaceComm,ier)
506 : : call xmpi_alltoallv(gwavef_alltoall1,recvcountsloc,rdisplsloc,gwavef,&
507 : 15208 : & sendcountsloc,sdisplsloc,spaceComm,ier)
508 : : else
509 [ + - ]: 15692 : if (sij_opt==1) then
510 : : call xmpi_alltoallv(swavef_alltoall2,recvcountsloc,rdisplsloc,swavef,&
511 : 15692 : & sendcountsloc,sdisplsloc,spaceComm,ier)
512 : : end if
513 [ + - ]: 15692 : if (.not.local_gvnlxc) call xmpi_alltoallv(gvnlxc_alltoall2,recvcountsloc,rdisplsloc,gvnlxc,&
514 : 15692 : & sendcountsloc,sdisplsloc,spaceComm,ier)
515 : : call xmpi_alltoallv(gwavef_alltoall2,recvcountsloc,rdisplsloc,gwavef,&
516 : 15692 : & sendcountsloc,sdisplsloc,spaceComm,ier)
517 : : end if
518 : :
519 : 30900 : call timab(545,2,tsec)
520 : : else
521 [ + + ]: 28232 : if(sij_opt == 1) then
522 [ - + - - : 508 : call DCOPY(2*ndatarecv*my_nspinor*bandpp, swavef_alltoall2, 1, swavef, 1)
- - - - -
+ - - -
- ]
523 : : end if
524 [ + + - + : 28232 : if (.not.local_gvnlxc) call DCOPY(2*ndatarecv*my_nspinor*bandpp, gvnlxc_alltoall2, 1, gvnlxc, 1)
- - - - -
- - + - -
- - ]
525 [ - + - - : 28232 : call DCOPY(2*ndatarecv*my_nspinor*bandpp, gwavef_alltoall2, 1, gwavef, 1)
- - - - -
+ - - -
- ]
526 : : end if
527 : :
528 : : !====================================================================
529 [ + + ]: 59132 : if (flag_inv_sym) then
530 : 3274 : gs_hamk%istwf_k = 2
531 : : end if
532 : : !====================================================================
533 [ - + ]: 59132 : ABI_FREE(sendcountsloc)
534 [ - + ]: 59132 : ABI_FREE(sdisplsloc)
535 [ - + ]: 59132 : ABI_FREE(recvcountsloc)
536 [ - + ]: 59132 : ABI_FREE(rdisplsloc)
537 : :
538 [ + - ]: 59132 : if(gs_hamk%gpu_option==ABI_GPU_KOKKOS) then
539 : : #if defined HAVE_GPU && defined HAVE_YAKL
540 : : ABI_FREE_MANAGED(cwavef_alltoall2)
541 : : ABI_FREE_MANAGED(gwavef_alltoall2)
542 : : ABI_FREE_MANAGED(gvnlxc_alltoall2)
543 : : ABI_FREE_MANAGED(swavef_alltoall2)
544 : : #endif
545 : : else
546 [ - + ]: 59132 : ABI_FREE(cwavef_alltoall2)
547 [ - + ]: 59132 : ABI_FREE(gwavef_alltoall2)
548 [ - + ]: 59132 : ABI_FREE(gvnlxc_alltoall2)
549 [ - + ]: 59132 : ABI_FREE(swavef_alltoall2)
550 : : end if
551 : :
552 : : if ( ((.not.flag_inv_sym) .and. bandpp==1 .and. mpi_enreg%paral_spinor==0 .and. my_nspinor==2 ).or. &
553 [ + + + + : 59132 : & ((.not.flag_inv_sym) .and. bandpp>1) .or. flag_inv_sym ) then
+ - + - +
+ ]
554 [ - + ]: 43440 : ABI_FREE(cwavef_alltoall1)
555 [ - + ]: 43440 : ABI_FREE(gwavef_alltoall1)
556 [ - + ]: 43440 : ABI_FREE(gvnlxc_alltoall1)
557 [ - + ]: 43440 : ABI_FREE(swavef_alltoall1)
558 : : end if
559 : :
560 : 59132 : call timab(630,2,tsec)
561 : :
562 [ - + - + : 118264 : end subroutine prep_getghc
- + - + -
+ - + - +
- + - + -
+ - + ]
563 : : !!***
564 : :
565 : : !!****f* abinit/prep_nonlop
566 : : !! NAME
567 : : !! prep_nonlop
568 : : !!
569 : : !! FUNCTION
570 : : !! this routine prepares the data to the call of nonlop.
571 : : !!
572 : : !! INPUTS
573 : : !! choice: chooses possible output:
574 : : !! choice=1 => a non-local energy contribution
575 : : !! =2 => a gradient with respect to atomic position(s)
576 : : !! =3 => a gradient with respect to strain(s)
577 : : !! =23=> a gradient with respect to atm. pos. and strain(s)
578 : : !! =4 => a 2nd derivative with respect to atomic pos.
579 : : !! =24=> a gradient and 2nd derivative with respect to atomic pos.
580 : : !! =5 => a gradient with respect to k wavevector
581 : : !! =6 => 2nd derivatives with respect to strain and atm. pos.
582 : : !! =7 => no operator, just projections
583 : : !! blocksize= size of block for FFT
584 : : !! cpopt=flag defining the status of cwaveprj=<Proj_i|Cnk> scalars (see below, side effects)
585 : : !! cwavef(2,npw*my_nspinor*blocksize)=planewave coefficients of wavefunction.
586 : : !! gvnlxc=matrix elements <G|Vnonlocal+VFockACE|C>
587 : : !! hamk <type(gs_hamiltonian_type)>=data defining the Hamiltonian at a given k (NL part involved here)
588 : : !! idir=direction of the - atom to be moved in the case (choice=2,signs=2),
589 : : !! - k point direction in the case (choice=5,signs=2)
590 : : !! - strain component (1:6) in the case (choice=2,signs=2) or (choice=6,signs=1)
591 : : !! lambdablock(blocksize)=factor to be used when computing (Vln-lambda.S) - only for paw_opt=2
592 : : !! mpi_enreg=information about mpi parallelization
593 : : !! nnlout=dimension of enlout (when signs=1):
594 : : !! ntypat=number of types of atoms in cell
595 : : !! paw_opt= define the nonlocal operator concerned with
596 : : !! signs= if 1, get contracted elements (energy, forces, stress, ...)
597 : : !! if 2, applies the non-local operator to a function in reciprocal space
598 : : !! tim_nonlop=timing code of the calling routine (can be set to 0 if not attributed)
599 : : !! vectproj(2,nprojs,my_nspinor*ndat)=Optional, vector to be used instead of cprjin%cp when provided
600 : : !!
601 : : !! OUTPUT
602 : : !! ==== if (signs==1) ====
603 : : !! enlout_block(nnlout)=
604 : : !! if paw_opt==0, 1 or 2: contribution of this block of states to the nl part of various properties
605 : : !! if paw_opt==3: contribution of this block of states to <c|S|c> (where S=overlap when PAW)
606 : : !! ==== if (signs==2) ====
607 : : !! if paw_opt==0, 1, 2 or 4:
608 : : !! gvnlc(2,my_nspinor*npw)=result of the application of the nl operator
609 : : !! or one of its derivative to the input vect.
610 : : !! if paw_opt==3 or 4:
611 : : !! gsc(2,my_nspinor*npw*(paw_opt/3))=result of the aplication of (I+S)
612 : : !! to the input vect. (where S=overlap when PAW)
613 : : !!
614 : : !! SIDE EFFECTS
615 : : !! ==== ONLY IF useylm=1
616 : : !! cwaveprj(natom,my_nspinor) <type(pawcprj_type)>=projected input wave function |c> on non-local projector
617 : : !! =<p_lmn|c> and derivatives
618 : : !! Treatment depends on cpopt parameter:
619 : : !! if cpopt=-1, <p_lmn|in> (and derivatives)
620 : : !! have to be computed (and not saved)
621 : : !! if cpopt= 0, <p_lmn|in> have to be computed and saved
622 : : !! derivatives are eventually computed but not saved
623 : : !! if cpopt= 1, <p_lmn|in> and first derivatives have to be computed and saved
624 : : !! other derivatives are eventually computed but not saved
625 : : !! if cpopt= 2 <p_lmn|in> are already in memory;
626 : : !! only derivatives are computed here and not saved
627 : : !! (if useylm=0, should have cpopt=-1)
628 : : !!
629 : : !! NOTES
630 : : !! cprj (as well as cg) is distributed over band processors.
631 : : !! Only the mod((iband-1)/mpi_enreg%bandpp,mpi_enreg%nproc_band) projected WFs are stored on each proc.
632 : : !!
633 : : !! SOURCE
634 : :
635 [ - + ]: 266542 : subroutine prep_nonlop(choice,cpopt,cwaveprj,enlout_block,hamk,idir,lambdablock,&
636 [ - + - + : 266542 : blocksize,mpi_enreg,nnlout,paw_opt,signs,gsc, tim_nonlop,cwavef,gvnlc, &
- + ]
637 [ + + + - ]: 266542 : already_transposed,gpu_option,vectproj,cwavef_tr,select_k) ! optional
638 : :
639 : : !Arguments ------------------------------------
640 : : integer, intent(in) :: blocksize,choice,cpopt,idir,signs,nnlout,paw_opt
641 : : integer,optional,intent(in) :: select_k
642 : : logical,optional,intent(in) :: already_transposed
643 : : integer,optional,intent(in) :: gpu_option
644 : : real(dp), intent(in) :: lambdablock(blocksize)
645 : : real(dp), intent(out) , target :: enlout_block(nnlout*blocksize),gvnlc(:,:),gsc(:,:)
646 : : real(dp), intent(inout), target :: cwavef(:,:)
647 : : real(dp),intent(inout), pointer,optional :: cwavef_tr(:,:)
648 : : real(dp),contiguous, optional,intent(inout) :: vectproj(:,:,:)
649 : : type(gs_hamiltonian_type),intent(in) :: hamk
650 : : type(mpi_type),intent(in) :: mpi_enreg
651 : : type(pawcprj_type),intent(inout) :: cwaveprj(:,:)
652 : :
653 : : !Local variables-------------------------------
654 : : !scalars
655 : : integer :: bandpp,ier,ikpt_this_proc,my_nspinor,ndatarecv,nproc_band,npw,nspinortot
656 : : integer :: spaceComm=0,tim_nonlop
657 : : logical :: do_transpose,transfer_cwavef
658 : : integer :: l_gpu_option
659 : : #ifdef HAVE_OPENMP_OFFLOAD
660 : : integer :: ipw,ind,iibandpp
661 : : #endif
662 : : !character(len=500) :: msg
663 : : !arrays
664 : 266542 : integer, allocatable :: index_wavef_band(:)
665 : 266542 : integer, allocatable :: rdisplsloc(:),recvcountsloc(:),sdisplsloc(:),sendcountsloc(:)
666 : 266542 : integer,contiguous, pointer :: rdispls(:),recvcounts(:),sdispls(:),sendcounts(:)
667 : 533084 : real(dp) :: lambda_nonlop(mpi_enreg%bandpp)
668 : : real(dp) :: tsec(2)
669 : :
670 : : #if defined HAVE_GPU && defined HAVE_YAKL
671 : : real(c_double), contiguous, pointer :: cwavef_alltoall2(:,:) => null()
672 : : real(c_double), contiguous, pointer :: gvnlc_alltoall2(:,:) => null()
673 : : real(c_double), contiguous, pointer :: gsc_alltoall2(:,:) => null()
674 : : integer(kind=C_SIZE_T) :: cwavef_alltoall2_size
675 : : integer(kind=C_SIZE_T) :: gvnlc_alltoall2_size
676 : : integer(kind=C_SIZE_T) :: gsc_alltoall2_size
677 : : #else
678 : 266542 : real(dp), pointer :: cwavef_alltoall2(:,:)
679 : 266542 : real(dp), allocatable :: gvnlc_alltoall2(:,:)
680 : 266542 : real(dp), allocatable :: gsc_alltoall2(:,:)
681 : : #endif
682 : :
683 : 266542 : real(dp), allocatable :: cwavef_alltoall1(:,:)
684 : 266542 : real(dp), allocatable :: gvnlc_alltoall1(:,:)
685 : 266542 : real(dp), allocatable :: gsc_alltoall1(:,:)
686 : 266542 : real(dp), allocatable :: enlout(:)
687 : :
688 : : #if defined HAVE_GPU && defined HAVE_YAKL
689 : : ! this buffer is necessary to avoid mixing "managed memory" buffer with "regular memory" buffer in MPI calls
690 : : ! Just to be clear:
691 : : ! - managed memory means memory allocated using ABI_MALLOC_MANAGED
692 : : ! - regular memory means memory allocated using either ABI_MALLOC or ABI_MALLOC_CUDA
693 : : !
694 : : ! here we chose to use a GPU buffer, would it be better to use a CPU buffer ? To be checked.
695 : : !real(dp), allocatable :: cwavef_mpi(:,:)
696 : : type(c_ptr) :: cwavef_mpi_c
697 : : real(c_double),pointer :: cwavef_mpi(:,:)
698 : : #endif
699 : :
700 : : ! *************************************************************************
701 : :
702 : : DBG_ENTER('COLL')
703 : :
704 : 266542 : call timab(570,1,tsec)
705 : :
706 : 266542 : do_transpose = .true.
707 : 266542 : bandpp = mpi_enreg%bandpp
708 [ + - ]: 266542 : if(present(already_transposed)) then
709 [ + + ]: 266542 : if(already_transposed) then
710 : 124920 : do_transpose = .false.
711 : 124920 : bandpp = blocksize
712 : : end if
713 : : end if
714 : :
715 : 266542 : l_gpu_option=ABI_GPU_DISABLED
716 [ + + ]: 266542 : if (present(gpu_option)) then
717 : 132454 : l_gpu_option = gpu_option
718 : : end if
719 : :
720 : 266542 : nproc_band = mpi_enreg%nproc_band
721 : :
722 : 266542 : spaceComm=mpi_enreg%comm_fft
723 [ + - ]: 266542 : if(mpi_enreg%paral_kgb==1) spaceComm=mpi_enreg%comm_band
724 : 266542 : my_nspinor=max(1,hamk%nspinor/mpi_enreg%nproc_spinor)
725 : 266542 : nspinortot=hamk%nspinor
726 : :
727 : : !Check sizes
728 [ + + ]: 266542 : npw=hamk%npw_k;if (.not.do_transpose) npw=hamk%npw_fft_k
729 [ + + - + ]: 799626 : if (size(cwavef)/=2*npw*my_nspinor*blocksize) then
730 : 0 : ABI_BUG('Incorrect size for cwavef!')
731 : : end if
732 [ + + + + ]: 266542 : if(choice/=0.and.signs==2) then
733 [ - + ]: 62460 : if (paw_opt/=3) then
734 [ # # # # ]: 0 : if (size(gvnlc)/=2*npw*my_nspinor*blocksize) then
735 : 0 : ABI_BUG('Incorrect size for gvnlc!')
736 : : end if
737 : : end if
738 [ + - ]: 62460 : if(paw_opt>=3) then
739 [ + + - + ]: 187380 : if (size(gsc)/=2*npw*my_nspinor*blocksize) then
740 : 0 : ABI_BUG('Incorrect size for gsc!')
741 : : end if
742 : : end if
743 : : end if
744 [ + + + + ]: 266542 : if(cpopt>=0.and. .not. present(vectproj)) then
745 [ + + - + ]: 231504 : if (size(cwaveprj)/=hamk%natom*my_nspinor*bandpp) then
746 : 0 : ABI_BUG('Incorrect size for cwaveprj!')
747 : : end if
748 : : end if
749 : :
750 : 266542 : transfer_cwavef=.false.
751 : : #ifdef HAVE_OPENMP_OFFLOAD
752 : : if(.not. xomp_target_is_present(c_loc(cwavef)) .and. hamk%gpu_option==ABI_GPU_OPENMP) then
753 : : transfer_cwavef=.true.
754 : : end if
755 : : !$OMP TARGET ENTER DATA MAP(to:cwavef) IF(transfer_cwavef)
756 : : #endif
757 : :
758 [ - + + - : 799626 : ABI_MALLOC(sendcountsloc,(nproc_band))
+ - + - -
+ - + ]
759 [ - + + - : 533084 : ABI_MALLOC(sdisplsloc ,(nproc_band))
- + - + ]
760 [ - + + - : 533084 : ABI_MALLOC(recvcountsloc,(nproc_band))
- + - + ]
761 [ - + + - : 533084 : ABI_MALLOC(rdisplsloc ,(nproc_band))
- + - + ]
762 : :
763 : 266542 : ikpt_this_proc=bandfft_kpt_get_ikpt()
764 : :
765 : 266542 : recvcounts => bandfft_kpt(ikpt_this_proc)%recvcounts(:)
766 : 266542 : sendcounts => bandfft_kpt(ikpt_this_proc)%sendcounts(:)
767 : 266542 : rdispls => bandfft_kpt(ikpt_this_proc)%rdispls (:)
768 : 266542 : sdispls => bandfft_kpt(ikpt_this_proc)%sdispls (:)
769 : 266542 : ndatarecv = bandfft_kpt(ikpt_this_proc)%ndatarecv
770 : :
771 [ + - ]: 266542 : if(hamk%gpu_option==ABI_GPU_KOKKOS) then
772 : : #if defined HAVE_GPU && defined HAVE_YAKL
773 : : ABI_MALLOC_MANAGED(cwavef_alltoall2, (/2,ndatarecv*my_nspinor*bandpp/))
774 : : cwavef_alltoall2_size = 2*ndatarecv*my_nspinor*bandpp*dp
775 : :
776 : : if (paw_opt >= 0 .and. paw_opt < 3) then
777 : : gsc_alltoall2 => null()
778 : : gsc_alltoall2_size = 0
779 : : else
780 : : ABI_MALLOC_MANAGED(gsc_alltoall2, (/2,ndatarecv*my_nspinor*(paw_opt/3)*bandpp/))
781 : : gsc_alltoall2_size = 2*ndatarecv*my_nspinor*(paw_opt/3)*bandpp*dp
782 : : endif
783 : :
784 : : ABI_MALLOC_MANAGED(gvnlc_alltoall2, (/2,ndatarecv*my_nspinor*bandpp/))
785 : : gvnlc_alltoall2_size = 2*ndatarecv*my_nspinor*bandpp*dp
786 : : #endif
787 : : else
788 [ + + ]: 266542 : if(present(cwavef_tr)) then
789 [ + + + - : 192 : if(any(shape(cwavef_tr)/=(/2,ndatarecv*my_nspinor*bandpp/))) then
- + ]
790 : 0 : ABI_BUG('cwavef_tr does not have the correct shape')
791 : : endif
792 : 64 : cwavef_alltoall2 => cwavef_tr
793 : : else
794 [ - + + - : 799434 : ABI_MALLOC(cwavef_alltoall2, (2,ndatarecv*my_nspinor*bandpp))
+ - + - -
+ - + ]
795 : : endif
796 [ + + + - : 799626 : ABI_MALLOC(gsc_alltoall2, (2,ndatarecv*my_nspinor*(paw_opt/3)*bandpp))
+ - + + -
+ - + -
+ ]
797 [ - + + - : 799626 : ABI_MALLOC(gvnlc_alltoall2, (2,ndatarecv*my_nspinor*bandpp))
+ - + - -
+ - + -
+ ]
798 : : end if
799 : :
800 : : #ifdef HAVE_OPENMP_OFFLOAD
801 : : !$OMP TARGET ENTER DATA MAP(alloc:cwavef_alltoall2) IF(hamk%gpu_option==ABI_GPU_OPENMP)
802 : : #endif
803 : :
804 [ + + + + : 266542 : if(do_transpose .and. (bandpp/=1 .or. (bandpp==1 .and. mpi_enreg%paral_spinor==0.and.nspinortot==2)))then
+ + + + ]
805 [ - + + - : 62286 : ABI_MALLOC(cwavef_alltoall1,(2,ndatarecv*my_nspinor*bandpp))
+ - + - -
+ - + ]
806 : : #ifdef HAVE_OPENMP_OFFLOAD
807 : : !$OMP TARGET ENTER DATA MAP(alloc:cwavef_alltoall1) IF(hamk%gpu_option==ABI_GPU_OPENMP)
808 : : #endif
809 [ - + ]: 20762 : if(signs==2)then
810 [ # # ]: 0 : if (paw_opt/=3) then
811 [ # # # # : 0 : ABI_MALLOC(gvnlc_alltoall1,(2,ndatarecv*my_nspinor*bandpp))
# # # # ]
812 : : end if
813 [ # # ]: 0 : if (paw_opt==3.or.paw_opt==4) then
814 [ # # # # : 0 : ABI_MALLOC(gsc_alltoall1,(2,ndatarecv*my_nspinor*bandpp))
# # # # ]
815 : : end if
816 : : end if
817 : : end if
818 : :
819 [ + + + - : 799626 : ABI_MALLOC(enlout,(nnlout*bandpp))
+ - + + -
+ - + -
+ ]
820 [ + + ]: 1350182 : enlout = zero
821 : :
822 [ + + ]: 927100 : recvcountsloc(:)=recvcounts(:)*2*my_nspinor*bandpp
823 [ + + ]: 927100 : rdisplsloc(:)=rdispls(:)*2*my_nspinor*bandpp
824 [ + + ]: 927100 : sendcountsloc(:)=sendcounts(:)*2*my_nspinor
825 [ + + ]: 927100 : sdisplsloc(:)=sdispls(:)*2*my_nspinor
826 : :
827 [ + + ]: 266542 : if(do_transpose) then
828 : 141622 : call timab(581,1,tsec)
829 [ + + + + : 141622 : if (bandpp/=1 .or. (bandpp==1 .and. mpi_enreg%paral_spinor==0.and.nspinortot==2)) then
+ + ]
830 [ + - ]: 20762 : if (l_gpu_option == ABI_GPU_KOKKOS) then
831 : : #if defined HAVE_GPU && defined HAVE_YAKL
832 : : ABI_MALLOC_CUDA(cwavef_mpi_c, INT(2, c_size_t) * npw * my_nspinor * blocksize * dp)
833 : : call c_f_pointer(cwavef_mpi_c, cwavef_mpi, (/2, npw * my_nspinor * blocksize/))
834 : :
835 : : ! use cwavef_mpi instead of cwavef (don't use managed memory in MPI calls)
836 : : call copy_gpu_to_gpu(cwavef_mpi_c, C_LOC(cwavef), INT(2, c_size_t) * npw * my_nspinor * blocksize * dp)
837 : :
838 : : call xmpi_alltoallv(cwavef_mpi,sendcountsloc,sdisplsloc,cwavef_alltoall1,&
839 : : & recvcountsloc,rdisplsloc,spaceComm,ier)
840 : :
841 : : ABI_FREE_CUDA(cwavef_mpi_c)
842 : : #endif
843 : : else
844 : : call xmpi_alltoallv(cwavef,sendcountsloc,sdisplsloc,cwavef_alltoall1,&
845 : 20762 : & recvcountsloc,rdisplsloc,spaceComm,ier,use_omp_map=(hamk%gpu_option==ABI_GPU_OPENMP))
846 : : end if
847 : :
848 : :
849 : : else
850 : : call xmpi_alltoallv(cwavef,sendcountsloc,sdisplsloc,cwavef_alltoall2,&
851 : 120860 : & recvcountsloc,rdisplsloc,spaceComm,ier,use_omp_map=(hamk%gpu_option==ABI_GPU_OPENMP))
852 : : end if
853 : 141622 : call timab(581,2,tsec)
854 : : else
855 : : ! Here, we cheat, and use DCOPY to bypass some compiler's overzealous bound-checking
856 : : ! (ndatarecv*my_nspinor*bandpp might be greater than the declared size of cwavef)
857 [ + - ]: 124920 : if (l_gpu_option == ABI_GPU_KOKKOS) then
858 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS)
859 : : call copy_gpu_to_gpu(C_LOC(cwavef_alltoall2), C_LOC(cwavef), INT(2, c_size_t) * ndatarecv * my_nspinor * bandpp * dp)
860 : : #endif
861 [ - + ]: 124920 : else if (hamk%gpu_option == ABI_GPU_OPENMP) then
862 [ # # # # : 0 : call gpu_copy(cwavef_alltoall2, cwavef, int(2,c_size_t)*ndatarecv*my_nspinor*bandpp)
# # # # #
# # # ]
863 : : else
864 [ - + - - : 124920 : call DCOPY(2*ndatarecv*my_nspinor*bandpp,cwavef,1,cwavef_alltoall2,1)
- - - - -
+ - - - -
- + ]
865 : : end if
866 : : end if
867 : :
868 : : #ifdef HAVE_OPENMP_OFFLOAD
869 : : !$OMP TARGET EXIT DATA MAP(delete:cwavef) IF(transfer_cwavef)
870 : : #endif
871 : :
872 : : !=====================================================================
873 [ + + ]: 266542 : if (bandpp==1) then
874 : :
875 : :
876 [ + + + + : 129188 : if (do_transpose .and. mpi_enreg%paral_spinor==0.and.nspinortot==2) then !Sort WF by spin
+ + ]
877 : 7368 : call prep_sort_wavef_spin(nproc_band,my_nspinor,ndatarecv,recvcounts,rdispls,index_wavef_band)
878 [ + - ]: 7368 : if(hamk%gpu_option==ABI_GPU_OPENMP) then
879 : : #ifdef HAVE_OPENMP_OFFLOAD
880 : : !$OMP TARGET TEAMS DISTRIBUTE &
881 : : !$OMP& PRIVATE(iibandpp) MAP(to:cwavef_alltoall1,cwavef_alltoall2,index_wavef_band)
882 : : do iibandpp=1,bandpp*my_nspinor
883 : : !$OMP PARALLEL DO PRIVATE(ind,ipw)
884 : : do ipw = 1 ,ndatarecv
885 : : ind=index_wavef_band(ipw + ndatarecv*(iibandpp-1))
886 : : cwavef_alltoall2(1,ipw + ndatarecv*(iibandpp-1)) = cwavef_alltoall1(1,ind)
887 : : cwavef_alltoall2(2,ipw + ndatarecv*(iibandpp-1)) = cwavef_alltoall1(2,ind)
888 : : end do
889 : : end do
890 : : #endif
891 : : else
892 [ + + + + ]: 6364632 : cwavef_alltoall2(:, :) = cwavef_alltoall1(:,index_wavef_band)
893 : : end if
894 : : end if
895 : :
896 [ + + ]: 129188 : if (paw_opt==2) then
897 : 71168 : lambda_nonlop(1)=lambdablock(mpi_enreg%me_band+1)
898 : : end if
899 : : call nonlop(choice,cpopt,cwaveprj,enlout,hamk,idir,lambda_nonlop,mpi_enreg,1,nnlout,paw_opt,&
900 [ + + + + ]: 257416 : & signs,gsc_alltoall2,tim_nonlop,cwavef_alltoall2,gvnlc_alltoall2,vectproj=vectproj,select_k=select_k)
901 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_YAKL)
902 : : !call gpu_device_synchronize()
903 : : !call gpu_data_prefetch_async_f(C_LOC(cwavef_alltoall2), cwavef_alltoall2_size, CPU_DEVICE_ID)
904 : : !call gpu_data_prefetch_async_f(C_LOC(gvnlc_alltoall2), gvnlc_alltoall2_size, CPU_DEVICE_ID)
905 : : !call gpu_data_prefetch_async_f(C_LOC(gsc_alltoall2), gsc_alltoall2_size, CPU_DEVICE_ID)
906 : : #endif
907 : :
908 [ + + + + : 129188 : if (do_transpose .and. mpi_enreg%paral_spinor == 0 .and. nspinortot==2.and.signs==2) then
+ + - + ]
909 [ # # # # : 0 : if (paw_opt/=3) gvnlc_alltoall1(:,index_wavef_band)=gvnlc_alltoall2(:,:)
# # ]
910 [ # # # # : 0 : if (paw_opt==3.or.paw_opt==4) gsc_alltoall1(:,index_wavef_band)=gsc_alltoall2(:,:)
# # ]
911 : : end if
912 : :
913 : : else ! bandpp/=1
914 : :
915 : : ! -------------------------------------------------------------
916 : : ! Computation of the index used to sort the waves functions below bandpp
917 : : ! -------------------------------------------------------------
918 [ + + ]: 137354 : if(do_transpose) then
919 : : call prep_index_wavef_bandpp(nproc_band,bandpp,&
920 : 13394 : & my_nspinor,ndatarecv,recvcounts,rdispls,index_wavef_band)
921 : :
922 : : ! -------------------------------------------------------
923 : : ! Sorting of the waves functions below bandpp
924 : : ! -------------------------------------------------------
925 [ + - ]: 13394 : if(hamk%gpu_option==ABI_GPU_OPENMP) then
926 : : #ifdef HAVE_OPENMP_OFFLOAD
927 : : !$OMP TARGET TEAMS DISTRIBUTE &
928 : : !$OMP& PRIVATE(iibandpp) MAP(to:cwavef_alltoall1,cwavef_alltoall2,index_wavef_band)
929 : : do iibandpp=1,bandpp*my_nspinor
930 : : !$OMP PARALLEL DO PRIVATE(ind,ipw)
931 : : do ipw = 1 ,ndatarecv
932 : : ind=index_wavef_band(ipw + ndatarecv*(iibandpp-1))
933 : : cwavef_alltoall2(1,ipw + ndatarecv*(iibandpp-1)) = cwavef_alltoall1(1,ind)
934 : : cwavef_alltoall2(2,ipw + ndatarecv*(iibandpp-1)) = cwavef_alltoall1(2,ind)
935 : : end do
936 : : end do
937 : : #endif
938 : : else
939 [ + + + + ]: 133129886 : cwavef_alltoall2(:,:) = cwavef_alltoall1(:,index_wavef_band)
940 : : end if
941 : : end if
942 : :
943 : : ! -------------------------------------------------------
944 : : ! Call nonlop
945 : : ! -------------------------------------------------------
946 [ + + ]: 137354 : if(paw_opt == 2) then
947 [ + + ]: 46604 : lambda_nonlop(1:bandpp) = lambdablock((mpi_enreg%me_band*bandpp)+1:((mpi_enreg%me_band+1)*bandpp))
948 : : end if
949 : : call nonlop(choice,cpopt,cwaveprj,enlout,hamk,idir,lambda_nonlop,mpi_enreg,bandpp,nnlout,paw_opt,&
950 [ + + + + ]: 150748 : & signs,gsc_alltoall2,tim_nonlop,cwavef_alltoall2,gvnlc_alltoall2,vectproj=vectproj,select_k=select_k)
951 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_YAKL)
952 : : if(hamk%gpu_option==ABI_GPU_KOKKOS) call gpu_device_synchronize()
953 : : !call gpu_data_prefetch_async_f(C_LOC(cwavef_alltoall2), cwavef_alltoall2_size, CPU_DEVICE_ID)
954 : : !call gpu_data_prefetch_async_f(C_LOC(gvnlc_alltoall2), gvnlc_alltoall2_size, CPU_DEVICE_ID)
955 : : !if (associated(gsc_alltoall2)) then
956 : : ! call gpu_data_prefetch_async_f(C_LOC(gsc_alltoall2), gsc_alltoall2_size, CPU_DEVICE_ID)
957 : : !end if
958 : : #endif
959 : :
960 : : ! -----------------------------------------------------
961 : : ! Sorting of waves functions below the processors
962 : : ! -----------------------------------------------------
963 [ + + - + ]: 137354 : if(do_transpose.and.signs==2) then
964 [ # # # # : 0 : if (paw_opt/=3) gvnlc_alltoall1(:,index_wavef_band)=gvnlc_alltoall2(:,:)
# # ]
965 [ # # # # : 0 : if (paw_opt==3.or.paw_opt==4) gsc_alltoall1(:,index_wavef_band)=gsc_alltoall2(:,:)
# # ]
966 : : end if
967 : :
968 : : end if
969 : :
970 : : !=====================================================================
971 : : ! -------------------------------------------------------
972 : : ! Deallocation
973 : : ! -------------------------------------------------------
974 [ + + ]: 266542 : if (allocated(index_wavef_band)) then
975 : 20762 : ABI_FREE(index_wavef_band)
976 : : end if
977 : :
978 : : !Transpose the gsc_alltoall or gvlnc_alltoall tabs
979 : : !according to the paw_opt and signs values
980 [ + + ]: 266542 : if(do_transpose) then
981 [ - + ]: 141622 : if (signs==2) then
982 : 0 : call timab(581,1,tsec)
983 [ # # # # : 0 : if(bandpp/=1 .or. (bandpp==1 .and. mpi_enreg%paral_spinor==0.and.nspinortot==2))then
# # ]
984 [ # # ]: 0 : if (paw_opt/=3) then
985 : : call xmpi_alltoallv(gvnlc_alltoall1,recvcountsloc,rdisplsloc,gvnlc,&
986 : 0 : & sendcountsloc,sdisplsloc,spaceComm,ier)
987 : : end if
988 [ # # ]: 0 : if (paw_opt==3.or.paw_opt==4) then
989 : : call xmpi_alltoallv(gsc_alltoall1,recvcountsloc,rdisplsloc,gsc,&
990 : 0 : & sendcountsloc,sdisplsloc,spaceComm,ier)
991 : : end if
992 : : else
993 [ # # ]: 0 : if (paw_opt/=3) then
994 : : call xmpi_alltoallv(gvnlc_alltoall2,recvcountsloc,rdisplsloc,gvnlc,&
995 : 0 : & sendcountsloc,sdisplsloc,spaceComm,ier)
996 : : end if
997 [ # # ]: 0 : if (paw_opt==3.or.paw_opt==4) then
998 : : call xmpi_alltoallv(gsc_alltoall2,recvcountsloc,rdisplsloc,gsc,&
999 : 0 : & sendcountsloc,sdisplsloc,spaceComm,ier)
1000 : : end if
1001 : : end if
1002 : 0 : call timab(581,2,tsec)
1003 : : end if
1004 : : else
1005 : : ! TODO check other usages, maybe
1006 [ + - ]: 124920 : if (l_gpu_option == ABI_GPU_KOKKOS) then
1007 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS)
1008 : : call copy_gpu_to_gpu(C_LOC(gsc), C_LOC(gsc_alltoall2), INT(2, c_size_t) * ndatarecv * my_nspinor * bandpp * dp)
1009 : : #endif
1010 : : else
1011 [ - + - - : 124920 : call DCOPY(2*ndatarecv*my_nspinor*bandpp, gsc_alltoall2, 1, gsc, 1)
- - - - -
+ - - -
- ]
1012 : : end if
1013 : : end if
1014 : :
1015 [ + + ]: 266542 : if (nnlout>0) then
1016 : 107492 : call xmpi_allgather(enlout,nnlout*bandpp,enlout_block,spaceComm,ier)
1017 : : end if
1018 [ - + ]: 266542 : ABI_FREE(enlout)
1019 [ - + ]: 266542 : ABI_FREE(sendcountsloc)
1020 [ - + ]: 266542 : ABI_FREE(sdisplsloc)
1021 [ - + ]: 266542 : ABI_FREE(recvcountsloc)
1022 [ - + ]: 266542 : ABI_FREE(rdisplsloc)
1023 : :
1024 : : #ifdef HAVE_OPENMP_OFFLOAD
1025 : : if(present(cwavef_tr)) then
1026 : : !$OMP TARGET UPDATE FROM(cwavef_alltoall2) IF(hamk%gpu_option==ABI_GPU_OPENMP)
1027 : : end if
1028 : : !$OMP TARGET EXIT DATA MAP(delete:cwavef_alltoall2) IF(hamk%gpu_option==ABI_GPU_OPENMP)
1029 : : #endif
1030 : :
1031 [ + - ]: 266542 : if(hamk%gpu_option==ABI_GPU_KOKKOS) then
1032 : : #if defined HAVE_GPU && defined HAVE_YAKL
1033 : : ABI_FREE_MANAGED(cwavef_alltoall2)
1034 : : ABI_FREE_MANAGED(gvnlc_alltoall2)
1035 : : if (paw_opt >= 3) then
1036 : : ABI_FREE_MANAGED(gsc_alltoall2)
1037 : : end if
1038 : : #endif
1039 : : else
1040 [ + + ]: 266542 : if(present(cwavef_tr)) then
1041 : 64 : nullify(cwavef_alltoall2)
1042 : : else
1043 [ - + ]: 266478 : ABI_FREE(cwavef_alltoall2)
1044 : : endif
1045 [ - + ]: 266542 : ABI_FREE(gvnlc_alltoall2)
1046 [ - + ]: 266542 : ABI_FREE(gsc_alltoall2)
1047 : : end if
1048 : :
1049 [ + + + + : 266542 : if(do_transpose .and. (bandpp/=1 .or. (bandpp==1 .and. mpi_enreg%paral_spinor==0.and.nspinortot==2)))then
+ + + + ]
1050 : : #ifdef HAVE_OPENMP_OFFLOAD
1051 : : !$OMP TARGET EXIT DATA MAP(delete:cwavef_alltoall1) IF(hamk%gpu_option==ABI_GPU_OPENMP)
1052 : : #endif
1053 [ - + ]: 20762 : ABI_FREE(cwavef_alltoall1)
1054 [ - + ]: 20762 : if(signs==2)then
1055 [ # # ]: 0 : if (paw_opt/=3) then
1056 [ # # ]: 0 : ABI_FREE(gvnlc_alltoall1)
1057 : : end if
1058 [ # # ]: 0 : if (paw_opt==3.or.paw_opt==4) then
1059 [ # # ]: 0 : ABI_FREE(gsc_alltoall1)
1060 : : end if
1061 : : end if
1062 : : end if
1063 : :
1064 : 266542 : call timab(570,2,tsec)
1065 : :
1066 : : DBG_EXIT('COLL')
1067 : :
1068 [ - + - + : 533084 : end subroutine prep_nonlop
- + - + -
+ - + -
+ ]
1069 : : !!***
1070 : :
1071 : : !!****f* ABINIT/prep_fourwf
1072 : : !! NAME
1073 : : !! prep_fourwf
1074 : : !!
1075 : : !! FUNCTION
1076 : : !! this routine prepares the data to the call of fourwf.
1077 : : !!
1078 : : !! INPUTS
1079 : : !! blocksize= size of block for FFT
1080 : : !! cwavef(2,npw*ndat)=planewave coefficients of wavefunction (one spinorial component?).
1081 : : !! dtfil <type(datafiles_type)>=variables related to files
1082 : : !! kg_k(3,npw_k)=reduced planewave coordinates.
1083 : : !! lmnmax=if useylm=1, max number of (l,m,n) comp. over all type of psps
1084 : : !! =if useylm=0, max number of (l,n) comp. over all type of psps
1085 : : !! mgfft=maximum size of 1d ffts
1086 : : !! mpi_enreg=information about mpi parallelization
1087 : : !! mpsang= 1+maximum angular momentum for nonlocal pseudopotentials
1088 : : !! mpssoang= 1+maximum (spin*angular momentum) for nonlocal pseudopotentials
1089 : : !! natom=number of atoms in cell.
1090 : : !! nband_k=number of bands at this k point for that spin polarization
1091 : : !! ndat=number of FFT to do in //
1092 : : !! ngfft(18)= contain all needed information about 3D FFT
1093 : : !! npw_k=number of plane waves at this k point
1094 : : !! nspinor=number of spinorial components of the wavefunctions
1095 : : !! ntypat=number of types of atoms in unit cell.
1096 : : !! n4,n5,n6 used for dimensionning of vlocal
1097 : : !! option_fourwf=option for fourwf (see fourwf.F90)
1098 : : !! prtvol=control print volume and debugging output
1099 : : !! ucvol=unit cell volume
1100 : : !! nfft_blocks=number of blocks fourwf is split into.
1101 : : !! [bandfft_kpt_tab]= (optional) if present, contains tabs used to implement
1102 : : !! the "band-fft" parallelism
1103 : : !! if not present, the bandfft_kpt global variable is used
1104 : : !! [gpu_option] = GPU implementation to use, i.e. cuda, openMP, ... (0=not using GPU)
1105 : : !!
1106 : : !! OUTPUT
1107 : : !! gwavef=(2,npw*ndat)=matrix elements <G|H|C>.
1108 : : !!
1109 : : !! SIDE EFFECTS
1110 : : !!
1111 : : !! SOURCE
1112 : :
1113 [ - + ]: 185482 : subroutine prep_fourwf(rhoaug,blocksize,cwavef,wfraug,iblock,istwf_k,mgfft,&
1114 : 185482 : & mpi_enreg,nband_k,ndat,ngfft,npw_k,n4,n5,n6,occ_k,option_fourwf,ucvol,wtk,&
1115 : : & nfft_blocks,&
1116 : : & bandfft_kpt_tab,gpu_option) ! Optional arguments
1117 : :
1118 : : !Arguments ------------------------------------
1119 : : !scalars
1120 : : integer,intent(in) :: blocksize,iblock,istwf_k,mgfft,n4,n5,n6,nband_k,ndat,npw_k,nfft_blocks
1121 : : integer,intent(in) :: option_fourwf
1122 : : integer,intent(in),optional :: gpu_option
1123 : : real(dp),intent(in) :: ucvol,wtk
1124 : : type(bandfft_kpt_type),optional,target,intent(in) :: bandfft_kpt_tab
1125 : : type(mpi_type),intent(in) :: mpi_enreg
1126 : : !arrays
1127 : : integer,intent(in) :: ngfft(18)
1128 : : real(dp),intent(in) :: occ_k(nband_k)
1129 : : real(dp),intent(out) :: rhoaug(n4,n5,n6)
1130 : : real(dp),intent(in), target :: cwavef(2,npw_k*blocksize)
1131 : : real(dp),target,intent(inout) :: wfraug(:,:,:,:) !(2,n4,n5,n6*{ndat,ndat/nfft_blocks+mod(ndat,nfft_blocks))
1132 : :
1133 : : !Local variables-------------------------------
1134 : : !scalars
1135 : : integer :: bandpp,bandpp_sym,ier,iibandpp,ikpt_this_proc,ind_occ,ind_occ1,ind_occ2,ipw
1136 : : integer :: istwf_k_,jjbandpp,me_fft,nd3,nproc_band,nproc_fft,npw_fft
1137 : : integer :: nfft_blocks_sym,nband_fftblock
1138 : : integer :: spaceComm=0,tim_fourwf,gpu_option_
1139 : : integer,pointer :: idatarecv0,ndatarecv,ndatarecv_tot,ndatasend_sym
1140 : : logical :: flag_inv_sym,have_to_reequilibrate,transfer_cwavef
1141 : : real(dp) :: weight,weight1,weight2
1142 : : type(bandfft_kpt_type),pointer :: bandfft_kpt_ptr
1143 : : !arrays
1144 : 185482 : integer,contiguous, pointer :: indices_pw_fft(:),kg_k_fft(:,:),kg_k_gather(:,:),kg_k_gather_sym(:,:)
1145 : 185482 : integer,contiguous, pointer :: rdispls(:),rdispls_sym(:)
1146 : 185482 : integer,contiguous, pointer :: recvcounts(:),recvcount_fft(:),recvcounts_sym(:),recvcounts_sym_tot(:)
1147 : 185482 : integer,contiguous, pointer :: recvdisp_fft(:),sdispls(:),sdispls_sym(:)
1148 : 185482 : integer,contiguous, pointer :: sendcounts(:),sendcount_fft(:),sendcounts_sym(:),sendcounts_sym_all(:)
1149 : 185482 : integer,contiguous, pointer :: senddisp_fft(:),tab_proc(:)
1150 : 185482 : integer,allocatable :: rdisplsloc(:)
1151 : 185482 : integer,allocatable :: recvcountsloc(:),sdisplsloc(:)
1152 : 185482 : integer,allocatable :: sendcountsloc(:)
1153 : 185482 : integer,allocatable :: index_wavef_band(:),index_wavef_send(:)
1154 : 185482 : integer,pointer :: gbound_(:,:)
1155 : : real(dp) :: dummy(2,1),tsec(2)
1156 : 185482 : real(dp),allocatable :: buff_wf(:,:)
1157 : :
1158 : : #if defined HAVE_GPU && defined HAVE_YAKL
1159 : : real(c_double), contiguous, pointer :: cwavef_alltoall1(:,:) => null()
1160 : : #else
1161 : 185482 : real(dp),allocatable :: cwavef_alltoall1(:,:)
1162 : : #endif
1163 : 185482 : real(dp),allocatable :: cwavef_alltoall2(:,:)
1164 : 185482 : real(dp),allocatable :: cwavef_fft(:,:), cwavef_fft_tr(:,:)
1165 : 185482 : real(dp),allocatable :: weight_t(:),weight1_t(:),weight2_t(:)
1166 : 185482 : real(dp),pointer :: ewavef_alltoall_sym(:,:),ewavef_alltoall_sym_work(:,:),wfraug_ptr(:,:,:,:)
1167 : :
1168 : : #if defined HAVE_GPU && defined HAVE_YAKL
1169 : : ! this buffer is necessary to avoid mixing "managed memory" buffer with "regular memory" buffer in MPI calls
1170 : : ! Just to be clear:
1171 : : ! - managed memory means memory allocated using ABI_MALLOC_MANAGED
1172 : : ! - regular memory means memory allocated using either ABI_MALLOC or ABI_MALLOC_CUDA
1173 : : !
1174 : : ! here we chose to use a CPU buffer, would it be better to use a GPU buffer ? To be checked.
1175 : : real(dp), allocatable :: cwavef_mpi(:,:)
1176 : : !type(c_ptr) :: cwavef_mpi_c
1177 : : !real(c_double),pointer :: cwavef_mpi(:,:)
1178 : : #endif
1179 : :
1180 : : ! *************************************************************************
1181 : :
1182 : 0 : ABI_CHECK((option_fourwf/=3),'Option=3 (FFT r->g) not implemented')
1183 [ - + ]: 185482 : ABI_CHECK((nfft_blocks>0),'nfft_blocks is null')
1184 [ - + ]: 185482 : ABI_CHECK((mpi_enreg%bandpp==ndat),'BUG: bandpp/=ndat')
1185 : :
1186 : 185482 : spaceComm=mpi_enreg%comm_band
1187 : 185482 : nproc_band = mpi_enreg%nproc_band
1188 : 185482 : nproc_fft = mpi_enreg%nproc_fft
1189 : 185482 : bandpp = mpi_enreg%bandpp
1190 : 185482 : me_fft = mpi_enreg%me_fft
1191 : :
1192 [ + - ]: 185482 : gpu_option_=ABI_GPU_DISABLED;if (present(gpu_option)) gpu_option_=gpu_option
1193 : :
1194 [ + - ]: 185482 : if (present(bandfft_kpt_tab)) then
1195 : : bandfft_kpt_ptr => bandfft_kpt_tab
1196 : : else
1197 : 185482 : ikpt_this_proc=bandfft_kpt_get_ikpt()
1198 : 185482 : bandfft_kpt_ptr => bandfft_kpt(ikpt_this_proc)
1199 : : end if
1200 : :
1201 : 185482 : have_to_reequilibrate = bandfft_kpt_ptr%have_to_reequilibrate
1202 : 185482 : istwf_k_=istwf_k
1203 [ + + + + : 185482 : flag_inv_sym = (istwf_k_==2 .and. any(ngfft(7) == [401,402,312,512]))
+ - ]
1204 [ - + ]: 185482 : if (option_fourwf==0) flag_inv_sym=((flag_inv_sym).and.(gpu_option_==ABI_GPU_DISABLED))
1205 : :
1206 [ + + ]: 185482 : if (flag_inv_sym) then
1207 : 6622 : istwf_k_ = 1
1208 [ + + ]: 6622 : if (modulo(bandpp,2)==0) then
1209 : 2494 : bandpp_sym = bandpp/2
1210 [ + - ]: 2494 : nfft_blocks_sym = nfft_blocks/2; if(modulo(nfft_blocks,2)/=0) nfft_blocks_sym=nfft_blocks_sym+1
1211 : : else
1212 : 4128 : bandpp_sym = bandpp
1213 : 4128 : nfft_blocks_sym = nfft_blocks
1214 : : end if
1215 : : end if
1216 : :
1217 [ - + ]: 185482 : if(have_to_reequilibrate .and. gpu_option_==ABI_GPU_OPENMP) then
1218 : 0 : ABI_BUG("Reequilibrating FFT isn't supported with OpenMP GPU yet !")
1219 : : end if
1220 : :
1221 : : !====================================================================================
1222 [ - + + - : 556446 : ABI_MALLOC(sendcountsloc,(nproc_band))
+ - + - -
+ - + ]
1223 [ - + + - : 370964 : ABI_MALLOC(sdisplsloc ,(nproc_band))
- + - + ]
1224 [ - + + - : 370964 : ABI_MALLOC(recvcountsloc,(nproc_band))
- + - + ]
1225 [ - + + - : 370964 : ABI_MALLOC(rdisplsloc ,(nproc_band))
- + - + ]
1226 : :
1227 : 185482 : recvcounts =>bandfft_kpt_ptr%recvcounts(:)
1228 : 185482 : sendcounts =>bandfft_kpt_ptr%sendcounts(:)
1229 : 185482 : rdispls =>bandfft_kpt_ptr%rdispls (:)
1230 : 185482 : sdispls =>bandfft_kpt_ptr%sdispls (:)
1231 : 185482 : ndatarecv =>bandfft_kpt_ptr%ndatarecv
1232 : :
1233 : 185482 : kg_k_gather =>bandfft_kpt_ptr%kg_k_gather(:,:)
1234 : 185482 : gbound_ =>bandfft_kpt_ptr%gbound(:,:)
1235 : :
1236 [ + + ]: 185482 : if (flag_inv_sym ) then
1237 : 6622 : idatarecv0 =>bandfft_kpt_ptr%idatarecv0
1238 : 6622 : ndatarecv_tot =>bandfft_kpt_ptr%ndatarecv_tot
1239 : 6622 : ndatasend_sym =>bandfft_kpt_ptr%ndatasend_sym
1240 : 6622 : kg_k_gather_sym =>bandfft_kpt_ptr%kg_k_gather_sym(:,:)
1241 : 6622 : rdispls_sym =>bandfft_kpt_ptr%rdispls_sym(:)
1242 : 6622 : recvcounts_sym =>bandfft_kpt_ptr%recvcounts_sym(:)
1243 : 6622 : recvcounts_sym_tot =>bandfft_kpt_ptr%recvcounts_sym_tot(:)
1244 : 6622 : sdispls_sym =>bandfft_kpt_ptr%sdispls_sym(:)
1245 : 6622 : sendcounts_sym =>bandfft_kpt_ptr%sendcounts_sym(:)
1246 : 6622 : sendcounts_sym_all =>bandfft_kpt_ptr%sendcounts_sym_all(:)
1247 : 6622 : tab_proc =>bandfft_kpt_ptr%tab_proc(:)
1248 : : end if
1249 : :
1250 [ - + + - : 556446 : ABI_MALLOC(cwavef_alltoall2,(2,ndatarecv*bandpp))
+ - + - -
+ - + ]
1251 [ + + + + ]: 185482 : if ( ((.not.flag_inv_sym) .and. (bandpp>1) ) .or. flag_inv_sym )then
1252 [ + - ]: 37386 : if(gpu_option_==ABI_GPU_KOKKOS) then
1253 : : #if defined HAVE_GPU && defined HAVE_YAKL
1254 : : ABI_MALLOC_MANAGED(cwavef_alltoall1,(/2,ndatarecv*bandpp/))
1255 : : #endif
1256 : : else
1257 [ - + + - : 222868 : ABI_MALLOC(cwavef_alltoall1,(2,ndatarecv*bandpp))
- + - + ]
1258 : : end if
1259 : : end if
1260 : :
1261 [ + + ]: 605876 : recvcountsloc(:)=recvcounts(:)*2*bandpp
1262 [ + + ]: 605876 : rdisplsloc(:)=rdispls(:)*2*bandpp
1263 [ + + ]: 605876 : sendcountsloc(:)=sendcounts(:)*2
1264 [ + + ]: 605876 : sdisplsloc(:)=sdispls(:)*2
1265 : :
1266 : 185482 : transfer_cwavef=.false.
1267 : 185482 : if(.not. xomp_target_is_present(c_loc(cwavef)) .and. gpu_option_==ABI_GPU_OPENMP) then
1268 : 185482 : transfer_cwavef=.true.
1269 : : end if
1270 : :
1271 : : #ifdef HAVE_OPENMP_OFFLOAD
1272 : : !$OMP TARGET ENTER DATA MAP(to:cwavef) IF(transfer_cwavef)
1273 : : !$OMP TARGET ENTER DATA MAP(alloc:cwavef_alltoall2) IF(gpu_option_==ABI_GPU_OPENMP)
1274 : : if ( ((.not.flag_inv_sym) .and. (bandpp>1) ) .or. flag_inv_sym )then
1275 : : !$OMP TARGET ENTER DATA MAP(alloc:cwavef_alltoall1) IF(gpu_option_==ABI_GPU_OPENMP)
1276 : : end if
1277 : : #endif
1278 : :
1279 : 185482 : call timab(547,1,tsec)
1280 [ + - ]: 185482 : if(gpu_option_==ABI_GPU_KOKKOS) then
1281 : : #if defined HAVE_GPU && defined HAVE_YAKL
1282 : : ABI_MALLOC(cwavef_mpi,(2,npw_k*blocksize))
1283 : :
1284 : : call gpu_data_prefetch_async(C_LOC(cwavef), INT(2, c_size_t)*npw_k*blocksize, CPU_DEVICE_ID)
1285 : : call gpu_device_synchronize()
1286 : :
1287 : : cwavef_mpi(:,:) = cwavef(:,:)
1288 : :
1289 : : call xmpi_alltoallv(cwavef_mpi,sendcountsloc,sdisplsloc,cwavef_alltoall2,&
1290 : : & recvcountsloc,rdisplsloc,spaceComm,ier)
1291 : : ABI_FREE(cwavef_mpi)
1292 : : #endif
1293 : : else
1294 : : call xmpi_alltoallv(cwavef,sendcountsloc,sdisplsloc,cwavef_alltoall2,&
1295 : 185482 : & recvcountsloc,rdisplsloc,spaceComm,ier,use_omp_map=(gpu_option_==ABI_GPU_OPENMP))
1296 : : end if
1297 : :
1298 : : #ifdef HAVE_OPENMP_OFFLOAD
1299 : : !$OMP TARGET EXIT DATA MAP(delete:cwavef) IF(transfer_cwavef)
1300 : : #endif
1301 : 185482 : call timab(547,2,tsec)
1302 : :
1303 : 185482 : tim_fourwf=16
1304 : :
1305 : : !Eventually adjust load balancing for FFT (by changing FFT distrib)
1306 [ - + ]: 185482 : if(have_to_reequilibrate) then
1307 : 0 : npw_fft = bandfft_kpt_ptr%npw_fft
1308 : 0 : sendcount_fft => bandfft_kpt_ptr%sendcount_fft(:)
1309 : 0 : recvcount_fft => bandfft_kpt_ptr%recvcount_fft(:)
1310 : 0 : senddisp_fft => bandfft_kpt_ptr%senddisp_fft(:)
1311 : 0 : recvdisp_fft => bandfft_kpt_ptr%recvdisp_fft(:)
1312 : 0 : indices_pw_fft => bandfft_kpt_ptr%indices_pw_fft(:)
1313 : 0 : kg_k_fft => bandfft_kpt_ptr%kg_k_fft(:,:)
1314 [ # # # # : 0 : ABI_MALLOC( buff_wf, (2,ndatarecv*bandpp) ) ! for sorting cgwavef
# # # # #
# # # ]
1315 [ # # # # : 0 : ABI_MALLOC( cwavef_fft, (2,npw_fft*bandpp) )
# # # # #
# # # ]
1316 [ # # ]: 0 : if(bandpp>1) then
1317 [ # # # # : 0 : ABI_MALLOC( cwavef_fft_tr, (2,npw_fft*bandpp))
# # # # ]
1318 : : end if
1319 : : end if
1320 : :
1321 [ - + - - : 185482 : if (option_fourwf==0) wfraug(:,:,:,:)=zero
- - - - -
- ]
1322 : :
1323 : : !====================================================================
1324 [ + + ]: 185482 : if ((.not.(flag_inv_sym)) .and. (bandpp==1)) then
1325 : :
1326 : : ! Compute the index of the band
1327 : 148096 : ind_occ = (iblock-1)*blocksize + mpi_enreg%me_band + 1
1328 : :
1329 [ + + - + ]: 148096 : if(abs(occ_k(ind_occ))>=tol8.or.option_fourwf==0) then
1330 : :
1331 : : ! Compute the weight of the band
1332 : 76198 : weight=occ_k(ind_occ)*wtk/ucvol
1333 : :
1334 [ - + ]: 76198 : if(have_to_reequilibrate) then
1335 : : ! filling of sorted send buffers before exchange
1336 [ # # ]: 0 : do ipw = 1 ,ndatarecv
1337 [ # # ]: 0 : buff_wf(1:2, indices_pw_fft(ipw) ) = cwavef_alltoall2(1:2,ipw)
1338 : : end do
1339 : : call xmpi_alltoallv(buff_wf,2*sendcount_fft,2*senddisp_fft, &
1340 [ # # # # : 0 : & cwavef_fft,2*recvcount_fft, 2*recvdisp_fft, mpi_enreg%comm_fft,ier)
# # # # #
# # # # #
# # ]
1341 : : call fourwf(1,rhoaug,cwavef_fft,dummy,wfraug,gbound_,gbound_,&
1342 : : & istwf_k_,kg_k_fft,kg_k_fft,mgfft,mpi_enreg,1,&
1343 : : & ngfft,npw_fft,1,n4,n5,n6,option_fourwf,tim_fourwf,weight,weight,&
1344 [ # # # # : 0 : & gpu_option=gpu_option_)
# # # # #
# # # # #
# # # # #
# # # # #
# # ]
1345 : : else
1346 : : call fourwf(1,rhoaug,cwavef_alltoall2,dummy,wfraug,gbound_,gbound_,&
1347 : : & istwf_k_,kg_k_gather,kg_k_gather,mgfft,mpi_enreg,1,&
1348 : : & ngfft,ndatarecv,1,n4,n5,n6,option_fourwf,tim_fourwf,weight,weight,&
1349 [ - + - - : 76198 : & gpu_option=gpu_option_)
- - - - -
- - - - +
- - - - -
- - - - +
- + ]
1350 : : end if
1351 [ - + - - ]: 76198 : if (option_fourwf==0.and.nproc_fft>1) then
1352 [ # # ]: 0 : if (me_fft>0) then
1353 : 0 : nd3=(ngfft(3)-1)/nproc_fft+1
1354 [ # # # # : 0 : wfraug(:,:,:,me_fft*nd3+1:me_fft*nd3+nd3)=wfraug(:,:,:,1:nd3)
# # # # #
# # # # #
# # # # ]
1355 [ # # # # : 0 : wfraug(:,:,:,1:nd3)=zero
# # # # ]
1356 : : end if
1357 : 0 : call xmpi_sum(wfraug,mpi_enreg%comm_fft,ier)
1358 : : end if
1359 : : end if
1360 : :
1361 : : !====================================================================
1362 [ + + ]: 37386 : else if ((.not.(flag_inv_sym)) .and. (bandpp>1) ) then
1363 : :
1364 : : ! -------------------------------------------------------------
1365 : : ! Computation of the index to class the waves functions below bandpp
1366 : : ! -------------------------------------------------------------
1367 : : call prep_index_wavef_bandpp(nproc_band,bandpp,&
1368 : : & 1,ndatarecv,&
1369 : : & recvcounts,rdispls,&
1370 : 30764 : & index_wavef_band)
1371 : :
1372 : : ! -------------------------------------------------------
1373 : : ! Sorting of the wave functions below bandpp
1374 : : ! -------------------------------------------------------
1375 [ + - ]: 30764 : if(gpu_option_==ABI_GPU_OPENMP) then
1376 : : !FIXME Having this OpenMP kernel here cause a seg fault, so I put it in a separate routine
1377 : : ! I can't figure what cause the issue...
1378 : : call omp_copy_sort_wf(bandpp,ndatarecv,cwavef_alltoall1,cwavef_alltoall2,index_wavef_band,flag_inv_sym)
1379 : : else
1380 [ + + + + ]: 105444488 : cwavef_alltoall1(:,:) = cwavef_alltoall2(:,index_wavef_band)
1381 : : end if
1382 : :
1383 : 30764 : if(have_to_reequilibrate) then
1384 : : ! filling of sorted send buffers before exchange
1385 [ # # ]: 0 : do iibandpp=1,bandpp
1386 [ # # ]: 0 : do ipw = 1 ,ndatarecv
1387 : : buff_wf(1:2, iibandpp + bandpp*(indices_pw_fft(ipw)-1)) = &
1388 [ # # ]: 0 : & cwavef_alltoall1(1:2,ipw + ndatarecv*(iibandpp-1))
1389 : : end do
1390 : : end do
1391 : : call xmpi_alltoallv(buff_wf,2*bandpp*sendcount_fft,2*bandpp*senddisp_fft, &
1392 [ # # # # : 0 : & cwavef_fft_tr,2*bandpp*recvcount_fft, 2*bandpp*recvdisp_fft, mpi_enreg%comm_fft,ier)
# # # # #
# # # # #
# # ]
1393 [ # # ]: 0 : do iibandpp=1,bandpp
1394 [ # # ]: 0 : do ipw = 1 ,npw_fft
1395 [ # # ]: 0 : cwavef_fft(1:2, ipw + npw_fft*(iibandpp -1)) = cwavef_fft_tr(1:2, iibandpp + bandpp*(ipw-1))
1396 : : end do
1397 : : end do
1398 : : end if
1399 : :
1400 : : ! -------------------
1401 : : ! Fourier calculation
1402 : : ! -------------------
1403 : : ! GPU version
1404 [ - + ]: 30764 : if(gpu_option_/=ABI_GPU_DISABLED) then
1405 [ # # # # : 0 : ABI_MALLOC(weight_t,(bandpp))
# # # # #
# ]
1406 [ # # ]: 0 : do iibandpp=1,bandpp
1407 : : ! Compute the index of the band
1408 : 0 : ind_occ = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + iibandpp
1409 : : ! Compute the weight of the band
1410 : 0 : weight_t(iibandpp)=occ_k(ind_occ)*wtk/ucvol
1411 [ # # ]: 0 : if(abs(occ_k(ind_occ)) < tol8) weight_t(iibandpp) = zero
1412 : : end do
1413 : :
1414 : 0 : nband_fftblock=bandpp/nfft_blocks+modulo(bandpp,nfft_blocks)
1415 : : call fourwf_optmem(1,rhoaug,&
1416 : : & cwavef_alltoall1,&
1417 : : & dummy,wfraug(:,:,:,1:n6*nband_fftblock),gbound_,gbound_,&
1418 : : & istwf_k_,kg_k_gather,kg_k_gather,mgfft,mpi_enreg,bandpp,nfft_blocks,&
1419 : : & ngfft,ndatarecv,1,n4,n5,n6,option_fourwf,tim_fourwf,weight,weight,&
1420 : : & weight_array_r=weight_t,weight_array_i=weight_t,&
1421 [ # # # # ]: 0 : & gpu_option=gpu_option_)
1422 [ # # ]: 0 : ABI_FREE(weight_t)
1423 : :
1424 : : ! Standard version
1425 : : else
1426 [ + + ]: 146244 : do iibandpp=1,bandpp
1427 : : ! Compute the index of the band
1428 : 115480 : ind_occ = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + iibandpp
1429 : : ! Compute the weight of the band
1430 : 115480 : weight=occ_k(ind_occ)*wtk/ucvol
1431 [ - + ]: 115480 : if (option_fourwf==0) then
1432 : 0 : wfraug_ptr => wfraug(:,:,:,(iibandpp-1)*n6+1:iibandpp*n6)
1433 : : else
1434 : 115480 : wfraug_ptr => wfraug
1435 : : end if
1436 [ + + - + ]: 146244 : if (abs(occ_k(ind_occ)) >=tol8.or.option_fourwf==0) then
1437 [ - + ]: 55811 : if(have_to_reequilibrate) then
1438 : : call fourwf(1,rhoaug, &
1439 : : & cwavef_fft(:,(npw_fft*(iibandpp-1))+1:(npw_fft*iibandpp)), &
1440 : : & dummy,wfraug_ptr,gbound_,gbound_,&
1441 : : & istwf_k_,kg_k_fft,kg_k_fft,mgfft,mpi_enreg,1,&
1442 : : & ngfft,npw_fft,1,n4,n5,n6,option_fourwf,tim_fourwf,weight,weight,&
1443 [ # # # # : 0 : & gpu_option=gpu_option_)
# # ]
1444 : : else
1445 : : call fourwf(1,rhoaug,&
1446 : : & cwavef_alltoall1(:,(ndatarecv*(iibandpp-1))+1:(ndatarecv*iibandpp)),&
1447 : : & dummy,wfraug_ptr,gbound_,gbound_,&
1448 : : & istwf_k_,kg_k_gather,kg_k_gather,mgfft,mpi_enreg,1,&
1449 : : & ngfft,ndatarecv,1,n4,n5,n6,option_fourwf,&
1450 [ - + - + : 55811 : & tim_fourwf,weight,weight)
- + ]
1451 : : end if
1452 [ - + - - ]: 55811 : if (option_fourwf==0.and.nproc_fft>1) then
1453 [ # # ]: 0 : if (me_fft>0) then
1454 : 0 : nd3=(ngfft(3)-1)/nproc_fft+1
1455 [ # # # # : 0 : wfraug_ptr(:,:,:,me_fft*nd3+1:me_fft*nd3+nd3)=wfraug_ptr(:,:,:,1:nd3)
# # # # #
# # # # #
# # # # ]
1456 [ # # # # : 0 : wfraug_ptr(:,:,:,1:nd3)=zero
# # # # ]
1457 : : end if
1458 : 0 : call xmpi_sum(wfraug_ptr,mpi_enreg%comm_fft,ier)
1459 : : end if
1460 : : end if
1461 : : end do
1462 : : end if ! (gpu_option/=0)
1463 : :
1464 : : ! -----------------------------------------------------
1465 : : ! Sorting waves functions below the processors
1466 : : ! -----------------------------------------------------
1467 : : ! cwavef_alltoall(:,index_wavef_band) = cwavef_alltoall(:,:) ! NOT NEEDED
1468 [ - + ]: 30764 : ABI_FREE(index_wavef_band)
1469 : :
1470 : : !====================================================================
1471 [ + - ]: 6622 : else if (flag_inv_sym) then
1472 : :
1473 : : ! -------------------------------------------------------------
1474 : : ! Computation of the index to class the waves functions below bandpp
1475 : : ! -------------------------------------------------------------
1476 : : call prep_index_wavef_bandpp(nproc_band,bandpp,&
1477 : : & 1,ndatarecv,&
1478 : : & recvcounts,rdispls,&
1479 : 6622 : & index_wavef_band)
1480 : :
1481 : : ! -------------------------------------------------------
1482 : : ! Sorting the wave functions below bandpp
1483 : : ! -------------------------------------------------------
1484 [ + - ]: 6622 : if(gpu_option_==ABI_GPU_OPENMP) then
1485 : : !FIXME Having this OpenMP kernel here cause a seg fault, so I put it in a separate routine
1486 : : ! I can't figure what cause the issue...
1487 : : call omp_copy_sort_wf(bandpp,ndatarecv,cwavef_alltoall1,cwavef_alltoall2,index_wavef_band,flag_inv_sym)
1488 : : else
1489 [ + + + + ]: 75855454 : cwavef_alltoall1(:,:) = cwavef_alltoall2(:,index_wavef_band)
1490 : : end if
1491 : :
1492 : : ! ------------------------------------------------------------
1493 : : ! We associate the waves functions by two
1494 : : ! ------------------------------------------------------------
1495 : : call prep_wavef_sym_do(mpi_enreg,bandpp,1,&
1496 : : & ndatarecv,&
1497 : : & ndatarecv_tot,ndatasend_sym,tab_proc,&
1498 : : & cwavef_alltoall1,&
1499 : : & sendcounts_sym,sdispls_sym,&
1500 : : & recvcounts_sym,rdispls_sym,&
1501 : : & ewavef_alltoall_sym,&
1502 : 6622 : & index_wavef_send)
1503 : :
1504 : : ! ------------------------------------------------------------
1505 : : ! Fourier calculation
1506 : : ! ------------------------------------------------------------
1507 : : ! GPU version
1508 [ - + ]: 6622 : if (gpu_option_/=ABI_GPU_DISABLED) then
1509 [ # # # # : 0 : ABI_MALLOC(weight1_t,(bandpp_sym))
# # # # #
# # # ]
1510 [ # # # # : 0 : ABI_MALLOC(weight2_t,(bandpp_sym))
# # # # ]
1511 [ # # ]: 0 : do iibandpp=1,bandpp_sym
1512 [ # # ]: 0 : if (bandpp/=1) then
1513 : 0 : ind_occ1 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + (2*iibandpp-1)
1514 : 0 : ind_occ2 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + (2*iibandpp )
1515 : : else
1516 : 0 : ind_occ1 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + 1
1517 : 0 : ind_occ2 = ind_occ1
1518 : : end if
1519 : 0 : weight1_t(iibandpp) = occ_k(ind_occ1)*wtk/ucvol
1520 : 0 : weight2_t(iibandpp) = occ_k(ind_occ2)*wtk/ucvol
1521 : : end do
1522 : :
1523 : 0 : nband_fftblock=bandpp_sym/nfft_blocks_sym+modulo(bandpp_sym,nfft_blocks_sym)
1524 : : call fourwf_optmem(1,rhoaug,&
1525 : : & ewavef_alltoall_sym,&
1526 : : & dummy,wfraug(:,:,:,1:n6*nband_fftblock),gbound_,gbound_,&
1527 : : & istwf_k_,kg_k_gather_sym,kg_k_gather_sym,mgfft,mpi_enreg,bandpp_sym,nfft_blocks_sym,&
1528 : : & ngfft,ndatarecv_tot,1,n4,n5,n6,option_fourwf,tim_fourwf,weight,weight,&
1529 : : & weight_array_r=weight1_t,weight_array_i=weight2_t,&
1530 [ # # # # : 0 : & gpu_option=gpu_option_)
# # ]
1531 : :
1532 [ # # ]: 0 : ABI_FREE(weight1_t)
1533 [ # # ]: 0 : ABI_FREE(weight2_t)
1534 : :
1535 : : ! Standard version
1536 : : else
1537 [ - + - - ]: 6622 : if (option_fourwf==0.and.bandpp>1) then
1538 [ # # # # : 0 : ABI_MALLOC(wfraug_ptr,(2,n4,n5,n6))
# # # # #
# # # # #
# # # # #
# ]
1539 : : else
1540 : 6622 : wfraug_ptr => wfraug
1541 : : end if
1542 [ + + ]: 23676 : do iibandpp=1,bandpp_sym
1543 [ + + ]: 17054 : if (bandpp/=1) then
1544 : 12926 : ind_occ1 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + (2*iibandpp-1)
1545 : 12926 : ind_occ2 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + (2*iibandpp )
1546 : : else
1547 : 4128 : ind_occ1 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + 1
1548 : 4128 : ind_occ2 = ind_occ1
1549 : : end if
1550 : 17054 : weight1 = occ_k(ind_occ1)*wtk/ucvol
1551 : 17054 : weight2 = occ_k(ind_occ2)*wtk/ucvol
1552 : 17054 : ewavef_alltoall_sym_work => ewavef_alltoall_sym(:,(ndatarecv_tot*(iibandpp-1))+1:(ndatarecv_tot*iibandpp))
1553 : : call fourwf(1,rhoaug,&
1554 : : & ewavef_alltoall_sym_work,&
1555 : : & dummy,wfraug_ptr,gbound_,gbound_,&
1556 : : & istwf_k_,kg_k_gather_sym,kg_k_gather_sym,mgfft,mpi_enreg,1,&
1557 : : & ngfft,ndatarecv_tot,1,n4,n5,n6,option_fourwf,&
1558 [ - + - + : 17054 : & tim_fourwf,weight1,weight2)
- + - + ]
1559 [ - + ]: 23676 : if (option_fourwf==0) then
1560 [ # # ]: 0 : if (modulo(bandpp,2)==0) then
1561 : 0 : jjbandpp=2*iibandpp-1
1562 [ # # # # : 0 : wfraug(1,:,:,(jjbandpp-1)*n6+1:jjbandpp*n6)=wfraug_ptr(1,:,:,1:n6)
# # # # #
# # # #
# ]
1563 [ # # # # : 0 : wfraug(1,:,:,(jjbandpp)*n6+1:(jjbandpp+1)*n6)=wfraug_ptr(2,:,:,1:n6)
# # # # #
# # # ]
1564 [ # # ]: 0 : else if (bandpp>1) then
1565 [ # # # # : 0 : wfraug(1,:,:,(iibandpp-1)*n6+1:iibandpp*n6)=wfraug_ptr(1,:,:,1:n6)
# # # # #
# # # #
# ]
1566 : : end if
1567 [ # # ]: 0 : if (nproc_fft>1) then
1568 [ # # ]: 0 : if (me_fft>0) then
1569 : 0 : nd3=(ngfft(3)-1)/nproc_fft+1
1570 [ # # # # : 0 : wfraug(1,:,:,me_fft*nd3+1:me_fft*nd3+nd3)=wfraug(1,:,:,1:nd3)
# # # # #
# # # #
# ]
1571 [ # # # # : 0 : wfraug(1,:,:,1:nd3)=zero
# # ]
1572 : : end if
1573 : 0 : call xmpi_sum(wfraug,mpi_enreg%comm_fft,ier)
1574 : : end if
1575 : : end if
1576 : : end do
1577 [ - + - - ]: 6622 : if (option_fourwf==0.and.bandpp>1) then
1578 [ # # ]: 0 : ABI_FREE(wfraug_ptr)
1579 : : end if
1580 : : end if ! (gpu_option/=ABI_GPU_DISABLED)
1581 : :
1582 : : ! ------------------------------------------------------------
1583 : : ! We dissociate each wave function in two waves functions
1584 : : ! gwavef is classed below of bandpp
1585 : : ! ------------------------------------------------------------
1586 : : call prep_wavef_sym_undo(mpi_enreg,bandpp,1,&
1587 : : & ndatarecv,&
1588 : : & ndatarecv_tot,ndatasend_sym,idatarecv0,&
1589 : : & cwavef_alltoall1,&
1590 : : & sendcounts_sym,sdispls_sym,&
1591 : : & recvcounts_sym,rdispls_sym,&
1592 : : & ewavef_alltoall_sym,&
1593 : 6622 : & index_wavef_send)
1594 : :
1595 [ - + ]: 6622 : ABI_FREE(ewavef_alltoall_sym)
1596 [ - + ]: 6622 : ABI_FREE(index_wavef_send)
1597 : :
1598 : : ! -------------------------------------------------------
1599 : : ! Sorting waves functions below the processors
1600 : : ! -------------------------------------------------------
1601 : : ! cwavef_alltoall(:,index_wavef_band) = cwavef_alltoall(:,:) ! NOT NEEDED
1602 : :
1603 [ - + ]: 6622 : ABI_FREE(index_wavef_band)
1604 : :
1605 : : end if
1606 : :
1607 : : !====================================================================
1608 [ - + ]: 185482 : if(have_to_reequilibrate) then
1609 [ # # ]: 0 : ABI_FREE(buff_wf)
1610 [ # # ]: 0 : ABI_FREE(cwavef_fft)
1611 [ # # ]: 0 : if(bandpp > 1) then
1612 [ # # ]: 0 : ABI_FREE(cwavef_fft_tr)
1613 : : end if
1614 : : end if
1615 : : #ifdef HAVE_OPENMP_OFFLOAD
1616 : : !$OMP TARGET EXIT DATA MAP(delete:cwavef_alltoall2) IF(gpu_option_==ABI_GPU_OPENMP)
1617 : : if ( ((.not.flag_inv_sym) .and. (bandpp>1) ) .or. flag_inv_sym ) then
1618 : : !$OMP TARGET EXIT DATA MAP(delete:cwavef_alltoall1) IF(gpu_option_==ABI_GPU_OPENMP)
1619 : : end if
1620 : : #endif
1621 [ - + ]: 185482 : ABI_FREE(sendcountsloc)
1622 [ - + ]: 185482 : ABI_FREE(sdisplsloc)
1623 [ - + ]: 185482 : ABI_FREE(recvcountsloc)
1624 [ - + ]: 185482 : ABI_FREE(rdisplsloc)
1625 [ - + ]: 185482 : ABI_FREE(cwavef_alltoall2)
1626 [ - + + + : 408350 : if ( ((.not.flag_inv_sym) .and. (bandpp>1) ) .or. flag_inv_sym ) then
+ + - + -
+ - + - +
- + - + ]
1627 [ + - ]: 37386 : if(gpu_option_==ABI_GPU_KOKKOS) then
1628 : : #if defined HAVE_GPU && defined HAVE_YAKL
1629 : : ABI_FREE_MANAGED(cwavef_alltoall1)
1630 : : #endif
1631 : : else
1632 [ - + ]: 37386 : ABI_FREE(cwavef_alltoall1)
1633 : : end if
1634 : : end if
1635 : :
1636 : : contains
1637 [ - + ]: 37386 : subroutine omp_copy_sort_wf(bandpp,ndatarecv,cwavef_alltoall_out,cwavef_alltoall_in,index_wavef_band,flag_inv_sym)
1638 : : !Arguments ------------------------------------
1639 : : !scalars
1640 : : integer,intent(in) :: bandpp,ndatarecv
1641 : : logical,intent(in) :: flag_inv_sym
1642 : : !arrays
1643 : : integer,intent(in) :: index_wavef_band(bandpp*ndatarecv)
1644 : : real(dp),intent(in) :: cwavef_alltoall_in(2,bandpp*ndatarecv)
1645 : : real(dp),intent(out) :: cwavef_alltoall_out(2,bandpp*ndatarecv)
1646 : :
1647 : : !Locals ---------------------------------------
1648 : : integer :: ind,iibandpp,ipw
1649 : :
1650 : : ! *************************************************************************
1651 : :
1652 : : #ifdef HAVE_OPENMP_OFFLOAD
1653 : : !FIXME This doesn't run fine when flag_invsym is on, why !!???
1654 : : !$OMP TARGET UPDATE FROM(cwavef_alltoall_in) IF(flag_inv_sym)
1655 : :
1656 : : !$OMP TARGET TEAMS DISTRIBUTE &
1657 : : !$OMP& PRIVATE(iibandpp) MAP(to:cwavef_alltoall_out,cwavef_alltoall_in,index_wavef_band) &
1658 : : !$OMP& IF(.not. flag_inv_sym)
1659 : : do iibandpp=1,bandpp
1660 : : !$OMP PARALLEL DO PRIVATE(ind,ipw)
1661 : : do ipw = 1 ,ndatarecv
1662 : : ind=index_wavef_band(ipw + ndatarecv*(iibandpp-1))
1663 : : cwavef_alltoall_out(1,ipw + ndatarecv*(iibandpp-1)) = cwavef_alltoall_in(1,ind)
1664 : : cwavef_alltoall_out(2,ipw + ndatarecv*(iibandpp-1)) = cwavef_alltoall_in(2,ind)
1665 : : end do
1666 : : end do
1667 : :
1668 : : !$OMP TARGET UPDATE TO(cwavef_alltoall_out) IF(flag_inv_sym)
1669 : :
1670 : : #else
1671 : : ABI_UNUSED((/ind,iibandpp,ipw/))
1672 : : ABI_UNUSED((/bandpp,ndatarecv/))
1673 : : ABI_UNUSED(flag_inv_sym)
1674 : : ABI_UNUSED(cwavef_alltoall_in)
1675 : : ABI_UNUSED(cwavef_alltoall_out)
1676 : : ABI_UNUSED(index_wavef_band)
1677 : : #endif
1678 : : end subroutine omp_copy_sort_wf
1679 : :
1680 : : end subroutine prep_fourwf
1681 : : !!***
1682 : :
1683 : : !!****f* ABINIT/prep_wavef_sym_do
1684 : : !! NAME
1685 : : !! prep_wavef_sym_do
1686 : : !!
1687 : : !! FUNCTION
1688 : : !! this routine associates waves functions by two as following
1689 : : !! E(G) = C(G) + D(G)
1690 : : !! E(-G) = C*(G) + iD*(G)
1691 : : !! the values are distributed on the processors in function of
1692 : : !! the value of mpi_enreg%distribfft%tab_fftwf2_distrib( (-kg_k_gather(2,i) )
1693 : : !!
1694 : : !! INPUTS
1695 : : !! mpi_enreg = information about mpi parallelization
1696 : : !! bandpp = number of couple of waves functions
1697 : : !! nspinor = number of spin
1698 : : !! ndatarecv = number of values received by the processor and sended
1699 : : !! by the other processors band
1700 : : !! ndatarecv_tot = total number of received values
1701 : : !! (ndatarecv + number of received opposited planewave coordinates)
1702 : : !! ndatasend_sym = number of sended values to the processors fft to create opposited
1703 : : !! planewave coordinates
1704 : : !! tab_proc = positions of opposited planewave coordinates in the list of the
1705 : : !! processors fft
1706 : : !! cwavef_alltoall = planewave coefficients of wavefunction
1707 : : !! ( initial of the processor + sended by other processors band)
1708 : : !! sendcounts_sym = number of sended values by the processor to each processor fft
1709 : : !! sdispls_sym = postions of the sended values by the processor to each processor fft
1710 : : !!
1711 : : !! recvcounts_sym = number of the received values by the processor from each processor fft
1712 : : !! rdispls_sym = postions of the received values by the processor from each processor fft
1713 : : !!
1714 : : !! OUTPUT
1715 : : !! ewavef_alltoall_sym = planewave coefficients of wavefunction
1716 : : !! initial of the processor +
1717 : : !! sended by other processors band +
1718 : : !! sended by other processors fft +
1719 : : !! and compisited if bandpp >1
1720 : : !! index_wavef_send = index to send the values in blocks to the other processor fft
1721 : : !!
1722 : : !! SIDE EFFECTS
1723 : : !!
1724 : : !! SOURCE
1725 : :
1726 : 9896 : subroutine prep_wavef_sym_do(mpi_enreg,bandpp,nspinor,&
1727 : : & ndatarecv,&
1728 : 19792 : & ndatarecv_tot,ndatasend_sym,tab_proc,&
1729 [ - + ]: 9896 : & cwavef_alltoall,&
1730 [ - + ]: 9896 : & sendcounts_sym,sdispls_sym,&
1731 [ - + - + : 9896 : & recvcounts_sym,rdispls_sym,&
- + ]
1732 : : & ewavef_alltoall_sym,&
1733 : : & index_wavef_send)
1734 : :
1735 : : !Arguments ------------------------------------
1736 : : !scalars
1737 : : integer,intent(in) :: bandpp,ndatarecv,ndatarecv_tot,ndatasend_sym
1738 : : integer,intent(in) :: nspinor
1739 : : type(mpi_type),intent(in) :: mpi_enreg
1740 : : !arrays
1741 : : integer,allocatable,intent(out) :: index_wavef_send(:)
1742 : : integer,intent(in) :: rdispls_sym(:),recvcounts_sym(:)
1743 : : integer,intent(in) :: sdispls_sym(:),sendcounts_sym(:)
1744 : : integer,intent(in) :: tab_proc(:)
1745 : : real(dp),intent(inout) :: cwavef_alltoall(2,ndatarecv*nspinor*bandpp)
1746 : : real(dp),pointer :: ewavef_alltoall_sym(:,:)
1747 : :
1748 : : !Local variables-------------------------------
1749 : : !scalars
1750 : : integer :: bandpp_sym,ibandpp,idatarecv,ideb_loc,idebc,idebd,idebe
1751 : : integer :: ier,ifin_loc,ifinc,ifind,ifine,iproc,jbandpp,jsendloc
1752 : : integer :: kbandpp,newspacecomm,nproc_fft
1753 : : logical :: flag_compose
1754 : : !arrays
1755 : : integer,allocatable :: rdispls_sym_loc(:),recvcounts_sym_loc(:)
1756 : : integer,allocatable :: sdispls_sym_loc(:),sendcounts_sym_loc(:)
1757 : 9896 : real(dp),allocatable :: ewavef_alltoall_loc(:,:),ewavef_alltoall_send(:,:)
1758 : :
1759 : : ! *********************************************************************
1760 : :
1761 : : !DEBUG
1762 : : !write(std_out,*)' prep_wavef_sym_do : enter '
1763 : : !ENDDEBUG
1764 : :
1765 : : !---------------------------------------------
1766 : : !Initialisation
1767 : : !---------------------------------------------
1768 : 9896 : nproc_fft = mpi_enreg%nproc_fft
1769 : :
1770 : 9896 : newspacecomm = mpi_enreg%comm_fft
1771 : :
1772 [ + + ]: 9896 : if (modulo(bandpp,2)==0) then
1773 : 4808 : bandpp_sym = bandpp/2
1774 : 4808 : flag_compose = .TRUE.
1775 : : else
1776 : : bandpp_sym = bandpp
1777 : : flag_compose = .FALSE.
1778 : : end if
1779 : :
1780 : : !---------------------------------------------
1781 : : !Allocation
1782 : : !---------------------------------------------
1783 [ - + + - : 29688 : ABI_MALLOC(ewavef_alltoall_sym ,(2,ndatarecv_tot*bandpp_sym))
+ - + - -
+ - + ]
1784 [ - + + - : 29688 : ABI_MALLOC(ewavef_alltoall_loc ,(2,ndatarecv *bandpp_sym))
+ - + - -
+ - + ]
1785 [ - + + - : 29688 : ABI_MALLOC(ewavef_alltoall_send ,(2,ndatasend_sym*bandpp_sym))
+ - + - -
+ - + ]
1786 [ - + + - : 29688 : ABI_MALLOC(index_wavef_send ,( ndatasend_sym*bandpp_sym))
+ - + - -
+ - + -
+ ]
1787 : :
1788 [ - + + - : 29688 : ABI_MALLOC(sendcounts_sym_loc ,(nproc_fft))
+ - + - -
+ - + ]
1789 [ - + + - : 19792 : ABI_MALLOC(sdispls_sym_loc ,(nproc_fft))
- + - + ]
1790 [ - + + - : 19792 : ABI_MALLOC(recvcounts_sym_loc ,(nproc_fft))
- + - + ]
1791 [ - + + - : 19792 : ABI_MALLOC(rdispls_sym_loc ,(nproc_fft))
- + - + ]
1792 : :
1793 : :
1794 : : !Initialisation
1795 : : !--------------
1796 [ + + + + ]: 106566374 : ewavef_alltoall_sym(:,:) =0.
1797 [ + + + + ]: 53319068 : ewavef_alltoall_loc(:,:) =0.
1798 : :
1799 [ + + ]: 21958 : sendcounts_sym_loc(:) =0
1800 [ + + ]: 21958 : sdispls_sym_loc(:) =0
1801 [ + + ]: 21958 : recvcounts_sym_loc(:) =0
1802 [ + + ]: 21958 : rdispls_sym_loc(:) =0
1803 : :
1804 [ + + ]: 17758998 : index_wavef_send(:) =0
1805 : :
1806 : :
1807 : : !-------------------------------------------------
1808 : : !We are bandpp blocks which we want to :
1809 : : !associate by two (band_sym==bandpp/2)
1810 : : !or not associate by two (band_sym==bandpp)
1811 : : !
1812 : : !So We'll have got bandpp_sym blocks
1813 : : !So we loop on the bandpp_sym blocks
1814 : : !--------------------------------------------------
1815 : :
1816 [ + + ]: 32494 : do kbandpp=1,bandpp_sym
1817 : :
1818 : : ! position of the two blocks
1819 : : ! --------------------------
1820 : 22598 : ibandpp = (kbandpp-1) * 2
1821 : 22598 : jbandpp = ibandpp + 1
1822 : :
1823 : 22598 : idebe = (kbandpp-1) * ndatarecv_tot + 1
1824 : 22598 : ifine = idebe + ndatarecv - 1
1825 : :
1826 : 22598 : idebc = ibandpp * ndatarecv + 1
1827 : 22598 : ifinc = idebc + ndatarecv - 1
1828 : :
1829 : 22598 : idebd = jbandpp * ndatarecv + 1
1830 : 22598 : ifind = idebd + ndatarecv - 1
1831 : :
1832 : 22598 : ideb_loc = (kbandpp-1) * ndatarecv + 1
1833 : 22598 : ifin_loc = ideb_loc + ndatarecv - 1
1834 : :
1835 : :
1836 [ + + ]: 32494 : if (flag_compose) then
1837 : :
1838 : :
1839 : : ! calcul ewavef(G)
1840 : : ! ----------------
1841 : : ewavef_alltoall_sym(1,idebe:ifine) = &
1842 : : & cwavef_alltoall(1,idebc:ifinc) &
1843 [ + + ]: 16760514 : & - cwavef_alltoall(2,idebd:ifind)
1844 : :
1845 : : ewavef_alltoall_sym(2,idebe:ifine) = &
1846 : : & cwavef_alltoall(2,idebc:ifinc) &
1847 [ + + ]: 16760514 : & + cwavef_alltoall(1,idebd:ifind)
1848 : :
1849 : : ! calcul ewavef_loc(-G)
1850 : : ! ---------------------
1851 : : ewavef_alltoall_loc(1,ideb_loc:ifin_loc) = &
1852 : : & cwavef_alltoall(1,idebc:ifinc) &
1853 [ + + ]: 16760514 : & + cwavef_alltoall(2,idebd:ifind)
1854 : :
1855 : : ewavef_alltoall_loc(2,ideb_loc:ifin_loc) = &
1856 : : & - cwavef_alltoall(2,idebc:ifinc) &
1857 [ + + ]: 16760514 : & + cwavef_alltoall(1,idebd:ifind)
1858 : : else
1859 : :
1860 : : ! calcul ewavef(G)
1861 : : ! ----------------
1862 [ + + ]: 1031808 : ewavef_alltoall_sym(1,idebe:ifine) = cwavef_alltoall(1,idebc:ifinc)
1863 [ + + ]: 1031808 : ewavef_alltoall_sym(2,idebe:ifine) = cwavef_alltoall(2,idebc:ifinc)
1864 : :
1865 : : ! calcul ewavef_loc(-G)
1866 : : ! ---------------------
1867 [ + + ]: 1031808 : ewavef_alltoall_loc(1,ideb_loc:ifin_loc) = cwavef_alltoall(1,idebc:ifinc)
1868 [ + + ]: 1031808 : ewavef_alltoall_loc(2,ideb_loc:ifin_loc) = - cwavef_alltoall(2,idebc:ifinc)
1869 : :
1870 : : end if
1871 : :
1872 : : end do
1873 : :
1874 : :
1875 : :
1876 : : !------------------------------------------------------------------------
1877 : : !Creation of datas blocks for each processor fft from ewavef_alltoall_loc
1878 : : !to send datas by blocks with a alltoall...
1879 : : !------------------------------------------------------------------------
1880 : :
1881 : : !Position of the blocks
1882 : : !----------------------
1883 : : jsendloc=0
1884 [ + + ]: 32494 : do ibandpp=1,bandpp_sym
1885 [ + + ]: 59044 : do iproc=1,nproc_fft
1886 [ + + ]: 20112644 : do idatarecv=1,ndatarecv
1887 [ + + ]: 20090046 : if (tab_proc(idatarecv)==(iproc-1)) then
1888 : 17749102 : jsendloc=jsendloc+1
1889 : 17749102 : index_wavef_send(jsendloc) = idatarecv + ndatarecv * (ibandpp-1)
1890 : : end if
1891 : : end do
1892 : : end do
1893 : : end do
1894 : :
1895 : : !Classment
1896 : : !----------
1897 [ + + + + ]: 53257202 : ewavef_alltoall_send(:,:)=ewavef_alltoall_loc(:,index_wavef_send)
1898 : :
1899 : :
1900 : : !-------------------------------------------------
1901 : : !Calcul of the number of received and sended datas
1902 : : !-------------------------------------------------
1903 [ + - - + : 31854 : sendcounts_sym_loc = sendcounts_sym*2
- - + + ]
1904 [ + - - + : 31854 : recvcounts_sym_loc = recvcounts_sym*2
- - + + ]
1905 : :
1906 : : !------------------------------------------
1907 : : !Exchange of the datas ewavef_allto_all_loc
1908 : : !------------------------------------------
1909 [ + + ]: 32494 : do ibandpp=1,bandpp_sym
1910 : :
1911 : : ! ------------------------------------------------
1912 : : ! Deplacment of the sended datas because of bandpp
1913 : : ! ------------------------------------------------
1914 [ + + ]: 49148 : sdispls_sym_loc(:) = sdispls_sym(:) + ndatasend_sym * (ibandpp-1)
1915 [ + + ]: 49148 : sdispls_sym_loc = sdispls_sym_loc *2
1916 : :
1917 : : ! --------------------------------------------------
1918 : : ! Deplacment of the received datas because of bandpp
1919 : : ! --------------------------------------------------
1920 [ + + ]: 49148 : rdispls_sym_loc(:) = rdispls_sym(:) + ndatarecv_tot * (ibandpp-1)
1921 [ + + ]: 49148 : rdispls_sym_loc = rdispls_sym_loc *2
1922 : :
1923 : :
1924 : : call xmpi_alltoallv(&
1925 : : & ewavef_alltoall_send(:,:) ,sendcounts_sym_loc,sdispls_sym_loc,&
1926 : : & ewavef_alltoall_sym(:,:) ,recvcounts_sym_loc,rdispls_sym_loc,&
1927 : 32494 : & newspacecomm,ier)
1928 : :
1929 : : end do
1930 : :
1931 : : !-----------------------
1932 : : !Desallocation
1933 : : !-----------------------
1934 : :
1935 [ - + ]: 9896 : ABI_FREE(sendcounts_sym_loc)
1936 [ - + ]: 9896 : ABI_FREE(recvcounts_sym_loc)
1937 [ - + ]: 9896 : ABI_FREE(sdispls_sym_loc)
1938 [ - + ]: 9896 : ABI_FREE(rdispls_sym_loc)
1939 : :
1940 [ - + ]: 9896 : ABI_FREE(ewavef_alltoall_loc)
1941 [ - + ]: 9896 : ABI_FREE(ewavef_alltoall_send)
1942 : :
1943 : 9896 : end subroutine prep_wavef_sym_do
1944 : : !!***
1945 : :
1946 : : !!****f* ABINIT/prep_wavef_sym_undo
1947 : : !! NAME
1948 : : !! prep_wavef_sym_undo
1949 : : !!
1950 : : !! FUNCTION
1951 : : !! this routine dissociates each wave function in two waves functions as following
1952 : : !! C(G) = ( E*(-G) + E(G))/2
1953 : : !! D(G) = i*( E*(-G) - E(G))/2
1954 : : !! the values are redistributed on the processors in function of
1955 : : !! the value of mpi_enreg%distribfft%tab_fftwf2_distrib( (-kg_k_gather(2,i) )
1956 : : !!
1957 : : !! INPUTS
1958 : : !! mpi_enreg = information about mpi parallelization
1959 : : !! bandpp = number of groups of couple of waves functions
1960 : : !! nspinor = number of spin
1961 : : !! ndatarecv = number of values received by the processor and sended
1962 : : !! by the other processors band
1963 : : !! ndatarecv_tot = total number of received values
1964 : : !! (ndatarecv + number of received opposited planewave coordinates)
1965 : : !! ndatasend_sym = number of sended values to the processors fft to create opposited
1966 : : !! planewave coordinates
1967 : : !! idatarecv0 = position of the planewave coordinates (0,0,0)
1968 : : !! sendcounts_sym = number of sended values by the processor to each processor fft
1969 : : !! sdispls_sym = postions of the sended values by the processor to each processor fft
1970 : : !!
1971 : : !! recvcounts_sym = number of the received values by the processor to each processor fft
1972 : : !!! rdispls_sym = postions of the received values by the processor to each processor fft
1973 : : !!
1974 : : !! gwavef_alltoall_sym = planewave coefficients of wavefunction
1975 : : !! initial of the processor +
1976 : : !! sended by other processors band +
1977 : : !! sended by other processors fft +
1978 : : !! and composited if bandpp >1
1979 : : !! index_wavef_send = index to send the values by block to the other processor fft
1980 : : !!
1981 : : !! OUTPUT
1982 : : !! gwavef_alltoall = planewave coefficients of wavefunction
1983 : : !! ( for of the processor + to send to other processors band)
1984 : : !!
1985 : : !! SOURCE
1986 : :
1987 : 12576 : subroutine prep_wavef_sym_undo(mpi_enreg,bandpp,nspinor,&
1988 : : & ndatarecv,&
1989 : : & ndatarecv_tot,ndatasend_sym,idatarecv0,&
1990 : 12576 : & gwavef_alltoall,&
1991 [ - + ]: 12576 : & sendcounts_sym,sdispls_sym,&
1992 [ - + - + : 12576 : & recvcounts_sym,rdispls_sym,&
- + ]
1993 [ - + ]: 12576 : & gwavef_alltoall_sym,&
1994 [ - + ]: 12576 : & index_wavef_send)
1995 : :
1996 : : !Arguments ------------------------------------
1997 : : !scalars
1998 : : integer,intent(in) :: bandpp,idatarecv0,ndatarecv,ndatarecv_tot,ndatasend_sym
1999 : : integer,intent(in) :: nspinor
2000 : : type(mpi_type),intent(in) :: mpi_enreg
2001 : : !arrays
2002 : : integer,intent(in) :: index_wavef_send(:),rdispls_sym(:),recvcounts_sym(:)
2003 : : integer,intent(in) :: sdispls_sym(:),sendcounts_sym(:)
2004 : : real(dp),intent(inout) :: gwavef_alltoall(2,ndatarecv*nspinor*bandpp)
2005 : : real(dp),intent(inout) :: gwavef_alltoall_sym(:,:)
2006 : :
2007 : : !Local variables-------------------------------
2008 : : !scalars
2009 : : integer :: bandpp_sym,ibandpp,ideb_loc,idebc,idebd
2010 : : integer :: idebe,ier,ifin_loc,ifinc,ifind,ifine
2011 : : integer :: jbandpp,kbandpp,newspacecomm,nproc_fft
2012 : : logical :: flag_compose
2013 : : !arrays
2014 : : integer,allocatable :: rdispls_sym_loc(:),recvcounts_sym_loc(:)
2015 : : integer,allocatable :: sdispls_sym_loc(:),sendcounts_sym_loc(:)
2016 : 12576 : real(dp),allocatable :: gwavef_alltoall_loc(:,:),gwavef_alltoall_rcv(:,:)
2017 : :
2018 : : ! *********************************************************************
2019 : :
2020 : : !DEBUG
2021 : : !write(std_out,*)' prep_wavef_sym_undo : enter '
2022 : : !ENDDEBUG
2023 : :
2024 : :
2025 : : !---------------------------------------------
2026 : : !Initialisation
2027 : : !---------------------------------------------
2028 : 12576 : nproc_fft = mpi_enreg%nproc_fft
2029 : :
2030 : 12576 : newspacecomm = mpi_enreg%comm_fft
2031 : :
2032 [ + + ]: 12576 : if (modulo(bandpp,2)==0) then
2033 : 7488 : bandpp_sym = bandpp/2
2034 : 7488 : flag_compose = .TRUE.
2035 : : else
2036 : : bandpp_sym = bandpp
2037 : : flag_compose = .FALSE.
2038 : : end if
2039 : :
2040 : : !---------------------------------------------
2041 : : !Allocation
2042 : : !---------------------------------------------
2043 [ - + + - : 37728 : ABI_MALLOC(gwavef_alltoall_loc ,(2,ndatarecv *bandpp_sym))
+ - + - -
+ - + ]
2044 [ - + + - : 37728 : ABI_MALLOC(gwavef_alltoall_rcv ,(2,ndatasend_sym *bandpp_sym))
+ - + - -
+ - + ]
2045 : :
2046 [ - + + - : 37728 : ABI_MALLOC(sendcounts_sym_loc ,(nproc_fft))
+ - + - -
+ - + ]
2047 [ - + + - : 25152 : ABI_MALLOC(sdispls_sym_loc ,(nproc_fft))
- + - + ]
2048 [ - + + - : 25152 : ABI_MALLOC(recvcounts_sym_loc ,(nproc_fft))
- + - + ]
2049 [ - + + - : 25152 : ABI_MALLOC(rdispls_sym_loc ,(nproc_fft))
- + - + ]
2050 : :
2051 : :
2052 : : !---------------------------------------------
2053 : : !Initialisation
2054 : : !---------------------------------------------
2055 [ + + + + ]: 67046016 : gwavef_alltoall_loc(:,:) =0.
2056 : :
2057 [ + + ]: 27966 : sendcounts_sym_loc(:) =0
2058 [ + + ]: 27966 : sdispls_sym_loc(:) =0
2059 [ + + ]: 27966 : recvcounts_sym_loc(:) =0
2060 [ + + ]: 27966 : rdispls_sym_loc(:) =0
2061 : :
2062 : :
2063 : : !-------------------------------------------------
2064 : : !Calcul of number of the sended and received datas
2065 : : !-------------------------------------------------
2066 [ + - - + : 40542 : sendcounts_sym_loc = sendcounts_sym*2
- - + + ]
2067 [ + - - + : 40542 : recvcounts_sym_loc = recvcounts_sym*2
- - + + ]
2068 : :
2069 : : !----------------------------------------------------
2070 : : !Sending of the values
2071 : : !----------------------------------------------------
2072 [ + + ]: 39930 : do ibandpp = 1,bandpp_sym
2073 : :
2074 : : ! -------------------------------------------------
2075 : : ! Deplacment of the sended values because of bandpp
2076 : : ! -------------------------------------------------
2077 [ + + ]: 60632 : sdispls_sym_loc(:) = sdispls_sym(:) + ndatasend_sym * (ibandpp-1)
2078 [ + + ]: 60632 : sdispls_sym_loc = sdispls_sym_loc *2
2079 : :
2080 : : ! ---------------------------------------------------
2081 : : ! Deplacment of the received values because of bandpp
2082 : : ! ---------------------------------------------------
2083 [ + + ]: 60632 : rdispls_sym_loc(:) = rdispls_sym(:) + ndatarecv_tot * (ibandpp-1)
2084 [ + + ]: 60632 : rdispls_sym_loc = rdispls_sym_loc *2
2085 : :
2086 : :
2087 : : call xmpi_alltoallv(&
2088 : : & gwavef_alltoall_sym(:,:) ,recvcounts_sym_loc,rdispls_sym_loc,&
2089 : : & gwavef_alltoall_rcv(:,:) ,sendcounts_sym_loc,sdispls_sym_loc,&
2090 : 39930 : & newspacecomm,ier)
2091 : :
2092 : : end do
2093 : :
2094 : :
2095 : : !----------------------
2096 : : !Dispatching the blocks
2097 : : !----------------------
2098 [ + + + + ]: 66972840 : gwavef_alltoall_loc(:,index_wavef_send(:)) = gwavef_alltoall_rcv(:,:)
2099 : :
2100 : : !----------------------
2101 : : !Case -kg = [0 0 0]
2102 : : !----------------------
2103 [ + + ]: 12576 : if (idatarecv0/=-1) then
2104 [ + + ]: 35561 : do kbandpp=1,bandpp_sym
2105 : : gwavef_alltoall_loc(:,(kbandpp-1)*ndatarecv + idatarecv0)= &
2106 [ + + ]: 84345 : gwavef_alltoall_sym(:,(kbandpp-1)*ndatarecv_tot + idatarecv0)
2107 : : end do
2108 : : end if
2109 : :
2110 : : !---------------------------------------------------
2111 : : !Build of hwavef_alltoall
2112 : : !
2113 : : !We have got :
2114 : : !bandpp_sym blocks to dissociate
2115 : : !or bandpp_sym blokcs to not dissociate
2116 : : !--------------------------------------------------
2117 [ + + ]: 39930 : do kbandpp=1,bandpp_sym
2118 : :
2119 : : ! position of the 2 blocks
2120 : : ! ----------------------------------
2121 : 27354 : ibandpp = (kbandpp-1) * 2
2122 : 27354 : jbandpp = ibandpp + 1
2123 : :
2124 : 27354 : idebe = (kbandpp-1) * ndatarecv_tot + 1
2125 : 27354 : ifine = idebe + ndatarecv - 1
2126 : :
2127 : 27354 : idebc = ibandpp * ndatarecv + 1
2128 : 27354 : ifinc = idebc + ndatarecv - 1
2129 : :
2130 : 27354 : idebd = jbandpp * ndatarecv + 1
2131 : 27354 : ifind = idebd + ndatarecv - 1
2132 : :
2133 : 27354 : ideb_loc = (kbandpp-1) * ndatarecv + 1
2134 : 27354 : ifin_loc = ideb_loc + ndatarecv - 1
2135 : :
2136 : :
2137 [ + + ]: 39930 : if (flag_compose) then
2138 : :
2139 : : ! calcul cwavef(G)
2140 : : ! ----------------
2141 : : gwavef_alltoall(1,idebc:ifinc) = gwavef_alltoall_sym(1,idebe:ifine) &
2142 [ + + ]: 21340026 : & + gwavef_alltoall_loc(1,ideb_loc:ifin_loc)
2143 : : gwavef_alltoall(2,idebc:ifinc) = gwavef_alltoall_sym(2,idebe:ifine) &
2144 [ + + ]: 21340026 : & - gwavef_alltoall_loc(2,ideb_loc:ifin_loc)
2145 : :
2146 : : ! calcul dwavef(G)
2147 : : ! ------------------
2148 : : gwavef_alltoall(1,idebd:ifind) = gwavef_alltoall_sym(2,idebe:ifine) &
2149 [ + + ]: 21340026 : & + gwavef_alltoall_loc(2,ideb_loc:ifin_loc)
2150 : : gwavef_alltoall(2,idebd:ifind) = - gwavef_alltoall_sym(1,idebe:ifine) &
2151 [ + + ]: 21340026 : & + gwavef_alltoall_loc(1,ideb_loc:ifin_loc)
2152 : : else
2153 : :
2154 : : ! calcul cwavef(G)
2155 : : ! ----------------
2156 : : gwavef_alltoall(1,idebc:ifinc) = gwavef_alltoall_sym(1,idebe:ifine) &
2157 [ + + ]: 1031808 : & + gwavef_alltoall_loc(1,ideb_loc:ifin_loc)
2158 : : gwavef_alltoall(2,idebc:ifinc) = gwavef_alltoall_sym(2,idebe:ifine) &
2159 [ + + ]: 1031808 : & - gwavef_alltoall_loc(2,ideb_loc:ifin_loc)
2160 : : end if
2161 : :
2162 : : end do
2163 : :
2164 : : !We divise by two
2165 [ + + + + ]: 130999296 : gwavef_alltoall(:,:) = gwavef_alltoall(:,:)/2
2166 : :
2167 : : !-----------------------
2168 : : !Desallocation
2169 : : !-----------------------
2170 : :
2171 [ - + ]: 12576 : ABI_FREE(sendcounts_sym_loc)
2172 [ - + ]: 12576 : ABI_FREE(recvcounts_sym_loc)
2173 [ - + ]: 12576 : ABI_FREE(sdispls_sym_loc)
2174 [ - + ]: 12576 : ABI_FREE(rdispls_sym_loc)
2175 : :
2176 [ - + ]: 12576 : ABI_FREE(gwavef_alltoall_loc)
2177 [ - + ]: 12576 : ABI_FREE(gwavef_alltoall_rcv)
2178 : :
2179 : 12576 : end subroutine prep_wavef_sym_undo
2180 : : !!***
2181 : :
2182 : : !!****f* ABINIT/prep_index_wavef_bandpp
2183 : : !! NAME
2184 : : !! prep_index_wavef_bandpp
2185 : : !!
2186 : : !! FUNCTION
2187 : : !! this routine sorts the waves functions by bandpp and by processors
2188 : : !! after the alltoall
2189 : : !!
2190 : : !! INPUTS
2191 : : !! nproc_band = number of processors below the band
2192 : : !! bandpp = number of groups of couple of waves functions
2193 : : !! nspinor = number of spin
2194 : : !! ndatarecv = total number of values received by the processor and sended
2195 : : !! by the other processors band
2196 : : !! recvcounts = number of values sended by each processor band and received
2197 : : !! by the processor
2198 : : !! rdispls = positions of the values received by the processor and
2199 : : !! sended by each processor band
2200 : : !!
2201 : : !! OUTPUT
2202 : : !! index_wavef_band = position of the sorted values
2203 : : !!
2204 : : !! SOURCE
2205 : :
2206 : 72596 : subroutine prep_index_wavef_bandpp(nproc_band,bandpp,&
2207 : : nspinor,ndatarecv,&
2208 : 72596 : recvcounts,rdispls,&
2209 : : index_wavef_band)
2210 : :
2211 : : !Arguments ------------------------------------
2212 : : !scalars
2213 : : integer,intent(in) :: bandpp,ndatarecv,nproc_band,nspinor
2214 : : !arrays
2215 : : integer,intent(in) :: rdispls(nproc_band),recvcounts(nproc_band)
2216 : : integer,allocatable,intent(out) :: index_wavef_band(:)
2217 : :
2218 : : !Local variables-------------------------------
2219 : : !scalars
2220 : : integer :: delta,idebc,idebe,ifinc,ifine,iindex,iproc,kbandpp,nb
2221 : :
2222 : : ! *********************************************************************
2223 : :
2224 : : !DEBUG
2225 : : !write(std_out,*)' prep_index_wavef_banpp : enter '
2226 : : !write(std_out,*) 'ndatarecv = ', ndatarecv
2227 : : !write(std_out,*) 'rdispls(:) = ', rdispls(:)
2228 : : !write(std_out,*) 'recvcounts(:) = ', recvcounts(:)
2229 : : !ENDDEBUG
2230 : :
2231 : :
2232 : : !---------------------------------------------
2233 : : !Allocation
2234 : : !---------------------------------------------
2235 [ - + + - : 217788 : ABI_MALLOC(index_wavef_band ,(bandpp*nspinor*ndatarecv))
+ - + - -
+ - + -
+ ]
2236 [ + + ]: 161806596 : index_wavef_band(:) =0
2237 : :
2238 : : !---------------------------------------------
2239 : : !Calcul : loops on bandpp and processors band
2240 : : !---------------------------------------------
2241 [ + + ]: 315916 : nb = sum(recvcounts(1:nproc_band))
2242 [ + + ]: 365306 : do kbandpp=1,bandpp
2243 : :
2244 [ + + ]: 1215584 : do iproc=1,nproc_band
2245 : :
2246 : 850278 : idebe = (rdispls(iproc) + 1) + (kbandpp-1) * ndatarecv*nspinor
2247 : 850278 : ifine = idebe + recvcounts(iproc) -1
2248 : :
2249 [ + + ]: 850278 : if (iproc==1) then
2250 : 292710 : idebc = (kbandpp-1)* recvcounts(iproc)*nspinor + 1
2251 : : else
2252 : : idebc = (bandpp) * sum(recvcounts(1:iproc-1))*nspinor &
2253 : : + (kbandpp-1)* recvcounts(iproc)*nspinor &
2254 [ + + ]: 1540848 : + 1
2255 : : end if
2256 : 850278 : ifinc = idebc + recvcounts(iproc) -1
2257 [ + - + + : 320070700 : index_wavef_band(idebe:ifine) = (/( iindex,iindex=idebc,ifinc)/)
+ + ]
2258 : 850278 : delta=ifine-idebe
2259 [ + + ]: 1142988 : if (nspinor==2) then
2260 [ + + + + ]: 5250560 : index_wavef_band(idebe+nb :idebe+nb +delta)=(/( iindex,iindex=ifinc+1,ifinc+1+delta)/)
2261 : : end if
2262 : : end do
2263 : : end do
2264 : :
2265 : 72596 : end subroutine prep_index_wavef_bandpp
2266 : : !!***
2267 : :
2268 : : !!****f* ABINIT/prep_sort_wavef_spin
2269 : : !! NAME
2270 : : !! prep_sort_wavef_spin
2271 : : !!
2272 : : !! FUNCTION
2273 : : !! Compute index used to sort a spinorial wave-function by spin
2274 : : !! Sort to have all nspinor=1 fisrt, then all nspinor=2
2275 : : !!
2276 : : !! INPUTS
2277 : : !! nproc_band=size of "band" communicator
2278 : : !! nspinor=number of spinorial components of the wavefunction
2279 : : !! ndatarecv=total number of values on all processors
2280 : : !! recvcounts(nproc_band)= number of received values by the processor
2281 : : !! rdispls(nproc_band)= offsets of the received values by the processor
2282 : : !!
2283 : : !! OUTPUT
2284 : : !! index_wavef(:)=array containing the sorted indexes (pointer, allocated in this routine)
2285 : : !!
2286 : : !! SOURCE
2287 : :
2288 : 7368 : subroutine prep_sort_wavef_spin(nproc_band,nspinor,ndatarecv,recvcounts,rdispls,index_wavef)
2289 : :
2290 : : !Arguments ------------------------------------
2291 : : !scalars
2292 : : integer,intent(in) :: ndatarecv,nproc_band,nspinor
2293 : : !arrays
2294 : : integer,intent(in) :: rdispls(nproc_band),recvcounts(nproc_band)
2295 : : integer,allocatable,intent(out) :: index_wavef(:)
2296 : :
2297 : : !Local variables-------------------------------
2298 : : !scalars
2299 : : integer :: isft,isft1,iproc,iindex
2300 : : !arrays
2301 : 7368 : integer,allocatable :: recvcountsloc(:),rdisplsloc(:)
2302 : :
2303 : : ! *********************************************************************
2304 : :
2305 [ - + + - : 22104 : ABI_MALLOC(index_wavef,(ndatarecv*nspinor))
+ - + - -
+ - + -
+ ]
2306 : :
2307 [ - + + - : 22104 : ABI_MALLOC(recvcountsloc,(nproc_band))
+ - + - -
+ - + ]
2308 [ - + + - : 14736 : ABI_MALLOC(rdisplsloc,(nproc_band))
- + + - ]
2309 [ + + ]: 22584 : recvcountsloc(:)=recvcounts(:)*2*nspinor
2310 [ + + ]: 22584 : rdisplsloc(:)=rdispls(:)*2*nspinor
2311 : :
2312 : : !---------------------------------------------
2313 : : !Loops on bandpp and processors band
2314 : : !---------------------------------------------
2315 : 22584 : isft=0
2316 [ + + ]: 22584 : do iproc=1,nproc_band
2317 : :
2318 : : ! ===== Spin up
2319 [ + + ]: 15216 : if (iproc==1) then
2320 : : isft= 0
2321 : : else
2322 [ + + ]: 23544 : isft= sum(recvcounts(1: (iproc-1)))
2323 : : end if
2324 : 15216 : isft1 = 0.5*rdisplsloc(iproc)
2325 : :
2326 : : index_wavef(1+isft:isft+recvcounts(iproc))= &
2327 [ + - + + : 2149520 : & (/(iindex,iindex=isft1+1,isft1+recvcounts(iproc))/)
+ + ]
2328 : :
2329 : : ! =====Spin down
2330 [ + + ]: 15216 : if (iproc==1)then
2331 [ + + ]: 22584 : isft=sum(recvcounts(1:nproc_band))
2332 : : else
2333 : : isft=sum(recvcounts(1:nproc_band)) &
2334 [ + + + + ]: 54936 : & +sum(recvcounts(1:iproc-1))
2335 : : end if
2336 : 15216 : isft1 = 0.5 * rdisplsloc(iproc) + recvcounts(iproc)
2337 : :
2338 : : index_wavef(1+isft:isft+recvcounts(iproc))= &
2339 [ + + + + ]: 2156888 : & (/(iindex,iindex=isft1+1,isft1+ recvcounts(iproc))/)
2340 : :
2341 : : end do
2342 : :
2343 [ - + ]: 7368 : ABI_FREE(recvcountsloc)
2344 [ - + ]: 7368 : ABI_FREE(rdisplsloc)
2345 : :
2346 : 7368 : end subroutine prep_sort_wavef_spin
2347 : : !!***
2348 : :
2349 : : end module m_prep_kgb
2350 : : !!***
|
