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), ABI_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 : !!***
|
