Line data Source code
1 : !!****m* ABINIT/m_gemm_nonlop_projectors
2 : !! NAME
3 : !! m_gemm_nonlop
4 : !!
5 : !! FUNCTION
6 : !! This module provides functions to compute the nonlocal operator by means of the BLAS GEMM
7 : !! routine. By treating ndat simultaneous wavefunctions, it is able to exploit BLAS3 routines,
8 : !! which leads to excellent CPU efficiency and OpenMP scalability.
9 : !!
10 : !! COPYRIGHT
11 : !! Copyright (C) 2014-2026 ABINIT group (AL)
12 : !! This file is distributed under the terms of the
13 : !! GNU General Public License, see ~abinit/COPYING
14 : !! or http://www.gnu.org/copyleft/gpl.txt .
15 : !!
16 : !! SOURCE
17 :
18 : ! TODO list :
19 : ! Don't allocate the full nkpt structures, only those that are treated by this proc: use same init as in m_bandfft_kpt
20 : ! support more options (forces & stresses mostly)
21 : ! Support RF/other computations (only GS right now)
22 : ! handle the case where nloalg(2) < 0, ie no precomputation of ph3d
23 : ! more systematic checking of the workflow (right now, only works if init/make/gemm/destroy, no multiple makes, etc)
24 : ! Avoid allocating the complex matrix when istwfk > 1
25 : ! Merge with chebfi's invovl
26 :
27 :
28 : #if defined HAVE_CONFIG_H
29 : #include "config.h"
30 : #endif
31 :
32 : #include "abi_common.h"
33 :
34 : module m_gemm_nonlop_projectors
35 :
36 : use defs_basis
37 : use m_errors
38 : use m_abicore
39 : use m_xomp
40 : use m_xmpi
41 : use m_fstrings, only : itoa, ftoa, sjoin
42 : use m_gputk
43 : use m_abi_linalg
44 :
45 : use defs_abitypes, only : MPI_type
46 : use m_kg, only : mkkpg
47 : use m_hamiltonian, only : KPRIME_H_K, K_H_KPRIME, K_H_K, KPRIME_H_KPRIME
48 :
49 : #if defined(HAVE_GPU)
50 : use m_gpu_toolbox
51 : #endif
52 :
53 : #if defined(HAVE_GPU_CUDA)
54 : use m_alloc_hamilt_gpu, only : gemm_nonlop_gpu_data
55 : #endif
56 :
57 : use, intrinsic :: iso_c_binding, only : c_int32_t, c_int64_t, c_float, c_double, c_size_t, c_loc, c_ptr
58 :
59 : implicit none
60 :
61 : private
62 :
63 : public :: init_gemm_nonlop
64 : public :: destroy_gemm_nonlop
65 : public :: reset_gemm_nonlop
66 : public :: set_gemm_nonlop_ikpt
67 : public :: prep_projectors
68 : public :: prep_dprojectors
69 : public :: refresh_projectors
70 :
71 : !!***
72 :
73 : !----------------------------------------------------------------------
74 :
75 : !!****t* m_gemm_nonlop_projectors/gemm_nonlop_type
76 : !! NAME
77 : !! gemm_nonlop_type
78 : !!
79 : !! FUNCTION
80 : !! Contains information needed to apply the nonlocal operator
81 : !!
82 : !! SOURCE
83 : type,public :: gemm_nonlop_type
84 :
85 : integer :: npw
86 : integer :: nprojs
87 : integer :: ngrads
88 : integer :: ngrads2
89 :
90 : integer :: nprojs_blk
91 : integer :: nprojs_last_blk
92 :
93 : real(dp), allocatable :: projs(:, :, :)
94 : ! (2, npw, nprojs)
95 :
96 : real(dp), allocatable :: projs_r(:, :, :)
97 : ! (1, npw, nprojs)
98 :
99 : real(dp), allocatable :: projs_i(:, :, :)
100 : ! (1, npw, nprojs)
101 :
102 : real(dp), allocatable :: dprojs(:, :, :)
103 : ! (2, npw, nprojs*ngrads)
104 : real(dp), allocatable :: dprojs_r(:, :, :)
105 : ! (1, npw, nprojs*ngrads)
106 : real(dp), allocatable :: dprojs_i(:, :, :)
107 : ! (1, npw, nprojs*ngrads)
108 :
109 : real(dp), allocatable :: d2projs(:, :, :)
110 : ! (2, npw, nprojs*ngrads)
111 :
112 : integer :: idir
113 : integer :: ikpt
114 : integer :: choice
115 :
116 : end type gemm_nonlop_type
117 : !!***
118 :
119 : type(gemm_nonlop_type), save, public, target :: gemm_nonlop_kpt(2)
120 : !(nkpt)
121 :
122 : integer, save, public :: gemm_nonlop_ikpt_this_proc_being_treated
123 : !! This is oh so very crude, but I can't find any other way to do it without passing ikpt deep down to nonlop
124 :
125 : logical, save, public :: gemm_nonlop_use_gemm = .false.
126 : ! Public variable indicating whether we should call gemm_nonlop or fall back to the usual nonlop. Set to false
127 : ! in order not to interfere with non-GS calls to nonlop.
128 :
129 : logical, save, public :: gemm_nonlop_is_distributed = .false.
130 : ! Public variable indicating whether we should gemm_nonlop operated in a distributed manner. Set to false by default
131 : ! but might be enabled by memory constraints or forced by user through parameters.
132 :
133 : logical, save, public :: gemm_nonlop_split_choice23 = .false.
134 : ! Public variable indicating whether choice 23 computation should be splitted. Set to false by default
135 : ! but might be enabled by memory constraints or forced by user through parameters.
136 :
137 : integer, save :: gemm_nonlop_nblocks = 1
138 : ! How many blocks of MPI tasks should the projs arrays be ditributed.
139 :
140 : integer, save, public :: gemm_nonlop_block_comm = xmpi_comm_null
141 : ! MPI communicator for MPI tasks processing the same gemm_nonlop block for projs array distribution
142 :
143 : integer, save, public :: gemm_nonlop_block_size = 0
144 : ! Public variable indicating size of a block (ie: number of MPI tasks in gemm_nonlop_block_comm)
145 : ! Default size 0 indicates no distribution at all.
146 :
147 : integer, save, public :: gemm_nonlop_choice = -1
148 :
149 : integer, save, public :: gemm_nonlop_gpu_option = ABI_GPU_DISABLED
150 :
151 : real(dp),save, allocatable, target :: atom_projs(:,:,:)
152 : real(dp),save, allocatable, target :: atom_dprojs(:,:,:,:)
153 : real(dp),save, allocatable, target :: atom_d2projs(:,:,:,:)
154 : integer,save, allocatable, target :: scal(:)
155 : integer,save, allocatable, target :: lmn_parity(:)
156 : integer, save :: mod__lmnmax, mod__npw, mod__ndprojs, mod__nd2projs
157 : ! Work arrays for prep_*projectors functions. Sized after mod__lmnmax, mod__npw and mod__ndprojs.
158 :
159 : #if defined(HAVE_FC_ISO_C_BINDING) && defined(HAVE_GPU_CUDA)
160 :
161 : type, bind(c), public :: gemm_nonlop_gpu_type
162 :
163 : integer(kind=c_int32_t) :: npw
164 : integer(kind=c_int32_t) :: nprojs
165 :
166 : ! array of double on GPU, dimensions are (2, npw, nprojs)
167 : type(c_ptr) :: projs
168 :
169 : ! array of double on GPU, dimensions are (1, npw, nprojs)
170 : type(c_ptr) :: projs_r
171 :
172 : ! array of double on GPU, dimensions are (1, npw, nprojs)
173 : type(c_ptr) :: projs_i
174 :
175 : end type gemm_nonlop_gpu_type
176 :
177 : !! array of size nkpt of sobjects of type gemm_nonlop_gpu_type, array size is nkpt
178 : type(gemm_nonlop_gpu_type), save, public, target :: gemm_nonlop_kpt_gpu(2)
179 : !(nkpt)
180 :
181 : #endif
182 :
183 : !!***
184 :
185 : !----------------------------------------------------------------------
186 :
187 : contains
188 :
189 : !----------------------------------------------------------------------
190 :
191 : !!****f* m_gemm_nonlop_projectors/init_gemm_nonlop
192 : !! NAME
193 : !! init_gemm_nonlop
194 : !!
195 : !! FUNCTION
196 : !! Initalization of the gemm_nonlop_kpt array
197 : !!
198 : !! INPUTS
199 : !! nkpt= number of k-points
200 : !!
201 : !! SOURCE
202 29 : subroutine init_gemm_nonlop(gpu_option)
203 :
204 : integer,intent(in) :: gpu_option
205 :
206 : ! *************************************************************************
207 :
208 87 : gemm_nonlop_kpt(:)%npw = -1
209 87 : gemm_nonlop_kpt(:)%nprojs = -1
210 87 : gemm_nonlop_kpt(:)%ngrads = -1
211 87 : gemm_nonlop_kpt(:)%ngrads2 = -1
212 87 : gemm_nonlop_kpt(:)%choice = -1
213 87 : gemm_nonlop_kpt(:)%idir = -1
214 87 : gemm_nonlop_kpt(:)%ikpt = -1
215 :
216 29 : if(gpu_option == ABI_GPU_LEGACY .or. gpu_option == ABI_GPU_KOKKOS) then
217 : #ifdef HAVE_GPU_CUDA
218 : gemm_nonlop_kpt_gpu(:)%npw = -1
219 : gemm_nonlop_kpt_gpu(:)%nprojs = -1
220 : gemm_nonlop_gpu_data % allocated = .false.
221 : #endif
222 : end if
223 :
224 29 : gemm_nonlop_block_comm=xmpi_comm_null
225 29 : gemm_nonlop_block_size=0
226 29 : gemm_nonlop_nblocks=1
227 29 : gemm_nonlop_gpu_option=gpu_option
228 :
229 29 : end subroutine init_gemm_nonlop
230 : !!***
231 :
232 : !----------------------------------------------------------------------
233 :
234 : !!****f* m_gemm_nonlop_projectors/destroy_gemm_nonlop
235 : !! NAME
236 : !! destroy_gemm_nonlop
237 : !!
238 : !! FUNCTION
239 : !! Destruction of the gemm_nonlop_kpt array
240 : !!
241 : !! INPUTS
242 : !! nkpt= number of k-points
243 : !!
244 : !! SOURCE
245 29 : subroutine destroy_gemm_nonlop(gpu_option)
246 :
247 : integer,intent(in) :: gpu_option
248 :
249 : ! *************************************************************************
250 :
251 29 : call free_gemm_nonlop_ikpt(1,gpu_option)
252 29 : call free_gemm_nonlop_ikpt(2,gpu_option)
253 29 : call destroy_work_arrays(gpu_option)
254 29 : if(gemm_nonlop_block_comm/=xmpi_comm_null) call xmpi_comm_free(gemm_nonlop_block_comm)
255 :
256 29 : end subroutine destroy_gemm_nonlop
257 : !!***
258 :
259 : !----------------------------------------------------------------------
260 :
261 : !!****f* m_gemm_nonlop_projectors/alloc_work_arrays
262 : !! NAME
263 : !! alloc_work_arrays
264 : !!
265 : !! FUNCTION
266 : !! Allocation of work arrays
267 : !!
268 : !! INPUTS
269 : !! gpu_option = which GPU code path is used
270 : !!
271 : !! SOURCE
272 1120 : subroutine alloc_work_arrays(lmnmax,npw,ndprojs,nd2projs,gpu_option)
273 :
274 : integer,intent(in) :: lmnmax,npw,ndprojs,nd2projs,gpu_option
275 :
276 : ! *************************************************************************
277 :
278 : !FIXME Would be nice to not allocate/reallocate at each call, but it seem troublesome in practice
279 : !if(mod__lmnmax>=lmnmax .and. mod__npw>=npw .and. mod__ndprojs>=ndprojs .and. mod__nd2projs>=nd2projs) then
280 : ! return ! Nothing to do
281 : !end if
282 :
283 1120 : call destroy_work_arrays(gpu_option)
284 :
285 4480 : ABI_MALLOC(atom_projs, (2, npw, lmnmax))
286 3360 : ABI_MALLOC(scal, (lmnmax))
287 2240 : ABI_MALLOC(lmn_parity, (lmnmax))
288 : #ifdef HAVE_OPENMP_OFFLOAD
289 : !$OMP TARGET ENTER DATA MAP(alloc:atom_projs,scal,lmn_parity) IF(gpu_option==ABI_GPU_OPENMP)
290 : #endif
291 :
292 1120 : if(ndprojs>0) then
293 1060 : ABI_MALLOC(atom_dprojs, (2, npw, ndprojs, lmnmax))
294 : #ifdef HAVE_OPENMP_OFFLOAD
295 : !$OMP TARGET ENTER DATA MAP(alloc:atom_dprojs) IF(gpu_option==ABI_GPU_OPENMP)
296 : #endif
297 : end if
298 :
299 1120 : if(nd2projs>0) then
300 0 : ABI_MALLOC(atom_d2projs, (2, npw, nd2projs, lmnmax))
301 : #ifdef HAVE_OPENMP_OFFLOAD
302 : !$OMP TARGET ENTER DATA MAP(alloc:atom_d2projs) IF(gpu_option==ABI_GPU_OPENMP)
303 : #endif
304 : end if
305 :
306 1120 : mod__lmnmax=lmnmax
307 1120 : mod__npw=npw
308 1120 : mod__ndprojs=ndprojs
309 1120 : mod__nd2projs=nd2projs
310 :
311 1120 : end subroutine alloc_work_arrays
312 : !!***
313 :
314 : !----------------------------------------------------------------------
315 :
316 : !!****f* m_gemm_nonlop_projectors/destroy_work_arrays
317 : !! NAME
318 : !! destroy_work_arrays
319 : !!
320 : !! FUNCTION
321 : !! Destruction of work arrays
322 : !!
323 : !! INPUTS
324 : !! gpu_option = which GPU code path is used
325 : !!
326 : !! SOURCE
327 1149 : subroutine destroy_work_arrays(gpu_option)
328 :
329 : integer,intent(in) :: gpu_option
330 :
331 : ! *************************************************************************
332 : ABI_UNUSED((/gpu_option/))
333 :
334 1149 : if(allocated(atom_projs)) then
335 : #ifdef HAVE_OPENMP_OFFLOAD
336 : !$OMP TARGET EXIT DATA MAP(delete:atom_projs) IF(gpu_option==ABI_GPU_OPENMP)
337 : #endif
338 1120 : ABI_FREE(atom_projs)
339 : end if
340 1149 : if(allocated(atom_dprojs)) then
341 : #ifdef HAVE_OPENMP_OFFLOAD
342 : !$OMP TARGET EXIT DATA MAP(delete:atom_dprojs) IF(gpu_option==ABI_GPU_OPENMP)
343 : #endif
344 212 : ABI_FREE(atom_dprojs)
345 : end if
346 1149 : if(allocated(atom_d2projs)) then
347 : #ifdef HAVE_OPENMP_OFFLOAD
348 : !$OMP TARGET EXIT DATA MAP(delete:atom_d2projs) IF(gpu_option==ABI_GPU_OPENMP)
349 : #endif
350 0 : ABI_FREE(atom_d2projs)
351 : end if
352 1149 : if(allocated(scal)) then
353 : #ifdef HAVE_OPENMP_OFFLOAD
354 : !$OMP TARGET EXIT DATA MAP(delete:scal) IF(gpu_option==ABI_GPU_OPENMP)
355 : #endif
356 1120 : ABI_FREE(scal)
357 : end if
358 1149 : if(allocated(lmn_parity)) then
359 : #ifdef HAVE_OPENMP_OFFLOAD
360 : !$OMP TARGET EXIT DATA MAP(delete:lmn_parity) IF(gpu_option==ABI_GPU_OPENMP)
361 : #endif
362 1120 : ABI_FREE(lmn_parity)
363 : end if
364 1149 : mod__lmnmax=0
365 1149 : mod__npw=0
366 1149 : mod__ndprojs=0
367 1149 : mod__nd2projs=0
368 :
369 1149 : end subroutine destroy_work_arrays
370 : !!***
371 :
372 : !----------------------------------------------------------------------
373 :
374 : !!****f* m_gemm_nonlop_projectors/free_gemm_nonlop_ikpt
375 : !! NAME
376 : !! free_destroy_gemm_nonlop_ikpt
377 : !!
378 : !! FUNCTION
379 : !! Release memory for one kpt value of the gemm_nonlop_kpt array
380 : !!
381 : !! INPUTS
382 : !! ikpt= index of gemm_nonlop_kptto be released
383 : !!
384 : !! SOURCE
385 1430 : subroutine free_gemm_nonlop_ikpt(ik, gpu_option)
386 :
387 : integer,intent(in) :: ik, gpu_option
388 :
389 : ! *************************************************************************
390 :
391 : if(gpu_option == ABI_GPU_LEGACY .or. gpu_option == ABI_GPU_KOKKOS) then
392 : #ifdef HAVE_GPU_CUDA
393 : if(gemm_nonlop_kpt_gpu(ik)%nprojs /= -1) then
394 : ! deallocate arrays projs, projs_r and projs_i
395 : if (allocated(gemm_nonlop_kpt(ik)%projs)) then
396 : call dealloc_on_gpu(gemm_nonlop_kpt_gpu(ik)%projs)
397 : end if
398 : if (allocated(gemm_nonlop_kpt(ik)%projs_r)) then
399 : call dealloc_on_gpu(gemm_nonlop_kpt_gpu(ik)%projs_r)
400 : end if
401 : if (allocated(gemm_nonlop_kpt(ik)%projs_i)) then
402 : call dealloc_on_gpu(gemm_nonlop_kpt_gpu(ik)%projs_i)
403 : end if
404 : gemm_nonlop_kpt_gpu(ik)%nprojs = -1
405 : gemm_nonlop_kpt_gpu(ik)%npw = -1
406 : end if
407 : #endif
408 : end if
409 :
410 : if(gpu_option == ABI_GPU_OPENMP) then
411 : #ifdef HAVE_OPENMP_OFFLOAD
412 : call free_ompgpu_current_ikpt(ik)
413 : #endif
414 : end if
415 :
416 1430 : if(gemm_nonlop_kpt(ik)%nprojs /= -1) then
417 1372 : if (allocated(gemm_nonlop_kpt(ik)%projs)) then
418 1086 : ABI_FREE(gemm_nonlop_kpt(ik)%projs)
419 : end if
420 1372 : if (allocated(gemm_nonlop_kpt(ik)%projs_r)) then
421 286 : ABI_FREE(gemm_nonlop_kpt(ik)%projs_r)
422 : end if
423 1372 : if (allocated(gemm_nonlop_kpt(ik)%projs_i)) then
424 286 : ABI_FREE(gemm_nonlop_kpt(ik)%projs_i)
425 : end if
426 1372 : gemm_nonlop_kpt(ik)%nprojs = -1
427 1372 : if(gemm_nonlop_kpt(ik)%ngrads /= -1) then
428 108 : if (allocated(gemm_nonlop_kpt(ik)%dprojs)) then
429 87 : ABI_FREE(gemm_nonlop_kpt(ik)%dprojs)
430 : end if
431 108 : if (allocated(gemm_nonlop_kpt(ik)%dprojs_r)) then
432 21 : ABI_FREE(gemm_nonlop_kpt(ik)%dprojs_r)
433 : end if
434 108 : if (allocated(gemm_nonlop_kpt(ik)%dprojs_i)) then
435 21 : ABI_FREE(gemm_nonlop_kpt(ik)%dprojs_i)
436 : end if
437 108 : gemm_nonlop_kpt(ik)%ngrads = -1
438 : end if
439 1372 : if(gemm_nonlop_kpt(ik)%ngrads2 /= -1) then
440 0 : if (allocated(gemm_nonlop_kpt(ik)%d2projs)) then
441 0 : ABI_FREE(gemm_nonlop_kpt(ik)%d2projs)
442 : end if
443 0 : gemm_nonlop_kpt(ik)%ngrads2 = -1
444 : end if
445 : end if
446 1430 : gemm_nonlop_kpt(ik)%choice = -1
447 1430 : gemm_nonlop_kpt(ik)%idir = -1
448 1430 : gemm_nonlop_kpt(ik)%ikpt = -1
449 :
450 1430 : if(gemm_nonlop_is_distributed) then
451 0 : gemm_nonlop_kpt(ik)%nprojs_blk = -1
452 0 : gemm_nonlop_kpt(ik)%nprojs_last_blk = -1
453 : end if
454 :
455 1430 : end subroutine free_gemm_nonlop_ikpt
456 : !!***
457 :
458 : !----------------------------------------------------------------------
459 :
460 : subroutine free_ompgpu_current_ikpt(ik)
461 :
462 : integer,intent(in) :: ik
463 :
464 : #ifdef HAVE_OPENMP_OFFLOAD
465 :
466 : call free_ompgpu_current_ikpt_projs(ik)
467 : call free_ompgpu_current_ikpt_dprojs(ik)
468 :
469 : #else
470 : ABI_UNUSED((/ik/))
471 : #endif
472 : end subroutine free_ompgpu_current_ikpt
473 :
474 : !----------------------------------------------------------------------
475 :
476 : subroutine free_ompgpu_current_ikpt_projs(ik)
477 :
478 : integer,intent(in) :: ik
479 : #ifdef HAVE_OPENMP_OFFLOAD
480 : !NOTE: Those pointers exists to be served to OpenMP TARGET directives to hide
481 : ! the datastructure gemm_nonlop_kpt which is not supported in GCC, LLVM and Cray
482 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_projs_ompptr(:,:,:)
483 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_projs_r_ompptr(:,:,:)
484 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_projs_i_ompptr(:,:,:)
485 :
486 :
487 : if(xomp_target_is_present(c_loc(gemm_nonlop_kpt(ik)%projs))) then
488 : gemm_nonlop_kpt_projs_ompptr => gemm_nonlop_kpt(ik)%projs
489 : !$OMP TARGET EXIT DATA MAP(delete:gemm_nonlop_kpt_projs_ompptr)
490 : end if
491 : if(xomp_target_is_present(c_loc(gemm_nonlop_kpt(ik)%projs_r))) then
492 : gemm_nonlop_kpt_projs_r_ompptr => gemm_nonlop_kpt(ik)%projs_r
493 : gemm_nonlop_kpt_projs_i_ompptr => gemm_nonlop_kpt(ik)%projs_i
494 : !$OMP TARGET EXIT DATA MAP(delete:gemm_nonlop_kpt_projs_r_ompptr)
495 : !$OMP TARGET EXIT DATA MAP(delete:gemm_nonlop_kpt_projs_i_ompptr)
496 : end if
497 :
498 : #else
499 : ABI_UNUSED((/ik/))
500 : #endif
501 : end subroutine free_ompgpu_current_ikpt_projs
502 :
503 : !----------------------------------------------------------------------
504 :
505 :
506 : subroutine free_ompgpu_current_ikpt_dprojs(ik)
507 :
508 : integer,intent(in) :: ik
509 : #ifdef HAVE_OPENMP_OFFLOAD
510 : !NOTE: Those pointers exists to be served to OpenMP TARGET directives to hide
511 : ! the datastructure gemm_nonlop_kpt which is not supported in GCC, LLVM and Cray
512 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_dprojs_ompptr(:,:,:)
513 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_dprojs_r_ompptr(:,:,:)
514 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_dprojs_i_ompptr(:,:,:)
515 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_d2projs_ompptr(:,:,:)
516 :
517 : if(xomp_target_is_present(c_loc(gemm_nonlop_kpt(ik)%dprojs))) then
518 : gemm_nonlop_kpt_dprojs_ompptr => gemm_nonlop_kpt(ik)%dprojs
519 : !$OMP TARGET EXIT DATA MAP(delete:gemm_nonlop_kpt_dprojs_ompptr)
520 : end if
521 : if(xomp_target_is_present(c_loc(gemm_nonlop_kpt(ik)%dprojs_i))) then
522 : gemm_nonlop_kpt_dprojs_r_ompptr => gemm_nonlop_kpt(ik)%dprojs_r
523 : gemm_nonlop_kpt_dprojs_i_ompptr => gemm_nonlop_kpt(ik)%dprojs_i
524 : !$OMP TARGET EXIT DATA MAP(delete:gemm_nonlop_kpt_dprojs_r_ompptr)
525 : !$OMP TARGET EXIT DATA MAP(delete:gemm_nonlop_kpt_dprojs_i_ompptr)
526 : end if
527 :
528 : if(xomp_target_is_present(c_loc(gemm_nonlop_kpt(ik)%d2projs))) then
529 : gemm_nonlop_kpt_d2projs_ompptr => gemm_nonlop_kpt(ik)%d2projs
530 : !$OMP TARGET EXIT DATA MAP(delete:gemm_nonlop_kpt_d2projs_ompptr)
531 : end if
532 :
533 : #else
534 : ABI_UNUSED((/ik/))
535 : #endif
536 : end subroutine free_ompgpu_current_ikpt_dprojs
537 :
538 : !----------------------------------------------------------------------
539 :
540 : !!****f* m_gemm_nonlop_projectors/set_gemm_nonlop_ikpt
541 : !! NAME
542 : !! set_gemm_nonlop_ikpt
543 : !!
544 : !! FUNCTION
545 : !! Set the K-point upon which projectors will be computed and
546 : !! pre-allocate projectors buffers.
547 : !!
548 : !! INPUTS
549 : !! ikpt= K-point id
550 : !! npw= number of plane-wave
551 : !! istwf_k=option parameter that describes the storage of wfs
552 : !! indlmn(6,nlmn)= array giving l,m,n,lm,ln,s for i=lmn
553 : !! ntypat=number of atoms types
554 : !! nattyp(ntypat)=number of atoms of each type
555 : !! gpu_option=which variant of GEMM nonlop is used
556 : !!
557 : !! SOURCE
558 768 : subroutine set_gemm_nonlop_ikpt(ikpt,npw,istwf_k,indlmn,ntypat,nattyp,gpu_option)
559 :
560 : integer,intent(in) :: ikpt,istwf_k,npw,ntypat,gpu_option
561 : integer,intent(in) :: indlmn(:,:,:), nattyp(ntypat)
562 :
563 : integer :: nprojs, itypat
564 :
565 : ! *************************************************************************
566 :
567 768 : gemm_nonlop_ikpt_this_proc_being_treated=ikpt
568 :
569 768 : nprojs=0
570 1880 : do itypat=1,ntypat
571 17656 : nprojs = nprojs + count(indlmn(3,:,itypat)>0)*nattyp(itypat)
572 : end do
573 : ! Call a "dummy" refresh of projectors buffers
574 : ! This is mostly a work-around in GPU workloads to ensure there
575 : ! is a buffer allocated in GPU memory.
576 768 : call refresh_projectors(npw,istwf_k,nprojs,0,0,.false.,gpu_option)
577 :
578 768 : end subroutine set_gemm_nonlop_ikpt
579 : !!***
580 :
581 : !!****f* m_gemm_nonlop_projectors/reset_gemm_nonlop
582 : !! NAME
583 : !! reset_gemm_nonlop
584 : !!
585 : !! FUNCTION
586 : !! Reset projectors to trigger their recomputation
587 : !!
588 : !! INPUTS
589 : !!
590 : !! SOURCE
591 137 : subroutine reset_gemm_nonlop()
592 :
593 : ! *************************************************************************
594 :
595 411 : gemm_nonlop_kpt(:)%ikpt = -1
596 411 : gemm_nonlop_kpt(:)%choice = -1
597 411 : gemm_nonlop_kpt(:)%idir = -1
598 :
599 137 : end subroutine reset_gemm_nonlop
600 : !!***
601 :
602 : !----------------------------------------------------------------------
603 :
604 : !!****f* m_gemm_nonlop_projectors/refresh_projectors
605 : !! NAME
606 : !! prep_projectors
607 : !!
608 : !! FUNCTION
609 : !! Check allocation of projectors arrays for GEMM nonlop, and resize them if need be
610 : !!
611 : !! INPUTS
612 : !!
613 : !! SOURCE
614 24960 : subroutine refresh_projectors(npw,istwf_k,nprojs,ndgxdt,nd2gxdt,&
615 : & is_kprime,gpu_option)
616 : integer,intent(in) :: npw,istwf_k,nprojs,ndgxdt,nd2gxdt,gpu_option
617 : logical,intent(in) :: is_kprime
618 : integer :: ik,rank,nprojs_blk,nprojs_my_blk,nprojs_last_blk,ierr,nprocs
619 : logical :: is_last_rank
620 : #ifdef HAVE_OPENMP_OFFLOAD
621 : !NOTE: Those pointers exists to be served to OpenMP TARGET directives to hide
622 : ! the datastructure gemm_nonlop_kpt which is not supported in GCC, LLVM and Cray
623 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_projs_ompptr(:,:,:)
624 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_projs_r_ompptr(:,:,:)
625 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_projs_i_ompptr(:,:,:)
626 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_dprojs_ompptr(:,:,:)
627 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_dprojs_r_ompptr(:,:,:)
628 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_dprojs_i_ompptr(:,:,:)
629 : real(dp), ABI_CONTIGUOUS pointer :: gemm_nonlop_kpt_d2projs_ompptr(:,:,:)
630 : #endif
631 :
632 24960 : ik=1; if(is_kprime) ik=2
633 : if(gemm_nonlop_kpt(ik)%ikpt/=gemm_nonlop_ikpt_this_proc_being_treated &
634 24960 : & .or. npw/=gemm_nonlop_kpt(ik)%npw .or. nprojs/=gemm_nonlop_kpt(ik)%nprojs) then
635 1372 : call free_gemm_nonlop_ikpt(ik, gpu_option)
636 : end if
637 :
638 24960 : if(gemm_nonlop_is_distributed) then
639 0 : nprocs = xmpi_comm_size(xmpi_world)
640 : ! If split size has changed, reset array and init MPI communicator
641 0 : if(gemm_nonlop_block_comm==xmpi_comm_null .or. gemm_nonlop_nblocks /= nprocs/gemm_nonlop_block_size) then
642 0 : call free_gemm_nonlop_ikpt(ik, gpu_option)
643 0 : if(gemm_nonlop_block_comm/=xmpi_comm_null) call xmpi_comm_free(gemm_nonlop_block_comm)
644 0 : rank = xmpi_comm_rank(xmpi_world);
645 0 : gemm_nonlop_nblocks=nprocs/gemm_nonlop_block_size
646 0 : write(std_out,'(A,I3,A,I3,A)') "Splitting GEMM nonlop projectors on ",&
647 0 : & gemm_nonlop_nblocks, " blocks of ", gemm_nonlop_block_size, " MPI tasks..."
648 0 : call xmpi_comm_split(xmpi_world, rank/gemm_nonlop_block_size, rank, gemm_nonlop_block_comm, ierr)
649 0 : if(ierr/=0) ABI_BUG("MPI_comm_split failed!")
650 : end if
651 : end if
652 :
653 24960 : nprojs_last_blk = nprojs
654 24960 : nprojs_my_blk = nprojs
655 24960 : nprojs_blk = nprojs
656 24960 : rank = 0; is_last_rank = .true.
657 :
658 24960 : if(gemm_nonlop_block_size > 1) then
659 0 : nprojs_blk = nprojs / gemm_nonlop_block_size
660 0 : nprojs_last_blk = nprojs_blk + modulo(nprojs,nprojs_blk)
661 :
662 0 : if(gemm_nonlop_is_distributed) then
663 0 : rank = xmpi_comm_rank(gemm_nonlop_block_comm);
664 0 : is_last_rank = (rank==gemm_nonlop_block_size-1)
665 : if(is_last_rank) then
666 24960 : nprojs_my_blk = nprojs_last_blk
667 : else
668 0 : nprojs_my_blk = nprojs_blk
669 : end if
670 : end if
671 : end if
672 :
673 :
674 : ! Allocation of buffers for 1st and 2nd order derivatives of projectors
675 : ! NOTE: those are allocated, if needed, before regular projectors buffers
676 : ! for optimization purposes, regarding GPU memory pool.
677 24960 : if(nprojs>0) then
678 24960 : if(ndgxdt>0) then
679 : if(npw/=gemm_nonlop_kpt(ik)%npw .or. nprojs/=gemm_nonlop_kpt(ik)%nprojs &
680 424 : & .or. ndgxdt /= gemm_nonlop_kpt(ik)%ngrads .or. nd2gxdt /= gemm_nonlop_kpt(ik)%ngrads2) then
681 : if(gpu_option == ABI_GPU_OPENMP) call free_ompgpu_current_ikpt_dprojs(ik)
682 424 : ABI_SFREE(gemm_nonlop_kpt(ik)%dprojs)
683 424 : ABI_SFREE(gemm_nonlop_kpt(ik)%dprojs_r)
684 424 : ABI_SFREE(gemm_nonlop_kpt(ik)%dprojs_i)
685 424 : ABI_SFREE(gemm_nonlop_kpt(ik)%d2projs)
686 424 : gemm_nonlop_kpt(ik)%ngrads = -1
687 424 : gemm_nonlop_kpt(ik)%ngrads2 = -1
688 :
689 424 : if(istwf_k <= 1) then
690 1216 : ABI_MALLOC(gemm_nonlop_kpt(ik)%dprojs, (2, npw, nprojs_last_blk*ndgxdt))
691 : #ifdef HAVE_OPENMP_OFFLOAD
692 : gemm_nonlop_kpt_dprojs_ompptr => gemm_nonlop_kpt(ik)%dprojs
693 : !$OMP TARGET ENTER DATA MAP(alloc:gemm_nonlop_kpt_dprojs_ompptr) IF(gpu_option==ABI_GPU_OPENMP)
694 : #endif
695 304 : if(nd2gxdt>0) then
696 0 : ABI_MALLOC(gemm_nonlop_kpt(ik)%d2projs, (2, npw, nprojs_last_blk*nd2gxdt))
697 : #ifdef HAVE_OPENMP_OFFLOAD
698 : gemm_nonlop_kpt_d2projs_ompptr => gemm_nonlop_kpt(ik)%d2projs
699 : !$OMP TARGET ENTER DATA MAP(alloc:gemm_nonlop_kpt_d2projs_ompptr) IF(gpu_option==ABI_GPU_OPENMP)
700 : #endif
701 : end if
702 : else
703 480 : ABI_MALLOC(gemm_nonlop_kpt(ik)%dprojs_r, (1, npw, nprojs_last_blk*ndgxdt))
704 360 : ABI_MALLOC(gemm_nonlop_kpt(ik)%dprojs_i, (1, npw, nprojs_last_blk*ndgxdt))
705 : #ifdef HAVE_OPENMP_OFFLOAD
706 : gemm_nonlop_kpt_dprojs_r_ompptr => gemm_nonlop_kpt(ik)%dprojs_r
707 : gemm_nonlop_kpt_dprojs_i_ompptr => gemm_nonlop_kpt(ik)%dprojs_i
708 : !$OMP TARGET ENTER DATA MAP(alloc:gemm_nonlop_kpt_dprojs_r_ompptr) IF(gpu_option==ABI_GPU_OPENMP)
709 : !$OMP TARGET ENTER DATA MAP(alloc:gemm_nonlop_kpt_dprojs_i_ompptr) IF(gpu_option==ABI_GPU_OPENMP)
710 : #endif
711 : end if
712 : end if
713 : end if
714 :
715 :
716 : ! Allocation of projectors buffers
717 24960 : if(npw/=gemm_nonlop_kpt(ik)%npw .or. nprojs/=gemm_nonlop_kpt(ik)%nprojs) then
718 1372 : if(istwf_k <= 1) then
719 4344 : ABI_MALLOC(gemm_nonlop_kpt(ik)%projs, (2, npw, nprojs_last_blk))
720 : #ifdef HAVE_OPENMP_OFFLOAD
721 : gemm_nonlop_kpt_projs_ompptr => gemm_nonlop_kpt(ik)%projs
722 : !$OMP TARGET ENTER DATA MAP(alloc:gemm_nonlop_kpt_projs_ompptr) IF(gpu_option==ABI_GPU_OPENMP)
723 : #endif
724 : else
725 1144 : ABI_MALLOC(gemm_nonlop_kpt(ik)%projs_r, (1, npw, nprojs_last_blk))
726 858 : ABI_MALLOC(gemm_nonlop_kpt(ik)%projs_i, (1, npw, nprojs_last_blk))
727 : #ifdef HAVE_OPENMP_OFFLOAD
728 : gemm_nonlop_kpt_projs_r_ompptr => gemm_nonlop_kpt(ik)%projs_r
729 : gemm_nonlop_kpt_projs_i_ompptr => gemm_nonlop_kpt(ik)%projs_i
730 : !$OMP TARGET ENTER DATA MAP(alloc:gemm_nonlop_kpt_projs_r_ompptr) IF(gpu_option==ABI_GPU_OPENMP)
731 : !$OMP TARGET ENTER DATA MAP(alloc:gemm_nonlop_kpt_projs_i_ompptr) IF(gpu_option==ABI_GPU_OPENMP)
732 : #endif
733 : end if
734 : end if
735 : end if
736 :
737 :
738 24960 : if (nprojs>0) gemm_nonlop_kpt(ik)%nprojs = nprojs
739 24960 : if (nprojs>0) gemm_nonlop_kpt(ik)%npw = npw
740 24960 : if (ndgxdt>0) gemm_nonlop_kpt(ik)%ngrads = ndgxdt
741 24960 : if (nd2gxdt>0) gemm_nonlop_kpt(ik)%ngrads2 = nd2gxdt
742 24960 : if(gemm_nonlop_block_size > 1) then
743 0 : if(nprojs_blk>0) gemm_nonlop_kpt(ik)%nprojs_blk = nprojs_blk
744 0 : if(nprojs_last_blk>0) gemm_nonlop_kpt(ik)%nprojs_last_blk = nprojs_last_blk
745 : end if
746 :
747 : !!!!! CUDA stuff
748 : if(gpu_option == ABI_GPU_LEGACY .or. gpu_option == ABI_GPU_KOKKOS) then
749 : #ifdef HAVE_GPU_CUDA
750 : if(gemm_nonlop_kpt_gpu(ik)%nprojs==-1) then
751 : gemm_nonlop_kpt_gpu(ik)%npw = npw
752 : gemm_nonlop_kpt_gpu(ik)%nprojs = nprojs
753 :
754 : #ifdef DEBUG_VERBOSE_GPU
755 : if(xmpi_comm_rank(xmpi_world) == 0) then
756 : call check_gpu_mem("refresh_projectors begin")
757 : call wrtout(std_out,sjoin(" npw .......", itoa(npw)), 'COLL')
758 : call wrtout(std_out,sjoin(" nprojs .......", itoa(nprojs)), 'COLL')
759 : end if
760 : #endif
761 :
762 : if(istwf_k <= 1) then
763 : call alloc_on_gpu(gemm_nonlop_kpt_gpu(ik)%projs, INT(2,c_size_t)*npw*nprojs*dp)
764 : ! TODO : gradients
765 : else
766 : call alloc_on_gpu(gemm_nonlop_kpt_gpu(ik)%projs_r, INT(1, c_size_t)*npw*nprojs*dp)
767 : call alloc_on_gpu(gemm_nonlop_kpt_gpu(ik)%projs_i, INT(1, c_size_t)*npw*nprojs*dp)
768 : ! TODO : gradients
769 : end if
770 :
771 : #ifdef DEBUG_VERBOSE_GPU
772 : if(xmpi_comm_rank(xmpi_world) == 0) then
773 : call check_gpu_mem("refresh_projectors end ")
774 : end if
775 : #endif
776 : end if
777 :
778 : #endif
779 : end if
780 24960 : end subroutine refresh_projectors
781 : !!***
782 :
783 : !----------------------------------------------------------------------
784 :
785 : !!****f* m_gemm_nonlop_projectors/prep_projectors
786 : !! NAME
787 : !! prep_projectors
788 : !!
789 : !! FUNCTION
790 : !! Prepare projectors array for GEMM nonlop (choice=={0,1})
791 : !!
792 : !! INPUTS
793 : !!
794 : !! SOURCE
795 696 : subroutine prep_projectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
796 696 : & ucvol,ffnl,ph3d,dimffnl,matblk,&
797 : & nprojs,is_kprime,gpu_option,&
798 : & iblock)
799 :
800 : integer, intent(in) :: npw,lmnmax,ntypat,dimffnl,matblk
801 : integer, intent(in) :: istwf_k
802 : integer, intent(in) :: nprojs,gpu_option,iblock
803 : logical, intent(in) :: is_kprime
804 : real(dp), intent(in) :: ucvol
805 : ! arrays
806 : integer, intent(in) :: indlmn(6,lmnmax,ntypat),nattyp(ntypat)
807 : real(dp),intent(in),target :: ffnl(npw,dimffnl,lmnmax,ntypat)
808 : real(dp),intent(in),target :: ph3d(2,npw,matblk)
809 :
810 : logical :: map_ffnl,map_ph3d,is_last_rank
811 : integer :: il, ipw, ik, ilmn_p, nlmn_p
812 : integer :: itypat, ilmn, nlmn, ia, iaph3d, shift, nprojs_my_blk
813 : integer :: lmn_beg,ibeg,iend,shift_do,nlmn_o
814 : real(dp):: wt,tmp
815 696 : real(dp), ABI_CONTIGUOUS pointer :: projs(:,:,:)
816 696 : real(dp), ABI_CONTIGUOUS pointer :: projs_r(:,:,:)
817 696 : real(dp), ABI_CONTIGUOUS pointer :: projs_i(:,:,:)
818 :
819 0 : ik=1; if(is_kprime) ik=2
820 696 : if(istwf_k <= 1) then
821 551 : projs => gemm_nonlop_kpt(ik)%projs
822 : else
823 145 : projs_r => gemm_nonlop_kpt(ik)%projs_r
824 145 : projs_i => gemm_nonlop_kpt(ik)%projs_i
825 : end if
826 :
827 696 : if(gpu_option==ABI_GPU_OPENMP) then
828 :
829 0 : if(istwf_k <= 1) then
830 0 : call gpu_set_to_zero(projs,int(2,c_size_t)*npw*nprojs)
831 : else
832 0 : call gpu_set_to_zero(projs_r,int(npw,c_size_t)*nprojs)
833 0 : call gpu_set_to_zero(projs_i,int(npw,c_size_t)*nprojs)
834 : end if
835 :
836 : else
837 :
838 696 : if(istwf_k <= 1) then
839 330194679 : projs(:,:,:) = zero
840 : else
841 36945149 : projs_r(:,:,:) = zero
842 36945149 : projs_i(:,:,:) = zero
843 : end if
844 :
845 : end if
846 :
847 696 : iaph3d = 1
848 696 : wt=four_pi/sqrt(ucvol)
849 :
850 : ! Allocate atom_projs and other work arrays if need be
851 696 : call alloc_work_arrays(lmnmax,npw,-1,-1,gpu_option)
852 :
853 696 : map_ph3d=.false.; map_ffnl=.false.
854 : #ifdef HAVE_OPENMP_OFFLOAD
855 : if(.not. xomp_target_is_present(c_loc(ffnl))) map_ffnl = .true.
856 : if(.not. xomp_target_is_present(c_loc(ph3d))) map_ph3d = .true.
857 : !$OMP TARGET ENTER DATA MAP(to:ffnl) IF(gpu_option==ABI_GPU_OPENMP .and. map_ffnl)
858 : !$OMP TARGET ENTER DATA MAP(to:ph3d) IF(gpu_option==ABI_GPU_OPENMP .and. map_ph3d)
859 : #endif
860 :
861 696 : shift = 0
862 696 : lmn_beg = 1
863 :
864 696 : if(gemm_nonlop_block_size > 1) then
865 0 : is_last_rank = (iblock==gemm_nonlop_block_size)
866 0 : nprojs_my_blk=gemm_nonlop_kpt(ik)%nprojs_blk
867 0 : if(is_last_rank) nprojs_my_blk=gemm_nonlop_kpt(ik)%nprojs_last_blk
868 0 : ibeg = (iblock-1)*gemm_nonlop_kpt(ik)%nprojs_blk+1
869 0 : iend = ibeg+nprojs_my_blk
870 0 : shift_do = 0
871 : end if
872 :
873 1736 : do itypat = 1, ntypat
874 16240 : nlmn = count(indlmn(3,:,itypat)>0)
875 1040 : nlmn_o = nlmn
876 :
877 4224 : do ia = 1, nattyp(itypat)
878 :
879 : ! In distributed mode, loops are skipped until reach the section
880 : ! of "ilmn" to be stored by local rank
881 3184 : if(gemm_nonlop_block_size > 1) then
882 0 : if(shift_do+nlmn < ibeg) then
883 0 : shift_do = shift_do + nlmn
884 0 : iaph3d = iaph3d + 1
885 0 : cycle
886 : end if
887 :
888 0 : lmn_beg = max(1,ibeg-shift_do)
889 0 : if(shift_do+nlmn > iend - 1) nlmn = iend - shift_do - 1
890 : end if
891 :
892 : !! build atom_projs, from opernlb
893 : !! P = 4pi/sqrt(ucvol)* conj(diag(ph3d)) * ffnl * diag(parity), with parity = (-i)^l
894 :
895 : ! start from 4pi/sqrt(ucvol)*ffnl
896 : ! atom_projs(1, :, lmn_beg:nlmn) = four_pi/sqrt(ham%ucvol) * ham%ffnl(:, 1, lmn_beg:nlmn)
897 : ! TODO vectorize (DCOPY with stride)
898 3184 : if(gpu_option==ABI_GPU_OPENMP) then
899 0 : call gpu_set_to_zero(atom_projs,int(2,c_size_t)*npw*lmnmax)
900 : else
901 470590796 : atom_projs(:,:,:) = zero
902 : end if
903 : #ifdef HAVE_OPENMP_OFFLOAD
904 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) &
905 : !$OMP& PRIVATE(ilmn,ipw) MAP(to:atom_projs,ffnl) &
906 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
907 : #endif
908 45776 : do ilmn=1,nlmn_o
909 128567056 : do ipw=1, npw
910 128563872 : atom_projs(1,ipw, ilmn) = wt * ffnl(ipw, 1, ilmn, itypat)
911 : end do
912 : end do
913 :
914 : nlmn_p=0
915 : ! multiply by (-i)^l
916 45776 : do ilmn=1,nlmn_o
917 42592 : il=mod(indlmn(1,ilmn, itypat),4);
918 42592 : if(.not. (mod(il,2)==0)) then
919 19104 : nlmn_p=nlmn_p+1
920 19104 : lmn_parity(nlmn_p)=ilmn
921 : end if
922 45776 : scal(ilmn)=1; if(il>1) scal(ilmn)=-1
923 : end do
924 : ! multiply by -1
925 3184 : if(gpu_option==ABI_GPU_OPENMP) then
926 : #ifdef HAVE_OPENMP_OFFLOAD
927 : !$OMP TARGET UPDATE TO(scal,lmn_parity)
928 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) &
929 : !$OMP& PRIVATE(ipw,ilmn) MAP(to:atom_projs,scal) &
930 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
931 : do ilmn=1,nlmn_o
932 : do ipw=1,npw
933 : atom_projs(1,ipw,ilmn) = atom_projs(1,ipw,ilmn) * scal(ilmn)
934 : atom_projs(2,ipw,ilmn) = atom_projs(2,ipw,ilmn) * scal(ilmn)
935 : end do
936 : end do
937 : #endif
938 : else
939 45776 : do ilmn=1,nlmn_o
940 385609616 : atom_projs(:,:,ilmn) = atom_projs(:,:,ilmn) * scal(ilmn)
941 : end do
942 : end if
943 : ! multiply by -i
944 : #ifdef HAVE_OPENMP_OFFLOAD
945 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) &
946 : !$OMP& MAP(to:atom_projs,lmn_parity) PRIVATE(ilmn,ilmn_p,ipw,tmp) &
947 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
948 : #endif
949 22288 : do ilmn_p=1,nlmn_p
950 53289268 : do ipw=1,npw
951 53266980 : ilmn=lmn_parity(ilmn_p)
952 53266980 : tmp = atom_projs(2,ipw,ilmn)
953 53266980 : atom_projs(2,ipw,ilmn) = -atom_projs(1,ipw,ilmn)
954 53286084 : atom_projs(1,ipw,ilmn) = tmp
955 : end do
956 : end do
957 :
958 : ! multiply by conj(ph3d)
959 : #ifdef HAVE_OPENMP_OFFLOAD
960 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) &
961 : !$OMP& PRIVATE(ilmn,ipw,tmp) MAP(to:atom_projs,ph3d) &
962 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
963 : #endif
964 45776 : do ilmn=1,nlmn_o
965 128567056 : do ipw=1,npw
966 128521280 : tmp = atom_projs(1, ipw, ilmn)
967 : atom_projs(1, ipw, ilmn) = atom_projs(1, ipw, ilmn) * ph3d(1, ipw, iaph3d) &
968 128521280 : & + atom_projs(2, ipw, ilmn) * ph3d(2, ipw, iaph3d)
969 : atom_projs(2, ipw, ilmn) = atom_projs(2, ipw, ilmn) * ph3d(1, ipw, iaph3d) &
970 128563872 : & - tmp * ph3d(2, ipw, iaph3d)
971 : end do
972 : end do
973 :
974 : !! atom_projs is built, copy to projs
975 :
976 3184 : if(gpu_option==ABI_GPU_OPENMP) then
977 : #ifdef HAVE_OPENMP_OFFLOAD
978 : if(istwf_k <= 1) then
979 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) &
980 : !$OMP& PRIVATE(ipw,ilmn) MAP(to:projs,atom_projs)
981 : do ilmn=1,nlmn-(lmn_beg-1)
982 : do ipw=1,npw
983 : projs(1, ipw, shift+ilmn) = atom_projs(1, ipw, ilmn+(lmn_beg-1))
984 : projs(2, ipw, shift+ilmn) = atom_projs(2, ipw, ilmn+(lmn_beg-1))
985 : end do
986 : end do
987 : else ! istwf_k>1
988 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) &
989 : !$OMP& PRIVATE(ipw,ilmn) MAP(to:projs_r,projs_i,atom_projs)
990 : do ilmn=1,nlmn-(lmn_beg-1)
991 : do ipw=1,npw
992 : projs_r(1, ipw, shift+ilmn) = atom_projs(1, ipw, ilmn+(lmn_beg-1))
993 : projs_i(1, ipw, shift+ilmn) = atom_projs(2, ipw, ilmn+(lmn_beg-1))
994 : end do
995 : end do
996 : end if
997 : #endif
998 : else
999 3184 : if(istwf_k <= 1) then
1000 660390702 : projs(1:2, :, shift+1:shift+(nlmn-(lmn_beg-1))) = atom_projs(:, :, lmn_beg:nlmn)
1001 : else ! istwf_k>1
1002 36956146 : projs_r(1, :, shift+1:shift+(nlmn-(lmn_beg-1))) = atom_projs(1, :, lmn_beg:nlmn)
1003 36956146 : projs_i(1, :, shift+1:shift+(nlmn-(lmn_beg-1))) = atom_projs(2, :, lmn_beg:nlmn)
1004 : end if
1005 : end if
1006 :
1007 3184 : iaph3d = iaph3d + 1
1008 4224 : if(gemm_nonlop_block_size > 1) then
1009 0 : shift = shift + nlmn - (lmn_beg-1)
1010 0 : shift_do = shift_do + nlmn
1011 0 : if(shift_do >= iend - 1) exit
1012 : else
1013 3184 : shift = shift + nlmn
1014 : end if
1015 : end do
1016 1736 : if(gemm_nonlop_block_size > 1 .and. shift_do >= iend - 1) exit
1017 : end do
1018 :
1019 :
1020 : !!!!! CUDA stuff
1021 : if(gpu_option == ABI_GPU_LEGACY .or. gpu_option == ABI_GPU_KOKKOS) then
1022 : #ifdef HAVE_GPU_CUDA
1023 : ! upload data to gpu memory
1024 : if(istwf_k <= 1) then
1025 : call copy_on_gpu(gemm_nonlop_kpt(ik)%projs, gemm_nonlop_kpt_gpu(ik)%projs, INT(2, c_size_t)*npw*nprojs*dp)
1026 : ! TODO : gradients
1027 : else
1028 : call copy_on_gpu(gemm_nonlop_kpt(ik)%projs_r, gemm_nonlop_kpt_gpu(ik)%projs_r, &
1029 : & INT(1, c_size_t)*npw*nprojs*dp)
1030 : call copy_on_gpu(gemm_nonlop_kpt(ik)%projs_i, gemm_nonlop_kpt_gpu(ik)%projs_i, &
1031 : & INT(1, c_size_t)*npw*nprojs*dp)
1032 : end if
1033 : #endif
1034 : end if
1035 :
1036 : #ifdef HAVE_OPENMP_OFFLOAD
1037 : !$OMP TARGET EXIT DATA MAP(delete:ffnl) IF(gpu_option==ABI_GPU_OPENMP .and. map_ffnl)
1038 : !$OMP TARGET EXIT DATA MAP(delete:ph3d) IF(gpu_option==ABI_GPU_OPENMP .and. map_ph3d)
1039 : #endif
1040 :
1041 696 : end subroutine prep_projectors
1042 : !!***
1043 :
1044 : !----------------------------------------------------------------------
1045 :
1046 : !!****f* m_gemm_nonlop_projectors/prep_dprojectors
1047 : !! NAME
1048 : !! prep_projectors
1049 : !!
1050 : !! FUNCTION
1051 : !! Prepare projectors' derivatives for given choice (choice={2,3,4,5,51,54,55})
1052 : !!
1053 : !! INPUTS
1054 : !!
1055 : !! SOURCE
1056 424 : subroutine prep_dprojectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
1057 424 : & ucvol,ffnl,ph3d,kpg,nkpg,dimffnl,matblk,&
1058 : & nprojs,ngrads,ngrads2,choice,signs,idir_pert,&
1059 : & is_kprime,gpu_option,iblock)
1060 :
1061 : integer, intent(in) :: npw,lmnmax,ntypat,nkpg,dimffnl,matblk
1062 : integer, intent(in) :: istwf_k,iblock
1063 : integer, intent(in) :: nprojs,ngrads,ngrads2,choice,signs,idir_pert,gpu_option
1064 : logical, intent(in) :: is_kprime
1065 : real(dp), intent(in) :: ucvol
1066 : ! arrays
1067 : integer, intent(in) :: indlmn(6,lmnmax,ntypat),nattyp(ntypat)
1068 : real(dp), intent(in), target :: ffnl(npw,dimffnl,lmnmax,ntypat)
1069 : real(dp), intent(in), target :: ph3d(2,npw,matblk)
1070 : real(dp), intent(in), target :: kpg(npw,nkpg)
1071 :
1072 : logical :: map_ffnl,map_ph3d,is_last_rank
1073 : integer,parameter :: alpha(6)=(/1,2,3,3,3,2/),beta(6)=(/1,2,3,2,1,1/)
1074 : integer,parameter :: gamma(3,3)=reshape((/1,6,5,6,2,4,5,4,3/),(/3,3/))
1075 : integer :: ndprojs, nd2projs, ilmn_p, nlmn_p
1076 : integer :: il, ipw, ik, idir, idir1, idir2, jdir1, jdir2, kdir1, kdir2, ldir1, ldir2, ldir3, ldir4, ffnl_dir
1077 : integer :: itypat, ilmn, nlmn, ia, iaph3d, igrad, shift, shift_grad, shift_grad2
1078 : integer :: lmn_beg,ibeg,iend,shift_do,nlmn_o,lmn_grad_beg,nprojs_my_blk
1079 : real(dp), parameter :: two_pi2=two_pi*two_pi
1080 : real(dp):: wt,tmp
1081 424 : real(dp), ABI_CONTIGUOUS pointer :: projs (:,:,:) ,dprojs(:,:,:) ,d2projs(:,:,:)
1082 424 : real(dp), ABI_CONTIGUOUS pointer :: projs_r(:,:,:),dprojs_r(:,:,:)
1083 424 : real(dp), ABI_CONTIGUOUS pointer :: projs_i(:,:,:),dprojs_i(:,:,:)
1084 :
1085 :
1086 0 : ik=1; if(is_kprime) ik=2
1087 424 : if(istwf_k <= 1) then
1088 304 : projs => gemm_nonlop_kpt(ik)%projs
1089 304 : dprojs => gemm_nonlop_kpt(ik)%dprojs
1090 304 : if(ngrads2 > 0) then
1091 0 : d2projs => gemm_nonlop_kpt(ik)%d2projs
1092 : end if
1093 : else
1094 120 : projs_r => gemm_nonlop_kpt(ik)%projs_r
1095 120 : projs_i => gemm_nonlop_kpt(ik)%projs_i
1096 120 : dprojs_r => gemm_nonlop_kpt(ik)%dprojs_r
1097 120 : dprojs_i => gemm_nonlop_kpt(ik)%dprojs_i
1098 : end if
1099 :
1100 424 : if(gpu_option==ABI_GPU_OPENMP) then
1101 :
1102 0 : if(istwf_k <= 1) then
1103 0 : call gpu_set_to_zero(dprojs,int(2,c_size_t)*npw*nprojs*ngrads)
1104 0 : if(ngrads2 > 0) then
1105 0 : call gpu_set_to_zero(d2projs,int(2,c_size_t)*npw*nprojs*ngrads2)
1106 : end if
1107 : else
1108 0 : call gpu_set_to_zero(dprojs_r,int(npw,c_size_t)*nprojs*ngrads)
1109 0 : call gpu_set_to_zero(dprojs_i,int(npw,c_size_t)*nprojs*ngrads)
1110 : end if
1111 :
1112 : else
1113 :
1114 424 : if(istwf_k <= 1) then
1115 395458864 : dprojs(:,:,:) = zero
1116 304 : if(ngrads2 > 0) then
1117 0 : d2projs(:,:,:) = zero
1118 : end if
1119 : else
1120 47384544 : dprojs_r(:,:,:) = zero
1121 47384544 : dprojs_i(:,:,:) = zero
1122 : end if
1123 :
1124 : end if
1125 :
1126 424 : iaph3d = 1
1127 424 : wt=four_pi/sqrt(ucvol)
1128 424 : ffnl_dir=1; if(dimffnl>2) ffnl_dir=idir_pert
1129 :
1130 424 : ndprojs = 0
1131 424 : if (signs==1 .and. (choice==3 .or. choice==23 .or. choice==54 .or. choice==55 .or. choice==6)) then
1132 212 : ndprojs = 3
1133 212 : else if(signs==2 .and. (choice==5 .or. choice==51 .or. choice==3)) then
1134 0 : ndprojs = 1
1135 : end if
1136 424 : nd2projs = 0
1137 424 : if(signs==1 .and. choice==54) then
1138 0 : nd2projs = 3
1139 424 : else if(signs==1 .and. choice==55) then
1140 0 : nd2projs = 6
1141 424 : else if(signs==1 .and. choice==6) then
1142 0 : nd2projs = 10
1143 : end if
1144 :
1145 : ! Allocate atom_projs and other work arrays if need be
1146 424 : call alloc_work_arrays(lmnmax,npw,ndprojs,nd2projs,gpu_option)
1147 :
1148 424 : map_ph3d=.false.; map_ffnl=.false.
1149 424 : if(.not. xomp_target_is_present(c_loc(ffnl))) map_ffnl = .true.
1150 424 : if(.not. xomp_target_is_present(c_loc(ph3d))) map_ph3d = .true.
1151 : #ifdef HAVE_OPENMP_OFFLOAD
1152 : !$OMP TARGET ENTER DATA MAP(to:ffnl) IF(gpu_option==ABI_GPU_OPENMP .and. map_ffnl)
1153 : !$OMP TARGET ENTER DATA MAP(to:ph3d) IF(gpu_option==ABI_GPU_OPENMP .and. map_ph3d)
1154 : !!$OMP TARGET ENTER DATA MAP(to:kpg) IF(gpu_option==ABI_GPU_OPENMP)
1155 : #endif
1156 :
1157 424 : shift = 0 ; shift_grad = 0; shift_grad2 = 0
1158 424 : lmn_beg = 1
1159 :
1160 424 : if(gemm_nonlop_block_size > 1) then
1161 0 : is_last_rank = (iblock==gemm_nonlop_block_size)
1162 0 : nprojs_my_blk=gemm_nonlop_kpt(ik)%nprojs_blk
1163 0 : if(is_last_rank) nprojs_my_blk=gemm_nonlop_kpt(ik)%nprojs_last_blk
1164 0 : ibeg = (iblock-1)*gemm_nonlop_kpt(ik)%nprojs_blk+1
1165 0 : iend = ibeg+nprojs_my_blk
1166 0 : shift_do = 0
1167 0 : lmn_grad_beg = -1
1168 :
1169 : end if
1170 :
1171 1136 : do itypat = 1, ntypat
1172 12168 : nlmn = count(indlmn(3,:,itypat)>0)
1173 712 : nlmn_o = nlmn
1174 :
1175 2008 : do ia = 1, nattyp(itypat)
1176 :
1177 : ! In distributed mode, loops are skipped until reach the section
1178 : ! of "ilmn" to be stored by local rank
1179 1296 : if(gemm_nonlop_block_size > 1) then
1180 0 : if(shift_do+nlmn < ibeg) then
1181 0 : shift_do = shift_do + nlmn
1182 0 : iaph3d = iaph3d + 1
1183 0 : cycle
1184 : end if
1185 :
1186 0 : lmn_beg = max(1,ibeg-shift_do)
1187 : if(lmn_grad_beg==-1) lmn_grad_beg = (lmn_beg-1)*ngrads
1188 0 : if(shift_do+nlmn > iend - 1) nlmn = iend - shift_do - 1
1189 : end if
1190 :
1191 : !! build atom_dprojs, from opernlb
1192 : !! P = 4pi/sqrt(ucvol)* conj(diag(ph3d)) * ffnl * diag(parity), with parity = (-i)^l
1193 :
1194 : ! start from 4pi/sqrt(ucvol)*ffnl
1195 1296 : if(gpu_option==ABI_GPU_OPENMP) then
1196 0 : if (ndprojs>0) then
1197 0 : call gpu_set_to_zero(atom_dprojs,int(2,c_size_t)*npw*ndprojs*lmnmax)
1198 0 : if(ngrads2>0) then
1199 0 : call gpu_set_to_zero(atom_d2projs,int(2,c_size_t)*npw*nd2projs*lmnmax)
1200 : end if
1201 : end if
1202 : else
1203 186112580 : if (ndprojs>0) atom_dprojs(:,:,:,:) = zero
1204 1296 : if (nd2projs>0) atom_d2projs(:,:,:,:) = zero
1205 : end if
1206 1296 : if (signs==1 .and. (choice==3 .or. choice==23 .or. choice==54 .or. choice==55 .or. choice==6)) then
1207 : #ifdef HAVE_OPENMP_OFFLOAD
1208 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1209 : !$OMP& PRIVATE(ilmn,idir,ipw) MAP(to:atom_dprojs,ffnl) &
1210 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1211 : #endif
1212 9992 : do ilmn=1,nlmn_o
1213 38024 : do idir=1,ndprojs
1214 51863912 : do ipw=1, npw
1215 51854568 : atom_dprojs(1, ipw, idir, ilmn) = wt * ffnl(ipw, idir+1, ilmn, itypat)
1216 : end do
1217 : end do
1218 : end do
1219 : end if
1220 1296 : if (signs==2 .and. (choice==3 .or. choice==5 .or. choice==51)) then
1221 : #ifdef HAVE_OPENMP_OFFLOAD
1222 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) &
1223 : !$OMP& PRIVATE(ilmn,ipw) MAP(to:atom_dprojs,ffnl) &
1224 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1225 : #endif
1226 0 : do ilmn=1,nlmn_o
1227 1296 : do ipw=1, npw
1228 0 : atom_dprojs(1, ipw, 1, ilmn) = wt * ffnl(ipw, 1+ffnl_dir, ilmn, itypat)
1229 : end do
1230 : end do
1231 : end if
1232 1296 : if(signs==1 .and. choice==54) then
1233 : #ifdef HAVE_OPENMP_OFFLOAD
1234 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1235 : !$OMP& PRIVATE(ilmn,idir,ipw) MAP(to:atom_d2projs,ffnl) &
1236 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1237 : #endif
1238 0 : do ilmn=1,nlmn_o
1239 0 : do idir=1,nd2projs
1240 0 : do ipw=1, npw
1241 0 : atom_d2projs(1, ipw, idir, ilmn) = wt * ffnl(ipw, idir+1, ilmn, itypat)
1242 : end do
1243 : end do
1244 : end do
1245 1296 : else if(signs==1 .and. choice==55) then
1246 : #ifdef HAVE_OPENMP_OFFLOAD
1247 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1248 : !$OMP& PRIVATE(ilmn,idir,ipw) MAP(to:atom_d2projs,ffnl) &
1249 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1250 : #endif
1251 0 : do ilmn=1,nlmn_o
1252 0 : do idir=1,nd2projs
1253 0 : do ipw=1, npw
1254 0 : atom_d2projs(1, ipw, idir, ilmn) = wt * ffnl(ipw, idir+4, ilmn, itypat)
1255 : end do
1256 : end do
1257 : end do
1258 1296 : else if(signs==1 .and. choice==6) then
1259 : #ifdef HAVE_OPENMP_OFFLOAD
1260 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1261 : !$OMP& PRIVATE(ilmn,idir,ipw) MAP(to:atom_d2projs,ffnl) &
1262 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1263 : #endif
1264 0 : do ilmn=1,nlmn_o
1265 0 : do idir=1,10
1266 0 : do ipw=1, npw
1267 0 : atom_d2projs(1, ipw, idir, ilmn) = wt * ffnl(ipw, idir, ilmn, itypat)
1268 : end do
1269 : end do
1270 : end do
1271 : end if
1272 :
1273 1296 : nlmn_p=0
1274 : ! multiply by (-i)^l
1275 1296 : if (ndprojs>0) then
1276 9992 : do ilmn=1,nlmn_o
1277 9344 : il=mod(indlmn(1,ilmn, itypat),4);
1278 9344 : if(.not. (mod(il,2)==0)) then
1279 3888 : nlmn_p=nlmn_p+1
1280 3888 : lmn_parity(nlmn_p)=ilmn
1281 : end if
1282 9992 : scal(ilmn)=1; if(il>1) scal(ilmn)=-1
1283 : end do
1284 : ! multiply by -1
1285 648 : if(gpu_option==ABI_GPU_OPENMP) then
1286 : #ifdef HAVE_OPENMP_OFFLOAD
1287 : !$OMP TARGET UPDATE TO(scal,lmn_parity)
1288 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1289 : !$OMP& PRIVATE(ilmn,idir,ipw) MAP(to:atom_dprojs,scal) &
1290 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1291 : do ilmn=1,nlmn_o
1292 : do idir=1,ndprojs
1293 : do ipw=1,npw
1294 : atom_dprojs(1,ipw,idir,ilmn) = atom_dprojs(1,ipw,idir,ilmn) * scal(ilmn)
1295 : atom_dprojs(2,ipw,idir,ilmn) = atom_dprojs(2,ipw,idir,ilmn) * scal(ilmn)
1296 : end do
1297 : end do
1298 : end do
1299 : if(nd2projs>0) then
1300 : !$OMP TARGET TEAMS DISTRIBUTE &
1301 : !$OMP& PRIVATE(ilmn) MAP(to:atom_d2projs,scal) &
1302 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1303 : do ilmn=1,nlmn_o
1304 : !$OMP PARALLEL DO COLLAPSE(2) PRIVATE(idir,ipw)
1305 : do idir=1,nd2projs
1306 : do ipw=1,npw
1307 : atom_d2projs(1,ipw,idir,ilmn) = atom_d2projs(1,ipw,idir,ilmn) * scal(ilmn)
1308 : atom_d2projs(2,ipw,idir,ilmn) = atom_d2projs(2,ipw,idir,ilmn) * scal(ilmn)
1309 : end do
1310 : end do
1311 : end do
1312 : end if
1313 : #endif
1314 : else
1315 9992 : do ilmn=1,nlmn_o
1316 155517632 : atom_dprojs(:,:,:,ilmn) = atom_dprojs(:,:,:,ilmn) * scal(ilmn)
1317 : end do
1318 648 : if(nd2projs>0) then
1319 0 : do ilmn=1,nlmn_o
1320 0 : atom_d2projs(:,:,:,ilmn) = atom_d2projs(:,:,:,ilmn) * scal(ilmn)
1321 : end do
1322 : end if
1323 : end if
1324 : ! multiply by -i
1325 : #ifdef HAVE_OPENMP_OFFLOAD
1326 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1327 : !$OMP MAP(to:atom_dprojs,lmn_parity) PRIVATE(ilmn,ilmn_p,idir,ipw,tmp) &
1328 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1329 : #endif
1330 4536 : do ilmn_p=1,nlmn_p
1331 16200 : do idir=1,ndprojs
1332 21157164 : do ipw=1,npw
1333 21141612 : ilmn=lmn_parity(ilmn_p)
1334 21141612 : tmp = atom_dprojs(2,ipw,idir,ilmn)
1335 21141612 : atom_dprojs(2,ipw,idir,ilmn) = -atom_dprojs(1,ipw,idir,ilmn)
1336 21153276 : atom_dprojs(1,ipw,idir,ilmn) = tmp
1337 : end do
1338 : end do
1339 : end do
1340 648 : if(ngrads2>0) then
1341 : #ifdef HAVE_OPENMP_OFFLOAD
1342 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1343 : !$OMP& MAP(to:atom_d2projs) PRIVATE(idir,ipw,ilmn) &
1344 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1345 : #endif
1346 0 : do ilmn=1,nlmn_o
1347 0 : do idir=1,nd2projs
1348 0 : do ipw=1,npw
1349 0 : atom_d2projs(1,ipw,idir,ilmn) = -atom_d2projs(1,ipw,idir,ilmn)
1350 0 : atom_d2projs(2,ipw,idir,ilmn) = -atom_d2projs(2,ipw,idir,ilmn)
1351 : end do
1352 : end do
1353 : end do
1354 : #ifdef HAVE_OPENMP_OFFLOAD
1355 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1356 : !$OMP& MAP(to:atom_d2projs,lmn_parity) PRIVATE(ilmn,ilmn_p,idir,ipw,tmp) &
1357 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1358 : #endif
1359 0 : do ilmn_p=1,nlmn_p
1360 0 : do idir=1,nd2projs
1361 0 : do ipw=1,npw
1362 0 : ilmn=lmn_parity(ilmn_p)
1363 0 : tmp = atom_d2projs(2,ipw,idir,ilmn)
1364 0 : atom_d2projs(2,ipw,idir,ilmn) = -atom_d2projs(1,ipw,idir,ilmn)
1365 0 : atom_d2projs(1,ipw,idir,ilmn) = tmp
1366 : end do
1367 : end do
1368 : end do
1369 : end if
1370 : end if
1371 :
1372 : ! multiply by conj(ph3d)
1373 1296 : if (ndprojs>0) then
1374 : #ifdef HAVE_OPENMP_OFFLOAD
1375 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1376 : !$OMP& PRIVATE(ilmn,tmp,ipw,idir) MAP(to:atom_dprojs,ph3d) &
1377 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1378 : #endif
1379 9992 : do ilmn=1,nlmn_o
1380 38024 : do idir=1,ndprojs
1381 51863912 : do ipw=1,npw
1382 51826536 : tmp = atom_dprojs(1, ipw, idir,ilmn)
1383 : atom_dprojs(1, ipw, idir,ilmn) = atom_dprojs(1, ipw, idir,ilmn) * ph3d(1, ipw, iaph3d) &
1384 51826536 : & + atom_dprojs(2, ipw, idir,ilmn) * ph3d(2, ipw, iaph3d)
1385 : atom_dprojs(2, ipw, idir,ilmn) = atom_dprojs(2, ipw, idir,ilmn) * ph3d(1, ipw, iaph3d) &
1386 51854568 : & - tmp * ph3d(2, ipw, iaph3d)
1387 : end do
1388 : end do
1389 : end do
1390 : end if
1391 1296 : if (nd2projs>0) then
1392 : #ifdef HAVE_OPENMP_OFFLOAD
1393 : !$OMP TARGET TEAMS DISTRIBUTE &
1394 : !$OMP& PRIVATE(ilmn,idir,ipw,tmp) MAP(to:atom_d2projs,ph3d) &
1395 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1396 : #endif
1397 0 : do ilmn=1,nlmn_o
1398 : !$OMP PARALLEL DO COLLAPSE(2) PRIVATE(ipw,idir,tmp)
1399 0 : do idir=1,nd2projs
1400 0 : do ipw=1,npw
1401 0 : tmp = atom_d2projs(1, ipw, idir,ilmn)
1402 : atom_d2projs(1, ipw, idir,ilmn) = atom_d2projs(1, ipw, idir,ilmn) * ph3d(1, ipw, iaph3d) &
1403 0 : & + atom_d2projs(2, ipw, idir,ilmn) * ph3d(2, ipw, iaph3d)
1404 : atom_d2projs(2, ipw, idir,ilmn) = atom_d2projs(2, ipw, idir,ilmn) * ph3d(1, ipw, iaph3d) &
1405 0 : & - tmp * ph3d(2, ipw, iaph3d)
1406 : end do
1407 : end do
1408 : end do
1409 : end if
1410 :
1411 : !! Handling dprojs
1412 :
1413 1296 : if(signs==1 .and. (choice==3 .or. choice==23 .or. choice==55 .or. choice==6)) then
1414 648 : if(istwf_k <= 1) then
1415 : #ifdef HAVE_OPENMP_OFFLOAD
1416 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1417 : !$OMP& PRIVATE(ilmn,ipw,idir,idir1,idir2) MAP(to:atom_dprojs,dprojs,kpg) &
1418 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1419 : #endif
1420 7088 : do ilmn=lmn_beg,nlmn
1421 46784 : do idir=1,6
1422 87912760 : do ipw=1,npw
1423 87866448 : idir1=alpha(idir);idir2=beta(idir)
1424 : dprojs(1, ipw, shift_grad+(ilmn-lmn_beg)*ngrads+idir) = &
1425 : & -half*(atom_dprojs(1, ipw, idir1, ilmn)*kpg(ipw,idir2) &
1426 87866448 : & +atom_dprojs(1, ipw, idir2, ilmn)*kpg(ipw,idir1))
1427 : dprojs(2, ipw, shift_grad+(ilmn-lmn_beg)*ngrads+idir) = &
1428 : & -half*(atom_dprojs(2, ipw, idir1, ilmn)*kpg(ipw,idir2) &
1429 87906144 : & +atom_dprojs(2, ipw, idir2, ilmn)*kpg(ipw,idir1))
1430 : end do
1431 : end do
1432 : end do
1433 : else ! istwf_k>1
1434 : #ifdef HAVE_OPENMP_OFFLOAD
1435 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1436 : !$OMP& PRIVATE(ilmn,ipw,idir,idir1,idir2) MAP(to:atom_dprojs,dprojs_r,dprojs_i,kpg) &
1437 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1438 : #endif
1439 2904 : do ilmn=lmn_beg,nlmn
1440 19272 : do idir=1,6
1441 15805720 : do ipw=1,npw
1442 15786624 : idir1=alpha(idir);idir2=beta(idir)
1443 : dprojs_r(1, ipw, shift_grad+(ilmn-lmn_beg)*ngrads+idir) = &
1444 : & -half*(atom_dprojs(1, ipw, idir1, ilmn)*kpg(ipw,idir2) &
1445 15786624 : & +atom_dprojs(1, ipw, idir2, ilmn)*kpg(ipw,idir1))
1446 :
1447 : dprojs_i(1, ipw, shift_grad+(ilmn-lmn_beg)*ngrads+idir) = &
1448 : & -half*(atom_dprojs(2, ipw, idir1, ilmn)*kpg(ipw,idir2) &
1449 15802992 : & +atom_dprojs(2, ipw, idir2, ilmn)*kpg(ipw,idir1))
1450 : end do
1451 : end do
1452 : end do
1453 : end if
1454 : end if
1455 :
1456 :
1457 1296 : if(signs==1 .and. (choice==2 .or. choice==23 .or. choice==4 .or. choice==54 .or. choice==6)) then
1458 648 : igrad=0; if(choice==23 .or. choice==6) igrad=6
1459 648 : if(istwf_k <= 1) then
1460 : #ifdef HAVE_OPENMP_OFFLOAD
1461 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1462 : !$OMP& PRIVATE(ilmn,ipw,idir) MAP(to:projs,dprojs,kpg) &
1463 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1464 : #endif
1465 7088 : do ilmn=lmn_beg,nlmn
1466 26936 : do idir=1,3
1467 43959688 : do ipw=1,npw
1468 : dprojs(1, ipw, shift_grad+(ilmn-lmn_beg)*ngrads+igrad+idir) = &
1469 43933224 : & +projs(2, ipw, shift+ilmn-lmn_beg+1)*kpg(ipw,idir)*two_pi
1470 : dprojs(2, ipw, shift_grad+(ilmn-lmn_beg)*ngrads+igrad+idir) = &
1471 43953072 : & -projs(1, ipw, shift+ilmn-lmn_beg+1)*kpg(ipw,idir)*two_pi
1472 : end do
1473 : end do
1474 : end do
1475 : else ! istwf_k>1
1476 : #ifdef HAVE_OPENMP_OFFLOAD
1477 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1478 : !$OMP& PRIVATE(ilmn,ipw,idir) MAP(to:projs_r,projs_i,dprojs_r,dprojs_i,kpg) &
1479 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1480 : #endif
1481 2904 : do ilmn=lmn_beg,nlmn
1482 11088 : do idir=1,3
1483 7904224 : do ipw=1,npw
1484 : dprojs_r(1, ipw, shift_grad+(ilmn-lmn_beg)*ngrads+igrad+idir) = &
1485 7893312 : & +projs_i(1, ipw, shift+ilmn-lmn_beg+1)*kpg(ipw,idir)*two_pi
1486 : dprojs_i(1, ipw, shift_grad+(ilmn-lmn_beg)*ngrads+igrad+idir) = &
1487 7901496 : & -projs_r(1, ipw, shift+ilmn-lmn_beg+1)*kpg(ipw,idir)*two_pi
1488 : end do
1489 : end do
1490 : end do
1491 : end if
1492 : end if
1493 :
1494 1296 : if(signs==1 .and. (choice==5 .or. choice==51 .or. choice==54 .or. choice==55)) then
1495 0 : igrad=0; if(choice==54) igrad=3; if(choice==55) igrad=6
1496 : #ifdef HAVE_OPENMP_OFFLOAD
1497 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1498 : !$OMP& PRIVATE(ilmn,ipw,idir) MAP(to:atom_dprojs,dprojs) &
1499 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1500 : #endif
1501 0 : do ilmn=lmn_beg,nlmn
1502 0 : do idir=1,3
1503 0 : do ipw=1,npw
1504 : dprojs(1, ipw, shift_grad+(ilmn-lmn_beg)*ngrads+igrad+idir) = &
1505 0 : & +atom_dprojs(1, ipw, idir, ilmn)
1506 : dprojs(2, ipw, shift_grad+(ilmn-lmn_beg)*ngrads+igrad+idir) = &
1507 0 : & +atom_dprojs(2, ipw, idir, ilmn)
1508 : end do
1509 : end do
1510 : end do
1511 : end if
1512 :
1513 1296 : if(signs==2 .and. (choice==2)) then
1514 : #ifdef HAVE_OPENMP_OFFLOAD
1515 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) &
1516 : !$OMP& PRIVATE(ilmn,ipw) MAP(to:projs,dprojs,kpg) &
1517 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1518 : #endif
1519 0 : do ilmn=lmn_beg,nlmn
1520 0 : do ipw=1,npw
1521 : dprojs(1, ipw, shift_grad+ilmn) = &
1522 0 : & projs(2, ipw, shift+ilmn-lmn_beg+1)*kpg(ipw,idir_pert)*two_pi
1523 : dprojs(2, ipw, shift_grad+ilmn) = &
1524 0 : & -projs(1, ipw, shift+ilmn-lmn_beg+1)*kpg(ipw,idir_pert)*two_pi
1525 : end do
1526 : end do
1527 : end if
1528 :
1529 :
1530 1296 : if(signs==2 .and. (choice==3)) then
1531 : #ifdef HAVE_OPENMP_OFFLOAD
1532 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) &
1533 : !$OMP& PRIVATE(ilmn,ipw) MAP(to:atom_dprojs,dprojs) &
1534 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1535 : #endif
1536 0 : do ilmn=lmn_beg,nlmn
1537 0 : do ipw=1,npw
1538 : dprojs(1, ipw, shift_grad+ilmn) = &
1539 0 : & -atom_dprojs(1, ipw, 1, ilmn)
1540 : dprojs(2, ipw, shift_grad+ilmn) = &
1541 0 : & -atom_dprojs(2, ipw, 1, ilmn)
1542 : end do
1543 : end do
1544 : end if
1545 :
1546 :
1547 1296 : if(signs==2 .and. (choice==5 .or. choice==51)) then
1548 : #ifdef HAVE_OPENMP_OFFLOAD
1549 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) &
1550 : !$OMP& PRIVATE(ilmn,ipw) MAP(to:atom_dprojs,dprojs) &
1551 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1552 : #endif
1553 0 : do ilmn=lmn_beg,nlmn
1554 0 : do ipw=1,npw
1555 : dprojs(1, ipw, shift_grad+ilmn) = &
1556 0 : & +atom_dprojs(1, ipw, 1, ilmn)
1557 : dprojs(2, ipw, shift_grad+ilmn) = &
1558 0 : & +atom_dprojs(2, ipw, 1, ilmn)
1559 : end do
1560 : end do
1561 : end if
1562 :
1563 :
1564 :
1565 : ! Handling d2projs
1566 :
1567 1296 : if(signs==1 .and. choice==4) then
1568 : #ifdef HAVE_OPENMP_OFFLOAD
1569 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) &
1570 : !$OMP& PRIVATE(idir,ilmn,ipw) MAP(to:projs,d2projs,kpg) &
1571 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1572 : #endif
1573 0 : do ilmn=lmn_beg,nlmn
1574 0 : do idir=1,6
1575 0 : do ipw=1,npw
1576 : d2projs(1, ipw, shift_grad2+(ilmn-lmn_beg)*ngrads2+idir) = &
1577 0 : & -projs(1, ipw, shift+ilmn-lmn_beg+1)*kpg(ipw,idir+3)*two_pi2
1578 : d2projs(2, ipw, shift_grad2+(ilmn-lmn_beg)*ngrads2+idir) = &
1579 0 : & -projs(2, ipw, shift+ilmn-lmn_beg+1)*kpg(ipw,idir+3)*two_pi2
1580 : end do
1581 : end do
1582 : end do
1583 : end if
1584 :
1585 1296 : if(signs==1 .and. choice==54) then
1586 : #ifdef HAVE_OPENMP_OFFLOAD
1587 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(4) &
1588 : !$OMP& PRIVATE(idir1,idir2,ilmn,ipw) MAP(to:atom_d2projs,d2projs,kpg) &
1589 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1590 : #endif
1591 0 : do ilmn=lmn_beg,nlmn
1592 0 : do idir1=1,3
1593 0 : do idir2=1,3
1594 0 : do ipw=1,npw
1595 : d2projs(1, ipw, shift_grad2+(ilmn-lmn_beg)*ngrads2+(idir1-1)*3+idir2) = &
1596 0 : & -atom_d2projs(2, ipw, idir2, ilmn)*kpg(ipw,idir1)*two_pi
1597 : d2projs(2, ipw, shift_grad2+(ilmn-lmn_beg)*ngrads2+(idir1-1)*3+idir2) = &
1598 0 : & +atom_d2projs(1, ipw, idir2, ilmn)*kpg(ipw,idir1)*two_pi
1599 : end do
1600 : end do
1601 : end do
1602 : end do
1603 : end if
1604 :
1605 1296 : if(signs==1 .and. choice==55) then
1606 : #ifdef HAVE_OPENMP_OFFLOAD
1607 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(4) &
1608 : !$OMP& PRIVATE(idir1,idir2,ilmn,ipw) MAP(to:atom_d2projs,d2projs,kpg) &
1609 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1610 : #endif
1611 0 : do ilmn=lmn_beg,nlmn
1612 0 : do idir1=1,6
1613 0 : do idir2=1,3
1614 0 : do ipw=1,npw
1615 : d2projs(1, ipw, shift_grad2+(ilmn-lmn_beg)*ngrads2+(idir1-1)*3+idir2) = &
1616 0 : & +atom_d2projs(1, ipw, idir1, ilmn)*kpg(ipw,idir2)
1617 : d2projs(2, ipw, shift_grad2+(ilmn-lmn_beg)*ngrads2+(idir1-1)*3+idir2) = &
1618 0 : & +atom_d2projs(2, ipw, idir1, ilmn)*kpg(ipw,idir2)
1619 : end do
1620 : end do
1621 : end do
1622 : end do
1623 : end if
1624 :
1625 1296 : if(signs==1 .and. choice==6) then
1626 : #ifdef HAVE_OPENMP_OFFLOAD
1627 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(4) &
1628 : !$OMP& PRIVATE(idir1,idir2,jdir1,jdir2,kdir1,kdir2,ldir1,ldir2,ldir3,ldir4,ilmn,ipw) &
1629 : !$OMP& MAP(to:atom_d2projs,d2projs,kpg) &
1630 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1631 : #endif
1632 0 : do ilmn=lmn_beg,nlmn
1633 0 : do idir1=1,6
1634 0 : do idir2=1,6
1635 0 : do ipw=1,npw
1636 0 : jdir1=alpha(idir1);jdir2=beta(idir1)
1637 0 : kdir1=alpha(idir2);kdir2=beta(idir2)
1638 0 : ldir1=gamma(jdir1,kdir1)
1639 0 : ldir2=gamma(jdir2,kdir1)
1640 0 : ldir3=gamma(jdir1,kdir2)
1641 0 : ldir4=gamma(jdir2,kdir2)
1642 : d2projs(1, ipw, shift_grad2+(ilmn-lmn_beg)*ngrads2+(idir1-1)*6+idir2) = &
1643 : & -atom_d2projs(1, ipw, 4+ldir1, ilmn)*kpg(ipw,jdir2)*kpg(ipw,kdir2)*quarter &
1644 : & -atom_d2projs(1, ipw, 4+ldir2, ilmn)*kpg(ipw,jdir1)*kpg(ipw,kdir2)*quarter &
1645 : & -atom_d2projs(1, ipw, 4+ldir3, ilmn)*kpg(ipw,jdir2)*kpg(ipw,kdir1)*quarter &
1646 0 : & -atom_d2projs(1, ipw, 4+ldir4, ilmn)*kpg(ipw,jdir1)*kpg(ipw,kdir1)*quarter
1647 : d2projs(2, ipw, shift_grad2+(ilmn-lmn_beg)*ngrads2+(idir1-1)*6+idir2) = &
1648 : & -atom_d2projs(2, ipw, 4+ldir1, ilmn)*kpg(ipw,jdir2)*kpg(ipw,kdir2)*quarter &
1649 : & -atom_d2projs(2, ipw, 4+ldir2, ilmn)*kpg(ipw,jdir1)*kpg(ipw,kdir2)*quarter &
1650 : & -atom_d2projs(2, ipw, 4+ldir3, ilmn)*kpg(ipw,jdir2)*kpg(ipw,kdir1)*quarter &
1651 0 : & -atom_d2projs(2, ipw, 4+ldir4, ilmn)*kpg(ipw,jdir1)*kpg(ipw,kdir1)*quarter
1652 : end do
1653 : end do
1654 : end do
1655 : end do
1656 0 : igrad=36
1657 : #ifdef HAVE_OPENMP_OFFLOAD
1658 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(4) &
1659 : !$OMP& PRIVATE(idir1,idir2,jdir1,jdir2,ilmn,ipw) &
1660 : !$OMP& MAP(to:atom_d2projs,d2projs,kpg) &
1661 : !$OMP& IF(gpu_option==ABI_GPU_OPENMP)
1662 : #endif
1663 0 : do ilmn=lmn_beg,nlmn
1664 0 : do idir1=1,6 !mub
1665 0 : do idir2=1,3 !mua
1666 0 : do ipw=1,npw
1667 0 : jdir1=alpha(idir1);jdir2=beta(idir1)
1668 : d2projs(1, ipw, shift_grad2+(ilmn-lmn_beg)*ngrads2+igrad+(idir1-1)*3+idir2) = &
1669 : & +kpg(ipw,idir2)*(atom_d2projs(2,ipw,1+jdir1,ilmn)*kpg(ipw,jdir2)*pi &
1670 0 : & + atom_d2projs(2,ipw,1+jdir2,ilmn)*kpg(ipw,jdir1)*pi)
1671 :
1672 : d2projs(2, ipw, shift_grad2+(ilmn-lmn_beg)*ngrads2+igrad+(idir1-1)*3+idir2) = &
1673 : & -kpg(ipw,idir2)*(atom_d2projs(1,ipw,1+jdir1,ilmn)*kpg(ipw,jdir2)*pi &
1674 0 : & + atom_d2projs(1,ipw,1+jdir2,ilmn)*kpg(ipw,jdir1)*pi)
1675 : end do
1676 : end do
1677 : end do
1678 : end do
1679 : end if
1680 :
1681 1296 : iaph3d = iaph3d + 1
1682 1296 : shift_grad2 = shift_grad2 + ngrads2*(nlmn-lmn_beg+1)
1683 1296 : shift_grad = shift_grad + ngrads*(nlmn-lmn_beg+1)
1684 :
1685 2008 : if(gemm_nonlop_block_size > 1) then
1686 0 : shift = shift + nlmn - (lmn_beg-1)
1687 0 : shift_do = shift_do + nlmn
1688 0 : if(shift_do >= iend - 1) exit
1689 : else
1690 1296 : shift = shift + nlmn
1691 : end if
1692 :
1693 : end do
1694 1136 : if(gemm_nonlop_block_size > 1 .and. shift_do >= iend - 1) exit
1695 : end do
1696 :
1697 : #ifdef HAVE_OPENMP_OFFLOAD
1698 : !!$OMP TARGET EXIT DATA MAP(delete:kpg) IF(gpu_option==ABI_GPU_OPENMP)
1699 : !$OMP TARGET EXIT DATA MAP(delete:ffnl) IF(gpu_option==ABI_GPU_OPENMP .and. map_ffnl)
1700 : !$OMP TARGET EXIT DATA MAP(delete:ph3d) IF(gpu_option==ABI_GPU_OPENMP .and. map_ph3d)
1701 : #endif
1702 :
1703 424 : end subroutine prep_dprojectors
1704 : !!***
1705 :
1706 0 : end module m_gemm_nonlop_projectors
1707 : !!***
|
