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 : : call refresh_projectors(npw,istwf_k,nprojs_all,ndgxdt,nd2gxdt,&
405 : 13290 : & is_kprime,gpu_option)
406 [ - + ]: 13290 : if(nprojs_all/=gemm_nonlop_kpt(ik)%nprojs) ABI_BUG("Problem")
407 : :
408 : : ! refresh_projectors may have reduced gemm_nonlop_block_size
409 : 13290 : iblock=1; nblocks=1
410 : 13290 : use_sliced_gemms=.false.
411 [ - + - - ]: 13290 : if(gemm_nonlop_block_size>1 .and. .not. use_distrib) then
412 : 0 : nblocks=gemm_nonlop_block_size
413 : 0 : use_sliced_gemms=.true.
414 : : end if
415 : :
416 : 13290 : nprojs_blk = nprojs
417 : 13290 : nprojs_last_blk = nprojs
418 [ - + ]: 13290 : if(use_distrib) then
419 : 0 : rank=xmpi_comm_rank(gemm_nonlop_block_comm); nprocs=xmpi_comm_size(gemm_nonlop_block_comm)
420 : 0 : nprojs_blk = gemm_nonlop_kpt(ik)%nprojs_blk
421 : 0 : nprojs_last_blk = gemm_nonlop_kpt(ik)%nprojs_last_blk
422 : 0 : iblock=rank+1
423 [ - + ]: 13290 : else if(gemm_nonlop_block_size>1) then
424 : 0 : nprojs_blk = nprojs / gemm_nonlop_block_size
425 : 0 : nprojs_last_blk = nprojs - (gemm_nonlop_block_size-1)*nprojs_blk
426 : : end if
427 : :
428 : 13290 : projs_beg=1; projs_end=nprojs;
429 : 13290 : dprojs_beg=1; dprojs_end=max(1,nprojs*ndgxdt)
430 : 13290 : d2projs_beg=1; d2projs_end=max(1,nprojs*nd2gxdt)
431 [ + + - + ]: 13290 : if((choice==2 .and. signs==2)) then
432 : 0 : projs_beg=atom_proj_shift+1
433 : 0 : projs_end=projs_beg+nprojs-1
434 : 0 : dprojs_beg=atom_proj_shift*ndgxdt+1
435 : 0 : dprojs_end=dprojs_beg+nprojs*ndgxdt-1
436 : 0 : d2projs_beg=atom_proj_shift*nd2gxdt+1
437 : 0 : d2projs_end=dprojs_beg+nprojs*nd2gxdt-1
438 : : end if
439 : :
440 [ - + ]: 13290 : if(gemm_nonlop_block_size>1) then
441 : 0 : projs_beg=1; projs_end=nprojs_last_blk;
442 : 0 : dprojs_beg=1; dprojs_end=max(1,nprojs_last_blk*ndgxdt)
443 : 0 : d2projs_beg=1; d2projs_end=max(1,nprojs_last_blk*nd2gxdt)
444 : : end if
445 : :
446 [ + + ]: 13290 : if(istwf_k == 1) then
447 : 9179 : projs => gemm_nonlop_kpt(ik)%projs(:,:,projs_beg:projs_end)
448 [ + + ]: 9179 : if(ndgxdt>0) dprojs => gemm_nonlop_kpt(ik)%dprojs(:,:,dprojs_beg:dprojs_end)
449 [ - + ]: 9179 : if(nd2gxdt>0) d2projs => gemm_nonlop_kpt(ik)%d2projs(:,:,d2projs_beg:d2projs_end)
450 : : else
451 : 4111 : projs_r => gemm_nonlop_kpt(ik)%projs_r(:,:,projs_beg:projs_end)
452 : 4111 : projs_i => gemm_nonlop_kpt(ik)%projs_i(:,:,projs_beg:projs_end)
453 [ + + ]: 4111 : if(ndgxdt>0) dprojs_r => gemm_nonlop_kpt(ik)%dprojs_r(:,:,dprojs_beg:dprojs_end)
454 [ + + ]: 4111 : if(ndgxdt>0) dprojs_i => gemm_nonlop_kpt(ik)%dprojs_i(:,:,dprojs_beg:dprojs_end)
455 : : end if
456 : :
457 [ + + - + ]: 13290 : if(gemm_nonlop_kpt(ik)%ikpt/=gemm_nonlop_ikpt_this_proc_being_treated .or. use_sliced_gemms) then
458 : : call prep_projectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
459 : : & ucvol,ffnl,ph3d,dimffnl,matblk,&
460 : 696 : & nprojs_last_blk,is_kprime,gpu_option,iblock)
461 : 696 : gemm_nonlop_kpt(ik)%ikpt=gemm_nonlop_ikpt_this_proc_being_treated
462 : : end if
463 [ + + + - ]: 13290 : if(choice>1 .and. ndgxdt>0) then
464 : : if( nd2gxdt/=gemm_nonlop_kpt(ik)%ngrads2 &
465 : : & .or. ndgxdt/=gemm_nonlop_kpt(ik)%ngrads &
466 : : & .or. choice/=gemm_nonlop_kpt(ik)%choice &
467 : : & .or. idir/=gemm_nonlop_kpt(ik)%idir &
468 [ + - + - : 424 : & .or. use_sliced_gemms) then
- + - - -
- ]
469 : : call prep_dprojectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
470 : : & ucvol,ffnl,ph3d,kpg,nkpg,dimffnl,matblk,&
471 : : & nprojs_last_blk,ndgxdt,nd2gxdt,choice,signs,idir,&
472 : 424 : & is_kprime,gpu_option,iblock)
473 : 424 : gemm_nonlop_kpt(ik)%choice = choice
474 : 424 : gemm_nonlop_kpt(ik)%idir = idir
475 : : end if
476 : : end if
477 : :
478 [ + + ]: 13290 : if(cplex == 2) then
479 [ + + ]: 18358 : do i=1,nblocks
480 : :
481 [ - + ]: 9179 : if(use_sliced_gemms .and. i>1) then
482 : : call prep_projectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
483 : : & ucvol,ffnl,ph3d,dimffnl,matblk,&
484 : 0 : & nprojs_last_blk,is_kprime,gpu_option,i)
485 : 0 : gemm_nonlop_kpt(ik)%ikpt=gemm_nonlop_ikpt_this_proc_being_treated
486 [ # # # # ]: 0 : if(choice>1 .and. ndgxdt>0) then
487 : : call prep_dprojectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
488 : : & ucvol,ffnl,ph3d,kpg,nkpg,dimffnl,matblk,&
489 : : & nprojs_last_blk,ndgxdt,nd2gxdt,choice,signs,idir,&
490 : 0 : & is_kprime,gpu_option,i)
491 : 0 : gemm_nonlop_kpt(ik)%choice = choice
492 : 0 : gemm_nonlop_kpt(ik)%idir = idir
493 : : end if
494 : : end if
495 : :
496 : 9179 : nprojs_cur_blk=nprojs
497 [ - + ]: 9179 : if(use_sliced_gemms) then
498 [ # # ]: 0 : if(i<nblocks) then
499 : 0 : nprojs_cur_blk=nprojs_blk
500 : : else
501 : 0 : nprojs_cur_blk=nprojs_last_blk
502 : : end if
503 : : end if
504 : :
505 [ + - ]: 9179 : if(cpopt<=1) then
506 : : call opernla_xgemm(cplex, 'C', 'N', nprojs_cur_blk, ndat*nspinor, npw, cone, &
507 : : & projs, npw,&
508 : : & vectin, npw, czero, gx, nprojs,&
509 : : & rank, nprocs,&
510 : : & nprojs_blk, nprojs_last_blk, i,&
511 : 9179 : & gpu_option, use_distrib, use_sliced_gemms)
512 : : end if
513 : :
514 [ + + + - ]: 9179 : if(ndgxdt>0 .and. cpopt<=3) then
515 : : call opernla_xgemm(cplex, 'C', 'N', ndgxdt*nprojs_cur_blk, ndat*nspinor, npw, cone, &
516 : : & dprojs, npw,&
517 : : & vectin, npw, czero, dgxdt, ndgxdt*nprojs,&
518 : : & rank, nprocs,&
519 : : & ndgxdt*nprojs_blk, ndgxdt*nprojs_last_blk, i,&
520 : 304 : & gpu_option, use_distrib, use_sliced_gemms)
521 : : end if
522 : :
523 [ - + ]: 18358 : if(nd2gxdt>0) then
524 : : call opernla_xgemm(cplex, 'C', 'N', nd2gxdt*nprojs_cur_blk, ndat*nspinor, npw, cone, &
525 : : & d2projs, npw,&
526 : : & vectin, npw, czero, d2gxdt, nd2gxdt*nprojs,&
527 : : & rank, nprocs,&
528 : : & nd2gxdt*nprojs_blk, nd2gxdt*nprojs_last_blk, i,&
529 : 0 : & gpu_option, use_distrib, use_sliced_gemms)
530 : : end if
531 : : end do
532 : :
533 : : else ! cplex==1
534 : :
535 : : ! only compute real part of gx = P^* psi => gx_r = P_r^T psi_r + P_i^T psi_i
536 [ + - ]: 4111 : if(gpu_option == ABI_GPU_DISABLED) then
537 [ + - + + : 20938194 : temp_realvec_r(1:npw*nspinor*ndat) = vectin(1,1:npw*nspinor*ndat)
+ + ]
538 : : else if(gpu_option == ABI_GPU_OPENMP) then
539 : : #ifdef HAVE_OPENMP_OFFLOAD
540 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO &
541 : : !$OMP& MAP(to:temp_realvec_r,vectin) PRIVATE(i)
542 : : do i=1, npw*nspinor*ndat
543 : : temp_realvec_r(i) = vectin(1,i)
544 : : end do
545 : : #endif
546 : : end if
547 : :
548 [ + + + + ]: 4111 : if(istwf_k == 2 .and. mpi_enreg%me_g0_fft == 1) then
549 [ + - ]: 2040 : if(gpu_option == ABI_GPU_DISABLED) then
550 [ + + ]: 8558 : do idat=1, ndat*nspinor
551 : 8558 : temp_realvec_r(1+(idat-1)*npw) = temp_realvec_r(1+(idat-1)*npw)/2
552 : : end do
553 : : else if(gpu_option == ABI_GPU_OPENMP) then
554 : : #ifdef HAVE_OPENMP_OFFLOAD
555 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO &
556 : : !$OMP& MAP(to:temp_realvec_r) PRIVATE(idat)
557 : : do idat=1, ndat*nspinor
558 : : temp_realvec_r(1+(idat-1)*npw) = temp_realvec_r(1+(idat-1)*npw)/2
559 : : end do
560 : : #endif
561 : : end if
562 : : end if
563 : :
564 : : ! Same with imaginary part
565 [ + - ]: 4111 : if(gpu_option == ABI_GPU_DISABLED) then
566 [ + - + + : 20938194 : temp_realvec_i(1:npw*nspinor*ndat) = vectin(2,1:npw*nspinor*ndat)
+ + ]
567 : : else if(gpu_option == ABI_GPU_OPENMP) then
568 : : #ifdef HAVE_OPENMP_OFFLOAD
569 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO &
570 : : !$OMP& MAP(to:temp_realvec_i,vectin) PRIVATE(i)
571 : : do i=1, npw*nspinor*ndat
572 : : temp_realvec_i(i) = vectin(2,i)
573 : : end do
574 : : #endif
575 : : end if
576 : :
577 [ + + + + ]: 4111 : if(istwf_k == 2 .and. mpi_enreg%me_g0_fft == 1) then
578 [ + - ]: 2040 : if(gpu_option == ABI_GPU_DISABLED) then
579 [ + + ]: 8558 : do idat=1, ndat*nspinor
580 : 8558 : temp_realvec_i(1+(idat-1)*npw) = zero
581 : : end do
582 : : else if(gpu_option == ABI_GPU_OPENMP) then
583 : : #ifdef HAVE_OPENMP_OFFLOAD
584 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO &
585 : : !$OMP& MAP(to:temp_realvec_i) PRIVATE(idat)
586 : : do idat=1, ndat*nspinor
587 : : temp_realvec_i(1+(idat-1)*npw) = zero
588 : : end do
589 : : #endif
590 : : end if
591 : : end if
592 : :
593 [ + + ]: 8222 : do i=1,nblocks
594 : :
595 [ - + - - ]: 4111 : if(use_sliced_gemms .and. i>1) then
596 : : call prep_projectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
597 : : & ucvol,ffnl,ph3d,dimffnl,matblk,&
598 : 0 : & nprojs_last_blk,is_kprime,gpu_option,i)
599 : 0 : gemm_nonlop_kpt(ik)%ikpt=gemm_nonlop_ikpt_this_proc_being_treated
600 [ # # # # ]: 0 : if(choice>1 .and. ndgxdt>0) then
601 : : call prep_dprojectors(npw,lmnmax,ntypat,indlmn,nattyp,istwf_k,&
602 : : & ucvol,ffnl,ph3d,kpg,nkpg,dimffnl,matblk,&
603 : : & nprojs_last_blk,ndgxdt,nd2gxdt,choice,signs,idir,&
604 : 0 : & is_kprime,gpu_option,i)
605 : 0 : gemm_nonlop_kpt(ik)%choice = choice
606 : 0 : gemm_nonlop_kpt(ik)%idir = idir
607 : : end if
608 : : end if
609 : :
610 : 4111 : nprojs_cur_blk=nprojs
611 [ - + ]: 4111 : if(use_sliced_gemms) then
612 [ # # ]: 0 : if(i<nblocks) then
613 : 0 : nprojs_cur_blk=nprojs_blk
614 : : else
615 : 0 : nprojs_cur_blk=nprojs_last_blk
616 : : end if
617 : : end if
618 : :
619 : : ! Real part
620 [ + - ]: 4111 : if(cpopt<=1) then
621 : : call opernla_xgemm(cplex, 'T', 'N', nprojs_cur_blk, ndat*nspinor, npw, cone, &
622 : : & projs_r, npw, &
623 : : & temp_realvec_r, npw, czero, gx, nprojs,&
624 : : & rank, nprocs,&
625 : : & nprojs_blk, nprojs_last_blk, i,&
626 [ - + - - : 4111 : & gpu_option, use_distrib, use_sliced_gemms)
- - - + ]
627 : : end if
628 [ + + + - ]: 4111 : if(ndgxdt>0 .and. cpopt<=3) then
629 : : call opernla_xgemm(cplex, 'T', 'N', ndgxdt*nprojs_cur_blk, ndat*nspinor, npw, cone, &
630 : : & dprojs_r, npw, &
631 : : & temp_realvec_r, npw, czero, dgxdt, ndgxdt*nprojs,&
632 : : & rank, nprocs,&
633 : : & ndgxdt*nprojs_blk, ndgxdt*nprojs_last_blk, i,&
634 [ - + - - : 120 : & gpu_option, use_distrib, use_sliced_gemms)
- - - + ]
635 : : end if
636 : :
637 : : ! Imaginary part
638 [ + - ]: 4111 : if(cpopt<=1) then
639 : : call opernla_xgemm(cplex, 'T', 'N', nprojs_cur_blk, ndat*nspinor, npw, cone, &
640 : : & projs_i, npw, &
641 : : & temp_realvec_i, npw, cone , gx, nprojs,&
642 : : & rank, nprocs,&
643 : : & nprojs_blk, nprojs_last_blk, i,&
644 [ - + - - : 4111 : & gpu_option, use_distrib, use_sliced_gemms)
- - - + ]
645 : : end if
646 [ + + ]: 8222 : if(ndgxdt>0) then
647 : : call opernla_xgemm(cplex, 'T', 'N', ndgxdt*nprojs_cur_blk, ndat*nspinor, npw, cone, &
648 : : & dprojs_i, npw, &
649 : : & temp_realvec_i, npw, cone , dgxdt, ndgxdt*nprojs,&
650 : : & rank, nprocs,&
651 : : & ndgxdt*nprojs_blk, ndgxdt*nprojs_last_blk, i,&
652 [ - + - - : 120 : & gpu_option, use_distrib, use_sliced_gemms)
- - - + ]
653 : : end if
654 : :
655 : : end do
656 : :
657 : : ! Scale gx
658 [ + - ]: 4111 : if(cpopt<=1) then
659 [ + - ]: 4111 : if(gpu_option == ABI_GPU_DISABLED) then
660 [ + + + + : 1629101 : gx = gx * 2
+ + ]
661 [ # # ]: 0 : else if(gpu_option == ABI_GPU_OPENMP) then
662 : 0 : call abi_xscal(nprojs*nspinor*ndat, ctwo, gx, 1, x_cplx=cplex, gpu_option=gpu_option)
663 : : end if
664 : : end if
665 : :
666 : : ! Scale dgxdt
667 [ + + + - ]: 4111 : if(ndgxdt>0 .and. cpopt<=3) then
668 [ + - ]: 120 : if(gpu_option == ABI_GPU_DISABLED) then
669 [ + + + + : 264988 : dgxdt = dgxdt * 2
+ + ]
670 [ # # ]: 0 : else if(gpu_option == ABI_GPU_OPENMP) then
671 : 0 : call abi_xscal(ndgxdt*nprojs*nspinor*ndat, ctwo, dgxdt, 1, x_cplx=cplex, gpu_option=gpu_option)
672 : : end if
673 : : end if
674 : :
675 : : end if ! cplex == 2
676 [ + - ]: 13290 : if(gpu_option == ABI_GPU_DISABLED) then
677 : 13290 : call xmpi_sum(gx,mpi_enreg%comm_fft,ierr)
678 [ + + ]: 13290 : if (choice>1) then
679 : 424 : call xmpi_sum(dgxdt,mpi_enreg%comm_fft,ierr)
680 : : end if
681 : : end if
682 : :
683 [ - + ]: 13290 : end subroutine opernla_gemm
684 : : !!***
685 : :
686 : : end module m_opernla_gemm
687 : : !!***
|
