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 : !!***
|
