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 264842 : subroutine prep_nonlop(choice,cpopt,cwaveprj,enlout_block,hamk,idir,lambdablock,&
636 264842 : blocksize,mpi_enreg,nnlout,paw_opt,signs,gsc, tim_nonlop,cwavef,gvnlc, &
637 264842 : 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 264842 : integer, allocatable :: index_wavef_band(:)
664 264842 : integer, allocatable :: rdisplsloc(:),recvcountsloc(:),sdisplsloc(:),sendcountsloc(:)
665 264842 : integer,contiguous, pointer :: rdispls(:),recvcounts(:),sdispls(:),sendcounts(:)
666 529684 : 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 264842 : real(dp), pointer :: cwavef_alltoall2(:,:)
678 264842 : real(dp), allocatable :: gvnlc_alltoall2(:,:)
679 264842 : real(dp), allocatable :: gsc_alltoall2(:,:)
680 : #endif
681 :
682 264842 : real(dp), allocatable :: cwavef_alltoall1(:,:)
683 264842 : real(dp), allocatable :: gvnlc_alltoall1(:,:)
684 264842 : real(dp), allocatable :: gsc_alltoall1(:,:)
685 264842 : 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 264842 : call timab(570,1,tsec)
704 :
705 264842 : do_transpose = .true.
706 264842 : bandpp = mpi_enreg%bandpp
707 264842 : if(present(already_transposed)) then
708 264842 : if(already_transposed) then
709 124920 : do_transpose = .false.
710 124920 : bandpp = blocksize
711 : end if
712 : end if
713 :
714 264842 : l_gpu_option=ABI_GPU_DISABLED
715 264842 : if (present(gpu_option)) then
716 132454 : l_gpu_option = gpu_option
717 : end if
718 :
719 264842 : nproc_band = mpi_enreg%nproc_band
720 :
721 264842 : spaceComm=mpi_enreg%comm_fft
722 264842 : if(mpi_enreg%paral_kgb==1) spaceComm=mpi_enreg%comm_band
723 264842 : my_nspinor=max(1,hamk%nspinor/mpi_enreg%nproc_spinor)
724 264842 : nspinortot=hamk%nspinor
725 :
726 : !Check sizes
727 264842 : npw=hamk%npw_k;if (.not.do_transpose) npw=hamk%npw_fft_k
728 794526 : if (size(cwavef)/=2*npw*my_nspinor*blocksize) then
729 0 : ABI_BUG('Incorrect size for cwavef!')
730 : end if
731 264842 : 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 264842 : if(cpopt>=0.and. .not. present(vectproj)) then
744 226404 : 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 264842 : 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 794526 : ABI_MALLOC(sendcountsloc,(nproc_band))
758 529684 : ABI_MALLOC(sdisplsloc ,(nproc_band))
759 529684 : ABI_MALLOC(recvcountsloc,(nproc_band))
760 529684 : ABI_MALLOC(rdisplsloc ,(nproc_band))
761 :
762 264842 : ikpt_this_proc=bandfft_kpt_get_ikpt()
763 :
764 264842 : recvcounts => bandfft_kpt(ikpt_this_proc)%recvcounts(:)
765 264842 : sendcounts => bandfft_kpt(ikpt_this_proc)%sendcounts(:)
766 264842 : rdispls => bandfft_kpt(ikpt_this_proc)%rdispls (:)
767 264842 : sdispls => bandfft_kpt(ikpt_this_proc)%sdispls (:)
768 264842 : ndatarecv = bandfft_kpt(ikpt_this_proc)%ndatarecv
769 :
770 264842 : 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 264842 : 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 794334 : ABI_MALLOC(cwavef_alltoall2, (2,ndatarecv*my_nspinor*bandpp))
794 : endif
795 794526 : ABI_MALLOC(gsc_alltoall2, (2,ndatarecv*my_nspinor*(paw_opt/3)*bandpp))
796 794526 : 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 264842 : if(do_transpose .and. (bandpp/=1 .or. (bandpp==1 .and. mpi_enreg%paral_spinor==0.and.nspinortot==2)))then
804 61986 : 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 20662 : 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 794526 : ABI_MALLOC(enlout,(nnlout*bandpp))
819 1348482 : enlout = zero
820 :
821 923600 : recvcountsloc(:)=recvcounts(:)*2*my_nspinor*bandpp
822 923600 : rdisplsloc(:)=rdispls(:)*2*my_nspinor*bandpp
823 923600 : sendcountsloc(:)=sendcounts(:)*2*my_nspinor
824 923600 : sdisplsloc(:)=sdispls(:)*2*my_nspinor
825 :
826 264842 : if(do_transpose) then
827 139922 : call timab(581,1,tsec)
828 139922 : if (bandpp/=1 .or. (bandpp==1 .and. mpi_enreg%paral_spinor==0.and.nspinortot==2)) then
829 20662 : 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 20662 : & 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 139922 : 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 264842 : 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 137254 : if(do_transpose) then
918 : call prep_index_wavef_bandpp(nproc_band,bandpp,&
919 13294 : & my_nspinor,ndatarecv,recvcounts,rdispls,index_wavef_band)
920 :
921 : ! -------------------------------------------------------
922 : ! Sorting of the waves functions below bandpp
923 : ! -------------------------------------------------------
924 13294 : 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 126494086 : cwavef_alltoall2(:,:) = cwavef_alltoall1(:,index_wavef_band)
939 : end if
940 : end if
941 :
942 : ! -------------------------------------------------------
943 : ! Call nonlop
944 : ! -------------------------------------------------------
945 137254 : if(paw_opt == 2) then
946 46304 : 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 150548 : & 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 137254 : 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 264842 : if (allocated(index_wavef_band)) then
974 20662 : 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 264842 : if(do_transpose) then
980 139922 : 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 264842 : if (nnlout>0) then
1015 107492 : call xmpi_allgather(enlout,nnlout*bandpp,enlout_block,spaceComm,ier)
1016 : end if
1017 264842 : ABI_FREE(enlout)
1018 264842 : ABI_FREE(sendcountsloc)
1019 264842 : ABI_FREE(sdisplsloc)
1020 264842 : ABI_FREE(recvcountsloc)
1021 264842 : 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 264842 : 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 264842 : if(present(cwavef_tr)) then
1040 64 : nullify(cwavef_alltoall2)
1041 : else
1042 264778 : ABI_FREE(cwavef_alltoall2)
1043 : endif
1044 264842 : ABI_FREE(gvnlc_alltoall2)
1045 264842 : ABI_FREE(gsc_alltoall2)
1046 : end if
1047 :
1048 264842 : 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 20662 : ABI_FREE(cwavef_alltoall1)
1053 20662 : 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 264842 : call timab(570,2,tsec)
1064 :
1065 : DBG_EXIT('COLL')
1066 :
1067 529684 : 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 182182 : subroutine prep_fourwf(rhoaug,blocksize,cwavef,wfraug,iblock,istwf_k,mgfft,&
1113 182182 : & 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 182182 : integer,contiguous, pointer :: indices_pw_fft(:),kg_k_fft(:,:),kg_k_gather(:,:),kg_k_gather_sym(:,:)
1144 182182 : integer,contiguous, pointer :: rdispls(:),rdispls_sym(:)
1145 182182 : integer,contiguous, pointer :: recvcounts(:),recvcount_fft(:),recvcounts_sym(:),recvcounts_sym_tot(:)
1146 182182 : integer,contiguous, pointer :: recvdisp_fft(:),sdispls(:),sdispls_sym(:)
1147 182182 : integer,contiguous, pointer :: sendcounts(:),sendcount_fft(:),sendcounts_sym(:),sendcounts_sym_all(:)
1148 182182 : integer,contiguous, pointer :: senddisp_fft(:),tab_proc(:)
1149 182182 : integer,allocatable :: rdisplsloc(:)
1150 182182 : integer,allocatable :: recvcountsloc(:),sdisplsloc(:)
1151 182182 : integer,allocatable :: sendcountsloc(:)
1152 182182 : integer,allocatable :: index_wavef_band(:),index_wavef_send(:)
1153 182182 : integer,pointer :: gbound_(:,:)
1154 : real(dp) :: dummy(2,1),tsec(2)
1155 182182 : 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 182182 : real(dp),allocatable :: cwavef_alltoall1(:,:)
1161 : #endif
1162 182182 : real(dp),allocatable :: cwavef_alltoall2(:,:)
1163 182182 : real(dp),allocatable :: cwavef_fft(:,:), cwavef_fft_tr(:,:)
1164 182182 : real(dp),allocatable :: weight_t(:),weight1_t(:),weight2_t(:)
1165 182182 : 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 182182 : ABI_CHECK((nfft_blocks>0),'nfft_blocks is null')
1183 182182 : ABI_CHECK((mpi_enreg%bandpp==ndat),'BUG: bandpp/=ndat')
1184 :
1185 182182 : spaceComm=mpi_enreg%comm_band
1186 182182 : nproc_band = mpi_enreg%nproc_band
1187 182182 : nproc_fft = mpi_enreg%nproc_fft
1188 182182 : bandpp = mpi_enreg%bandpp
1189 182182 : me_fft = mpi_enreg%me_fft
1190 :
1191 182182 : gpu_option_=ABI_GPU_DISABLED;if (present(gpu_option)) gpu_option_=gpu_option
1192 :
1193 182182 : if (present(bandfft_kpt_tab)) then
1194 : bandfft_kpt_ptr => bandfft_kpt_tab
1195 : else
1196 182182 : ikpt_this_proc=bandfft_kpt_get_ikpt()
1197 182182 : bandfft_kpt_ptr => bandfft_kpt(ikpt_this_proc)
1198 : end if
1199 :
1200 182182 : have_to_reequilibrate = bandfft_kpt_ptr%have_to_reequilibrate
1201 182182 : istwf_k_=istwf_k
1202 182182 : flag_inv_sym = (istwf_k_==2 .and. any(ngfft(7) == [401,402,312,512]))
1203 182182 : if (option_fourwf==0) flag_inv_sym=((flag_inv_sym).and.(gpu_option_==ABI_GPU_DISABLED))
1204 :
1205 182182 : 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 182182 : 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 546546 : ABI_MALLOC(sendcountsloc,(nproc_band))
1222 364364 : ABI_MALLOC(sdisplsloc ,(nproc_band))
1223 364364 : ABI_MALLOC(recvcountsloc,(nproc_band))
1224 364364 : ABI_MALLOC(rdisplsloc ,(nproc_band))
1225 :
1226 182182 : recvcounts =>bandfft_kpt_ptr%recvcounts(:)
1227 182182 : sendcounts =>bandfft_kpt_ptr%sendcounts(:)
1228 182182 : rdispls =>bandfft_kpt_ptr%rdispls (:)
1229 182182 : sdispls =>bandfft_kpt_ptr%sdispls (:)
1230 182182 : ndatarecv =>bandfft_kpt_ptr%ndatarecv
1231 :
1232 182182 : kg_k_gather =>bandfft_kpt_ptr%kg_k_gather(:,:)
1233 182182 : gbound_ =>bandfft_kpt_ptr%gbound(:,:)
1234 :
1235 182182 : 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 546546 : ABI_MALLOC(cwavef_alltoall2,(2,ndatarecv*bandpp))
1250 182182 : if ( ((.not.flag_inv_sym) .and. (bandpp>1) ) .or. flag_inv_sym )then
1251 37286 : 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 219468 : ABI_MALLOC(cwavef_alltoall1,(2,ndatarecv*bandpp))
1257 : end if
1258 : end if
1259 :
1260 599176 : recvcountsloc(:)=recvcounts(:)*2*bandpp
1261 599176 : rdisplsloc(:)=rdispls(:)*2*bandpp
1262 599176 : sendcountsloc(:)=sendcounts(:)*2
1263 599176 : sdisplsloc(:)=sdispls(:)*2
1264 :
1265 182182 : transfer_cwavef=.false.
1266 182182 : if(.not. xomp_target_is_present(c_loc(cwavef)) .and. gpu_option_==ABI_GPU_OPENMP) then
1267 182182 : 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 182182 : call timab(547,1,tsec)
1279 182182 : 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 182182 : & 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 182182 : call timab(547,2,tsec)
1301 :
1302 182182 : tim_fourwf=16
1303 :
1304 : !Eventually adjust load balancing for FFT (by changing FFT distrib)
1305 182182 : 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 182182 : if (option_fourwf==0) wfraug(:,:,:,:)=zero
1321 :
1322 : !====================================================================
1323 182182 : 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 37286 : 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 30664 : & index_wavef_band)
1370 :
1371 : ! -------------------------------------------------------
1372 : ! Sorting of the wave functions below bandpp
1373 : ! -------------------------------------------------------
1374 30664 : 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 98808688 : cwavef_alltoall1(:,:) = cwavef_alltoall2(:,index_wavef_band)
1380 : end if
1381 :
1382 30664 : 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 30664 : 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 145944 : do iibandpp=1,bandpp
1426 : ! Compute the index of the band
1427 115280 : ind_occ = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + iibandpp
1428 : ! Compute the weight of the band
1429 115280 : weight=occ_k(ind_occ)*wtk/ucvol
1430 115280 : if (option_fourwf==0) then
1431 0 : wfraug_ptr => wfraug(:,:,:,(iibandpp-1)*n6+1:iibandpp*n6)
1432 : else
1433 115280 : wfraug_ptr => wfraug
1434 : end if
1435 145944 : if (abs(occ_k(ind_occ)) >=tol8.or.option_fourwf==0) then
1436 55611 : 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 55611 : & tim_fourwf,weight,weight)
1450 : end if
1451 55611 : 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 30664 : 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 182182 : 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 182182 : ABI_FREE(sendcountsloc)
1621 182182 : ABI_FREE(sdisplsloc)
1622 182182 : ABI_FREE(recvcountsloc)
1623 182182 : ABI_FREE(rdisplsloc)
1624 182182 : ABI_FREE(cwavef_alltoall2)
1625 401650 : if ( ((.not.flag_inv_sym) .and. (bandpp>1) ) .or. flag_inv_sym ) then
1626 37286 : 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 37286 : ABI_FREE(cwavef_alltoall1)
1632 : end if
1633 : end if
1634 :
1635 : contains
1636 37286 : 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 72396 : subroutine prep_index_wavef_bandpp(nproc_band,bandpp,&
2206 : nspinor,ndatarecv,&
2207 72396 : 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 217188 : ABI_MALLOC(index_wavef_band ,(bandpp*nspinor*ndatarecv))
2235 157382596 : index_wavef_band(:) =0
2236 :
2237 : !---------------------------------------------
2238 : !Calcul : loops on bandpp and processors band
2239 : !---------------------------------------------
2240 315316 : nb = sum(recvcounts(1:nproc_band))
2241 364706 : do kbandpp=1,bandpp
2242 :
2243 1214184 : do iproc=1,nproc_band
2244 :
2245 849478 : idebe = (rdispls(iproc) + 1) + (kbandpp-1) * ndatarecv*nspinor
2246 849478 : ifine = idebe + recvcounts(iproc) -1
2247 :
2248 849478 : if (iproc==1) then
2249 292310 : 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 1540048 : + 1
2254 : end if
2255 849478 : ifinc = idebc + recvcounts(iproc) -1
2256 311221500 : index_wavef_band(idebe:ifine) = (/( iindex,iindex=idebc,ifinc)/)
2257 849478 : delta=ifine-idebe
2258 1141788 : 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 72396 : 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 : !!***
|
