(git:3f0d10b)
Loading...
Searching...
No Matches
nnp_gpu.F
Go to the documentation of this file.
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!> \brief GPU offload driver for the NNP force evaluation. Hosts the
10!> ISO_C_BINDING interfaces to the CUDA launchers in src/nnp/gpu/ and
11!> the whole-step device pipeline called from nnp_force. All GPU
12!> entry points compile to inert stubs unless the build enables the
13!> NNP GPU backend (__OFFLOAD without __NO_OFFLOAD_NNP).
14!> \author Dhruv Sharma (ds2173@cam.ac.uk)
15! **************************************************************************************************
16MODULE nnp_gpu
17
18 USE iso_c_binding, ONLY: c_char,&
19 c_double,&
20 c_int,&
21 c_null_ptr
22 USE input_constants, ONLY: nnp_gpu_prec_mixed
23 USE kinds, ONLY: dp
24 USE nnp_acsf, ONLY: nnp_build_radial_splines
25 USE nnp_environment_types, ONLY: nnp_type
26 USE offload_api, ONLY: offload_activate_chosen_device,&
27 offload_get_device_count,&
28 offload_init
29#include "./base/base_uses.f90"
30
31 IMPLICIT NONE
32 PRIVATE
33
34 CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'nnp_gpu'
35
36 ! Initial per-element neighbour-list width for the device walk; grown and
37 ! retried on overflow, then reused across MD steps.
38 INTEGER, PARAMETER, PRIVATE :: gpu_walk_init_max_n_neigh = 96
39
40 ! nnp_gpu_run_walk, nnp_gpu_build_descriptor and nnp_gpu_predict stay private:
41 ! each is only valid after the previous one has run for this step, and
42 ! nnp_gpu_energy_force is the entry point that guarantees that order.
43 PUBLIC :: nnp_gpu_available
44 PUBLIC :: nnp_gpu_init_device
45 PUBLIC :: nnp_gpu_model_capacity_ok
46 PUBLIC :: nnp_gpu_upload_descriptor_statics
47 PUBLIC :: nnp_gpu_upload_network
48 PUBLIC :: nnp_gpu_energy_force
49
50 ! Every device entry point below is compiled only for accelerator builds.
51 ! The guard tests the backend macros the build system defines globally
52 ! (__OFFLOAD_CUDA / __OFFLOAD_HIP) alongside the __OFFLOAD umbrella, so it
53 ! reads the same in every translation unit that carries it, including
54 ! cp2k_info.F, which does not include base_uses.f90 and therefore never
55 ! sees a Fortran-side __OFFLOAD definition.
56#if (defined(__OFFLOAD) || defined(__OFFLOAD_CUDA) || defined(__OFFLOAD_HIP)) && !defined(__NO_OFFLOAD_NNP)
57 INTERFACE
58 FUNCTION nnp_gpu_set_precision_c(use_fp32) &
59 result(rc) BIND(C, NAME="nnp_gpu_set_precision")
60 USE iso_c_binding, ONLY: c_int
61 INTEGER(KIND=c_int), VALUE :: use_fp32
62 INTEGER(KIND=c_int) :: rc
63
64 END FUNCTION nnp_gpu_set_precision_c
65
66 ! Counts this environment as a user of the per-process device state and
67 ! returns the number now live.
68 FUNCTION nnp_gpu_register_env_c() &
69 result(n_live) BIND(C, NAME="nnp_gpu_register_env")
70 USE iso_c_binding, ONLY: c_int
71 INTEGER(KIND=c_int) :: n_live
72
73 END FUNCTION nnp_gpu_register_env_c
74
75 ! Undoes a registration the caller decided not to honour, leaving the
76 ! state the owning environment is using untouched.
77 SUBROUTINE nnp_gpu_unregister_env_c() &
78 BIND(C, NAME="nnp_gpu_unregister_env")
79 END SUBROUTINE nnp_gpu_unregister_env_c
80
81 ! Compiled-in capacities of the per-(group, element) device buffers and
82 ! of the run-constant parameter cache, for the model check at init.
83 FUNCTION nnp_gpu_max_se_c() &
84 result(cap) BIND(C, NAME="nnp_gpu_max_se_c")
85 USE iso_c_binding, ONLY: c_int
86 INTEGER(KIND=c_int) :: cap
87
88 END FUNCTION nnp_gpu_max_se_c
89
90 FUNCTION nnp_gpu_spline_cache_max_c() &
91 result(cap) BIND(C, NAME="nnp_gpu_spline_cache_max_c")
92 USE iso_c_binding, ONLY: c_int
93 INTEGER(KIND=c_int) :: cap
94
95 END FUNCTION nnp_gpu_spline_cache_max_c
96
97 FUNCTION nnp_gpu_upload_positions_step_c(num_atoms, h_coord) &
98 result(rc) BIND(C, NAME="nnp_gpu_upload_positions_step_c")
99 USE iso_c_binding, ONLY: c_double, c_int
100 INTEGER(KIND=c_int), VALUE :: num_atoms
101 REAL(KIND=c_double), DIMENSION(*), INTENT(IN) :: h_coord
102 INTEGER(KIND=c_int) :: rc
103
104 END FUNCTION nnp_gpu_upload_positions_step_c
105
106 FUNCTION nnp_gpu_upload_image_table_c(n_images, h_image_table) &
107 result(rc) BIND(C, NAME="nnp_gpu_upload_image_table_c")
108 USE iso_c_binding, ONLY: c_double, c_int
109 INTEGER(KIND=c_int), VALUE :: n_images
110 REAL(KIND=c_double), DIMENSION(*), INTENT(IN) :: h_image_table
111 INTEGER(KIND=c_int) :: rc
112
113 END FUNCTION nnp_gpu_upload_image_table_c
114
115 FUNCTION nnp_gpu_upload_cell_list_c(n_images, n_cells, h_image_atom, &
116 h_image_shift, h_head, h_next) &
117 result(rc) BIND(C, NAME="nnp_gpu_upload_cell_list_c")
118 USE iso_c_binding, ONLY: c_int
119 INTEGER(KIND=c_int), VALUE :: n_images, n_cells
120 INTEGER(KIND=c_int), DIMENSION(*), INTENT(IN) :: h_image_atom, h_image_shift, h_head, &
121 h_next
122 INTEGER(KIND=c_int) :: rc
123
124 END FUNCTION nnp_gpu_upload_cell_list_c
125
126 FUNCTION nnp_gpu_upload_pair_map_c(n_ele, max_pair_grps, max_n_radgrp, &
127 max_n_anggrp, h_n_rad, h_n_ang1, h_n_ang2, &
128 h_pair_max_cutoff, h_rad_groups, h_ang1_groups, &
129 h_ang2_groups, h_rad_cutoff, h_ang_cutoff) &
130 result(rc) BIND(C, NAME="nnp_gpu_upload_pair_map_c")
131 USE iso_c_binding, ONLY: c_double, c_int
132 INTEGER(KIND=c_int), VALUE :: n_ele, max_pair_grps, max_n_radgrp, &
133 max_n_anggrp
134 INTEGER(KIND=c_int), DIMENSION(*), INTENT(IN) :: h_n_rad, h_n_ang1, h_n_ang2
135 REAL(KIND=c_double), DIMENSION(*), INTENT(IN) :: h_pair_max_cutoff
136 INTEGER(KIND=c_int), DIMENSION(*), INTENT(IN) :: h_rad_groups, h_ang1_groups, &
137 h_ang2_groups
138 REAL(KIND=c_double), DIMENSION(*), INTENT(IN) :: h_rad_cutoff, h_ang_cutoff
139 INTEGER(KIND=c_int) :: rc
140
141 END FUNCTION nnp_gpu_upload_pair_map_c
142
143 FUNCTION nnp_gpu_cell_list_walk_c(num_atoms_global, n_ele, n_images, n_cells, &
144 n_atoms_in_e_max, max_n_neigh_rad, &
145 max_n_neigh_ang1, max_n_neigh_ang2, max_n_radgrp, &
146 max_n_anggrp, h_nbin, h_bin_lower, h_bin_width, &
147 h_bin_span, h_perd, h_exact_pbc_copies, &
148 h_coord_scaled, h_ele_ind_global, &
149 h_atoms_global_per_e, h_n_atoms_in_e) &
150 result(rc) BIND(C, NAME="nnp_gpu_cell_list_walk_c")
151 USE iso_c_binding, ONLY: c_double, c_int
152 INTEGER(KIND=c_int), VALUE :: num_atoms_global, n_ele, n_images, n_cells, n_atoms_in_e_max, &
153 max_n_neigh_rad, max_n_neigh_ang1, max_n_neigh_ang2, max_n_radgrp, max_n_anggrp
154 INTEGER(KIND=c_int), DIMENSION(*), INTENT(IN) :: h_nbin
155 REAL(KIND=c_double), DIMENSION(*), INTENT(IN) :: h_bin_lower, h_bin_width
156 INTEGER(KIND=c_int), DIMENSION(*), INTENT(IN) :: h_bin_span, h_perd, h_exact_pbc_copies
157 REAL(KIND=c_double), DIMENSION(*), INTENT(IN) :: h_coord_scaled
158 INTEGER(KIND=c_int), DIMENSION(*), INTENT(IN) :: h_ele_ind_global, h_atoms_global_per_e, &
159 h_n_atoms_in_e
160 INTEGER(KIND=c_int) :: rc
161
162 END FUNCTION nnp_gpu_cell_list_walk_c
163
164 FUNCTION nnp_gpu_walk_counts_d2h_c(n_atoms_in_e_max, max_n_radgrp, max_n_anggrp, &
165 n_ele, h_n_rad_pa, h_n_ang1_pa, h_n_ang2_pa) &
166 result(rc) BIND(C, NAME="nnp_gpu_walk_counts_d2h_c")
167 USE iso_c_binding, ONLY: c_int
168 INTEGER(KIND=c_int), VALUE :: n_atoms_in_e_max, max_n_radgrp, &
169 max_n_anggrp, n_ele
170 INTEGER(KIND=c_int), DIMENSION(*), INTENT(INOUT) :: h_n_rad_pa, h_n_ang1_pa, h_n_ang2_pa
171 INTEGER(KIND=c_int) :: rc
172
173 END FUNCTION nnp_gpu_walk_counts_d2h_c
174
175 FUNCTION nnp_gpu_begin_descriptor_step_c(n_atoms, n_input_max) &
176 result(rc) BIND(C, NAME="nnp_gpu_begin_descriptor_step")
177 USE iso_c_binding, ONLY: c_int
178 INTEGER(KIND=c_int), VALUE :: n_atoms, n_input_max
179 INTEGER(KIND=c_int) :: rc
180
181 END FUNCTION nnp_gpu_begin_descriptor_step_c
182
183 FUNCTION nnp_gpu_upload_scaling_tables_c(n_ele, n_input_max, h_loc_av, h_loc_min, &
184 h_loc_max, h_sigma, h_n_input_nodes) &
185 result(rc) BIND(C, NAME="nnp_gpu_upload_scaling_tables_c")
186 USE iso_c_binding, ONLY: c_double, c_int
187 INTEGER(KIND=c_int), VALUE :: n_ele, n_input_max
188 REAL(KIND=c_double), DIMENSION(*), INTENT(IN) :: h_loc_av, h_loc_min, h_loc_max, h_sigma
189 INTEGER(KIND=c_int), DIMENSION(*), INTENT(IN) :: h_n_input_nodes
190 INTEGER(KIND=c_int) :: rc
191
192 END FUNCTION nnp_gpu_upload_scaling_tables_c
193
194 FUNCTION nnp_gpu_upload_se_static_tables_c(n_ele, max_n_radgrp, max_n_anggrp, &
195 total_rad_symf, total_ang_symf, &
196 h_rad_packed, h_rad_offsets, &
197 h_ang_packed, h_ang_offsets) &
198 result(rc) BIND(C, NAME="nnp_gpu_upload_se_static_tables_c")
199 USE iso_c_binding, ONLY: c_int
200 INTEGER(KIND=c_int), VALUE :: n_ele, max_n_radgrp, max_n_anggrp, &
201 total_rad_symf, total_ang_symf
202 INTEGER(KIND=c_int), DIMENSION(*), INTENT(IN) :: h_rad_packed, h_rad_offsets, &
203 h_ang_packed, h_ang_offsets
204 INTEGER(KIND=c_int) :: rc
205
206 END FUNCTION nnp_gpu_upload_se_static_tables_c
207
208 FUNCTION nnp_gpu_scale_g_dev_c(n_atoms, n_input_max, n_ele, scaling_mode, &
209 scmin, scmax, h_ele_ind_for_atoms) &
210 result(rc) BIND(C, NAME="nnp_gpu_scale_G_dev_c")
211 USE iso_c_binding, ONLY: c_double, c_int
212 INTEGER(KIND=c_int), VALUE :: n_atoms, n_input_max, n_ele, scaling_mode
213 REAL(KIND=c_double), VALUE :: scmin, scmax
214 INTEGER(KIND=c_int), DIMENSION(*), INTENT(IN) :: h_ele_ind_for_atoms
215 INTEGER(KIND=c_int) :: rc
216
217 END FUNCTION nnp_gpu_scale_g_dev_c
218
219 ! Whether the last scaling pass saw a descriptor outside its training
220 ! range; valid once the step's device work has been synchronized.
221 FUNCTION nnp_gpu_extrapolation_flag_c() &
222 result(flag) BIND(C, NAME="nnp_gpu_extrapolation_flag_c")
223 USE iso_c_binding, ONLY: c_int
224 INTEGER(KIND=c_int) :: flag
225
226 END FUNCTION nnp_gpu_extrapolation_flag_c
227
228 FUNCTION nnp_gpu_radial_sf_to_g_dev_c(n_atoms_in_e, max_n_neigh_per_se, n_symf, &
229 spline_n, spline_dx, spline_dx_inv, &
230 spline_x_max, h_spline_y, h_spline_dy, &
231 s_0based, e_0based, max_n_radgrp_walk, &
232 n_atoms_in_e_max_walk, max_n_neigh_walk, &
233 istart_1based, n_input_max, se_idx) &
234 result(rc) BIND(C, NAME="nnp_gpu_radial_sf_to_G_dev_c")
235 USE iso_c_binding, ONLY: c_double, c_int
236 INTEGER(KIND=c_int), VALUE :: n_atoms_in_e, max_n_neigh_per_se, &
237 n_symf, spline_n
238 REAL(KIND=c_double), VALUE :: spline_dx, spline_dx_inv, spline_x_max
239 REAL(KIND=c_double), DIMENSION(*), INTENT(IN) :: h_spline_y, h_spline_dy
240 INTEGER(KIND=c_int), VALUE :: s_0based, e_0based, max_n_radgrp_walk, n_atoms_in_e_max_walk, &
241 max_n_neigh_walk, istart_1based, n_input_max, se_idx
242 INTEGER(KIND=c_int) :: rc
243
244 END FUNCTION nnp_gpu_radial_sf_to_g_dev_c
245
246 FUNCTION nnp_gpu_angular_sf_to_g_dev_c(homo, cut_type, max_n_ang1_per_se, &
247 max_n_ang2_per_se, max_n_kk_per_se, n_symf, &
248 n_atoms_in_e, cutoff_s, cutoff_sqr, &
249 h_pack_lam, h_pack_zeta, h_pack_eta, &
250 h_pack_prefzeta, h_pack_izeta, &
251 h_pack_use_int_zeta, s_0based, e_0based, &
252 max_n_anggrp_walk, n_atoms_in_e_max_walk, &
253 max_n_neigh_ang1_walk, max_n_neigh_ang2_walk, &
254 istart_1based, n_rad_for_e, n_input_max, se_idx) &
255 result(rc) BIND(C, NAME="nnp_gpu_angular_sf_to_G_dev_c")
256 USE iso_c_binding, ONLY: c_char, c_double, c_int
257 INTEGER(KIND=c_int), VALUE :: homo, cut_type, max_n_ang1_per_se, &
258 max_n_ang2_per_se, max_n_kk_per_se, &
259 n_symf, n_atoms_in_e
260 REAL(KIND=c_double), VALUE :: cutoff_s, cutoff_sqr
261 REAL(KIND=c_double), DIMENSION(*), INTENT(IN) :: h_pack_lam, h_pack_zeta, h_pack_eta, &
262 h_pack_prefzeta
263 INTEGER(KIND=c_int), DIMENSION(*), INTENT(IN) :: h_pack_izeta
264 CHARACTER(KIND=c_char), DIMENSION(*), INTENT(IN) :: h_pack_use_int_zeta
265 INTEGER(KIND=c_int), VALUE :: s_0based, e_0based, max_n_anggrp_walk, n_atoms_in_e_max_walk, &
266 max_n_neigh_ang1_walk, max_n_neigh_ang2_walk, istart_1based, n_rad_for_e, n_input_max, &
267 se_idx
268 INTEGER(KIND=c_int) :: rc
269
270 END FUNCTION nnp_gpu_angular_sf_to_g_dev_c
271
272 FUNCTION nnp_gpu_upload_nn_packed_c(n_committee, n_layer, max_n_nodes, n_ele, &
273 h_n_nodes, h_act_fnct, h_atom_energies, &
274 h_W_packed, h_b_packed) &
275 result(rc) BIND(C, NAME="nnp_gpu_upload_nn_packed_c")
276 USE iso_c_binding, ONLY: c_double, c_int
277 INTEGER(KIND=c_int), VALUE :: n_committee, n_layer, max_n_nodes, n_ele
278 INTEGER(KIND=c_int), DIMENSION(*), INTENT(IN) :: h_n_nodes, h_act_fnct
279 REAL(KIND=c_double), DIMENSION(*), INTENT(IN) :: h_atom_energies, h_w_packed, h_b_packed
280 INTEGER(KIND=c_int) :: rc
281
282 END FUNCTION nnp_gpu_upload_nn_packed_c
283
284 FUNCTION nnp_gpu_nn_atomwise_with_g_dev_c(n_atoms, n_committee, n_layer, max_n_nodes, &
285 n_input_max, normnodes, n_ele, h_ele_ind, &
286 h_atomic_energy, h_dE_dG) &
287 result(rc) BIND(C, NAME="nnp_gpu_nn_atomwise_with_G_dev_c")
288 USE iso_c_binding, ONLY: c_double, c_int, c_ptr
289 INTEGER(KIND=c_int), VALUE :: n_atoms, n_committee, n_layer, &
290 max_n_nodes, n_input_max, normnodes, &
291 n_ele
292 INTEGER(KIND=c_int), DIMENSION(*), INTENT(IN) :: h_ele_ind
293 REAL(KIND=c_double), DIMENSION(*), INTENT(INOUT) :: h_atomic_energy
294 TYPE(c_ptr), VALUE :: h_dE_dG
295 INTEGER(KIND=c_int) :: rc
296
297! Optional descriptor-gradient read-back; C_NULL_PTR leaves the
298! gradient on the device for the force scatter to consume.
299
300 END FUNCTION nnp_gpu_nn_atomwise_with_g_dev_c
301
302 FUNCTION nnp_gpu_force_scatter_begin_c(num_atoms_global, n_committee, h_force) &
303 result(rc) BIND(C, NAME="nnp_gpu_force_scatter_begin_c")
304 USE iso_c_binding, ONLY: c_double, c_int
305 INTEGER(KIND=c_int), VALUE :: num_atoms_global, n_committee
306 REAL(KIND=c_double), DIMENSION(*), INTENT(IN) :: h_force
307 INTEGER(KIND=c_int) :: rc
308
309 END FUNCTION nnp_gpu_force_scatter_begin_c
310
311 FUNCTION nnp_gpu_force_scatter_end_c(num_atoms_global, n_committee, h_force) &
312 result(rc) BIND(C, NAME="nnp_gpu_force_scatter_end_c")
313 USE iso_c_binding, ONLY: c_double, c_int
314 INTEGER(KIND=c_int), VALUE :: num_atoms_global, n_committee
315 REAL(KIND=c_double), DIMENSION(*), INTENT(INOUT) :: h_force
316 INTEGER(KIND=c_int) :: rc
317
318 END FUNCTION nnp_gpu_force_scatter_end_c
319
320 ! D2H the per-committee analytical virial accumulator (9 * n_committee
321 ! doubles, layout [9*c + 3*a + b]). Call after nnp_gpu_force_scatter_end_c,
322 ! whose stream sync has committed the scatter kernels' atomic adds.
323 FUNCTION nnp_gpu_force_scatter_virial_d2h_c(n_committee, h_virial) &
324 result(rc) BIND(C, NAME="nnp_gpu_force_scatter_virial_d2h_c")
325 USE iso_c_binding, ONLY: c_double, c_int
326 INTEGER(KIND=c_int), VALUE :: n_committee
327 REAL(KIND=c_double), DIMENSION(*), INTENT(INOUT) :: h_virial
328 INTEGER(KIND=c_int) :: rc
329
330 END FUNCTION nnp_gpu_force_scatter_virial_d2h_c
331
332 FUNCTION nnp_gpu_force_scatter_radial_c(n_atoms_in_e, n_symf, max_n_neigh_per_se, &
333 n_committee, n_input_max, num_atoms_global, &
334 ele_idx, scaling_mode, scmin, scmax, s_0based, &
335 max_n_radgrp_walk, n_atoms_in_e_max_walk, &
336 max_n_neigh_walk, istart_1based, se_idx, &
337 accumulate_virial) &
338 result(rc) BIND(C, NAME="nnp_gpu_force_scatter_radial_c")
339 USE iso_c_binding, ONLY: c_double, c_int
340 INTEGER(KIND=c_int), VALUE :: n_atoms_in_e, n_symf, &
341 max_n_neigh_per_se, n_committee, &
342 n_input_max, num_atoms_global, &
343 ele_idx, scaling_mode
344 REAL(KIND=c_double), VALUE :: scmin, scmax
345 INTEGER(KIND=c_int), VALUE :: s_0based, max_n_radgrp_walk, n_atoms_in_e_max_walk, &
346 max_n_neigh_walk, istart_1based, se_idx, accumulate_virial
347 INTEGER(KIND=c_int) :: rc
348
349 END FUNCTION nnp_gpu_force_scatter_radial_c
350
351 FUNCTION nnp_gpu_force_scatter_angular_c(n_atoms_in_e, n_symf, max_n_ang1_per_se, &
352 max_n_kk_per_se, n_committee, n_input_max, &
353 num_atoms_global, ele_idx, n_rad_offset, &
354 scaling_mode, homo, scmin, scmax, s_0based, &
355 max_n_anggrp_walk, n_atoms_in_e_max_walk, &
356 max_n_neigh_ang1_walk, max_n_neigh_ang2_walk, &
357 istart_1based, se_idx, accumulate_virial) &
358 result(rc) BIND(C, NAME="nnp_gpu_force_scatter_angular_c")
359 USE iso_c_binding, ONLY: c_double, c_int
360 INTEGER(KIND=c_int), VALUE :: n_atoms_in_e, n_symf, max_n_ang1_per_se, max_n_kk_per_se, &
361 n_committee, n_input_max, num_atoms_global, ele_idx, n_rad_offset, scaling_mode, homo
362 REAL(KIND=c_double), VALUE :: scmin, scmax
363 INTEGER(KIND=c_int), VALUE :: s_0based, max_n_anggrp_walk, n_atoms_in_e_max_walk, &
364 max_n_neigh_ang1_walk, max_n_neigh_ang2_walk, istart_1based, se_idx, accumulate_virial
365 INTEGER(KIND=c_int) :: rc
366
367 END FUNCTION nnp_gpu_force_scatter_angular_c
368 END INTERFACE
369#endif
370
371CONTAINS
372
373! **************************************************************************************************
374!> \brief Report whether the NNP GPU backend can run: the build must enable
375!> it and at least one offload device must be visible. Builds without
376!> the backend always report .FALSE., so USE_GPU AUTO resolves to the
377!> CPU path there.
378!> \return .TRUE. when the GPU code path is usable on this rank
379! **************************************************************************************************
380 FUNCTION nnp_gpu_available() RESULT(available)
381 LOGICAL :: available
382
383#if (defined(__OFFLOAD) || defined(__OFFLOAD_CUDA) || defined(__OFFLOAD_HIP)) && !defined(__NO_OFFLOAD_NNP)
384 CALL offload_init()
385 available = (offload_get_device_count() >= 1)
386#else
387 available = .false.
388#endif
389
390 END FUNCTION nnp_gpu_available
391
392! **************************************************************************************************
393!> \brief Claim the offload device for this NNP environment and push the
394!> requested compute precision to the device side.
395!>
396!> The device itself is chosen once per process at startup (see
397!> f77_interface), which is also the id DBCSR is given; this routine only
398!> activates it, exactly as every other CP2K offload consumer does.
399!> Re-selecting here would leave the two out of step for the rest of the
400!> run.
401!> \param gpu_precision resolved GPU_PRECISION keyword value
402!> \param accepted .FALSE. when another NNP environment already holds the
403!> device, in which case nothing was initialized and the caller
404!> must fall back to the CPU path
405! **************************************************************************************************
406 SUBROUTINE nnp_gpu_init_device(gpu_precision, accepted)
407 INTEGER, INTENT(IN) :: gpu_precision
408 LOGICAL, INTENT(OUT) :: accepted
409
410 INTEGER :: handle
411
412#if (defined(__OFFLOAD) || defined(__OFFLOAD_CUDA) || defined(__OFFLOAD_HIP)) && !defined(__NO_OFFLOAD_NNP)
413 CHARACTER(len=*), PARAMETER :: routinen = 'nnp_gpu_init_device'
414 INTEGER :: ierr
415
416 CALL timeset(routinen, handle)
417
418 CALL offload_activate_chosen_device()
419
420 ! The device state is one per process, and its once-per-run uploads
421 ! (network weights, pair map, scaling and symmetry-function tables) are
422 ! keyed by shape rather than by environment. Two environments live at the
423 ! same time and sharing an architecture would therefore evaluate the first
424 ! one's weights against the second one's descriptor tables, with no
425 ! diagnostic. Decline the device for the second one instead and let it run
426 ! on the CPU. Sequential environments are fine: releasing one resets the
427 ! upload flags.
428 IF (nnp_gpu_register_env_c() > 1) THEN
429 CALL nnp_gpu_unregister_env_c()
430 accepted = .false.
431 CALL timestop(handle)
432 RETURN
433 END IF
434 accepted = .true.
435
436 ierr = nnp_gpu_set_precision_c(merge(1, 0, gpu_precision == nnp_gpu_prec_mixed))
437 IF (ierr /= 0) THEN
438 cpabort("NNP GPU init failed: cannot push the precision selector to the device.")
439 END IF
440
441 CALL timestop(handle)
442#else
443 mark_used(gpu_precision)
444 accepted = .false.
445 mark_used(handle)
446 cpabort("NNP GPU support is not compiled into this binary.")
447#endif
448
449 END SUBROUTINE nnp_gpu_init_device
450
451! **************************************************************************************************
452!> \brief Test whether the model fits the compiled-in device-table capacities:
453!> NNP_MAX_SE (group, element) slots per symmetry-function kind, and
454!> NNP_SPLINE_CACHE_MAX run-constant parameter tables, of which each
455!> radial group holds two (spline values and derivatives) and each
456!> angular group four (the packed parameter tables). Called before the
457!> device is claimed, so an oversized model is turned away before any
458!> upload rather than in the first force evaluation.
459!> \param nnp the NNP environment with its radial and angular groups populated
460!> \param reason on a .FALSE. result, a phrase naming the exceeded capacity
461!> \return .TRUE. when the model fits the device tables
462! **************************************************************************************************
463 FUNCTION nnp_gpu_model_capacity_ok(nnp, reason) RESULT(fits)
464 TYPE(nnp_type), INTENT(IN), POINTER :: nnp
465 CHARACTER(LEN=*), INTENT(OUT) :: reason
466 LOGICAL :: fits
467
468#if (defined(__OFFLOAD) || defined(__OFFLOAD_CUDA) || defined(__OFFLOAD_HIP)) && !defined(__NO_OFFLOAD_NNP)
469 INTEGER :: cache_needed, e, mag, mrg, se_cap
470
471 fits = .true.
472 reason = ""
473
474 CALL nnp_gpu_max_symfgrp(nnp, mrg, mag)
475 se_cap = nnp_gpu_max_se_c()
476 IF (mrg*nnp%n_ele > se_cap .OR. mag*nnp%n_ele > se_cap) THEN
477 fits = .false.
478 WRITE (reason, '(A,I0,A,I0)') "the model needs ", &
479 max(mrg*nnp%n_ele, mag*nnp%n_ele), &
480 " (group, element) device slots per kind and this build holds "// &
481 "NNP_MAX_SE = ", se_cap
482 RETURN
483 END IF
484
485 cache_needed = 0
486 DO e = 1, nnp%n_ele
487 cache_needed = cache_needed + 2*nnp%rad(e)%n_symfgrp + 4*nnp%ang(e)%n_symfgrp
488 END DO
489 IF (cache_needed > nnp_gpu_spline_cache_max_c()) THEN
490 fits = .false.
491 WRITE (reason, '(A,I0,A,I0)') "the model needs ", cache_needed, &
492 " device parameter-cache entries and this build holds "// &
493 "NNP_SPLINE_CACHE_MAX = ", nnp_gpu_spline_cache_max_c()
494 END IF
495#else
496 mark_used(nnp)
497 fits = .true.
498 reason = ""
499#endif
500
501 END FUNCTION nnp_gpu_model_capacity_ok
502
503! **************************************************************************************************
504!> \brief Largest per-element radial and angular symmetry-function group
505!> counts, floored to one. Every device-side (group, element) buffer
506!> shape and flattened slot index derives from these two strides, so
507!> every caller must take them from this one definition. The floor keeps
508!> the buffer shapes and kernel launch grids nonzero for a model with no
509!> groups of one kind; the padding slot does no work because the
510!> per-element group counts still drive the iteration.
511!> \param nnp the NNP environment with its radial and angular groups populated
512!> \param max_rad largest radial group count over the elements, at least one
513!> \param max_ang largest angular group count over the elements, at least one
514! **************************************************************************************************
515 SUBROUTINE nnp_gpu_max_symfgrp(nnp, max_rad, max_ang)
516 TYPE(nnp_type), INTENT(IN), POINTER :: nnp
517 INTEGER, INTENT(OUT) :: max_rad, max_ang
518
519 INTEGER :: e
520
521 max_rad = 1
522 max_ang = 1
523 DO e = 1, nnp%n_ele
524 max_rad = max(max_rad, nnp%rad(e)%n_symfgrp)
525 max_ang = max(max_ang, nnp%ang(e)%n_symfgrp)
526 END DO
527
528 END SUBROUTINE nnp_gpu_max_symfgrp
529
530! **************************************************************************************************
531!> \brief Build the device-resident neighbour lists for one MD step: upload the
532!> positions, image table and cell-list chain, pack and upload the
533!> species-pair routing table, then run the on-device cell-list walk over
534!> this rank's atoms and read back the per-(atom-slot, group, element)
535!> neighbour counts. On a per-list overflow the neighbour-list widths are
536!> doubled and the walk is retried.
537!> \param nnp the NNP environment (its cell-list cache and pair map must already
538!> be prepared for this step)
539!> \param istart first rank-local atom (1-based, global index)
540!> \param mecalc number of rank-local atoms
541! **************************************************************************************************
542 SUBROUTINE nnp_gpu_run_walk(nnp, istart, mecalc)
543 TYPE(nnp_type), INTENT(INOUT), POINTER :: nnp
544 INTEGER, INTENT(IN) :: istart, mecalc
545
546 INTEGER :: handle
547
548#if (defined(__OFFLOAD) || defined(__OFFLOAD_CUDA) || defined(__OFFLOAD_HIP)) && !defined(__NO_OFFLOAD_NNP)
549 CHARACTER(len=*), PARAMETER :: routinen = 'nnp_gpu_run_walk'
550 INTEGER :: attempts, e_w, ec, gg, ie, ii, je, mag, mpg, mrg, n_atoms_in_e_max, ne, &
551 pflat, rc_walk, slot_w, ss
552 INTEGER(KIND=c_int), DIMENSION(3) :: epbc_buf, nbin_buf, perd_buf, span_buf
553 REAL(kind=dp), DIMENSION(3) :: lower_buf, width_buf
554 INTEGER(KIND=c_int), ALLOCATABLE, DIMENSION(:) :: atoms_global_per_e_buf, &
555 ele_ind_global_buf, n_atoms_in_e_buf, &
556 pm_ang1_g, pm_ang2_g, pm_n_ang1, &
557 pm_n_ang2, pm_n_rad, pm_rad_g
558 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: pm_ang_cut, pm_max_cut, pm_rad_cut
559
560 CALL timeset(routinen, handle)
561
562 IF (.NOT. ALLOCATED(nnp%gpu_host_state)) ALLOCATE (nnp%gpu_host_state)
563
564 ne = nnp%n_ele
565 CALL nnp_gpu_max_symfgrp(nnp, mrg, mag)
566
567 ! Positions every step; the image table only when the translations have
568 ! been regenerated since the copy the device is holding. The comparison is
569 ! against the cache's rebuild counter rather than a per-step flag: this
570 ! routine is reached only on the steps that take the GPU path, while the
571 ! cache is rebuilt on any force evaluation, energy-only ones included.
572 ec = nnp_gpu_upload_positions_step_c(int(nnp%num_atoms, c_int), &
573 nnp%cell_list_cache%coord_primary)
574 IF (ec /= 0) cpabort("NNP GPU: position upload failed.")
575
576 IF (nnp%cell_list_cache%n_images > 0) THEN
577 IF (nnp%gpu_host_state%image_table_generation /= &
578 nnp%cell_list_cache%rebuild_count) THEN
579 ec = nnp_gpu_upload_image_table_c(int(nnp%cell_list_cache%n_images, c_int), &
580 nnp%cell_list_cache%image_translation)
581 IF (ec /= 0) cpabort("NNP GPU: image-table upload failed.")
582 nnp%gpu_host_state%image_table_generation = nnp%cell_list_cache%rebuild_count
583 END IF
584 IF (nnp%cell_list_cache%n_cells > 0) THEN
585 ec = nnp_gpu_upload_cell_list_c(int(nnp%cell_list_cache%n_images, c_int), &
586 int(nnp%cell_list_cache%n_cells, c_int), &
587 int(nnp%cell_list_cache%image_atom, c_int), &
588 int(nnp%cell_list_cache%image_shift, c_int), &
589 int(nnp%cell_list_cache%head, c_int), &
590 int(nnp%cell_list_cache%next, c_int))
591 IF (ec /= 0) cpabort("NNP GPU: cell-list upload failed.")
592 END IF
593 END IF
594
595 ! Pack and upload the species-pair routing table. The launcher performs
596 ! the host-to-device copy once and returns immediately thereafter, so it
597 ! is safe to call every step. The walk applies the true symmetry-function
598 ! cutoffs (it runs every step), so no Verlet-skin widening is needed.
599 mpg = 1
600 DO ie = 1, ne
601 DO je = 1, ne
602 mpg = max(mpg, nnp%neighbor_interface_state%pair_map(ie, je)%n_rad, &
603 nnp%neighbor_interface_state%pair_map(ie, je)%n_ang1, &
604 nnp%neighbor_interface_state%pair_map(ie, je)%n_ang2)
605 END DO
606 END DO
607 ALLOCATE (pm_n_rad(ne*ne), pm_n_ang1(ne*ne), pm_n_ang2(ne*ne), pm_max_cut(ne*ne))
608 ALLOCATE (pm_rad_g(mpg*ne*ne), pm_ang1_g(mpg*ne*ne), pm_ang2_g(mpg*ne*ne))
609 ALLOCATE (pm_rad_cut(mrg*ne), pm_ang_cut(mag*ne))
610 pm_rad_g(:) = 0
611 pm_ang1_g(:) = 0
612 pm_ang2_g(:) = 0
613 pm_rad_cut(:) = 0.0_dp
614 pm_ang_cut(:) = 0.0_dp
615 DO je = 1, ne
616 DO ie = 1, ne
617 pflat = (ie - 1) + ne*(je - 1)
618 associate(pmap => nnp%neighbor_interface_state%pair_map(ie, je))
619 pm_n_rad(pflat + 1) = int(pmap%n_rad, c_int)
620 pm_n_ang1(pflat + 1) = int(pmap%n_ang1, c_int)
621 pm_n_ang2(pflat + 1) = int(pmap%n_ang2, c_int)
622 pm_max_cut(pflat + 1) = pmap%max_relevant_cutoff
623 DO gg = 1, pmap%n_rad
624 pm_rad_g((gg - 1) + mpg*pflat + 1) = int(pmap%rad_groups(gg), c_int)
625 END DO
626 DO gg = 1, pmap%n_ang1
627 pm_ang1_g((gg - 1) + mpg*pflat + 1) = int(pmap%ang1_groups(gg), c_int)
628 END DO
629 DO gg = 1, pmap%n_ang2
630 pm_ang2_g((gg - 1) + mpg*pflat + 1) = int(pmap%ang2_groups(gg), c_int)
631 END DO
632 END associate
633 END DO
634 END DO
635 DO ie = 1, ne
636 DO ss = 1, nnp%rad(ie)%n_symfgrp
637 pm_rad_cut((ss - 1) + mrg*(ie - 1) + 1) = nnp%rad(ie)%symfgrp(ss)%cutoff
638 END DO
639 DO ss = 1, nnp%ang(ie)%n_symfgrp
640 pm_ang_cut((ss - 1) + mag*(ie - 1) + 1) = nnp%ang(ie)%symfgrp(ss)%cutoff
641 END DO
642 END DO
643 ec = nnp_gpu_upload_pair_map_c(int(ne, c_int), int(mpg, c_int), int(mrg, c_int), &
644 int(mag, c_int), pm_n_rad, pm_n_ang1, pm_n_ang2, &
645 pm_max_cut, pm_rad_g, pm_ang1_g, pm_ang2_g, &
646 pm_rad_cut, pm_ang_cut)
647 IF (ec /= 0) cpabort("NNP GPU: pair-map upload failed.")
648
649 ! Count rank-local atoms per element and build the (slot, element) to
650 ! global-atom map in the same order the descriptor kernels expect.
651 ALLOCATE (n_atoms_in_e_buf(ne))
652 n_atoms_in_e_buf(:) = 0
653 DO ii = istart, istart + mecalc - 1
654 e_w = nnp%ele_ind(ii)
655 n_atoms_in_e_buf(e_w) = n_atoms_in_e_buf(e_w) + 1
656 END DO
657 n_atoms_in_e_max = maxval(n_atoms_in_e_buf)
658
659 ! Publish the per-element counts before any early return: a rank whose
660 ! atom window is empty (more ranks than atoms) still has to leave the host
661 ! state in a well-defined state, because the caller checks it rather than
662 ! recomputing the window.
663 nnp%gpu_host_state%n_atoms_in_e_max = n_atoms_in_e_max
664 IF (ALLOCATED(nnp%gpu_host_state%n_atoms_in_e)) THEN
665 DEALLOCATE (nnp%gpu_host_state%n_atoms_in_e)
666 END IF
667 ALLOCATE (nnp%gpu_host_state%n_atoms_in_e(ne))
668 nnp%gpu_host_state%n_atoms_in_e(:) = n_atoms_in_e_buf(:)
669
670 IF (n_atoms_in_e_max <= 0) THEN
671 DEALLOCATE (n_atoms_in_e_buf, pm_n_rad, pm_n_ang1, pm_n_ang2, pm_max_cut, &
672 pm_rad_g, pm_ang1_g, pm_ang2_g, pm_rad_cut, pm_ang_cut)
673 CALL timestop(handle)
674 RETURN
675 END IF
676 ALLOCATE (atoms_global_per_e_buf(n_atoms_in_e_max*ne))
677 atoms_global_per_e_buf(:) = 0
678 n_atoms_in_e_buf(:) = 0
679 DO ii = istart, istart + mecalc - 1
680 e_w = nnp%ele_ind(ii)
681 n_atoms_in_e_buf(e_w) = n_atoms_in_e_buf(e_w) + 1
682 slot_w = n_atoms_in_e_buf(e_w)
683 atoms_global_per_e_buf((slot_w - 1) + n_atoms_in_e_max*(e_w - 1) + 1) = int(ii, c_int)
684 END DO
685 ALLOCATE (ele_ind_global_buf(nnp%num_atoms))
686 DO ii = 1, nnp%num_atoms
687 ele_ind_global_buf(ii) = int(nnp%ele_ind(ii), c_int)
688 END DO
689
690 nbin_buf = int(nnp%cell_list_cache%nbin, c_int)
691 span_buf = int(nnp%cell_list_cache%bin_span, c_int)
692 perd_buf = int(nnp%cell_list_cache%perd, c_int)
693 epbc_buf = int(nnp%cell_list_cache%exact_pbc_copies, c_int)
694 lower_buf = nnp%cell_list_cache%lower
695 width_buf = nnp%cell_list_cache%bin_width
696
697 IF (.NOT. nnp%gpu_host_state%walk_initialised) THEN
698 nnp%gpu_host_state%max_n_neigh_rad = gpu_walk_init_max_n_neigh
699 nnp%gpu_host_state%max_n_neigh_ang1 = gpu_walk_init_max_n_neigh
700 nnp%gpu_host_state%max_n_neigh_ang2 = gpu_walk_init_max_n_neigh
701 nnp%gpu_host_state%walk_initialised = .true.
702 END IF
703
704 attempts = 0
705 DO
706 attempts = attempts + 1
707 rc_walk = nnp_gpu_cell_list_walk_c(int(nnp%num_atoms, c_int), int(ne, c_int), &
708 int(nnp%cell_list_cache%n_images, c_int), &
709 int(nnp%cell_list_cache%n_cells, c_int), &
710 int(n_atoms_in_e_max, c_int), &
711 int(nnp%gpu_host_state%max_n_neigh_rad, c_int), &
712 int(nnp%gpu_host_state%max_n_neigh_ang1, c_int), &
713 int(nnp%gpu_host_state%max_n_neigh_ang2, c_int), &
714 int(mrg, c_int), int(mag, c_int), nbin_buf, &
715 lower_buf, width_buf, span_buf, perd_buf, epbc_buf, &
716 nnp%cell_list_cache%coord_scaled, ele_ind_global_buf, &
717 atoms_global_per_e_buf, n_atoms_in_e_buf)
718 IF (rc_walk == 0) EXIT
719 IF (rc_walk == -2 .AND. attempts < 4) THEN
720 nnp%gpu_host_state%max_n_neigh_rad = nnp%gpu_host_state%max_n_neigh_rad*2
721 nnp%gpu_host_state%max_n_neigh_ang1 = nnp%gpu_host_state%max_n_neigh_ang1*2
722 nnp%gpu_host_state%max_n_neigh_ang2 = nnp%gpu_host_state%max_n_neigh_ang2*2
723 cycle
724 END IF
725 IF (rc_walk == -2) THEN
726 cpabort("NNP GPU: a per-atom neighbour list still overflows after three width doublings.")
727 END IF
728 cpabort("NNP GPU: cell-list walk failed.")
729 END DO
730
731 ! Read back the per-(slot, group, element) neighbour counts, used by the
732 ! descriptor kernels to size their per-atom neighbour buffers.
733 IF (ALLOCATED(nnp%gpu_host_state%n_rad_pa)) DEALLOCATE (nnp%gpu_host_state%n_rad_pa)
734 IF (ALLOCATED(nnp%gpu_host_state%n_ang1_pa)) DEALLOCATE (nnp%gpu_host_state%n_ang1_pa)
735 IF (ALLOCATED(nnp%gpu_host_state%n_ang2_pa)) DEALLOCATE (nnp%gpu_host_state%n_ang2_pa)
736 ALLOCATE (nnp%gpu_host_state%n_rad_pa(n_atoms_in_e_max, mrg, ne))
737 ALLOCATE (nnp%gpu_host_state%n_ang1_pa(n_atoms_in_e_max, mag, ne))
738 ALLOCATE (nnp%gpu_host_state%n_ang2_pa(n_atoms_in_e_max, mag, ne))
739 nnp%gpu_host_state%n_rad_pa(:, :, :) = 0
740 nnp%gpu_host_state%n_ang1_pa(:, :, :) = 0
741 nnp%gpu_host_state%n_ang2_pa(:, :, :) = 0
742 ec = nnp_gpu_walk_counts_d2h_c(int(n_atoms_in_e_max, c_int), int(mrg, c_int), &
743 int(mag, c_int), int(ne, c_int), &
744 nnp%gpu_host_state%n_rad_pa, nnp%gpu_host_state%n_ang1_pa, &
745 nnp%gpu_host_state%n_ang2_pa)
746 IF (ec /= 0) cpabort("NNP GPU: walk-count read-back failed.")
747
748 DEALLOCATE (n_atoms_in_e_buf, atoms_global_per_e_buf, ele_ind_global_buf, &
749 pm_n_rad, pm_n_ang1, pm_n_ang2, pm_max_cut, pm_rad_g, pm_ang1_g, pm_ang2_g, &
750 pm_rad_cut, pm_ang_cut)
751
752 CALL timestop(handle)
753#else
754 mark_used(nnp)
755 mark_used(istart)
756 mark_used(mecalc)
757 mark_used(handle)
758 cpabort("NNP GPU support is not compiled into this binary.")
759#endif
760
761 END SUBROUTINE nnp_gpu_run_walk
762
763! **************************************************************************************************
764!> \brief Upload the model-static descriptor tables to the device once at NNP
765!> initialization: the per-element symmetry-function scaling parameters
766!> (average, minimum, maximum, sigma) and the per-(group, element)
767!> symmetry-function to descriptor-row maps. Both are immutable after the
768!> input parse, so the device copies are valid for the whole environment
769!> lifetime and the per-step descriptor kernels read them without any
770!> further host-to-device traffic.
771!> \param nnp the NNP environment, with its radial and angular groups populated
772! **************************************************************************************************
773 SUBROUTINE nnp_gpu_upload_descriptor_statics(nnp)
774 TYPE(nnp_type), INTENT(IN), POINTER :: nnp
775
776 INTEGER :: handle
777
778#if (defined(__OFFLOAD) || defined(__OFFLOAD_CUDA) || defined(__OFFLOAD_HIP)) && !defined(__NO_OFFLOAD_NNP)
779 CHARACTER(len=*), PARAMETER :: routinen = 'nnp_gpu_upload_descriptor_statics'
780 INTEGER :: e, ec, m, max_n_anggrp, max_n_radgrp, n_ang_se_total, n_input_max, &
781 n_rad_se_total, n_symf_g, s, se_flat, se_off, sf, total_ang_symf, total_rad_symf
782 INTEGER(KIND=c_int), ALLOCATABLE, DIMENSION(:) :: ang_offsets, ang_packed, n_in_buf, &
783 rad_offsets, rad_packed
784 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: loc_av_buf, loc_max_buf, loc_min_buf, &
785 sigma_buf
786
787 CALL timeset(routinen, handle)
788
789 associate(n_ele => nnp%n_ele)
790
791 ! Per-element scaling tables. The device layout is one slab of
792 ! n_input_max doubles per element, with the element's radial groups in
793 ! the leading slots and its angular groups next, matching the layout
794 ! the scaling kernel indexes.
795 n_input_max = 0
796 DO e = 1, n_ele
797 n_input_max = max(n_input_max, nnp%arc(e)%n_nodes(1))
798 END DO
799
800 ALLOCATE (loc_av_buf(n_ele*n_input_max))
801 ALLOCATE (loc_min_buf(n_ele*n_input_max))
802 ALLOCATE (loc_max_buf(n_ele*n_input_max))
803 ALLOCATE (sigma_buf(n_ele*n_input_max))
804 ALLOCATE (n_in_buf(n_ele))
805 loc_av_buf = 0.0_dp
806 loc_min_buf = 0.0_dp
807 loc_max_buf = 0.0_dp
808 sigma_buf = 0.0_dp
809
810 DO e = 1, n_ele
811 n_in_buf(e) = int(nnp%arc(e)%n_nodes(1), c_int)
812 DO m = 1, nnp%n_rad(e)
813 loc_av_buf((e - 1)*n_input_max + m) = nnp%rad(e)%loc_av(m)
814 loc_min_buf((e - 1)*n_input_max + m) = nnp%rad(e)%loc_min(m)
815 loc_max_buf((e - 1)*n_input_max + m) = nnp%rad(e)%loc_max(m)
816 sigma_buf((e - 1)*n_input_max + m) = nnp%rad(e)%sigma(m)
817 END DO
818 DO m = 1, nnp%n_ang(e)
819 loc_av_buf((e - 1)*n_input_max + nnp%n_rad(e) + m) = nnp%ang(e)%loc_av(m)
820 loc_min_buf((e - 1)*n_input_max + nnp%n_rad(e) + m) = nnp%ang(e)%loc_min(m)
821 loc_max_buf((e - 1)*n_input_max + nnp%n_rad(e) + m) = nnp%ang(e)%loc_max(m)
822 sigma_buf((e - 1)*n_input_max + nnp%n_rad(e) + m) = nnp%ang(e)%sigma(m)
823 END DO
824 END DO
825
826 ec = nnp_gpu_upload_scaling_tables_c(int(n_ele, c_int), int(n_input_max, c_int), &
827 loc_av_buf, loc_min_buf, loc_max_buf, sigma_buf, &
828 n_in_buf)
829 IF (ec /= 0) cpabort("NNP GPU: descriptor scaling-table upload failed.")
830
831 DEALLOCATE (loc_av_buf, loc_min_buf, loc_max_buf, sigma_buf, n_in_buf)
832
833 ! Per-(group, element) symmetry-function maps, packed radial then
834 ! angular with prefix-sum offsets keyed se_flat = (s-1) + max_grp*(e-1)
835 ! (element outermost). Padding slots for elements with fewer groups get
836 ! zero-length entries; the final slot carries the sentinel total.
837 CALL nnp_gpu_max_symfgrp(nnp, max_n_radgrp, max_n_anggrp)
838
839 n_rad_se_total = max_n_radgrp*n_ele
840 n_ang_se_total = max_n_anggrp*n_ele
841
842 total_rad_symf = 0
843 DO e = 1, n_ele
844 DO s = 1, nnp%rad(e)%n_symfgrp
845 total_rad_symf = total_rad_symf + nnp%rad(e)%symfgrp(s)%n_symf
846 END DO
847 END DO
848 total_ang_symf = 0
849 DO e = 1, n_ele
850 DO s = 1, nnp%ang(e)%n_symfgrp
851 total_ang_symf = total_ang_symf + nnp%ang(e)%symfgrp(s)%n_symf
852 END DO
853 END DO
854
855 ALLOCATE (rad_offsets(max(1, n_rad_se_total + 1)))
856 ALLOCATE (ang_offsets(max(1, n_ang_se_total + 1)))
857 ALLOCATE (rad_packed(max(1, total_rad_symf)))
858 ALLOCATE (ang_packed(max(1, total_ang_symf)))
859 rad_offsets = 0_c_int
860 ang_offsets = 0_c_int
861 rad_packed = 0_c_int
862 ang_packed = 0_c_int
863
864 se_off = 0
865 DO e = 1, n_ele
866 DO s = 1, max_n_radgrp
867 se_flat = (s - 1) + max_n_radgrp*(e - 1)
868 rad_offsets(se_flat + 1) = int(se_off, c_int)
869 IF (s <= nnp%rad(e)%n_symfgrp) THEN
870 n_symf_g = nnp%rad(e)%symfgrp(s)%n_symf
871 DO sf = 1, n_symf_g
872 rad_packed(se_off + sf) = int(nnp%rad(e)%symfgrp(s)%symf(sf), c_int)
873 END DO
874 se_off = se_off + n_symf_g
875 END IF
876 END DO
877 END DO
878 IF (n_rad_se_total > 0) rad_offsets(n_rad_se_total + 1) = int(se_off, c_int)
879
880 se_off = 0
881 DO e = 1, n_ele
882 DO s = 1, max_n_anggrp
883 se_flat = (s - 1) + max_n_anggrp*(e - 1)
884 ang_offsets(se_flat + 1) = int(se_off, c_int)
885 IF (s <= nnp%ang(e)%n_symfgrp) THEN
886 n_symf_g = nnp%ang(e)%symfgrp(s)%n_symf
887 DO sf = 1, n_symf_g
888 ang_packed(se_off + sf) = int(nnp%ang(e)%symfgrp(s)%symf(sf), c_int)
889 END DO
890 se_off = se_off + n_symf_g
891 END IF
892 END DO
893 END DO
894 IF (n_ang_se_total > 0) ang_offsets(n_ang_se_total + 1) = int(se_off, c_int)
895
896 ec = nnp_gpu_upload_se_static_tables_c(int(n_ele, c_int), int(max_n_radgrp, c_int), &
897 int(max_n_anggrp, c_int), &
898 int(total_rad_symf, c_int), &
899 int(total_ang_symf, c_int), &
900 rad_packed, rad_offsets, ang_packed, ang_offsets)
901 IF (ec /= 0) cpabort("NNP GPU: descriptor symf-map upload failed.")
902
903 DEALLOCATE (rad_packed, rad_offsets, ang_packed, ang_offsets)
904
905 END associate
906
907 CALL timestop(handle)
908#else
909 mark_used(nnp)
910 mark_used(handle)
911 cpabort("NNP GPU support is not compiled into this binary.")
912#endif
913
914 END SUBROUTINE nnp_gpu_upload_descriptor_statics
915
916! **************************************************************************************************
917!> \brief Build the device-resident descriptor for one MD step: zero G_dev,
918!> dispatch the radial then the angular symmetry-function kernels for
919!> every populated (group, element) slot, then scale G_dev into the
920!> network input. The kernels read the neighbour lists, atom maps and
921!> symmetry-function maps from device buffers, so this routine only
922!> supplies per-group sizes, model parameters and index scalars. It must
923!> be called after nnp_gpu_run_walk has populated the host walk-count
924!> state for this step.
925!> \param nnp the NNP environment
926!> \param istart first rank-local atom (1-based, global index)
927!> \param mecalc number of rank-local atoms
928!> \param n_input_max descriptor-row stride of G_dev (max over elements of the
929!> network input-layer width)
930! **************************************************************************************************
931 SUBROUTINE nnp_gpu_build_descriptor(nnp, istart, mecalc, n_input_max)
932 TYPE(nnp_type), INTENT(INOUT), POINTER :: nnp
933 INTEGER, INTENT(IN) :: istart, mecalc, n_input_max
934
935 INTEGER :: handle
936
937#if (defined(__OFFLOAD) || defined(__OFFLOAD_CUDA) || defined(__OFFLOAD_HIP)) && !defined(__NO_OFFLOAD_NNP)
938 CHARACTER(len=*), PARAMETER :: routinen = 'nnp_gpu_build_descriptor'
939 INTEGER :: cut_int, e, ec, i, i_local, max_n_ang1, max_n_ang2, max_n_anggrp, max_n_kk, &
940 max_n_neigh, max_n_radgrp, n_atoms_e, s, scaling_mode
941 LOGICAL :: homo, splines_ready
942 REAL(kind=dp) :: cutoff_s, cutoff_sqr
943 INTEGER(KIND=c_int), ALLOCATABLE, DIMENSION(:) :: ele_ind_buf, izeta_buf
944 CHARACTER(KIND=c_char), ALLOCATABLE, DIMENSION(:) :: use_int_zeta_buf
945
946 CALL timeset(routinen, handle)
947
948 ! Group-count strides used to flatten the (group, element) slot index the
949 ! device buffers are keyed by.
950 CALL nnp_gpu_max_symfgrp(nnp, max_n_radgrp, max_n_anggrp)
951
952 ! Zero G_dev for this step's rank-local atoms.
953 ec = nnp_gpu_begin_descriptor_step_c(int(mecalc, c_int), int(n_input_max, c_int))
954 IF (ec /= 0) cpabort("NNP GPU: descriptor-step begin failed.")
955
956 ! The radial kernels sample the Hermite splines built from the model's
957 ! radial symmetry-function parameters; build them once if a fresh
958 ! environment has not done so yet.
959 splines_ready = .true.
960 DO e = 1, nnp%n_ele
961 DO s = 1, nnp%rad(e)%n_symfgrp
962 IF (.NOT. nnp%rad(e)%symfgrp(s)%spline_built) splines_ready = .false.
963 END DO
964 END DO
965 IF (.NOT. splines_ready) CALL nnp_build_radial_splines(nnp)
966
967 associate(hstate => nnp%gpu_host_state)
968
969 ! Radial symmetry functions into G_dev, one dispatch per populated
970 ! (group, element) slot.
971 DO e = 1, nnp%n_ele
972 n_atoms_e = hstate%n_atoms_in_e(e)
973 IF (n_atoms_e <= 0) cycle
974 DO s = 1, nnp%rad(e)%n_symfgrp
975 associate(symfgrp => nnp%rad(e)%symfgrp(s))
976 max_n_neigh = maxval(hstate%n_rad_pa(1:n_atoms_e, s, e))
977 IF (max_n_neigh <= 0) cycle
978 ec = nnp_gpu_radial_sf_to_g_dev_c( &
979 int(n_atoms_e, c_int), int(max_n_neigh, c_int), &
980 int(symfgrp%n_symf, c_int), int(symfgrp%spline_n, c_int), &
981 REAL(symfgrp%spline_dx, c_double), &
982 REAL(symfgrp%spline_dx_inv, c_double), &
983 REAL(symfgrp%spline_x_max, c_double), &
984 symfgrp%spline_y, symfgrp%spline_dy, &
985 int(s - 1, c_int), int(e - 1, c_int), int(max_n_radgrp, c_int), &
986 int(hstate%n_atoms_in_e_max, c_int), &
987 int(hstate%max_n_neigh_rad, c_int), int(istart, c_int), &
988 int(n_input_max, c_int), int((s - 1) + max_n_radgrp*(e - 1), c_int))
989 IF (ec /= 0) cpabort("NNP GPU: radial descriptor kernel failed.")
990 END associate
991 END DO
992 END DO
993
994 ! Angular symmetry functions into G_dev. max_n_ang1 / max_n_ang2 are
995 ! the per-slot high-water neighbour counts from the walk; max_n_ang2 is
996 ! zero for homonuclear groups (their second neighbour list is unused)
997 ! and the kernel's neighbour-pair width max_n_kk is the ang1 list width
998 ! for homonuclear groups and the ang2 list width otherwise.
999 cut_int = nnp%cut_type - 1
1000 DO e = 1, nnp%n_ele
1001 n_atoms_e = hstate%n_atoms_in_e(e)
1002 IF (n_atoms_e <= 0) cycle
1003 DO s = 1, nnp%ang(e)%n_symfgrp
1004 associate(symfgrp => nnp%ang(e)%symfgrp(s))
1005 homo = (symfgrp%ele(1) == symfgrp%ele(2))
1006 cutoff_s = symfgrp%cutoff
1007 cutoff_sqr = cutoff_s*cutoff_s
1008
1009 max_n_ang1 = maxval(hstate%n_ang1_pa(1:n_atoms_e, s, e))
1010 IF (max_n_ang1 <= 0) cycle
1011 IF (homo) THEN
1012 max_n_ang2 = 0
1013 max_n_kk = max(1, max_n_ang1)
1014 ELSE
1015 max_n_ang2 = maxval(hstate%n_ang2_pa(1:n_atoms_e, s, e))
1016 ! No atom of e has an ele(2) neighbour in range: the pair
1017 ! list is empty and the group's contribution is exactly
1018 ! zero (G_dev is zeroed each step), so skip it, mirroring
1019 ! the max_n_ang1 guard above. Happens e.g. for ion-ion
1020 ! angular groups in dilute electrolytes whenever no ion
1021 ! pair is within the cutoff. The force-scatter loop
1022 ! below applies the SAME guard: a skipped group leaves
1023 ! its per-(s,e) device buffers unpopulated, and
1024 ! launching the scatter on them would read stale data.
1025 IF (max_n_ang2 <= 0) cycle
1026 max_n_kk = max(1, max_n_ang2)
1027 END IF
1028
1029 ALLOCATE (izeta_buf(symfgrp%n_symf))
1030 ALLOCATE (use_int_zeta_buf(symfgrp%n_symf))
1031 DO i = 1, symfgrp%n_symf
1032 izeta_buf(i) = int(symfgrp%pack_izeta(i), c_int)
1033 use_int_zeta_buf(i) = merge(achar(1), achar(0), symfgrp%pack_use_int_zeta(i))
1034 END DO
1035
1036 ec = nnp_gpu_angular_sf_to_g_dev_c( &
1037 int(merge(1, 0, homo), c_int), int(cut_int, c_int), &
1038 int(max(1, max_n_ang1), c_int), int(max_n_ang2, c_int), &
1039 int(max_n_kk, c_int), int(symfgrp%n_symf, c_int), &
1040 int(n_atoms_e, c_int), real(cutoff_s, c_double), &
1041 REAL(cutoff_sqr, c_double), &
1042 symfgrp%pack_lam, symfgrp%pack_zeta, symfgrp%pack_eta, &
1043 symfgrp%pack_prefzeta, izeta_buf, use_int_zeta_buf, &
1044 int(s - 1, c_int), int(e - 1, c_int), int(max_n_anggrp, c_int), &
1045 int(hstate%n_atoms_in_e_max, c_int), &
1046 int(hstate%max_n_neigh_ang1, c_int), &
1047 int(hstate%max_n_neigh_ang2, c_int), int(istart, c_int), &
1048 int(nnp%n_rad(e), c_int), int(n_input_max, c_int), &
1049 int((s - 1) + max_n_anggrp*(e - 1), c_int))
1050 IF (ec /= 0) cpabort("NNP GPU: angular descriptor kernel failed.")
1051
1052 DEALLOCATE (izeta_buf, use_int_zeta_buf)
1053 END associate
1054 END DO
1055 END DO
1056
1057 END associate
1058
1059 ! Scale the aggregated descriptor into the network input. The scaling mode
1060 ! selects which of centering, scaling and sigma normalization the model
1061 ! requested; the per-atom 0-based element index routes each atom to its
1062 ! element's scaling parameters.
1063 IF (nnp%scale_sigma_acsf) THEN
1064 scaling_mode = 4
1065 ELSE IF (nnp%center_acsf .AND. nnp%scale_acsf) THEN
1066 scaling_mode = 3
1067 ELSE IF (nnp%center_acsf) THEN
1068 scaling_mode = 2
1069 ELSE IF (nnp%scale_acsf) THEN
1070 scaling_mode = 1
1071 ELSE
1072 scaling_mode = 0
1073 END IF
1074
1075 ALLOCATE (ele_ind_buf(max(1, mecalc)))
1076 DO i = istart, istart + mecalc - 1
1077 i_local = i - istart + 1
1078 ele_ind_buf(i_local) = int(nnp%ele_ind(i) - 1, c_int)
1079 END DO
1080
1081 ec = nnp_gpu_scale_g_dev_c(int(mecalc, c_int), int(n_input_max, c_int), &
1082 int(nnp%n_ele, c_int), int(scaling_mode, c_int), &
1083 REAL(nnp%scmin, c_double), REAL(nnp%scmax, c_double), &
1084 ele_ind_buf)
1085 IF (ec /= 0) cpabort("NNP GPU: descriptor scaling failed.")
1086
1087 DEALLOCATE (ele_ind_buf)
1088
1089 CALL timestop(handle)
1090#else
1091 mark_used(nnp)
1092 mark_used(istart)
1093 mark_used(mecalc)
1094 mark_used(n_input_max)
1095 mark_used(handle)
1096 cpabort("NNP GPU support is not compiled into this binary.")
1097#endif
1098
1099 END SUBROUTINE nnp_gpu_build_descriptor
1100
1101! **************************************************************************************************
1102!> \brief Upload the committee network to the device once at NNP initialization:
1103!> the per-element layer topology, the shared activation-function codes,
1104!> the per-element atomic reference energies and the packed per-committee
1105!> weights and biases. These are immutable after the weights are read, so
1106!> the device copies serve every MD step and the per-step network launcher
1107!> needs no further weight traffic. The weights and biases are packed into
1108!> zero-padded max_n_nodes-wide slabs so a single stride addresses every
1109!> (element, layer, committee) block regardless of that block's true size.
1110!> \param nnp the NNP environment, with its architecture and committee weights
1111!> already read
1112! **************************************************************************************************
1113 SUBROUTINE nnp_gpu_upload_network(nnp)
1114 TYPE(nnp_type), INTENT(IN), POINTER :: nnp
1115
1116 INTEGER :: handle
1117
1118#if (defined(__OFFLOAD) || defined(__OFFLOAD_CUDA) || defined(__OFFLOAD_HIP)) && !defined(__NO_OFFLOAD_NNP)
1119 CHARACTER(len=*), PARAMETER :: routinen = 'nnp_gpu_upload_network'
1120 INTEGER :: c, e, ec, i, idx_b, idx_w, j, l, max_n_nodes, n_com, n_ele, n_in, n_layer, n_out
1121 INTEGER(KIND=c_int), ALLOCATABLE, DIMENSION(:) :: act_fnct_buf, n_nodes_buf
1122 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: b_packed_buf, w_packed_buf
1123
1124 CALL timeset(routinen, handle)
1125
1126 n_com = nnp%n_committee
1127 n_layer = nnp%n_layer
1128 n_ele = nnp%n_ele
1129
1130 ! Slab width for the packed weights and biases: the largest layer node
1131 ! count over all elements and layers. The kernel reads only the leading
1132 ! n_nodes(L-1) x n_nodes(L) corner of each slab.
1133 max_n_nodes = 0
1134 DO e = 1, n_ele
1135 DO l = 1, n_layer
1136 IF (nnp%arc(e)%n_nodes(l) > max_n_nodes) max_n_nodes = nnp%arc(e)%n_nodes(l)
1137 END DO
1138 END DO
1139
1140 ! Layer node counts, flattened layer-fastest then element.
1141 ALLOCATE (n_nodes_buf(n_layer*n_ele))
1142 DO e = 1, n_ele
1143 DO l = 1, n_layer
1144 n_nodes_buf((l - 1) + n_layer*(e - 1) + 1) = int(nnp%arc(e)%n_nodes(l), c_int)
1145 END DO
1146 END DO
1147
1148 ! Activation-function codes for layers 2..n_layer (shared across elements).
1149 ALLOCATE (act_fnct_buf(n_layer - 1))
1150 DO l = 2, n_layer
1151 act_fnct_buf(l - 1) = int(nnp%actfnct(l - 1), c_int)
1152 END DO
1153
1154 ! Weights and biases packed into zero-padded max_n_nodes-wide slabs, one
1155 ! slab per (element, layer, committee). Input node i is fastest in the
1156 ! weight slab; the bias slab is one node column per (element, layer,
1157 ! committee). Layout matches nnp_gpu_nn_atomwise_with_G_dev_c and the
1158 ! legacy host packing in nnp_model.F.
1159 ALLOCATE (w_packed_buf(max_n_nodes*max_n_nodes*n_com*(n_layer - 1)*n_ele))
1160 ALLOCATE (b_packed_buf(max_n_nodes*n_com*(n_layer - 1)*n_ele))
1161 w_packed_buf = 0.0_dp
1162 b_packed_buf = 0.0_dp
1163 DO e = 1, n_ele
1164 DO l = 2, n_layer
1165 n_in = nnp%arc(e)%n_nodes(l - 1)
1166 n_out = nnp%arc(e)%n_nodes(l)
1167 DO c = 1, n_com
1168 DO j = 1, n_out
1169 DO i = 1, n_in
1170 idx_w = (i - 1) + max_n_nodes*(j - 1) + &
1171 max_n_nodes*max_n_nodes*(c - 1) + &
1172 max_n_nodes*max_n_nodes*n_com*(l - 2) + &
1173 max_n_nodes*max_n_nodes*n_com*(n_layer - 1)*(e - 1) + 1
1174 w_packed_buf(idx_w) = nnp%arc(e)%layer(l)%weights(i, j, c)
1175 END DO
1176 idx_b = (j - 1) + max_n_nodes*(c - 1) + &
1177 max_n_nodes*n_com*(l - 2) + &
1178 max_n_nodes*n_com*(n_layer - 1)*(e - 1) + 1
1179 b_packed_buf(idx_b) = nnp%arc(e)%layer(l)%bweights(j, c)
1180 END DO
1181 END DO
1182 END DO
1183 END DO
1184
1185 ec = nnp_gpu_upload_nn_packed_c(int(n_com, c_int), int(n_layer, c_int), &
1186 int(max_n_nodes, c_int), int(n_ele, c_int), &
1187 n_nodes_buf, act_fnct_buf, nnp%atom_energies, &
1188 w_packed_buf, b_packed_buf)
1189 IF (ec /= 0) cpabort("NNP GPU: committee network upload failed.")
1190
1191 DEALLOCATE (n_nodes_buf, act_fnct_buf, w_packed_buf, b_packed_buf)
1192
1193 CALL timestop(handle)
1194#else
1195 mark_used(nnp)
1196 mark_used(handle)
1197 cpabort("NNP GPU support is not compiled into this binary.")
1198#endif
1199
1200 END SUBROUTINE nnp_gpu_upload_network
1201
1202! **************************************************************************************************
1203!> \brief Evaluate the committee network for one MD step from the scaled
1204!> device-resident descriptor. The network topology and weights were
1205!> uploaded once by nnp_gpu_upload_network, so this launcher only supplies
1206!> the per-atom element index and receives the per-atom per-committee
1207!> energies back. The reference atomic energies are added inside the
1208!> kernel, so the returned energy is the complete per-atom contribution
1209!> and is copied straight into nnp%atomic_energy without any further
1210!> offset. The descriptor gradient is left on the device for the force
1211!> scatter; its host buffer is only allocated and passed to satisfy the
1212!> launcher signature. Must be called after nnp_gpu_build_descriptor.
1213!> \param nnp the NNP environment; nnp%atomic_energy is filled for this window
1214!> \param istart first rank-local atom (1-based, global index)
1215!> \param mecalc number of rank-local atoms
1216!> \param n_input_max descriptor-row stride of the device descriptor (max over
1217!> elements of the network input-layer width)
1218! **************************************************************************************************
1219 SUBROUTINE nnp_gpu_predict(nnp, istart, mecalc, n_input_max)
1220 TYPE(nnp_type), INTENT(INOUT), POINTER :: nnp
1221 INTEGER, INTENT(IN) :: istart, mecalc, n_input_max
1222
1223 INTEGER :: handle
1224
1225#if (defined(__OFFLOAD) || defined(__OFFLOAD_CUDA) || defined(__OFFLOAD_HIP)) && !defined(__NO_OFFLOAD_NNP)
1226 CHARACTER(len=*), PARAMETER :: routinen = 'nnp_gpu_predict'
1227 INTEGER :: atom_local, c, e, ec, i, l, max_n_nodes, n_com, n_ele, n_layer
1228 INTEGER(KIND=c_int), ALLOCATABLE, DIMENSION(:) :: ele_ind_buf
1229 REAL(kind=c_double), ALLOCATABLE, DIMENSION(:) :: h_atomic_energy
1230
1231 CALL timeset(routinen, handle)
1232
1233 n_com = nnp%n_committee
1234 n_layer = nnp%n_layer
1235 n_ele = nnp%n_ele
1236
1237 IF (mecalc <= 0) THEN
1238 CALL timestop(handle)
1239 RETURN
1240 END IF
1241
1242 ! Slab width the packed weights were uploaded with; the launcher reads the
1243 ! device weights with this stride.
1244 max_n_nodes = 0
1245 DO e = 1, n_ele
1246 DO l = 1, n_layer
1247 IF (nnp%arc(e)%n_nodes(l) > max_n_nodes) max_n_nodes = nnp%arc(e)%n_nodes(l)
1248 END DO
1249 END DO
1250
1251 ! Per-atom 0-based element index over this rank's window [istart, ...].
1252 ALLOCATE (ele_ind_buf(mecalc))
1253 DO i = istart, istart + mecalc - 1
1254 atom_local = i - istart + 1
1255 ele_ind_buf(atom_local) = int(nnp%ele_ind(i) - 1, c_int)
1256 END DO
1257
1258 ! Network outputs in kernel layout: energy[c + n_com*atom_local]. The
1259 ! descriptor gradient stays on the device, where the force scatter reads
1260 ! it, so no host buffer is offered for it.
1261 ALLOCATE (h_atomic_energy(n_com*mecalc))
1262 h_atomic_energy = 0.0_c_double
1263
1264 ec = nnp_gpu_nn_atomwise_with_g_dev_c(int(mecalc, c_int), int(n_com, c_int), &
1265 int(n_layer, c_int), int(max_n_nodes, c_int), &
1266 int(n_input_max, c_int), &
1267 merge(1_c_int, 0_c_int, nnp%normnodes), &
1268 int(n_ele, c_int), ele_ind_buf, &
1269 h_atomic_energy, c_null_ptr)
1270 IF (ec /= 0) cpabort("NNP GPU: committee network evaluation failed.")
1271
1272 ! Store the per-atom per-committee energy. The kernel already added the
1273 ! reference atomic energy atom_energies(ele), so this is the full per-atom
1274 ! contribution and is copied without any further offset (mirroring
1275 ! nnp_model.F nnp_predict_gradient_atomwise_gpu, which stores E_out_flat
1276 ! directly, and nnp_force.F, which copies atomic_energy_batch verbatim).
1277 DO atom_local = 1, mecalc
1278 i = istart + atom_local - 1
1279 DO c = 1, n_com
1280 nnp%atomic_energy(i, c) = h_atomic_energy((c - 1) + n_com*(atom_local - 1) + 1)
1281 END DO
1282 END DO
1283
1284 DEALLOCATE (ele_ind_buf, h_atomic_energy)
1285
1286 CALL timestop(handle)
1287#else
1288 mark_used(nnp)
1289 mark_used(istart)
1290 mark_used(mecalc)
1291 mark_used(n_input_max)
1292 mark_used(handle)
1293 cpabort("NNP GPU support is not compiled into this binary.")
1294#endif
1295
1296 END SUBROUTINE nnp_gpu_predict
1297
1298! **************************************************************************************************
1299!> \brief Run the whole per-step NNP force evaluation on the device: the
1300!> cell-list walk, the descriptor build, the committee network and the
1301!> force scatter, in that order. The walk builds the device neighbour
1302!> lists, the descriptor build fills the scaled device descriptor, the
1303!> network fills nnp%atomic_energy and leaves its descriptor gradient on
1304!> the device, and the scatter multiplies that gradient by the descriptor
1305!> derivatives and reads the total forces back into nnp%myforce. The
1306!> scatter accumulator is seeded from nnp%myforce, which the caller has
1307!> zeroed, so on return it holds the complete NNP forces for this window.
1308!> \param nnp the NNP environment; nnp%atomic_energy and nnp%myforce are filled
1309!> \param istart first rank-local atom (1-based, global index)
1310!> \param mecalc number of rank-local atoms
1311!> \param n_input_max descriptor-row stride of the device descriptor (max over
1312!> elements of the network input-layer width)
1313!> \param calc_stress .TRUE. when the caller wants the analytic virial, which
1314!> the force-scatter kernels then accumulate alongside the
1315!> forces
1316! **************************************************************************************************
1317 SUBROUTINE nnp_gpu_energy_force(nnp, istart, mecalc, n_input_max, calc_stress)
1318 TYPE(nnp_type), INTENT(INOUT), POINTER :: nnp
1319 INTEGER, INTENT(IN) :: istart, mecalc, n_input_max
1320 LOGICAL, INTENT(IN) :: calc_stress
1321
1322 INTEGER :: handle
1323
1324#if (defined(__OFFLOAD) || defined(__OFFLOAD_CUDA) || defined(__OFFLOAD_HIP)) && !defined(__NO_OFFLOAD_NNP)
1325 CHARACTER(len=*), PARAMETER :: routinen = 'nnp_gpu_energy_force'
1326 INTEGER :: accumulate_virial, e, ec, iva, ivb, ivc, max_n_ang1, max_n_ang2, max_n_anggrp, &
1327 max_n_kk, max_n_neigh_per_se, max_n_radgrp, n_atoms_e, s, scaling_mode
1328 LOGICAL :: homo
1329 REAL(kind=c_double), ALLOCATABLE :: c_virial(:)
1330
1331 CALL timeset(routinen, handle)
1332
1333 ! A rank can legitimately own no atoms (more ranks than atoms), in which
1334 ! case it contributes nothing: the caller has already zeroed nnp%myforce
1335 ! and nnp%atomic_energy for this window. The device kernels all require a
1336 ! non-empty window, so return before any of them is launched.
1337 IF (mecalc <= 0) THEN
1338 CALL timestop(handle)
1339 RETURN
1340 END IF
1341
1342 ! When a stress tensor is requested, the from-state scatter kernels also
1343 ! accumulate the analytical virial on device; the flag keeps the kernel
1344 ! virial branch off on energy/force-only steps.
1345 accumulate_virial = merge(1, 0, calc_stress)
1346
1347 ! Device neighbour lists, then the scaled device descriptor, then the
1348 ! committee energy (which also leaves its descriptor gradient on device).
1349 CALL nnp_gpu_run_walk(nnp, istart, mecalc)
1350 CALL nnp_gpu_build_descriptor(nnp, istart, mecalc, n_input_max)
1351 CALL nnp_gpu_predict(nnp, istart, mecalc, n_input_max)
1352
1353 ! Force scatter. The scaling mode selects which of centering, scaling and
1354 ! sigma normalization the model requested; it must match the mode the
1355 ! descriptor build used, so the kernels undo the same transform when they
1356 ! propagate the descriptor gradient back to the atoms.
1357 IF (nnp%scale_sigma_acsf) THEN
1358 scaling_mode = 4
1359 ELSE IF (nnp%center_acsf .AND. nnp%scale_acsf) THEN
1360 scaling_mode = 3
1361 ELSE IF (nnp%center_acsf) THEN
1362 scaling_mode = 2
1363 ELSE IF (nnp%scale_acsf) THEN
1364 scaling_mode = 1
1365 ELSE
1366 scaling_mode = 0
1367 END IF
1368
1369 ! Group-count strides used to flatten the (group, element) slot index the
1370 ! device walk and scatter buffers are keyed by.
1371 CALL nnp_gpu_max_symfgrp(nnp, max_n_radgrp, max_n_anggrp)
1372
1373 ! Seed the device force accumulator from the host force array (zeroed by
1374 ! the caller).
1375 ec = nnp_gpu_force_scatter_begin_c(int(nnp%num_atoms, c_int), &
1376 int(nnp%n_committee, c_int), nnp%myforce)
1377 IF (ec /= 0) cpabort("NNP GPU: force-scatter begin failed.")
1378
1379 associate(hstate => nnp%gpu_host_state)
1380
1381 ! Radial scatter, one dispatch per populated (group, element) slot.
1382 DO e = 1, nnp%n_ele
1383 n_atoms_e = hstate%n_atoms_in_e(e)
1384 IF (n_atoms_e <= 0) cycle
1385 DO s = 1, nnp%rad(e)%n_symfgrp
1386 associate(symfgrp => nnp%rad(e)%symfgrp(s))
1387 max_n_neigh_per_se = maxval(hstate%n_rad_pa(1:n_atoms_e, s, e))
1388 IF (max_n_neigh_per_se <= 0) cycle
1389 ec = nnp_gpu_force_scatter_radial_c( &
1390 int(n_atoms_e, c_int), int(symfgrp%n_symf, c_int), &
1391 int(max_n_neigh_per_se, c_int), int(nnp%n_committee, c_int), &
1392 int(n_input_max, c_int), int(nnp%num_atoms, c_int), &
1393 int(e - 1, c_int), int(scaling_mode, c_int), &
1394 REAL(nnp%scmin, c_double), REAL(nnp%scmax, c_double), &
1395 int(s - 1, c_int), int(max_n_radgrp, c_int), &
1396 int(hstate%n_atoms_in_e_max, c_int), &
1397 int(hstate%max_n_neigh_rad, c_int), int(istart, c_int), &
1398 int((s - 1) + max_n_radgrp*(e - 1), c_int), &
1399 int(accumulate_virial, c_int))
1400 IF (ec /= 0) cpabort("NNP GPU: radial force-scatter kernel failed.")
1401 END associate
1402 END DO
1403 END DO
1404
1405 ! Angular scatter, one dispatch per populated (group, element) slot. The
1406 ! per-slot high-water neighbour counts and the neighbour-pair width
1407 ! max_n_kk are sized exactly as the angular descriptor build: max_n_ang2
1408 ! is unused for a homonuclear group (its second neighbour list is empty)
1409 ! and the pair width max_n_kk is the ang1 list width for homonuclear
1410 ! groups and the ang2 list width otherwise.
1411 DO e = 1, nnp%n_ele
1412 n_atoms_e = hstate%n_atoms_in_e(e)
1413 IF (n_atoms_e <= 0) cycle
1414 DO s = 1, nnp%ang(e)%n_symfgrp
1415 associate(symfgrp => nnp%ang(e)%symfgrp(s))
1416 homo = (symfgrp%ele(1) == symfgrp%ele(2))
1417 max_n_ang1 = maxval(hstate%n_ang1_pa(1:n_atoms_e, s, e))
1418 IF (max_n_ang1 <= 0) cycle
1419 IF (homo) THEN
1420 max_n_ang2 = 0
1421 max_n_kk = max(1, max_n_ang1)
1422 ELSE
1423 max_n_ang2 = maxval(hstate%n_ang2_pa(1:n_atoms_e, s, e))
1424 ! Empty pair list: the SF stage skipped this group (same
1425 ! guard there), so its per-(s,e) device buffers are
1426 ! unpopulated, so launching the scatter would read stale
1427 ! data and produce silently wrong forces. Contribution
1428 ! is exactly zero; skip.
1429 IF (max_n_ang2 <= 0) cycle
1430 max_n_kk = max(1, max_n_ang2)
1431 END IF
1432 ec = nnp_gpu_force_scatter_angular_c( &
1433 int(n_atoms_e, c_int), int(symfgrp%n_symf, c_int), &
1434 int(max(1, max_n_ang1), c_int), int(max_n_kk, c_int), &
1435 int(nnp%n_committee, c_int), int(n_input_max, c_int), &
1436 int(nnp%num_atoms, c_int), int(e - 1, c_int), &
1437 int(nnp%n_rad(e), c_int), int(scaling_mode, c_int), &
1438 int(merge(1, 0, homo), c_int), real(nnp%scmin, c_double), &
1439 REAL(nnp%scmax, c_double), int(s - 1, c_int), &
1440 int(max_n_anggrp, c_int), int(hstate%n_atoms_in_e_max, c_int), &
1441 int(hstate%max_n_neigh_ang1, c_int), &
1442 int(hstate%max_n_neigh_ang2, c_int), int(istart, c_int), &
1443 int((s - 1) + max_n_anggrp*(e - 1), c_int), &
1444 int(accumulate_virial, c_int))
1445 IF (ec /= 0) cpabort("NNP GPU: angular force-scatter kernel failed.")
1446 END associate
1447 END DO
1448 END DO
1449
1450 END associate
1451
1452 ! Read the accumulated forces back into the host force array.
1453 ec = nnp_gpu_force_scatter_end_c(int(nnp%num_atoms, c_int), &
1454 int(nnp%n_committee, c_int), nnp%myforce)
1455 IF (ec /= 0) cpabort("NNP GPU: force-scatter end failed.")
1456
1457 ! Stress: the scatter kernels accumulated the per-committee analytical
1458 ! virial into the device buffer when accumulate_virial was set. D2H it and
1459 ! fold into committee_stress (zeroed by the caller), so the committee
1460 ! reduction in nnp_calc_energy_force forms virial%pv_virial unchanged. The
1461 ! device layout virial_dev[9*c + 3*a + b] (a = force component, b = bond
1462 ! component) maps to committee_stress(b, a, c) (column-major 3x3).
1463 IF (accumulate_virial /= 0) THEN
1464 ALLOCATE (c_virial(9*nnp%n_committee))
1465 ec = nnp_gpu_force_scatter_virial_d2h_c(int(nnp%n_committee, c_int), c_virial)
1466 IF (ec /= 0) cpabort("NNP GPU: force-scatter virial read-back failed.")
1467 DO ivc = 1, nnp%n_committee
1468 DO iva = 1, 3
1469 DO ivb = 1, 3
1470 nnp%committee_stress(ivb, iva, ivc) = &
1471 nnp%committee_stress(ivb, iva, ivc) + &
1472 c_virial(9*(ivc - 1) + 3*(iva - 1) + (ivb - 1) + 1)
1473 END DO
1474 END DO
1475 END DO
1476 DEALLOCATE (c_virial)
1477 END IF
1478
1479 ! Out-of-training-range detection, mirroring nnp_check_extrapolation on
1480 ! the CPU path: the scaling kernel raised the flag on the device, the
1481 ! force read-back above synchronized the stream, and the caller resets
1482 ! output_expol before dispatch and sums it over the ranks.
1483 nnp%output_expol = (nnp_gpu_extrapolation_flag_c() /= 0)
1484
1485 CALL timestop(handle)
1486#else
1487 mark_used(nnp)
1488 mark_used(istart)
1489 mark_used(mecalc)
1490 mark_used(n_input_max)
1491 mark_used(calc_stress)
1492 mark_used(handle)
1493 cpabort("NNP GPU support is not compiled into this binary.")
1494#endif
1495
1496 END SUBROUTINE nnp_gpu_energy_force
1497
1498END MODULE nnp_gpu
input_constants
collects all constants needed in input so that they can be used without circular dependencies
Definition input_constants.F:18
input_constants::nnp_gpu_prec_mixed
integer, parameter, public nnp_gpu_prec_mixed
Definition input_constants.F:72
kinds
Defines the basic variable types.
Definition kinds.F:23
kinds::dp
integer, parameter, public dp
Definition kinds.F:34
nnp_acsf
Functionality for atom centered symmetry functions for neural network potentials.
Definition nnp_acsf.F:16
nnp_acsf::nnp_build_radial_splines
subroutine, public nnp_build_radial_splines(nnp)
Build sf-first (n_symf, n_grid) Hermite cubic spline tables for radial SFs. Uses a custom sf-first la...
Definition nnp_acsf.F:1059
nnp_environment_types
Data types for neural network potentials.
Definition nnp_environment_types.F:14
nnp_gpu
GPU offload driver for the NNP force evaluation. Hosts the ISO_C_BINDING interfaces to the CUDA launc...
Definition nnp_gpu.F:16
nnp_gpu::nnp_gpu_init_device
subroutine, public nnp_gpu_init_device(gpu_precision, accepted)
Claim the offload device for this NNP environment and push the requested compute precision to the dev...
Definition nnp_gpu.F:407
nnp_gpu::nnp_gpu_model_capacity_ok
logical function, public nnp_gpu_model_capacity_ok(nnp, reason)
Test whether the model fits the compiled-in device-table capacities: NNP_MAX_SE (group,...
Definition nnp_gpu.F:464
nnp_gpu::nnp_gpu_upload_descriptor_statics
subroutine, public nnp_gpu_upload_descriptor_statics(nnp)
Upload the model-static descriptor tables to the device once at NNP initialization: the per-element s...
Definition nnp_gpu.F:774
nnp_gpu::nnp_gpu_energy_force
subroutine, public nnp_gpu_energy_force(nnp, istart, mecalc, n_input_max, calc_stress)
Run the whole per-step NNP force evaluation on the device: the cell-list walk, the descriptor build,...
Definition nnp_gpu.F:1318
nnp_gpu::nnp_gpu_available
logical function, public nnp_gpu_available()
Report whether the NNP GPU backend can run: the build must enable it and at least one offload device ...
Definition nnp_gpu.F:381
nnp_gpu::nnp_gpu_upload_network
subroutine, public nnp_gpu_upload_network(nnp)
Upload the committee network to the device once at NNP initialization: the per-element layer topology...
Definition nnp_gpu.F:1114
offload_api
Fortran API for the offload package, which is written in C.
Definition offload_api.F:12
offload_api::offload_activate_chosen_device
subroutine, public offload_activate_chosen_device()
Activates the device selected via offload_set_chosen_device().
Definition offload_api.F:175
offload_api::offload_get_device_count
integer function, public offload_get_device_count()
Returns the number of available devices.
Definition offload_api.F:112
offload_api::offload_init
subroutine, public offload_init()
Initialize runtime.
Definition offload_api.F:96
nnp_environment_types::nnp_type
Main data type collecting all relevant data for neural network potentials.
Definition nnp_environment_types.F:97