Line data Source code
1 : !!****m* ABINIT/m_opernlb_gemm
2 : !! NAME
3 : !! m_opernlb_gemm
4 : !!
5 : !! FUNCTION
6 : !!
7 : !! COPYRIGHT
8 : !! Copyright (C) 2008-2026 ABINIT group (MT)
9 : !! This file is distributed under the terms of the
10 : !! GNU General Public License, see ~abinit/COPYING
11 : !! or http://www.gnu.org/copyleft/gpl.txt .
12 : !!
13 : !! SOURCE
14 :
15 : #if defined HAVE_CONFIG_H
16 : #include "config.h"
17 : #endif
18 :
19 : #include "abi_common.h"
20 :
21 : module m_opernlb_gemm
22 :
23 : use defs_basis
24 : use m_abicore
25 : use m_errors
26 : USE_MPI
27 : use m_xmpi
28 : use m_abi_linalg
29 : use m_gputk
30 : use m_gemm_nonlop_projectors
31 :
32 : use defs_abitypes, only : MPI_type
33 : use m_time, only : timab
34 :
35 : use, intrinsic :: iso_c_binding, only : c_ptr,c_loc,c_size_t,c_f_pointer
36 :
37 : implicit none
38 :
39 : private
40 : !!***
41 :
42 : public :: opernlb_gemm
43 :
44 : contains
45 : !!***
46 :
47 :
48 : !----------------------------------------------------------------------
49 :
50 : !!****f* m_opernlb_gemm/opernlb_gemm_distributed
51 : !! NAME
52 : !! opernlb_gemm_distributed
53 : !!
54 : !! FUNCTION
55 : !! Distributed version of "opernlb" GEMM called in gemm_nonlop.
56 : !!
57 : !! INPUTS
58 : !!
59 : !! SOURCE
60 0 : subroutine opernlb_gemm_distributed(rank,nprocs,npw,ndat,&
61 : & transa,transb,&
62 : & nprojs,nprojs_blk,nprojs_last_blk,cplex,&
63 0 : & projs_local,projections,vectout,gpu_option)
64 : integer, intent(in) :: rank,nprocs,npw,ndat,gpu_option
65 : integer, intent(in) :: nprojs,nprojs_blk,nprojs_last_blk,cplex
66 : character(len=1),intent(in) :: transa,transb
67 : real(dp), intent(in), target :: projs_local(cplex,npw,nprojs_last_blk)
68 : real(dp), intent(in), target :: projections(cplex,nprojs,ndat)
69 : real(dp), intent(out), target :: vectout(cplex,npw,ndat)
70 :
71 : !Local variables
72 : integer :: iblock,ibeg,req(2),ierr,nprojs_cur_blk,rank_prev,rank_next
73 : complex(dp) :: beta
74 0 : real(dp), ABI_CONTIGUOUS pointer :: recv_buf(:,:,:), work_buf(:,:,:)
75 0 : real(dp), allocatable, target :: projs_recv(:,:,:)
76 0 : real(dp), ABI_CONTIGUOUS pointer :: projections_1d(:)
77 : type(c_ptr) :: projections_cptr
78 : ! *************************************************************************
79 :
80 0 : ABI_MALLOC(projs_recv, (cplex, npw, nprojs_last_blk))
81 : #ifdef HAVE_OPENMP_OFFLOAD
82 : !$OMP TARGET ENTER DATA MAP(alloc:projs_recv) IF(gpu_option==ABI_GPU_OPENMP)
83 : #endif
84 :
85 0 : rank_next=modulo(rank + 1,nprocs)
86 0 : rank_prev=rank - 1
87 0 : if(rank_prev == -1) rank_prev = nprocs - 1
88 :
89 0 : beta = czero
90 :
91 0 : do iblock=1,nprocs
92 :
93 0 : if(rank+iblock == nprocs) then
94 0 : nprojs_cur_blk = nprojs_last_blk
95 : else
96 0 : nprojs_cur_blk = nprojs_blk
97 : end if
98 :
99 0 : if(modulo(iblock,2)==1) then
100 : ! XG20241028 : This coding confused the gnu_8.5 compiler of buda2_gnu_8.5_cuda, wrt the contiguous character of the target.
101 : ! It declared an error. Make it simple !
102 : ! work_buf => projs_local(1:cplex,1:npw,1:nprojs_last_blk)
103 : ! recv_buf => projs_recv(1:cplex,1:npw,1:nprojs_last_blk)
104 0 : work_buf => projs_local
105 0 : recv_buf => projs_recv
106 : else
107 : ! XG20241028 : Same as above
108 : ! work_buf => projs_recv(1:cplex,1:npw,1:nprojs_last_blk)
109 : ! recv_buf => projs_local(1:cplex,1:npw,1:nprojs_last_blk)
110 0 : work_buf => projs_recv
111 0 : recv_buf => projs_local
112 : end if
113 :
114 0 : if(gpu_option == ABI_GPU_DISABLED) then
115 0 : call xmpi_isend(work_buf,rank_prev,iblock,gemm_nonlop_block_comm,req(1),ierr)
116 0 : call xmpi_irecv(recv_buf,rank_next,iblock,gemm_nonlop_block_comm,req(2),ierr)
117 : else if(gpu_option == ABI_GPU_OPENMP) then
118 : #ifdef HAVE_OPENMP_OFFLOAD
119 : #ifndef HAVE_GPU_MPI
120 :
121 : ! GPU-aware MPI not available : perform MPI comms on CPU
122 : !$OMP TARGET UPDATE FROM(work_buf) if(iblock==1)
123 : call xmpi_isend(work_buf,rank_prev,iblock,gemm_nonlop_block_comm,req(1),ierr)
124 : call xmpi_irecv(recv_buf,rank_next,iblock,gemm_nonlop_block_comm,req(2),ierr)
125 :
126 : #else
127 :
128 : ! GPU-aware MPI available : pass GPU buffers to MPI
129 : call xmpi_isend(work_buf,rank_prev,iblock,gemm_nonlop_block_comm,req(1),ierr,use_omp_map=.true.)
130 : call xmpi_irecv(recv_buf,rank_next,iblock,gemm_nonlop_block_comm,req(2),ierr,use_omp_map=.true.)
131 :
132 : #endif
133 : #endif
134 : end if
135 :
136 0 : ibeg = 1 + modulo(rank+iblock-1,nprocs)*nprojs_blk
137 :
138 : ! Small trickery here:
139 : ! The multiplication is performed over a slice of projectors in 'work_buf' matrix.
140 : ! In that case, 'projections' matrix is still sized by all projectors but only
141 : ! a slice of the projectors present in work_buf is used.
142 : !
143 : ! Therefore, we need to provide GEMM with the start of the slice
144 : ! with 'projections' rather than the start of 'projections' as usual.
145 : !
146 : ! For that reason, we turn 'projections' matrix into single-rank in order
147 : ! to pass the right starting row.
148 0 : projections_cptr = c_loc(projections)
149 0 : call c_f_pointer(projections_cptr, projections_1d, [cplex * nprojs * ndat])
150 :
151 : call abi_xgemm(transa,transb,&
152 : & npw, ndat, nprojs_cur_blk, cone, &
153 : & work_buf, npw,&
154 : & projections_1d(cplex*ibeg-(cplex-1):cplex*nprojs*ndat),nprojs,&
155 : & beta, &
156 : & vectout, npw, &
157 0 : & x_cplx=cplex,gpu_option=gpu_option)
158 :
159 0 : beta = cone
160 :
161 0 : call xmpi_wait(req(1),ierr)
162 0 : call xmpi_wait(req(2),ierr)
163 : !call xmpi_waitall(req,ierr)
164 :
165 : #ifdef HAVE_OPENMP_OFFLOAD
166 : #ifndef HAVE_GPU_MPI
167 : ! If MPI is not GPU-aware, push received data to GPU
168 : !$OMP TARGET UPDATE TO(recv_buf) IF(gpu_option==ABI_GPU_OPENMP)
169 : #endif
170 : #endif
171 :
172 : end do
173 :
174 0 : if(modulo(iblock,2)==1) then
175 0 : if(gpu_option == ABI_GPU_DISABLED) then
176 0 : call DCOPY(cplex*npw*nprojs_cur_blk, recv_buf, 1, work_buf, 1)
177 : else if(gpu_option == ABI_GPU_OPENMP) then
178 : #ifdef HAVE_OPENMP_OFFLOAD
179 : !$OMP TARGET DATA USE_DEVICE_ADDR(work_buf,recv_buf)
180 : call copy_gpu_to_gpu(c_loc(work_buf), c_loc(recv_buf), INT(cplex, c_size_t)*npw*nprojs_last_blk*dp)
181 : !$OMP END TARGET DATA
182 : #endif
183 : end if
184 : end if
185 :
186 : #ifdef HAVE_OPENMP_OFFLOAD
187 : !$OMP TARGET EXIT DATA MAP(delete:projs_recv) IF(gpu_option==ABI_GPU_OPENMP)
188 : #endif
189 0 : ABI_FREE(projs_recv)
190 :
191 0 : end subroutine opernlb_gemm_distributed
192 : !!***
193 :
194 : !----------------------------------------------------------------------
195 :
196 25082 : subroutine opernlb_xgemm(cplex,transa,transb,npw,ndat,nprojs,alpha,a,lda,b,ldb,beta,c,ldc,&
197 : & rank, nprocs,&
198 : & nprojs_blk, nprojs_last_blk,&
199 : & iblock,&
200 : & gpu_option,use_distrib,use_sliced_gemms)
201 :
202 : !Arguments ------------------------------------
203 : integer,intent(in) :: cplex,lda,ldb,ldc,npw,ndat,nprojs,gpu_option
204 : integer,intent(in) :: rank,nprocs,nprojs_blk,nprojs_last_blk
205 : integer,intent(in) :: iblock
206 : logical,intent(in) :: use_distrib,use_sliced_gemms
207 : complex(dp),intent(in) :: alpha,beta
208 : character(len=1),intent(in) :: transa,transb
209 : real(dp),target,intent(in) :: a(cplex,lda,nprojs), b(cplex,ldb,ndat)
210 : real(dp),target,intent(inout) :: c(cplex,ldc,ndat)
211 :
212 : integer :: ibeg
213 25082 : real(dp), ABI_CONTIGUOUS pointer :: b_1d(:)
214 : type(c_ptr) :: b_cptr
215 : ! *********************************************************************
216 :
217 : #ifdef FC_NVHPC
218 : ! Silly fix for NVHPC
219 : if(cplex == -42) write(100,*) cplex
220 : #endif
221 :
222 25082 : if(use_distrib) then
223 : call opernlb_gemm_distributed(rank,nprocs,npw,ndat,&
224 : & transa,transb,&
225 : & nprojs,&
226 : & nprojs_blk,&
227 : & nprojs_last_blk,&
228 : & cplex,&
229 : & a,&
230 0 : & b,c,gpu_option)
231 : else
232 : ! Small trickery here:
233 : ! When use_sliced_gemm is on, the multiplication is performed over
234 : ! a slice of projectors in A matrix.
235 : ! In that case, B matrix is still sized by all projectors but only
236 : ! a slice matching the projectors present in A is used.
237 : !
238 : ! Therefore, we need to provide GEMM with the start of the slice
239 : ! of relevant projectors within B rather than the start of B as usual.
240 : !
241 : ! For that reason, we turn B matrix into single-rank in order
242 : ! to pass the right starting row.
243 25082 : ibeg = 1
244 25082 : if(use_sliced_gemms) ibeg = 1 + (iblock-1)*nprojs_blk
245 25082 : b_cptr = c_loc(b)
246 50164 : call c_f_pointer(b_cptr, b_1d, [cplex * ldb * ndat])
247 :
248 : call abi_xgemm(transa,transb,npw,ndat,nprojs,alpha,&
249 : & a,lda,&
250 : & b_1d(cplex*ibeg-(cplex-1):cplex*ldb*ndat),ldb,beta,&
251 : & c,ldc,&
252 25082 : & x_cplx=cplex,gpu_option=gpu_option)
253 : end if
254 :
255 25082 : end subroutine opernlb_xgemm
256 :
257 : !----------------------------------------------------------------------
258 :
259 : !!****f* m_opernlb_gemm/opernlb_gemm
260 : !! NAME
261 : !! opernlb_gemm
262 : !!
263 : !! FUNCTION
264 : !! For a given wave-function |c>, get all projected scalars
265 : !! <p_lmn|c> where |p_lmn> are non-local projectors
266 : !! With:
267 : !! <p_lmn|c>=4pi/sqrt(vol) (i)^l Sum_g[c(g).f_nl(g).Y_lm(g).exp(2pi.i.g.R)]
268 : !!
269 : !! INPUTS
270 : !! choice=chooses possible output:
271 : !! if choice>=0: compute projected scalars
272 : !! if choice<0: same as choice>0 but use already computed projected scalars
273 : !! if ABS(choice)>1, then compute additional quantities:
274 : !! 2: compute projected scalars and derivatives wrt atm pos.
275 : !! 3: compute projected scalars and derivatives wrt strains
276 : !! 22: compute projected scalars and 2nd derivatives wrt atm pos. and q-vector.
277 : !! 23: compute projected scalars, derivatives wrt atm pos. and derivatives wrt strains
278 : !! 25: compute projected scalars and 3rd derivatives wrt atm pos. and two q-vectors.
279 : !! 4, 24: compute projected scalars, derivatives wrt atm pos.
280 : !! and 2nd derivatives wrt atm pos.
281 : !! 33: compute projected scalars and 2nd derivatives wrt strain and q-vector.
282 : !! 5,51,52: compute projected scalars and derivatives wrt wave vector k
283 : !! 53: compute projected scalars and derivatives wrt wave vector k in direction idir+1 and idir+2 mod 3
284 : !! 54: compute projected scalars, deriv. wrt atm pos., deriv. wrt wave vector k
285 : !! and 2nd derivatives wrt right wave vector k and atm pos.
286 : !! 55: compute projected scalars, deriv. strains, deriv. wrt wave vector k
287 : !! and 2nd derivatives wrt right wave vector k and strain
288 : !! 6: compute projected scalars, derivatives wrt atm pos., derivatives wrt strains,
289 : !! 2nd derivatives wrt 2 strains and derivatives wrt strain and atm pos.
290 : !! 7: not available
291 : !! 8: compute projected scalars, derivatives wrt wave vector k
292 : !! and 2nd derivatives wrt 2 wave vectors k
293 : !! cplex=1 if <p_lmn|c> scalars are real or pure imaginary (equivalent to istwfk>1)
294 : !! 2 if <p_lmn|c> scalars are complex
295 : !! dimffnl=second dimension of ffnl
296 : !! ffnl(npw,dimffnl,nlmn)= nonlocal quantities containing nonlocal form factors
297 : !! ia3=gives the number of the first atom in the subset presently treated
298 : !! idir=direction of the - atom to be moved in the case (choice=2,signs=2) or (choice=22,signs=2)
299 : !! - k point direction in the case (choice=5,signs=2)
300 : !! - strain component (1:6) in the case (choice=3,signs=2) or (choice=6,signs=1)
301 : !! - strain component (1:9) in the case (choice=33,signs=2)
302 : !! - (1:9) components to specify the atom to be moved and the second q-gradient
303 : !! direction in the case (choice=25,signs=2)
304 : !! indlmn(6,nlmn)= array giving l,m,n,lm,ln,s for i=lmn
305 : !! istwf_k=option parameter that describes the storage of wfs
306 : !! kpg(npw,nkpg)=(k+G) components for ikpg=1...3 (if nkpg=3 or 9)
307 : !! [(k+G)_a].[(k+G)_b] quantities for ikpg=4...9 (if nkpg=9)
308 : !! (k+G) Cartesian components for choice==33
309 : !! matblk=dimension of the array ph3d
310 : !! mpi_enreg=information about MPI parallelization
311 : !! ndgxdt=second dimension of dgxdt
312 : !! nd2gxdt=second dimension of d2gxdt
313 : !! nincat=number of atoms in the subset here treated
314 : !! nkpg=second dimension of array kpg (0, 3 or 9)
315 : !! nlmn=number of (l,m,n) numbers for current type of atom
316 : !! nloalg(3)=governs the choice of the algorithm for non-local operator.
317 : !! npw=number of plane waves in reciprocal space
318 : !! nspinor=number of spinorial components of the wavefunctions (on current proc)
319 : !! ph3d(2,npw,matblk)=three-dimensional phase factors
320 : !! [qdir]= optional, direction of the q-gradient (only for choice=22 choice=25 and choice=33)
321 : !! signs=chooses possible output:
322 : !! signs=1: compute derivatives in all directions
323 : !! signs=2: compute derivative in direction IDIR only
324 : !! compatible only with 1st-order derivatives and "single" derivatives
325 : !! ucvol=unit cell volume (bohr^3)
326 : !! vect(2,npw*my_nspinor)=starting vector in reciprocal space
327 : !!
328 : !! OUTPUT
329 : !! if (choice>1) dgxdt(cplex,ndgxdt,nlmn,nincat,nspinor)=
330 : !! gradients of projected scalars wrt coords (choice=2, 23, 4, 54, 6)
331 : !! wrt strains (choice=3, 23, 55)
332 : !! wrt k wave vect. (choice=5, 51, 52, 53, 54, 55, 8)
333 : !! wrt coords and q vect (choice=22)
334 : !! wrt coords and two q vects (choice=25)
335 : !! wrt strains and q vect (choice=33)
336 : !! if (choice=4, 24, 33, 54, 55, 6, 8) d2gxdt(cplex,nd2gxdt,nlmn,nincat,nspinor)=
337 : !! 2nd grads of projected scalars wrt 2 coords (choice=4 or 24)
338 : !! wrt coords & k wave vect. (choice=54)
339 : !! wrt strains & k wave vect. (choice=55)
340 : !! wrt coords & strains (choice=6)
341 : !! wrt 2 strains (choice=6)
342 : !! wrt 2 k wave vect. (choice=8)
343 : !! wrt strains and q vect (choice=33)
344 : !! only compatible with signs=1
345 : !! cplex_dgxdt(ndgxdt) = used only when cplex = 1
346 : !! cplex_dgxdt(i) = 1 if dgxdt(1,i,:,:) is real, 2 if it is pure imaginary
347 : !! cplex_d2gxdt(nd2gxdt) = used only when cplex = 1
348 : !! cplex_d2gxdt(i) = 1 if d2gxdt(1,i,:,:) is real, 2 if it is pure imaginary
349 : !!
350 : !! SIDE EFFECTS
351 : !! gx(cplex,nlmn,nincat,nspinor)= projected scalars - input if choice<0, output if choice>=0
352 : !!
353 : !! NOTES
354 : !! 1-The openMP version is different from the standard version:
355 : !! the standard version is more effifient on one CPU core.
356 : !! 2-Operate for one type of atom, and within this given type of atom,
357 : !! for a subset of at most nincat atoms.
358 : !!
359 : !! SOURCE
360 10902 : subroutine opernlb_gemm(choice,cplex,cplex_dgxdt,cplex_d2gxdt,cplex_fac,&
361 32706 : & d2gxdtfac,d2gxdtfac_sij,dgxdtfac,dgxdtfac_sij,&
362 21804 : & dimffnl,ffnl,gxfac,gxfac_sij,&
363 10902 : & idir,indlmn,kpg,matblk,istwf_k,&
364 : & nd2gxdt,nd2gxdtfac,ndgxdt,ndgxdtfac,&
365 10902 : & nkpg,npw,nspinor,signs,ucvol,ndat,ntypat,lmnmax,nattyp,&
366 : & is_kprime,iatom_only,atom_proj_shift,&
367 10902 : & paw_opt,ph3d,&
368 : & nprojs,&
369 21804 : & vectin,vectout,svectout,&
370 11232 : & temp_realvec_r,temp_realvec_i,&
371 : & gpu_option,use_distrib)
372 :
373 : !Arguments ------------------------------------
374 : !scalars
375 : integer,intent(in) :: choice,cplex,cplex_fac,idir,istwf_k,nd2gxdt,nd2gxdtfac
376 : integer,intent(in) :: ndgxdt,dimffnl,nkpg,lmnmax,ntypat,ndgxdtfac,matblk,npw,nspinor
377 : integer,intent(in) :: paw_opt,signs,ndat,iatom_only,atom_proj_shift
378 : integer,intent(in) :: nprojs
379 : real(dp),intent(in) :: ucvol
380 : integer,intent(in) :: gpu_option
381 : logical,intent(in) :: use_distrib,is_kprime
382 : !arrays
383 : integer,intent(in) :: indlmn(6,lmnmax,ntypat),nattyp(ntypat)
384 : integer,intent(in) :: cplex_dgxdt(ndgxdt),cplex_d2gxdt(nd2gxdt)
385 : real(dp),intent(in) :: ffnl(npw,dimffnl,lmnmax,ntypat),kpg(npw,nkpg)
386 : real(dp),intent(in) :: ph3d(2,npw,matblk)
387 : real(dp),target,intent(in) :: vectin(:,:)
388 : real(dp),target,intent(inout) :: vectout(:,:),svectout(:,:)
389 : real(dp),target,intent(in) :: d2gxdtfac(cplex_fac,nd2gxdtfac,nprojs,ndat*nspinor)
390 : real(dp),target,intent(in) :: d2gxdtfac_sij(cplex,nd2gxdtfac,nprojs,ndat*nspinor)
391 : real(dp),target,intent(inout) :: dgxdtfac(cplex_fac,ndgxdtfac*nprojs,ndat*nspinor)
392 : real(dp),target,intent(inout) :: dgxdtfac_sij(cplex,ndgxdtfac*nprojs,ndat*nspinor)
393 : real(dp),target,intent(in) :: gxfac(cplex_fac,nprojs,ndat*nspinor)
394 : real(dp),target,intent(in) :: gxfac_sij(cplex,nprojs,ndat*nspinor)
395 : real(dp),target,intent(out) :: temp_realvec_r(:),temp_realvec_i(:)
396 :
397 : !Local variables-------------------------------
398 : integer :: idat,i,ik,nprojs_all,iproj,iplex
399 : integer :: projs_beg,projs_end,dprojs_beg,dprojs_end
400 : integer :: nprojs_blk,nprojs_last_blk,nprojs_cur_blk,rank,nprocs,iblock,nblocks
401 : logical :: use_sliced_gemms
402 : complex(dp) :: beta
403 10902 : real(dp), ABI_CONTIGUOUS pointer :: projs(:,:,:),projs_r(:,:,:),projs_i(:,:,:)
404 10902 : real(dp), ABI_CONTIGUOUS pointer :: dprojs(:,:,:),dprojs_r(:,:,:),dprojs_i(:,:,:)
405 :
406 0 : ik=1; if(is_kprime) ik=2
407 : #ifndef HAVE_OPENMP_OFFLOAD
408 : ABI_UNUSED((/iproj,idat,iplex/))
409 : #endif
410 : ABI_UNUSED(cplex_dgxdt)
411 : ABI_UNUSED(cplex_d2gxdt)
412 : ABI_UNUSED(d2gxdtfac)
413 : ABI_UNUSED(d2gxdtfac_sij)
414 :
415 10902 : nprojs_all=nprojs
416 10902 : if(iatom_only>0) then
417 0 : nprojs_all=0
418 0 : do i=1,ntypat
419 0 : nprojs_all = nprojs_all + count(indlmn(3,:,i)>0)*nattyp(i)
420 : end do
421 : end if
422 10902 : nprojs_last_blk=nprojs_all
423 10902 : iblock=1; nblocks=1
424 10902 : use_sliced_gemms=.false.
425 10902 : if(gemm_nonlop_block_size>1 .and. .not. use_distrib) then
426 0 : nblocks=gemm_nonlop_block_size
427 0 : use_sliced_gemms=.true.
428 : end if
429 :
430 : call refresh_projectors(npw,istwf_k,nprojs_all,ndgxdt,nd2gxdt,&
431 10902 : & is_kprime,gpu_option)
432 10902 : if(nprojs_all/=gemm_nonlop_kpt(ik)%nprojs) ABI_BUG("Problem")
433 10902 : nprojs_blk = nprojs
434 10902 : nprojs_last_blk = nprojs
435 10902 : if(use_distrib) then
436 0 : rank=xmpi_comm_rank(gemm_nonlop_block_comm); nprocs=xmpi_comm_size(gemm_nonlop_block_comm)
437 0 : nprojs_blk = gemm_nonlop_kpt(ik)%nprojs_blk
438 0 : nprojs_last_blk = gemm_nonlop_kpt(ik)%nprojs_last_blk
439 0 : iblock=rank+1
440 10902 : else if(gemm_nonlop_block_size>1) then
441 0 : nprojs_blk = nprojs / gemm_nonlop_block_size
442 0 : nprojs_last_blk = nprojs_blk + modulo(nprojs,nprojs_blk)
443 : end if
444 :
445 10902 : projs_beg=1; projs_end=nprojs;
446 10902 : dprojs_beg=1; dprojs_end=max(1,nprojs*ndgxdt)
447 10902 : if((choice==2 .and. signs==2)) then
448 0 : projs_beg=atom_proj_shift+1
449 0 : projs_end=projs_beg+nprojs-1
450 0 : dprojs_beg=atom_proj_shift*ndgxdt+1
451 0 : dprojs_end=dprojs_beg+nprojs*ndgxdt-1
452 : end if
453 :
454 10902 : if(gemm_nonlop_block_size>1) then
455 0 : projs_beg=1; projs_end=nprojs_last_blk;
456 0 : dprojs_beg=1; dprojs_end=max(1,nprojs_last_blk*ndgxdt)
457 : end if
458 :
459 10902 : if(istwf_k == 1) then
460 7689 : projs => gemm_nonlop_kpt(ik)%projs(:,:,projs_beg:projs_end)
461 7689 : if(ndgxdt>0) dprojs => gemm_nonlop_kpt(ik)%dprojs(:,:,dprojs_beg:dprojs_end)
462 : else
463 3213 : projs_r => gemm_nonlop_kpt(ik)%projs_r(:,:,projs_beg:projs_end)
464 3213 : projs_i => gemm_nonlop_kpt(ik)%projs_i(:,:,projs_beg:projs_end)
465 : if(ndgxdt>0) dprojs_r => gemm_nonlop_kpt(ik)%dprojs_r(:,:,dprojs_beg:dprojs_end)
466 : if(ndgxdt>0) dprojs_i => gemm_nonlop_kpt(ik)%dprojs_i(:,:,dprojs_beg:dprojs_end)
467 : end if
468 :
469 10902 : if(gemm_nonlop_kpt(ik)%ikpt/=gemm_nonlop_ikpt_this_proc_being_treated .or. use_sliced_gemms) then
470 : call prep_projectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
471 : & ucvol,ffnl,ph3d,dimffnl,matblk,&
472 0 : & nprojs_last_blk,is_kprime,gpu_option,iblock)
473 0 : gemm_nonlop_kpt(ik)%ikpt=gemm_nonlop_ikpt_this_proc_being_treated
474 : end if
475 10902 : if(choice>1 .and. ndgxdt>0) then
476 : if( nd2gxdt/=gemm_nonlop_kpt(ik)%ngrads2 &
477 : & .or. ndgxdt/=gemm_nonlop_kpt(ik)%ngrads &
478 : & .or. choice/=gemm_nonlop_kpt(ik)%choice &
479 : & .or. idir/=gemm_nonlop_kpt(ik)%idir &
480 0 : & .or. use_sliced_gemms) then
481 : call prep_dprojectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
482 : & ucvol,ffnl,ph3d,kpg,nkpg,dimffnl,matblk,&
483 : & nprojs_last_blk,ndgxdt,nd2gxdt,choice,signs,idir,&
484 0 : & is_kprime,gpu_option,iblock)
485 0 : gemm_nonlop_kpt(ik)%choice = choice
486 0 : gemm_nonlop_kpt(ik)%idir = idir
487 : end if
488 : end if
489 :
490 10902 : if(paw_opt == 3 .or. paw_opt == 4) then
491 :
492 : ! Get svectout from gxfac_sij
493 8822 : if(cplex == 2) then
494 :
495 : ! With many blocks, GEMM results will be summed using beta=cone
496 6101 : beta = czero
497 :
498 12202 : do i=1,nblocks
499 6101 : if(use_sliced_gemms .and. i>1) then
500 : call prep_projectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
501 : & ucvol,ffnl,ph3d,dimffnl,matblk,&
502 0 : & nprojs_last_blk,is_kprime,gpu_option,i)
503 0 : gemm_nonlop_kpt(ik)%ikpt=gemm_nonlop_ikpt_this_proc_being_treated
504 0 : if(choice>1 .and. ndgxdt>0) then
505 : call prep_dprojectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
506 : & ucvol,ffnl,ph3d,kpg,nkpg,dimffnl,matblk,&
507 : & nprojs_last_blk,ndgxdt,nd2gxdt,choice,signs,idir,&
508 0 : & is_kprime,gpu_option,i)
509 0 : gemm_nonlop_kpt(ik)%choice = choice
510 0 : gemm_nonlop_kpt(ik)%idir = idir
511 : end if
512 : end if
513 :
514 6101 : nprojs_cur_blk=nprojs
515 6101 : if(use_sliced_gemms) then
516 0 : if(i<nblocks) then
517 0 : nprojs_cur_blk=nprojs_blk
518 : else
519 0 : nprojs_cur_blk=nprojs_last_blk
520 : end if
521 : end if
522 :
523 6101 : if(choice==1 .or. choice==7) then
524 : call opernlb_xgemm(cplex, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
525 : & projs, npw,&
526 : & gxfac_sij, nprojs, beta, svectout, npw,&
527 : & rank, nprocs,&
528 : & nprojs_blk, nprojs_last_blk, i,&
529 6101 : & gpu_option, use_distrib, use_sliced_gemms)
530 : else if(choice==2) then
531 : call opernlb_xgemm(cplex, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
532 : & dprojs, npw, &
533 : & gxfac_sij, nprojs, beta, svectout, npw,&
534 : & rank, nprocs,&
535 : & nprojs_blk, nprojs_last_blk, i,&
536 0 : & gpu_option, use_distrib, use_sliced_gemms)
537 : call opernlb_xgemm(cplex, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
538 : & projs, npw, &
539 : & dgxdtfac_sij, nprojs, cone, svectout, npw,&
540 : & rank, nprocs,&
541 : & nprojs_blk, nprojs_last_blk, i,&
542 0 : & gpu_option, use_distrib, use_sliced_gemms)
543 : else if(choice==3) then
544 0 : if(idir<=3) then
545 0 : if(gpu_option == ABI_GPU_DISABLED) then
546 0 : dgxdtfac_sij(:,:,:) = dgxdtfac_sij(:,:,:) - gxfac_sij(:,:,:)
547 : else if(gpu_option == ABI_GPU_OPENMP) then
548 : #ifdef HAVE_OPENMP_OFFLOAD
549 : !$OMP TARGET TEAMS DISTRIBUTE &
550 : !$OMP& MAP(to:gxfac_sij,dgxdtfac_sij) PRIVATE(idat)
551 : do idat=1,ndat
552 : !$OMP PARALLEL DO PRIVATE(iproj,iplex) COLLAPSE(2)
553 : do iproj=1,nprojs
554 : do iplex=1,cplex
555 : dgxdtfac_sij(iplex,iproj,idat) = dgxdtfac_sij(iplex,iproj,idat) - gxfac_sij(iplex,iproj,idat)
556 : end do
557 : end do
558 : end do
559 : #endif
560 : end if
561 : end if
562 : call opernlb_xgemm(cplex, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
563 : & projs, npw, &
564 : & dgxdtfac_sij, nprojs, beta, svectout, npw,&
565 : & rank, nprocs,&
566 : & nprojs_blk, nprojs_last_blk, i,&
567 0 : & gpu_option, use_distrib, use_sliced_gemms)
568 : call opernlb_xgemm(cplex, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
569 : & dprojs, npw, &
570 : & gxfac_sij, nprojs, cone, svectout, npw,&
571 : & rank, nprocs,&
572 : & nprojs_blk, nprojs_last_blk, i,&
573 0 : & gpu_option, use_distrib, use_sliced_gemms)
574 : else if(choice==5) then
575 : call opernlb_xgemm(cplex, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
576 : & projs, npw, &
577 : & dgxdtfac_sij, nprojs, beta, svectout, npw,&
578 : & rank, nprocs,&
579 : & nprojs_blk, nprojs_last_blk, i,&
580 0 : & gpu_option, use_distrib, use_sliced_gemms)
581 : call opernlb_xgemm(cplex, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
582 : & dprojs, npw, &
583 : & gxfac_sij, nprojs, cone, svectout, npw,&
584 : & rank, nprocs,&
585 : & nprojs_blk, nprojs_last_blk, i,&
586 0 : & gpu_option, use_distrib, use_sliced_gemms)
587 : else if(choice==51) then
588 : call opernlb_xgemm(cplex, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
589 : & projs, npw, &
590 : & dgxdtfac_sij, nprojs, beta, svectout, npw,&
591 : & rank, nprocs,&
592 : & nprojs_blk, nprojs_last_blk, i,&
593 0 : & gpu_option, use_distrib, use_sliced_gemms)
594 : end if
595 :
596 12202 : beta = cone
597 : end do
598 : else
599 :
600 : ! With many blocks, GEMM results will be summed using beta=cone
601 2721 : beta = czero
602 :
603 5442 : do i=1,nblocks
604 2721 : if(use_sliced_gemms .and. i>1) then
605 : call prep_projectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
606 : & ucvol,ffnl,ph3d,dimffnl,matblk,&
607 0 : & nprojs_last_blk,is_kprime,gpu_option,i)
608 0 : gemm_nonlop_kpt(ik)%ikpt=gemm_nonlop_ikpt_this_proc_being_treated
609 0 : if(choice>1 .and. ndgxdt>0) then
610 : call prep_dprojectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
611 : & ucvol,ffnl,ph3d,kpg,nkpg,dimffnl,matblk,&
612 : & nprojs_last_blk,ndgxdt,nd2gxdt,choice,signs,idir,&
613 0 : & is_kprime,gpu_option,i)
614 0 : gemm_nonlop_kpt(ik)%choice = choice
615 0 : gemm_nonlop_kpt(ik)%idir = idir
616 : end if
617 : end if
618 :
619 2721 : nprojs_cur_blk=nprojs
620 2721 : if(use_sliced_gemms) then
621 0 : if(i<nblocks) then
622 0 : nprojs_cur_blk=nprojs_blk
623 : else
624 0 : nprojs_cur_blk=nprojs_last_blk
625 : end if
626 : end if
627 :
628 : call opernlb_xgemm(cplex, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
629 : & projs_r, npw, &
630 : & gxfac_sij, nprojs, beta, temp_realvec_r, npw,&
631 : & rank, nprocs,&
632 : & nprojs_blk, nprojs_last_blk, i,&
633 2721 : & gpu_option, use_distrib, use_sliced_gemms)
634 : call opernlb_xgemm(cplex, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
635 : & projs_i, npw,&
636 : & gxfac_sij, nprojs, beta, temp_realvec_i, npw,&
637 : & rank, nprocs,&
638 : & nprojs_blk, nprojs_last_blk, i,&
639 2721 : & gpu_option, use_distrib, use_sliced_gemms)
640 :
641 5442 : beta=cone
642 :
643 : end do
644 :
645 2721 : if(gpu_option == ABI_GPU_DISABLED) then
646 2919694 : svectout(1,1:npw*nspinor*ndat) = temp_realvec_r(1:npw*nspinor*ndat)
647 2919694 : svectout(2,1:npw*nspinor*ndat) = temp_realvec_i(1:npw*nspinor*ndat)
648 : else if(gpu_option == ABI_GPU_OPENMP) then
649 : #ifdef HAVE_OPENMP_OFFLOAD
650 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO &
651 : !$OMP& MAP(to:temp_realvec_r,temp_realvec_i,svectout) PRIVATE(i)
652 : do i=1, npw*nspinor*ndat
653 : svectout(1,i) = temp_realvec_r(i)
654 : svectout(2,i) = temp_realvec_i(i)
655 : end do
656 : #endif
657 : end if
658 :
659 : end if ! cplex = 2
660 :
661 8822 : if(choice /= 7 .and. choice /= 5 .and. choice/=51 .and. choice/=2 .and. choice/=3) then
662 8246 : if(gpu_option == ABI_GPU_DISABLED) then
663 47661172 : svectout = svectout + vectin ! TODO understand this
664 0 : else if(gpu_option == ABI_GPU_OPENMP) then
665 : call abi_xaxpy( 2*npw*nspinor*ndat, cone, &
666 0 : & vectin, 1, svectout, 1, x_cplx=1, gpu_option=gpu_option)
667 : end if
668 : end if
669 : end if ! (paw_opt == 3 .or. paw_opt == 4)
670 :
671 10902 : if(paw_opt == 0 .or. paw_opt == 1 .or. paw_opt == 4) then
672 : ! Get vectout from gxfac
673 10326 : if(cplex_fac == 2) then
674 :
675 : ! With many blocks, GEMM results will be summed using beta=cone
676 7113 : beta = czero
677 :
678 14226 : do i=1,nblocks
679 7113 : if(use_sliced_gemms .and. (i>1 .or. paw_opt == 3 .or. paw_opt == 4)) then
680 : call prep_projectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
681 : & ucvol,ffnl,ph3d,dimffnl,matblk,&
682 0 : & nprojs_last_blk,is_kprime,gpu_option,i)
683 0 : gemm_nonlop_kpt(ik)%ikpt=gemm_nonlop_ikpt_this_proc_being_treated
684 0 : if(choice>1 .and. ndgxdt>0) then
685 : call prep_dprojectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
686 : & ucvol,ffnl,ph3d,kpg,nkpg,dimffnl,matblk,&
687 : & nprojs_last_blk,ndgxdt,nd2gxdt,choice,signs,idir,&
688 0 : & is_kprime,gpu_option,i)
689 0 : gemm_nonlop_kpt(ik)%choice = choice
690 0 : gemm_nonlop_kpt(ik)%idir = idir
691 : end if
692 : end if
693 :
694 7113 : nprojs_cur_blk=nprojs
695 7113 : if(use_sliced_gemms) then
696 0 : if(i<nblocks) then
697 0 : nprojs_cur_blk=nprojs_blk
698 : else
699 0 : nprojs_cur_blk=nprojs_last_blk
700 : end if
701 : end if
702 :
703 7113 : if(choice==1 .or. choice==7) then
704 : call opernlb_xgemm(cplex_fac, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
705 : & projs, npw, &
706 : & gxfac, nprojs, beta, vectout, npw,&
707 : & rank, nprocs,&
708 : & nprojs_blk, nprojs_last_blk, i,&
709 7113 : & gpu_option, use_distrib, use_sliced_gemms)
710 : else if(choice==2) then
711 : call opernlb_xgemm(cplex_fac, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
712 : & dprojs, npw, &
713 : & gxfac, nprojs, beta, vectout, npw,&
714 : & rank, nprocs,&
715 : & nprojs_blk, nprojs_last_blk, i,&
716 0 : & gpu_option, use_distrib, use_sliced_gemms)
717 : call opernlb_xgemm(cplex_fac, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
718 : & projs, npw, &
719 : & dgxdtfac, nprojs, cone, vectout, npw,&
720 : & rank, nprocs,&
721 : & nprojs_blk, nprojs_last_blk, i,&
722 0 : & gpu_option, use_distrib, use_sliced_gemms)
723 : else if(choice==3) then
724 0 : if(idir<=3) then
725 0 : if(gpu_option == ABI_GPU_DISABLED) then
726 0 : dgxdtfac(:,:,:) = dgxdtfac(:,:,:) - gxfac(:,:,:)
727 : else if(gpu_option == ABI_GPU_OPENMP) then
728 : #ifdef HAVE_OPENMP_OFFLOAD
729 : !$OMP TARGET TEAMS DISTRIBUTE &
730 : !$OMP& MAP(to:gxfac,dgxdtfac) PRIVATE(idat)
731 : do idat=1,ndat
732 : !$OMP PARALLEL DO PRIVATE(iproj,iplex) COLLAPSE(2)
733 : do iproj=1,nprojs
734 : do iplex=1,cplex_fac
735 : dgxdtfac(iplex,iproj,idat) = dgxdtfac(iplex,iproj,idat) - gxfac(iplex,iproj,idat)
736 : end do
737 : end do
738 : end do
739 : #endif
740 : end if
741 : end if
742 : call opernlb_xgemm(cplex_fac, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
743 : & projs, npw, &
744 : & dgxdtfac, nprojs, beta, vectout, npw,&
745 : & rank, nprocs,&
746 : & nprojs_blk, nprojs_last_blk, i,&
747 0 : & gpu_option, use_distrib, use_sliced_gemms)
748 : call opernlb_xgemm(cplex_fac, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
749 : & dprojs, npw, &
750 : & gxfac, nprojs, cone, vectout, npw,&
751 : & rank, nprocs,&
752 : & nprojs_blk, nprojs_last_blk, i,&
753 0 : & gpu_option, use_distrib, use_sliced_gemms)
754 : else if(choice==5) then
755 : call opernlb_xgemm(cplex_fac, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
756 : & projs, npw, &
757 : & dgxdtfac, nprojs, beta, vectout, npw,&
758 : & rank, nprocs,&
759 : & nprojs_blk, nprojs_last_blk, i,&
760 0 : & gpu_option, use_distrib, use_sliced_gemms)
761 : call opernlb_xgemm(cplex_fac, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
762 : & dprojs, npw, &
763 : & gxfac, nprojs, cone, vectout, npw,&
764 : & rank, nprocs,&
765 : & nprojs_blk, nprojs_last_blk, i,&
766 0 : & gpu_option, use_distrib, use_sliced_gemms)
767 : else if(choice==51) then
768 : call opernlb_xgemm(cplex_fac, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
769 : & projs, npw, &
770 : & dgxdtfac, nprojs, beta, vectout, npw,&
771 : & rank, nprocs,&
772 : & nprojs_blk, nprojs_last_blk, i,&
773 0 : & gpu_option, use_distrib, use_sliced_gemms)
774 : end if
775 14226 : beta = cone
776 : end do
777 : else
778 :
779 : ! With many blocks, GEMM results will be summed using beta=cone
780 3213 : beta = czero
781 :
782 6426 : do i=1,nblocks
783 3213 : if(use_sliced_gemms .and. (i>1 .or. paw_opt == 3 .or. paw_opt == 4)) then
784 : call prep_projectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
785 : & ucvol,ffnl,ph3d,dimffnl,matblk,&
786 0 : & nprojs_last_blk,is_kprime,gpu_option,i)
787 0 : gemm_nonlop_kpt(ik)%ikpt=gemm_nonlop_ikpt_this_proc_being_treated
788 0 : if(choice>1 .and. ndgxdt>0) then
789 : call prep_dprojectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
790 : & ucvol,ffnl,ph3d,kpg,nkpg,dimffnl,matblk,&
791 : & nprojs_last_blk,ndgxdt,nd2gxdt,choice,signs,idir,&
792 0 : & is_kprime,gpu_option,i)
793 0 : gemm_nonlop_kpt(ik)%choice = choice
794 0 : gemm_nonlop_kpt(ik)%idir = idir
795 : end if
796 : end if
797 :
798 3213 : nprojs_cur_blk=nprojs
799 3213 : if(use_sliced_gemms) then
800 0 : if(i<nblocks) then
801 0 : nprojs_cur_blk=nprojs_blk
802 : else
803 0 : nprojs_cur_blk=nprojs_last_blk
804 : end if
805 : end if
806 :
807 : call opernlb_xgemm(cplex_fac, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
808 : & projs_r, npw, &
809 : & gxfac, nprojs, beta, temp_realvec_r, npw,&
810 : & rank, nprocs,&
811 : & nprojs_blk, nprojs_last_blk, i,&
812 3213 : & gpu_option, use_distrib, use_sliced_gemms)
813 : call opernlb_xgemm(cplex_fac, 'N', 'N', npw, ndat*nspinor, nprojs_cur_blk, cone, &
814 : & projs_i, npw, &
815 : & gxfac, nprojs, beta, temp_realvec_i, npw,&
816 : & rank, nprocs,&
817 : & nprojs_blk, nprojs_last_blk, i,&
818 3213 : & gpu_option, use_distrib, use_sliced_gemms)
819 :
820 6426 : beta=cone
821 :
822 : end do
823 :
824 3213 : if(gpu_option == ABI_GPU_DISABLED) then
825 18660350 : vectout(1,1:npw*nspinor*ndat) = temp_realvec_r(1:npw*nspinor*ndat)
826 18660350 : vectout(2,1:npw*nspinor*ndat) = temp_realvec_i(1:npw*nspinor*ndat)
827 : else if(gpu_option == ABI_GPU_OPENMP) then
828 : #ifdef HAVE_OPENMP_OFFLOAD
829 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO &
830 : !$OMP& MAP(to:temp_realvec_r,temp_realvec_i,vectout) PRIVATE(i)
831 : do i=1, npw*nspinor*ndat
832 : vectout(1,i) = temp_realvec_r(i)
833 : vectout(2,i) = temp_realvec_i(i)
834 : end do
835 : #endif
836 : end if
837 : end if ! cplex_fac == 2
838 : end if ! (paw_opt == 0 .or. paw_opt == 1 .or. paw_opt == 4)
839 :
840 10902 : end subroutine opernlb_gemm
841 : !!***
842 :
843 : end module m_opernlb_gemm
844 : !!***
|