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 [ - + ]: 264942 : subroutine prep_nonlop(choice,cpopt,cwaveprj,enlout_block,hamk,idir,lambdablock,&
636 [ - + - + : 264942 : blocksize,mpi_enreg,nnlout,paw_opt,signs,gsc, tim_nonlop,cwavef,gvnlc, &
- + ]
637 [ + + + - ]: 264942 : already_transposed,gpu_option,vectproj,cwavef_tr) ! optional
638 : :
639 : : !Arguments ------------------------------------
640 : : integer, intent(in) :: blocksize,choice,cpopt,idir,signs,nnlout,paw_opt
641 : : logical,optional,intent(in) :: already_transposed
642 : : integer,optional,intent(in) :: gpu_option
643 : : real(dp), intent(in) :: lambdablock(blocksize)
644 : : real(dp), intent(out) , target :: enlout_block(nnlout*blocksize),gvnlc(:,:),gsc(:,:)
645 : : real(dp), intent(inout), target :: cwavef(:,:)
646 : : real(dp),intent(inout), pointer,optional :: cwavef_tr(:,:)
647 : : real(dp),contiguous, optional,intent(inout) :: vectproj(:,:,:)
648 : : type(gs_hamiltonian_type),intent(in) :: hamk
649 : : type(mpi_type),intent(in) :: mpi_enreg
650 : : type(pawcprj_type),intent(inout) :: cwaveprj(:,:)
651 : :
652 : : !Local variables-------------------------------
653 : : !scalars
654 : : integer :: bandpp,ier,ikpt_this_proc,my_nspinor,ndatarecv,nproc_band,npw,nspinortot
655 : : integer :: spaceComm=0,tim_nonlop
656 : : logical :: do_transpose,transfer_cwavef
657 : : integer :: l_gpu_option
658 : : #ifdef HAVE_OPENMP_OFFLOAD
659 : : integer :: ipw,ind,iibandpp
660 : : #endif
661 : : !character(len=500) :: msg
662 : : !arrays
663 : 264942 : integer, allocatable :: index_wavef_band(:)
664 : 264942 : integer, allocatable :: rdisplsloc(:),recvcountsloc(:),sdisplsloc(:),sendcountsloc(:)
665 : 264942 : integer,contiguous, pointer :: rdispls(:),recvcounts(:),sdispls(:),sendcounts(:)
666 : 529884 : real(dp) :: lambda_nonlop(mpi_enreg%bandpp)
667 : : real(dp) :: tsec(2)
668 : :
669 : : #if defined HAVE_GPU && defined HAVE_YAKL
670 : : real(c_double), contiguous, pointer :: cwavef_alltoall2(:,:) => null()
671 : : real(c_double), contiguous, pointer :: gvnlc_alltoall2(:,:) => null()
672 : : real(c_double), contiguous, pointer :: gsc_alltoall2(:,:) => null()
673 : : integer(kind=C_SIZE_T) :: cwavef_alltoall2_size
674 : : integer(kind=C_SIZE_T) :: gvnlc_alltoall2_size
675 : : integer(kind=C_SIZE_T) :: gsc_alltoall2_size
676 : : #else
677 : 264942 : real(dp), pointer :: cwavef_alltoall2(:,:)
678 : 264942 : real(dp), allocatable :: gvnlc_alltoall2(:,:)
679 : 264942 : real(dp), allocatable :: gsc_alltoall2(:,:)
680 : : #endif
681 : :
682 : 264942 : real(dp), allocatable :: cwavef_alltoall1(:,:)
683 : 264942 : real(dp), allocatable :: gvnlc_alltoall1(:,:)
684 : 264942 : real(dp), allocatable :: gsc_alltoall1(:,:)
685 : 264942 : real(dp), allocatable :: enlout(:)
686 : :
687 : : #if defined HAVE_GPU && defined HAVE_YAKL
688 : : ! this buffer is necessary to avoid mixing "managed memory" buffer with "regular memory" buffer in MPI calls
689 : : ! Just to be clear:
690 : : ! - managed memory means memory allocated using ABI_MALLOC_MANAGED
691 : : ! - regular memory means memory allocated using either ABI_MALLOC or ABI_MALLOC_CUDA
692 : : !
693 : : ! here we chose to use a GPU buffer, would it be better to use a CPU buffer ? To be checked.
694 : : !real(dp), allocatable :: cwavef_mpi(:,:)
695 : : type(c_ptr) :: cwavef_mpi_c
696 : : real(c_double),pointer :: cwavef_mpi(:,:)
697 : : #endif
698 : :
699 : : ! *************************************************************************
700 : :
701 : : DBG_ENTER('COLL')
702 : :
703 : 264942 : call timab(570,1,tsec)
704 : :
705 : 264942 : do_transpose = .true.
706 : 264942 : bandpp = mpi_enreg%bandpp
707 [ + - ]: 264942 : if(present(already_transposed)) then
708 [ + + ]: 264942 : if(already_transposed) then
709 : 124920 : do_transpose = .false.
710 : 124920 : bandpp = blocksize
711 : : end if
712 : : end if
713 : :
714 : 264942 : l_gpu_option=ABI_GPU_DISABLED
715 [ + + ]: 264942 : if (present(gpu_option)) then
716 : 132454 : l_gpu_option = gpu_option
717 : : end if
718 : :
719 : 264942 : nproc_band = mpi_enreg%nproc_band
720 : :
721 : 264942 : spaceComm=mpi_enreg%comm_fft
722 [ + - ]: 264942 : if(mpi_enreg%paral_kgb==1) spaceComm=mpi_enreg%comm_band
723 : 264942 : my_nspinor=max(1,hamk%nspinor/mpi_enreg%nproc_spinor)
724 : 264942 : nspinortot=hamk%nspinor
725 : :
726 : : !Check sizes
727 [ + + ]: 264942 : npw=hamk%npw_k;if (.not.do_transpose) npw=hamk%npw_fft_k
728 [ + + - + ]: 794826 : if (size(cwavef)/=2*npw*my_nspinor*blocksize) then
729 : 0 : ABI_BUG('Incorrect size for cwavef!')
730 : : end if
731 [ + + + + ]: 264942 : if(choice/=0.and.signs==2) then
732 [ - + ]: 62460 : if (paw_opt/=3) then
733 [ # # # # ]: 0 : if (size(gvnlc)/=2*npw*my_nspinor*blocksize) then
734 : 0 : ABI_BUG('Incorrect size for gvnlc!')
735 : : end if
736 : : end if
737 [ + - ]: 62460 : if(paw_opt>=3) then
738 [ + + - + ]: 187380 : if (size(gsc)/=2*npw*my_nspinor*blocksize) then
739 : 0 : ABI_BUG('Incorrect size for gsc!')
740 : : end if
741 : : end if
742 : : end if
743 [ + + + + ]: 264942 : if(cpopt>=0.and. .not. present(vectproj)) then
744 [ + + - + ]: 226704 : if (size(cwaveprj)/=hamk%natom*my_nspinor*bandpp) then
745 : 0 : ABI_BUG('Incorrect size for cwaveprj!')
746 : : end if
747 : : end if
748 : :
749 : 264942 : transfer_cwavef=.false.
750 : : #ifdef HAVE_OPENMP_OFFLOAD
751 : : if(.not. xomp_target_is_present(c_loc(cwavef)) .and. hamk%gpu_option==ABI_GPU_OPENMP) then
752 : : transfer_cwavef=.true.
753 : : end if
754 : : !$OMP TARGET ENTER DATA MAP(to:cwavef) IF(transfer_cwavef)
755 : : #endif
756 : :
757 [ - + + - : 794826 : ABI_MALLOC(sendcountsloc,(nproc_band))
+ - + - -
+ - + ]
758 [ - + + - : 529884 : ABI_MALLOC(sdisplsloc ,(nproc_band))
- + - + ]
759 [ - + + - : 529884 : ABI_MALLOC(recvcountsloc,(nproc_band))
- + - + ]
760 [ - + + - : 529884 : ABI_MALLOC(rdisplsloc ,(nproc_band))
- + - + ]
761 : :
762 : 264942 : ikpt_this_proc=bandfft_kpt_get_ikpt()
763 : :
764 : 264942 : recvcounts => bandfft_kpt(ikpt_this_proc)%recvcounts(:)
765 : 264942 : sendcounts => bandfft_kpt(ikpt_this_proc)%sendcounts(:)
766 : 264942 : rdispls => bandfft_kpt(ikpt_this_proc)%rdispls (:)
767 : 264942 : sdispls => bandfft_kpt(ikpt_this_proc)%sdispls (:)
768 : 264942 : ndatarecv = bandfft_kpt(ikpt_this_proc)%ndatarecv
769 : :
770 [ + - ]: 264942 : if(hamk%gpu_option==ABI_GPU_KOKKOS) then
771 : : #if defined HAVE_GPU && defined HAVE_YAKL
772 : : ABI_MALLOC_MANAGED(cwavef_alltoall2, (/2,ndatarecv*my_nspinor*bandpp/))
773 : : cwavef_alltoall2_size = 2*ndatarecv*my_nspinor*bandpp*dp
774 : :
775 : : if (paw_opt >= 0 .and. paw_opt < 3) then
776 : : gsc_alltoall2 => null()
777 : : gsc_alltoall2_size = 0
778 : : else
779 : : ABI_MALLOC_MANAGED(gsc_alltoall2, (/2,ndatarecv*my_nspinor*(paw_opt/3)*bandpp/))
780 : : gsc_alltoall2_size = 2*ndatarecv*my_nspinor*(paw_opt/3)*bandpp*dp
781 : : endif
782 : :
783 : : ABI_MALLOC_MANAGED(gvnlc_alltoall2, (/2,ndatarecv*my_nspinor*bandpp/))
784 : : gvnlc_alltoall2_size = 2*ndatarecv*my_nspinor*bandpp*dp
785 : : #endif
786 : : else
787 [ + + ]: 264942 : if(present(cwavef_tr)) then
788 [ + + + - : 192 : if(any(shape(cwavef_tr)/=(/2,ndatarecv*my_nspinor*bandpp/))) then
- + ]
789 : 0 : ABI_BUG('cwavef_tr does not have the correct shape')
790 : : endif
791 : 64 : cwavef_alltoall2 => cwavef_tr
792 : : else
793 [ - + + - : 794634 : ABI_MALLOC(cwavef_alltoall2, (2,ndatarecv*my_nspinor*bandpp))
+ - + - -
+ - + ]
794 : : endif
795 [ + + + - : 794826 : ABI_MALLOC(gsc_alltoall2, (2,ndatarecv*my_nspinor*(paw_opt/3)*bandpp))
+ - + + -
+ - + -
+ ]
796 [ - + + - : 794826 : ABI_MALLOC(gvnlc_alltoall2, (2,ndatarecv*my_nspinor*bandpp))
+ - + - -
+ - + -
+ ]
797 : : end if
798 : :
799 : : #ifdef HAVE_OPENMP_OFFLOAD
800 : : !$OMP TARGET ENTER DATA MAP(alloc:cwavef_alltoall2) IF(hamk%gpu_option==ABI_GPU_OPENMP)
801 : : #endif
802 : :
803 [ + + + + : 264942 : if(do_transpose .and. (bandpp/=1 .or. (bandpp==1 .and. mpi_enreg%paral_spinor==0.and.nspinortot==2)))then
+ + + + ]
804 [ - + + - : 62286 : ABI_MALLOC(cwavef_alltoall1,(2,ndatarecv*my_nspinor*bandpp))
+ - + - -
+ - + ]
805 : : #ifdef HAVE_OPENMP_OFFLOAD
806 : : !$OMP TARGET ENTER DATA MAP(alloc:cwavef_alltoall1) IF(hamk%gpu_option==ABI_GPU_OPENMP)
807 : : #endif
808 [ - + ]: 20762 : if(signs==2)then
809 [ # # ]: 0 : if (paw_opt/=3) then
810 [ # # # # : 0 : ABI_MALLOC(gvnlc_alltoall1,(2,ndatarecv*my_nspinor*bandpp))
# # # # ]
811 : : end if
812 [ # # ]: 0 : if (paw_opt==3.or.paw_opt==4) then
813 [ # # # # : 0 : ABI_MALLOC(gsc_alltoall1,(2,ndatarecv*my_nspinor*bandpp))
# # # # ]
814 : : end if
815 : : end if
816 : : end if
817 : :
818 [ + + + - : 794826 : ABI_MALLOC(enlout,(nnlout*bandpp))
+ - + + -
+ - + -
+ ]
819 [ + + ]: 1348582 : enlout = zero
820 : :
821 [ + + ]: 923900 : recvcountsloc(:)=recvcounts(:)*2*my_nspinor*bandpp
822 [ + + ]: 923900 : rdisplsloc(:)=rdispls(:)*2*my_nspinor*bandpp
823 [ + + ]: 923900 : sendcountsloc(:)=sendcounts(:)*2*my_nspinor
824 [ + + ]: 923900 : sdisplsloc(:)=sdispls(:)*2*my_nspinor
825 : :
826 [ + + ]: 264942 : if(do_transpose) then
827 : 140022 : call timab(581,1,tsec)
828 [ + + + + : 140022 : if (bandpp/=1 .or. (bandpp==1 .and. mpi_enreg%paral_spinor==0.and.nspinortot==2)) then
+ + ]
829 [ + - ]: 20762 : if (l_gpu_option == ABI_GPU_KOKKOS) then
830 : : #if defined HAVE_GPU && defined HAVE_YAKL
831 : : ABI_MALLOC_CUDA(cwavef_mpi_c, INT(2, c_size_t) * npw * my_nspinor * blocksize * dp)
832 : : call c_f_pointer(cwavef_mpi_c, cwavef_mpi, (/2, npw * my_nspinor * blocksize/))
833 : :
834 : : ! use cwavef_mpi instead of cwavef (don't use managed memory in MPI calls)
835 : : call copy_gpu_to_gpu(cwavef_mpi_c, C_LOC(cwavef), INT(2, c_size_t) * npw * my_nspinor * blocksize * dp)
836 : :
837 : : call xmpi_alltoallv(cwavef_mpi,sendcountsloc,sdisplsloc,cwavef_alltoall1,&
838 : : & recvcountsloc,rdisplsloc,spaceComm,ier)
839 : :
840 : : ABI_FREE_CUDA(cwavef_mpi_c)
841 : : #endif
842 : : else
843 : : call xmpi_alltoallv(cwavef,sendcountsloc,sdisplsloc,cwavef_alltoall1,&
844 : 20762 : & recvcountsloc,rdisplsloc,spaceComm,ier,use_omp_map=(hamk%gpu_option==ABI_GPU_OPENMP))
845 : : end if
846 : :
847 : :
848 : : else
849 : : call xmpi_alltoallv(cwavef,sendcountsloc,sdisplsloc,cwavef_alltoall2,&
850 : 119260 : & recvcountsloc,rdisplsloc,spaceComm,ier,use_omp_map=(hamk%gpu_option==ABI_GPU_OPENMP))
851 : : end if
852 : 140022 : call timab(581,2,tsec)
853 : : else
854 : : ! Here, we cheat, and use DCOPY to bypass some compiler's overzealous bound-checking
855 : : ! (ndatarecv*my_nspinor*bandpp might be greater than the declared size of cwavef)
856 [ + - ]: 124920 : if (l_gpu_option == ABI_GPU_KOKKOS) then
857 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS)
858 : : call copy_gpu_to_gpu(C_LOC(cwavef_alltoall2), C_LOC(cwavef), INT(2, c_size_t) * ndatarecv * my_nspinor * bandpp * dp)
859 : : #endif
860 [ - + ]: 124920 : else if (hamk%gpu_option == ABI_GPU_OPENMP) then
861 [ # # # # : 0 : call gpu_copy(cwavef_alltoall2, cwavef, int(2,c_size_t)*ndatarecv*my_nspinor*bandpp)
# # # # #
# # # ]
862 : : else
863 [ - + - - : 124920 : call DCOPY(2*ndatarecv*my_nspinor*bandpp,cwavef,1,cwavef_alltoall2,1)
- - - - -
+ - - - -
- + ]
864 : : end if
865 : : end if
866 : :
867 : : #ifdef HAVE_OPENMP_OFFLOAD
868 : : !$OMP TARGET EXIT DATA MAP(delete:cwavef) IF(transfer_cwavef)
869 : : #endif
870 : :
871 : : !=====================================================================
872 [ + + ]: 264942 : if (bandpp==1) then
873 : :
874 : :
875 [ + + + + : 127588 : if (do_transpose .and. mpi_enreg%paral_spinor==0.and.nspinortot==2) then !Sort WF by spin
+ + ]
876 : 7368 : call prep_sort_wavef_spin(nproc_band,my_nspinor,ndatarecv,recvcounts,rdispls,index_wavef_band)
877 [ + - ]: 7368 : if(hamk%gpu_option==ABI_GPU_OPENMP) then
878 : : #ifdef HAVE_OPENMP_OFFLOAD
879 : : !$OMP TARGET TEAMS DISTRIBUTE &
880 : : !$OMP& PRIVATE(iibandpp) MAP(to:cwavef_alltoall1,cwavef_alltoall2,index_wavef_band)
881 : : do iibandpp=1,bandpp*my_nspinor
882 : : !$OMP PARALLEL DO PRIVATE(ind,ipw)
883 : : do ipw = 1 ,ndatarecv
884 : : ind=index_wavef_band(ipw + ndatarecv*(iibandpp-1))
885 : : cwavef_alltoall2(1,ipw + ndatarecv*(iibandpp-1)) = cwavef_alltoall1(1,ind)
886 : : cwavef_alltoall2(2,ipw + ndatarecv*(iibandpp-1)) = cwavef_alltoall1(2,ind)
887 : : end do
888 : : end do
889 : : #endif
890 : : else
891 [ + + + + ]: 6364632 : cwavef_alltoall2(:, :) = cwavef_alltoall1(:,index_wavef_band)
892 : : end if
893 : : end if
894 : :
895 [ + + ]: 127588 : if (paw_opt==2) then
896 : 69568 : lambda_nonlop(1)=lambdablock(mpi_enreg%me_band+1)
897 : : end if
898 : : call nonlop(choice,cpopt,cwaveprj,enlout,hamk,idir,lambda_nonlop,mpi_enreg,1,nnlout,paw_opt,&
899 [ + + + + ]: 254216 : & signs,gsc_alltoall2,tim_nonlop,cwavef_alltoall2,gvnlc_alltoall2,vectproj=vectproj)
900 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_YAKL)
901 : : !call gpu_device_synchronize()
902 : : !call gpu_data_prefetch_async_f(C_LOC(cwavef_alltoall2), cwavef_alltoall2_size, CPU_DEVICE_ID)
903 : : !call gpu_data_prefetch_async_f(C_LOC(gvnlc_alltoall2), gvnlc_alltoall2_size, CPU_DEVICE_ID)
904 : : !call gpu_data_prefetch_async_f(C_LOC(gsc_alltoall2), gsc_alltoall2_size, CPU_DEVICE_ID)
905 : : #endif
906 : :
907 [ + + + + : 127588 : if (do_transpose .and. mpi_enreg%paral_spinor == 0 .and. nspinortot==2.and.signs==2) then
+ + - + ]
908 [ # # # # : 0 : if (paw_opt/=3) gvnlc_alltoall1(:,index_wavef_band)=gvnlc_alltoall2(:,:)
# # ]
909 [ # # # # : 0 : if (paw_opt==3.or.paw_opt==4) gsc_alltoall1(:,index_wavef_band)=gsc_alltoall2(:,:)
# # ]
910 : : end if
911 : :
912 : : else ! bandpp/=1
913 : :
914 : : ! -------------------------------------------------------------
915 : : ! Computation of the index used to sort the waves functions below bandpp
916 : : ! -------------------------------------------------------------
917 [ + + ]: 137354 : if(do_transpose) then
918 : : call prep_index_wavef_bandpp(nproc_band,bandpp,&
919 : 13394 : & my_nspinor,ndatarecv,recvcounts,rdispls,index_wavef_band)
920 : :
921 : : ! -------------------------------------------------------
922 : : ! Sorting of the waves functions below bandpp
923 : : ! -------------------------------------------------------
924 [ + - ]: 13394 : if(hamk%gpu_option==ABI_GPU_OPENMP) then
925 : : #ifdef HAVE_OPENMP_OFFLOAD
926 : : !$OMP TARGET TEAMS DISTRIBUTE &
927 : : !$OMP& PRIVATE(iibandpp) MAP(to:cwavef_alltoall1,cwavef_alltoall2,index_wavef_band)
928 : : do iibandpp=1,bandpp*my_nspinor
929 : : !$OMP PARALLEL DO PRIVATE(ind,ipw)
930 : : do ipw = 1 ,ndatarecv
931 : : ind=index_wavef_band(ipw + ndatarecv*(iibandpp-1))
932 : : cwavef_alltoall2(1,ipw + ndatarecv*(iibandpp-1)) = cwavef_alltoall1(1,ind)
933 : : cwavef_alltoall2(2,ipw + ndatarecv*(iibandpp-1)) = cwavef_alltoall1(2,ind)
934 : : end do
935 : : end do
936 : : #endif
937 : : else
938 [ + + + + ]: 133129886 : cwavef_alltoall2(:,:) = cwavef_alltoall1(:,index_wavef_band)
939 : : end if
940 : : end if
941 : :
942 : : ! -------------------------------------------------------
943 : : ! Call nonlop
944 : : ! -------------------------------------------------------
945 [ + + ]: 137354 : if(paw_opt == 2) then
946 [ + + ]: 46604 : lambda_nonlop(1:bandpp) = lambdablock((mpi_enreg%me_band*bandpp)+1:((mpi_enreg%me_band+1)*bandpp))
947 : : end if
948 : : call nonlop(choice,cpopt,cwaveprj,enlout,hamk,idir,lambda_nonlop,mpi_enreg,bandpp,nnlout,paw_opt,&
949 [ + + + + ]: 150748 : & signs,gsc_alltoall2,tim_nonlop,cwavef_alltoall2,gvnlc_alltoall2,vectproj=vectproj)
950 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_YAKL)
951 : : if(hamk%gpu_option==ABI_GPU_KOKKOS) call gpu_device_synchronize()
952 : : !call gpu_data_prefetch_async_f(C_LOC(cwavef_alltoall2), cwavef_alltoall2_size, CPU_DEVICE_ID)
953 : : !call gpu_data_prefetch_async_f(C_LOC(gvnlc_alltoall2), gvnlc_alltoall2_size, CPU_DEVICE_ID)
954 : : !if (associated(gsc_alltoall2)) then
955 : : ! call gpu_data_prefetch_async_f(C_LOC(gsc_alltoall2), gsc_alltoall2_size, CPU_DEVICE_ID)
956 : : !end if
957 : : #endif
958 : :
959 : : ! -----------------------------------------------------
960 : : ! Sorting of waves functions below the processors
961 : : ! -----------------------------------------------------
962 [ + + - + ]: 137354 : if(do_transpose.and.signs==2) then
963 [ # # # # : 0 : if (paw_opt/=3) gvnlc_alltoall1(:,index_wavef_band)=gvnlc_alltoall2(:,:)
# # ]
964 [ # # # # : 0 : if (paw_opt==3.or.paw_opt==4) gsc_alltoall1(:,index_wavef_band)=gsc_alltoall2(:,:)
# # ]
965 : : end if
966 : :
967 : : end if
968 : :
969 : : !=====================================================================
970 : : ! -------------------------------------------------------
971 : : ! Deallocation
972 : : ! -------------------------------------------------------
973 [ + + ]: 264942 : if (allocated(index_wavef_band)) then
974 : 20762 : ABI_FREE(index_wavef_band)
975 : : end if
976 : :
977 : : !Transpose the gsc_alltoall or gvlnc_alltoall tabs
978 : : !according to the paw_opt and signs values
979 [ + + ]: 264942 : if(do_transpose) then
980 [ - + ]: 140022 : if (signs==2) then
981 : 0 : call timab(581,1,tsec)
982 [ # # # # : 0 : if(bandpp/=1 .or. (bandpp==1 .and. mpi_enreg%paral_spinor==0.and.nspinortot==2))then
# # ]
983 [ # # ]: 0 : if (paw_opt/=3) then
984 : : call xmpi_alltoallv(gvnlc_alltoall1,recvcountsloc,rdisplsloc,gvnlc,&
985 : 0 : & sendcountsloc,sdisplsloc,spaceComm,ier)
986 : : end if
987 [ # # ]: 0 : if (paw_opt==3.or.paw_opt==4) then
988 : : call xmpi_alltoallv(gsc_alltoall1,recvcountsloc,rdisplsloc,gsc,&
989 : 0 : & sendcountsloc,sdisplsloc,spaceComm,ier)
990 : : end if
991 : : else
992 [ # # ]: 0 : if (paw_opt/=3) then
993 : : call xmpi_alltoallv(gvnlc_alltoall2,recvcountsloc,rdisplsloc,gvnlc,&
994 : 0 : & sendcountsloc,sdisplsloc,spaceComm,ier)
995 : : end if
996 [ # # ]: 0 : if (paw_opt==3.or.paw_opt==4) then
997 : : call xmpi_alltoallv(gsc_alltoall2,recvcountsloc,rdisplsloc,gsc,&
998 : 0 : & sendcountsloc,sdisplsloc,spaceComm,ier)
999 : : end if
1000 : : end if
1001 : 0 : call timab(581,2,tsec)
1002 : : end if
1003 : : else
1004 : : ! TODO check other usages, maybe
1005 [ + - ]: 124920 : if (l_gpu_option == ABI_GPU_KOKKOS) then
1006 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS)
1007 : : call copy_gpu_to_gpu(C_LOC(gsc), C_LOC(gsc_alltoall2), INT(2, c_size_t) * ndatarecv * my_nspinor * bandpp * dp)
1008 : : #endif
1009 : : else
1010 [ - + - - : 124920 : call DCOPY(2*ndatarecv*my_nspinor*bandpp, gsc_alltoall2, 1, gsc, 1)
- - - - -
+ - - -
- ]
1011 : : end if
1012 : : end if
1013 : :
1014 [ + + ]: 264942 : if (nnlout>0) then
1015 : 107492 : call xmpi_allgather(enlout,nnlout*bandpp,enlout_block,spaceComm,ier)
1016 : : end if
1017 [ - + ]: 264942 : ABI_FREE(enlout)
1018 [ - + ]: 264942 : ABI_FREE(sendcountsloc)
1019 [ - + ]: 264942 : ABI_FREE(sdisplsloc)
1020 [ - + ]: 264942 : ABI_FREE(recvcountsloc)
1021 [ - + ]: 264942 : ABI_FREE(rdisplsloc)
1022 : :
1023 : : #ifdef HAVE_OPENMP_OFFLOAD
1024 : : if(present(cwavef_tr)) then
1025 : : !$OMP TARGET UPDATE FROM(cwavef_alltoall2) IF(hamk%gpu_option==ABI_GPU_OPENMP)
1026 : : end if
1027 : : !$OMP TARGET EXIT DATA MAP(delete:cwavef_alltoall2) IF(hamk%gpu_option==ABI_GPU_OPENMP)
1028 : : #endif
1029 : :
1030 [ + - ]: 264942 : if(hamk%gpu_option==ABI_GPU_KOKKOS) then
1031 : : #if defined HAVE_GPU && defined HAVE_YAKL
1032 : : ABI_FREE_MANAGED(cwavef_alltoall2)
1033 : : ABI_FREE_MANAGED(gvnlc_alltoall2)
1034 : : if (paw_opt >= 3) then
1035 : : ABI_FREE_MANAGED(gsc_alltoall2)
1036 : : end if
1037 : : #endif
1038 : : else
1039 [ + + ]: 264942 : if(present(cwavef_tr)) then
1040 : 64 : nullify(cwavef_alltoall2)
1041 : : else
1042 [ - + ]: 264878 : ABI_FREE(cwavef_alltoall2)
1043 : : endif
1044 [ - + ]: 264942 : ABI_FREE(gvnlc_alltoall2)
1045 [ - + ]: 264942 : ABI_FREE(gsc_alltoall2)
1046 : : end if
1047 : :
1048 [ + + + + : 264942 : if(do_transpose .and. (bandpp/=1 .or. (bandpp==1 .and. mpi_enreg%paral_spinor==0.and.nspinortot==2)))then
+ + + + ]
1049 : : #ifdef HAVE_OPENMP_OFFLOAD
1050 : : !$OMP TARGET EXIT DATA MAP(delete:cwavef_alltoall1) IF(hamk%gpu_option==ABI_GPU_OPENMP)
1051 : : #endif
1052 [ - + ]: 20762 : ABI_FREE(cwavef_alltoall1)
1053 [ - + ]: 20762 : if(signs==2)then
1054 [ # # ]: 0 : if (paw_opt/=3) then
1055 [ # # ]: 0 : ABI_FREE(gvnlc_alltoall1)
1056 : : end if
1057 [ # # ]: 0 : if (paw_opt==3.or.paw_opt==4) then
1058 [ # # ]: 0 : ABI_FREE(gsc_alltoall1)
1059 : : end if
1060 : : end if
1061 : : end if
1062 : :
1063 : 264942 : call timab(570,2,tsec)
1064 : :
1065 : : DBG_EXIT('COLL')
1066 : :
1067 [ - + - + : 529884 : end subroutine prep_nonlop
- + - + -
+ - + -
+ ]
1068 : : !!***
1069 : :
1070 : : !!****f* ABINIT/prep_fourwf
1071 : : !! NAME
1072 : : !! prep_fourwf
1073 : : !!
1074 : : !! FUNCTION
1075 : : !! this routine prepares the data to the call of fourwf.
1076 : : !!
1077 : : !! INPUTS
1078 : : !! blocksize= size of block for FFT
1079 : : !! cwavef(2,npw*ndat)=planewave coefficients of wavefunction (one spinorial component?).
1080 : : !! dtfil <type(datafiles_type)>=variables related to files
1081 : : !! kg_k(3,npw_k)=reduced planewave coordinates.
1082 : : !! lmnmax=if useylm=1, max number of (l,m,n) comp. over all type of psps
1083 : : !! =if useylm=0, max number of (l,n) comp. over all type of psps
1084 : : !! mgfft=maximum size of 1d ffts
1085 : : !! mpi_enreg=information about mpi parallelization
1086 : : !! mpsang= 1+maximum angular momentum for nonlocal pseudopotentials
1087 : : !! mpssoang= 1+maximum (spin*angular momentum) for nonlocal pseudopotentials
1088 : : !! natom=number of atoms in cell.
1089 : : !! nband_k=number of bands at this k point for that spin polarization
1090 : : !! ndat=number of FFT to do in //
1091 : : !! ngfft(18)= contain all needed information about 3D FFT
1092 : : !! npw_k=number of plane waves at this k point
1093 : : !! nspinor=number of spinorial components of the wavefunctions
1094 : : !! ntypat=number of types of atoms in unit cell.
1095 : : !! n4,n5,n6 used for dimensionning of vlocal
1096 : : !! option_fourwf=option for fourwf (see fourwf.F90)
1097 : : !! prtvol=control print volume and debugging output
1098 : : !! ucvol=unit cell volume
1099 : : !! nfft_blocks=number of blocks fourwf is split into.
1100 : : !! [bandfft_kpt_tab]= (optional) if present, contains tabs used to implement
1101 : : !! the "band-fft" parallelism
1102 : : !! if not present, the bandfft_kpt global variable is used
1103 : : !! [gpu_option] = GPU implementation to use, i.e. cuda, openMP, ... (0=not using GPU)
1104 : : !!
1105 : : !! OUTPUT
1106 : : !! gwavef=(2,npw*ndat)=matrix elements <G|H|C>.
1107 : : !!
1108 : : !! SIDE EFFECTS
1109 : : !!
1110 : : !! SOURCE
1111 : :
1112 [ - + ]: 182282 : subroutine prep_fourwf(rhoaug,blocksize,cwavef,wfraug,iblock,istwf_k,mgfft,&
1113 : 182282 : & mpi_enreg,nband_k,ndat,ngfft,npw_k,n4,n5,n6,occ_k,option_fourwf,ucvol,wtk,&
1114 : : & nfft_blocks,&
1115 : : & bandfft_kpt_tab,gpu_option) ! Optional arguments
1116 : :
1117 : : !Arguments ------------------------------------
1118 : : !scalars
1119 : : integer,intent(in) :: blocksize,iblock,istwf_k,mgfft,n4,n5,n6,nband_k,ndat,npw_k,nfft_blocks
1120 : : integer,intent(in) :: option_fourwf
1121 : : integer,intent(in),optional :: gpu_option
1122 : : real(dp),intent(in) :: ucvol,wtk
1123 : : type(bandfft_kpt_type),optional,target,intent(in) :: bandfft_kpt_tab
1124 : : type(mpi_type),intent(in) :: mpi_enreg
1125 : : !arrays
1126 : : integer,intent(in) :: ngfft(18)
1127 : : real(dp),intent(in) :: occ_k(nband_k)
1128 : : real(dp),intent(out) :: rhoaug(n4,n5,n6)
1129 : : real(dp),intent(in), target :: cwavef(2,npw_k*blocksize)
1130 : : real(dp),target,intent(inout) :: wfraug(:,:,:,:) !(2,n4,n5,n6*{ndat,ndat/nfft_blocks+mod(ndat,nfft_blocks))
1131 : :
1132 : : !Local variables-------------------------------
1133 : : !scalars
1134 : : integer :: bandpp,bandpp_sym,ier,iibandpp,ikpt_this_proc,ind_occ,ind_occ1,ind_occ2,ipw
1135 : : integer :: istwf_k_,jjbandpp,me_fft,nd3,nproc_band,nproc_fft,npw_fft
1136 : : integer :: nfft_blocks_sym,nband_fftblock
1137 : : integer :: spaceComm=0,tim_fourwf,gpu_option_
1138 : : integer,pointer :: idatarecv0,ndatarecv,ndatarecv_tot,ndatasend_sym
1139 : : logical :: flag_inv_sym,have_to_reequilibrate,transfer_cwavef
1140 : : real(dp) :: weight,weight1,weight2
1141 : : type(bandfft_kpt_type),pointer :: bandfft_kpt_ptr
1142 : : !arrays
1143 : 182282 : integer,contiguous, pointer :: indices_pw_fft(:),kg_k_fft(:,:),kg_k_gather(:,:),kg_k_gather_sym(:,:)
1144 : 182282 : integer,contiguous, pointer :: rdispls(:),rdispls_sym(:)
1145 : 182282 : integer,contiguous, pointer :: recvcounts(:),recvcount_fft(:),recvcounts_sym(:),recvcounts_sym_tot(:)
1146 : 182282 : integer,contiguous, pointer :: recvdisp_fft(:),sdispls(:),sdispls_sym(:)
1147 : 182282 : integer,contiguous, pointer :: sendcounts(:),sendcount_fft(:),sendcounts_sym(:),sendcounts_sym_all(:)
1148 : 182282 : integer,contiguous, pointer :: senddisp_fft(:),tab_proc(:)
1149 : 182282 : integer,allocatable :: rdisplsloc(:)
1150 : 182282 : integer,allocatable :: recvcountsloc(:),sdisplsloc(:)
1151 : 182282 : integer,allocatable :: sendcountsloc(:)
1152 : 182282 : integer,allocatable :: index_wavef_band(:),index_wavef_send(:)
1153 : 182282 : integer,pointer :: gbound_(:,:)
1154 : : real(dp) :: dummy(2,1),tsec(2)
1155 : 182282 : real(dp),allocatable :: buff_wf(:,:)
1156 : :
1157 : : #if defined HAVE_GPU && defined HAVE_YAKL
1158 : : real(c_double), contiguous, pointer :: cwavef_alltoall1(:,:) => null()
1159 : : #else
1160 : 182282 : real(dp),allocatable :: cwavef_alltoall1(:,:)
1161 : : #endif
1162 : 182282 : real(dp),allocatable :: cwavef_alltoall2(:,:)
1163 : 182282 : real(dp),allocatable :: cwavef_fft(:,:), cwavef_fft_tr(:,:)
1164 : 182282 : real(dp),allocatable :: weight_t(:),weight1_t(:),weight2_t(:)
1165 : 182282 : real(dp),pointer :: ewavef_alltoall_sym(:,:),ewavef_alltoall_sym_work(:,:),wfraug_ptr(:,:,:,:)
1166 : :
1167 : : #if defined HAVE_GPU && defined HAVE_YAKL
1168 : : ! this buffer is necessary to avoid mixing "managed memory" buffer with "regular memory" buffer in MPI calls
1169 : : ! Just to be clear:
1170 : : ! - managed memory means memory allocated using ABI_MALLOC_MANAGED
1171 : : ! - regular memory means memory allocated using either ABI_MALLOC or ABI_MALLOC_CUDA
1172 : : !
1173 : : ! here we chose to use a CPU buffer, would it be better to use a GPU buffer ? To be checked.
1174 : : real(dp), allocatable :: cwavef_mpi(:,:)
1175 : : !type(c_ptr) :: cwavef_mpi_c
1176 : : !real(c_double),pointer :: cwavef_mpi(:,:)
1177 : : #endif
1178 : :
1179 : : ! *************************************************************************
1180 : :
1181 : 0 : ABI_CHECK((option_fourwf/=3),'Option=3 (FFT r->g) not implemented')
1182 [ - + ]: 182282 : ABI_CHECK((nfft_blocks>0),'nfft_blocks is null')
1183 [ - + ]: 182282 : ABI_CHECK((mpi_enreg%bandpp==ndat),'BUG: bandpp/=ndat')
1184 : :
1185 : 182282 : spaceComm=mpi_enreg%comm_band
1186 : 182282 : nproc_band = mpi_enreg%nproc_band
1187 : 182282 : nproc_fft = mpi_enreg%nproc_fft
1188 : 182282 : bandpp = mpi_enreg%bandpp
1189 : 182282 : me_fft = mpi_enreg%me_fft
1190 : :
1191 [ + - ]: 182282 : gpu_option_=ABI_GPU_DISABLED;if (present(gpu_option)) gpu_option_=gpu_option
1192 : :
1193 [ + - ]: 182282 : if (present(bandfft_kpt_tab)) then
1194 : : bandfft_kpt_ptr => bandfft_kpt_tab
1195 : : else
1196 : 182282 : ikpt_this_proc=bandfft_kpt_get_ikpt()
1197 : 182282 : bandfft_kpt_ptr => bandfft_kpt(ikpt_this_proc)
1198 : : end if
1199 : :
1200 : 182282 : have_to_reequilibrate = bandfft_kpt_ptr%have_to_reequilibrate
1201 : 182282 : istwf_k_=istwf_k
1202 [ + + + + : 182282 : flag_inv_sym = (istwf_k_==2 .and. any(ngfft(7) == [401,402,312,512]))
+ - ]
1203 [ - + ]: 182282 : if (option_fourwf==0) flag_inv_sym=((flag_inv_sym).and.(gpu_option_==ABI_GPU_DISABLED))
1204 : :
1205 [ + + ]: 182282 : if (flag_inv_sym) then
1206 : 6622 : istwf_k_ = 1
1207 [ + + ]: 6622 : if (modulo(bandpp,2)==0) then
1208 : 2494 : bandpp_sym = bandpp/2
1209 [ + - ]: 2494 : nfft_blocks_sym = nfft_blocks/2; if(modulo(nfft_blocks,2)/=0) nfft_blocks_sym=nfft_blocks_sym+1
1210 : : else
1211 : 4128 : bandpp_sym = bandpp
1212 : 4128 : nfft_blocks_sym = nfft_blocks
1213 : : end if
1214 : : end if
1215 : :
1216 [ - + ]: 182282 : if(have_to_reequilibrate .and. gpu_option_==ABI_GPU_OPENMP) then
1217 : 0 : ABI_BUG("Reequilibrating FFT isn't supported with OpenMP GPU yet !")
1218 : : end if
1219 : :
1220 : : !====================================================================================
1221 [ - + + - : 546846 : ABI_MALLOC(sendcountsloc,(nproc_band))
+ - + - -
+ - + ]
1222 [ - + + - : 364564 : ABI_MALLOC(sdisplsloc ,(nproc_band))
- + - + ]
1223 [ - + + - : 364564 : ABI_MALLOC(recvcountsloc,(nproc_band))
- + - + ]
1224 [ - + + - : 364564 : ABI_MALLOC(rdisplsloc ,(nproc_band))
- + - + ]
1225 : :
1226 : 182282 : recvcounts =>bandfft_kpt_ptr%recvcounts(:)
1227 : 182282 : sendcounts =>bandfft_kpt_ptr%sendcounts(:)
1228 : 182282 : rdispls =>bandfft_kpt_ptr%rdispls (:)
1229 : 182282 : sdispls =>bandfft_kpt_ptr%sdispls (:)
1230 : 182282 : ndatarecv =>bandfft_kpt_ptr%ndatarecv
1231 : :
1232 : 182282 : kg_k_gather =>bandfft_kpt_ptr%kg_k_gather(:,:)
1233 : 182282 : gbound_ =>bandfft_kpt_ptr%gbound(:,:)
1234 : :
1235 [ + + ]: 182282 : if (flag_inv_sym ) then
1236 : 6622 : idatarecv0 =>bandfft_kpt_ptr%idatarecv0
1237 : 6622 : ndatarecv_tot =>bandfft_kpt_ptr%ndatarecv_tot
1238 : 6622 : ndatasend_sym =>bandfft_kpt_ptr%ndatasend_sym
1239 : 6622 : kg_k_gather_sym =>bandfft_kpt_ptr%kg_k_gather_sym(:,:)
1240 : 6622 : rdispls_sym =>bandfft_kpt_ptr%rdispls_sym(:)
1241 : 6622 : recvcounts_sym =>bandfft_kpt_ptr%recvcounts_sym(:)
1242 : 6622 : recvcounts_sym_tot =>bandfft_kpt_ptr%recvcounts_sym_tot(:)
1243 : 6622 : sdispls_sym =>bandfft_kpt_ptr%sdispls_sym(:)
1244 : 6622 : sendcounts_sym =>bandfft_kpt_ptr%sendcounts_sym(:)
1245 : 6622 : sendcounts_sym_all =>bandfft_kpt_ptr%sendcounts_sym_all(:)
1246 : 6622 : tab_proc =>bandfft_kpt_ptr%tab_proc(:)
1247 : : end if
1248 : :
1249 [ - + + - : 546846 : ABI_MALLOC(cwavef_alltoall2,(2,ndatarecv*bandpp))
+ - + - -
+ - + ]
1250 [ + + + + ]: 182282 : if ( ((.not.flag_inv_sym) .and. (bandpp>1) ) .or. flag_inv_sym )then
1251 [ + - ]: 37386 : if(gpu_option_==ABI_GPU_KOKKOS) then
1252 : : #if defined HAVE_GPU && defined HAVE_YAKL
1253 : : ABI_MALLOC_MANAGED(cwavef_alltoall1,(/2,ndatarecv*bandpp/))
1254 : : #endif
1255 : : else
1256 [ - + + - : 219668 : ABI_MALLOC(cwavef_alltoall1,(2,ndatarecv*bandpp))
- + - + ]
1257 : : end if
1258 : : end if
1259 : :
1260 [ + + ]: 599476 : recvcountsloc(:)=recvcounts(:)*2*bandpp
1261 [ + + ]: 599476 : rdisplsloc(:)=rdispls(:)*2*bandpp
1262 [ + + ]: 599476 : sendcountsloc(:)=sendcounts(:)*2
1263 [ + + ]: 599476 : sdisplsloc(:)=sdispls(:)*2
1264 : :
1265 : 182282 : transfer_cwavef=.false.
1266 : 182282 : if(.not. xomp_target_is_present(c_loc(cwavef)) .and. gpu_option_==ABI_GPU_OPENMP) then
1267 : 182282 : transfer_cwavef=.true.
1268 : : end if
1269 : :
1270 : : #ifdef HAVE_OPENMP_OFFLOAD
1271 : : !$OMP TARGET ENTER DATA MAP(to:cwavef) IF(transfer_cwavef)
1272 : : !$OMP TARGET ENTER DATA MAP(alloc:cwavef_alltoall2) IF(gpu_option_==ABI_GPU_OPENMP)
1273 : : if ( ((.not.flag_inv_sym) .and. (bandpp>1) ) .or. flag_inv_sym )then
1274 : : !$OMP TARGET ENTER DATA MAP(alloc:cwavef_alltoall1) IF(gpu_option_==ABI_GPU_OPENMP)
1275 : : end if
1276 : : #endif
1277 : :
1278 : 182282 : call timab(547,1,tsec)
1279 [ + - ]: 182282 : if(gpu_option_==ABI_GPU_KOKKOS) then
1280 : : #if defined HAVE_GPU && defined HAVE_YAKL
1281 : : ABI_MALLOC(cwavef_mpi,(2,npw_k*blocksize))
1282 : :
1283 : : call gpu_data_prefetch_async(C_LOC(cwavef), INT(2, c_size_t)*npw_k*blocksize, CPU_DEVICE_ID)
1284 : : call gpu_device_synchronize()
1285 : :
1286 : : cwavef_mpi(:,:) = cwavef(:,:)
1287 : :
1288 : : call xmpi_alltoallv(cwavef_mpi,sendcountsloc,sdisplsloc,cwavef_alltoall2,&
1289 : : & recvcountsloc,rdisplsloc,spaceComm,ier)
1290 : : ABI_FREE(cwavef_mpi)
1291 : : #endif
1292 : : else
1293 : : call xmpi_alltoallv(cwavef,sendcountsloc,sdisplsloc,cwavef_alltoall2,&
1294 : 182282 : & recvcountsloc,rdisplsloc,spaceComm,ier,use_omp_map=(gpu_option_==ABI_GPU_OPENMP))
1295 : : end if
1296 : :
1297 : : #ifdef HAVE_OPENMP_OFFLOAD
1298 : : !$OMP TARGET EXIT DATA MAP(delete:cwavef) IF(transfer_cwavef)
1299 : : #endif
1300 : 182282 : call timab(547,2,tsec)
1301 : :
1302 : 182282 : tim_fourwf=16
1303 : :
1304 : : !Eventually adjust load balancing for FFT (by changing FFT distrib)
1305 [ - + ]: 182282 : if(have_to_reequilibrate) then
1306 : 0 : npw_fft = bandfft_kpt_ptr%npw_fft
1307 : 0 : sendcount_fft => bandfft_kpt_ptr%sendcount_fft(:)
1308 : 0 : recvcount_fft => bandfft_kpt_ptr%recvcount_fft(:)
1309 : 0 : senddisp_fft => bandfft_kpt_ptr%senddisp_fft(:)
1310 : 0 : recvdisp_fft => bandfft_kpt_ptr%recvdisp_fft(:)
1311 : 0 : indices_pw_fft => bandfft_kpt_ptr%indices_pw_fft(:)
1312 : 0 : kg_k_fft => bandfft_kpt_ptr%kg_k_fft(:,:)
1313 [ # # # # : 0 : ABI_MALLOC( buff_wf, (2,ndatarecv*bandpp) ) ! for sorting cgwavef
# # # # #
# # # ]
1314 [ # # # # : 0 : ABI_MALLOC( cwavef_fft, (2,npw_fft*bandpp) )
# # # # #
# # # ]
1315 [ # # ]: 0 : if(bandpp>1) then
1316 [ # # # # : 0 : ABI_MALLOC( cwavef_fft_tr, (2,npw_fft*bandpp))
# # # # ]
1317 : : end if
1318 : : end if
1319 : :
1320 [ - + - - : 182282 : if (option_fourwf==0) wfraug(:,:,:,:)=zero
- - - - -
- ]
1321 : :
1322 : : !====================================================================
1323 [ + + ]: 182282 : if ((.not.(flag_inv_sym)) .and. (bandpp==1)) then
1324 : :
1325 : : ! Compute the index of the band
1326 : 144896 : ind_occ = (iblock-1)*blocksize + mpi_enreg%me_band + 1
1327 : :
1328 [ + + - + ]: 144896 : if(abs(occ_k(ind_occ))>=tol8.or.option_fourwf==0) then
1329 : :
1330 : : ! Compute the weight of the band
1331 : 73382 : weight=occ_k(ind_occ)*wtk/ucvol
1332 : :
1333 [ - + ]: 73382 : if(have_to_reequilibrate) then
1334 : : ! filling of sorted send buffers before exchange
1335 [ # # ]: 0 : do ipw = 1 ,ndatarecv
1336 [ # # ]: 0 : buff_wf(1:2, indices_pw_fft(ipw) ) = cwavef_alltoall2(1:2,ipw)
1337 : : end do
1338 : : call xmpi_alltoallv(buff_wf,2*sendcount_fft,2*senddisp_fft, &
1339 [ # # # # : 0 : & cwavef_fft,2*recvcount_fft, 2*recvdisp_fft, mpi_enreg%comm_fft,ier)
# # # # #
# # # # #
# # ]
1340 : : call fourwf(1,rhoaug,cwavef_fft,dummy,wfraug,gbound_,gbound_,&
1341 : : & istwf_k_,kg_k_fft,kg_k_fft,mgfft,mpi_enreg,1,&
1342 : : & ngfft,npw_fft,1,n4,n5,n6,option_fourwf,tim_fourwf,weight,weight,&
1343 [ # # # # : 0 : & gpu_option=gpu_option_)
# # # # #
# # # # #
# # # # #
# # # # #
# # ]
1344 : : else
1345 : : call fourwf(1,rhoaug,cwavef_alltoall2,dummy,wfraug,gbound_,gbound_,&
1346 : : & istwf_k_,kg_k_gather,kg_k_gather,mgfft,mpi_enreg,1,&
1347 : : & ngfft,ndatarecv,1,n4,n5,n6,option_fourwf,tim_fourwf,weight,weight,&
1348 [ - + - - : 73382 : & gpu_option=gpu_option_)
- - - - -
- - - - +
- - - - -
- - - - +
- + ]
1349 : : end if
1350 [ - + - - ]: 73382 : if (option_fourwf==0.and.nproc_fft>1) then
1351 [ # # ]: 0 : if (me_fft>0) then
1352 : 0 : nd3=(ngfft(3)-1)/nproc_fft+1
1353 [ # # # # : 0 : wfraug(:,:,:,me_fft*nd3+1:me_fft*nd3+nd3)=wfraug(:,:,:,1:nd3)
# # # # #
# # # # #
# # # # ]
1354 [ # # # # : 0 : wfraug(:,:,:,1:nd3)=zero
# # # # ]
1355 : : end if
1356 : 0 : call xmpi_sum(wfraug,mpi_enreg%comm_fft,ier)
1357 : : end if
1358 : : end if
1359 : :
1360 : : !====================================================================
1361 [ + + ]: 37386 : else if ((.not.(flag_inv_sym)) .and. (bandpp>1) ) then
1362 : :
1363 : : ! -------------------------------------------------------------
1364 : : ! Computation of the index to class the waves functions below bandpp
1365 : : ! -------------------------------------------------------------
1366 : : call prep_index_wavef_bandpp(nproc_band,bandpp,&
1367 : : & 1,ndatarecv,&
1368 : : & recvcounts,rdispls,&
1369 : 30764 : & index_wavef_band)
1370 : :
1371 : : ! -------------------------------------------------------
1372 : : ! Sorting of the wave functions below bandpp
1373 : : ! -------------------------------------------------------
1374 [ + - ]: 30764 : if(gpu_option_==ABI_GPU_OPENMP) then
1375 : : !FIXME Having this OpenMP kernel here cause a seg fault, so I put it in a separate routine
1376 : : ! I can't figure what cause the issue...
1377 : : call omp_copy_sort_wf(bandpp,ndatarecv,cwavef_alltoall1,cwavef_alltoall2,index_wavef_band,flag_inv_sym)
1378 : : else
1379 [ + + + + ]: 105444488 : cwavef_alltoall1(:,:) = cwavef_alltoall2(:,index_wavef_band)
1380 : : end if
1381 : :
1382 : 30764 : if(have_to_reequilibrate) then
1383 : : ! filling of sorted send buffers before exchange
1384 [ # # ]: 0 : do iibandpp=1,bandpp
1385 [ # # ]: 0 : do ipw = 1 ,ndatarecv
1386 : : buff_wf(1:2, iibandpp + bandpp*(indices_pw_fft(ipw)-1)) = &
1387 [ # # ]: 0 : & cwavef_alltoall1(1:2,ipw + ndatarecv*(iibandpp-1))
1388 : : end do
1389 : : end do
1390 : : call xmpi_alltoallv(buff_wf,2*bandpp*sendcount_fft,2*bandpp*senddisp_fft, &
1391 [ # # # # : 0 : & cwavef_fft_tr,2*bandpp*recvcount_fft, 2*bandpp*recvdisp_fft, mpi_enreg%comm_fft,ier)
# # # # #
# # # # #
# # ]
1392 [ # # ]: 0 : do iibandpp=1,bandpp
1393 [ # # ]: 0 : do ipw = 1 ,npw_fft
1394 [ # # ]: 0 : cwavef_fft(1:2, ipw + npw_fft*(iibandpp -1)) = cwavef_fft_tr(1:2, iibandpp + bandpp*(ipw-1))
1395 : : end do
1396 : : end do
1397 : : end if
1398 : :
1399 : : ! -------------------
1400 : : ! Fourier calculation
1401 : : ! -------------------
1402 : : ! GPU version
1403 [ - + ]: 30764 : if(gpu_option_/=ABI_GPU_DISABLED) then
1404 [ # # # # : 0 : ABI_MALLOC(weight_t,(bandpp))
# # # # #
# ]
1405 [ # # ]: 0 : do iibandpp=1,bandpp
1406 : : ! Compute the index of the band
1407 : 0 : ind_occ = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + iibandpp
1408 : : ! Compute the weight of the band
1409 : 0 : weight_t(iibandpp)=occ_k(ind_occ)*wtk/ucvol
1410 [ # # ]: 0 : if(abs(occ_k(ind_occ)) < tol8) weight_t(iibandpp) = zero
1411 : : end do
1412 : :
1413 : 0 : nband_fftblock=bandpp/nfft_blocks+modulo(bandpp,nfft_blocks)
1414 : : call fourwf_optmem(1,rhoaug,&
1415 : : & cwavef_alltoall1,&
1416 : : & dummy,wfraug(:,:,:,1:n6*nband_fftblock),gbound_,gbound_,&
1417 : : & istwf_k_,kg_k_gather,kg_k_gather,mgfft,mpi_enreg,bandpp,nfft_blocks,&
1418 : : & ngfft,ndatarecv,1,n4,n5,n6,option_fourwf,tim_fourwf,weight,weight,&
1419 : : & weight_array_r=weight_t,weight_array_i=weight_t,&
1420 [ # # # # ]: 0 : & gpu_option=gpu_option_)
1421 [ # # ]: 0 : ABI_FREE(weight_t)
1422 : :
1423 : : ! Standard version
1424 : : else
1425 [ + + ]: 146244 : do iibandpp=1,bandpp
1426 : : ! Compute the index of the band
1427 : 115480 : ind_occ = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + iibandpp
1428 : : ! Compute the weight of the band
1429 : 115480 : weight=occ_k(ind_occ)*wtk/ucvol
1430 [ - + ]: 115480 : if (option_fourwf==0) then
1431 : 0 : wfraug_ptr => wfraug(:,:,:,(iibandpp-1)*n6+1:iibandpp*n6)
1432 : : else
1433 : 115480 : wfraug_ptr => wfraug
1434 : : end if
1435 [ + + - + ]: 146244 : if (abs(occ_k(ind_occ)) >=tol8.or.option_fourwf==0) then
1436 [ - + ]: 55811 : if(have_to_reequilibrate) then
1437 : : call fourwf(1,rhoaug, &
1438 : : & cwavef_fft(:,(npw_fft*(iibandpp-1))+1:(npw_fft*iibandpp)), &
1439 : : & dummy,wfraug_ptr,gbound_,gbound_,&
1440 : : & istwf_k_,kg_k_fft,kg_k_fft,mgfft,mpi_enreg,1,&
1441 : : & ngfft,npw_fft,1,n4,n5,n6,option_fourwf,tim_fourwf,weight,weight,&
1442 [ # # # # : 0 : & gpu_option=gpu_option_)
# # ]
1443 : : else
1444 : : call fourwf(1,rhoaug,&
1445 : : & cwavef_alltoall1(:,(ndatarecv*(iibandpp-1))+1:(ndatarecv*iibandpp)),&
1446 : : & dummy,wfraug_ptr,gbound_,gbound_,&
1447 : : & istwf_k_,kg_k_gather,kg_k_gather,mgfft,mpi_enreg,1,&
1448 : : & ngfft,ndatarecv,1,n4,n5,n6,option_fourwf,&
1449 [ - + - + : 55811 : & tim_fourwf,weight,weight)
- + ]
1450 : : end if
1451 [ - + - - ]: 55811 : if (option_fourwf==0.and.nproc_fft>1) then
1452 [ # # ]: 0 : if (me_fft>0) then
1453 : 0 : nd3=(ngfft(3)-1)/nproc_fft+1
1454 [ # # # # : 0 : wfraug_ptr(:,:,:,me_fft*nd3+1:me_fft*nd3+nd3)=wfraug_ptr(:,:,:,1:nd3)
# # # # #
# # # # #
# # # # ]
1455 [ # # # # : 0 : wfraug_ptr(:,:,:,1:nd3)=zero
# # # # ]
1456 : : end if
1457 : 0 : call xmpi_sum(wfraug_ptr,mpi_enreg%comm_fft,ier)
1458 : : end if
1459 : : end if
1460 : : end do
1461 : : end if ! (gpu_option/=0)
1462 : :
1463 : : ! -----------------------------------------------------
1464 : : ! Sorting waves functions below the processors
1465 : : ! -----------------------------------------------------
1466 : : ! cwavef_alltoall(:,index_wavef_band) = cwavef_alltoall(:,:) ! NOT NEEDED
1467 [ - + ]: 30764 : ABI_FREE(index_wavef_band)
1468 : :
1469 : : !====================================================================
1470 [ + - ]: 6622 : else if (flag_inv_sym) then
1471 : :
1472 : : ! -------------------------------------------------------------
1473 : : ! Computation of the index to class the waves functions below bandpp
1474 : : ! -------------------------------------------------------------
1475 : : call prep_index_wavef_bandpp(nproc_band,bandpp,&
1476 : : & 1,ndatarecv,&
1477 : : & recvcounts,rdispls,&
1478 : 6622 : & index_wavef_band)
1479 : :
1480 : : ! -------------------------------------------------------
1481 : : ! Sorting the wave functions below bandpp
1482 : : ! -------------------------------------------------------
1483 [ + - ]: 6622 : if(gpu_option_==ABI_GPU_OPENMP) then
1484 : : !FIXME Having this OpenMP kernel here cause a seg fault, so I put it in a separate routine
1485 : : ! I can't figure what cause the issue...
1486 : : call omp_copy_sort_wf(bandpp,ndatarecv,cwavef_alltoall1,cwavef_alltoall2,index_wavef_band,flag_inv_sym)
1487 : : else
1488 [ + + + + ]: 75855454 : cwavef_alltoall1(:,:) = cwavef_alltoall2(:,index_wavef_band)
1489 : : end if
1490 : :
1491 : : ! ------------------------------------------------------------
1492 : : ! We associate the waves functions by two
1493 : : ! ------------------------------------------------------------
1494 : : call prep_wavef_sym_do(mpi_enreg,bandpp,1,&
1495 : : & ndatarecv,&
1496 : : & ndatarecv_tot,ndatasend_sym,tab_proc,&
1497 : : & cwavef_alltoall1,&
1498 : : & sendcounts_sym,sdispls_sym,&
1499 : : & recvcounts_sym,rdispls_sym,&
1500 : : & ewavef_alltoall_sym,&
1501 : 6622 : & index_wavef_send)
1502 : :
1503 : : ! ------------------------------------------------------------
1504 : : ! Fourier calculation
1505 : : ! ------------------------------------------------------------
1506 : : ! GPU version
1507 [ - + ]: 6622 : if (gpu_option_/=ABI_GPU_DISABLED) then
1508 [ # # # # : 0 : ABI_MALLOC(weight1_t,(bandpp_sym))
# # # # #
# # # ]
1509 [ # # # # : 0 : ABI_MALLOC(weight2_t,(bandpp_sym))
# # # # ]
1510 [ # # ]: 0 : do iibandpp=1,bandpp_sym
1511 [ # # ]: 0 : if (bandpp/=1) then
1512 : 0 : ind_occ1 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + (2*iibandpp-1)
1513 : 0 : ind_occ2 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + (2*iibandpp )
1514 : : else
1515 : 0 : ind_occ1 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + 1
1516 : 0 : ind_occ2 = ind_occ1
1517 : : end if
1518 : 0 : weight1_t(iibandpp) = occ_k(ind_occ1)*wtk/ucvol
1519 : 0 : weight2_t(iibandpp) = occ_k(ind_occ2)*wtk/ucvol
1520 : : end do
1521 : :
1522 : 0 : nband_fftblock=bandpp_sym/nfft_blocks_sym+modulo(bandpp_sym,nfft_blocks_sym)
1523 : : call fourwf_optmem(1,rhoaug,&
1524 : : & ewavef_alltoall_sym,&
1525 : : & dummy,wfraug(:,:,:,1:n6*nband_fftblock),gbound_,gbound_,&
1526 : : & istwf_k_,kg_k_gather_sym,kg_k_gather_sym,mgfft,mpi_enreg,bandpp_sym,nfft_blocks_sym,&
1527 : : & ngfft,ndatarecv_tot,1,n4,n5,n6,option_fourwf,tim_fourwf,weight,weight,&
1528 : : & weight_array_r=weight1_t,weight_array_i=weight2_t,&
1529 [ # # # # : 0 : & gpu_option=gpu_option_)
# # ]
1530 : :
1531 [ # # ]: 0 : ABI_FREE(weight1_t)
1532 [ # # ]: 0 : ABI_FREE(weight2_t)
1533 : :
1534 : : ! Standard version
1535 : : else
1536 [ - + - - ]: 6622 : if (option_fourwf==0.and.bandpp>1) then
1537 [ # # # # : 0 : ABI_MALLOC(wfraug_ptr,(2,n4,n5,n6))
# # # # #
# # # # #
# # # # #
# ]
1538 : : else
1539 : 6622 : wfraug_ptr => wfraug
1540 : : end if
1541 [ + + ]: 23676 : do iibandpp=1,bandpp_sym
1542 [ + + ]: 17054 : if (bandpp/=1) then
1543 : 12926 : ind_occ1 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + (2*iibandpp-1)
1544 : 12926 : ind_occ2 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + (2*iibandpp )
1545 : : else
1546 : 4128 : ind_occ1 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + 1
1547 : 4128 : ind_occ2 = ind_occ1
1548 : : end if
1549 : 17054 : weight1 = occ_k(ind_occ1)*wtk/ucvol
1550 : 17054 : weight2 = occ_k(ind_occ2)*wtk/ucvol
1551 : 17054 : ewavef_alltoall_sym_work => ewavef_alltoall_sym(:,(ndatarecv_tot*(iibandpp-1))+1:(ndatarecv_tot*iibandpp))
1552 : : call fourwf(1,rhoaug,&
1553 : : & ewavef_alltoall_sym_work,&
1554 : : & dummy,wfraug_ptr,gbound_,gbound_,&
1555 : : & istwf_k_,kg_k_gather_sym,kg_k_gather_sym,mgfft,mpi_enreg,1,&
1556 : : & ngfft,ndatarecv_tot,1,n4,n5,n6,option_fourwf,&
1557 [ - + - + : 17054 : & tim_fourwf,weight1,weight2)
- + - + ]
1558 [ - + ]: 23676 : if (option_fourwf==0) then
1559 [ # # ]: 0 : if (modulo(bandpp,2)==0) then
1560 : 0 : jjbandpp=2*iibandpp-1
1561 [ # # # # : 0 : wfraug(1,:,:,(jjbandpp-1)*n6+1:jjbandpp*n6)=wfraug_ptr(1,:,:,1:n6)
# # # # #
# # # #
# ]
1562 [ # # # # : 0 : wfraug(1,:,:,(jjbandpp)*n6+1:(jjbandpp+1)*n6)=wfraug_ptr(2,:,:,1:n6)
# # # # #
# # # ]
1563 [ # # ]: 0 : else if (bandpp>1) then
1564 [ # # # # : 0 : wfraug(1,:,:,(iibandpp-1)*n6+1:iibandpp*n6)=wfraug_ptr(1,:,:,1:n6)
# # # # #
# # # #
# ]
1565 : : end if
1566 [ # # ]: 0 : if (nproc_fft>1) then
1567 [ # # ]: 0 : if (me_fft>0) then
1568 : 0 : nd3=(ngfft(3)-1)/nproc_fft+1
1569 [ # # # # : 0 : wfraug(1,:,:,me_fft*nd3+1:me_fft*nd3+nd3)=wfraug(1,:,:,1:nd3)
# # # # #
# # # #
# ]
1570 [ # # # # : 0 : wfraug(1,:,:,1:nd3)=zero
# # ]
1571 : : end if
1572 : 0 : call xmpi_sum(wfraug,mpi_enreg%comm_fft,ier)
1573 : : end if
1574 : : end if
1575 : : end do
1576 [ - + - - ]: 6622 : if (option_fourwf==0.and.bandpp>1) then
1577 [ # # ]: 0 : ABI_FREE(wfraug_ptr)
1578 : : end if
1579 : : end if ! (gpu_option/=ABI_GPU_DISABLED)
1580 : :
1581 : : ! ------------------------------------------------------------
1582 : : ! We dissociate each wave function in two waves functions
1583 : : ! gwavef is classed below of bandpp
1584 : : ! ------------------------------------------------------------
1585 : : call prep_wavef_sym_undo(mpi_enreg,bandpp,1,&
1586 : : & ndatarecv,&
1587 : : & ndatarecv_tot,ndatasend_sym,idatarecv0,&
1588 : : & cwavef_alltoall1,&
1589 : : & sendcounts_sym,sdispls_sym,&
1590 : : & recvcounts_sym,rdispls_sym,&
1591 : : & ewavef_alltoall_sym,&
1592 : 6622 : & index_wavef_send)
1593 : :
1594 [ - + ]: 6622 : ABI_FREE(ewavef_alltoall_sym)
1595 [ - + ]: 6622 : ABI_FREE(index_wavef_send)
1596 : :
1597 : : ! -------------------------------------------------------
1598 : : ! Sorting waves functions below the processors
1599 : : ! -------------------------------------------------------
1600 : : ! cwavef_alltoall(:,index_wavef_band) = cwavef_alltoall(:,:) ! NOT NEEDED
1601 : :
1602 [ - + ]: 6622 : ABI_FREE(index_wavef_band)
1603 : :
1604 : : end if
1605 : :
1606 : : !====================================================================
1607 [ - + ]: 182282 : if(have_to_reequilibrate) then
1608 [ # # ]: 0 : ABI_FREE(buff_wf)
1609 [ # # ]: 0 : ABI_FREE(cwavef_fft)
1610 [ # # ]: 0 : if(bandpp > 1) then
1611 [ # # ]: 0 : ABI_FREE(cwavef_fft_tr)
1612 : : end if
1613 : : end if
1614 : : #ifdef HAVE_OPENMP_OFFLOAD
1615 : : !$OMP TARGET EXIT DATA MAP(delete:cwavef_alltoall2) IF(gpu_option_==ABI_GPU_OPENMP)
1616 : : if ( ((.not.flag_inv_sym) .and. (bandpp>1) ) .or. flag_inv_sym ) then
1617 : : !$OMP TARGET EXIT DATA MAP(delete:cwavef_alltoall1) IF(gpu_option_==ABI_GPU_OPENMP)
1618 : : end if
1619 : : #endif
1620 [ - + ]: 182282 : ABI_FREE(sendcountsloc)
1621 [ - + ]: 182282 : ABI_FREE(sdisplsloc)
1622 [ - + ]: 182282 : ABI_FREE(recvcountsloc)
1623 [ - + ]: 182282 : ABI_FREE(rdisplsloc)
1624 [ - + ]: 182282 : ABI_FREE(cwavef_alltoall2)
1625 [ - + + + : 401950 : if ( ((.not.flag_inv_sym) .and. (bandpp>1) ) .or. flag_inv_sym ) then
+ + - + -
+ - + - +
- + - + ]
1626 [ + - ]: 37386 : if(gpu_option_==ABI_GPU_KOKKOS) then
1627 : : #if defined HAVE_GPU && defined HAVE_YAKL
1628 : : ABI_FREE_MANAGED(cwavef_alltoall1)
1629 : : #endif
1630 : : else
1631 [ - + ]: 37386 : ABI_FREE(cwavef_alltoall1)
1632 : : end if
1633 : : end if
1634 : :
1635 : : contains
1636 [ - + ]: 37386 : subroutine omp_copy_sort_wf(bandpp,ndatarecv,cwavef_alltoall_out,cwavef_alltoall_in,index_wavef_band,flag_inv_sym)
1637 : : !Arguments ------------------------------------
1638 : : !scalars
1639 : : integer,intent(in) :: bandpp,ndatarecv
1640 : : logical,intent(in) :: flag_inv_sym
1641 : : !arrays
1642 : : integer,intent(in) :: index_wavef_band(bandpp*ndatarecv)
1643 : : real(dp),intent(in) :: cwavef_alltoall_in(2,bandpp*ndatarecv)
1644 : : real(dp),intent(out) :: cwavef_alltoall_out(2,bandpp*ndatarecv)
1645 : :
1646 : : !Locals ---------------------------------------
1647 : : integer :: ind,iibandpp,ipw
1648 : :
1649 : : ! *************************************************************************
1650 : :
1651 : : #ifdef HAVE_OPENMP_OFFLOAD
1652 : : !FIXME This doesn't run fine when flag_invsym is on, why !!???
1653 : : !$OMP TARGET UPDATE FROM(cwavef_alltoall_in) IF(flag_inv_sym)
1654 : :
1655 : : !$OMP TARGET TEAMS DISTRIBUTE &
1656 : : !$OMP& PRIVATE(iibandpp) MAP(to:cwavef_alltoall_out,cwavef_alltoall_in,index_wavef_band) &
1657 : : !$OMP& IF(.not. flag_inv_sym)
1658 : : do iibandpp=1,bandpp
1659 : : !$OMP PARALLEL DO PRIVATE(ind,ipw)
1660 : : do ipw = 1 ,ndatarecv
1661 : : ind=index_wavef_band(ipw + ndatarecv*(iibandpp-1))
1662 : : cwavef_alltoall_out(1,ipw + ndatarecv*(iibandpp-1)) = cwavef_alltoall_in(1,ind)
1663 : : cwavef_alltoall_out(2,ipw + ndatarecv*(iibandpp-1)) = cwavef_alltoall_in(2,ind)
1664 : : end do
1665 : : end do
1666 : :
1667 : : !$OMP TARGET UPDATE TO(cwavef_alltoall_out) IF(flag_inv_sym)
1668 : :
1669 : : #else
1670 : : ABI_UNUSED((/ind,iibandpp,ipw/))
1671 : : ABI_UNUSED((/bandpp,ndatarecv/))
1672 : : ABI_UNUSED(flag_inv_sym)
1673 : : ABI_UNUSED(cwavef_alltoall_in)
1674 : : ABI_UNUSED(cwavef_alltoall_out)
1675 : : ABI_UNUSED(index_wavef_band)
1676 : : #endif
1677 : : end subroutine omp_copy_sort_wf
1678 : :
1679 : : end subroutine prep_fourwf
1680 : : !!***
1681 : :
1682 : : !!****f* ABINIT/prep_wavef_sym_do
1683 : : !! NAME
1684 : : !! prep_wavef_sym_do
1685 : : !!
1686 : : !! FUNCTION
1687 : : !! this routine associates waves functions by two as following
1688 : : !! E(G) = C(G) + D(G)
1689 : : !! E(-G) = C*(G) + iD*(G)
1690 : : !! the values are distributed on the processors in function of
1691 : : !! the value of mpi_enreg%distribfft%tab_fftwf2_distrib( (-kg_k_gather(2,i) )
1692 : : !!
1693 : : !! INPUTS
1694 : : !! mpi_enreg = information about mpi parallelization
1695 : : !! bandpp = number of couple of waves functions
1696 : : !! nspinor = number of spin
1697 : : !! ndatarecv = number of values received by the processor and sended
1698 : : !! by the other processors band
1699 : : !! ndatarecv_tot = total number of received values
1700 : : !! (ndatarecv + number of received opposited planewave coordinates)
1701 : : !! ndatasend_sym = number of sended values to the processors fft to create opposited
1702 : : !! planewave coordinates
1703 : : !! tab_proc = positions of opposited planewave coordinates in the list of the
1704 : : !! processors fft
1705 : : !! cwavef_alltoall = planewave coefficients of wavefunction
1706 : : !! ( initial of the processor + sended by other processors band)
1707 : : !! sendcounts_sym = number of sended values by the processor to each processor fft
1708 : : !! sdispls_sym = postions of the sended values by the processor to each processor fft
1709 : : !!
1710 : : !! recvcounts_sym = number of the received values by the processor from each processor fft
1711 : : !! rdispls_sym = postions of the received values by the processor from each processor fft
1712 : : !!
1713 : : !! OUTPUT
1714 : : !! ewavef_alltoall_sym = planewave coefficients of wavefunction
1715 : : !! initial of the processor +
1716 : : !! sended by other processors band +
1717 : : !! sended by other processors fft +
1718 : : !! and compisited if bandpp >1
1719 : : !! index_wavef_send = index to send the values in blocks to the other processor fft
1720 : : !!
1721 : : !! SIDE EFFECTS
1722 : : !!
1723 : : !! SOURCE
1724 : :
1725 : 9896 : subroutine prep_wavef_sym_do(mpi_enreg,bandpp,nspinor,&
1726 : : & ndatarecv,&
1727 : 19792 : & ndatarecv_tot,ndatasend_sym,tab_proc,&
1728 [ - + ]: 9896 : & cwavef_alltoall,&
1729 [ - + ]: 9896 : & sendcounts_sym,sdispls_sym,&
1730 [ - + - + : 9896 : & recvcounts_sym,rdispls_sym,&
- + ]
1731 : : & ewavef_alltoall_sym,&
1732 : : & index_wavef_send)
1733 : :
1734 : : !Arguments ------------------------------------
1735 : : !scalars
1736 : : integer,intent(in) :: bandpp,ndatarecv,ndatarecv_tot,ndatasend_sym
1737 : : integer,intent(in) :: nspinor
1738 : : type(mpi_type),intent(in) :: mpi_enreg
1739 : : !arrays
1740 : : integer,allocatable,intent(out) :: index_wavef_send(:)
1741 : : integer,intent(in) :: rdispls_sym(:),recvcounts_sym(:)
1742 : : integer,intent(in) :: sdispls_sym(:),sendcounts_sym(:)
1743 : : integer,intent(in) :: tab_proc(:)
1744 : : real(dp),intent(inout) :: cwavef_alltoall(2,ndatarecv*nspinor*bandpp)
1745 : : real(dp),pointer :: ewavef_alltoall_sym(:,:)
1746 : :
1747 : : !Local variables-------------------------------
1748 : : !scalars
1749 : : integer :: bandpp_sym,ibandpp,idatarecv,ideb_loc,idebc,idebd,idebe
1750 : : integer :: ier,ifin_loc,ifinc,ifind,ifine,iproc,jbandpp,jsendloc
1751 : : integer :: kbandpp,newspacecomm,nproc_fft
1752 : : logical :: flag_compose
1753 : : !arrays
1754 : : integer,allocatable :: rdispls_sym_loc(:),recvcounts_sym_loc(:)
1755 : : integer,allocatable :: sdispls_sym_loc(:),sendcounts_sym_loc(:)
1756 : 9896 : real(dp),allocatable :: ewavef_alltoall_loc(:,:),ewavef_alltoall_send(:,:)
1757 : :
1758 : : ! *********************************************************************
1759 : :
1760 : : !DEBUG
1761 : : !write(std_out,*)' prep_wavef_sym_do : enter '
1762 : : !ENDDEBUG
1763 : :
1764 : : !---------------------------------------------
1765 : : !Initialisation
1766 : : !---------------------------------------------
1767 : 9896 : nproc_fft = mpi_enreg%nproc_fft
1768 : :
1769 : 9896 : newspacecomm = mpi_enreg%comm_fft
1770 : :
1771 [ + + ]: 9896 : if (modulo(bandpp,2)==0) then
1772 : 4808 : bandpp_sym = bandpp/2
1773 : 4808 : flag_compose = .TRUE.
1774 : : else
1775 : : bandpp_sym = bandpp
1776 : : flag_compose = .FALSE.
1777 : : end if
1778 : :
1779 : : !---------------------------------------------
1780 : : !Allocation
1781 : : !---------------------------------------------
1782 [ - + + - : 29688 : ABI_MALLOC(ewavef_alltoall_sym ,(2,ndatarecv_tot*bandpp_sym))
+ - + - -
+ - + ]
1783 [ - + + - : 29688 : ABI_MALLOC(ewavef_alltoall_loc ,(2,ndatarecv *bandpp_sym))
+ - + - -
+ - + ]
1784 [ - + + - : 29688 : ABI_MALLOC(ewavef_alltoall_send ,(2,ndatasend_sym*bandpp_sym))
+ - + - -
+ - + ]
1785 [ - + + - : 29688 : ABI_MALLOC(index_wavef_send ,( ndatasend_sym*bandpp_sym))
+ - + - -
+ - + -
+ ]
1786 : :
1787 [ - + + - : 29688 : ABI_MALLOC(sendcounts_sym_loc ,(nproc_fft))
+ - + - -
+ - + ]
1788 [ - + + - : 19792 : ABI_MALLOC(sdispls_sym_loc ,(nproc_fft))
- + - + ]
1789 [ - + + - : 19792 : ABI_MALLOC(recvcounts_sym_loc ,(nproc_fft))
- + - + ]
1790 [ - + + - : 19792 : ABI_MALLOC(rdispls_sym_loc ,(nproc_fft))
- + - + ]
1791 : :
1792 : :
1793 : : !Initialisation
1794 : : !--------------
1795 [ + + + + ]: 106566374 : ewavef_alltoall_sym(:,:) =0.
1796 [ + + + + ]: 53319068 : ewavef_alltoall_loc(:,:) =0.
1797 : :
1798 [ + + ]: 21958 : sendcounts_sym_loc(:) =0
1799 [ + + ]: 21958 : sdispls_sym_loc(:) =0
1800 [ + + ]: 21958 : recvcounts_sym_loc(:) =0
1801 [ + + ]: 21958 : rdispls_sym_loc(:) =0
1802 : :
1803 [ + + ]: 17758998 : index_wavef_send(:) =0
1804 : :
1805 : :
1806 : : !-------------------------------------------------
1807 : : !We are bandpp blocks which we want to :
1808 : : !associate by two (band_sym==bandpp/2)
1809 : : !or not associate by two (band_sym==bandpp)
1810 : : !
1811 : : !So We'll have got bandpp_sym blocks
1812 : : !So we loop on the bandpp_sym blocks
1813 : : !--------------------------------------------------
1814 : :
1815 [ + + ]: 32494 : do kbandpp=1,bandpp_sym
1816 : :
1817 : : ! position of the two blocks
1818 : : ! --------------------------
1819 : 22598 : ibandpp = (kbandpp-1) * 2
1820 : 22598 : jbandpp = ibandpp + 1
1821 : :
1822 : 22598 : idebe = (kbandpp-1) * ndatarecv_tot + 1
1823 : 22598 : ifine = idebe + ndatarecv - 1
1824 : :
1825 : 22598 : idebc = ibandpp * ndatarecv + 1
1826 : 22598 : ifinc = idebc + ndatarecv - 1
1827 : :
1828 : 22598 : idebd = jbandpp * ndatarecv + 1
1829 : 22598 : ifind = idebd + ndatarecv - 1
1830 : :
1831 : 22598 : ideb_loc = (kbandpp-1) * ndatarecv + 1
1832 : 22598 : ifin_loc = ideb_loc + ndatarecv - 1
1833 : :
1834 : :
1835 [ + + ]: 32494 : if (flag_compose) then
1836 : :
1837 : :
1838 : : ! calcul ewavef(G)
1839 : : ! ----------------
1840 : : ewavef_alltoall_sym(1,idebe:ifine) = &
1841 : : & cwavef_alltoall(1,idebc:ifinc) &
1842 [ + + ]: 16760514 : & - cwavef_alltoall(2,idebd:ifind)
1843 : :
1844 : : ewavef_alltoall_sym(2,idebe:ifine) = &
1845 : : & cwavef_alltoall(2,idebc:ifinc) &
1846 [ + + ]: 16760514 : & + cwavef_alltoall(1,idebd:ifind)
1847 : :
1848 : : ! calcul ewavef_loc(-G)
1849 : : ! ---------------------
1850 : : ewavef_alltoall_loc(1,ideb_loc:ifin_loc) = &
1851 : : & cwavef_alltoall(1,idebc:ifinc) &
1852 [ + + ]: 16760514 : & + cwavef_alltoall(2,idebd:ifind)
1853 : :
1854 : : ewavef_alltoall_loc(2,ideb_loc:ifin_loc) = &
1855 : : & - cwavef_alltoall(2,idebc:ifinc) &
1856 [ + + ]: 16760514 : & + cwavef_alltoall(1,idebd:ifind)
1857 : : else
1858 : :
1859 : : ! calcul ewavef(G)
1860 : : ! ----------------
1861 [ + + ]: 1031808 : ewavef_alltoall_sym(1,idebe:ifine) = cwavef_alltoall(1,idebc:ifinc)
1862 [ + + ]: 1031808 : ewavef_alltoall_sym(2,idebe:ifine) = cwavef_alltoall(2,idebc:ifinc)
1863 : :
1864 : : ! calcul ewavef_loc(-G)
1865 : : ! ---------------------
1866 [ + + ]: 1031808 : ewavef_alltoall_loc(1,ideb_loc:ifin_loc) = cwavef_alltoall(1,idebc:ifinc)
1867 [ + + ]: 1031808 : ewavef_alltoall_loc(2,ideb_loc:ifin_loc) = - cwavef_alltoall(2,idebc:ifinc)
1868 : :
1869 : : end if
1870 : :
1871 : : end do
1872 : :
1873 : :
1874 : :
1875 : : !------------------------------------------------------------------------
1876 : : !Creation of datas blocks for each processor fft from ewavef_alltoall_loc
1877 : : !to send datas by blocks with a alltoall...
1878 : : !------------------------------------------------------------------------
1879 : :
1880 : : !Position of the blocks
1881 : : !----------------------
1882 : : jsendloc=0
1883 [ + + ]: 32494 : do ibandpp=1,bandpp_sym
1884 [ + + ]: 59044 : do iproc=1,nproc_fft
1885 [ + + ]: 20112644 : do idatarecv=1,ndatarecv
1886 [ + + ]: 20090046 : if (tab_proc(idatarecv)==(iproc-1)) then
1887 : 17749102 : jsendloc=jsendloc+1
1888 : 17749102 : index_wavef_send(jsendloc) = idatarecv + ndatarecv * (ibandpp-1)
1889 : : end if
1890 : : end do
1891 : : end do
1892 : : end do
1893 : :
1894 : : !Classment
1895 : : !----------
1896 [ + + + + ]: 53257202 : ewavef_alltoall_send(:,:)=ewavef_alltoall_loc(:,index_wavef_send)
1897 : :
1898 : :
1899 : : !-------------------------------------------------
1900 : : !Calcul of the number of received and sended datas
1901 : : !-------------------------------------------------
1902 [ + - - + : 31854 : sendcounts_sym_loc = sendcounts_sym*2
- - + + ]
1903 [ + - - + : 31854 : recvcounts_sym_loc = recvcounts_sym*2
- - + + ]
1904 : :
1905 : : !------------------------------------------
1906 : : !Exchange of the datas ewavef_allto_all_loc
1907 : : !------------------------------------------
1908 [ + + ]: 32494 : do ibandpp=1,bandpp_sym
1909 : :
1910 : : ! ------------------------------------------------
1911 : : ! Deplacment of the sended datas because of bandpp
1912 : : ! ------------------------------------------------
1913 [ + + ]: 49148 : sdispls_sym_loc(:) = sdispls_sym(:) + ndatasend_sym * (ibandpp-1)
1914 [ + + ]: 49148 : sdispls_sym_loc = sdispls_sym_loc *2
1915 : :
1916 : : ! --------------------------------------------------
1917 : : ! Deplacment of the received datas because of bandpp
1918 : : ! --------------------------------------------------
1919 [ + + ]: 49148 : rdispls_sym_loc(:) = rdispls_sym(:) + ndatarecv_tot * (ibandpp-1)
1920 [ + + ]: 49148 : rdispls_sym_loc = rdispls_sym_loc *2
1921 : :
1922 : :
1923 : : call xmpi_alltoallv(&
1924 : : & ewavef_alltoall_send(:,:) ,sendcounts_sym_loc,sdispls_sym_loc,&
1925 : : & ewavef_alltoall_sym(:,:) ,recvcounts_sym_loc,rdispls_sym_loc,&
1926 : 32494 : & newspacecomm,ier)
1927 : :
1928 : : end do
1929 : :
1930 : : !-----------------------
1931 : : !Desallocation
1932 : : !-----------------------
1933 : :
1934 [ - + ]: 9896 : ABI_FREE(sendcounts_sym_loc)
1935 [ - + ]: 9896 : ABI_FREE(recvcounts_sym_loc)
1936 [ - + ]: 9896 : ABI_FREE(sdispls_sym_loc)
1937 [ - + ]: 9896 : ABI_FREE(rdispls_sym_loc)
1938 : :
1939 [ - + ]: 9896 : ABI_FREE(ewavef_alltoall_loc)
1940 [ - + ]: 9896 : ABI_FREE(ewavef_alltoall_send)
1941 : :
1942 : 9896 : end subroutine prep_wavef_sym_do
1943 : : !!***
1944 : :
1945 : : !!****f* ABINIT/prep_wavef_sym_undo
1946 : : !! NAME
1947 : : !! prep_wavef_sym_undo
1948 : : !!
1949 : : !! FUNCTION
1950 : : !! this routine dissociates each wave function in two waves functions as following
1951 : : !! C(G) = ( E*(-G) + E(G))/2
1952 : : !! D(G) = i*( E*(-G) - E(G))/2
1953 : : !! the values are redistributed on the processors in function of
1954 : : !! the value of mpi_enreg%distribfft%tab_fftwf2_distrib( (-kg_k_gather(2,i) )
1955 : : !!
1956 : : !! INPUTS
1957 : : !! mpi_enreg = information about mpi parallelization
1958 : : !! bandpp = number of groups of couple of waves functions
1959 : : !! nspinor = number of spin
1960 : : !! ndatarecv = number of values received by the processor and sended
1961 : : !! by the other processors band
1962 : : !! ndatarecv_tot = total number of received values
1963 : : !! (ndatarecv + number of received opposited planewave coordinates)
1964 : : !! ndatasend_sym = number of sended values to the processors fft to create opposited
1965 : : !! planewave coordinates
1966 : : !! idatarecv0 = position of the planewave coordinates (0,0,0)
1967 : : !! sendcounts_sym = number of sended values by the processor to each processor fft
1968 : : !! sdispls_sym = postions of the sended values by the processor to each processor fft
1969 : : !!
1970 : : !! recvcounts_sym = number of the received values by the processor to each processor fft
1971 : : !!! rdispls_sym = postions of the received values by the processor to each processor fft
1972 : : !!
1973 : : !! gwavef_alltoall_sym = planewave coefficients of wavefunction
1974 : : !! initial of the processor +
1975 : : !! sended by other processors band +
1976 : : !! sended by other processors fft +
1977 : : !! and composited if bandpp >1
1978 : : !! index_wavef_send = index to send the values by block to the other processor fft
1979 : : !!
1980 : : !! OUTPUT
1981 : : !! gwavef_alltoall = planewave coefficients of wavefunction
1982 : : !! ( for of the processor + to send to other processors band)
1983 : : !!
1984 : : !! SOURCE
1985 : :
1986 : 12576 : subroutine prep_wavef_sym_undo(mpi_enreg,bandpp,nspinor,&
1987 : : & ndatarecv,&
1988 : : & ndatarecv_tot,ndatasend_sym,idatarecv0,&
1989 : 12576 : & gwavef_alltoall,&
1990 [ - + ]: 12576 : & sendcounts_sym,sdispls_sym,&
1991 [ - + - + : 12576 : & recvcounts_sym,rdispls_sym,&
- + ]
1992 [ - + ]: 12576 : & gwavef_alltoall_sym,&
1993 [ - + ]: 12576 : & index_wavef_send)
1994 : :
1995 : : !Arguments ------------------------------------
1996 : : !scalars
1997 : : integer,intent(in) :: bandpp,idatarecv0,ndatarecv,ndatarecv_tot,ndatasend_sym
1998 : : integer,intent(in) :: nspinor
1999 : : type(mpi_type),intent(in) :: mpi_enreg
2000 : : !arrays
2001 : : integer,intent(in) :: index_wavef_send(:),rdispls_sym(:),recvcounts_sym(:)
2002 : : integer,intent(in) :: sdispls_sym(:),sendcounts_sym(:)
2003 : : real(dp),intent(inout) :: gwavef_alltoall(2,ndatarecv*nspinor*bandpp)
2004 : : real(dp),intent(inout) :: gwavef_alltoall_sym(:,:)
2005 : :
2006 : : !Local variables-------------------------------
2007 : : !scalars
2008 : : integer :: bandpp_sym,ibandpp,ideb_loc,idebc,idebd
2009 : : integer :: idebe,ier,ifin_loc,ifinc,ifind,ifine
2010 : : integer :: jbandpp,kbandpp,newspacecomm,nproc_fft
2011 : : logical :: flag_compose
2012 : : !arrays
2013 : : integer,allocatable :: rdispls_sym_loc(:),recvcounts_sym_loc(:)
2014 : : integer,allocatable :: sdispls_sym_loc(:),sendcounts_sym_loc(:)
2015 : 12576 : real(dp),allocatable :: gwavef_alltoall_loc(:,:),gwavef_alltoall_rcv(:,:)
2016 : :
2017 : : ! *********************************************************************
2018 : :
2019 : : !DEBUG
2020 : : !write(std_out,*)' prep_wavef_sym_undo : enter '
2021 : : !ENDDEBUG
2022 : :
2023 : :
2024 : : !---------------------------------------------
2025 : : !Initialisation
2026 : : !---------------------------------------------
2027 : 12576 : nproc_fft = mpi_enreg%nproc_fft
2028 : :
2029 : 12576 : newspacecomm = mpi_enreg%comm_fft
2030 : :
2031 [ + + ]: 12576 : if (modulo(bandpp,2)==0) then
2032 : 7488 : bandpp_sym = bandpp/2
2033 : 7488 : flag_compose = .TRUE.
2034 : : else
2035 : : bandpp_sym = bandpp
2036 : : flag_compose = .FALSE.
2037 : : end if
2038 : :
2039 : : !---------------------------------------------
2040 : : !Allocation
2041 : : !---------------------------------------------
2042 [ - + + - : 37728 : ABI_MALLOC(gwavef_alltoall_loc ,(2,ndatarecv *bandpp_sym))
+ - + - -
+ - + ]
2043 [ - + + - : 37728 : ABI_MALLOC(gwavef_alltoall_rcv ,(2,ndatasend_sym *bandpp_sym))
+ - + - -
+ - + ]
2044 : :
2045 [ - + + - : 37728 : ABI_MALLOC(sendcounts_sym_loc ,(nproc_fft))
+ - + - -
+ - + ]
2046 [ - + + - : 25152 : ABI_MALLOC(sdispls_sym_loc ,(nproc_fft))
- + - + ]
2047 [ - + + - : 25152 : ABI_MALLOC(recvcounts_sym_loc ,(nproc_fft))
- + - + ]
2048 [ - + + - : 25152 : ABI_MALLOC(rdispls_sym_loc ,(nproc_fft))
- + - + ]
2049 : :
2050 : :
2051 : : !---------------------------------------------
2052 : : !Initialisation
2053 : : !---------------------------------------------
2054 [ + + + + ]: 67046016 : gwavef_alltoall_loc(:,:) =0.
2055 : :
2056 [ + + ]: 27966 : sendcounts_sym_loc(:) =0
2057 [ + + ]: 27966 : sdispls_sym_loc(:) =0
2058 [ + + ]: 27966 : recvcounts_sym_loc(:) =0
2059 [ + + ]: 27966 : rdispls_sym_loc(:) =0
2060 : :
2061 : :
2062 : : !-------------------------------------------------
2063 : : !Calcul of number of the sended and received datas
2064 : : !-------------------------------------------------
2065 [ + - - + : 40542 : sendcounts_sym_loc = sendcounts_sym*2
- - + + ]
2066 [ + - - + : 40542 : recvcounts_sym_loc = recvcounts_sym*2
- - + + ]
2067 : :
2068 : : !----------------------------------------------------
2069 : : !Sending of the values
2070 : : !----------------------------------------------------
2071 [ + + ]: 39930 : do ibandpp = 1,bandpp_sym
2072 : :
2073 : : ! -------------------------------------------------
2074 : : ! Deplacment of the sended values because of bandpp
2075 : : ! -------------------------------------------------
2076 [ + + ]: 60632 : sdispls_sym_loc(:) = sdispls_sym(:) + ndatasend_sym * (ibandpp-1)
2077 [ + + ]: 60632 : sdispls_sym_loc = sdispls_sym_loc *2
2078 : :
2079 : : ! ---------------------------------------------------
2080 : : ! Deplacment of the received values because of bandpp
2081 : : ! ---------------------------------------------------
2082 [ + + ]: 60632 : rdispls_sym_loc(:) = rdispls_sym(:) + ndatarecv_tot * (ibandpp-1)
2083 [ + + ]: 60632 : rdispls_sym_loc = rdispls_sym_loc *2
2084 : :
2085 : :
2086 : : call xmpi_alltoallv(&
2087 : : & gwavef_alltoall_sym(:,:) ,recvcounts_sym_loc,rdispls_sym_loc,&
2088 : : & gwavef_alltoall_rcv(:,:) ,sendcounts_sym_loc,sdispls_sym_loc,&
2089 : 39930 : & newspacecomm,ier)
2090 : :
2091 : : end do
2092 : :
2093 : :
2094 : : !----------------------
2095 : : !Dispatching the blocks
2096 : : !----------------------
2097 [ + + + + ]: 66972840 : gwavef_alltoall_loc(:,index_wavef_send(:)) = gwavef_alltoall_rcv(:,:)
2098 : :
2099 : : !----------------------
2100 : : !Case -kg = [0 0 0]
2101 : : !----------------------
2102 [ + + ]: 12576 : if (idatarecv0/=-1) then
2103 [ + + ]: 35561 : do kbandpp=1,bandpp_sym
2104 : : gwavef_alltoall_loc(:,(kbandpp-1)*ndatarecv + idatarecv0)= &
2105 [ + + ]: 84345 : gwavef_alltoall_sym(:,(kbandpp-1)*ndatarecv_tot + idatarecv0)
2106 : : end do
2107 : : end if
2108 : :
2109 : : !---------------------------------------------------
2110 : : !Build of hwavef_alltoall
2111 : : !
2112 : : !We have got :
2113 : : !bandpp_sym blocks to dissociate
2114 : : !or bandpp_sym blokcs to not dissociate
2115 : : !--------------------------------------------------
2116 [ + + ]: 39930 : do kbandpp=1,bandpp_sym
2117 : :
2118 : : ! position of the 2 blocks
2119 : : ! ----------------------------------
2120 : 27354 : ibandpp = (kbandpp-1) * 2
2121 : 27354 : jbandpp = ibandpp + 1
2122 : :
2123 : 27354 : idebe = (kbandpp-1) * ndatarecv_tot + 1
2124 : 27354 : ifine = idebe + ndatarecv - 1
2125 : :
2126 : 27354 : idebc = ibandpp * ndatarecv + 1
2127 : 27354 : ifinc = idebc + ndatarecv - 1
2128 : :
2129 : 27354 : idebd = jbandpp * ndatarecv + 1
2130 : 27354 : ifind = idebd + ndatarecv - 1
2131 : :
2132 : 27354 : ideb_loc = (kbandpp-1) * ndatarecv + 1
2133 : 27354 : ifin_loc = ideb_loc + ndatarecv - 1
2134 : :
2135 : :
2136 [ + + ]: 39930 : if (flag_compose) then
2137 : :
2138 : : ! calcul cwavef(G)
2139 : : ! ----------------
2140 : : gwavef_alltoall(1,idebc:ifinc) = gwavef_alltoall_sym(1,idebe:ifine) &
2141 [ + + ]: 21340026 : & + gwavef_alltoall_loc(1,ideb_loc:ifin_loc)
2142 : : gwavef_alltoall(2,idebc:ifinc) = gwavef_alltoall_sym(2,idebe:ifine) &
2143 [ + + ]: 21340026 : & - gwavef_alltoall_loc(2,ideb_loc:ifin_loc)
2144 : :
2145 : : ! calcul dwavef(G)
2146 : : ! ------------------
2147 : : gwavef_alltoall(1,idebd:ifind) = gwavef_alltoall_sym(2,idebe:ifine) &
2148 [ + + ]: 21340026 : & + gwavef_alltoall_loc(2,ideb_loc:ifin_loc)
2149 : : gwavef_alltoall(2,idebd:ifind) = - gwavef_alltoall_sym(1,idebe:ifine) &
2150 [ + + ]: 21340026 : & + gwavef_alltoall_loc(1,ideb_loc:ifin_loc)
2151 : : else
2152 : :
2153 : : ! calcul cwavef(G)
2154 : : ! ----------------
2155 : : gwavef_alltoall(1,idebc:ifinc) = gwavef_alltoall_sym(1,idebe:ifine) &
2156 [ + + ]: 1031808 : & + gwavef_alltoall_loc(1,ideb_loc:ifin_loc)
2157 : : gwavef_alltoall(2,idebc:ifinc) = gwavef_alltoall_sym(2,idebe:ifine) &
2158 [ + + ]: 1031808 : & - gwavef_alltoall_loc(2,ideb_loc:ifin_loc)
2159 : : end if
2160 : :
2161 : : end do
2162 : :
2163 : : !We divise by two
2164 [ + + + + ]: 130999296 : gwavef_alltoall(:,:) = gwavef_alltoall(:,:)/2
2165 : :
2166 : : !-----------------------
2167 : : !Desallocation
2168 : : !-----------------------
2169 : :
2170 [ - + ]: 12576 : ABI_FREE(sendcounts_sym_loc)
2171 [ - + ]: 12576 : ABI_FREE(recvcounts_sym_loc)
2172 [ - + ]: 12576 : ABI_FREE(sdispls_sym_loc)
2173 [ - + ]: 12576 : ABI_FREE(rdispls_sym_loc)
2174 : :
2175 [ - + ]: 12576 : ABI_FREE(gwavef_alltoall_loc)
2176 [ - + ]: 12576 : ABI_FREE(gwavef_alltoall_rcv)
2177 : :
2178 : 12576 : end subroutine prep_wavef_sym_undo
2179 : : !!***
2180 : :
2181 : : !!****f* ABINIT/prep_index_wavef_bandpp
2182 : : !! NAME
2183 : : !! prep_index_wavef_bandpp
2184 : : !!
2185 : : !! FUNCTION
2186 : : !! this routine sorts the waves functions by bandpp and by processors
2187 : : !! after the alltoall
2188 : : !!
2189 : : !! INPUTS
2190 : : !! nproc_band = number of processors below the band
2191 : : !! bandpp = number of groups of couple of waves functions
2192 : : !! nspinor = number of spin
2193 : : !! ndatarecv = total number of values received by the processor and sended
2194 : : !! by the other processors band
2195 : : !! recvcounts = number of values sended by each processor band and received
2196 : : !! by the processor
2197 : : !! rdispls = positions of the values received by the processor and
2198 : : !! sended by each processor band
2199 : : !!
2200 : : !! OUTPUT
2201 : : !! index_wavef_band = position of the sorted values
2202 : : !!
2203 : : !! SOURCE
2204 : :
2205 : 72596 : subroutine prep_index_wavef_bandpp(nproc_band,bandpp,&
2206 : : nspinor,ndatarecv,&
2207 : 72596 : recvcounts,rdispls,&
2208 : : index_wavef_band)
2209 : :
2210 : : !Arguments ------------------------------------
2211 : : !scalars
2212 : : integer,intent(in) :: bandpp,ndatarecv,nproc_band,nspinor
2213 : : !arrays
2214 : : integer,intent(in) :: rdispls(nproc_band),recvcounts(nproc_band)
2215 : : integer,allocatable,intent(out) :: index_wavef_band(:)
2216 : :
2217 : : !Local variables-------------------------------
2218 : : !scalars
2219 : : integer :: delta,idebc,idebe,ifinc,ifine,iindex,iproc,kbandpp,nb
2220 : :
2221 : : ! *********************************************************************
2222 : :
2223 : : !DEBUG
2224 : : !write(std_out,*)' prep_index_wavef_banpp : enter '
2225 : : !write(std_out,*) 'ndatarecv = ', ndatarecv
2226 : : !write(std_out,*) 'rdispls(:) = ', rdispls(:)
2227 : : !write(std_out,*) 'recvcounts(:) = ', recvcounts(:)
2228 : : !ENDDEBUG
2229 : :
2230 : :
2231 : : !---------------------------------------------
2232 : : !Allocation
2233 : : !---------------------------------------------
2234 [ - + + - : 217788 : ABI_MALLOC(index_wavef_band ,(bandpp*nspinor*ndatarecv))
+ - + - -
+ - + -
+ ]
2235 [ + + ]: 161806596 : index_wavef_band(:) =0
2236 : :
2237 : : !---------------------------------------------
2238 : : !Calcul : loops on bandpp and processors band
2239 : : !---------------------------------------------
2240 [ + + ]: 315916 : nb = sum(recvcounts(1:nproc_band))
2241 [ + + ]: 365306 : do kbandpp=1,bandpp
2242 : :
2243 [ + + ]: 1215584 : do iproc=1,nproc_band
2244 : :
2245 : 850278 : idebe = (rdispls(iproc) + 1) + (kbandpp-1) * ndatarecv*nspinor
2246 : 850278 : ifine = idebe + recvcounts(iproc) -1
2247 : :
2248 [ + + ]: 850278 : if (iproc==1) then
2249 : 292710 : idebc = (kbandpp-1)* recvcounts(iproc)*nspinor + 1
2250 : : else
2251 : : idebc = (bandpp) * sum(recvcounts(1:iproc-1))*nspinor &
2252 : : + (kbandpp-1)* recvcounts(iproc)*nspinor &
2253 [ + + ]: 1540848 : + 1
2254 : : end if
2255 : 850278 : ifinc = idebc + recvcounts(iproc) -1
2256 [ + - + + : 320070700 : index_wavef_band(idebe:ifine) = (/( iindex,iindex=idebc,ifinc)/)
+ + ]
2257 : 850278 : delta=ifine-idebe
2258 [ + + ]: 1142988 : if (nspinor==2) then
2259 [ + + + + ]: 5250560 : index_wavef_band(idebe+nb :idebe+nb +delta)=(/( iindex,iindex=ifinc+1,ifinc+1+delta)/)
2260 : : end if
2261 : : end do
2262 : : end do
2263 : :
2264 : 72596 : end subroutine prep_index_wavef_bandpp
2265 : : !!***
2266 : :
2267 : : !!****f* ABINIT/prep_sort_wavef_spin
2268 : : !! NAME
2269 : : !! prep_sort_wavef_spin
2270 : : !!
2271 : : !! FUNCTION
2272 : : !! Compute index used to sort a spinorial wave-function by spin
2273 : : !! Sort to have all nspinor=1 fisrt, then all nspinor=2
2274 : : !!
2275 : : !! INPUTS
2276 : : !! nproc_band=size of "band" communicator
2277 : : !! nspinor=number of spinorial components of the wavefunction
2278 : : !! ndatarecv=total number of values on all processors
2279 : : !! recvcounts(nproc_band)= number of received values by the processor
2280 : : !! rdispls(nproc_band)= offsets of the received values by the processor
2281 : : !!
2282 : : !! OUTPUT
2283 : : !! index_wavef(:)=array containing the sorted indexes (pointer, allocated in this routine)
2284 : : !!
2285 : : !! SOURCE
2286 : :
2287 : 7368 : subroutine prep_sort_wavef_spin(nproc_band,nspinor,ndatarecv,recvcounts,rdispls,index_wavef)
2288 : :
2289 : : !Arguments ------------------------------------
2290 : : !scalars
2291 : : integer,intent(in) :: ndatarecv,nproc_band,nspinor
2292 : : !arrays
2293 : : integer,intent(in) :: rdispls(nproc_band),recvcounts(nproc_band)
2294 : : integer,allocatable,intent(out) :: index_wavef(:)
2295 : :
2296 : : !Local variables-------------------------------
2297 : : !scalars
2298 : : integer :: isft,isft1,iproc,iindex
2299 : : !arrays
2300 : 7368 : integer,allocatable :: recvcountsloc(:),rdisplsloc(:)
2301 : :
2302 : : ! *********************************************************************
2303 : :
2304 [ - + + - : 22104 : ABI_MALLOC(index_wavef,(ndatarecv*nspinor))
+ - + - -
+ - + -
+ ]
2305 : :
2306 [ - + + - : 22104 : ABI_MALLOC(recvcountsloc,(nproc_band))
+ - + - -
+ - + ]
2307 [ - + + - : 14736 : ABI_MALLOC(rdisplsloc,(nproc_band))
- + + - ]
2308 [ + + ]: 22584 : recvcountsloc(:)=recvcounts(:)*2*nspinor
2309 [ + + ]: 22584 : rdisplsloc(:)=rdispls(:)*2*nspinor
2310 : :
2311 : : !---------------------------------------------
2312 : : !Loops on bandpp and processors band
2313 : : !---------------------------------------------
2314 : 22584 : isft=0
2315 [ + + ]: 22584 : do iproc=1,nproc_band
2316 : :
2317 : : ! ===== Spin up
2318 [ + + ]: 15216 : if (iproc==1) then
2319 : : isft= 0
2320 : : else
2321 [ + + ]: 23544 : isft= sum(recvcounts(1: (iproc-1)))
2322 : : end if
2323 : 15216 : isft1 = 0.5*rdisplsloc(iproc)
2324 : :
2325 : : index_wavef(1+isft:isft+recvcounts(iproc))= &
2326 [ + - + + : 2149520 : & (/(iindex,iindex=isft1+1,isft1+recvcounts(iproc))/)
+ + ]
2327 : :
2328 : : ! =====Spin down
2329 [ + + ]: 15216 : if (iproc==1)then
2330 [ + + ]: 22584 : isft=sum(recvcounts(1:nproc_band))
2331 : : else
2332 : : isft=sum(recvcounts(1:nproc_band)) &
2333 [ + + + + ]: 54936 : & +sum(recvcounts(1:iproc-1))
2334 : : end if
2335 : 15216 : isft1 = 0.5 * rdisplsloc(iproc) + recvcounts(iproc)
2336 : :
2337 : : index_wavef(1+isft:isft+recvcounts(iproc))= &
2338 [ + + + + ]: 2156888 : & (/(iindex,iindex=isft1+1,isft1+ recvcounts(iproc))/)
2339 : :
2340 : : end do
2341 : :
2342 [ - + ]: 7368 : ABI_FREE(recvcountsloc)
2343 [ - + ]: 7368 : ABI_FREE(rdisplsloc)
2344 : :
2345 : 7368 : end subroutine prep_sort_wavef_spin
2346 : : !!***
2347 : :
2348 : : end module m_prep_kgb
2349 : : !!***
|
