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