Line data Source code
1 : !!****m* ABINIT/m_invovl
2 : !! NAME
3 : !! m_invovl
4 : !!
5 : !! FUNCTION
6 : !! Provides functions to invert the overlap matrix S. Used by Chebyshev in PAW
7 : !! See paper by A. Levitt and M. Torrent for details
8 : !! S = 1 + projs * s_projs * projs'
9 : !! S^-1 = 1 + projs * inv_s_projs * projs', with
10 : !! inv_s_projs = - (s_projs^-1 + projs'*projs)^-1
11 : !!
12 : !! COPYRIGHT
13 : !! Copyright (C) 2013-2026 ABINIT group (AL)
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 : !!
18 : !! SOURCE
19 :
20 : #if defined HAVE_CONFIG_H
21 : #include "config.h"
22 : #endif
23 :
24 : #include "abi_common.h"
25 :
26 : ! nvtx related macro definition
27 : #include "nvtx_macros.h"
28 :
29 : MODULE m_invovl
30 :
31 : use defs_basis
32 : use m_errors
33 : use m_xmpi
34 : use m_xomp
35 : use m_abicore
36 : use m_abi_linalg
37 :
38 : use defs_abitypes, only : mpi_type
39 : use m_time, only : timab
40 : use m_hamiltonian, only : gs_hamiltonian_type
41 : use m_bandfft_kpt, only : bandfft_kpt_get_ikpt
42 : use m_pawcprj, only : pawcprj_type, pawcprj_alloc, pawcprj_free, pawcprj_axpby
43 : use m_gemm_nonlop_projectors, only : gemm_nonlop_use_gemm
44 : use m_nonlop, only : nonlop
45 : use m_prep_kgb, only : prep_nonlop
46 :
47 : #ifdef HAVE_FC_ISO_C_BINDING
48 : ! FIXME Don't know what's wrong with GCC when OpenMP GPU Offload is enabled here...
49 : #ifdef FC_GNU
50 : use, intrinsic :: iso_c_binding, only : c_int32_t, c_int64_t, c_float, c_double, c_size_t, c_loc
51 : #else
52 : use, intrinsic :: iso_c_binding, only : c_ptr, c_int32_t, c_int64_t, c_float, c_double, c_size_t, c_loc
53 : #endif
54 : #endif
55 :
56 : #if defined(HAVE_GPU_MARKERS)
57 : use m_nvtx_data
58 : #endif
59 :
60 : #ifdef HAVE_GPU
61 : use m_gpu_toolbox
62 : #endif
63 :
64 : #ifdef HAVE_KOKKOS
65 : use m_manage_kokkos, only : add_array_kokkos
66 : #endif
67 :
68 : implicit none
69 :
70 : private
71 :
72 : !public procedures.
73 : public :: init_invovl
74 : public :: make_invovl
75 : public :: apply_invovl
76 : public :: destroy_invovl
77 :
78 : ! Those routines are here to assess memory requirements
79 : public :: invovl_ompgpu_work_mem
80 : public :: invovl_ompgpu_static_mem
81 : !!***
82 :
83 : !!****t* m_invovl/invovl_kpt_type
84 : !! NAME
85 : !! invovl_kpt_type
86 : !!
87 : !! FUNCTION
88 : !! Contains information needed to invert the overlap matrix S
89 : !!
90 : !! SOURCE
91 :
92 : #if defined(HAVE_FC_ISO_C_BINDING) && defined(HAVE_GPU_CUDA)
93 :
94 : type, public :: invovl_kpt_type
95 :
96 : integer(kind=c_int32_t) :: nprojs
97 : ! total number of projectors
98 : ! nlmn for a specific atom = count(indlmn(3,:,itypat)>0)
99 : ! A value of -1 means that the following arrays are not allocated
100 :
101 : real(kind=c_double), allocatable :: gram_projs(:,:,:)
102 : ! gram_projs(cplx, nprojs, nprojs)
103 : ! projs' * projs
104 :
105 : real(kind=c_double), allocatable :: inv_sij(:,:,:,:)
106 : ! inv_sij(cplx, lmnmax, lmnmax, ntypat)
107 : ! inverse of ham%sij
108 :
109 : real(kind=c_double), allocatable :: inv_s_approx(:,:,:,:)
110 : ! inv_s_approx(cplx, lmnmax, lmnmax, ntypat)
111 : ! preconditionner
112 :
113 : end type invovl_kpt_type
114 :
115 : !> companion type to invovl_kpt_type to pass data to gpu/cuda
116 : type, bind(c), public :: invovl_kpt_gpu_type
117 :
118 : integer(kind=c_int32_t) :: nprojs
119 : ! total number of projectors
120 : ! nlmn for a specific atom = count(indlmn(3,:,itypat)>0)
121 : ! A value of -1 means that the following arrays are not allocated
122 :
123 : type(c_ptr) :: gram_projs
124 : ! gram_projs(cplx, nprojs, nprojs)
125 : ! projs' * projs
126 :
127 : integer(kind=c_int32_t) :: gram_projs_dim(3)
128 :
129 : type(c_ptr) :: inv_sij
130 : ! inv_sij(cplx, lmnmax, lmnmax, ntypat)
131 : ! inverse of ham%sij
132 :
133 : integer(kind=c_int32_t) :: inv_sij_dim(4)
134 :
135 : type(c_ptr) :: inv_s_approx
136 : ! inv_s_approx(cplx, lmnmax, lmnmax, ntypat)
137 : ! preconditionner
138 :
139 : integer(kind=c_int32_t) :: inv_s_approx_dim(4)
140 :
141 : end type invovl_kpt_gpu_type
142 :
143 : #else
144 :
145 : type, public :: invovl_kpt_type
146 :
147 : integer :: nprojs
148 : ! total number of projectors
149 : ! nlmn for a specific atom = count(indlmn(3,:,itypat)>0)
150 : ! A value of -1 means that the following arrays are not allocated
151 :
152 : real(dp), allocatable :: gram_projs(:,:,:)
153 : ! gram_projs(cplx, nprojs, nprojs)
154 : ! projs' * projs
155 :
156 : real(dp), allocatable :: inv_sij(:,:,:,:)
157 : ! inv_sij(cplx, lmnmax, lmnmax, ntypat)
158 : ! inverse of ham%sij
159 :
160 : real(dp), allocatable :: inv_s_approx(:,:,:,:)
161 : ! inv_s_approx(cplx, lmnmax, lmnmax, ntypat)
162 : ! preconditionner
163 :
164 : end type invovl_kpt_type
165 :
166 : #endif
167 :
168 : !!***
169 :
170 : integer, parameter :: tim_nonlop = 13
171 :
172 : integer, parameter :: &
173 : & timer_apply_inv_ovl_opernla = 1630, &
174 : & timer_apply_inv_ovl_opernlb = 1631, &
175 : & timer_apply_inv_ovl_inv_s = 1632
176 :
177 : type(invovl_kpt_type), public,save,allocatable, target :: invovl_kpt(:)
178 : #ifdef HAVE_OPENMP_OFFLOAD
179 : real(dp), ABI_CONTIGUOUS pointer :: current_gram_projs(:,:,:)
180 : real(dp), ABI_CONTIGUOUS pointer :: current_inv_sij(:,:,:,:)
181 : real(dp), ABI_CONTIGUOUS pointer :: current_inv_s_approx(:,:,:,:)
182 : real(dp),allocatable, target :: proj_ompgpu(:,:,:)
183 : real(dp),allocatable, target :: sm1proj_ompgpu(:,:,:)
184 : real(dp),allocatable, target :: PtPsm1proj_ompgpu(:,:,:)
185 : !Module variable keeping track of which K-point data is so=tored on GPU
186 : integer, save :: current_ikpt_in_gpu=-1
187 : integer, save :: gpu_initialized=0
188 : integer, save :: mod__cplx=0, mod__nprojs=0
189 : #endif
190 :
191 : #if defined(HAVE_GPU_CUDA)
192 :
193 : !> this interface is only useful when gpu is enabled
194 : !! these functions are defined in 46_manage_gpu/gpu_apply_invovl_inner.cu
195 : !! these functions are entry point for calling cuda implemented functions
196 : interface
197 :
198 : !> allocate GPU workspace for apply_invovl
199 : subroutine f_gpu_apply_invovl_inner_alloc(proj_dim, ntypat, realloc) bind(c, name='gpu_apply_invovl_inner_alloc')
200 : use, intrinsic :: iso_c_binding
201 : integer(kind=c_int32_t), intent(in) :: proj_dim(3)
202 : integer(kind=c_int32_t), value, intent(in) :: ntypat
203 : integer(kind=c_int32_t), value, intent(in) :: realloc
204 : end subroutine f_gpu_apply_invovl_inner_alloc
205 :
206 : !> deallocate GPU workspace for apply_invovl
207 : subroutine f_gpu_apply_invovl_inner_dealloc() bind(c, name='gpu_apply_invovl_inner_dealloc')
208 : end subroutine f_gpu_apply_invovl_inner_dealloc
209 :
210 : !> allocate GPU workspace for make_invovl (sij and s_approx)
211 : subroutine f_gpu_apply_invovl_matrix_alloc(cplx, nprojs, lmnmax, ntypat, realloc) bind(c, name='gpu_apply_invovl_matrix_alloc')
212 : use, intrinsic :: iso_c_binding
213 : integer(kind=c_int32_t), value, intent(in) :: cplx
214 : integer(kind=c_int32_t), value, intent(in) :: nprojs
215 : integer(kind=c_int32_t), value, intent(in) :: ntypat
216 : integer(kind=c_int32_t), value, intent(in) :: lmnmax
217 : integer(kind=c_int32_t), value, intent(in) :: realloc
218 : end subroutine f_gpu_apply_invovl_matrix_alloc
219 :
220 : !> deallocate GPU workspace for make_invovl (sij and s_approx)
221 : subroutine f_gpu_apply_invovl_matrix_dealloc() bind(c, name='gpu_apply_invovl_matrix_dealloc')
222 : end subroutine f_gpu_apply_invovl_matrix_dealloc
223 :
224 : !> init data for GPU
225 : subroutine f_gpu_init_invovl_data(indlmn_dim, indlmn_ptr) bind(c, name='init_invovl_data')
226 : use, intrinsic :: iso_c_binding
227 : integer(kind=c_int32_t), intent(in) :: indlmn_dim(3)
228 : type(c_ptr) , value :: indlmn_ptr
229 : end subroutine f_gpu_init_invovl_data
230 :
231 : !> upload inverse overlap matrices
232 : subroutine f_upload_inverse_overlap(invovl_gpu, cplx, nprojs, lmnmax, ntypat) bind(c, name='upload_inverse_overlap')
233 : use, intrinsic :: iso_c_binding
234 : import invovl_kpt_gpu_type
235 : type(invovl_kpt_gpu_type), value, intent(in) :: invovl_gpu
236 : integer(kind=c_int32_t), value, intent(in) :: cplx
237 : integer(kind=c_int32_t), value, intent(in) :: nprojs
238 : integer(kind=c_int32_t), value, intent(in) :: lmnmax
239 : integer(kind=c_int32_t), value, intent(in) :: ntypat
240 : end subroutine f_upload_inverse_overlap
241 :
242 : !> solve_inner on GPU
243 : subroutine f_solve_inner_gpu(proj_dim, proj_ptr, sm1proj_ptr, ptp_sm1proj_ptr, &
244 : & nattyp_dim, nattyp_ptr, ntypat, lmnmax, cplx, block_sliced) bind(c, name='solve_inner_gpu')
245 :
246 : use, intrinsic :: iso_c_binding
247 : integer(kind=c_int32_t), intent(in) :: proj_dim(3)
248 : type(c_ptr) , value :: proj_ptr
249 : type(c_ptr) , value :: sm1proj_ptr
250 : type(c_ptr) , value :: ptp_sm1proj_ptr
251 : integer(kind=c_int32_t), value, intent(in) :: nattyp_dim
252 : type(c_ptr) , value :: nattyp_ptr
253 : integer(kind=c_int32_t), value, intent(in) :: ntypat
254 : integer(kind=c_int32_t), value, intent(in) :: lmnmax
255 : integer(kind=c_int32_t), value, intent(in) :: cplx
256 : integer(kind=c_int32_t), value, intent(in) :: block_sliced
257 : end subroutine f_solve_inner_gpu
258 :
259 : end interface
260 :
261 : #endif
262 :
263 : CONTAINS
264 :
265 : #if defined(HAVE_FC_ISO_C_BINDING) && defined(HAVE_GPU_CUDA)
266 : !!****f* m_invovl/make_invovl_kpt_gpu
267 : !! NAME
268 : !! make_invovl_kpt
269 : !!
270 : !! FUNCTION
271 : !! Create a invovl_pkt_gpu_type from a cpu counter part for cuda interoperability
272 : !! SOURCE
273 : function make_invovl_kpt_gpu(invovl) result(invovl_gpu)
274 : type(invovl_kpt_type), intent(inout),target :: invovl
275 : type(invovl_kpt_gpu_type) :: invovl_gpu
276 :
277 : invovl_gpu%nprojs = invovl%nprojs
278 :
279 : invovl_gpu%gram_projs = c_loc(invovl%gram_projs(1,1,1))
280 : invovl_gpu%gram_projs_dim = (/ &
281 : & size(invovl%gram_projs,1), &
282 : & size(invovl%gram_projs,2), &
283 : & size(invovl%gram_projs,3) &
284 : & /)
285 :
286 : invovl_gpu%inv_sij = c_loc(invovl%inv_sij(1,1,1,1))
287 : invovl_gpu%inv_sij_dim = (/ &
288 : & size(invovl%inv_sij,1), &
289 : & size(invovl%inv_sij,2), &
290 : & size(invovl%inv_sij,3), &
291 : & size(invovl%inv_sij,4) &
292 : & /)
293 :
294 : invovl_gpu%inv_s_approx = c_loc(invovl%inv_s_approx(1,1,1,1))
295 : invovl_gpu%inv_s_approx_dim = (/ &
296 : & size(invovl%inv_s_approx,1), &
297 : & size(invovl%inv_s_approx,2), &
298 : & size(invovl%inv_s_approx,3), &
299 : & size(invovl%inv_s_approx,4) &
300 : & /)
301 :
302 : end function make_invovl_kpt_gpu
303 : !!***
304 : #endif
305 :
306 : #ifdef HAVE_OPENMP_OFFLOAD
307 :
308 : subroutine alloc_ompgpu_buffers(cplx,nprojs,nspinor,ndat)
309 : integer,intent(in) :: cplx,nprojs,nspinor,ndat
310 :
311 : if(gpu_initialized == 0 .or. mod__cplx/=cplx .or. mod__nprojs/=nprojs) then
312 :
313 : if(gpu_initialized==1) then
314 : ABI_FREE(proj_ompgpu)
315 : ABI_FREE(sm1proj_ompgpu)
316 : ABI_FREE(PtPsm1proj_ompgpu)
317 : end if
318 :
319 : ABI_MALLOC(proj_ompgpu, (cplx,nprojs,nspinor*ndat))
320 : ABI_MALLOC(sm1proj_ompgpu, (cplx,nprojs,nspinor*ndat))
321 : ABI_MALLOC(PtPsm1proj_ompgpu, (cplx,nprojs,nspinor*ndat))
322 :
323 : mod__cplx=cplx; mod__nprojs=nprojs
324 :
325 : !FIXME Smater buffer management ?
326 : !!$OMP TARGET ENTER DATA MAP(alloc:proj_ompgpu,sm1proj_ompgpu,PtPsm1proj_ompgpu)
327 : gpu_initialized=1
328 : end if
329 :
330 : end subroutine alloc_ompgpu_buffers
331 :
332 : subroutine refresh_invovl_ompgpu_kpt(ikpt)
333 : integer,intent(in) :: ikpt
334 :
335 : ! *************************************************************************
336 :
337 : if(ikpt < 0) then
338 : ABI_BUG("Requested GPU upload of a negative K-point index !")
339 : end if
340 :
341 : if(current_ikpt_in_gpu /= -1) then
342 : !$OMP TARGET EXIT DATA MAP(delete:current_gram_projs)
343 : !$OMP TARGET EXIT DATA MAP(delete:current_inv_sij)
344 : !$OMP TARGET EXIT DATA MAP(delete:current_inv_s_approx)
345 : end if
346 :
347 : current_gram_projs => invovl_kpt(ikpt)%gram_projs
348 : current_inv_sij => invovl_kpt(ikpt)%inv_sij
349 : current_inv_s_approx => invovl_kpt(ikpt)%inv_s_approx
350 :
351 : !$OMP TARGET ENTER DATA MAP(to:current_gram_projs)
352 : !$OMP TARGET ENTER DATA MAP(to:current_inv_sij)
353 : !$OMP TARGET ENTER DATA MAP(to:current_inv_s_approx)
354 : current_ikpt_in_gpu=ikpt
355 :
356 : end subroutine refresh_invovl_ompgpu_kpt
357 : #endif
358 :
359 : !!****f* m_invovl/init_invovl
360 : !! NAME
361 : !! init_invovl
362 : !!
363 : !! FUNCTION
364 : !! Initalization of the invovl_kpt array
365 : !!
366 : !! INPUTS
367 : !! nkpt= number of k-points
368 : !!
369 : !! SOURCE
370 :
371 109 : subroutine init_invovl(nkpt)
372 :
373 : integer, intent(in) :: nkpt
374 : integer :: ikpt
375 :
376 : ! *************************************************************************
377 :
378 1633 : ABI_MALLOC(invovl_kpt, (nkpt))
379 : ! TODO add cycling if kpt parallelism
380 1415 : do ikpt=1,nkpt
381 1415 : invovl_kpt(ikpt)%nprojs = -1
382 : end do
383 :
384 109 : end subroutine init_invovl
385 : !!***
386 :
387 : !!****f* m_invovl/destroy_invovl_ikpt
388 : !! NAME
389 : !! destroy_invovl_ikpt
390 : !!
391 : !! FUNCTION
392 : !! Destruction of the ikpt-th member of invovl array
393 : !!
394 : !! INPUTS
395 : !! ikpt= index of k-point
396 : !!
397 : !! SOURCE
398 16611 : subroutine destroy_invovl_ikpt(ikpt, gpu_option)
399 :
400 : integer, intent(in) :: ikpt
401 : integer, intent(in) :: gpu_option
402 :
403 : ! *************************************************************************
404 :
405 : if(gpu_option==ABI_GPU_OPENMP) then
406 : #ifdef HAVE_OPENMP_OFFLOAD
407 : if(gpu_initialized==1 .and. current_ikpt_in_gpu == ikpt) then
408 : !$OMP TARGET EXIT DATA MAP(delete:current_gram_projs)
409 : !$OMP TARGET EXIT DATA MAP(delete:current_inv_sij)
410 : !$OMP TARGET EXIT DATA MAP(delete:current_inv_s_approx)
411 : nullify(current_gram_projs)
412 : nullify(current_inv_sij)
413 : nullify(current_inv_s_approx)
414 : current_ikpt_in_gpu = -1
415 : !FIXME Smater buffer management ?
416 : !!$OMP TARGET EXIT DATA MAP(delete:proj_ompgpu,sm1proj_ompgpu,PtPsm1proj_ompgpu)
417 : ABI_FREE(proj_ompgpu)
418 : ABI_FREE(sm1proj_ompgpu)
419 : ABI_FREE(PtPsm1proj_ompgpu)
420 : gpu_initialized = 0
421 : end if
422 : #endif
423 : end if
424 :
425 16611 : ABI_FREE(invovl_kpt(ikpt)%gram_projs)
426 16611 : ABI_FREE(invovl_kpt(ikpt)%inv_sij)
427 16611 : ABI_FREE(invovl_kpt(ikpt)%inv_s_approx)
428 16611 : invovl_kpt(ikpt)%nprojs = -1
429 :
430 : #if defined(HAVE_GPU_CUDA) && defined(HAVE_FC_ISO_C_BINDING)
431 : if (gpu_option == ABI_GPU_LEGACY .or. gpu_option == ABI_GPU_KOKKOS) then
432 : call f_gpu_apply_invovl_inner_dealloc()
433 : call f_gpu_apply_invovl_matrix_dealloc()
434 : end if
435 : #endif
436 :
437 16611 : end subroutine destroy_invovl_ikpt
438 : !!***
439 :
440 : !!****f* m_invovl/destroy_invovl
441 : !! NAME
442 : !! destroy_invovl
443 : !!
444 : !! FUNCTION
445 : !! Destruction of the invovl_kpt array
446 : !!
447 : !! INPUTS
448 : !! nkpt= number of k-points
449 : !!
450 : !! SOURCE
451 108 : subroutine destroy_invovl(nkpt, gpu_option)
452 :
453 : integer, intent(in) :: nkpt
454 : integer, intent(in) :: gpu_option
455 : integer :: ikpt
456 :
457 : ! *************************************************************************
458 :
459 : ! TODO add cycling if kpt parallelism
460 1386 : do ikpt=1,nkpt
461 1278 : if(invovl_kpt(ikpt)%nprojs == -1) then
462 : ! write(0, *) 'ERROR invovl_kpt is unallocated'
463 : cycle
464 : end if
465 1386 : call destroy_invovl_ikpt(ikpt, gpu_option)
466 : end do
467 :
468 1386 : ABI_FREE(invovl_kpt)
469 :
470 108 : end subroutine destroy_invovl
471 : !!***
472 :
473 : !!****f* m_invovl/make_invovl
474 : !! NAME
475 : !! make_invovl
476 : !!
477 : !! FUNCTION
478 : !! Builds of the invovl structure
479 : !!
480 : !! INPUTS
481 : !!
482 : !! SOURCE
483 :
484 16639 : subroutine make_invovl(ham, dimffnl, ffnl, ph3d, mpi_enreg)
485 :
486 : use m_abi_linalg
487 :
488 : type(gs_hamiltonian_type),intent(in), target :: ham
489 : integer, intent(in) :: dimffnl
490 : real(dp),intent(in) :: ffnl(ham%npw_k,dimffnl,ham%lmnmax,ham%ntypat)
491 : real(dp),intent(in) :: ph3d(2,ham%npw_k,ham%matblk)
492 : type(mpi_type) :: mpi_enreg
493 :
494 33278 : real(dp) :: atom_projs(2, ham%npw_k, ham%lmnmax)
495 33278 : real(dp) :: temp(ham%npw_k)
496 16639 : complex(dp), allocatable :: work(:)
497 16639 : real(dp), allocatable,target :: projs(:,:,:)
498 16639 : real(dp), allocatable :: gram_proj(:,:,:)
499 16639 : integer, allocatable :: ipiv(:)
500 :
501 : integer :: itypat, ilmn, nlmn, jlmn, ia, iaph3d, shift
502 : integer :: il, ilm, jlm, ipw, info, ierr, cplx
503 : integer :: ikpt_this_proc,cplex_dij
504 : logical :: parity
505 : real(dp) :: tsec(2)
506 : #if defined(HAVE_FC_ISO_C_BINDING) && defined(HAVE_GPU_CUDA)
507 : character(len=500) :: message
508 : #endif
509 : character :: blas_transpose
510 :
511 : type(invovl_kpt_type), pointer :: invovl
512 : #if defined(HAVE_FC_ISO_C_BINDING) && defined(HAVE_GPU_CUDA)
513 : type(invovl_kpt_gpu_type) :: invovl_gpu
514 : #endif
515 33278 : integer :: array_nprojs_pp(mpi_enreg%nproc_fft)
516 : integer :: iproc, slice_size
517 16639 : real(dp), allocatable :: gramwork(:,:,:)
518 : #ifdef HAVE_OPENMP_OFFLOAD
519 : real(dp), ABI_CONTIGUOUS pointer :: invovl_gram_projs(:,:,:)
520 : #endif
521 :
522 : integer, parameter :: timer_mkinvovl = 1620, timer_mkinvovl_build_d = 1621, timer_mkinvovl_build_ptp = 1622
523 :
524 : ! *************************************************************************
525 :
526 : !! S = 1 + PDP', so
527 : !! S^-1 = 1 + P inv_s_projs P', with
528 : !! inv_s_projs = - (D^-1 + P'*P)^-1
529 :
530 : ABI_NVTX_START_RANGE(NVTX_MAKE_INVOVL)
531 :
532 16639 : if(ham%istwf_k == 1) then
533 16511 : cplx = 2
534 16511 : blas_transpose = 'c'
535 : else
536 128 : cplx = 1
537 128 : blas_transpose = 't'
538 : end if
539 :
540 16639 : ikpt_this_proc=bandfft_kpt_get_ikpt()
541 16639 : invovl => invovl_kpt(ikpt_this_proc)
542 :
543 16639 : if(invovl%nprojs /= -1) then
544 : ! We have been here before, cleanup before remaking
545 15736 : call destroy_invovl_ikpt(ikpt_this_proc, ham%gpu_option)
546 : end if
547 :
548 16639 : iaph3d = 1
549 :
550 16639 : call timab(timer_mkinvovl,1,tsec)
551 16639 : call timab(timer_mkinvovl_build_d,1,tsec)
552 :
553 : ! build nprojs
554 16639 : invovl%nprojs = 0
555 33736 : do itypat=1,ham%ntypat
556 170882 : invovl%nprojs = invovl%nprojs + count(ham%indlmn(3,:,itypat)>0)*ham%nattyp(itypat)
557 : end do
558 :
559 66556 : ABI_MALLOC(projs, (2, ham%npw_k, invovl%nprojs))
560 99834 : ABI_MALLOC(invovl%inv_sij, (cplx, ham%lmnmax, ham%lmnmax, ham%ntypat))
561 99834 : ABI_MALLOC(invovl%inv_s_approx, (cplx, ham%lmnmax, ham%lmnmax, ham%ntypat))
562 : ! workspace for inversion
563 49917 : ABI_MALLOC(ipiv, (ham%lmnmax))
564 49917 : ABI_MALLOC(work, (64*ham%lmnmax))
565 :
566 132942015 : projs = zero
567 3465982 : invovl%inv_sij = zero
568 3465982 : invovl%inv_s_approx = zero
569 :
570 : shift = 0
571 33736 : do itypat = 1, ham%ntypat
572 154243 : nlmn = count(ham%indlmn(3,:,itypat)>0)
573 17097 : if (size(ham%sij(:,itypat))==ham%lmnmax*(ham%lmnmax+1)/2) then
574 : cplex_dij = 1
575 176 : else if (size(ham%sij(:,itypat))==ham%lmnmax*(ham%lmnmax+1)) then
576 : cplex_dij = 2
577 : else
578 0 : ABI_ERROR('sij size not recognize')
579 : end if
580 : !! unpack ham%sij into inv_sij
581 154223 : do jlmn = 1, nlmn
582 137126 : if (cplex_dij==1) then
583 135718 : invovl%inv_sij(1, jlmn, jlmn, itypat) = ham%sij(jlmn*(jlmn-1)/2 + jlmn, itypat)
584 : else
585 1408 : invovl%inv_sij(1, jlmn, jlmn, itypat) = ham%sij(jlmn*(jlmn-1) + 2*jlmn-1, itypat)
586 : end if
587 637314 : do ilmn = 1, jlmn-1
588 620217 : if (cplex_dij==1) then
589 478163 : invovl%inv_sij(1, ilmn, jlmn, itypat) = ham%sij(jlmn*(jlmn-1)/2 + ilmn, itypat)
590 478163 : invovl%inv_sij(1, jlmn, ilmn, itypat) = ham%sij(jlmn*(jlmn-1)/2 + ilmn, itypat)
591 : else
592 4928 : invovl%inv_sij(1, ilmn, jlmn, itypat) = ham%sij(jlmn*(jlmn-1) + 2*ilmn-1, itypat)
593 4928 : invovl%inv_sij(1, jlmn, ilmn, itypat) = ham%sij(jlmn*(jlmn-1) + 2*ilmn-1, itypat)
594 :
595 : end if
596 : end do
597 : end do
598 :
599 : ! Invert sij
600 17097 : if(cplx == 2) then
601 16841 : call ZHETRF('U', nlmn, invovl%inv_sij(:,:,:,itypat), ham%lmnmax, ipiv, work, (64*nlmn), info)
602 16841 : call ZHETRI('U', nlmn, invovl%inv_sij(:,:,:,itypat), ham%lmnmax, ipiv, work, info)
603 : else
604 256 : call DSYTRF('U', nlmn, invovl%inv_sij(:,:,:,itypat), ham%lmnmax, ipiv, work, (64*nlmn), info)
605 256 : call DSYTRI('U', nlmn, invovl%inv_sij(:,:,:,itypat), ham%lmnmax, ipiv, work, info)
606 : end if
607 : ! complete the matrix
608 154223 : do ilm=1, nlmn
609 637314 : do jlm=1, ilm-1
610 620217 : invovl%inv_sij(1,ilm,jlm,itypat) = invovl%inv_sij(1,jlm,ilm,itypat)
611 : end do
612 : end do
613 :
614 : !! loop on atoms to build atom_projs and fill projs, s_projs
615 67339 : do ia = 1, ham%nattyp(itypat)
616 :
617 : !! build atom_projs, from opernlb
618 : !! P = 4pi/sqrt(ucvol)* conj(diag(ph3d)) * ffnl * diag(parity), with parity = (-i)^l
619 133024979 : atom_projs(:,:,:) = zero
620 :
621 : ! start from 4pi/sqrt(ucvol)*ffnl
622 : ! atom_projs(1, :, 1:nlmn) = four_pi/sqrt(ham%ucvol) * ffnl(:, 1, 1:nlmn)
623 : ! TODO vectorize (DCOPY with stride)
624 5074316 : do ipw=1, ham%npw_k
625 49291890 : atom_projs(1,ipw, 1:nlmn) = four_pi/sqrt(ham%ucvol) * ffnl(ipw, 1, 1:nlmn, itypat)
626 : end do
627 :
628 : ! multiply by (-i)^l
629 306257 : do ilmn=1,nlmn
630 272654 : il=mod(ham%indlmn(1,ilmn, itypat),4);
631 272654 : parity=(mod(il,2)==0)
632 272654 : if (il>1) then
633 : ! multiply by -1
634 11679440 : atom_projs(:,:,ilmn) = -atom_projs(:,:,ilmn)
635 : end if
636 306257 : if(.not. parity) then
637 : ! multiply by -i
638 30445896 : temp = atom_projs(2,:,ilmn)
639 30445896 : atom_projs(2,:,ilmn) = -atom_projs(1,:,ilmn)
640 30445896 : atom_projs(1,:,ilmn) = temp
641 : end if
642 : end do
643 :
644 : ! multiply by conj(ph3d)
645 306257 : do ilmn=1,nlmn
646 44490228 : temp = atom_projs(1, :, ilmn)
647 44490228 : atom_projs(1, :, ilmn) = atom_projs(1, :, ilmn) * ph3d(1, :, iaph3d) + atom_projs(2, :, ilmn) * ph3d(2, :, iaph3d)
648 44523831 : atom_projs(2, :, ilmn) = atom_projs(2, :, ilmn) * ph3d(1, :, iaph3d) - temp * ph3d(2, :, iaph3d)
649 : end do
650 :
651 : ! me_g0 trick
652 33603 : if(ham%istwf_k == 2 .and. mpi_enreg%me_g0 == 1) then
653 126 : atom_projs(1,1,:) = atom_projs(1,1,:) / sqrt2
654 126 : atom_projs(2,1,:) = zero
655 : end if
656 33603 : if(ham%istwf_k > 1) then
657 198240 : atom_projs(:,:,:) = atom_projs(:,:,:) * sqrt2
658 : end if
659 :
660 :
661 : !! atom_projs and typat_s_projs are built, copy them to projs and inv_s_projs
662 132958979 : projs(:, :, shift+1:shift+nlmn) = atom_projs(:, :, 1:nlmn)
663 33603 : shift = shift + nlmn
664 :
665 : ! build inv_s_approx = (D^-1+PtP)^-1 restricted to a single atom block
666 : ! can be optimized (real, build directly from ffnl)
667 33603 : if(ia == 1) then
668 : ! D^-1
669 1257487 : invovl%inv_s_approx(1, :, :, itypat) = invovl%inv_sij(1, :, :, itypat)
670 : ! + PtP
671 85445 : ABI_MALLOC(gram_proj, (cplx, nlmn, nlmn))
672 : call abi_xgemm(blas_transpose,'N', nlmn, nlmn, (3-cplx)*ham%npw_k, cone, atom_projs(:,:,1), (3-cplx)*ham%npw_k, &
673 17089 : & atom_projs(:,:,1), (3-cplx)*ham%npw_k, czero, gram_proj(:,:,1), nlmn,x_cplx=cplx)
674 17089 : call xmpi_sum(gram_proj,mpi_enreg%comm_bandspinorfft,ierr)
675 3446155 : invovl%inv_s_approx(:,1:nlmn,1:nlmn,itypat) = invovl%inv_s_approx(:,1:nlmn,1:nlmn,itypat) + gram_proj(:,:,:)
676 17089 : ABI_FREE(gram_proj)
677 : ! ^-1
678 17089 : if(cplx == 2) then
679 16833 : call ZHETRF('U', nlmn, invovl%inv_s_approx(:,:,:,itypat), ham%lmnmax, ipiv, work, (64*nlmn), info)
680 16833 : call ZHETRI('U', nlmn, invovl%inv_s_approx(:,:,:,itypat), ham%lmnmax, ipiv, work, info)
681 : else
682 256 : call DSYTRF('U', nlmn, invovl%inv_s_approx(:,:,:,itypat), ham%lmnmax, ipiv, work, (64*nlmn), info)
683 256 : call DSYTRI('U', nlmn, invovl%inv_s_approx(:,:,:,itypat), ham%lmnmax, ipiv, work, info)
684 : end if
685 : ! complete lower triangle of matrix
686 171240 : do ilm=1, nlmn
687 637018 : do jlm=1, ilm-1
688 482867 : invovl%inv_s_approx(1, ilm, jlm, itypat) = invovl%inv_s_approx(1, jlm, ilm, itypat)
689 619929 : if(cplx == 2) then
690 475699 : invovl%inv_s_approx(2, ilm, jlm, itypat) = -invovl%inv_s_approx(2, jlm, ilm, itypat)
691 : end if
692 : end do
693 : end do
694 : end if
695 :
696 50700 : iaph3d = iaph3d + 1
697 : end do
698 : end do
699 16639 : ABI_FREE(ipiv)
700 16639 : ABI_FREE(work)
701 :
702 16639 : call timab(timer_mkinvovl_build_d, 2, tsec)
703 16639 : call timab(timer_mkinvovl_build_ptp, 1, tsec)
704 :
705 : ! Compute P'P one column slice at a time (might be too big to fit in one proc)
706 16639 : if(mpi_enreg%paral_kgb == 1) then
707 : ! Split the work evenly the fft processors
708 13464 : array_nprojs_pp(:) = invovl%nprojs / mpi_enreg%nproc_fft
709 : ! not enough work, there's MOD(nprojs,mpi_enreg%nproc_fft) tasks left
710 : ! assign them to the first ones
711 6728 : array_nprojs_pp(1:MOD(invovl%nprojs,mpi_enreg%nproc_fft)) = array_nprojs_pp(1:MOD(invovl%nprojs,mpi_enreg%nproc_fft)) + 1
712 : else
713 19822 : array_nprojs_pp = invovl%nprojs
714 : end if
715 83195 : ABI_MALLOC(invovl%gram_projs, (cplx,invovl%nprojs,array_nprojs_pp(mpi_enreg%me_fft+1)))
716 16639 : shift = 0
717 16639 : if (ham%gpu_option==ABI_GPU_OPENMP .and. mpi_enreg%nproc_fft==1) then
718 : #ifdef HAVE_OPENMP_OFFLOAD
719 : ! compute gram_projs in one GEMM, only one FFT proc expected in GPU mode
720 : slice_size = array_nprojs_pp(1)
721 : invovl_gram_projs => invovl%gram_projs
722 : !$OMP TARGET ENTER DATA MAP(alloc:invovl_gram_projs)
723 : !$OMP TARGET ENTER DATA MAP(to:projs)
724 :
725 : call abi_xgemm(blas_transpose,'N', invovl%nprojs, slice_size, (3-cplx)*ham%npw_k, cone, &
726 : & projs, (3-cplx)*ham%npw_k, &
727 : & projs, (3-cplx)*ham%npw_k, czero, &
728 : & invovl_gram_projs, invovl%nprojs, &
729 : & x_cplx=cplx, gpu_option=ham%gpu_option)
730 : call xmpi_sum(invovl%gram_projs,mpi_enreg%comm_band,ierr,use_omp_map=.true.)
731 : !$OMP TARGET EXIT DATA MAP(from:invovl_gram_projs)
732 : !$OMP TARGET EXIT DATA MAP(delete:projs)
733 : #endif
734 : else
735 33286 : do iproc = 1, mpi_enreg%nproc_fft
736 : ! compute local contribution to slice iproc of gram_projs
737 16647 : slice_size = array_nprojs_pp(iproc)
738 83235 : ABI_MALLOC(gramwork, (cplx,invovl%nprojs,slice_size))
739 : call abi_xgemm(blas_transpose,'N', invovl%nprojs, slice_size, (3-cplx)*ham%npw_k, cone,&
740 : & projs(:,:,1), (3-cplx)*ham%npw_k, &
741 : & projs(:, :, shift+1), (3-cplx)*ham%npw_k, czero, &
742 : & gramwork(:,:,1), invovl%nprojs,&
743 16647 : & x_cplx=cplx)
744 16647 : shift = shift + slice_size
745 : ! reduce on proc i
746 16647 : call xmpi_sum_master(gramwork, iproc-1, mpi_enreg%comm_fft, ierr)
747 16647 : if(iproc == mpi_enreg%me_fft+1) then
748 21888376 : invovl%gram_projs = gramwork
749 : end if
750 33286 : ABI_FREE(gramwork)
751 : end do
752 16639 : call xmpi_sum(invovl%gram_projs,mpi_enreg%comm_band,ierr)
753 : end if
754 :
755 16639 : call timab(timer_mkinvovl_build_ptp, 2, tsec)
756 16639 : call timab(timer_mkinvovl,2,tsec)
757 :
758 16639 : ABI_FREE(projs)
759 :
760 : #if defined(HAVE_FC_ISO_C_BINDING) && defined(HAVE_GPU_CUDA)
761 :
762 : ! upload inverse overlap matrices (sij and s_approx) to GPU memory
763 : if (ham%gpu_option==ABI_GPU_LEGACY .or. ham%gpu_option==ABI_GPU_KOKKOS) then
764 : ! allocate memory for sij and s_approx on GPU
765 : write(message,'(a,a,i12,a,a,i6,a,a,i6,a,a,es12.4,a)') &
766 : & 'Allocate GPU memory for inverse overlap computations (sij and s_approx) : ',&
767 : & 'nprojs=',invovl%nprojs,ch10,&
768 : & 'nlmnax=',ham%lmnmax,ch10,&
769 : & 'ntypat=',ham%ntypat,ch10,&
770 : & 'gram_projs_gpu_size (GBytes)=',1e-9*cplx*invovl%nprojs*invovl%nprojs*dp,ch10
771 : call wrtout(std_out,message,'COLL')
772 : call f_gpu_apply_invovl_matrix_alloc(cplx, invovl%nprojs, ham%lmnmax, ham%ntypat, 0)
773 :
774 : invovl_gpu = make_invovl_kpt_gpu(invovl)
775 : call f_upload_inverse_overlap(invovl_gpu, cplx, invovl%nprojs, ham%lmnmax, ham%ntypat)
776 : write(message,*) 'Invovl uploaded to GPU memory'
777 : call wrtout(std_out,message,'COLL')
778 : end if
779 :
780 : #endif
781 :
782 : #ifdef HAVE_OPENMP_OFFLOAD
783 : if (ham%gpu_option==ABI_GPU_OPENMP) then
784 : call refresh_invovl_ompgpu_kpt(ikpt_this_proc)
785 : end if
786 : #endif
787 :
788 : ! LB-10/06/24: This message is too verbose on some cases (for example many k-points)
789 : ! write(message,*) 'Invovl built'
790 : ! call wrtout(std_out,message,'COLL')
791 :
792 : ABI_NVTX_END_RANGE()
793 :
794 16639 : end subroutine make_invovl
795 : !!***
796 :
797 : !!****f* m_invovl/apply_invovl
798 : !! NAME
799 : !! apply_invovl
800 : !!
801 : !! FUNCTION
802 : !! Applies the inverse of the overlap matrix to cwavef
803 : !!
804 : !! INPUTS
805 : !!
806 : !! SOURCE
807 :
808 160806 : subroutine apply_invovl(ham, cwavef, sm1cwavef, cwaveprj, npw, ndat, mpi_enreg, nspinor, block_sliced)
809 :
810 : #if defined(HAVE_FC_ISO_C_BINDING) && defined(HAVE_GPU_CUDA)
811 : use, intrinsic :: iso_c_binding
812 : #endif
813 :
814 : ! args
815 : type(gs_hamiltonian_type), intent(in), target :: ham
816 : integer, intent(in) :: npw, ndat
817 : integer, intent(in) :: nspinor
818 : integer, intent(in) :: block_sliced
819 : real(dp), intent(inout), target :: cwavef(2, npw*nspinor*ndat) ! TODO should be in, fix nonlop
820 : type(mpi_type) :: mpi_enreg
821 : real(dp), intent(inout), target :: sm1cwavef(2, npw*nspinor*ndat)
822 : type(pawcprj_type), intent(inout) :: cwaveprj(:,:)
823 :
824 160806 : real(dp),allocatable, target :: proj(:,:,:), sm1proj(:,:,:), PtPsm1proj(:,:,:)
825 :
826 : ! used to pass proj dimensions to cuda
827 : integer(kind=c_int32_t) :: proj_dim(3)
828 : integer(kind=c_int32_t) :: nattyp_dim
829 : integer(kind=c_int32_t) :: indlmn_dim(3)
830 : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
831 : integer(kind=c_int32_t) :: cwavef_size
832 : #endif
833 :
834 : integer :: idat, iatom, nlmn, shift
835 : real(dp) :: tsec(2)
836 :
837 : integer :: choice, cpopt, paw_opt , cplx
838 : character :: blas_transpose
839 160806 : type(pawcprj_type),allocatable :: cwaveprj_in(:,:)
840 :
841 : integer :: ikpt_this_proc
842 : ! dummies
843 321612 : real(dp) :: enlout(ndat), lambda_block(1), gvnlxc(1,1)
844 : integer, parameter :: nnlout = 0, idir = 0, signs = 2
845 :
846 : type(invovl_kpt_type), pointer :: invovl
847 :
848 : ! *************************************************************************
849 :
850 : if(ham%gpu_option==ABI_GPU_OPENMP) then
851 : #ifdef HAVE_OPENMP_OFFLOAD
852 : call apply_invovl_ompgpu(ham, cwavef, sm1cwavef, cwaveprj, npw, ndat, mpi_enreg, nspinor, block_sliced)
853 : return
854 : #endif
855 : end if
856 :
857 : ABI_NVTX_START_RANGE(NVTX_INVOVL_PREP)
858 160806 : ikpt_this_proc=bandfft_kpt_get_ikpt()
859 160806 : invovl => invovl_kpt(ikpt_this_proc)
860 :
861 160806 : if(ham%istwf_k == 1) then
862 143995 : cplx = 2
863 143995 : blas_transpose = 'c'
864 : else
865 16811 : cplx = 1
866 16811 : blas_transpose = 't'
867 : end if
868 :
869 804030 : ABI_MALLOC(proj, (cplx,invovl%nprojs,nspinor*ndat))
870 643224 : ABI_MALLOC(sm1proj, (cplx,invovl%nprojs,nspinor*ndat))
871 643224 : ABI_MALLOC(PtPsm1proj, (cplx,invovl%nprojs,nspinor*ndat))
872 65308502 : proj = zero
873 65308502 : sm1proj = zero
874 65308502 : PtPsm1proj = zero
875 :
876 160806 : proj_dim = (/ size(proj,1), size(proj,2), size(proj,3) /)
877 :
878 160806 : nattyp_dim = size(ham%nattyp)
879 :
880 160806 : indlmn_dim = (/ size(ham%indlmn,1), size(ham%indlmn,2), size(ham%indlmn,3) /)
881 :
882 : #if defined(HAVE_GPU_CUDA) && defined(HAVE_FC_ISO_C_BINDING)
883 :
884 : !! memory allocation of data used in solve_inner_gpu
885 : !! note : this is actually done only once
886 : if (ham%gpu_option==ABI_GPU_LEGACY .or. ham%gpu_option==ABI_GPU_KOKKOS) then
887 :
888 : #ifdef DEBUG_VERBOSE_GPU
889 : if(xmpi_comm_rank(xmpi_world) == 0) then
890 : call check_gpu_mem("gpu_apply_invovl_inner_alloc begin")
891 : end if
892 : #endif
893 :
894 : ! make sure to use sizes from apply_invovl
895 : call f_gpu_apply_invovl_inner_alloc(proj_dim, ham%ntypat, 0)
896 :
897 : #ifdef DEBUG_VERBOSE_GPU
898 : if(xmpi_comm_rank(xmpi_world) == 0) then
899 : call check_gpu_mem("gpu_apply_invovl_inner_alloc end")
900 : end if
901 : #endif
902 :
903 : ! TODO find a better place to put that initialization
904 : call f_gpu_init_invovl_data(indlmn_dim, c_loc(ham%indlmn(1,1,1)))
905 :
906 : end if
907 :
908 : #endif
909 :
910 :
911 160806 : call timab(timer_apply_inv_ovl_opernla, 1, tsec)
912 :
913 : ! cwaveprj may be dummy or unused if gemm nonlop is turned on
914 482418 : if((.not. gemm_nonlop_use_gemm) .or. size(cwaveprj) > 1) then
915 4200552 : ABI_MALLOC(cwaveprj_in, (ham%natom,nspinor*ndat))
916 160518 : call pawcprj_alloc(cwaveprj_in,0,ham%dimcprj)
917 : else
918 864 : ABI_MALLOC(cwaveprj_in, (1,1))
919 288 : call pawcprj_alloc(cwaveprj_in,0,(/1/))
920 : end if
921 : ABI_NVTX_END_RANGE()
922 :
923 : ! get the cprj
924 : ABI_NVTX_START_RANGE(NVTX_INVOVL_NONLOP1)
925 160806 : choice = 0 ! only compute cprj, nothing else
926 160806 : cpopt = 0 ! compute and save cprj
927 160806 : paw_opt = 3 ! S nonlocal operator
928 160806 : if (mpi_enreg%paral_kgb==1) then
929 : call prep_nonlop(choice,cpopt,cwaveprj_in,enlout,ham,idir,lambda_block,ndat,mpi_enreg,&
930 : & nnlout,paw_opt,signs,sm1cwavef,tim_nonlop,cwavef,gvnlxc,&
931 : & already_transposed=.true.,&
932 : & gpu_option=ham%gpu_option,&
933 62460 : & vectproj=proj)
934 : else
935 : call nonlop(choice,cpopt,cwaveprj_in,enlout,ham,idir,lambda_block,mpi_enreg,ndat,&
936 98346 : & nnlout,paw_opt,signs,sm1cwavef,tim_nonlop,cwavef,gvnlxc,vectproj=proj)
937 : end if
938 : ABI_NVTX_END_RANGE()
939 :
940 160806 : call timab(timer_apply_inv_ovl_opernla, 2, tsec)
941 160806 : call timab(timer_apply_inv_ovl_inv_s, 1, tsec)
942 :
943 : ! If using GEMM nonlop, proj array is used directly instead of writing in cwaveprj_in content, so skip this copy
944 160806 : if(.not. gemm_nonlop_use_gemm) then
945 : ! copy cwaveprj_in to proj(:,:)
946 1288502 : do idat=1, ndat*nspinor
947 1128272 : shift = 0
948 3713526 : do iatom = 1, ham%natom
949 2425024 : nlmn = cwaveprj_in(iatom, idat)%nlmn
950 66317440 : proj(1:cplx, shift+1:shift+nlmn, idat) = cwaveprj_in(iatom, idat)%cp(1:cplx, 1:nlmn)
951 3553296 : shift = shift + nlmn
952 : end do
953 : end do
954 : end if
955 :
956 : !multiply by S^1
957 : ABI_NVTX_START_RANGE(NVTX_INVOVL_INNER)
958 : ! TODO : when solve_inner_gpu is ready, update the following to activate GPU computation
959 160806 : if (ham%gpu_option == ABI_GPU_LEGACY .or. ham%gpu_option==ABI_GPU_KOKKOS) then
960 :
961 : #if defined(HAVE_FC_ISO_C_BINDING) && defined(HAVE_GPU_CUDA)
962 :
963 : if (mpi_enreg%nproc_fft /= 1) then
964 : ABI_ERROR("[66_wfs/m_invovl.F90:apply_invovl] nproc_fft must be 1, when GPU/CUDA is activated")
965 : end if
966 :
967 : call f_solve_inner_gpu(proj_dim, c_loc(proj(1,1,1)), &
968 : & c_loc(sm1proj(1,1,1)), c_loc(PtPsm1proj(1,1,1)), &
969 : & nattyp_dim, c_loc(ham%nattyp(1)), ham%ntypat, &
970 : & ham%lmnmax, cplx, block_sliced)
971 :
972 : #endif
973 :
974 : else
975 :
976 160806 : call solve_inner(invovl, ham, cplx, mpi_enreg, proj, ndat*nspinor, sm1proj, PtPsm1proj, block_sliced)
977 65308502 : sm1proj = - sm1proj
978 65308502 : PtPsm1proj = - PtPsm1proj
979 : end if
980 :
981 : ABI_NVTX_END_RANGE()
982 :
983 : ! If using GEMM nonlop, sm1proj array is used directly instead of reading cwaveprj content, so skip this copy
984 160806 : if(.not. gemm_nonlop_use_gemm) then
985 : ! copy sm1proj to cwaveprj(:,:)
986 1288502 : do idat=1, ndat*nspinor
987 1128272 : shift = 0
988 3713526 : do iatom = 1, ham%natom
989 2425024 : nlmn = cwaveprj(iatom, idat)%nlmn
990 66317440 : cwaveprj(iatom, idat)%cp(1:cplx, 1:nlmn) = sm1proj(1:cplx, shift+1:shift+nlmn, idat)
991 3553296 : shift = shift + nlmn
992 : end do
993 : end do
994 : end if
995 160806 : call timab(timer_apply_inv_ovl_inv_s, 2, tsec)
996 160806 : call timab(timer_apply_inv_ovl_opernlb, 1, tsec)
997 :
998 : ! get the corresponding wf
999 : ABI_NVTX_START_RANGE(NVTX_INVOVL_NONLOP2)
1000 160806 : cpopt = 2 ! reuse cprj
1001 160806 : choice = 7 ! get wf from cprj, without the application of S
1002 : paw_opt = 3
1003 160806 : if (mpi_enreg%paral_kgb==1) then
1004 : call prep_nonlop(choice,cpopt,cwaveprj,enlout,ham,idir,lambda_block,ndat,mpi_enreg,nnlout,&
1005 : & paw_opt,signs,sm1cwavef,tim_nonlop,cwavef,gvnlxc,already_transposed=.true.,&
1006 62460 : & gpu_option=ham%gpu_option,vectproj=sm1proj)
1007 : else
1008 : call nonlop(choice,cpopt,cwaveprj,enlout,ham,idir,lambda_block,mpi_enreg,ndat,nnlout,paw_opt,&
1009 98346 : & signs,sm1cwavef,tim_nonlop,cwavef,gvnlxc,vectproj=sm1proj)
1010 : end if
1011 : ABI_NVTX_END_RANGE()
1012 :
1013 160806 : call timab(timer_apply_inv_ovl_opernlb, 2, tsec)
1014 :
1015 : ABI_NVTX_START_RANGE(NVTX_INVOVL_POST2)
1016 482418 : if(size(cwaveprj) > 1) then
1017 : ! copy PtPSm1proj to cwaveprj(:,:)
1018 1290518 : do idat=1, ndat*nspinor
1019 1130000 : shift = 0
1020 3718998 : do iatom = 1, ham%natom
1021 2428480 : nlmn = cwaveprj(iatom, idat)%nlmn
1022 66403840 : cwaveprj(iatom, idat)%cp(1:cplx, 1:nlmn) = PtPsm1proj(1:cplx, shift+1:shift+nlmn, idat)
1023 3558480 : shift = shift + nlmn
1024 : end do
1025 : end do
1026 : !cwaveprj_in is empty if GEMM nonlop is being used, so populate it here
1027 160518 : if(gemm_nonlop_use_gemm) then
1028 2016 : do idat=1, ndat*nspinor
1029 1728 : shift = 0
1030 5472 : do iatom = 1, ham%natom
1031 3456 : nlmn = cwaveprj_in(iatom, idat)%nlmn
1032 86400 : cwaveprj_in(iatom, idat)%cp(1:cplx, 1:nlmn) = proj(1:cplx, shift+1:shift+nlmn, idat)
1033 5184 : shift = shift + nlmn
1034 : end do
1035 : end do
1036 : end if
1037 160518 : call pawcprj_axpby(one, one, cwaveprj_in, cwaveprj)
1038 : end if
1039 160806 : call pawcprj_free(cwaveprj_in)
1040 2589574 : ABI_FREE(cwaveprj_in)
1041 :
1042 160806 : if (ham%gpu_option == ABI_GPU_LEGACY .or. ham%gpu_option==ABI_GPU_KOKKOS) then
1043 : #if defined(HAVE_GPU_CUDA) && defined(HAVE_KOKKOS) && defined(HAVE_YAKL)
1044 : cwavef_size = 2*npw*nspinor*ndat
1045 : call add_array_kokkos(c_loc(sm1cwavef), c_loc(cwavef), cwavef_size)
1046 : #endif
1047 : else
1048 669394662 : sm1cwavef = cwavef + sm1cwavef
1049 : end if
1050 :
1051 : ABI_NVTX_END_RANGE()
1052 :
1053 160806 : ABI_FREE(proj)
1054 160806 : ABI_FREE(sm1proj)
1055 160806 : ABI_FREE(PtPsm1proj)
1056 :
1057 160806 : end subroutine apply_invovl
1058 : !!***
1059 :
1060 : !!****f* m_invovl/solve_inner
1061 : !! NAME
1062 : !! solve_inner
1063 : !!
1064 : !! FUNCTION
1065 : !! Helper function: iteratively solves the inner system
1066 : !!
1067 : !! INPUTS
1068 : !!
1069 : !! SOURCE
1070 160806 : subroutine solve_inner(invovl, ham, cplx, mpi_enreg, proj, ndat, sm1proj, PtPsm1proj, block_sliced)
1071 :
1072 : use m_abi_linalg
1073 :
1074 : integer,intent(in) :: ndat,cplx
1075 : type(invovl_kpt_type), intent(in) :: invovl
1076 : real(dp), intent(inout) :: proj(cplx, invovl%nprojs,ndat)
1077 : real(dp), intent(inout) :: sm1proj(cplx, invovl%nprojs, ndat)
1078 : real(dp), intent(inout) :: PtPsm1proj(cplx, invovl%nprojs, ndat)
1079 160806 : real(dp), allocatable :: temp_proj(:,:,:)
1080 : type(mpi_type), intent(in) :: mpi_enreg
1081 : type(gs_hamiltonian_type),intent(in) :: ham
1082 : integer, intent(in) :: block_sliced
1083 :
1084 321612 : integer :: array_nlmntot_pp(mpi_enreg%nproc_fft)
1085 : integer :: nlmntot_this_proc, ibeg, iend, ierr, i, nprojs
1086 321612 : real(dp) :: resid(cplx, invovl%nprojs,ndat), precondresid(cplx, invovl%nprojs,ndat)
1087 321612 : real(dp) :: normprojs(ndat), errs(ndat), maxerr, previous_maxerr
1088 : character(len=500) :: message
1089 :
1090 : real(dp), parameter :: precision = 1e-16 ! maximum relative error. TODO: use tolwfr ?
1091 : real(dp) :: convergence_rate
1092 : integer :: additional_steps_to_take
1093 :
1094 : ! *************************************************************************
1095 :
1096 160806 : nprojs = invovl%nprojs
1097 65308502 : normprojs = SUM(SUM(proj**2, 1),1)
1098 160806 : call xmpi_sum(normprojs, mpi_enreg%comm_fft, ierr)
1099 :
1100 : ! Compute work distribution : split nprojs evenly between the fft processors
1101 160806 : if(mpi_enreg%paral_kgb == 1) then
1102 125208 : array_nlmntot_pp(:) = nprojs / mpi_enreg%nproc_fft
1103 : ! not enough work, there's MOD(nprojs,mpi_enreg%nproc_fft) tasks left
1104 : ! assign them to the first ones
1105 62460 : array_nlmntot_pp(1:MOD(nprojs,mpi_enreg%nproc_fft)) = array_nlmntot_pp(1:MOD(nprojs,mpi_enreg%nproc_fft)) + 1
1106 62604 : ibeg = SUM(array_nlmntot_pp(1:mpi_enreg%me_fft)) + 1
1107 62460 : iend = ibeg + array_nlmntot_pp(1+mpi_enreg%me_fft) - 1
1108 62460 : nlmntot_this_proc = iend - ibeg + 1
1109 : else
1110 98346 : ibeg = 1
1111 98346 : iend = nprojs
1112 98346 : nlmntot_this_proc = nprojs
1113 : end if
1114 :
1115 804030 : ABI_MALLOC(temp_proj, (cplx, nlmntot_this_proc, ndat))
1116 :
1117 : ! first guess for sm1proj
1118 160806 : call apply_block(ham, cplx, invovl%inv_s_approx, nprojs, ndat, proj, sm1proj, block_sliced)
1119 :
1120 : ! Iterative refinement
1121 : ! TODO use a more efficient iterative algorithm than iterative refinement, use locking
1122 160806 : additional_steps_to_take = -1
1123 1476115 : do i=1, 30
1124 : ! compute resid = proj - (D^-1 + PtP)sm1proj
1125 1476115 : call apply_block(ham, cplx, invovl%inv_sij, nprojs, ndat, sm1proj, resid, block_sliced)
1126 724585084 : temp_proj = sm1proj(:,ibeg:iend,:)
1127 :
1128 : ! compute matrix multiplication : PtPsm1proj(:,:,1) = invovl%gram * temp_proj(:,:,1)
1129 : call abi_xgemm('N', 'N', nprojs, ndat, nlmntot_this_proc, cone, &
1130 : & invovl%gram_projs(:,:,1), nprojs, &
1131 : & temp_proj(:,:,1), nlmntot_this_proc, czero, &
1132 : & PtPsm1proj(:,:,1), nprojs, &
1133 1476115 : & x_cplx=cplx)
1134 1476115 : call xmpi_sum(PtPsm1proj, mpi_enreg%comm_fft, ierr)
1135 723593097 : resid = proj - resid - Ptpsm1proj
1136 : ! exit check
1137 723593097 : errs = SUM(SUM(resid**2, 1),1)
1138 1476115 : call xmpi_sum(errs, mpi_enreg%comm_fft, ierr)
1139 :
1140 13553758 : maxerr = sqrt(MAXVAL(errs/normprojs))
1141 1476115 : if(maxerr < precision .or. additional_steps_to_take == 1) then
1142 : exit
1143 : ! We might stall and never get to the specified precision because of machine errors.
1144 : ! If we got to 1e-10, extrapolate convergence rate and determine the number of additional
1145 : ! steps to take to reach precision
1146 1398471 : else if(maxerr < 1e-10 .and. additional_steps_to_take == -1) then
1147 160806 : convergence_rate = -LOG(1e-10) / i
1148 160806 : additional_steps_to_take = CEILING(-LOG(precision/1e-10)/convergence_rate) + 1
1149 1237665 : else if(additional_steps_to_take > 0) then
1150 525573 : if(previous_maxerr<maxerr)exit
1151 442411 : additional_steps_to_take = additional_steps_to_take - 1
1152 : end if
1153 1315309 : previous_maxerr=maxerr
1154 :
1155 : ! add preconditionned residual
1156 1315309 : call apply_block(ham, cplx, invovl%inv_s_approx, nprojs, ndat, resid, precondresid, block_sliced)
1157 658445401 : sm1proj = sm1proj + precondresid
1158 : end do
1159 :
1160 160806 : if(maxerr >= precision .and. maxerr >= 1e-10) then
1161 0 : write(message, *) 'In invovl, max error was', maxerr, ' after 30 iterations'
1162 0 : ABI_WARNING(message)
1163 : else
1164 : ! write(message,'(a,i2,a,es13.5)') 'Iterative solver in invovl finished in ', i, ' iterations, error', maxerr
1165 : ! call wrtout(std_out,message,'COLL')
1166 : end if
1167 :
1168 160806 : ABI_FREE(temp_proj)
1169 :
1170 160806 : end subroutine solve_inner
1171 : !!***
1172 :
1173 : !!****f* m_invovl/apply_block
1174 : !! NAME
1175 : !! apply_block
1176 : !!
1177 : !! FUNCTION
1178 : !! Helper function: applies a block-diagonal matrix mat(lmnmax, lmnmax, ntypat)
1179 : !!
1180 : !! INPUTS
1181 : !!
1182 : !! SOURCE
1183 2952230 : subroutine apply_block(ham, cplx, mat, nprojs, ndat, x, y, block_sliced)
1184 :
1185 : use m_abi_linalg
1186 :
1187 : integer,intent(in) :: ndat, nprojs, cplx
1188 : real(dp), intent(inout), target :: x(cplx, nprojs, ndat), y(cplx, nprojs, ndat)
1189 : type(gs_hamiltonian_type),intent(in) :: ham
1190 : real(dp), intent(in) :: mat(cplx, ham%lmnmax, ham%lmnmax, ham%ntypat)
1191 : integer, intent(in) :: block_sliced
1192 :
1193 : integer :: nlmn, shift, itypat, idat
1194 2952230 : real(dp),pointer :: work_x(:,:),work_y(:,:)
1195 :
1196 : ! *************************************************************************
1197 :
1198 2952230 : if (block_sliced == 1) then
1199 :
1200 24155286 : do idat = 1, ndat
1201 21203056 : shift = 1
1202 52113246 : do itypat=1, ham%ntypat
1203 271874380 : nlmn = count(ham%indlmn(3,:,itypat)>0)
1204 : !! apply mat to all atoms at once
1205 : ! perform natom multiplications of size nlmn
1206 : ! compute y = mat*x
1207 27957960 : work_x => x(:, shift:shift+nlmn*ham%nattyp(itypat)-1, idat)
1208 27957960 : work_y => y(:, shift:shift+nlmn*ham%nattyp(itypat)-1, idat)
1209 27957960 : if(cplx == 2) then
1210 : call ZHEMM('L','U', nlmn, ham%nattyp(itypat), cone, &
1211 : & mat(:, :, :, itypat), ham%lmnmax, &
1212 : & work_x, nlmn, czero, &
1213 25098432 : & work_y, nlmn)
1214 : else
1215 : call DSYMM('L','U', nlmn, ham%nattyp(itypat), one, &
1216 : & mat(:, :, :, itypat), ham%lmnmax, &
1217 : & work_x, nlmn, zero, &
1218 2859528 : & work_y, nlmn)
1219 : end if
1220 49161016 : shift = shift + nlmn*ham%nattyp(itypat)
1221 : end do
1222 : end do
1223 :
1224 : else ! block_sliced = 0
1225 :
1226 : shift = 1
1227 0 : do itypat=1, ham%ntypat
1228 0 : nlmn = count(ham%indlmn(3,:,itypat)>0)
1229 : !! apply mat to all atoms at once, all idat at once
1230 : ! perform natom multiplications of size nlmn
1231 : ! be careful here matrix extracted from x and y are not memory contiguous
1232 : ! ==> so in the GPU version we will need to adapt leading dimension
1233 0 : if(cplx == 2) then
1234 : call ZHEMM('L','U', nlmn, ham%nattyp(itypat)*ndat, cone, &
1235 : & mat(:, :, :, itypat), ham%lmnmax, &
1236 : & x(:, 1:nlmn*ham%nattyp(itypat), 1:ndat), nlmn, czero, &
1237 0 : & y(:, 1:shift+nlmn*ham%nattyp(itypat)-1, 1:ndat), nlmn)
1238 : else
1239 : call DSYMM('L','U', nlmn, ham%nattyp(itypat)*ndat, one, &
1240 : & mat(:, :, :, itypat), ham%lmnmax, &
1241 : & x(:, shift:shift+nlmn*ham%nattyp(itypat)-1, 1:ndat), nlmn, zero, &
1242 0 : & y(:, shift:shift+nlmn*ham%nattyp(itypat)-1, 1:ndat), nlmn)
1243 : end if
1244 0 : shift = shift + nlmn*ham%nattyp(itypat)
1245 : end do
1246 :
1247 : end if
1248 :
1249 2952230 : end subroutine apply_block
1250 : !!***
1251 :
1252 0 : function invovl_ompgpu_work_mem(ham, ndat) result(req_mem)
1253 :
1254 : type(gs_hamiltonian_type), intent(in) :: ham
1255 : integer, intent(in) :: ndat
1256 : integer :: nprojs, cplx, itypat
1257 : integer(kind=c_size_t) :: req_mem
1258 :
1259 0 : nprojs = 0
1260 0 : do itypat=1,ham%ntypat
1261 0 : nprojs = nprojs + count(ham%indlmn(3,:,itypat)>0)*ham%nattyp(itypat)
1262 : end do
1263 0 : cplx = 2; if(ham%istwf_k == 2) cplx = 1
1264 :
1265 0 : req_mem = 0
1266 0 : req_mem = req_mem + dp * cplx * int(nprojs, c_size_t) * int(ndat, c_size_t) ! proj
1267 : req_mem = req_mem + dp * cplx * int(nprojs, c_size_t) * int(ndat, c_size_t) ! sm1proj
1268 : req_mem = req_mem + dp * cplx * int(nprojs, c_size_t) * int(ndat, c_size_t) ! PtPsm1proj
1269 0 : req_mem = req_mem + dp * cplx * int(nprojs, c_size_t) * int(ndat, c_size_t) ! resid (solve_inner)
1270 0 : req_mem = req_mem + dp * cplx * int(nprojs, c_size_t) * int(ndat, c_size_t) ! precondresid (solve_inner)
1271 :
1272 0 : end function invovl_ompgpu_work_mem
1273 :
1274 0 : function invovl_ompgpu_static_mem(ham) result(req_mem)
1275 :
1276 : type(gs_hamiltonian_type), intent(in) :: ham
1277 : integer :: nprojs, cplx, itypat
1278 : integer(kind=c_size_t) :: req_mem
1279 :
1280 0 : nprojs = 0
1281 0 : do itypat=1,ham%ntypat
1282 0 : nprojs = nprojs + count(ham%indlmn(3,:,itypat)>0)*ham%nattyp(itypat)
1283 : end do
1284 0 : cplx = 2; if(ham%istwf_k > 1) cplx = 1
1285 :
1286 0 : req_mem = 0
1287 0 : req_mem = req_mem + dp * cplx * int(nprojs, c_size_t) * int(nprojs, c_size_t) ! gram_projs
1288 : req_mem = req_mem + dp * cplx * int(ham%lmnmax, c_size_t) * &
1289 0 : & int(ham%lmnmax, c_size_t) * int(ham%ntypat, c_size_t) ! inv_sij
1290 : req_mem = req_mem + dp * cplx * int(ham%lmnmax, c_size_t) * &
1291 0 : & int(ham%lmnmax, c_size_t) * int(ham%ntypat, c_size_t) ! inv_s_approx
1292 :
1293 0 : end function invovl_ompgpu_static_mem
1294 :
1295 : #ifdef HAVE_OPENMP_OFFLOAD
1296 : !*******************************************************************************************************************************!
1297 : ! OpenMP GPU routines !
1298 : !*******************************************************************************************************************************!
1299 :
1300 : !!****f* m_invovl/apply_invovl_ompgpu
1301 : !! NAME
1302 : !! apply_invovl_ompgpu
1303 : !!
1304 : !! FUNCTION
1305 : !! Applies the inverse of the overlap matrix to cwavef (OpenMP GPU implementation)
1306 : !!
1307 : !! INPUTS
1308 : !!
1309 : !! SOURCE
1310 :
1311 : subroutine apply_invovl_ompgpu(ham, cwavef, sm1cwavef, cwaveprj, npw, ndat, mpi_enreg, nspinor, block_sliced)
1312 :
1313 : #if defined(HAVE_FC_ISO_C_BINDING) && defined(HAVE_GPU)
1314 : use, intrinsic :: iso_c_binding
1315 : #endif
1316 :
1317 : ! args
1318 : type(gs_hamiltonian_type), intent(in), target :: ham
1319 : integer, intent(in) :: npw, ndat
1320 : integer, intent(in) :: nspinor
1321 : integer, intent(in) :: block_sliced
1322 : real(dp), intent(inout), target :: cwavef(2, npw*nspinor*ndat) ! TODO should be in, fix nonlop
1323 : type(mpi_type) :: mpi_enreg
1324 : real(dp), intent(inout), target :: sm1cwavef(2, npw*nspinor*ndat)
1325 : type(pawcprj_type), intent(inout) :: cwaveprj(:,:)
1326 : logical :: transfer_omp_args
1327 :
1328 : real(dp), ABI_CONTIGUOUS pointer :: proj(:,:,:),sm1proj(:,:,:),PtPsm1proj(:,:,:)
1329 :
1330 : integer :: idat, iatom, icplx, iproj, nprojs, nlmn, shift
1331 : real(dp) :: tsec(2)
1332 :
1333 : integer :: choice, cpopt, paw_opt , cplx, old_me_g0
1334 : type(pawcprj_type),allocatable :: cwaveprj_in(:,:)
1335 :
1336 : integer :: ikpt_this_proc
1337 : ! dummies
1338 : real(dp) :: enlout(ndat), lambda_block(1), gvnlxc(1,1)
1339 : integer, parameter :: nnlout = 0, idir = 0, signs = 2
1340 :
1341 : type(invovl_kpt_type), pointer :: invovl
1342 :
1343 : ! *************************************************************************
1344 :
1345 : ikpt_this_proc=bandfft_kpt_get_ikpt()
1346 : invovl => invovl_kpt(ikpt_this_proc)
1347 : nprojs=invovl%nprojs
1348 : if(ikpt_this_proc /= current_ikpt_in_gpu) call refresh_invovl_ompgpu_kpt(ikpt_this_proc)
1349 :
1350 : if(ham%istwf_k == 1) then
1351 : cplx = 2
1352 : else
1353 : cplx = 1
1354 : end if
1355 : call alloc_ompgpu_buffers(cplx,nprojs,nspinor,ndat)
1356 : proj => proj_ompgpu
1357 : sm1proj => sm1proj_ompgpu
1358 : PtPsm1proj => PtPsm1proj_ompgpu
1359 : !$OMP TARGET ENTER DATA MAP(alloc:proj,sm1proj,PtPsm1proj)
1360 :
1361 : transfer_omp_args = .not. ( xomp_target_is_present(c_loc(sm1cwavef)) &
1362 : .and. xomp_target_is_present(c_loc(cwavef)))
1363 : !$OMP TARGET ENTER DATA MAP(alloc:gvnlxc)
1364 : if(transfer_omp_args) then
1365 : !$OMP TARGET ENTER DATA MAP(alloc:sm1cwavef,cwavef)
1366 : !$OMP TARGET UPDATE TO(sm1cwavef,cwavef)
1367 : end if
1368 :
1369 : call timab(timer_apply_inv_ovl_opernla, 1, tsec)
1370 :
1371 : ! get the cprj
1372 : ABI_NVTX_START_RANGE(NVTX_INVOVL_NONLOP1)
1373 : choice = 0 ! only compute cprj, nothing else
1374 : cpopt = 0 ! compute and save cprj
1375 : paw_opt = 3 ! S nonlocal operator
1376 :
1377 : if(ham%istwf_k==2) then
1378 : old_me_g0=mpi_enreg%me_g0
1379 : if (mpi_enreg%me_fft==0) then
1380 : mpi_enreg%me_g0=1
1381 : else
1382 : mpi_enreg%me_g0=0
1383 : end if
1384 : end if
1385 : call nonlop(choice,cpopt,cwaveprj_in,enlout,ham,idir,lambda_block,mpi_enreg,ndat,nnlout,&
1386 : paw_opt,signs,sm1cwavef,tim_nonlop,cwavef,gvnlxc,vectproj=proj)
1387 : ABI_NVTX_END_RANGE()
1388 :
1389 : call timab(timer_apply_inv_ovl_opernla, 2, tsec)
1390 : call timab(timer_apply_inv_ovl_inv_s, 1, tsec)
1391 :
1392 : !multiply by S^1
1393 : ABI_NVTX_START_RANGE(NVTX_INVOVL_INNER)
1394 : call solve_inner_ompgpu(invovl, ham, cplx, mpi_enreg, proj, ndat*nspinor, sm1proj, PtPsm1proj, block_sliced)
1395 : !$OMP TARGET TEAMS DISTRIBUTE MAP(to:sm1proj,PtPsm1proj)
1396 : do idat =1, ndat*nspinor
1397 : !$OMP PARALLEL DO COLLAPSE(2) PRIVATE(iproj,icplx)
1398 : do iproj = 1, nprojs
1399 : do icplx = 1, cplx
1400 : sm1proj(icplx,iproj,idat) = - sm1proj(icplx,iproj,idat)
1401 : PtPsm1proj(icplx,iproj,idat) = - PtPsm1proj(icplx,iproj,idat)
1402 : end do
1403 : end do
1404 : end do
1405 :
1406 : ABI_NVTX_END_RANGE()
1407 :
1408 : call timab(timer_apply_inv_ovl_inv_s, 2, tsec)
1409 : call timab(timer_apply_inv_ovl_opernlb, 1, tsec)
1410 :
1411 : ! get the corresponding wf
1412 : ABI_NVTX_START_RANGE(NVTX_INVOVL_NONLOP2)
1413 : cpopt = 2 ! reuse cprj
1414 : choice = 7 ! get wf from cprj, without the application of S
1415 : paw_opt = 3
1416 : call nonlop(choice,cpopt,cwaveprj,enlout,ham,idir,lambda_block,mpi_enreg,ndat,nnlout,paw_opt,&
1417 : signs,sm1cwavef,tim_nonlop,cwavef,gvnlxc,vectproj=sm1proj)
1418 : ABI_NVTX_END_RANGE()
1419 :
1420 : call timab(timer_apply_inv_ovl_opernlb, 2, tsec)
1421 : if (ham%istwf_k==2) mpi_enreg%me_g0=old_me_g0
1422 :
1423 : if(size(cwaveprj) > 1) then
1424 : ABI_MALLOC(cwaveprj_in, (ham%natom,nspinor*ndat))
1425 : call pawcprj_alloc(cwaveprj_in,0,ham%dimcprj)
1426 : !$OMP TARGET UPDATE FROM(PtPsm1proj,proj)
1427 : ! copy PtPsm1proj to cwaveprj(:,:)
1428 : do idat=1, ndat*nspinor
1429 : shift = 0
1430 : do iatom = 1, ham%natom
1431 : nlmn = cwaveprj(iatom, idat)%nlmn
1432 : cwaveprj(iatom, idat)%cp(1:cplx, 1:nlmn) = PtPsm1proj(1:cplx, shift+1:shift+nlmn, idat)
1433 : shift = shift + nlmn
1434 : end do
1435 : end do
1436 : do idat=1, ndat*nspinor
1437 : shift = 0
1438 : do iatom = 1, ham%natom
1439 : nlmn = cwaveprj_in(iatom, idat)%nlmn
1440 : cwaveprj_in(iatom, idat)%cp(1:cplx, 1:nlmn) = proj(1:cplx, shift+1:shift+nlmn, idat)
1441 : shift = shift + nlmn
1442 : end do
1443 : end do
1444 : call pawcprj_axpby(one, one, cwaveprj_in, cwaveprj)
1445 : call pawcprj_free(cwaveprj_in)
1446 : ABI_FREE(cwaveprj_in)
1447 : end if
1448 :
1449 : call abi_xaxpy(2*npw*nspinor*ndat, cone, cwavef, 1, sm1cwavef, 1, gpu_option=ABI_GPU_OPENMP)
1450 :
1451 : if(transfer_omp_args) then
1452 : !$OMP TARGET UPDATE FROM(sm1cwavef,cwavef)
1453 : !$OMP TARGET EXIT DATA MAP(delete:sm1cwavef,cwavef)
1454 : end if
1455 :
1456 : !$OMP TARGET EXIT DATA MAP(delete:gvnlxc)
1457 : !$OMP TARGET EXIT DATA MAP(delete:proj,sm1proj,PtPsm1proj)
1458 :
1459 : end subroutine apply_invovl_ompgpu
1460 : !!***
1461 :
1462 : !!****f* m_invovl/solve_inner_ompgpu
1463 : !! NAME
1464 : !! solve_inner_ompgpu
1465 : !!
1466 : !! FUNCTION
1467 : !! Helper function: iteratively solves the inner system (OpenMP GPU offload implementation)
1468 : !!
1469 : !! INPUTS
1470 : !!
1471 : !! SOURCE
1472 : subroutine solve_inner_ompgpu(invovl, ham, cplx, mpi_enreg, proj, ndat, sm1proj, PtPsm1proj, block_sliced)
1473 :
1474 : use m_abi_linalg
1475 :
1476 : integer,intent(in) :: ndat,cplx
1477 : type(invovl_kpt_type), intent(in), target :: invovl
1478 : real(dp), intent(inout) :: proj(cplx, invovl%nprojs,ndat)
1479 : real(dp), intent(inout), target :: sm1proj(cplx, invovl%nprojs, ndat)
1480 : real(dp), intent(inout), target :: PtPsm1proj(cplx, invovl%nprojs, ndat)
1481 : type(mpi_type), intent(in) :: mpi_enreg
1482 : type(gs_hamiltonian_type),intent(in) :: ham
1483 : integer, intent(in) :: block_sliced
1484 :
1485 : integer :: array_nlmntot_pp(mpi_enreg%nproc_fft)
1486 : integer :: nlmntot_this_proc, ibeg, iend, ierr, i, nprojs
1487 : real(dp) :: resid(cplx, invovl%nprojs,ndat), precondresid(cplx, invovl%nprojs,ndat)
1488 : real(dp) :: normprojs(ndat), errs(ndat), maxerr, previous_maxerr
1489 : character(len=500) :: message
1490 :
1491 : real(dp), parameter :: precision = 1e-16 ! maximum relative error. TODO: use tolwfr ?
1492 : real(dp) :: convergence_rate,sum_tmp
1493 : integer :: additional_steps_to_take,idat,iproj,icplx
1494 : integer :: Ptsize(3)
1495 :
1496 : ! *************************************************************************
1497 :
1498 : Ptsize(1) = cplx
1499 : Ptsize(2) = invovl%nprojs
1500 : Ptsize(3) = ndat
1501 : nprojs = invovl%nprojs
1502 :
1503 : !$OMP TARGET ENTER DATA MAP(alloc:errs,precondresid,resid,normprojs)
1504 :
1505 : !FIXME LLVM has trouble with performing team reduction (AOMP 15.0.2)
1506 : #ifdef FC_LLVM
1507 : !$OMP TARGET UPDATE FROM(proj)
1508 : #else
1509 : !$OMP TARGET TEAMS DISTRIBUTE MAP(to:normprojs,proj) PRIVATE(idat,sum_tmp)
1510 : #endif
1511 : do idat = 1,ndat
1512 : sum_tmp=0
1513 : #ifndef FC_LLVM
1514 : !$OMP PARALLEL DO COLLAPSE(2) REDUCTION(+:sum_tmp) PRIVATE(iproj,icplx)
1515 : #endif
1516 : do iproj = 1,nprojs
1517 : do icplx = 1,cplx
1518 : sum_tmp = sum_tmp + proj(icplx,iproj,idat)**2
1519 : end do
1520 : end do
1521 : normprojs(idat)=sum_tmp
1522 : end do
1523 : #ifndef FC_LLVM
1524 : !$OMP TARGET UPDATE FROM(normprojs)
1525 : #endif
1526 :
1527 : ibeg = 1
1528 : iend = nprojs
1529 : nlmntot_this_proc = nprojs
1530 :
1531 : ! first guess for sm1proj
1532 : call apply_block_ompgpu(ham, cplx, invovl%inv_s_approx, nprojs, ndat, proj, sm1proj, block_sliced)
1533 :
1534 : ! Iterative refinement
1535 : ! TODO use a more efficient iterative algorithm than iterative refinement, use locking
1536 : additional_steps_to_take = -1
1537 : do i=1, 30
1538 : #ifdef FC_NVHPC
1539 : ! Silly fix for NVHPC 25.1
1540 : if(ndat == -42) write(100,*) ndat
1541 : #endif
1542 : ! compute resid = proj - (D^-1 + PtP)sm1proj
1543 : call apply_block_ompgpu(ham, cplx, invovl%inv_sij, nprojs, ndat, sm1proj, resid, block_sliced)
1544 :
1545 : ! compute matrix multiplication : PtPsm1proj(:,:,1) = invovl%gram * sm1proj(:,:,1)
1546 : ABI_NVTX_START_RANGE(NVTX_INVOVL_INNER_GEMM)
1547 : call abi_xgemm('N', 'N', nprojs, ndat, nlmntot_this_proc, cone, &
1548 : & current_gram_projs, nprojs,&
1549 : & sm1proj, nlmntot_this_proc, czero, &
1550 : & PtPsm1proj, nprojs, &
1551 : & x_cplx=cplx, gpu_option=ABI_GPU_OPENMP)
1552 :
1553 : !$OMP TARGET TEAMS DISTRIBUTE &
1554 : !$OMP& PRIVATE(idat) MAP(to:proj,resid,PtPsm1proj)
1555 : do idat =1, ndat
1556 : !$OMP PARALLEL DO COLLAPSE(2) PRIVATE(iproj,icplx)
1557 : do iproj =1, nprojs
1558 : do icplx = 1,cplx
1559 : resid(icplx, iproj, idat) = proj(icplx, iproj, idat) - resid(icplx, iproj, idat) - PtPsm1proj(icplx, iproj, idat)
1560 : end do
1561 : end do
1562 : end do
1563 :
1564 : ! exit check
1565 : #ifdef FC_LLVM
1566 : !FIXME LLVM has trouble with performing team reduction (v16.0.0 from AMD ROCm 5.6.0)
1567 : !$OMP TARGET UPDATE FROM(resid)
1568 : errs = SUM(SUM(resid**2, 1),1)
1569 : #else
1570 : !$OMP TARGET TEAMS DISTRIBUTE MAP(to:errs,resid) PRIVATE(idat,sum_tmp)
1571 : do idat = 1,ndat
1572 : sum_tmp=0
1573 : !$OMP PARALLEL DO COLLAPSE(2) REDUCTION(+:sum_tmp) PRIVATE(iproj,icplx)
1574 : do iproj = 1,nprojs
1575 : do icplx = 1,cplx
1576 : sum_tmp = sum_tmp + resid(icplx,iproj,idat)**2
1577 : end do
1578 : end do
1579 : errs(idat)=sum_tmp
1580 : end do
1581 : !$OMP TARGET UPDATE FROM(errs)
1582 : #endif
1583 :
1584 : ABI_NVTX_END_RANGE()
1585 :
1586 : maxerr = sqrt(MAXVAL(errs/normprojs))
1587 : if(maxerr < precision .or. additional_steps_to_take == 1) then
1588 : exit
1589 : ! We might stall and never get to the specified precision because of machine errors.
1590 : ! If we got to 1e-10, extrapolate convergence rate and determine the number of additional
1591 : ! steps to take to reach precision
1592 : else if(maxerr < 1e-10 .and. additional_steps_to_take == -1) then
1593 : convergence_rate = -LOG(1e-10) / i
1594 : additional_steps_to_take = CEILING(-LOG(precision/1e-10)/convergence_rate) + 1
1595 : else if(additional_steps_to_take > 0) then
1596 : if(previous_maxerr<maxerr)exit
1597 : additional_steps_to_take = additional_steps_to_take - 1
1598 : end if
1599 : previous_maxerr=maxerr
1600 :
1601 : ! add preconditionned residual
1602 : call apply_block_ompgpu(ham, cplx, invovl%inv_s_approx, nprojs, ndat, resid, precondresid, block_sliced)
1603 :
1604 : !$OMP TARGET TEAMS DISTRIBUTE &
1605 : !$OMP& PRIVATE(idat) MAP(to:sm1proj,precondresid)
1606 : do idat =1, ndat
1607 : !$OMP PARALLEL DO PRIVATE(iproj,icplx) COLLAPSE(2)
1608 : do iproj =1, nprojs
1609 : do icplx = 1,cplx
1610 : sm1proj(icplx, iproj, idat) = sm1proj(icplx, iproj, idat) + precondresid(icplx, iproj, idat)
1611 : end do
1612 : end do
1613 : end do
1614 : end do
1615 : !$OMP TARGET EXIT DATA MAP(delete:errs,resid,precondresid,normprojs)
1616 :
1617 : if(maxerr >= precision .and. maxerr >= 1e-10) then
1618 : write(message, *) 'In invovl, max error was', maxerr, ' after 30 iterations'
1619 : ABI_WARNING(message)
1620 : else
1621 : ! write(message,'(a,i2,a,es13.5)') 'Iterative solver in invovl finished in ', i, ' iterations, error', maxerr
1622 : ! call wrtout(std_out,message,'COLL')
1623 : end if
1624 :
1625 : end subroutine solve_inner_ompgpu
1626 : !!***
1627 :
1628 : !!****f* m_invovl/apply_block_ompgpu
1629 : !! NAME
1630 : !! apply_block_ompgpu
1631 : !!
1632 : !! FUNCTION
1633 : !! Helper function: applies a block-diagonal matrix mat(lmnmax, lmnmax, ntypat)
1634 : !!
1635 : !! INPUTS
1636 : !!
1637 : !! SOURCE
1638 : subroutine apply_block_ompgpu(ham, cplx, mat, nprojs, ndat, x, y, block_sliced)
1639 :
1640 : use m_abi_linalg
1641 :
1642 : integer,intent(in) :: ndat, nprojs, cplx
1643 : real(dp), intent(inout), target :: x(cplx, nprojs, ndat), y(cplx, nprojs, ndat)
1644 : type(gs_hamiltonian_type),intent(in) :: ham
1645 : real(dp), intent(in), target :: mat(cplx, ham%lmnmax, ham%lmnmax, ham%ntypat)
1646 : integer, intent(in) :: block_sliced
1647 :
1648 : integer :: nlmn, shift, itypat, idat
1649 : real(dp), ABI_CONTIGUOUS pointer :: x_ptr(:, :, :), y_ptr(:, :, :), mat_ptr(:,:,:)
1650 :
1651 : ! *************************************************************************
1652 :
1653 : if (block_sliced == 1) then
1654 :
1655 : do idat = 1, ndat
1656 : shift = 1
1657 : do itypat=1, ham%ntypat
1658 : nlmn = count(ham%indlmn(3,:,itypat)>0)
1659 : !! apply mat to all atoms at once
1660 : ! perform natom multiplications of size nlmn
1661 : ! compute y = mat*x
1662 : if(cplx == 2) then
1663 : !$OMP TARGET DATA USE_DEVICE_ADDR(mat,x,y)
1664 : call abi_gpu_zhemm('L','U', nlmn, ham%nattyp(itypat), cone, &
1665 : c_loc(mat(:, :, :, itypat)), ham%lmnmax, &
1666 : c_loc(x(:, shift:shift+nlmn*ham%nattyp(itypat)-1, idat)), nlmn, czero, &
1667 : c_loc(y(:, shift:shift+nlmn*ham%nattyp(itypat)-1, idat)), nlmn)
1668 : !$OMP END TARGET DATA
1669 : else
1670 : !$OMP TARGET DATA USE_DEVICE_ADDR(mat,x,y)
1671 : call abi_gpu_xsymm(cplx, 'L','U', nlmn, ham%nattyp(itypat), cone, &
1672 : c_loc(mat(:, :, :, itypat)), ham%lmnmax, &
1673 : c_loc(x(:, shift:shift+nlmn*ham%nattyp(itypat)-1, idat)), nlmn, czero, &
1674 : c_loc(y(:, shift:shift+nlmn*ham%nattyp(itypat)-1, idat)), nlmn)
1675 : !$OMP END TARGET DATA
1676 : end if
1677 : shift = shift + nlmn*ham%nattyp(itypat)
1678 : end do
1679 : end do
1680 :
1681 : else ! block_sliced = 0
1682 :
1683 : shift = 1
1684 : do itypat=1, ham%ntypat
1685 : nlmn = count(ham%indlmn(3,:,itypat)>0)
1686 : x_ptr => x(:, shift:shift+nlmn*ham%nattyp(itypat)-1, :)
1687 : y_ptr => y(:, shift:shift+nlmn*ham%nattyp(itypat)-1, :)
1688 : mat_ptr => mat(:, :, :, itypat)
1689 : !! apply mat to all atoms at once, all idat at once
1690 : ! perform natom multiplications of size nlmn
1691 : ! be careful here matrix extracted from x and y are not memory contiguous
1692 : ! ==> so in the GPU version we will need to adapt leading dimension
1693 : !$OMP TARGET DATA USE_DEVICE_ADDR(mat_ptr,x_ptr,y_ptr)
1694 : call abi_gpu_xgemm_strided(cplx, 'N','N', &
1695 : nlmn, ham%nattyp(itypat), nlmn, cone, &
1696 : c_loc(mat_ptr), ham%lmnmax, 0, &
1697 : c_loc(x_ptr), nlmn, nprojs, &
1698 : czero, &
1699 : c_loc(y_ptr), nlmn, nprojs, ndat)
1700 : !$OMP END TARGET DATA
1701 : shift = shift + nlmn*ham%nattyp(itypat)
1702 : end do
1703 :
1704 : end if
1705 :
1706 : end subroutine apply_block_ompgpu
1707 : !!***
1708 : #endif
1709 :
1710 0 : end MODULE m_invovl
1711 : !!***
|
