Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2026 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : ! **************************************************************************************************
9 : !> \note
10 : !> This module contains routines necessary to operate on plane waves on GPUs
11 : ! > independently of the GPU platform.
12 : !> \par History
13 : !> BGL (06-Mar-2008) : Created
14 : !> AG (18-May-2012) : Refacturing:
15 : !> - added explicit interfaces to C routines
16 : !> - enable double precision complex transformations
17 : !> AG (11-Sept-2012) : Modifications:
18 : !> - use pointers if precision mapping is not required
19 : !> - use OMP for mapping
20 : !> MT (Jan 2022) : Modifications
21 : !> - use a generic interface for fft calls to GPUs
22 : !> - Support both Nvidia and AMD GPUs. Other GPUs manufacturers
23 : !> can be added easily.
24 : !> \author Benjamin G. Levine
25 : ! **************************************************************************************************
26 : MODULE pw_gpu
27 : USE ISO_C_BINDING, ONLY: C_DOUBLE,&
28 : C_INT,&
29 : C_LOC,&
30 : C_PTR
31 : USE fft_tools, ONLY: &
32 : cube_transpose_1, cube_transpose_2, fft_scratch_sizes, fft_scratch_type, get_fft_scratch, &
33 : release_fft_scratch, x_to_yz, xz_to_yz, yz_to_x, yz_to_xz
34 : USE kinds, ONLY: dp
35 : USE mathconstants, ONLY: z_zero
36 : USE message_passing, ONLY: mp_cart_type
37 : USE pw_grid_types, ONLY: FULLSPACE
38 : USE pw_types, ONLY: pw_c1d_gs_type,&
39 : pw_r3d_rs_type
40 : #include "../base/base_uses.f90"
41 :
42 : IMPLICIT NONE
43 :
44 : PRIVATE
45 :
46 : PUBLIC :: pw_gpu_r3dc1d_3d
47 : PUBLIC :: pw_gpu_c1dr3d_3d
48 : PUBLIC :: pw_gpu_r3dc1d_3d_ps
49 : PUBLIC :: pw_gpu_c1dr3d_3d_ps
50 : PUBLIC :: pw_gpu_init, pw_gpu_finalize
51 :
52 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'pw_gpu'
53 : LOGICAL, PARAMETER, PRIVATE :: debug_this_module = .FALSE.
54 :
55 : CONTAINS
56 :
57 : ! **************************************************************************************************
58 : !> \brief Allocates resources on the gpu device for gpu fft acceleration
59 : !> \author Ole Schuett
60 : ! **************************************************************************************************
61 10486 : SUBROUTINE pw_gpu_init()
62 : INTEGER :: dummy
63 : INTERFACE
64 : SUBROUTINE pw_gpu_init_c() BIND(C, name="pw_gpu_init")
65 : END SUBROUTINE pw_gpu_init_c
66 : END INTERFACE
67 :
68 : MARK_USED(dummy) ! TODO: fix fpretty
69 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_PW)
70 : CALL pw_gpu_init_c()
71 : #else
72 : ! Nothing to do.
73 : #endif
74 10486 : END SUBROUTINE pw_gpu_init
75 :
76 : ! **************************************************************************************************
77 : !> \brief Releases resources on the gpu device for gpu fft acceleration
78 : !> \author Ole Schuett
79 : ! **************************************************************************************************
80 10486 : SUBROUTINE pw_gpu_finalize()
81 : INTEGER :: dummy
82 : INTERFACE
83 : SUBROUTINE pw_gpu_finalize_c() BIND(C, name="pw_gpu_finalize")
84 : END SUBROUTINE pw_gpu_finalize_c
85 : END INTERFACE
86 :
87 : MARK_USED(dummy) ! TODO: fix fpretty
88 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_PW)
89 : CALL pw_gpu_finalize_c()
90 : #else
91 : ! Nothing to do.
92 : #endif
93 10486 : END SUBROUTINE pw_gpu_finalize
94 :
95 : ! **************************************************************************************************
96 : !> \brief perform an fft followed by a gather on the gpu
97 : !> \param pw1 ...
98 : !> \param pw2 ...
99 : !> \author Benjamin G Levine
100 : ! **************************************************************************************************
101 0 : SUBROUTINE pw_gpu_r3dc1d_3d(pw1, pw2)
102 : TYPE(pw_r3d_rs_type), INTENT(IN) :: pw1
103 : TYPE(pw_c1d_gs_type), INTENT(INOUT) :: pw2
104 :
105 : CHARACTER(len=*), PARAMETER :: routineN = 'pw_gpu_r3dc1d_3d'
106 :
107 : COMPLEX(KIND=dp), POINTER :: ptr_pwout
108 : INTEGER :: handle, l1, l2, l3, ngpts
109 0 : INTEGER, DIMENSION(:), POINTER :: npts
110 : INTEGER, POINTER :: ptr_ghatmap
111 : REAL(KIND=dp) :: scale
112 : REAL(KIND=dp), POINTER :: ptr_pwin
113 : INTERFACE
114 : SUBROUTINE pw_gpu_cfffg_c(din, zout, ghatmap, npts, ngpts, scale) BIND(C, name="pw_gpu_cfffg")
115 : IMPORT
116 : TYPE(C_PTR), INTENT(IN), VALUE :: din
117 : TYPE(C_PTR), VALUE :: zout
118 : TYPE(C_PTR), INTENT(IN), VALUE :: ghatmap
119 : INTEGER(KIND=C_INT), DIMENSION(*), INTENT(IN):: npts
120 : INTEGER(KIND=C_INT), INTENT(IN), VALUE :: ngpts
121 : REAL(KIND=C_DOUBLE), INTENT(IN), VALUE :: scale
122 :
123 : END SUBROUTINE pw_gpu_cfffg_c
124 : END INTERFACE
125 :
126 0 : CALL timeset(routineN, handle)
127 :
128 0 : scale = 1.0_dp/REAL(pw1%pw_grid%ngpts, KIND=dp)
129 :
130 0 : ngpts = SIZE(pw2%pw_grid%gsq)
131 : l1 = LBOUND(pw1%array, 1)
132 : l2 = LBOUND(pw1%array, 2)
133 : l3 = LBOUND(pw1%array, 3)
134 0 : npts => pw1%pw_grid%npts
135 :
136 : ! pointers to data arrays
137 0 : ptr_pwin => pw1%array(l1, l2, l3)
138 0 : ptr_pwout => pw2%array(1)
139 :
140 : ! pointer to map array
141 0 : ptr_ghatmap => pw2%pw_grid%g_hatmap(1, 1)
142 :
143 : ! invoke the combined transformation
144 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_PW)
145 : CALL pw_gpu_cfffg_c(c_loc(ptr_pwin), c_loc(ptr_pwout), c_loc(ptr_ghatmap), npts, ngpts, scale)
146 : #else
147 0 : CPABORT("Compiled without pw offloading.")
148 : #endif
149 :
150 0 : CALL timestop(handle)
151 0 : END SUBROUTINE pw_gpu_r3dc1d_3d
152 :
153 : ! **************************************************************************************************
154 : !> \brief perform an scatter followed by a fft on the gpu
155 : !> \param pw1 ...
156 : !> \param pw2 ...
157 : !> \author Benjamin G Levine
158 : ! **************************************************************************************************
159 0 : SUBROUTINE pw_gpu_c1dr3d_3d(pw1, pw2)
160 : TYPE(pw_c1d_gs_type), INTENT(IN) :: pw1
161 : TYPE(pw_r3d_rs_type), INTENT(INOUT) :: pw2
162 :
163 : CHARACTER(len=*), PARAMETER :: routineN = 'pw_gpu_c1dr3d_3d'
164 :
165 : COMPLEX(KIND=dp), POINTER :: ptr_pwin
166 : INTEGER :: handle, l1, l2, l3, ngpts, nmaps
167 0 : INTEGER, DIMENSION(:), POINTER :: npts
168 : INTEGER, POINTER :: ptr_ghatmap
169 : REAL(KIND=dp) :: scale
170 : REAL(KIND=dp), POINTER :: ptr_pwout
171 : INTERFACE
172 : SUBROUTINE pw_gpu_sfffc_c(zin, dout, ghatmap, npts, ngpts, nmaps, scale) BIND(C, name="pw_gpu_sfffc")
173 : IMPORT
174 : TYPE(C_PTR), INTENT(IN), VALUE :: zin
175 : TYPE(C_PTR), VALUE :: dout
176 : TYPE(C_PTR), INTENT(IN), VALUE :: ghatmap
177 : INTEGER(KIND=C_INT), DIMENSION(*), INTENT(IN):: npts
178 : INTEGER(KIND=C_INT), INTENT(IN), VALUE :: ngpts, nmaps
179 : REAL(KIND=C_DOUBLE), INTENT(IN), VALUE :: scale
180 : END SUBROUTINE pw_gpu_sfffc_c
181 : END INTERFACE
182 :
183 0 : CALL timeset(routineN, handle)
184 :
185 0 : scale = 1.0_dp
186 :
187 0 : ngpts = SIZE(pw1%pw_grid%gsq)
188 : l1 = LBOUND(pw2%array, 1)
189 : l2 = LBOUND(pw2%array, 2)
190 : l3 = LBOUND(pw2%array, 3)
191 0 : npts => pw1%pw_grid%npts
192 :
193 : ! pointers to data arrays
194 0 : ptr_pwin => pw1%array(1)
195 0 : ptr_pwout => pw2%array(l1, l2, l3)
196 :
197 : ! pointer to map array
198 0 : nmaps = SIZE(pw1%pw_grid%g_hatmap, 2)
199 0 : ptr_ghatmap => pw1%pw_grid%g_hatmap(1, 1)
200 :
201 : ! invoke the combined transformation
202 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_PW)
203 : CALL pw_gpu_sfffc_c(c_loc(ptr_pwin), c_loc(ptr_pwout), c_loc(ptr_ghatmap), npts, ngpts, nmaps, scale)
204 : #else
205 0 : CPABORT("Compiled without pw offloading")
206 : #endif
207 :
208 0 : CALL timestop(handle)
209 0 : END SUBROUTINE pw_gpu_c1dr3d_3d
210 :
211 : ! **************************************************************************************************
212 : !> \brief perform an parallel fft followed by a gather on the gpu
213 : !> \param pw1 ...
214 : !> \param pw2 ...
215 : !> \author Andreas Gloess
216 : ! **************************************************************************************************
217 0 : SUBROUTINE pw_gpu_r3dc1d_3d_ps(pw1, pw2)
218 : TYPE(pw_r3d_rs_type), INTENT(IN) :: pw1
219 : TYPE(pw_c1d_gs_type), INTENT(INOUT) :: pw2
220 :
221 : CHARACTER(len=*), PARAMETER :: routineN = 'pw_gpu_r3dc1d_3d_ps'
222 :
223 0 : COMPLEX(KIND=dp), DIMENSION(:, :), POINTER :: grays, pbuf, qbuf, rbuf, sbuf
224 0 : COMPLEX(KIND=dp), DIMENSION(:, :, :), POINTER :: tbuf
225 : INTEGER :: g_pos, handle, lg, lmax, mg, mmax, mx2, &
226 : mz2, n1, n2, ngpts, nmax, numtask, rp
227 0 : INTEGER, ALLOCATABLE, DIMENSION(:) :: p2p
228 : INTEGER, DIMENSION(2) :: r_dim, r_pos
229 0 : INTEGER, DIMENSION(:), POINTER :: n, nloc, nyzray
230 0 : INTEGER, DIMENSION(:, :, :, :), POINTER :: bo
231 : REAL(KIND=dp) :: scale
232 0 : TYPE(fft_scratch_sizes) :: fft_scratch_size
233 : TYPE(fft_scratch_type), POINTER :: fft_scratch
234 0 : TYPE(mp_cart_type) :: rs_group
235 :
236 0 : CALL timeset(routineN, handle)
237 :
238 0 : scale = 1.0_dp/REAL(pw1%pw_grid%ngpts, KIND=dp)
239 :
240 : ! dimensions
241 0 : n => pw1%pw_grid%npts
242 0 : nloc => pw1%pw_grid%npts_local
243 0 : grays => pw1%pw_grid%grays
244 0 : ngpts = nloc(1)*nloc(2)*nloc(3)
245 :
246 : !..transform
247 0 : IF (pw1%pw_grid%para%ray_distribution) THEN
248 0 : rs_group = pw1%pw_grid%para%group
249 0 : nyzray => pw1%pw_grid%para%nyzray
250 0 : bo => pw1%pw_grid%para%bo
251 :
252 0 : g_pos = rs_group%mepos
253 0 : numtask = rs_group%num_pe
254 0 : r_dim = rs_group%num_pe_cart
255 0 : r_pos = rs_group%mepos_cart
256 :
257 0 : lg = SIZE(grays, 1)
258 0 : mg = SIZE(grays, 2)
259 0 : mmax = MAX(mg, 1)
260 0 : lmax = MAX(lg, (ngpts/mmax + 1))
261 :
262 0 : ALLOCATE (p2p(0:numtask - 1))
263 :
264 0 : CALL rs_group%rank_compare(rs_group, p2p)
265 :
266 0 : rp = p2p(g_pos)
267 0 : mx2 = bo(2, 1, rp, 2) - bo(1, 1, rp, 2) + 1
268 0 : mz2 = bo(2, 3, rp, 2) - bo(1, 3, rp, 2) + 1
269 0 : n1 = MAXVAL(bo(2, 1, :, 1) - bo(1, 1, :, 1) + 1)
270 0 : n2 = MAXVAL(bo(2, 2, :, 1) - bo(1, 2, :, 1) + 1)
271 0 : nmax = MAX((2*n2)/numtask, 2)*mx2*mz2
272 0 : nmax = MAX(nmax, n1*MAXVAL(nyzray))
273 :
274 0 : fft_scratch_size%nx = nloc(1)
275 0 : fft_scratch_size%ny = nloc(2)
276 0 : fft_scratch_size%nz = nloc(3)
277 0 : fft_scratch_size%lmax = lmax
278 0 : fft_scratch_size%mmax = mmax
279 0 : fft_scratch_size%mx1 = bo(2, 1, rp, 1) - bo(1, 1, rp, 1) + 1
280 0 : fft_scratch_size%mx2 = mx2
281 0 : fft_scratch_size%my1 = bo(2, 2, rp, 1) - bo(1, 2, rp, 1) + 1
282 0 : fft_scratch_size%mz2 = mz2
283 0 : fft_scratch_size%lg = lg
284 0 : fft_scratch_size%mg = mg
285 0 : fft_scratch_size%nbx = MAXVAL(bo(2, 1, :, 2))
286 0 : fft_scratch_size%nbz = MAXVAL(bo(2, 3, :, 2))
287 0 : fft_scratch_size%mcz1 = MAXVAL(bo(2, 3, :, 1) - bo(1, 3, :, 1) + 1)
288 0 : fft_scratch_size%mcx2 = MAXVAL(bo(2, 1, :, 2) - bo(1, 1, :, 2) + 1)
289 0 : fft_scratch_size%mcz2 = MAXVAL(bo(2, 3, :, 2) - bo(1, 3, :, 2) + 1)
290 0 : fft_scratch_size%nmax = nmax
291 0 : fft_scratch_size%nmray = MAXVAL(nyzray)
292 0 : fft_scratch_size%nyzray = nyzray(g_pos)
293 0 : fft_scratch_size%rs_group = rs_group
294 0 : fft_scratch_size%g_pos = g_pos
295 0 : fft_scratch_size%r_pos = r_pos
296 0 : fft_scratch_size%r_dim = r_dim
297 0 : fft_scratch_size%numtask = numtask
298 :
299 0 : IF (r_dim(2) > 1) THEN
300 : !
301 : ! real space is distributed over x and y coordinate
302 : ! we have two stages of communication
303 : !
304 0 : IF (r_dim(1) == 1) THEN
305 0 : CPABORT("This processor distribution is not supported.")
306 : END IF
307 :
308 0 : CALL get_fft_scratch(fft_scratch, tf_type=300, n=n, fft_sizes=fft_scratch_size)
309 :
310 : ! assign buffers
311 0 : qbuf => fft_scratch%p2buf
312 0 : rbuf => fft_scratch%p3buf
313 0 : pbuf => fft_scratch%p4buf
314 0 : sbuf => fft_scratch%p5buf
315 :
316 : ! FFT along z
317 0 : CALL pw_gpu_cf(pw1, qbuf)
318 :
319 : ! Exchange data ( transpose of matrix )
320 0 : CALL cube_transpose_2(qbuf, bo(:, :, :, 1), bo(:, :, :, 2), rbuf, fft_scratch)
321 :
322 : ! FFT along y
323 : ! use the inbuild fft-lib
324 : ! CALL fft_1dm(fft_scratch%fft_plan(2), rbuf, pbuf, 1.0_dp, stat)
325 : ! or cufft (works faster, but is only faster if plans are stored)
326 0 : CALL pw_gpu_f(rbuf, pbuf, +1, n(2), mx2*mz2)
327 :
328 : ! Exchange data ( transpose of matrix ) and sort
329 : CALL xz_to_yz(pbuf, rs_group, r_dim, g_pos, p2p, pw1%pw_grid%para%yzp, nyzray, &
330 0 : bo(:, :, :, 2), sbuf, fft_scratch)
331 :
332 : ! FFT along x
333 0 : CALL pw_gpu_fg(sbuf, pw2, scale)
334 :
335 0 : CALL release_fft_scratch(fft_scratch)
336 :
337 : ELSE
338 : !
339 : ! real space is only distributed over x coordinate
340 : ! we have one stage of communication, after the transform of
341 : ! direction x
342 : !
343 :
344 0 : CALL get_fft_scratch(fft_scratch, tf_type=200, n=n, fft_sizes=fft_scratch_size)
345 :
346 : ! assign buffers
347 0 : tbuf => fft_scratch%tbuf
348 0 : sbuf => fft_scratch%r1buf
349 :
350 : ! FFT along y and z
351 0 : CALL pw_gpu_cff(pw1, tbuf)
352 :
353 : ! Exchange data ( transpose of matrix ) and sort
354 : CALL yz_to_x(tbuf, rs_group, g_pos, p2p, pw1%pw_grid%para%yzp, nyzray, &
355 0 : bo(:, :, :, 2), sbuf, fft_scratch)
356 :
357 : ! FFT along x
358 0 : CALL pw_gpu_fg(sbuf, pw2, scale)
359 :
360 0 : CALL release_fft_scratch(fft_scratch)
361 :
362 : END IF
363 :
364 0 : DEALLOCATE (p2p)
365 :
366 : !--------------------------------------------------------------------------
367 : ELSE
368 0 : CPABORT("Not implemented (no ray_distr.) in: pw_gpu_r3dc1d_3d_ps.")
369 : END IF
370 :
371 0 : CALL timestop(handle)
372 0 : END SUBROUTINE pw_gpu_r3dc1d_3d_ps
373 :
374 : ! **************************************************************************************************
375 : !> \brief perform an parallel scatter followed by a fft on the gpu
376 : !> \param pw1 ...
377 : !> \param pw2 ...
378 : !> \author Andreas Gloess
379 : ! **************************************************************************************************
380 0 : SUBROUTINE pw_gpu_c1dr3d_3d_ps(pw1, pw2)
381 : TYPE(pw_c1d_gs_type), INTENT(IN) :: pw1
382 : TYPE(pw_r3d_rs_type), INTENT(INOUT) :: pw2
383 :
384 : CHARACTER(len=*), PARAMETER :: routineN = 'pw_gpu_c1dr3d_3d_ps'
385 :
386 0 : COMPLEX(KIND=dp), DIMENSION(:, :), POINTER :: grays, pbuf, qbuf, rbuf, sbuf
387 0 : COMPLEX(KIND=dp), DIMENSION(:, :, :), POINTER :: tbuf
388 : INTEGER :: g_pos, handle, lg, lmax, mg, mmax, mx2, &
389 : mz2, n1, n2, ngpts, nmax, numtask, rp
390 0 : INTEGER, ALLOCATABLE, DIMENSION(:) :: p2p
391 : INTEGER, DIMENSION(2) :: r_dim, r_pos
392 0 : INTEGER, DIMENSION(:), POINTER :: n, nloc, nyzray
393 0 : INTEGER, DIMENSION(:, :, :, :), POINTER :: bo
394 : REAL(KIND=dp) :: scale
395 0 : TYPE(fft_scratch_sizes) :: fft_scratch_size
396 : TYPE(fft_scratch_type), POINTER :: fft_scratch
397 0 : TYPE(mp_cart_type) :: rs_group
398 :
399 0 : CALL timeset(routineN, handle)
400 :
401 0 : scale = 1.0_dp
402 :
403 : ! dimensions
404 0 : n => pw1%pw_grid%npts
405 0 : nloc => pw1%pw_grid%npts_local
406 0 : grays => pw1%pw_grid%grays
407 0 : ngpts = nloc(1)*nloc(2)*nloc(3)
408 :
409 : !..transform
410 0 : IF (pw1%pw_grid%para%ray_distribution) THEN
411 0 : rs_group = pw1%pw_grid%para%group
412 0 : nyzray => pw1%pw_grid%para%nyzray
413 0 : bo => pw1%pw_grid%para%bo
414 :
415 0 : g_pos = rs_group%mepos
416 0 : numtask = rs_group%num_pe
417 0 : r_dim = rs_group%num_pe_cart
418 0 : r_pos = rs_group%mepos_cart
419 :
420 0 : lg = SIZE(grays, 1)
421 0 : mg = SIZE(grays, 2)
422 0 : mmax = MAX(mg, 1)
423 0 : lmax = MAX(lg, (ngpts/mmax + 1))
424 :
425 0 : ALLOCATE (p2p(0:numtask - 1))
426 :
427 0 : CALL rs_group%rank_compare(rs_group, p2p)
428 :
429 0 : rp = p2p(g_pos)
430 0 : mx2 = bo(2, 1, rp, 2) - bo(1, 1, rp, 2) + 1
431 0 : mz2 = bo(2, 3, rp, 2) - bo(1, 3, rp, 2) + 1
432 0 : n1 = MAXVAL(bo(2, 1, :, 1) - bo(1, 1, :, 1) + 1)
433 0 : n2 = MAXVAL(bo(2, 2, :, 1) - bo(1, 2, :, 1) + 1)
434 0 : nmax = MAX((2*n2)/numtask, 2)*mx2*mz2
435 0 : nmax = MAX(nmax, n1*MAXVAL(nyzray))
436 :
437 0 : fft_scratch_size%nx = nloc(1)
438 0 : fft_scratch_size%ny = nloc(2)
439 0 : fft_scratch_size%nz = nloc(3)
440 0 : fft_scratch_size%lmax = lmax
441 0 : fft_scratch_size%mmax = mmax
442 0 : fft_scratch_size%mx1 = bo(2, 1, rp, 1) - bo(1, 1, rp, 1) + 1
443 0 : fft_scratch_size%mx2 = mx2
444 0 : fft_scratch_size%my1 = bo(2, 2, rp, 1) - bo(1, 2, rp, 1) + 1
445 0 : fft_scratch_size%mz2 = mz2
446 0 : fft_scratch_size%lg = lg
447 0 : fft_scratch_size%mg = mg
448 0 : fft_scratch_size%nbx = MAXVAL(bo(2, 1, :, 2))
449 0 : fft_scratch_size%nbz = MAXVAL(bo(2, 3, :, 2))
450 0 : fft_scratch_size%mcz1 = MAXVAL(bo(2, 3, :, 1) - bo(1, 3, :, 1) + 1)
451 0 : fft_scratch_size%mcx2 = MAXVAL(bo(2, 1, :, 2) - bo(1, 1, :, 2) + 1)
452 0 : fft_scratch_size%mcz2 = MAXVAL(bo(2, 3, :, 2) - bo(1, 3, :, 2) + 1)
453 0 : fft_scratch_size%nmax = nmax
454 0 : fft_scratch_size%nmray = MAXVAL(nyzray)
455 0 : fft_scratch_size%nyzray = nyzray(g_pos)
456 0 : fft_scratch_size%rs_group = rs_group
457 0 : fft_scratch_size%g_pos = g_pos
458 0 : fft_scratch_size%r_pos = r_pos
459 0 : fft_scratch_size%r_dim = r_dim
460 0 : fft_scratch_size%numtask = numtask
461 :
462 0 : IF (r_dim(2) > 1) THEN
463 : !
464 : ! real space is distributed over x and y coordinate
465 : ! we have two stages of communication
466 : !
467 0 : IF (r_dim(1) == 1) THEN
468 0 : CPABORT("This processor distribution is not supported.")
469 : END IF
470 :
471 0 : CALL get_fft_scratch(fft_scratch, tf_type=300, n=n, fft_sizes=fft_scratch_size)
472 :
473 : ! assign buffers
474 0 : pbuf => fft_scratch%p7buf
475 0 : qbuf => fft_scratch%p4buf
476 0 : rbuf => fft_scratch%p3buf
477 0 : sbuf => fft_scratch%p2buf
478 :
479 : ! FFT along x
480 0 : CALL pw_gpu_sf(pw1, pbuf, scale)
481 :
482 : ! Exchange data ( transpose of matrix ) and sort
483 0 : IF (pw1%pw_grid%grid_span /= FULLSPACE) qbuf = z_zero
484 : CALL yz_to_xz(pbuf, rs_group, r_dim, g_pos, p2p, pw1%pw_grid%para%yzp, nyzray, &
485 0 : bo(:, :, :, 2), qbuf, fft_scratch)
486 :
487 : ! FFT along y
488 : ! use the inbuild fft-lib
489 : ! CALL fft_1dm(fft_scratch%fft_plan(5), qbuf, rbuf, 1.0_dp, stat)
490 : ! or cufft (works faster, but is only faster if plans are stored)
491 0 : CALL pw_gpu_f(qbuf, rbuf, -1, n(2), mx2*mz2)
492 :
493 : ! Exchange data ( transpose of matrix )
494 0 : IF (pw1%pw_grid%grid_span /= FULLSPACE) sbuf = z_zero
495 :
496 0 : CALL cube_transpose_1(rbuf, bo(:, :, :, 2), bo(:, :, :, 1), sbuf, fft_scratch)
497 :
498 : ! FFT along z
499 0 : CALL pw_gpu_fc(sbuf, pw2)
500 :
501 0 : CALL release_fft_scratch(fft_scratch)
502 :
503 : ELSE
504 : !
505 : ! real space is only distributed over x coordinate
506 : ! we have one stage of communication, after the transform of
507 : ! direction x
508 : !
509 :
510 0 : CALL get_fft_scratch(fft_scratch, tf_type=200, n=n, fft_sizes=fft_scratch_size)
511 :
512 : ! assign buffers
513 0 : sbuf => fft_scratch%r1buf
514 0 : tbuf => fft_scratch%tbuf
515 :
516 : ! FFT along x
517 0 : CALL pw_gpu_sf(pw1, sbuf, scale)
518 :
519 : ! Exchange data ( transpose of matrix ) and sort
520 0 : IF (pw1%pw_grid%grid_span /= FULLSPACE) tbuf = z_zero
521 : CALL x_to_yz(sbuf, rs_group, g_pos, p2p, pw1%pw_grid%para%yzp, nyzray, &
522 0 : bo(:, :, :, 2), tbuf, fft_scratch)
523 :
524 : ! FFT along y and z
525 0 : CALL pw_gpu_ffc(tbuf, pw2)
526 :
527 0 : CALL release_fft_scratch(fft_scratch)
528 :
529 : END IF
530 :
531 0 : DEALLOCATE (p2p)
532 :
533 : !--------------------------------------------------------------------------
534 : ELSE
535 0 : CPABORT("Not implemented (no ray_distr.) in: pw_gpu_c1dr3d_3d_ps.")
536 : END IF
537 :
538 0 : CALL timestop(handle)
539 0 : END SUBROUTINE pw_gpu_c1dr3d_3d_ps
540 :
541 : ! **************************************************************************************************
542 : !> \brief perform a parallel real_to_complex copy followed by a 2D-FFT on the gpu
543 : !> \param pw1 ...
544 : !> \param pwbuf ...
545 : !> \author Andreas Gloess
546 : ! **************************************************************************************************
547 0 : SUBROUTINE pw_gpu_cff(pw1, pwbuf)
548 : TYPE(pw_r3d_rs_type), INTENT(IN) :: pw1
549 : COMPLEX(KIND=dp), DIMENSION(:, :, :), &
550 : INTENT(INOUT), TARGET :: pwbuf
551 :
552 : CHARACTER(len=*), PARAMETER :: routineN = 'pw_gpu_cff'
553 :
554 : COMPLEX(KIND=dp), POINTER :: ptr_pwout
555 : INTEGER :: handle, l1, l2, l3
556 0 : INTEGER, DIMENSION(:), POINTER :: npts
557 : REAL(KIND=dp), POINTER :: ptr_pwin
558 : INTERFACE
559 : SUBROUTINE pw_gpu_cff_c(din, zout, npts) BIND(C, name="pw_gpu_cff")
560 : IMPORT
561 : TYPE(C_PTR), INTENT(IN), VALUE :: din
562 : TYPE(C_PTR), VALUE :: zout
563 : INTEGER(KIND=C_INT), DIMENSION(*), INTENT(IN):: npts
564 : END SUBROUTINE pw_gpu_cff_c
565 : END INTERFACE
566 :
567 0 : CALL timeset(routineN, handle)
568 :
569 : ! dimensions
570 0 : npts => pw1%pw_grid%npts_local
571 : l1 = LBOUND(pw1%array, 1)
572 : l2 = LBOUND(pw1%array, 2)
573 : l3 = LBOUND(pw1%array, 3)
574 :
575 : ! pointers to data arrays
576 0 : ptr_pwin => pw1%array(l1, l2, l3)
577 0 : ptr_pwout => pwbuf(1, 1, 1)
578 :
579 : ! invoke the combined transformation
580 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_PW)
581 : CALL pw_gpu_cff_c(c_loc(ptr_pwin), c_loc(ptr_pwout), npts)
582 : #else
583 0 : CPABORT("Compiled without pw offloading")
584 : #endif
585 :
586 0 : CALL timestop(handle)
587 0 : END SUBROUTINE pw_gpu_cff
588 :
589 : ! **************************************************************************************************
590 : !> \brief perform a parallel 2D-FFT followed by a complex_to_real copy on the gpu
591 : !> \param pwbuf ...
592 : !> \param pw2 ...
593 : !> \author Andreas Gloess
594 : ! **************************************************************************************************
595 0 : SUBROUTINE pw_gpu_ffc(pwbuf, pw2)
596 : COMPLEX(KIND=dp), DIMENSION(:, :, :), INTENT(IN), &
597 : TARGET :: pwbuf
598 : TYPE(pw_r3d_rs_type), INTENT(IN) :: pw2
599 :
600 : CHARACTER(len=*), PARAMETER :: routineN = 'pw_gpu_ffc'
601 :
602 : COMPLEX(KIND=dp), POINTER :: ptr_pwin
603 : INTEGER :: handle, l1, l2, l3
604 0 : INTEGER, DIMENSION(:), POINTER :: npts
605 : REAL(KIND=dp), POINTER :: ptr_pwout
606 : INTERFACE
607 : SUBROUTINE pw_gpu_ffc_c(zin, dout, npts) BIND(C, name="pw_gpu_ffc")
608 : IMPORT
609 : TYPE(C_PTR), INTENT(IN), VALUE :: zin
610 : TYPE(C_PTR), VALUE :: dout
611 : INTEGER(KIND=C_INT), DIMENSION(*), INTENT(IN):: npts
612 : END SUBROUTINE pw_gpu_ffc_c
613 : END INTERFACE
614 :
615 0 : CALL timeset(routineN, handle)
616 :
617 : ! dimensions
618 0 : npts => pw2%pw_grid%npts_local
619 : l1 = LBOUND(pw2%array, 1)
620 : l2 = LBOUND(pw2%array, 2)
621 : l3 = LBOUND(pw2%array, 3)
622 :
623 : ! pointers to data arrays
624 0 : ptr_pwin => pwbuf(1, 1, 1)
625 0 : ptr_pwout => pw2%array(l1, l2, l3)
626 :
627 : ! invoke the combined transformation
628 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_PW)
629 : CALL pw_gpu_ffc_c(c_loc(ptr_pwin), c_loc(ptr_pwout), npts)
630 : #else
631 0 : CPABORT("Compiled without pw offloading")
632 : #endif
633 :
634 0 : CALL timestop(handle)
635 0 : END SUBROUTINE pw_gpu_ffc
636 :
637 : ! **************************************************************************************************
638 : !> \brief perform a parallel real_to_complex copy followed by a 1D-FFT on the gpu
639 : !> \param pw1 ...
640 : !> \param pwbuf ...
641 : !> \author Andreas Gloess
642 : ! **************************************************************************************************
643 0 : SUBROUTINE pw_gpu_cf(pw1, pwbuf)
644 : TYPE(pw_r3d_rs_type), INTENT(IN) :: pw1
645 : COMPLEX(KIND=dp), DIMENSION(:, :), INTENT(INOUT), &
646 : TARGET :: pwbuf
647 :
648 : CHARACTER(len=*), PARAMETER :: routineN = 'pw_gpu_cf'
649 :
650 : COMPLEX(KIND=dp), POINTER :: ptr_pwout
651 : INTEGER :: handle, l1, l2, l3
652 0 : INTEGER, DIMENSION(:), POINTER :: npts
653 : REAL(KIND=dp), POINTER :: ptr_pwin
654 : INTERFACE
655 : SUBROUTINE pw_gpu_cf_c(din, zout, npts) BIND(C, name="pw_gpu_cf")
656 : IMPORT
657 : TYPE(C_PTR), INTENT(IN), VALUE :: din
658 : TYPE(C_PTR), VALUE :: zout
659 : INTEGER(KIND=C_INT), DIMENSION(*), INTENT(IN):: npts
660 : END SUBROUTINE pw_gpu_cf_c
661 : END INTERFACE
662 :
663 0 : CALL timeset(routineN, handle)
664 :
665 : ! dimensions
666 0 : npts => pw1%pw_grid%npts_local
667 : l1 = LBOUND(pw1%array, 1)
668 : l2 = LBOUND(pw1%array, 2)
669 : l3 = LBOUND(pw1%array, 3)
670 :
671 : ! pointers to data arrays
672 0 : ptr_pwin => pw1%array(l1, l2, l3)
673 0 : ptr_pwout => pwbuf(1, 1)
674 :
675 : ! invoke the combined transformation
676 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_PW)
677 : CALL pw_gpu_cf_c(c_loc(ptr_pwin), c_loc(ptr_pwout), npts)
678 : #else
679 0 : CPABORT("Compiled without pw offloading")
680 : #endif
681 0 : CALL timestop(handle)
682 0 : END SUBROUTINE pw_gpu_cf
683 :
684 : ! **************************************************************************************************
685 : !> \brief perform a parallel 1D-FFT followed by a complex_to_real copy on the gpu
686 : !> \param pwbuf ...
687 : !> \param pw2 ...
688 : !> \author Andreas Gloess
689 : ! **************************************************************************************************
690 0 : SUBROUTINE pw_gpu_fc(pwbuf, pw2)
691 : COMPLEX(KIND=dp), DIMENSION(:, :), INTENT(IN), &
692 : TARGET :: pwbuf
693 : TYPE(pw_r3d_rs_type), INTENT(IN) :: pw2
694 :
695 : CHARACTER(len=*), PARAMETER :: routineN = 'pw_gpu_fc'
696 :
697 : COMPLEX(KIND=dp), POINTER :: ptr_pwin
698 : INTEGER :: handle, l1, l2, l3
699 0 : INTEGER, DIMENSION(:), POINTER :: npts
700 : REAL(KIND=dp), POINTER :: ptr_pwout
701 : INTERFACE
702 : SUBROUTINE pw_gpu_fc_c(zin, dout, npts) BIND(C, name="pw_gpu_fc")
703 : IMPORT
704 : TYPE(C_PTR), INTENT(IN), VALUE :: zin
705 : TYPE(C_PTR), VALUE :: dout
706 : INTEGER(KIND=C_INT), DIMENSION(*), INTENT(IN):: npts
707 : END SUBROUTINE pw_gpu_fc_c
708 : END INTERFACE
709 :
710 0 : CALL timeset(routineN, handle)
711 :
712 0 : npts => pw2%pw_grid%npts_local
713 : l1 = LBOUND(pw2%array, 1)
714 : l2 = LBOUND(pw2%array, 2)
715 : l3 = LBOUND(pw2%array, 3)
716 :
717 : ! pointers to data arrays
718 0 : ptr_pwin => pwbuf(1, 1)
719 0 : ptr_pwout => pw2%array(l1, l2, l3)
720 :
721 : ! invoke the combined transformation
722 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_PW)
723 : CALL pw_gpu_fc_c(c_loc(ptr_pwin), c_loc(ptr_pwout), npts)
724 : #else
725 0 : CPABORT("Compiled without pw offloading")
726 : #endif
727 :
728 0 : CALL timestop(handle)
729 0 : END SUBROUTINE pw_gpu_fc
730 :
731 : ! **************************************************************************************************
732 : !> \brief perform a parallel 1D-FFT on the gpu
733 : !> \param pwbuf1 ...
734 : !> \param pwbuf2 ...
735 : !> \param dir ...
736 : !> \param n ...
737 : !> \param m ...
738 : !> \author Andreas Gloess
739 : ! **************************************************************************************************
740 0 : SUBROUTINE pw_gpu_f(pwbuf1, pwbuf2, dir, n, m)
741 : COMPLEX(KIND=dp), DIMENSION(:, :), INTENT(IN), &
742 : TARGET :: pwbuf1
743 : COMPLEX(KIND=dp), DIMENSION(:, :), INTENT(INOUT), &
744 : TARGET :: pwbuf2
745 : INTEGER, INTENT(IN) :: dir, n, m
746 :
747 : CHARACTER(len=*), PARAMETER :: routineN = 'pw_gpu_f'
748 :
749 : COMPLEX(KIND=dp), POINTER :: ptr_pwin, ptr_pwout
750 : INTEGER :: handle
751 : INTERFACE
752 : SUBROUTINE pw_gpu_f_c(zin, zout, dir, n, m) BIND(C, name="pw_gpu_f")
753 : IMPORT
754 : TYPE(C_PTR), INTENT(IN), VALUE :: zin
755 : TYPE(C_PTR), VALUE :: zout
756 : INTEGER(KIND=C_INT), INTENT(IN), VALUE :: dir, n, m
757 : END SUBROUTINE pw_gpu_f_c
758 : END INTERFACE
759 :
760 0 : CALL timeset(routineN, handle)
761 :
762 0 : IF (n*m /= 0) THEN
763 : ! pointers to data arrays
764 0 : ptr_pwin => pwbuf1(1, 1)
765 0 : ptr_pwout => pwbuf2(1, 1)
766 :
767 : ! invoke the combined transformation
768 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_PW)
769 : CALL pw_gpu_f_c(c_loc(ptr_pwin), c_loc(ptr_pwout), dir, n, m)
770 : #else
771 : MARK_USED(dir)
772 0 : CPABORT("Compiled without pw offloading")
773 : #endif
774 : END IF
775 :
776 0 : CALL timestop(handle)
777 0 : END SUBROUTINE pw_gpu_f
778 : ! **************************************************************************************************
779 : !> \brief perform a parallel 1D-FFT followed by a gather on the gpu
780 : !> \param pwbuf ...
781 : !> \param pw2 ...
782 : !> \param scale ...
783 : !> \author Andreas Gloess
784 : ! **************************************************************************************************
785 0 : SUBROUTINE pw_gpu_fg(pwbuf, pw2, scale)
786 : COMPLEX(KIND=dp), DIMENSION(:, :), INTENT(IN), &
787 : TARGET :: pwbuf
788 : TYPE(pw_c1d_gs_type), INTENT(IN) :: pw2
789 : REAL(KIND=dp), INTENT(IN) :: scale
790 :
791 : CHARACTER(len=*), PARAMETER :: routineN = 'pw_gpu_fg'
792 :
793 : COMPLEX(KIND=dp), POINTER :: ptr_pwin, ptr_pwout
794 : INTEGER :: handle, mg, mmax, ngpts
795 0 : INTEGER, DIMENSION(:), POINTER :: npts
796 : INTEGER, POINTER :: ptr_ghatmap
797 : INTERFACE
798 : SUBROUTINE pw_gpu_fg_c(zin, zout, ghatmap, npts, mmax, ngpts, scale) BIND(C, name="pw_gpu_fg")
799 : IMPORT
800 : TYPE(C_PTR), INTENT(IN), VALUE :: zin
801 : TYPE(C_PTR), VALUE :: zout
802 : TYPE(C_PTR), INTENT(IN), VALUE :: ghatmap
803 : INTEGER(KIND=C_INT), DIMENSION(*), INTENT(IN):: npts
804 : INTEGER(KIND=C_INT), INTENT(IN), VALUE :: mmax, ngpts
805 : REAL(KIND=C_DOUBLE), INTENT(IN), VALUE :: scale
806 :
807 : END SUBROUTINE pw_gpu_fg_c
808 : END INTERFACE
809 :
810 0 : CALL timeset(routineN, handle)
811 :
812 0 : ngpts = SIZE(pw2%pw_grid%gsq)
813 0 : npts => pw2%pw_grid%npts
814 :
815 0 : IF ((npts(1) /= 0) .AND. (ngpts /= 0)) THEN
816 0 : mg = SIZE(pw2%pw_grid%grays, 2)
817 0 : mmax = MAX(mg, 1)
818 :
819 : ! pointers to data arrays
820 0 : ptr_pwin => pwbuf(1, 1)
821 0 : ptr_pwout => pw2%array(1)
822 :
823 : ! pointer to map array
824 0 : ptr_ghatmap => pw2%pw_grid%g_hatmap(1, 1)
825 :
826 : ! invoke the combined transformation
827 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_PW)
828 : CALL pw_gpu_fg_c(c_loc(ptr_pwin), c_loc(ptr_pwout), c_loc(ptr_ghatmap), npts, mmax, ngpts, scale)
829 : #else
830 : MARK_USED(scale)
831 0 : CPABORT("Compiled without pw offloading")
832 : #endif
833 : END IF
834 :
835 0 : CALL timestop(handle)
836 0 : END SUBROUTINE pw_gpu_fg
837 :
838 : ! **************************************************************************************************
839 : !> \brief perform a parallel scatter followed by a 1D-FFT on the gpu
840 : !> \param pw1 ...
841 : !> \param pwbuf ...
842 : !> \param scale ...
843 : !> \author Andreas Gloess
844 : ! **************************************************************************************************
845 0 : SUBROUTINE pw_gpu_sf(pw1, pwbuf, scale)
846 : TYPE(pw_c1d_gs_type), INTENT(IN) :: pw1
847 : COMPLEX(KIND=dp), DIMENSION(:, :), INTENT(INOUT), &
848 : TARGET :: pwbuf
849 : REAL(KIND=dp), INTENT(IN) :: scale
850 :
851 : CHARACTER(len=*), PARAMETER :: routineN = 'pw_gpu_sf'
852 :
853 : COMPLEX(KIND=dp), POINTER :: ptr_pwin, ptr_pwout
854 : INTEGER :: handle, mg, mmax, ngpts, nmaps
855 0 : INTEGER, DIMENSION(:), POINTER :: npts
856 : INTEGER, POINTER :: ptr_ghatmap
857 : INTERFACE
858 : SUBROUTINE pw_gpu_sf_c(zin, zout, ghatmap, npts, mmax, ngpts, nmaps, scale) BIND(C, name="pw_gpu_sf")
859 : IMPORT
860 : TYPE(C_PTR), INTENT(IN), VALUE :: zin
861 : TYPE(C_PTR), VALUE :: zout
862 : TYPE(C_PTR), INTENT(IN), VALUE :: ghatmap
863 : INTEGER(KIND=C_INT), DIMENSION(*), INTENT(IN):: npts
864 : INTEGER(KIND=C_INT), INTENT(IN), VALUE :: mmax, ngpts, nmaps
865 : REAL(KIND=C_DOUBLE), INTENT(IN), VALUE :: scale
866 :
867 : END SUBROUTINE pw_gpu_sf_c
868 : END INTERFACE
869 :
870 0 : CALL timeset(routineN, handle)
871 :
872 0 : ngpts = SIZE(pw1%pw_grid%gsq)
873 0 : npts => pw1%pw_grid%npts
874 :
875 0 : IF ((npts(1) /= 0) .AND. (ngpts /= 0)) THEN
876 0 : mg = SIZE(pw1%pw_grid%grays, 2)
877 0 : mmax = MAX(mg, 1)
878 :
879 : ! pointers to data arrays
880 0 : ptr_pwin => pw1%array(1)
881 0 : ptr_pwout => pwbuf(1, 1)
882 :
883 : ! pointer to map array
884 0 : nmaps = SIZE(pw1%pw_grid%g_hatmap, 2)
885 0 : ptr_ghatmap => pw1%pw_grid%g_hatmap(1, 1)
886 :
887 : ! invoke the combined transformation
888 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_PW)
889 : CALL pw_gpu_sf_c(c_loc(ptr_pwin), c_loc(ptr_pwout), c_loc(ptr_ghatmap), npts, mmax, ngpts, nmaps, scale)
890 : #else
891 : MARK_USED(scale)
892 0 : CPABORT("Compiled without pw offloading")
893 : #endif
894 : END IF
895 :
896 0 : CALL timestop(handle)
897 0 : END SUBROUTINE pw_gpu_sf
898 :
899 : END MODULE pw_gpu
900 :
|
