Branch data Line data Source code
1 : : !!****m* ABINIT/m_ompgpu_fourwf
2 : : !! NAME
3 : : !! m_ompgpu_fourwf
4 : : !!
5 : : !! FUNCTION
6 : : !! Fake module to dupe the build system and allow it to include cuda files
7 : : !! in the chain of dependencies.
8 : : !!
9 : : !! COPYRIGHT
10 : : !! Copyright (C) 2000-2026 ABINIT group (MT)
11 : : !! This file is distributed under the terms of the
12 : : !! GNU General Public License, see ~abinit/COPYING
13 : : !! or http://www.gnu.org/copyleft/gpl.txt .
14 : : !!
15 : : !! SOURCE
16 : :
17 : : #if defined HAVE_CONFIG_H
18 : : #include "config.h"
19 : : #endif
20 : :
21 : : #include "abi_common.h"
22 : :
23 : : module m_ompgpu_fourwf
24 : :
25 : : use defs_basis
26 : : use m_abicore
27 : : use m_gputk
28 : : use m_abi_linalg
29 : : use m_errors
30 : : use m_xomp
31 : : use iso_c_binding
32 : :
33 : : #ifdef HAVE_GPU
34 : : use m_gpu_toolbox
35 : : #endif
36 : :
37 : : implicit none
38 : :
39 : : private
40 : :
41 : : #ifdef HAVE_OPENMP_OFFLOAD
42 : :
43 : : ! STATIC vars to avoid too much cuda overhead
44 : : integer :: fourdp_initialized=0
45 : : integer :: fourwf_initialized=0
46 : :
47 : : #define FOURDP_ID 1
48 : : #define FOURWF_ID 0
49 : :
50 : : ! FFT plans
51 : : integer :: fft_size_fourdp=-1
52 : : integer :: ndat_fourdp=-1
53 : : integer :: fft_size_fourwf=-1
54 : : integer :: ndat_fourwf=-1
55 : :
56 : : #endif
57 : :
58 : : public :: ompgpu_fourdp
59 : : !public :: alloc_ompgpu_fourdp
60 : : !public :: free_ompgpu_fourdp
61 : : public :: ompgpu_fourwf
62 : : public :: alloc_ompgpu_fourwf
63 : : public :: free_ompgpu_fourwf
64 : :
65 : : ! This routine is here to assess memory requirements
66 : : public :: ompgpu_fourwf_work_mem
67 : : contains
68 : :
69 [ # # ]: 0 : function ompgpu_fourwf_work_mem(ngfft, ndat) result(req_mem)
70 : :
71 : : integer, intent(in) :: ngfft(:), ndat
72 : : integer(kind=c_size_t) :: req_mem
73 : :
74 : 0 : req_mem = int(2, c_size_t) * dp * ngfft(1) * ngfft(2) * ngfft(3) * ndat
75 : :
76 : 0 : end function ompgpu_fourwf_work_mem
77 : :
78 : : !Tested usecases :
79 : : ! - Nvidia GPUs : FC_NVHPC + CUDA
80 : : ! - AMD GPUs : FC_LLVM + HIP
81 : : ! An eventual Intel implementation would use the OneAPI LLVM compiler.
82 : : ! Homemade CUDA/HIP interfaces would allow the use of GCC.
83 : : ! But it is likely that OpenMP performance won't be optimal outside GPU vendors compilers.
84 : : #ifdef HAVE_OPENMP_OFFLOAD
85 : :
86 : : subroutine ompgpu_fourdp(cplex,ngfft,ldx,ldy,ldz,ndat,isign,fofg,fofr)
87 : :
88 : : !Arguments ------------------------------------
89 : : !scalars
90 : : integer,intent(in) :: cplex,ngfft(18),ldx,ldy,ldz,ndat,isign
91 : : !arrays
92 : : real(dp),target,intent(inout) :: fofg(2*ldx*ldy*ldz*ndat)
93 : : real(dp),target,intent(inout) :: fofr(cplex*ldx*ldy*ldz*ndat)
94 : :
95 : : !Local variables-------------------------------
96 : : !scalars
97 : : integer :: n1,n2,n3,nfft_tot
98 : : complex(dp) :: norm
99 : : logical :: transfer_fofr, transfer_fofg
100 : : character(len=500) :: msg
101 : : ! *************************************************************************
102 : :
103 : : n1=ngfft(1);
104 : : n2=ngfft(2);
105 : : n3=ngfft(3);
106 : : nfft_tot=n1*n2*n3;
107 : : norm=dcmplx(one/(n1*n2*n3), 0.0_dp)
108 : :
109 : : !*********** CHECK some compatibilities **************
110 : : if( (n1/=ldx) .or. (n2/=ldy) .or. (n3/=ldz)) then
111 : : write(msg,"(a,a,i5,i5,i5,a,i5,i5,i5,a)") "FFT SIZE ERROR: \n when gpu mode is on the fft grid must not be augmented",&
112 : : "(n1,n2,n3) = (", n1,n2,n3,") whereas (ldx,ldy,ldz) = (", ldx,ldy,ldz, ")"
113 : : ABI_ERROR(msg)
114 : : end if
115 : :
116 : : ! *********** GPU ALLOCATIONS ***********************
117 : :
118 : : if (fourdp_initialized == 0) then
119 : : call alloc_ompgpu_fourdp(ngfft,ndat)
120 : : endif !end of initialisation
121 : :
122 : : ! If fft size has changed, we realloc our buffers
123 : : if((nfft_tot/=fft_size_fourdp) .or. (ndat/=ndat_fourdp)) then
124 : : call free_ompgpu_fourdp()
125 : : call alloc_ompgpu_fourdp(ngfft,ndat)
126 : : end if !end if "fft size changed"
127 : :
128 : : transfer_fofg= .not. xomp_target_is_present(c_loc(fofg))
129 : : transfer_fofr= .not. xomp_target_is_present(c_loc(fofr))
130 : :
131 : : !$OMP TARGET ENTER DATA MAP(alloc:fofg) IF(transfer_fofg)
132 : : !$OMP TARGET ENTER DATA MAP(alloc:fofr) IF(transfer_fofr)
133 : : !$OMP TARGET UPDATE TO(fofg) IF(transfer_fofg .and. isign==FFT_INVERSE)
134 : : !$OMP TARGET UPDATE TO(fofr) IF(transfer_fofr .and. isign==FFT_FORWARD)
135 : :
136 : : select case (cplex)
137 : : case (2)
138 : : ! Complex to Complex.
139 : : select case (isign)
140 : : case (FFT_INVERSE) ! +1
141 : : !$OMP TARGET DATA USE_DEVICE_ADDR(fofg,fofr)
142 : : call gpu_fft_exec_z2z(FOURDP_ID, c_loc(fofg), c_loc(fofr), FFT_INVERSE)
143 : : !$OMP END TARGET DATA
144 : : call gpu_fft_stream_synchronize(FOURDP_ID)
145 : : case (FFT_FORWARD) ! -1
146 : : !$OMP TARGET DATA USE_DEVICE_ADDR(fofg,fofr)
147 : : call gpu_fft_exec_z2z(FOURDP_ID, c_loc(fofr), c_loc(fofg), FFT_FORWARD)
148 : : !$OMP END TARGET DATA
149 : : call gpu_fft_stream_synchronize(FOURDP_ID)
150 : : ! Normalize here
151 : : call abi_xscal(ldx*ldy*ldz*ndat,norm,fofg,1,x_cplx=2,gpu_option=ABI_GPU_OPENMP)
152 : : case default
153 : : ABI_BUG("Wrong isign")
154 : : end select
155 : : case (1)
156 : : ! Real case.
157 : : ABI_BUG("Real case for OpenMP GPU fourdp not handled yet !")
158 : : !select case (isign)
159 : : !case (ABI_FFTW_BACKWARD) ! +1
160 : : ! ! +1; G --> R
161 : : ! call fftw3_c2r_op(nx,ny,nz,ldx,ldy,ldz,ndat,fofg,fofr)
162 : : !case (ABI_FFTW_FORWARD) ! -1
163 : : ! ! -1; R --> G
164 : : ! call fftw3_r2c_op(nx,ny,nz,ldx,ldy,ldz,ndat,fofr,fofg)
165 : : !case default
166 : : ! ABI_BUG("Wrong isign")
167 : : !end select
168 : : end select
169 : :
170 : : !$OMP TARGET UPDATE FROM(fofg) IF(transfer_fofg .and. isign==FFT_FORWARD)
171 : : !$OMP TARGET UPDATE FROM(fofr) IF(transfer_fofr .and. isign==FFT_INVERSE)
172 : : !$OMP TARGET EXIT DATA MAP(delete:fofg) IF(transfer_fofg)
173 : : !$OMP TARGET EXIT DATA MAP(delete:fofr) IF(transfer_fofr)
174 : :
175 : : end subroutine ompgpu_fourdp
176 : :
177 : : subroutine ompgpu_fourwf(cplex,denpot,fofgin,fofgout,fofr,gboundin,gboundout,istwf_k,&
178 : : & kg_kin,kg_kout,mgfft,me_g0,ndat,ngfft,npwin,npwout,ldx,ldy,ldz,option,weight_r,weight_i,&
179 : : & use_ndo,fofginb)
180 : :
181 : : !Arguments ------------------------------------
182 : : !scalars
183 : : integer,intent(in) :: cplex,istwf_k,ldx,ldy,ldz,me_g0,ndat,npwin,npwout,option,mgfft
184 : : real(dp),intent(in) :: weight_i(ndat),weight_r(ndat)
185 : : integer,intent(in),optional :: use_ndo
186 : : !arrays
187 : : integer,intent(in) :: gboundin(2*mgfft+8,2),gboundout(2*mgfft+8,2)
188 : : integer,intent(in) :: kg_kin(3,npwin),kg_kout(3,npwout),ngfft(18)
189 : : real(dp),target,intent(inout) :: denpot(cplex*ldx,ldy,ldz)
190 : : real(dp),target,intent(in) :: fofgin(2,npwin*ndat)
191 : : real(dp),target,intent(inout) :: fofr(2,ldx,ldy,ldz*ndat)
192 : : real(dp),target,intent(out) :: fofgout(2,npwout*ndat)
193 : : real(dp),target,intent(inout),optional :: fofginb(2,npwin*ndat)
194 : :
195 : : !Local variables-------------------------------
196 : : !scalars
197 : : character(len=500) :: msg
198 : : real(dp),allocatable,target :: fofrb(:,:,:,:)
199 : : logical :: l_use_ndo
200 : :
201 : : real(dp) :: xnorm,one,tmp
202 : :
203 : : integer :: n1,n2,n3,nfft_tot,npwmin,izd
204 : : integer :: cfft_size
205 : : integer :: shift_inv1,shift_inv2,shift_inv3
206 : : integer :: i1,i2,i3,ipw,idat;
207 : : integer :: i1inv,i2inv,i3inv
208 : : logical :: transfer_fofgin, transfer_fofginb, transfer_fofgout, transfer_denpot, transfer_fofr
209 : : integer(C_SIZE_T) :: byte_count
210 : : real(dp), contiguous, pointer :: work_gpu(:,:,:,:)
211 : :
212 : : npwmin=1; if(me_g0==1 .and. istwf_k==2) npwmin=2
213 : : l_use_ndo=.false.; if(present(use_ndo)) l_use_ndo=use_ndo==1
214 : :
215 : : n1=ngfft(1);
216 : : n2=ngfft(2);
217 : : n3=ngfft(3);
218 : : nfft_tot=n1*n2*n3;
219 : :
220 : : !*********** CHECK some compatibilities **************
221 : : if( (n1/=ldx) .or. (n2/=ldy) .or. (n3/=ldz)) then
222 : : write(msg,"(a,a,i5,i5,i5,a,i5,i5,i5,a)") "FFT SIZE ERROR: \n when gpu mode is on the fft grid must not be augmented",&
223 : : "(n1,n2,n3) = (", n1,n2,n3,") whereas (ldx,ldy,ldz) = (", ldx,ldy,ldz, ")"
224 : : ABI_ERROR(msg)
225 : : end if
226 : :
227 : : !*************** CUDA INITIALISATION STAGE ****
228 : : if (fourwf_initialized == 0) then
229 : : call alloc_ompgpu_fourwf(ngfft,ndat)
230 : : endif !end of initialisation
231 : :
232 : :
233 : : ! *********** GPU ALLOCATIONS ***********************
234 : :
235 : : ! If fft size has changed, we realloc our buffers
236 : : if((nfft_tot/=fft_size_fourwf) .or. (ndat/=ndat_fourwf)) then
237 : : call free_ompgpu_fourwf()
238 : : call alloc_ompgpu_fourwf(ngfft,ndat)
239 : : end if !end if "fft size changed"
240 : :
241 : : transfer_fofgin= .not. xomp_target_is_present(c_loc(fofgin)) .and. (option/=3)
242 : : transfer_fofgout=.not. xomp_target_is_present(c_loc(fofgout)) .and. (option==2 .or. option==3)
243 : : transfer_denpot =.not. xomp_target_is_present(c_loc(denpot)) .and. (option==2 .or. option==1)
244 : : transfer_fofr= .not. xomp_target_is_present(c_loc(fofr))
245 : : if(l_use_ndo) then
246 : : transfer_fofginb= .not. xomp_target_is_present(c_loc(fofginb)) .and. (option/=3)
247 : : end if
248 : :
249 : : !$OMP TARGET ENTER DATA MAP(to:fofgin) IF(transfer_fofgin)
250 : : !$OMP TARGET ENTER DATA MAP(alloc:fofgout) IF(transfer_fofgout)
251 : : !$OMP TARGET ENTER DATA MAP(alloc:denpot) IF(transfer_denpot)
252 : : !$OMP TARGET ENTER DATA MAP(alloc:fofr) IF(transfer_fofr)
253 : : if(l_use_ndo) then
254 : : !$OMP TARGET ENTER DATA MAP(to:fofginb) IF(transfer_fofginb)
255 : : ABI_MALLOC(fofrb, (2,n1,n2,n3*ndat))
256 : : !$OMP TARGET ENTER DATA MAP(alloc:fofrb)
257 : : end if
258 : :
259 : : if(option==3) then
260 : : ! We don't want to erase fofr initial value so we alloc a separate array.
261 : : ABI_MALLOC(work_gpu, (2,n1,n2,n3*ndat))
262 : : !$OMP TARGET ENTER DATA MAP(alloc:work_gpu)
263 : : else
264 : : work_gpu => fofr
265 : : end if
266 : :
267 : : if(option==3) then
268 : : !$OMP TARGET UPDATE TO(fofr) IF(transfer_fofr)
269 : : endif
270 : :
271 : : if(option==1 .or. option==2) then
272 : : ! We launch async transfert of denpot
273 : : !FIXME This async transfer might be better handled through CUDA/HIP after all...
274 : : ! Issues randomly occurs when using Cray compiler, seems fine with NVHPC.
275 : : #ifdef FC_CRAY
276 : : !$OMP TARGET UPDATE TO(denpot) IF(transfer_denpot)
277 : : #else
278 : : !$OMP TARGET UPDATE TO(denpot) NOWAIT IF(transfer_denpot)
279 : : #endif
280 : : if(option == 1) then
281 : : #ifdef FC_CRAY
282 : : !$OMP TARGET ENTER DATA MAP(to:weight_r,weight_i)
283 : : #else
284 : : !$OMP TARGET ENTER DATA MAP(to:weight_r,weight_i) NOWAIT
285 : : #endif
286 : : endif
287 : : endif
288 : :
289 : : if(option/=3) then
290 : :
291 : : ! GPU_SPHERE_IN
292 : :
293 : : cfft_size = 2*n1*n2*n3*ndat
294 : :
295 : : call gpu_set_to_zero(work_gpu,int(2,c_size_t)*n1*n2*n3*ndat)
296 : :
297 : : ! During GPU calculation we do some pre-calculation on symetries
298 : : if((istwf_k==2) .or. (istwf_k==4) .or. (istwf_k==6) .or. (istwf_k==8)) then
299 : : shift_inv1 = n1;
300 : : else
301 : : shift_inv1 = n1-1;
302 : : end if
303 : :
304 : : if((istwf_k>=2) .and. (istwf_k<=5)) then
305 : : shift_inv2 = n2;
306 : : else
307 : : shift_inv2 = n2-1;
308 : : end if
309 : :
310 : : if((istwf_k==2) .or. (istwf_k==3) .or. (istwf_k==6) .or. (istwf_k==7)) then
311 : : shift_inv3 = n3;
312 : : else
313 : : shift_inv3 = n3-1;
314 : : end if
315 : :
316 : :
317 : : ! !$OMP TARGET TEAMS DISTRIBUTE MAP(to:work_gpu,kg_kin,fofgin)
318 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) PRIVATE(i1,i2,i3) MAP(to:work_gpu,kg_kin,fofgin)
319 : : do idat = 1, ndat
320 : : !!! $OMP PARALLEL DO PRIVATE(i1,i2,i3)
321 : : do ipw = 1, npwin
322 : : i1=kg_kin(1,ipw); if(i1<0)i1=i1+n1; i1=i1+1
323 : : i2=kg_kin(2,ipw); if(i2<0)i2=i2+n2; i2=i2+1
324 : : i3=kg_kin(3,ipw); if(i3<0)i3=i3+n3; i3=i3+1
325 : : ! We write cfft(i1,i2,i3)
326 : : ! (double2): cfft[i1 + n1*(i2 + n2*(i3 + n3*idat))] = cg[ipw + npw*idat]
327 : : work_gpu(1, i1, i2, i3+n3*(idat-1)) = fofgin(1, (ipw + npwin*(idat-1)))
328 : : work_gpu(2, i1, i2, i3+n3*(idat-1)) = fofgin(2, (ipw + npwin*(idat-1)))
329 : : end do
330 : : end do
331 : :
332 : : if(istwf_k > 1) then
333 : : if (istwf_k==2 .and. me_g0==1) then
334 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO MAP(to:work_gpu,fofgin)
335 : : do idat = 1, ndat
336 : : work_gpu(1, 1, 1, 1+n3*(idat-1)) = fofgin(1, (1 + npwin*(idat-1)))
337 : : work_gpu(2, 1, 1, 1+n3*(idat-1)) = zero
338 : : end do
339 : : end if
340 : : ! !$OMP TARGET TEAMS DISTRIBUTE MAP(to:work_gpu,kg_kin,fofgin)
341 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) PRIVATE(i1,i2,i3,i1inv,i2inv,i3inv) MAP(to:work_gpu,kg_kin,fofgin)
342 : : do idat = 1, ndat
343 : : ! !$OMP PARALLEL DO PRIVATE(i1,i2,i3,i1inv,i2inv,i3inv)
344 : : do ipw = npwmin, npwin
345 : : i1=kg_kin(1,ipw); if(i1<0)i1=i1+n1;
346 : : i2=kg_kin(2,ipw); if(i2<0)i2=i2+n2;
347 : : i3=kg_kin(3,ipw); if(i3<0)i3=i3+n3;
348 : : #if defined HAVE_GPU_CUDA
349 : : i1inv = modulo(shift_inv1 - i1, n1) + 1
350 : : i2inv = modulo(shift_inv2 - i2, n2) + 1
351 : : i3inv = modulo(shift_inv3 - i3, n3) + 1
352 : : #elif defined HAVE_GPU_HIP
353 : : i1inv = (shift_inv1-i1) - ( ((shift_inv1-i1)/n1) * n1 ) + 1
354 : : i2inv = (shift_inv2-i2) - ( ((shift_inv2-i2)/n2) * n2 ) + 1
355 : : i3inv = (shift_inv3-i3) - ( ((shift_inv3-i3)/n3) * n3 ) + 1
356 : : #endif
357 : : work_gpu(1, i1inv, i2inv, i3inv+n3*(idat-1)) = fofgin(1, (ipw + npwin*(idat-1)))
358 : : work_gpu(2, i1inv, i2inv, i3inv+n3*(idat-1)) = -fofgin(2, (ipw + npwin*(idat-1)))
359 : : end do
360 : : end do
361 : : end if
362 : :
363 : : ! call backward fourier transform on gpu work_gpu => fofr_gpu
364 : : !$OMP TARGET DATA USE_DEVICE_ADDR(work_gpu,fofr)
365 : : call gpu_fft_exec_z2z(FOURWF_ID, c_loc(work_gpu), c_loc(fofr), FFT_INVERSE)
366 : : !$OMP END TARGET DATA
367 : : call gpu_fft_stream_synchronize(FOURWF_ID)
368 : :
369 : : ! In non-diagonal specific use-case, perform same operation with fofginb:
370 : : ! - box-to-sphere : fofginb => fofrb
371 : : ! - backward FFT : fofrb => fofrb
372 : : if(l_use_ndo) then
373 : :
374 : : call gpu_set_to_zero(fofrb,int(2,c_size_t)*n1*n2*n3*ndat)
375 : :
376 : : ! !$OMP TARGET TEAMS DISTRIBUTE MAP(to:fofrb,kg_kin,fofginb)
377 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) PRIVATE(i1,i2,i3) MAP(to:fofrb,kg_kin,fofginb)
378 : : do idat = 1, ndat
379 : : ! !$OMP PARALLEL DO PRIVATE(i1,i2,i3)
380 : : do ipw = 1, npwin
381 : : i1=kg_kin(1,ipw); if(i1<0)i1=i1+n1; i1=i1+1
382 : : i2=kg_kin(2,ipw); if(i2<0)i2=i2+n2; i2=i2+1
383 : : i3=kg_kin(3,ipw); if(i3<0)i3=i3+n3; i3=i3+1
384 : : ! We write cfft(i1,i2,i3)
385 : : ! (double2): cfft[i1 + n1*(i2 + n2*(i3 + n3*idat))] = cg[ipw + npw*idat]
386 : : fofrb(1, i1, i2, i3+n3*(idat-1)) = fofginb(1, (ipw + npwin*(idat-1)))
387 : : fofrb(2, i1, i2, i3+n3*(idat-1)) = fofginb(2, (ipw + npwin*(idat-1)))
388 : : end do
389 : : end do
390 : :
391 : : !$OMP TARGET DATA USE_DEVICE_ADDR(fofrb)
392 : : call gpu_fft_exec_z2z(FOURWF_ID, c_loc(fofrb), c_loc(fofrb), FFT_INVERSE)
393 : : !$OMP END TARGET DATA
394 : : call gpu_fft_stream_synchronize(FOURWF_ID)
395 : : end if
396 : : end if
397 : :
398 : : if(option==0) then
399 : : ! We copy back fofr
400 : : !$OMP TARGET UPDATE from(fofr) IF(transfer_fofr)
401 : : end if
402 : :
403 : : if(option==1) then
404 : : ! We finish denpot and weight transferts
405 : : !!$OMP TASKWAIT depend(in:denpot,weight_r,weight_i)
406 : : !$OMP TASKWAIT
407 : :
408 : : ! Perform density accumulation
409 : :
410 : : ! Non-diagonal case
411 : : if(l_use_ndo) then
412 : :
413 : : ! Collapse(4) would lead to race conditions here
414 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) MAP(to:fofr,fofrb,denpot,weight_r,weight_i)
415 : : do i3=1, n3
416 : : do i2=1, n2
417 : : do i1=1, n1
418 : : do idat = 1, ndat
419 : : denpot(i1,i2,i3) = denpot(i1,i2,i3) + weight_r(idat) *&
420 : : & (fofr(1, i1, i2, i3+(idat-1)*n3)*fofrb(1, i1, i2, i3+(idat-1)*n3)&
421 : : & +fofr(2, i1, i2, i3+(idat-1)*n3)*fofrb(2, i1, i2, i3+(idat-1)*n3))
422 : : denpot(i1,i2,i3) = denpot(i1,i2,i3) + weight_i(idat) *&
423 : : & (fofrb(2, i1, i2, i3+(idat-1)*n3)*fofr(1, i1, i2, i3+(idat-1)*n3)&
424 : : & -fofrb(1, i1, i2, i3+(idat-1)*n3)*fofr(2, i1, i2, i3+(idat-1)*n3))
425 : : end do
426 : : end do
427 : : end do
428 : : end do
429 : :
430 : : ! General case
431 : : else
432 : :
433 : : ! Collapse(4) would lead to race conditions here
434 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(3) MAP(to:fofr,denpot,weight_r,weight_i)
435 : : do i3=1, n3
436 : : do i2=1, n2
437 : : do i1=1, n1
438 : : do idat = 1, ndat
439 : : denpot(i1,i2,i3) = denpot(i1,i2,i3) + fofr(1, i1, i2, i3+(idat-1)*n3)*fofr(1, i1, i2, i3+(idat-1)*n3)*weight_r(idat)
440 : : denpot(i1,i2,i3) = denpot(i1,i2,i3) + fofr(2, i1, i2, i3+(idat-1)*n3)*fofr(2, i1, i2, i3+(idat-1)*n3)*weight_i(idat)
441 : : end do
442 : : end do
443 : : end do
444 : : end do
445 : : end if
446 : :
447 : : !$OMP TARGET UPDATE from(denpot) IF(transfer_denpot)
448 : :
449 : : end if
450 : :
451 : : if(option==2) then
452 : : ! We finished denpot transfert
453 : : !!$OMP TASKWAIT depend(in:denpot)
454 : : !$OMP TASKWAIT
455 : :
456 : : ! call gpu routine to Apply local potential
457 : : if(cplex==1) then
458 : : ! !$OMP TARGET TEAMS DISTRIBUTE MAP(to:denpot,fofr)
459 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(4) PRIVATE(izd, tmp) MAP(to:denpot,fofr)
460 : : do idat = 1, ndat
461 : : ! !$OMP PARALLEL DO PRIVATE(izd, tmp) COLLAPSE(3)
462 : : do i3=1, n3
463 : : do i2=1, n2
464 : : do i1=1, n1
465 : : izd = i3+(idat-1)*n3
466 : : tmp = denpot(i1,i2,i3)
467 : : fofr(1, i1, i2, izd) = fofr(1, i1, i2, izd) * tmp
468 : : fofr(2, i1, i2, izd) = fofr(2, i1, i2, izd) * tmp
469 : : end do
470 : : end do
471 : : end do
472 : : end do
473 : : else ! cplex==2
474 : :
475 : : ! !$OMP TARGET TEAMS DISTRIBUTE MAP(to:denpot,fofr)
476 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(4) PRIVATE(tmp) MAP(to:denpot,fofr)
477 : : do idat = 1, ndat
478 : : ! !$OMP PARALLEL DO COLLAPSE(3) PRIVATE(tmp)
479 : : do i3=1, n3
480 : : do i2=1, n2
481 : : do i1=1, n1
482 : : tmp = fofr(1, i1, i2, i3+(idat-1)*n3)
483 : : fofr(1, i1, i2, i3+(idat-1)*n3) = &
484 : : fofr(1, i1, i2, i3+(idat-1)*n3) * denpot(2*i1-1,i2,i3) &
485 : : - fofr(2, i1, i2, i3+(idat-1)*n3) * denpot(2*i1,i2,i3)
486 : :
487 : : fofr(2, i1, i2, i3+(idat-1)*n3) = &
488 : : fofr(2, i1, i2, i3+(idat-1)*n3) * denpot(2*i1-1,i2,i3) &
489 : : + tmp * denpot(2*i1,i2,i3)
490 : : end do
491 : : end do
492 : : end do
493 : : end do
494 : : end if
495 : : end if
496 : :
497 : : if(option==2 .or. option==3) then
498 : :
499 : : ! call forward fourier transform on gpu: fofr_gpu ==> work_gpu
500 : : !$OMP TARGET DATA USE_DEVICE_ADDR(work_gpu,fofr)
501 : : call gpu_fft_exec_z2z(FOURWF_ID, c_loc(fofr), c_loc(work_gpu), FFT_FORWARD)
502 : : !$OMP END TARGET DATA
503 : : call gpu_fft_stream_synchronize(FOURWF_ID)
504 : :
505 : : one=1
506 : : xnorm=one/dble(n1*n2*n3)
507 : : if (istwf_k==2 .and. me_g0==1) then
508 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO MAP(to:work_gpu,fofgout)
509 : : do idat = 1, ndat
510 : : fofgout (1, 1 + npwout*(idat-1)) = work_gpu(1, 1, 1, 1+n3*(idat-1)) * xnorm
511 : : fofgout (2, 1 + npwout*(idat-1)) = zero
512 : : end do
513 : : end if
514 : :
515 : : ! !$OMP TARGET TEAMS DISTRIBUTE MAP(to:work_gpu,kg_kout,fofgout)
516 : : !$OMP TARGET TEAMS DISTRIBUTE PARALLEL DO COLLAPSE(2) PRIVATE(i1,i2,i3) MAP(to:work_gpu,kg_kout,fofgout)
517 : : do idat = 1, ndat
518 : : ! !$OMP PARALLEL DO PRIVATE(i1,i2,i3)
519 : : do ipw = npwmin, npwout
520 : : i1=kg_kout(1,ipw); if(i1<0)i1=i1+n1; i1=i1+1
521 : : i2=kg_kout(2,ipw); if(i2<0)i2=i2+n2; i2=i2+1
522 : : i3=kg_kout(3,ipw); if(i3<0)i3=i3+n3; i3=i3+1
523 : :
524 : : ! We write cg(ig)
525 : : fofgout (1, ipw + npwout*(idat-1)) = work_gpu(1, i1, i2, i3+n3*(idat-1)) * xnorm
526 : : fofgout (2, ipw + npwout*(idat-1)) = work_gpu(2, i1, i2, i3+n3*(idat-1)) * xnorm
527 : : end do
528 : : end do
529 : :
530 : : !$OMP TARGET UPDATE FROM(fofgout) IF(transfer_fofgout)
531 : : end if
532 : :
533 : : if(option==1 .or. option==2) then
534 : : !$OMP TARGET EXIT DATA MAP(delete:denpot) IF(transfer_denpot)
535 : : if(option == 1) then
536 : : !$OMP TARGET EXIT DATA MAP(delete:weight_r,weight_i)
537 : : endif
538 : : endif
539 : :
540 : : if(option==3) then
541 : : !$OMP TARGET EXIT DATA MAP(delete:work_gpu)
542 : : ABI_FREE(work_gpu)
543 : : end if
544 : :
545 : : !$OMP TARGET EXIT DATA MAP(delete:fofgin) IF(transfer_fofgin)
546 : : !$OMP TARGET EXIT DATA MAP(delete:fofgout) IF(transfer_fofgout)
547 : : !$OMP TARGET EXIT DATA MAP(delete:fofr) IF(transfer_fofr)
548 : :
549 : : if(l_use_ndo) then
550 : : !$OMP TARGET EXIT DATA MAP(delete:fofginb) IF(transfer_fofginb)
551 : : !$OMP TARGET EXIT DATA MAP(delete:fofrb)
552 : : ABI_FREE(fofrb)
553 : : end if
554 : :
555 : : end subroutine ompgpu_fourwf
556 : : !!***
557 : :
558 : : !!****************************************************************************************************************
559 : : !!****************************************************************************************************************
560 : : !!****************************************************************************************************************
561 : :
562 : : ! Memory allocation routine
563 : : subroutine alloc_ompgpu_fourdp(ngfft, ndat)
564 : : integer, intent(in) :: ngfft(18), ndat
565 : :
566 : : integer :: n1,n2,n3, ldx, ldy, ldz
567 : : integer, target :: t_fft(3)
568 : :
569 : : fourdp_initialized = 1
570 : :
571 : : ! Size modification if too little memory allocation
572 : : n1=ngfft(1)
573 : : n2=ngfft(2)
574 : : n3=ngfft(3)
575 : : ldx = ngfft(4)
576 : : ldy = ngfft(5)
577 : : ldz = ngfft(6)
578 : : fft_size_fourdp=n1*n2*n3
579 : : ndat_fourdp = ndat
580 : :
581 : : ! Initialisation des plans FFT
582 : : t_fft(1) = n3;
583 : : t_fft(2) = n2;
584 : : t_fft(3) = n1;
585 : :
586 : : ! Creation du plan
587 : : call gpu_fft_plan_many(FOURDP_ID, 3, c_loc(t_fft), c_null_ptr, 1, ldx*ldy*ldz, c_null_ptr, 1, ldx*ldy*ldz, FFT_Z2Z, ndat);
588 : :
589 : : end subroutine alloc_ompgpu_fourdp
590 : :
591 : : subroutine free_ompgpu_fourdp()
592 : :
593 : : if(fourdp_initialized==0) return
594 : :
595 : : ! On detruit l'ancien plan
596 : : call gpu_fft_plan_destroy(FOURDP_ID)
597 : :
598 : : fourdp_initialized = 0;
599 : :
600 : : end subroutine free_ompgpu_fourdp
601 : :
602 : : subroutine alloc_ompgpu_fourwf(ngfft, ndat)
603 : : integer, intent(in) :: ngfft(18), ndat
604 : :
605 : : integer :: n1,n2,n3, ldx, ldy, ldz
606 : : integer, target :: t_fft(3),embed(3)
607 : :
608 : : fourwf_initialized = 1
609 : :
610 : : ! Size modification if too little memory allocation
611 : : n1=ngfft(1)
612 : : n2=ngfft(2)
613 : : n3=ngfft(3)
614 : : ldx = ngfft(4)
615 : : ldy = ngfft(5)
616 : : ldz = ngfft(6)
617 : : fft_size_fourwf=n1*n2*n3
618 : : ndat_fourwf = ndat
619 : :
620 : : ! Initialisation des plans FFT
621 : : t_fft(1) = n3;
622 : : t_fft(2) = n2;
623 : : t_fft(3) = n1;
624 : :
625 : : embed(1) = ldz
626 : : embed(2) = ldy
627 : : embed(3) = ldx
628 : :
629 : : ! Creation du plan
630 : : call gpu_fft_plan_many(FOURWF_ID, 3, c_loc(t_fft), c_loc(embed), 1, ldx*ldy*ldz, c_loc(embed), 1, ldx*ldy*ldz, FFT_Z2Z, ndat);
631 : :
632 : : end subroutine alloc_ompgpu_fourwf
633 : :
634 : : subroutine free_ompgpu_fourwf()
635 : :
636 : : if(fourwf_initialized==0) return
637 : :
638 : : ! On detruit l'ancien plan
639 : : call gpu_fft_plan_destroy(FOURWF_ID)
640 : :
641 : : fourwf_initialized = 0
642 : :
643 : : end subroutine free_ompgpu_fourwf
644 : :
645 : : #else
646 : : ! interface for unsupported compilers
647 : 0 : subroutine ompgpu_fourdp(cplex,ngfft,ldx,ldy,ldz,ndat,isign,fofg,fofr)
648 : :
649 : : !Arguments ------------------------------------
650 : : !scalars
651 : : integer,intent(in) :: cplex,ngfft(18),ldx,ldy,ldz,ndat,isign
652 : : !arrays
653 : : real(dp),intent(inout) :: fofg(2*ldx*ldy*ldz*ndat)
654 : : real(dp),intent(inout) :: fofr(cplex*ldx*ldy*ldz*ndat)
655 : :
656 : : ! *************************************************************************
657 : :
658 : : ABI_UNUSED((/cplex,ldx,ldy,ldz,ndat,isign/))
659 : : ABI_UNUSED((/ngfft/))
660 : : ABI_UNUSED((/fofg,fofr/))
661 : :
662 : 0 : end subroutine ompgpu_fourdp
663 : :
664 : 0 : subroutine ompgpu_fourwf(cplex,denpot,fofgin,fofgout,fofr,gboundin,gboundout,istwf_k,&
665 : 0 : & kg_kin,kg_kout,mgfft,ndat,ngfft,npwin,npwout,ldx,ldy,ldz,option,weight_r,weight_i)
666 : :
667 : : !Arguments ------------------------------------
668 : : !scalars
669 : : integer,intent(in) :: cplex,istwf_k,ldx,ldy,ldz,ndat,npwin,npwout,option,mgfft
670 : : real(dp),intent(in) :: weight_i(:),weight_r(:)
671 : : !arrays
672 : : integer,intent(in) :: gboundin(2*mgfft+8,2),gboundout(2*mgfft+8,2)
673 : : integer,intent(in) :: kg_kin(3,npwin),kg_kout(3,npwout),ngfft(18)
674 : : real(dp),target,intent(inout) :: denpot(cplex*ldx,ldy,ldz)
675 : : real(dp),intent(in) :: fofgin(2,npwin*ndat)
676 : : real(dp),intent(inout),target :: fofr(2,ldx,ldy,ldz*ndat)
677 : : real(dp),intent(out) :: fofgout(2,npwout*ndat)
678 : :
679 : : ! Marking unused variables to please code quality checks
680 : : ABI_UNUSED((/cplex,istwf_k,ldx,ldy,ldz,ndat,npwin,npwout,option,mgfft/))
681 : : ABI_UNUSED((/weight_i,weight_r/))
682 : : ABI_UNUSED((/gboundin,gboundout/))
683 : : ABI_UNUSED((/kg_kin,kg_kout,ngfft/))
684 : : ABI_UNUSED((/denpot,fofgin,fofgout,fofr/))
685 : 0 : ABI_BUG("Unhandled configuration for OpenMP GPU immplementation")
686 : 0 : end subroutine ompgpu_fourwf
687 : :
688 : 0 : subroutine alloc_ompgpu_fourwf(ngfft, ndat)
689 : : !Arguments ------------------------------------
690 : : integer, intent(in) :: ngfft(6), ndat
691 : :
692 : : ABI_UNUSED((/ngfft,ndat/))
693 : 0 : ABI_BUG("Unhandled configuration for OpenMP GPU immplementation")
694 : 0 : end subroutine alloc_ompgpu_fourwf
695 : :
696 : 0 : subroutine free_ompgpu_fourwf()
697 : 0 : ABI_BUG("Unhandled configuration for OpenMP GPU immplementation")
698 : 0 : end subroutine free_ompgpu_fourwf
699 : :
700 : : #endif
701 : :
702 : : end module m_ompgpu_fourwf
703 : : !!***
|