Line data Source code
1 : !!****f* ABINIT/m_lobpcgwf
2 : !! NAME
3 : !! m_lobpcgwf
4 : !!
5 : !! FUNCTION
6 : !! This routine updates the whole wave functions at a given k-point,
7 : !! using the lobpcg method
8 : !! for a given spin-polarization, from a fixed hamiltonian
9 : !! but might also simply compute eigenvectors and eigenvalues at this k point.
10 : !! it will also update the matrix elements of the hamiltonian.
11 : !!
12 : !! COPYRIGHT
13 : !! Copyright (C) 1998-2026 ABINIT group (JB)
14 : !! this file is distributed under the terms of the
15 : !! gnu general public license, see ~abinit/COPYING
16 : !! or http://www.gnu.org/copyleft/gpl.txt .
17 : !! for the initials of contributors, see ~abinit/doc/developers/contributors.txt .
18 : !!
19 : !! SOURCE
20 :
21 : #if defined HAVE_CONFIG_H
22 : #include "config.h"
23 : #endif
24 :
25 : #include "abi_common.h"
26 :
27 : ! nvtx related macro definition
28 : #include "nvtx_macros.h"
29 :
30 : module m_lobpcgwf
31 :
32 : use defs_basis
33 : use m_abicore
34 : use m_lobpcg
35 : use m_xmpi
36 : use m_errors
37 : use m_time
38 : use m_xomp
39 : use m_fstrings
40 : use m_xg
41 : use m_xgTransposer
42 : use m_lobpcg2
43 : use m_dtset
44 :
45 : use defs_abitypes, only : mpi_type
46 : use m_hamiltonian, only : gs_hamiltonian_type
47 : use m_pawcprj, only : pawcprj_type
48 : use m_nonlop, only : nonlop
49 : use m_prep_kgb, only : prep_getghc,prep_nonlop
50 : use m_getghc, only : multithreaded_getghc
51 :
52 : #if defined(HAVE_GPU)
53 : use m_gpu_toolbox
54 : #endif
55 :
56 : #if defined(HAVE_GPU_MARKERS)
57 : use m_nvtx_data
58 : #endif
59 :
60 : use, intrinsic :: iso_c_binding
61 :
62 : implicit none
63 : private
64 :
65 : integer, parameter :: l_tim_getghc=5
66 : double precision, parameter :: inv_sqrt2 = 1/sqrt2
67 :
68 : ! For use in getghc_gsc1
69 : integer,save :: l_cpopt
70 : logical,save :: l_paw
71 : integer,save :: l_prtvol
72 : integer,save :: l_sij_opt
73 : type(mpi_type),pointer,save :: l_mpi_enreg
74 : type(gs_hamiltonian_type),pointer,save :: l_gs_hamk
75 :
76 : public :: lobpcgwf2
77 :
78 : contains
79 :
80 24063 : subroutine lobpcgwf2(cg,dtset,eig,occ,enl_out,gs_hamk,isppol,ikpt,inonsc,istep,kinpw,mpi_enreg,&
81 24063 : & nband,npw,nspinor,prtvol,resid,nbdbuf)
82 :
83 :
84 :
85 : !Arguments ------------------------------------
86 : integer,intent(in) :: nband,npw,prtvol,nspinor
87 : integer,intent(in) :: isppol,ikpt,inonsc,istep,nbdbuf
88 : type(gs_hamiltonian_type),target,intent(inout) :: gs_hamk
89 : type(dataset_type) ,intent(in ) :: dtset
90 : type(mpi_type) ,target,intent(in) :: mpi_enreg
91 : real(dp) ,target,intent(inout) :: cg(2,nspinor*nband*npw)
92 : real(dp) ,intent(in ) :: kinpw(npw)
93 : real(dp) ,target,intent( out) :: resid(nband)
94 : real(dp) ,intent( out) :: enl_out(nband)
95 : real(dp) ,target,intent( out) :: eig(nband)
96 : real(dp) ,target,intent(in ) :: occ(nband)
97 :
98 : !Local variables-------------------------------
99 :
100 : type(xgBlock_t) :: xgx0
101 : type(xgBlock_t) :: xgeigen
102 : type(xgBlock_t) :: xgresidu
103 : type(xgBlock_t) :: xgocc
104 : type(xgBlock_t) :: xg_precond
105 24063 : type(lobpcg_t) :: lobpcg
106 :
107 : integer :: space, blockdim
108 :
109 : logical :: transfer_cg
110 :
111 : integer, parameter :: tim_lobpcgwf2 = 1640
112 : integer, parameter :: tim_enl = 1657
113 : double precision :: tsec(2)
114 :
115 : ! Important things for NC
116 : integer,parameter :: choice=1, paw_opt=0, signs=1
117 192504 : type(pawcprj_type) :: cprj_dum(1,1)
118 : integer :: iblock, shift, me_g0, me_g0_fft
119 : real(dp) :: gsc_dummy(0,0)
120 24063 : real(dp), allocatable :: gvnlxc(:,:)
121 24063 : real(dp), allocatable :: pcon(:),occ_tmp(:)
122 :
123 : ! *********************************************************************
124 :
125 24063 : call timab(tim_lobpcgwf2,1,tsec)
126 :
127 : ! Set module variables
128 24063 : l_paw = (gs_hamk%usepaw==1)
129 24063 : l_cpopt=-1;l_sij_opt=0;if (l_paw) l_sij_opt=1
130 24063 : l_prtvol = prtvol
131 24063 : l_mpi_enreg => mpi_enreg
132 24063 : l_gs_hamk => gs_hamk
133 :
134 : !Variables
135 24063 : blockdim=nband/dtset%nblock_lobpcg
136 24063 : if (blockdim/=mpi_enreg%nproc_band*mpi_enreg%bandpp) then ! without this check computation of enl_out can be wrong
137 0 : ABI_ERROR('blockdim is not consistent with nproc_band and bandpp')
138 : end if
139 :
140 : !Depends on istwfk
141 24063 : if ( gs_hamk%istwf_k > 1 ) then ! Real only
142 : ! SPACE_CR mean that we have complex numbers but no re*im terms only re*re
143 : ! and im*im so that a vector of complex is consider as a long vector of real
144 : ! therefore the number of data is (2*npw*nspinor)*nband
145 : ! This space is completely equivalent to SPACE_R but will correctly set and
146 : ! get the array data into the xgBlock
147 5036 : space = SPACE_CR
148 : else ! complex
149 19027 : space = SPACE_C
150 : end if
151 :
152 : !For preconditionning
153 72189 : ABI_MALLOC(pcon,(npw))
154 24063 : call build_pcon(pcon,kinpw,npw)
155 :
156 24063 : transfer_cg = .false.
157 : #ifdef HAVE_OPENMP_OFFLOAD
158 : if(gs_hamk%gpu_option==ABI_GPU_OPENMP) then
159 : !$OMP TARGET ENTER DATA MAP(to:eig,resid,occ,pcon)
160 : transfer_cg = .not. xomp_target_is_present(c_loc(cg))
161 : !$OMP TARGET ENTER DATA MAP(to:cg) IF(transfer_cg)
162 : end if
163 : #endif
164 :
165 : ! Local variables for lobpcg
166 24063 : me_g0 = -1
167 24063 : me_g0_fft = -1
168 24063 : if (space==SPACE_CR) then
169 5036 : me_g0 = 0
170 5036 : me_g0_fft = 0
171 5036 : if (gs_hamk%istwf_k == 2) then
172 2061 : if (mpi_enreg%me_g0 == 1) me_g0 = 1
173 2061 : if (mpi_enreg%me_g0_fft == 1) me_g0_fft = 1
174 : end if
175 : end if
176 : call xgBlock_map(xgx0,cg,space,npw*nspinor,nband,comm=mpi_enreg%comm_bandspinorfft,me_g0=me_g0,&
177 24063 : & gpu_option=dtset%gpu_option)
178 :
179 24063 : call xgBlock_map_1d(xg_precond,pcon,SPACE_R,npw,gpu_option=dtset%gpu_option)
180 :
181 24063 : call xgBlock_map_1d(xgeigen,eig,SPACE_R,nband,gpu_option=dtset%gpu_option)
182 :
183 24063 : call xgBlock_map_1d(xgresidu,resid,SPACE_R,nband,gpu_option=dtset%gpu_option)
184 :
185 : ! Occupancies in LOBPCG are used for convergence criteria only
186 24063 : if (dtset%nbdbuf==-101.and.nspinor==1.and.dtset%nsppol==1) then
187 0 : ABI_MALLOC(occ_tmp,(nband))
188 0 : occ_tmp(:) = half*occ(:)
189 : #ifdef HAVE_OPENMP_OFFLOAD
190 : !$OMP TARGET ENTER DATA MAP(to:occ_tmp) IF(gs_hamk%gpu_option==ABI_GPU_OPENMP)
191 : #endif
192 0 : call xgBlock_map_1d(xgocc,occ_tmp,SPACE_R,nband,gpu_option=dtset%gpu_option)
193 : else
194 24063 : call xgBlock_map_1d(xgocc,occ,SPACE_R,nband,gpu_option=dtset%gpu_option)
195 : end if
196 :
197 : call lobpcg_init(lobpcg,nband,npw*nspinor,blockdim,dtset%tolwfr_diago,dtset%nline,&
198 : space,mpi_enreg%comm_bandspinorfft,dtset%paral_kgb,mpi_enreg%comm_spinorfft,mpi_enreg%comm_band,&
199 24063 : me_g0,me_g0_fft,gs_hamk%gpu_option,gpu_thread_limit=dtset%gpu_thread_limit)
200 :
201 : ! Run lobpcg
202 24063 : call lobpcg_run(lobpcg,xgx0,getghc_gsc1,xg_precond,xgeigen,xgocc,xgresidu,prtvol,nspinor,isppol,ikpt,inonsc,istep,nbdbuf)
203 :
204 24063 : if (allocated(occ_tmp)) then
205 : #ifdef HAVE_OPENMP_OFFLOAD
206 : !$OMP TARGET EXIT DATA MAP(delete:occ_tmp) IF(gs_hamk%gpu_option==ABI_GPU_OPENMP)
207 : #endif
208 0 : ABI_FREE(occ_tmp)
209 : end if
210 : ! Free preconditionning since not needed anymore
211 : #ifdef HAVE_OPENMP_OFFLOAD
212 : !$OMP TARGET EXIT DATA MAP(delete:pcon) IF(gs_hamk%gpu_option==ABI_GPU_OPENMP)
213 : #endif
214 24063 : ABI_FREE(pcon)
215 :
216 24063 : if ( .not. l_paw ) then
217 8820 : call timab(tim_enl,1,tsec)
218 : #ifdef FC_CRAY
219 : ABI_MALLOC(gvnlxc,(1,1))
220 : #else
221 8820 : ABI_MALLOC(gvnlxc,(0,0))
222 : #endif
223 :
224 : !Call nonlop
225 8820 : if (dtset%paral_kgb==0) then
226 :
227 : call nonlop(choice,l_cpopt,cprj_dum,enl_out,l_gs_hamk,0,eig,mpi_enreg,nband,1,paw_opt,&
228 552 : & signs,gsc_dummy,l_tim_getghc,cg,gvnlxc)
229 :
230 : else
231 : #ifdef HAVE_OPENMP_OFFLOAD
232 : !$OMP TARGET UPDATE FROM(cg) IF(gs_hamk%gpu_option==ABI_GPU_OPENMP)
233 : #endif
234 53964 : do iblock=1,nband/blockdim
235 45696 : shift = (iblock-1)*blockdim*npw*nspinor
236 : call prep_nonlop(choice,l_cpopt,cprj_dum, &
237 : & enl_out((iblock-1)*blockdim+1:iblock*blockdim),gs_hamk,0,&
238 : & eig((iblock-1)*blockdim+1:iblock*blockdim),blockdim,mpi_enreg,1,paw_opt,signs,&
239 : & gsc_dummy,l_tim_getghc,cg(:,shift+1:shift+blockdim*npw*nspinor),gvnlxc(:,:),&
240 53964 : & already_transposed=.false.)
241 : end do
242 : end if
243 8820 : ABI_FREE(gvnlxc)
244 8820 : call timab(tim_enl,2,tsec)
245 : end if
246 :
247 : ! Free lobpcg
248 24063 : call lobpcg_free(lobpcg)
249 :
250 : #ifdef HAVE_OPENMP_OFFLOAD
251 : if(gs_hamk%gpu_option==ABI_GPU_OPENMP) then
252 : !$OMP TARGET EXIT DATA MAP(from:eig,resid,occ,pcon)
253 : !$OMP TARGET EXIT DATA MAP(from:cg) IF(transfer_cg)
254 : end if
255 : #endif
256 :
257 24063 : call timab(tim_lobpcgwf2,2,tsec)
258 :
259 : DBG_EXIT("COLL")
260 :
261 48126 : end subroutine lobpcgwf2
262 :
263 696724 : subroutine getghc_gsc1(X,AX,BX)
264 :
265 : !Arguments ------------------------------------
266 : type(xgBlock_t), intent(inout) :: X
267 : type(xgBlock_t), intent(inout) :: AX
268 : type(xgBlock_t), intent(inout) :: BX
269 :
270 : !Local variables-------------------------------
271 : !scalars
272 : integer :: blockdim
273 : integer :: spacedim
274 : real(dp) :: eval,dum
275 10331371 : type(pawcprj_type) :: cprj_dum(l_gs_hamk%natom,1)
276 : !arrays
277 696724 : real(dp), pointer :: cg(:,:)
278 696724 : real(dp), pointer :: ghc(:,:)
279 696724 : real(dp), pointer :: gsc(:,:)
280 696724 : real(dp), allocatable :: gvnlxc(:,:)
281 :
282 : ! *********************************************************************
283 :
284 : ABI_NVTX_START_RANGE(NVTX_GETGHC)
285 :
286 696724 : call xgBlock_getSize(X,spacedim,blockdim)
287 696724 : call xgBlock_check(X,AX)
288 696724 : call xgBlock_check(X,BX)
289 :
290 696724 : call xgBlock_reverseMap(X,cg,rows=1,cols=spacedim*blockdim)
291 696724 : call xgBlock_reverseMap(AX,ghc,rows=1,cols=spacedim*blockdim)
292 696724 : call xgBlock_reverseMap(BX,gsc,rows=1,cols=spacedim*blockdim)
293 :
294 : #ifdef FC_CRAY
295 : ABI_MALLOC(gvnlxc,(1,1))
296 : #else
297 696724 : ABI_MALLOC(gvnlxc,(0,0))
298 : #endif
299 :
300 696724 : if (l_mpi_enreg%nproc_fft==1.or.l_gs_hamk%istwf_k==1) then
301 : call multithreaded_getghc(l_cpopt,cg,cprj_dum,ghc,gsc,&
302 695068 : l_gs_hamk,gvnlxc,eval,l_mpi_enreg,blockdim,l_prtvol,l_sij_opt,l_tim_getghc,0)
303 1656 : else if (l_gs_hamk%istwf_k==2) then ! nproc_fft>1 and istwfk==2
304 : call prep_getghc(cg(:,1:blockdim*spacedim),l_gs_hamk,gvnlxc,ghc,gsc(:,1:blockdim*spacedim),dum,blockdim,&
305 1656 : l_mpi_enreg,l_prtvol,l_sij_opt,l_cpopt,cprj_dum,already_transposed=.true.)
306 : else ! nproc_fft>1 and istwfk>2
307 0 : ABI_ERROR('getghc in lobpcg not implemented for npfft>1 and istwfk>2')
308 : end if
309 :
310 696724 : ABI_FREE(gvnlxc)
311 :
312 : #if defined(HAVE_GPU_CUDA) && defined(HAVE_YAKL)
313 : call gpu_device_synchronize()
314 : #endif
315 :
316 696724 : if ( .not. l_paw ) call xgBlock_copy(X,BX)
317 :
318 : ABI_NVTX_END_RANGE()
319 :
320 2090172 : end subroutine getghc_gsc1
321 :
322 24063 : subroutine build_pcon(pcon,kinpw,npw)
323 :
324 : integer,intent(in) :: npw
325 : real(dp),intent(in) :: kinpw(:)
326 : real(dp),intent(out) :: pcon(:)
327 :
328 : integer :: ipw
329 :
330 : !$omp parallel do schedule(static), shared(pcon,kinpw)
331 4526833 : do ipw=1,npw
332 4526833 : if(kinpw(ipw)>huge(0.0_dp)*1.d-11) then
333 194498 : pcon(ipw)=0.d0
334 : else
335 : pcon(ipw) = (27+kinpw(ipw)*(18+kinpw(ipw)*(12+8*kinpw(ipw)))) &
336 4308272 : & / (27+kinpw(ipw)*(18+kinpw(ipw)*(12+8*kinpw(ipw))) + 16*kinpw(ipw)**4)
337 : end if
338 : end do
339 :
340 24063 : end subroutine build_pcon
341 :
342 : end module m_lobpcgwf
343 : !!***
|