Line data Source code
1 : !!****m* ABINIT/m_opernla_gemm
2 : !! NAME
3 : !! m_opernla_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_opernla_gemm
22 :
23 : use defs_basis
24 : use m_abicore
25 : use m_errors
26 : USE_MPI
27 : use m_xmpi
28 : use m_gputk
29 : use m_abi_linalg
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 :: opernla_gemm
43 :
44 : contains
45 : !!***
46 :
47 :
48 : !!****f* m_opernla_gemm/opernla_gemm_distributed
49 : !! NAME
50 : !! opernla_gemm_distributed
51 : !!
52 : !! FUNCTION
53 : !! Distributed version of "opernla" GEMM called in gemm_nonlop.
54 : !!
55 : !! INPUTS
56 : !!
57 : !! SOURCE
58 0 : subroutine opernla_gemm_distributed(rank,nprocs,npw,ndat,&
59 : & transa,transb,&
60 : & nprojs,nprojs_blk,nprojs_last_blk,cplex,beta,&
61 0 : & projs_local,vectin,projections,gpu_option)
62 : integer, intent(in) :: rank,nprocs,npw,ndat,gpu_option
63 : integer, intent(in) :: nprojs,nprojs_blk,nprojs_last_blk,cplex
64 : character(len=1),intent(in) :: transa,transb
65 : complex(dp), intent(in) :: beta
66 : real(dp), intent(in), target :: projs_local(cplex,npw,nprojs_last_blk)
67 : real(dp), intent(in), target :: vectin(2,npw,ndat)
68 : real(dp), intent(out), target :: projections(cplex,nprojs,ndat)
69 :
70 : !Local variables
71 : integer :: iblock,ibeg,req(2),ierr,nprojs_cur_blk,rank_prev,rank_next
72 0 : real(dp), ABI_CONTIGUOUS pointer :: recv_buf(:,:,:), work_buf(:,:,:)
73 0 : real(dp), allocatable, target :: projs_recv(:,:,:)
74 0 : real(dp), ABI_CONTIGUOUS pointer :: projections_1d(:)
75 : type(c_ptr) :: projections_cptr
76 : ! *************************************************************************
77 :
78 0 : ABI_MALLOC(projs_recv, (cplex, npw, nprojs_last_blk))
79 : #ifdef HAVE_OPENMP_OFFLOAD
80 : !$OMP TARGET ENTER DATA MAP(alloc:projs_recv) IF(gpu_option==ABI_GPU_OPENMP)
81 : #endif
82 :
83 0 : rank_next=modulo(rank + 1,nprocs)
84 0 : rank_prev=rank - 1
85 0 : if(rank_prev == -1) rank_prev = nprocs - 1
86 :
87 0 : do iblock=1,nprocs
88 :
89 0 : if(rank+iblock == nprocs) then
90 0 : nprojs_cur_blk = nprojs_last_blk
91 : else
92 0 : nprojs_cur_blk = nprojs_blk
93 : end if
94 :
95 0 : if(modulo(iblock,2)==1) then
96 : ! XG20241028 : This coding confused the gnu_8.5 compiler of buda2_gnu_8.5_cuda, wrt the contiguous character of the target.
97 : ! It declared an error. Make it simple !
98 : ! work_buf => projs_local(1:cplex,1:npw,1:nprojs_last_blk)
99 : ! recv_buf => projs_recv(1:cplex,1:npw,1:nprojs_last_blk)
100 0 : work_buf => projs_local
101 0 : recv_buf => projs_recv
102 :
103 : else
104 : ! XG20241028 : Same than above
105 : ! work_buf => projs_recv(1:cplex,1:npw,1:nprojs_last_blk)
106 : ! recv_buf => projs_local(1:cplex,1:npw,1:nprojs_last_blk)
107 0 : work_buf => projs_recv
108 0 : recv_buf => projs_local
109 : end if
110 :
111 0 : if(gpu_option == ABI_GPU_DISABLED) then
112 0 : call xmpi_isend(work_buf,rank_prev,iblock,gemm_nonlop_block_comm,req(1),ierr)
113 0 : call xmpi_irecv(recv_buf,rank_next,iblock,gemm_nonlop_block_comm,req(2),ierr)
114 : else if(gpu_option == ABI_GPU_OPENMP) then
115 : #ifdef HAVE_OPENMP_OFFLOAD
116 : #ifndef HAVE_GPU_MPI
117 :
118 : ! GPU-aware MPI not available : perform MPI comms on CPU
119 : !$OMP TARGET UPDATE FROM(work_buf) if(iblock==1)
120 : call xmpi_isend(work_buf,rank_prev,iblock,gemm_nonlop_block_comm,req(1),ierr)
121 : call xmpi_irecv(recv_buf,rank_next,iblock,gemm_nonlop_block_comm,req(2),ierr)
122 :
123 : #else
124 :
125 : ! GPU-aware MPI available : pass GPU buffers to MPI
126 : call xmpi_isend(work_buf,rank_prev,iblock,gemm_nonlop_block_comm,req(1),ierr,use_omp_map=.true.)
127 : call xmpi_irecv(recv_buf,rank_next,iblock,gemm_nonlop_block_comm,req(2),ierr,use_omp_map=.true.)
128 :
129 : #endif
130 : #endif
131 : end if
132 :
133 0 : ibeg = 1 + modulo(rank+iblock-1,nprocs)*nprojs_blk
134 :
135 : ! Small trickery here:
136 : ! The multiplication is performed over a slice of projectors in 'work_buf' matrix.
137 : ! In that case, 'projections' matrix is still sized by all projectors,
138 : ! and the result is a non-contiguous slice for the projectors contained in work_buf.
139 : !
140 : ! Therefore, we need to provide GEMM with the start of the slice
141 : ! with 'projections' rather than the start of 'projections' as usual.
142 : !
143 : ! For that reason, we turn 'projections' matrix into single-rank in order
144 : ! to pass the right starting row.
145 0 : projections_cptr = c_loc(projections)
146 0 : call c_f_pointer(projections_cptr, projections_1d, [cplex * nprojs * ndat])
147 :
148 : call abi_xgemm(transa,transb,&
149 : & nprojs_cur_blk, ndat, npw, cone,&
150 : & work_buf, npw,&
151 : & vectin, npw, &
152 : & beta, &
153 : & projections_1d(cplex*ibeg-(cplex-1):cplex*nprojs*ndat),nprojs,&
154 0 : & x_cplx=cplex,gpu_option=gpu_option)
155 :
156 :
157 0 : call xmpi_wait(req(1),ierr)
158 0 : call xmpi_wait(req(2),ierr)
159 : !call xmpi_waitall(req,ierr)
160 :
161 : #ifdef HAVE_OPENMP_OFFLOAD
162 : #ifndef HAVE_GPU_MPI
163 : ! If MPI is not GPU-aware, push received data to GPU
164 : !$OMP TARGET UPDATE TO(recv_buf) IF(gpu_option==ABI_GPU_OPENMP)
165 : #endif
166 : #endif
167 :
168 : end do
169 :
170 0 : if(modulo(iblock,2)==1) then
171 0 : if(gpu_option == ABI_GPU_DISABLED) then
172 0 : call DCOPY(cplex*npw*nprojs_cur_blk, recv_buf, 1, work_buf, 1)
173 : else if(gpu_option == ABI_GPU_OPENMP) then
174 : #ifdef HAVE_OPENMP_OFFLOAD
175 : !$OMP TARGET DATA USE_DEVICE_ADDR(work_buf,recv_buf)
176 : call copy_gpu_to_gpu(c_loc(work_buf), c_loc(recv_buf), INT(cplex, c_size_t)*npw*nprojs_last_blk*dp)
177 : !$OMP END TARGET DATA
178 : #endif
179 : end if
180 : end if
181 :
182 : #ifdef HAVE_OPENMP_OFFLOAD
183 : !$OMP TARGET EXIT DATA MAP(delete:projs_recv) IF(gpu_option==ABI_GPU_OPENMP)
184 : #endif
185 0 : ABI_FREE(projs_recv)
186 :
187 0 : end subroutine opernla_gemm_distributed
188 : !!***
189 :
190 : !----------------------------------------------------------------------
191 :
192 17945 : subroutine opernla_xgemm(cplex,transa,transb,nprojs,ndat,npw,alpha,a,lda,b,ldb,beta,c,ldc,&
193 : & rank, nprocs,&
194 : & nprojs_blk, nprojs_last_blk, &
195 : & iblock,&
196 : & gpu_option,use_distrib,use_sliced_gemms)
197 :
198 : !Arguments ------------------------------------
199 : integer,intent(in) :: cplex,lda,ldb,ldc,nprojs,ndat,npw,gpu_option
200 : integer,intent(in) :: rank,nprocs,nprojs_blk,nprojs_last_blk
201 : integer,intent(in) :: iblock
202 : logical,intent(in) :: use_distrib,use_sliced_gemms
203 : complex(dp),intent(in) :: alpha,beta
204 : character(len=1),intent(in) :: transa,transb
205 : real(dp),target,intent(in) :: a(cplex,lda,nprojs),b(cplex,ldb,ndat)
206 : real(dp),target,intent(inout) :: c(cplex,ldc,ndat)
207 :
208 : integer :: ibeg
209 17945 : real(dp), ABI_CONTIGUOUS pointer :: c_1d(:)
210 : type(c_ptr) :: c_cptr
211 : ! *********************************************************************
212 :
213 17945 : if(use_distrib) then
214 : call opernla_gemm_distributed(rank,nprocs,npw,ndat,&
215 : & transa,transb,&
216 : & nprojs,&
217 : & nprojs_blk,&
218 : & nprojs_last_blk,&
219 : & cplex,alpha,&
220 : & a,&
221 0 : & b,c,gpu_option)
222 : else
223 : ! Small trickery here:
224 : ! When use_sliced_gemm is on, the multiplication is performed
225 : ! over a slice of A matrix, with nprojs == {nprojs_blk,nprojs_last_blk}.
226 : ! In that case, C matrix is still fully sized (ldc == nprojs_all),
227 : ! and the result is a non-contiguous slice.
228 : ! Therefore, we need to provide GEMM with the start of the slice
229 : ! with C rather than the start of C as usual.
230 : !
231 : ! For that reason, we turn C matrix into single-rank in order
232 : ! to pass the right starting row.
233 : ! No buffer-overflow occurs as ldc (nprojs_all) is higher than
234 : ! nprojs (nprojs_blk or nprojs_last_blk).
235 17945 : ibeg = 1
236 17945 : if(use_sliced_gemms) ibeg = 1 + (iblock-1)*nprojs_blk
237 17945 : c_cptr = c_loc(c)
238 35890 : call c_f_pointer(c_cptr, c_1d, [cplex * ldc * ndat])
239 :
240 : call abi_xgemm(transa,transb,nprojs,ndat,npw,alpha,&
241 : & a,lda,&
242 : & b,ldb,beta,&
243 : & c_1d(cplex*ibeg-(cplex-1):cplex*ldc*ndat),ldc,&
244 17945 : & x_cplx=cplex,gpu_option=gpu_option)
245 : end if
246 :
247 17945 : end subroutine opernla_xgemm
248 :
249 : !----------------------------------------------------------------------
250 :
251 : !!****f* m_opernla_gemm/opernla_gemm
252 : !! NAME
253 : !! opernla_gemm
254 : !!
255 : !! FUNCTION
256 : !! For a given wave-function |c>, get all projected scalars
257 : !! <p_lmn|c> where |p_lmn> are non-local projectors
258 : !! With:
259 : !! <p_lmn|c>=4pi/sqrt(vol) (i)^l Sum_g[c(g).f_nl(g).Y_lm(g).exp(2pi.i.g.R)]
260 : !!
261 : !! INPUTS
262 : !! choice=chooses possible output:
263 : !! if choice>=0: compute projected scalars
264 : !! if choice<0: same as choice>0 but use already computed projected scalars
265 : !! if ABS(choice)>1, then compute additional quantities:
266 : !! 2: compute projected scalars and derivatives wrt atm pos.
267 : !! 3: compute projected scalars and derivatives wrt strains
268 : !! 22: compute projected scalars and 2nd derivatives wrt atm pos. and q-vector.
269 : !! 23: compute projected scalars, derivatives wrt atm pos. and derivatives wrt strains
270 : !! 25: compute projected scalars and 3rd derivatives wrt atm pos. and two q-vectors.
271 : !! 4, 24: compute projected scalars, derivatives wrt atm pos.
272 : !! and 2nd derivatives wrt atm pos.
273 : !! 33: compute projected scalars and 2nd derivatives wrt strain and q-vector.
274 : !! 5,51,52: compute projected scalars and derivatives wrt wave vector k
275 : !! 53: compute projected scalars and derivatives wrt wave vector k in direction idir+1 and idir+2 mod 3
276 : !! 54: compute projected scalars, deriv. wrt atm pos., deriv. wrt wave vector k
277 : !! and 2nd derivatives wrt right wave vector k and atm pos.
278 : !! 55: compute projected scalars, deriv. strains, deriv. wrt wave vector k
279 : !! and 2nd derivatives wrt right wave vector k and strain
280 : !! 6: compute projected scalars, derivatives wrt atm pos., derivatives wrt strains,
281 : !! 2nd derivatives wrt 2 strains and derivatives wrt strain and atm pos.
282 : !! 7: not available
283 : !! 8: compute projected scalars, derivatives wrt wave vector k
284 : !! and 2nd derivatives wrt 2 wave vectors k
285 : !! cplex=1 if <p_lmn|c> scalars are real or pure imaginary (equivalent to istwfk>1)
286 : !! 2 if <p_lmn|c> scalars are complex
287 : !! dimffnl=second dimension of ffnl
288 : !! ffnl(npw,dimffnl,nlmn)= nonlocal quantities containing nonlocal form factors
289 : !! ia3=gives the number of the first atom in the subset presently treated
290 : !! idir=direction of the - atom to be moved in the case (choice=2,signs=2) or (choice=22,signs=2)
291 : !! - k point direction in the case (choice=5,signs=2)
292 : !! - strain component (1:6) in the case (choice=3,signs=2) or (choice=6,signs=1)
293 : !! - strain component (1:9) in the case (choice=33,signs=2)
294 : !! - (1:9) components to specify the atom to be moved and the second q-gradient
295 : !! direction in the case (choice=25,signs=2)
296 : !! indlmn(6,nlmn)= array giving l,m,n,lm,ln,s for i=lmn
297 : !! istwf_k=option parameter that describes the storage of wfs
298 : !! kpg(npw,nkpg)=(k+G) components for ikpg=1...3 (if nkpg=3 or 9)
299 : !! [(k+G)_a].[(k+G)_b] quantities for ikpg=4...9 (if nkpg=9)
300 : !! (k+G) Cartesian components for choice==33
301 : !! matblk=dimension of the array ph3d
302 : !! mpi_enreg=information about MPI parallelization
303 : !! ndgxdt=second dimension of dgxdt
304 : !! nd2gxdt=second dimension of d2gxdt
305 : !! nincat=number of atoms in the subset here treated
306 : !! nkpg=second dimension of array kpg (0, 3 or 9)
307 : !! nlmn=number of (l,m,n) numbers for current type of atom
308 : !! nloalg(3)=governs the choice of the algorithm for non-local operator.
309 : !! npw=number of plane waves in reciprocal space
310 : !! nspinor=number of spinorial components of the wavefunctions (on current proc)
311 : !! ph3d(2,npw,matblk)=three-dimensional phase factors
312 : !! [qdir]= optional, direction of the q-gradient (only for choice=22 choice=25 and choice=33)
313 : !! signs=chooses possible output:
314 : !! signs=1: compute derivatives in all directions
315 : !! signs=2: compute derivative in direction IDIR only
316 : !! compatible only with 1st-order derivatives and "single" derivatives
317 : !! ucvol=unit cell volume (bohr^3)
318 : !! vect(2,npw*my_nspinor)=starting vector in reciprocal space
319 : !!
320 : !! OUTPUT
321 : !! if (choice>1) dgxdt(cplex,ndgxdt,nlmn,nincat,nspinor)=
322 : !! gradients of projected scalars wrt coords (choice=2, 23, 4, 54, 6)
323 : !! wrt strains (choice=3, 23, 55)
324 : !! wrt k wave vect. (choice=5, 51, 52, 53, 54, 55, 8)
325 : !! wrt coords and q vect (choice=22)
326 : !! wrt coords and two q vects (choice=25)
327 : !! wrt strains and q vect (choice=33)
328 : !! if (choice=4, 24, 33, 54, 55, 6, 8) d2gxdt(cplex,nd2gxdt,nlmn,nincat,nspinor)=
329 : !! 2nd grads of projected scalars wrt 2 coords (choice=4 or 24)
330 : !! wrt coords & k wave vect. (choice=54)
331 : !! wrt strains & k wave vect. (choice=55)
332 : !! wrt coords & strains (choice=6)
333 : !! wrt 2 strains (choice=6)
334 : !! wrt 2 k wave vect. (choice=8)
335 : !! wrt strains and q vect (choice=33)
336 : !! only compatible with signs=1
337 : !! cplex_dgxdt(ndgxdt) = used only when cplex = 1
338 : !! cplex_dgxdt(i) = 1 if dgxdt(1,i,:,:) is real, 2 if it is pure imaginary
339 : !! cplex_d2gxdt(nd2gxdt) = used only when cplex = 1
340 : !! cplex_d2gxdt(i) = 1 if d2gxdt(1,i,:,:) is real, 2 if it is pure imaginary
341 : !!
342 : !! SIDE EFFECTS
343 : !! gx(cplex,nlmn,nincat,nspinor)= projected scalars - input if choice<0, output if choice>=0
344 : !!
345 : !! NOTES
346 : !! 1-The openMP version is different from the standard version:
347 : !! the standard version is more effifient on one CPU core.
348 : !! 2-Operate for one type of atom, and within this given type of atom,
349 : !! for a subset of at most nincat atoms.
350 : !!
351 : !! SOURCE
352 13290 : subroutine opernla_gemm(choice,cplex,cplex_dgxdt,cplex_d2gxdt,dimffnl,&
353 13290 : & d2gxdt,dgxdt,ffnl,gx,&
354 13290 : & idir,indlmn,istwf_k,kpg,matblk,mpi_enreg,nd2gxdt,ndgxdt,nkpg,&
355 13290 : & npw,nspinor,ph3d,signs,ucvol,ndat,ntypat,lmnmax,nattyp,is_kprime,&
356 : & iatom_only,atom_proj_shift,cpopt,&
357 : & nprojs,&
358 13290 : & vectin,&
359 13290 : & temp_realvec_r,temp_realvec_i,&
360 : & gpu_option,use_distrib)
361 :
362 : !Arguments ------------------------------------
363 : !scalars
364 : integer,intent(in) :: choice,cplex,dimffnl,idir,istwf_k,matblk,nd2gxdt
365 : integer,intent(in) :: ndgxdt,nkpg,lmnmax,ntypat,npw,nspinor,signs,ndat
366 : integer,intent(in) :: cpopt,iatom_only,atom_proj_shift
367 : integer,intent(in) :: nprojs
368 : real(dp),intent(in) :: ucvol
369 : type(MPI_type),intent(in) :: mpi_enreg
370 : integer,intent(in) :: gpu_option
371 : logical,intent(in) :: use_distrib,is_kprime
372 : !arrays
373 : integer,intent(in) :: indlmn(6,lmnmax,ntypat),nattyp(ntypat)
374 : integer,intent(out) :: cplex_dgxdt(ndgxdt),cplex_d2gxdt(nd2gxdt)
375 : real(dp),intent(in) :: ffnl(npw,dimffnl,lmnmax,ntypat),kpg(npw,nkpg)
376 : real(dp),intent(in) :: ph3d(2,npw,matblk)
377 : real(dp),target,intent(in) :: vectin(2,npw*nspinor*ndat)
378 : real(dp),target,intent(inout) :: d2gxdt(cplex,nd2gxdt,nprojs,ndat*nspinor)
379 : real(dp),target,intent(inout) :: dgxdt(cplex,ndgxdt*nprojs,ndat*nspinor)
380 : real(dp),target,intent(inout) :: gx(cplex,nprojs,ndat*nspinor)
381 : real(dp),target,intent(inout) :: temp_realvec_r(:),temp_realvec_i(:)
382 :
383 : !Local variables-------------------------------
384 : integer :: idat,ierr,i,ik,nprojs_all
385 : integer :: projs_beg,projs_end,dprojs_beg,dprojs_end,d2projs_beg,d2projs_end
386 : integer :: nprojs_blk,nprojs_last_blk,nprojs_cur_blk,rank,nprocs,iblock,nblocks
387 : logical :: use_sliced_gemms
388 13290 : real(dp), ABI_CONTIGUOUS pointer :: projs(:,:,:),projs_r(:,:,:),projs_i(:,:,:)
389 13290 : real(dp), ABI_CONTIGUOUS pointer :: dprojs(:,:,:),dprojs_r(:,:,:),dprojs_i(:,:,:)
390 13290 : real(dp), ABI_CONTIGUOUS pointer :: d2projs(:,:,:)
391 :
392 0 : ik=1; if(is_kprime) ik=2
393 :
394 19849 : cplex_dgxdt(:) = 0 ; if (cplex == 1) cplex_dgxdt(:) = 1
395 17401 : cplex_d2gxdt(:) = 0 ; if (cplex == 1) cplex_d2gxdt(:) = 1
396 : ! When istwf_k > 1, gx derivatives can be real or pure imaginary
397 : ! cplex_dgxdt(i) = 1 if dgxdt(1,i,:,:) is real, 2 if it is pure imaginary
398 : ! cplex_d2gxdt(i) = 1 if d2gxdt(1,i,:,:) is real, 2 if it is pure imaginary
399 13290 : if(ndgxdt > 0 .and. cplex==1) then
400 120 : if (choice==5.or.choice==51) cplex_dgxdt(:) = 2
401 120 : if (choice==54.and.signs==1) cplex_dgxdt(4:6) = 2
402 : !if (choice==54.and.signs==2) cplex_dgxdt(:) = 2
403 120 : if (choice==55.and.signs==1) cplex_dgxdt(7:9) = 2
404 : end if
405 13290 : if(nd2gxdt > 0 .and. cplex==1) then
406 0 : if (choice==54) cplex_d2gxdt(:) = 2
407 0 : if (choice==55.and.signs==1) cplex_d2gxdt(1:18)= 2
408 : end if
409 :
410 13290 : nprojs_all=nprojs
411 13290 : if(iatom_only>0) then
412 0 : nprojs_all=0
413 0 : do i=1,ntypat
414 0 : nprojs_all = nprojs_all + count(indlmn(3,:,i)>0)*nattyp(i)
415 : end do
416 : end if
417 13290 : nprojs_last_blk=nprojs_all
418 13290 : iblock=1; nblocks=1
419 13290 : use_sliced_gemms=.false.
420 13290 : if(gemm_nonlop_block_size>1 .and. .not. use_distrib) then
421 0 : nblocks=gemm_nonlop_block_size
422 0 : use_sliced_gemms=.true.
423 : end if
424 :
425 : call refresh_projectors(npw,istwf_k,nprojs_all,ndgxdt,nd2gxdt,&
426 13290 : & is_kprime,gpu_option)
427 13290 : if(nprojs_all/=gemm_nonlop_kpt(ik)%nprojs) ABI_BUG("Problem")
428 13290 : nprojs_blk = nprojs
429 13290 : nprojs_last_blk = nprojs
430 13290 : if(use_distrib) then
431 0 : rank=xmpi_comm_rank(gemm_nonlop_block_comm); nprocs=xmpi_comm_size(gemm_nonlop_block_comm)
432 0 : nprojs_blk = gemm_nonlop_kpt(ik)%nprojs_blk
433 0 : nprojs_last_blk = gemm_nonlop_kpt(ik)%nprojs_last_blk
434 0 : iblock=rank+1
435 13290 : else if(gemm_nonlop_block_size>1) then
436 0 : nprojs_blk = nprojs / gemm_nonlop_block_size
437 0 : nprojs_last_blk = nprojs_blk + modulo(nprojs,nprojs_blk)
438 : end if
439 :
440 13290 : projs_beg=1; projs_end=nprojs;
441 13290 : dprojs_beg=1; dprojs_end=max(1,nprojs*ndgxdt)
442 13290 : d2projs_beg=1; d2projs_end=max(1,nprojs*nd2gxdt)
443 13290 : if((choice==2 .and. signs==2)) then
444 0 : projs_beg=atom_proj_shift+1
445 0 : projs_end=projs_beg+nprojs-1
446 0 : dprojs_beg=atom_proj_shift*ndgxdt+1
447 0 : dprojs_end=dprojs_beg+nprojs*ndgxdt-1
448 0 : d2projs_beg=atom_proj_shift*nd2gxdt+1
449 0 : d2projs_end=dprojs_beg+nprojs*nd2gxdt-1
450 : end if
451 :
452 13290 : if(gemm_nonlop_block_size>1) then
453 0 : projs_beg=1; projs_end=nprojs_last_blk;
454 0 : dprojs_beg=1; dprojs_end=max(1,nprojs_last_blk*ndgxdt)
455 0 : d2projs_beg=1; d2projs_end=max(1,nprojs_last_blk*nd2gxdt)
456 : end if
457 :
458 13290 : if(istwf_k == 1) then
459 9179 : projs => gemm_nonlop_kpt(ik)%projs(:,:,projs_beg:projs_end)
460 9179 : if(ndgxdt>0) dprojs => gemm_nonlop_kpt(ik)%dprojs(:,:,dprojs_beg:dprojs_end)
461 9179 : if(nd2gxdt>0) d2projs => gemm_nonlop_kpt(ik)%d2projs(:,:,d2projs_beg:d2projs_end)
462 : else
463 4111 : projs_r => gemm_nonlop_kpt(ik)%projs_r(:,:,projs_beg:projs_end)
464 4111 : projs_i => gemm_nonlop_kpt(ik)%projs_i(:,:,projs_beg:projs_end)
465 4111 : if(ndgxdt>0) dprojs_r => gemm_nonlop_kpt(ik)%dprojs_r(:,:,dprojs_beg:dprojs_end)
466 4111 : if(ndgxdt>0) dprojs_i => gemm_nonlop_kpt(ik)%dprojs_i(:,:,dprojs_beg:dprojs_end)
467 : end if
468 :
469 13290 : 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 696 : & nprojs_last_blk,is_kprime,gpu_option,iblock)
473 696 : gemm_nonlop_kpt(ik)%ikpt=gemm_nonlop_ikpt_this_proc_being_treated
474 : end if
475 13290 : 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 424 : & .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 424 : & is_kprime,gpu_option,iblock)
485 424 : gemm_nonlop_kpt(ik)%choice = choice
486 424 : gemm_nonlop_kpt(ik)%idir = idir
487 : end if
488 : end if
489 :
490 13290 : if(cplex == 2) then
491 18358 : do i=1,nblocks
492 :
493 9179 : if(use_sliced_gemms .and. i>1) then
494 : call prep_projectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
495 : & ucvol,ffnl,ph3d,dimffnl,matblk,&
496 0 : & nprojs_last_blk,is_kprime,gpu_option,i)
497 0 : gemm_nonlop_kpt(ik)%ikpt=gemm_nonlop_ikpt_this_proc_being_treated
498 0 : if(choice>1 .and. ndgxdt>0) then
499 : call prep_dprojectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
500 : & ucvol,ffnl,ph3d,kpg,nkpg,dimffnl,matblk,&
501 : & nprojs_last_blk,ndgxdt,nd2gxdt,choice,signs,idir,&
502 0 : & is_kprime,gpu_option,i)
503 0 : gemm_nonlop_kpt(ik)%choice = choice
504 0 : gemm_nonlop_kpt(ik)%idir = idir
505 : end if
506 : end if
507 :
508 9179 : nprojs_cur_blk=nprojs
509 9179 : if(use_sliced_gemms) then
510 0 : if(i<nblocks) then
511 0 : nprojs_cur_blk=nprojs_blk
512 : else
513 0 : nprojs_cur_blk=nprojs_last_blk
514 : end if
515 : end if
516 :
517 9179 : if(cpopt<=1) then
518 : call opernla_xgemm(cplex, 'C', 'N', nprojs_cur_blk, ndat*nspinor, npw, cone, &
519 : & projs, npw,&
520 : & vectin, npw, czero, gx, nprojs,&
521 : & rank, nprocs,&
522 : & nprojs_blk, nprojs_last_blk, i,&
523 9179 : & gpu_option, use_distrib, use_sliced_gemms)
524 : end if
525 :
526 9179 : if(ndgxdt>0 .and. cpopt<=3) then
527 : call opernla_xgemm(cplex, 'C', 'N', ndgxdt*nprojs_cur_blk, ndat*nspinor, npw, cone, &
528 : & dprojs, npw,&
529 : & vectin, npw, czero, dgxdt, ndgxdt*nprojs,&
530 : & rank, nprocs,&
531 : & ndgxdt*nprojs_blk, ndgxdt*nprojs_last_blk, i,&
532 304 : & gpu_option, use_distrib, use_sliced_gemms)
533 : end if
534 :
535 18358 : if(nd2gxdt>0) then
536 : call opernla_xgemm(cplex, 'C', 'N', nd2gxdt*nprojs_cur_blk, ndat*nspinor, npw, cone, &
537 : & d2projs, npw,&
538 : & vectin, npw, czero, d2gxdt, nd2gxdt*nprojs,&
539 : & rank, nprocs,&
540 : & nd2gxdt*nprojs_blk, nd2gxdt*nprojs_last_blk, i,&
541 0 : & gpu_option, use_distrib, use_sliced_gemms)
542 : end if
543 : end do
544 :
545 : else ! cplex==1
546 :
547 : ! only compute real part of gx = P^* psi => gx_r = P_r^T psi_r + P_i^T psi_i
548 4111 : if(gpu_option == ABI_GPU_DISABLED) then
549 20938194 : temp_realvec_r(1:npw*nspinor*ndat) = vectin(1,1:npw*nspinor*ndat)
550 : else if(gpu_option == ABI_GPU_OPENMP) then
551 : #ifdef HAVE_OPENMP_OFFLOAD
552 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO &
553 : !$OMP& MAP(to:temp_realvec_r,vectin) PRIVATE(i)
554 : do i=1, npw*nspinor*ndat
555 : temp_realvec_r(i) = vectin(1,i)
556 : end do
557 : #endif
558 : end if
559 :
560 4111 : if(istwf_k == 2 .and. mpi_enreg%me_g0_fft == 1) then
561 2040 : if(gpu_option == ABI_GPU_DISABLED) then
562 8558 : do idat=1, ndat*nspinor
563 8558 : temp_realvec_r(1+(idat-1)*npw) = temp_realvec_r(1+(idat-1)*npw)/2
564 : end do
565 : else if(gpu_option == ABI_GPU_OPENMP) then
566 : #ifdef HAVE_OPENMP_OFFLOAD
567 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO &
568 : !$OMP& MAP(to:temp_realvec_r) PRIVATE(idat)
569 : do idat=1, ndat*nspinor
570 : temp_realvec_r(1+(idat-1)*npw) = temp_realvec_r(1+(idat-1)*npw)/2
571 : end do
572 : #endif
573 : end if
574 : end if
575 :
576 : ! Same with imaginary part
577 4111 : if(gpu_option == ABI_GPU_DISABLED) then
578 20938194 : temp_realvec_i(1:npw*nspinor*ndat) = vectin(2,1:npw*nspinor*ndat)
579 : else if(gpu_option == ABI_GPU_OPENMP) then
580 : #ifdef HAVE_OPENMP_OFFLOAD
581 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO &
582 : !$OMP& MAP(to:temp_realvec_i,vectin) PRIVATE(i)
583 : do i=1, npw*nspinor*ndat
584 : temp_realvec_i(i) = vectin(2,i)
585 : end do
586 : #endif
587 : end if
588 :
589 4111 : if(istwf_k == 2 .and. mpi_enreg%me_g0_fft == 1) then
590 2040 : if(gpu_option == ABI_GPU_DISABLED) then
591 8558 : do idat=1, ndat*nspinor
592 8558 : temp_realvec_i(1+(idat-1)*npw) = zero
593 : end do
594 : else if(gpu_option == ABI_GPU_OPENMP) then
595 : #ifdef HAVE_OPENMP_OFFLOAD
596 : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO &
597 : !$OMP& MAP(to:temp_realvec_i) PRIVATE(idat)
598 : do idat=1, ndat*nspinor
599 : temp_realvec_i(1+(idat-1)*npw) = zero
600 : end do
601 : #endif
602 : end if
603 : end if
604 :
605 8222 : do i=1,nblocks
606 :
607 4111 : if(use_sliced_gemms .and. i>1) then
608 : call prep_projectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
609 : & ucvol,ffnl,ph3d,dimffnl,matblk,&
610 0 : & nprojs_last_blk,is_kprime,gpu_option,i)
611 0 : gemm_nonlop_kpt(ik)%ikpt=gemm_nonlop_ikpt_this_proc_being_treated
612 0 : if(choice>1 .and. ndgxdt>0) then
613 : call prep_dprojectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
614 : & ucvol,ffnl,ph3d,kpg,nkpg,dimffnl,matblk,&
615 : & nprojs_last_blk,ndgxdt,nd2gxdt,choice,signs,idir,&
616 0 : & is_kprime,gpu_option,i)
617 0 : gemm_nonlop_kpt(ik)%choice = choice
618 0 : gemm_nonlop_kpt(ik)%idir = idir
619 : end if
620 : end if
621 :
622 4111 : nprojs_cur_blk=nprojs
623 4111 : if(use_sliced_gemms) then
624 0 : if(i<nblocks) then
625 0 : nprojs_cur_blk=nprojs_blk
626 : else
627 0 : nprojs_cur_blk=nprojs_last_blk
628 : end if
629 : end if
630 :
631 : ! Real part
632 4111 : if(cpopt<=1) then
633 : call opernla_xgemm(cplex, 'T', 'N', nprojs_cur_blk, ndat*nspinor, npw, cone, &
634 : & projs_r, npw, &
635 : & temp_realvec_r, npw, czero, gx, nprojs,&
636 : & rank, nprocs,&
637 : & nprojs_blk, nprojs_last_blk, i,&
638 4111 : & gpu_option, use_distrib, use_sliced_gemms)
639 : end if
640 4111 : if(ndgxdt>0 .and. cpopt<=3) then
641 : call opernla_xgemm(cplex, 'T', 'N', ndgxdt*nprojs_cur_blk, ndat*nspinor, npw, cone, &
642 : & dprojs_r, npw, &
643 : & temp_realvec_r, npw, czero, dgxdt, ndgxdt*nprojs,&
644 : & rank, nprocs,&
645 : & ndgxdt*nprojs_blk, ndgxdt*nprojs_last_blk, i,&
646 120 : & gpu_option, use_distrib, use_sliced_gemms)
647 : end if
648 :
649 : ! Imaginary part
650 4111 : if(cpopt<=1) then
651 : call opernla_xgemm(cplex, 'T', 'N', nprojs_cur_blk, ndat*nspinor, npw, cone, &
652 : & projs_i, npw, &
653 : & temp_realvec_i, npw, cone , gx, nprojs,&
654 : & rank, nprocs,&
655 : & nprojs_blk, nprojs_last_blk, i,&
656 4111 : & gpu_option, use_distrib, use_sliced_gemms)
657 : end if
658 8222 : if(ndgxdt>0) then
659 : call opernla_xgemm(cplex, 'T', 'N', ndgxdt*nprojs_cur_blk, ndat*nspinor, npw, cone, &
660 : & dprojs_i, npw, &
661 : & temp_realvec_i, npw, cone , dgxdt, ndgxdt*nprojs,&
662 : & rank, nprocs,&
663 : & ndgxdt*nprojs_blk, ndgxdt*nprojs_last_blk, i,&
664 120 : & gpu_option, use_distrib, use_sliced_gemms)
665 : end if
666 :
667 : end do
668 :
669 : ! Scale gx
670 4111 : if(cpopt<=1) then
671 4111 : if(gpu_option == ABI_GPU_DISABLED) then
672 1629101 : gx = gx * 2
673 0 : else if(gpu_option == ABI_GPU_OPENMP) then
674 0 : call abi_xscal(nprojs*nspinor*ndat, ctwo, gx, 1, x_cplx=cplex, gpu_option=gpu_option)
675 : end if
676 : end if
677 :
678 : ! Scale dgxdt
679 4111 : if(ndgxdt>0 .and. cpopt<=3) then
680 120 : if(gpu_option == ABI_GPU_DISABLED) then
681 264988 : dgxdt = dgxdt * 2
682 0 : else if(gpu_option == ABI_GPU_OPENMP) then
683 0 : call abi_xscal(ndgxdt*nprojs*nspinor*ndat, ctwo, dgxdt, 1, x_cplx=cplex, gpu_option=gpu_option)
684 : end if
685 : end if
686 :
687 : end if ! cplex == 2
688 13290 : if(gpu_option == ABI_GPU_DISABLED) then
689 13290 : call xmpi_sum(gx,mpi_enreg%comm_fft,ierr)
690 13290 : if (choice>1) then
691 424 : call xmpi_sum(dgxdt,mpi_enreg%comm_fft,ierr)
692 : end if
693 : end if
694 :
695 13290 : end subroutine opernla_gemm
696 : !!***
697 :
698 : end module m_opernla_gemm
699 : !!***
|
