Branch data Line data Source code
1 : : !!****m* ABINIT/m_alloc_hamilt_gpu
2 : : !! NAME
3 : : !! m_alloc_hamilt_gpu
4 : : !!
5 : : !! FUNCTION
6 : : !!
7 : : !! COPYRIGHT
8 : : !! Copyright (C) 2000-2026 ABINIT group (MT, FDahm)
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_alloc_hamilt_gpu
22 : :
23 : : use defs_basis
24 : : use m_abicore
25 : : use m_errors
26 : : use m_xmpi
27 : : use m_dtset
28 : : use m_errors
29 : : use m_ompgpu_fourwf
30 : : #if defined HAVE_GPU
31 : : use m_gpu_toolbox
32 : : use m_gputk
33 : : #endif
34 : :
35 : : #ifdef HAVE_FC_ISO_C_BINDING
36 : : use, intrinsic :: iso_c_binding
37 : : #endif
38 : :
39 : : use defs_datatypes, only : pseudopotential_type
40 : : use defs_abitypes, only : MPI_type
41 : :
42 : : implicit none
43 : :
44 : : private
45 : : !!***
46 : :
47 : : public :: alloc_hamilt_gpu
48 : : public :: dealloc_hamilt_gpu
49 : : !!***
50 : :
51 : : private :: alloc_nonlop_gpu_data
52 : : private :: dealloc_nonlop_gpu_data
53 : :
54 : : !! data type to store pointers to data used on GPU, mostly in gemm nonlop_gpu
55 : : !! opernla/b/c
56 : : type, public :: gemm_nonlop_gpu_data_type
57 : :
58 : : logical :: allocated
59 : : type(c_ptr) :: projections_gpu
60 : : type(c_ptr) :: s_projections_gpu
61 : : type(c_ptr) :: vnl_projections_gpu
62 : :
63 : : type(c_ptr) :: vectin_gpu
64 : : type(c_ptr) :: vectout_gpu
65 : : type(c_ptr) :: svectout_gpu
66 : :
67 : : end type gemm_nonlop_gpu_data_type
68 : :
69 : : type(gemm_nonlop_gpu_data_type), save, public, target :: gemm_nonlop_gpu_data
70 : :
71 : : integer, save, public :: hamilt_gpu_nfft_blocks
72 : :
73 : : contains
74 : : !!***
75 : :
76 : : !!****f* ABINIT/alloc_hamilt_gpu
77 : : !! NAME
78 : : !! alloc_hamilt_gpu
79 : : !!
80 : : !! FUNCTION
81 : : !! allocate several memory pieces on a GPU device for the application
82 : : !! of Hamiltonian using a GPU
83 : : !!
84 : : !! INPUTS
85 : : !! atindx1(natom)=index table for atoms, inverse of atindx (see gstate.f)
86 : : !! dtset <type(dataset_type)>=all input variables for this dataset
87 : : !! gprimd(3,3)=dimensional reciprocal space primitive translations
88 : : !! mpi_enreg=information about MPI parallelization
89 : : !! nattyp(ntypat)= # atoms of each type.
90 : : !! option=0: allocate data for local operator (FFT)
91 : : !! 1: allocate data for nonlocal operator
92 : : !! 2: allocate both
93 : : !! psps <type(pseudopotential_type)>=variables related to pseudopotentials
94 : : !! gpu_option= GPU implementation to use, i.e. cuda, openMP, ... (0=not using GPU)
95 : : !!
96 : : !! OUTPUT
97 : : !! (no output - only allocation on GPU)
98 : : !!
99 : : !! SOURCE
100 : :
101 : 0 : subroutine alloc_hamilt_gpu(atindx1,dtset,gprimd,mpi_enreg,nattyp,npwarr,option,psps,gpu_option)
102 : :
103 : : !Arguments ------------------------------------
104 : : !scalars
105 : : integer,intent(in) :: option,gpu_option
106 : : type(dataset_type),intent(in) :: dtset
107 : : type(MPI_type),intent(in) :: mpi_enreg
108 : : type(pseudopotential_type),intent(in) :: psps
109 : : !arrays
110 : : integer,intent(in) :: atindx1(dtset%natom),nattyp(dtset%ntypat),npwarr(dtset%nkpt)
111 : : real(dp),intent(in) :: gprimd(3,3)
112 : :
113 : : !Local variables-------------------------------
114 : : !scalars
115 : : #if defined HAVE_GPU
116 : : integer :: dimekb1_max,dimekb2_max,dimffnl_max,ierr,ikpt,npw_max_loc,npw_max_nonloc
117 : : integer :: cplex
118 : : integer ::npwarr_tmp(dtset%nkpt)
119 : :
120 : : integer(kind=c_int32_t) :: proj_dim(3)
121 : : #endif
122 : :
123 : : ! *************************************************************************
124 : :
125 : : if (gpu_option==ABI_GPU_DISABLED) return
126 : :
127 : : #if defined(HAVE_GPU)
128 : : !=== Local Hamiltonian ===
129 : : if (option==0.or.option==2) then
130 : : ! Compute max of total planes waves
131 : : npw_max_loc=0
132 : : if(mpi_enreg%paral_kgb==1) then
133 : : npwarr_tmp=npwarr
134 : : call xmpi_sum(npwarr_tmp,mpi_enreg%comm_bandfft,ierr)
135 : : npw_max_loc =maxval(npwarr_tmp)
136 : : else
137 : : npw_max_loc=dtset%mpw
138 : : end if
139 : : ! Initialize gpu data needed in fourwf
140 : : ! ndat=bandpp when paral_kgb=1
141 : : ! no matter paral_kgb=0 or 1, we gathet all bands into a single call gpu_fourwf
142 : : if(gpu_option == ABI_GPU_LEGACY) then
143 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_FC_ISO_C_BINDING)
144 : : call alloc_gpu_fourwf(dtset%ngfft,dtset%bandpp,npw_max_loc,npw_max_loc)
145 : : #endif
146 : : else if (gpu_option == ABI_GPU_KOKKOS) then
147 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_FC_ISO_C_BINDING)
148 : : call alloc_gpu_fourwf_managed(dtset%ngfft,dtset%bandpp,npw_max_loc,npw_max_loc)
149 : : #endif
150 : : else if (gpu_option == ABI_GPU_OPENMP) then
151 : : !call alloc_ompgpu_fourwf(dtset%ngfft,dtset%bandpp)
152 : : end if
153 : :
154 : : end if
155 : : !=== Nonlocal Hamiltonian ===
156 : : if (option==1.or.option==2) then
157 : : ! Compute max of total planes waves
158 : : npw_max_nonloc=0
159 : : if(mpi_enreg%paral_kgb==1) then
160 : : npwarr_tmp=npwarr
161 : : call xmpi_sum(npwarr_tmp,mpi_enreg%comm_band,ierr)
162 : : npw_max_nonloc =maxval(npwarr_tmp)
163 : : else
164 : : npw_max_nonloc=dtset%mpw
165 : : end if
166 : : ! Initialize all gpu data needed in nonlop
167 : : dimffnl_max=4
168 : : ! if (abs(dtset%berryopt) == 5) dimffnl_max=4
169 : : dimekb1_max=psps%dimekb
170 : : if (dtset%usepaw==0) dimekb2_max=psps%ntypat
171 : : if (dtset%usepaw==1) dimekb2_max=dtset%natom
172 : :
173 : : if (gpu_option == ABI_GPU_KOKKOS .or. gpu_option == ABI_GPU_LEGACY) then
174 : :
175 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_FC_ISO_C_BINDING)
176 : : ! TODO (PK) : disable this allocation when Kokkos is available
177 : : ! to save memory on GPU side
178 : : call alloc_nonlop_gpu(npw_max_nonloc, &
179 : : & npw_max_nonloc,&
180 : : & dtset%nspinor,&
181 : : & dtset%natom,&
182 : : & dtset%ntypat,&
183 : : & psps%lmnmax,&
184 : : & psps%indlmn,&
185 : : & nattyp,&
186 : : & atindx1,&
187 : : & gprimd,&
188 : : & dimffnl_max,&
189 : : & dimffnl_max,&
190 : : & dimekb1_max,&
191 : : & dimekb2_max)
192 : :
193 : : if (dtset%use_gemm_nonlop == 1) then
194 : : call alloc_nonlop_gpu_data(dtset, &
195 : : & psps, &
196 : : & npw_max_nonloc,&
197 : : & npw_max_nonloc,&
198 : : & atindx1,&
199 : : & nattyp,&
200 : : & gpu_option)
201 : : end if
202 : : #endif
203 : :
204 : : end if
205 : :
206 : : end if ! option 1 or 2
207 : :
208 : : call xmpi_barrier(mpi_enreg%comm_cell)
209 : :
210 : : #else
211 : :
212 : : ABI_UNUSED(npwarr)
213 : : ABI_UNUSED_A(psps)
214 : : if (.false.) then
215 : : write(std_out,*) atindx1(1),dtset%natom,gprimd(1,1),mpi_enreg%me,nattyp(1),option
216 : : end if
217 : :
218 : : #endif
219 : :
220 : : end subroutine alloc_hamilt_gpu
221 : : !!***
222 : :
223 : : !!****f* ABINIT/dealloc_hamilt_gpu
224 : : !! NAME
225 : : !! dealloc_hamilt_gpu
226 : : !!
227 : : !! FUNCTION
228 : : !! deallocate several memory pieces on a GPU device used for the application
229 : : !! of Hamiltonian using a GPU
230 : : !!
231 : : !! INPUTS
232 : : !! option=0: deallocate data for local operator (FFT)
233 : : !! 1: deallocate data for nonlocal operator
234 : : !! 2: deallocate both
235 : : !! gpu_option= GPU implementation to use, i.e. cuda, openMP, ... (0=not using GPU)
236 : : !!
237 : : !! OUTPUT
238 : : !! (no output - only deallocation on GPU)
239 : : !!
240 : : !! SOURCE
241 : :
242 : 0 : subroutine dealloc_hamilt_gpu(option,gpu_option)
243 : :
244 : : !Arguments ------------------------------------
245 : : !scalars
246 : : integer,intent(in) :: option,gpu_option
247 : : !arrays
248 : :
249 : : !Local variables-------------------------------
250 : :
251 : : ! *************************************************************************
252 : :
253 [ # # ]: 0 : if (gpu_option==ABI_GPU_DISABLED) return
254 : :
255 [ # # ]: 0 : if (gpu_option == ABI_GPU_KOKKOS .or. gpu_option == ABI_GPU_LEGACY) then
256 : : #if defined(HAVE_GPU_CUDA) && defined(HAVE_FC_ISO_C_BINDING)
257 : : if (option==0.or.option==2) then
258 : : if (gpu_option == ABI_GPU_LEGACY) then
259 : : call free_gpu_fourwf()
260 : : else if(gpu_option == ABI_GPU_KOKKOS) then
261 : : call free_gpu_fourwf_managed()
262 : : end if
263 : : end if
264 : :
265 : : if (option==1.or.option==2) then
266 : : call free_nonlop_gpu()
267 : : call dealloc_nonlop_gpu_data()
268 : : end if ! option 1 or 2
269 : : #endif
270 [ # # ]: 0 : else if(gpu_option == ABI_GPU_OPENMP) then
271 : 0 : call free_ompgpu_fourwf()
272 : : end if
273 : :
274 : : if (.false.) then
275 : : write(std_out,*) option
276 : : end if
277 : :
278 : : end subroutine dealloc_hamilt_gpu
279 : : !!***
280 : :
281 : : !!****f* ABINIT/alloc_nonlop_gpu_data
282 : : !! NAME
283 : : !! alloc_hamilt_gpu
284 : : !!
285 : : !! FUNCTION
286 : : !! allocate several memory pieces on a GPU device for the application
287 : : !! of Hamiltonian using a GPU
288 : : !!
289 : : !! INPUTS
290 : : !! atindx1(natom)=index table for atoms, inverse of atindx (see gstate.f)
291 : : !! dtset <type(dataset_type)>=all input variables for this dataset
292 : : !! nattyp(ntypat)= # atoms of each type.
293 : : !! npwin is the number of plane waves for vectin
294 : : !! npwout is the number of plane waves for vectout
295 : : !! psps <type(pseudopotential_type)>=variables related to pseudopotentials
296 : : !! gpu_option= GPU implementation to use, i.e. cuda, openMP, ... (0=not using GPU)
297 : : !!
298 : : !! OUTPUT
299 : : !! (no output - only allocation on GPU, member of gemm_nonlop_gpu_data)
300 : : !!
301 : : !! SOURCE
302 : :
303 : : subroutine alloc_nonlop_gpu_data(dtset,&
304 : : & psps,&
305 : : & npwin,&
306 : : & npwout,&
307 : : & atindx1,&
308 : : & nattyp,&
309 : : & gpu_option)
310 : :
311 : : !Arguments ------------------------------------
312 : : !scalars
313 : : type(dataset_type), intent(in) :: dtset
314 : : type(pseudopotential_type),intent(in) :: psps
315 : : integer, intent(in) :: npwin
316 : : integer, intent(in) :: npwout
317 : : integer, intent(in) :: atindx1(dtset%natom)
318 : : !arrays
319 : : integer, intent(in) :: nattyp(dtset%ntypat)
320 : : !integer, intent(in) :: npwarr(dtset%nkpt)
321 : :
322 : : ! other
323 : : integer, intent(in) :: gpu_option
324 : :
325 : : !Local variables-------------------------------
326 : : !scalars
327 : : #if defined HAVE_GPU_CUDA
328 : : integer :: cplex
329 : : integer :: nprojs
330 : : integer :: itypat
331 : :
332 : : real(dp) :: allocated_size_bytes
333 : : character(len=500) :: message
334 : : #endif
335 : :
336 : : ! *************************************************************************
337 : :
338 : : #if defined HAVE_GPU_CUDA
339 : : allocated_size_bytes = 0.
340 : :
341 : : cplex=2; !if (istwf_k>1) cplex=1
342 : :
343 : : ! compute nprojs
344 : : nprojs = 0
345 : : do itypat = 1,dtset%ntypat
346 : : nprojs = nprojs + count(psps%indlmn(3,:,itypat)>0) * nattyp(itypat)
347 : : end do
348 : :
349 : :
350 : : !! allocate memory on device
351 : :
352 : : if (gemm_nonlop_gpu_data % allocated .eqv. .false.) then
353 : : ! These will store the non-local factors for vectin, svectout and vectout respectively
354 : : ABI_MALLOC_CUDA(gemm_nonlop_gpu_data% projections_gpu, INT(cplex, c_size_t) * nprojs * dtset%nspinor*dtset%bandpp * dp)
355 : : ABI_MALLOC_CUDA(gemm_nonlop_gpu_data% s_projections_gpu, INT(cplex, c_size_t) * nprojs * dtset%nspinor*dtset%bandpp * dp)
356 : : ABI_MALLOC_CUDA(gemm_nonlop_gpu_data%vnl_projections_gpu, INT(2 , c_size_t) * nprojs * dtset%nspinor*dtset%bandpp * dp)
357 : :
358 : : allocated_size_bytes = allocated_size_bytes + (2*cplex+2)*nprojs * dtset%nspinor*dtset%bandpp * dp
359 : :
360 : : ABI_MALLOC_CUDA(gemm_nonlop_gpu_data% vectin_gpu, INT(2, c_size_t) * npwin * dtset%nspinor*dtset%bandpp * dp)
361 : : ABI_MALLOC_CUDA(gemm_nonlop_gpu_data% vectout_gpu, INT(2, c_size_t) * npwout * dtset%nspinor*dtset%bandpp * dp)
362 : : ABI_MALLOC_CUDA(gemm_nonlop_gpu_data% svectout_gpu, INT(2, c_size_t) * npwout * dtset%nspinor*dtset%bandpp * dp)
363 : :
364 : : allocated_size_bytes = allocated_size_bytes + &
365 : : & 2 * (npwin+2*npwout) * dtset%nspinor * dtset%bandpp * dp
366 : :
367 : : gemm_nonlop_gpu_data % allocated = .true.
368 : :
369 : : write(message,*) ' alloc_nonlop_gpu_data allocated ', allocated_size_bytes*1e-6, ' MBytes on device memory'
370 : : call wrtout(std_out,message,'COLL')
371 : :
372 : : end if
373 : :
374 : : #else
375 : :
376 : : if (.false.) then
377 : : write(std_out,*) psps%indlmn(1,1,1),dtset%natom,npwin,npwout,atindx1(1),nattyp(1),gpu_option
378 : : end if
379 : :
380 : : #endif
381 : :
382 : : end subroutine alloc_nonlop_gpu_data
383 : : !!***
384 : :
385 : : !!****f* ABINIT/dealloc_nonlop_gpu_data
386 : : !! NAME
387 : : !! dealloc_nonlop_gpu_data
388 : : !!
389 : : !! FUNCTION
390 : : !! deallocate several memory pieces on a GPU device used for the application
391 : : !! of nonlop operators using Kokkos implementation
392 : : !!
393 : : !!
394 : : !! OUTPUT
395 : : !! (no output - only deallocation on GPU)
396 : : !!
397 : : !! SOURCE
398 : :
399 : : subroutine dealloc_nonlop_gpu_data()
400 : :
401 : : #if defined HAVE_GPU_CUDA
402 : : if (gemm_nonlop_gpu_data % allocated) then
403 : : ABI_FREE_CUDA(gemm_nonlop_gpu_data% projections_gpu)
404 : : ABI_FREE_CUDA(gemm_nonlop_gpu_data% s_projections_gpu)
405 : : ABI_FREE_CUDA(gemm_nonlop_gpu_data%vnl_projections_gpu)
406 : :
407 : : ABI_FREE_CUDA(gemm_nonlop_gpu_data% vectin_gpu)
408 : : ABI_FREE_CUDA(gemm_nonlop_gpu_data% vectout_gpu)
409 : : ABI_FREE_CUDA(gemm_nonlop_gpu_data%svectout_gpu)
410 : :
411 : : gemm_nonlop_gpu_data % allocated = .false.
412 : : end if
413 : : #endif
414 : :
415 : : end subroutine dealloc_nonlop_gpu_data
416 : : !!***
417 : :
418 : 0 : end module m_alloc_hamilt_gpu
419 : : !!***
|
