Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2026 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : ! **************************************************************************************************
9 : !> \brief Experimental CP2K-native GPW real-space-grid path for SKALA TorchScript models.
10 : ! **************************************************************************************************
11 : MODULE skala_gpw_functional
12 : USE OMP_LIB, ONLY: omp_get_max_threads,&
13 : omp_set_num_threads
14 : USE cell_types, ONLY: cell_type,&
15 : pbc
16 : USE cp_array_utils, ONLY: cp_3d_r_cp_type
17 : USE cp_log_handling, ONLY: cp_logger_get_default_io_unit
18 : USE input_section_types, ONLY: section_get_rval,&
19 : section_vals_get_subs_vals,&
20 : section_vals_get_subs_vals2,&
21 : section_vals_type,&
22 : section_vals_val_get
23 : USE kinds, ONLY: default_path_length,&
24 : dp,&
25 : int_8
26 : USE message_passing, ONLY: mp_comm_type
27 : USE offload_api, ONLY: offload_set_chosen_device
28 : USE particle_types, ONLY: particle_type
29 : USE pw_grid_types, ONLY: pw_grid_type
30 : USE pw_methods, ONLY: pw_scale,&
31 : pw_zero
32 : USE pw_pool_types, ONLY: pw_pool_type
33 : USE pw_types, ONLY: pw_c1d_gs_type,&
34 : pw_r3d_rs_type
35 : USE qs_grid_atom, ONLY: grid_atom_type
36 : USE skala_gpw_features, ONLY: &
37 : skala_gpw_atom_partition_hard, skala_gpw_atom_partition_smooth, &
38 : skala_gpw_atom_subchunk_count, skala_gpw_atom_subchunk_layout, skala_gpw_feature_build, &
39 : skala_gpw_feature_build_atom_subchunk_bounds, skala_gpw_feature_release, &
40 : skala_gpw_feature_type, skala_gpw_smooth_partition_derivatives, &
41 : smooth_partition_atomic_weight_scale_derivative
42 : USE skala_torch_api, ONLY: skala_torch_model_get_exc,&
43 : skala_torch_model_load,&
44 : skala_torch_model_release,&
45 : skala_torch_model_type
46 : USE string_utilities, ONLY: uppercase
47 : USE torch_api, ONLY: &
48 : torch_cuda_device_count, torch_cuda_is_available, torch_dict_create, torch_dict_insert, &
49 : torch_dict_release, torch_dict_type, torch_tensor_backward_scalar, torch_tensor_data_ptr, &
50 : torch_tensor_from_array, torch_tensor_grad, torch_tensor_grad_batch3, torch_tensor_narrow, &
51 : torch_tensor_release, torch_tensor_reset_from_array, torch_tensor_to_device_leaf, &
52 : torch_tensor_type, torch_use_cuda
53 : USE xc_input_constants, ONLY: skala_gapw_atom_composite_grid,&
54 : skala_gapw_common_grid,&
55 : skala_gapw_cp2k_default,&
56 : skala_gapw_direct_valence,&
57 : skala_gapw_paw_one_center,&
58 : skala_gapw_paw_one_center_split
59 : USE xc_rho_cflags_types, ONLY: xc_rho_cflags_type
60 : USE xc_rho_set_types, ONLY: xc_rho_set_create,&
61 : xc_rho_set_get,&
62 : xc_rho_set_release,&
63 : xc_rho_set_type,&
64 : xc_rho_set_update
65 : USE xc_util, ONLY: xc_pw_divergence,&
66 : xc_requires_tmp_g
67 : #include "../base/base_uses.f90"
68 :
69 : IMPLICIT NONE
70 :
71 : PRIVATE
72 :
73 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'skala_gpw_functional'
74 : INTEGER, PARAMETER, PRIVATE :: atom_chunk_auto_max_rows = 400000, &
75 : atom_chunk_auto_min_rows = 100000, &
76 : atom_chunk_auto_row_quantum = 100000, &
77 : ncollapsed_grad_per_point = 5, ngrad_per_point = 10
78 : INTEGER, PARAMETER, PUBLIC :: skala_gapw_density_partition_hard_minus_soft = 1, &
79 : skala_gapw_density_partition_hard_only = 2, &
80 : skala_gapw_density_partition_soft_only = 3, &
81 : skala_gapw_density_partition_none = 4
82 :
83 : TYPE skala_gapw_atom_cuda_tensor_cache_type
84 : TYPE(torch_tensor_type) :: atomic_coords_t, atomic_grid_weights_t, &
85 : density_t, grad_t, grid_coords_t, &
86 : grid_weights_t, kin_t
87 : END TYPE skala_gapw_atom_cuda_tensor_cache_type
88 :
89 : PUBLIC :: build_vxc_from_feature_grads, ensure_native_skala_grid_scope, get_gauxc_section, &
90 : skala_gapw_atom_composite_energy, skala_gapw_atom_vxc_of_r, &
91 : native_skala_gapw_atom_composite_requested, &
92 : native_skala_uses_atom_composite_grid, &
93 : native_skala_gapw_composite_direct_ao, native_skala_gapw_composite_reference, &
94 : native_skala_gapw_density_partition, skala_gpw_eval, skala_gpw_weight_derivative, &
95 : skala_gapw_representation, xc_section_uses_native_skala_evaluator, &
96 : xc_section_uses_native_skala_grid, xc_section_uses_gauxc_model
97 :
98 : TYPE(skala_torch_model_type), SAVE :: cached_model
99 : TYPE(skala_gapw_atom_cuda_tensor_cache_type), SAVE, TARGET :: cached_atom_cuda_tensors
100 : CHARACTER(len=default_path_length), SAVE :: cached_model_path = ""
101 : LOGICAL, SAVE :: cached_model_loaded = .FALSE.
102 : LOGICAL, SAVE :: cached_model_fp64 = .FALSE.
103 : INTEGER, SAVE :: cached_model_cuda_device = -3
104 : INTEGER, SAVE :: logged_cuda_device = -3, &
105 : logged_cuda_device_count = -1, &
106 : logged_cuda_nproc = -1, &
107 : logged_cuda_request = -3
108 :
109 : CONTAINS
110 :
111 : ! **************************************************************************************************
112 : !> \brief Return true if the GAUXC subsection requests the CP2K-native GPW grid path.
113 : !> \param xc_section ...
114 : !> \return ...
115 : ! **************************************************************************************************
116 308416 : FUNCTION xc_section_uses_native_skala_grid(xc_section) RESULT(uses_native_grid)
117 : TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
118 : LOGICAL :: uses_native_grid
119 :
120 : TYPE(section_vals_type), POINTER :: gauxc_section
121 :
122 308416 : uses_native_grid = .FALSE.
123 308416 : gauxc_section => get_gauxc_section(xc_section)
124 308416 : IF (ASSOCIATED(gauxc_section)) THEN
125 1814 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID", l_val=uses_native_grid)
126 : END IF
127 :
128 308416 : END FUNCTION xc_section_uses_native_skala_grid
129 :
130 : ! **************************************************************************************************
131 : !> \brief Return the pseudopotential GAPW representation selected for an active model.
132 : !> \param xc_section ...
133 : !> \return ...
134 : ! **************************************************************************************************
135 201638 : FUNCTION skala_gapw_representation(xc_section) RESULT(representation)
136 : TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
137 : INTEGER :: representation
138 :
139 : TYPE(section_vals_type), POINTER :: gauxc_section
140 :
141 : representation = skala_gapw_cp2k_default
142 200894 : IF (.NOT. xc_section_uses_gauxc_model(xc_section)) RETURN
143 :
144 744 : gauxc_section => get_gauxc_section(xc_section)
145 744 : CPASSERT(ASSOCIATED(gauxc_section))
146 : CALL section_vals_val_get(gauxc_section, "PSEUDOPOTENTIAL_GAPW_REPRESENTATION", &
147 744 : i_val=representation)
148 :
149 : SELECT CASE (representation)
150 : CASE (skala_gapw_direct_valence, skala_gapw_paw_one_center, &
151 : skala_gapw_cp2k_default, skala_gapw_paw_one_center_split)
152 0 : CONTINUE
153 : CASE DEFAULT
154 744 : CALL cp_abort(__LOCATION__, "Unknown pseudopotential GAPW representation.")
155 : END SELECT
156 :
157 : END FUNCTION skala_gapw_representation
158 :
159 : ! **************************************************************************************************
160 : !> \brief Return true when SKALA must be evaluated by the CP2K-native grid machinery.
161 : !> \param xc_section ...
162 : !> \return ...
163 : ! **************************************************************************************************
164 169711 : FUNCTION xc_section_uses_native_skala_evaluator(xc_section) RESULT(uses_native_evaluator)
165 : TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
166 : LOGICAL :: uses_native_evaluator
167 :
168 : uses_native_evaluator = xc_section_uses_native_skala_grid(xc_section) .OR. &
169 : skala_gapw_representation(xc_section) == &
170 169711 : skala_gapw_paw_one_center
171 :
172 169711 : END FUNCTION xc_section_uses_native_skala_evaluator
173 :
174 : ! **************************************************************************************************
175 : !> \brief Return true if native SKALA should use the full GAPW ORB density on one common grid.
176 : !> \param xc_section ...
177 : !> \return ...
178 : ! **************************************************************************************************
179 157284 : FUNCTION native_skala_gapw_composite_reference(xc_section) RESULT(use_composite_reference)
180 : TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
181 : LOGICAL :: use_composite_reference
182 :
183 : LOGICAL :: native_grid
184 : TYPE(section_vals_type), POINTER :: gauxc_section
185 :
186 157284 : use_composite_reference = .FALSE.
187 157284 : native_grid = .FALSE.
188 157284 : gauxc_section => get_gauxc_section(xc_section)
189 157284 : IF (ASSOCIATED(gauxc_section)) THEN
190 1760 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID", l_val=native_grid)
191 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_GAPW_COMPOSITE_REFERENCE", &
192 1760 : l_val=use_composite_reference)
193 : END IF
194 157284 : use_composite_reference = native_grid .AND. use_composite_reference
195 :
196 157284 : END FUNCTION native_skala_gapw_composite_reference
197 :
198 : ! **************************************************************************************************
199 : !> \brief Return true when the explicit atom-centered composite reference is requested.
200 : !> \param xc_section ...
201 : !> \return ...
202 : ! **************************************************************************************************
203 1866 : FUNCTION native_skala_gapw_atom_composite_requested(xc_section) RESULT(use_atom_composite)
204 : TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
205 : LOGICAL :: use_atom_composite
206 :
207 : LOGICAL :: native_grid
208 : TYPE(section_vals_type), POINTER :: gauxc_section
209 :
210 1866 : use_atom_composite = .FALSE.
211 1866 : native_grid = .FALSE.
212 1866 : gauxc_section => get_gauxc_section(xc_section)
213 1866 : IF (ASSOCIATED(gauxc_section)) THEN
214 352 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID", l_val=native_grid)
215 : CALL section_vals_val_get(gauxc_section, &
216 : "NATIVE_GRID_GAPW_ATOM_COMPOSITE_REFERENCE", &
217 352 : l_val=use_atom_composite)
218 : END IF
219 1866 : use_atom_composite = native_grid .AND. use_atom_composite
220 :
221 1866 : END FUNCTION native_skala_gapw_atom_composite_requested
222 :
223 : ! **************************************************************************************************
224 : !> \brief Return true when native Skala uses atom-centered grids.
225 : !> \param xc_section ...
226 : !> \return ...
227 : ! **************************************************************************************************
228 616 : FUNCTION native_skala_uses_atom_composite_grid(xc_section) RESULT(use_atom_composite)
229 : TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
230 : LOGICAL :: use_atom_composite
231 :
232 : INTEGER :: composite_grid
233 : TYPE(section_vals_type), POINTER :: gauxc_section
234 :
235 616 : use_atom_composite = .FALSE.
236 616 : composite_grid = skala_gapw_atom_composite_grid
237 616 : gauxc_section => get_gauxc_section(xc_section)
238 616 : IF (ASSOCIATED(gauxc_section)) THEN
239 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_LAYOUT", &
240 616 : i_val=composite_grid)
241 : END IF
242 : SELECT CASE (composite_grid)
243 : CASE (skala_gapw_common_grid)
244 552 : CONTINUE
245 : CASE (skala_gapw_atom_composite_grid)
246 552 : use_atom_composite = .TRUE.
247 : CASE DEFAULT
248 616 : CALL cp_abort(__LOCATION__, "Unknown native-grid layout.")
249 : END SELECT
250 :
251 616 : END FUNCTION native_skala_uses_atom_composite_grid
252 :
253 : ! **************************************************************************************************
254 : !> \brief Return true if the GAPW composite reference uses direct full-ORB collocation.
255 : !> \param xc_section ...
256 : !> \return ...
257 : ! **************************************************************************************************
258 8 : FUNCTION native_skala_gapw_composite_direct_ao(xc_section) RESULT(use_direct_ao)
259 : TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
260 : LOGICAL :: use_direct_ao
261 :
262 : TYPE(section_vals_type), POINTER :: gauxc_section
263 :
264 8 : use_direct_ao = .FALSE.
265 8 : gauxc_section => get_gauxc_section(xc_section)
266 8 : IF (ASSOCIATED(gauxc_section)) THEN
267 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_GAPW_COMPOSITE_DIRECT_AO", &
268 8 : l_val=use_direct_ao)
269 : END IF
270 8 : use_direct_ao = native_skala_gapw_composite_reference(xc_section) .AND. use_direct_ao
271 :
272 8 : END FUNCTION native_skala_gapw_composite_direct_ao
273 :
274 : ! **************************************************************************************************
275 : !> \brief Return true if the GAUXC subsection requests a model evaluation.
276 : !> \param xc_section ...
277 : !> \return ...
278 : ! **************************************************************************************************
279 235606 : FUNCTION xc_section_uses_gauxc_model(xc_section) RESULT(uses_gauxc_model)
280 : TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
281 : LOGICAL :: uses_gauxc_model
282 :
283 : CHARACTER(len=default_path_length) :: model_key, model_name, xc_key, xc_name
284 : TYPE(section_vals_type), POINTER :: gauxc_section
285 :
286 235606 : uses_gauxc_model = .FALSE.
287 235606 : gauxc_section => get_gauxc_section(xc_section)
288 235606 : IF (ASSOCIATED(gauxc_section)) THEN
289 1366 : CALL section_vals_val_get(gauxc_section, "MODEL", c_val=model_name)
290 1366 : CALL section_vals_val_get(gauxc_section, "FUNCTIONAL", c_val=xc_name)
291 1366 : model_key = ADJUSTL(model_name)
292 1366 : xc_key = ADJUSTL(xc_name)
293 1366 : CALL uppercase(model_key)
294 1366 : CALL uppercase(xc_key)
295 : uses_gauxc_model = (TRIM(model_key) /= "" .AND. TRIM(model_key) /= "NONE" .AND. &
296 1366 : TRIM(model_key) /= TRIM(xc_key))
297 : END IF
298 :
299 235606 : END FUNCTION xc_section_uses_gauxc_model
300 :
301 : ! **************************************************************************************************
302 : !> \brief Return the hard/soft GAPW one-center density partition for native SKALA.
303 : !> \param xc_section ...
304 : !> \return ...
305 : ! **************************************************************************************************
306 270 : FUNCTION native_skala_gapw_density_partition(xc_section) RESULT(partition)
307 : TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
308 : INTEGER :: partition
309 :
310 : TYPE(section_vals_type), POINTER :: gauxc_section
311 :
312 270 : partition = skala_gapw_density_partition_hard_minus_soft
313 270 : gauxc_section => get_gauxc_section(xc_section)
314 270 : IF (ASSOCIATED(gauxc_section)) THEN
315 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_GAPW_DENSITY_PARTITION", &
316 270 : i_val=partition)
317 : END IF
318 :
319 : SELECT CASE (partition)
320 : CASE (skala_gapw_density_partition_hard_minus_soft, &
321 : skala_gapw_density_partition_hard_only, &
322 : skala_gapw_density_partition_soft_only, &
323 : skala_gapw_density_partition_none)
324 0 : CONTINUE
325 : CASE DEFAULT
326 : CALL cp_abort(__LOCATION__, &
327 270 : "Unknown GAUXC%NATIVE_GRID_GAPW_DENSITY_PARTITION value.")
328 : END SELECT
329 :
330 270 : END FUNCTION native_skala_gapw_density_partition
331 :
332 : ! **************************************************************************************************
333 : !> \brief Enforce the currently implemented native SKALA GPW input scope.
334 : !> \param xc_section ...
335 : ! **************************************************************************************************
336 600 : SUBROUTINE ensure_native_skala_grid_scope(xc_section)
337 : TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
338 :
339 : CHARACTER(len=default_path_length) :: model_key, model_name
340 : INTEGER :: ifun, nfun
341 : LOGICAL :: native_grid
342 : TYPE(section_vals_type), POINTER :: functionals, gauxc_section, xc_fun
343 :
344 300 : NULLIFY (gauxc_section)
345 300 : IF (.NOT. ASSOCIATED(xc_section)) THEN
346 0 : CPABORT("Native SKALA GPW requires an XC section")
347 : END IF
348 :
349 300 : functionals => section_vals_get_subs_vals(xc_section, "XC_FUNCTIONAL")
350 300 : IF (.NOT. ASSOCIATED(functionals)) THEN
351 0 : CPABORT("Native SKALA GPW requires an XC_FUNCTIONAL section")
352 : END IF
353 :
354 300 : nfun = 0
355 300 : ifun = 0
356 : DO
357 600 : ifun = ifun + 1
358 600 : xc_fun => section_vals_get_subs_vals2(functionals, i_section=ifun)
359 600 : IF (.NOT. ASSOCIATED(xc_fun)) EXIT
360 300 : nfun = nfun + 1
361 600 : IF (xc_fun%section%name == "GAUXC") gauxc_section => xc_fun
362 : END DO
363 :
364 300 : IF (.NOT. ASSOCIATED(gauxc_section)) THEN
365 0 : CPABORT("Native SKALA GPW requires an XC_FUNCTIONAL%GAUXC section")
366 : END IF
367 300 : IF (nfun /= 1) THEN
368 0 : CPABORT("Native SKALA GPW requires GAUXC to be the only XC functional")
369 : END IF
370 :
371 300 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID", l_val=native_grid)
372 300 : IF (.NOT. native_grid .AND. &
373 0 : .NOT. xc_section_uses_native_skala_evaluator(xc_section)) RETURN
374 :
375 300 : CALL section_vals_val_get(gauxc_section, "MODEL", c_val=model_name)
376 300 : model_key = ADJUSTL(model_name)
377 300 : CALL uppercase(model_key)
378 300 : IF (TRIM(model_key) == "NONE" .OR. TRIM(model_key) == "") THEN
379 0 : CPABORT("Native SKALA GPW requires GAUXC%MODEL SKALA or a TorchScript model path")
380 : END IF
381 :
382 : END SUBROUTINE ensure_native_skala_grid_scope
383 :
384 : ! **************************************************************************************************
385 : !> \brief Evaluate SKALA energy and first derivatives on a CP2K GPW grid.
386 : !> \param vxc_rho ...
387 : !> \param vxc_tau ...
388 : !> \param exc ...
389 : !> \param rho_r ...
390 : !> \param rho_g ...
391 : !> \param tau ...
392 : !> \param xc_section ...
393 : !> \param weights ...
394 : !> \param pw_pool ...
395 : !> \param particle_set ...
396 : !> \param cell ...
397 : !> \param compute_virial ...
398 : !> \param virial_xc ...
399 : !> \param just_energy ...
400 : !> \param atom_force ...
401 : ! **************************************************************************************************
402 48 : SUBROUTINE skala_gpw_eval(vxc_rho, vxc_tau, exc, rho_r, rho_g, tau, xc_section, &
403 : weights, pw_pool, particle_set, cell, compute_virial, virial_xc, &
404 48 : just_energy, atom_force)
405 : TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: vxc_rho, vxc_tau
406 : REAL(KIND=dp), INTENT(OUT) :: exc
407 : TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: rho_r
408 : TYPE(pw_c1d_gs_type), DIMENSION(:), POINTER :: rho_g
409 : TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: tau
410 : TYPE(section_vals_type), POINTER :: xc_section
411 : TYPE(pw_r3d_rs_type), POINTER :: weights
412 : TYPE(pw_pool_type), POINTER :: pw_pool
413 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
414 : TYPE(cell_type), POINTER :: cell
415 : LOGICAL, INTENT(IN) :: compute_virial
416 : REAL(KIND=dp), DIMENSION(3, 3), INTENT(OUT) :: virial_xc
417 : LOGICAL, INTENT(IN), OPTIONAL :: just_energy
418 : REAL(KIND=dp), DIMENSION(:, :), INTENT(OUT), &
419 : OPTIONAL :: atom_force
420 :
421 : CHARACTER(len=default_path_length) :: model_path
422 : INTEGER :: i, ipt, ispin, iw, j, k, native_grid_atom_chunk_max_rows, &
423 : native_grid_atom_partition, native_grid_atom_subchunks, native_grid_cuda_device, nspins, &
424 : omp_max_threads_restore, phase_handle, selected_cuda_device, xc_deriv_method_id, &
425 : xc_rho_smooth_id
426 : INTEGER, DIMENSION(2, 3) :: bo
427 : LOGICAL :: has_atom_chunk_work, have_atom_coord_grad, lsd, my_just_energy, &
428 : native_grid_atom_chunk_routing, native_grid_atom_chunks, native_grid_diagnostics, &
429 : native_grid_use_cuda, needs_atom_force, use_atom_subchunks
430 : REAL(KIND=dp) :: density_contraction, tau_contraction
431 48 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: density_grad, kin_grad
432 48 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: grad_grad
433 : REAL(KIND=dp), DIMENSION(3, 3) :: virial_before
434 : TYPE(section_vals_type), POINTER :: gauxc_section
435 48 : TYPE(skala_gpw_feature_type) :: features
436 : TYPE(torch_tensor_type) :: atom_coord_grad_t, &
437 : atomic_grid_weight_grad_t, exc_tensor, &
438 : grid_coord_grad_t, grid_weight_grad_t
439 : TYPE(xc_rho_cflags_type) :: needs
440 : TYPE(xc_rho_set_type) :: rho_set
441 :
442 48 : virial_xc = 0.0_dp
443 48 : exc = 0.0_dp
444 48 : my_just_energy = .FALSE.
445 48 : IF (PRESENT(just_energy)) my_just_energy = just_energy
446 48 : needs_atom_force = PRESENT(atom_force)
447 112 : IF (needs_atom_force) atom_force = 0.0_dp
448 48 : have_atom_coord_grad = .FALSE.
449 48 : omp_max_threads_restore = omp_get_max_threads()
450 :
451 48 : IF (compute_virial .AND. my_just_energy) THEN
452 : CALL cp_abort(__LOCATION__, &
453 0 : "Native SKALA GPW stress/virial requires feature gradients.")
454 : END IF
455 48 : IF (.NOT. ASSOCIATED(rho_g)) THEN
456 : CALL cp_abort(__LOCATION__, &
457 0 : "Native SKALA GPW requires the reciprocal-space density to form density gradients.")
458 : END IF
459 48 : IF (.NOT. ASSOCIATED(tau)) THEN
460 : CALL cp_abort(__LOCATION__, &
461 0 : "Native SKALA GPW requires the kinetic-energy density.")
462 : END IF
463 :
464 48 : nspins = SIZE(rho_r)
465 48 : lsd = (nspins /= 1)
466 48 : CALL get_skala_model_path(xc_section, model_path)
467 48 : gauxc_section => get_gauxc_section(xc_section)
468 48 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_USE_CUDA", l_val=native_grid_use_cuda)
469 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_CUDA_DEVICE", &
470 48 : i_val=native_grid_cuda_device)
471 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNKS", &
472 48 : l_val=native_grid_atom_chunks)
473 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNK_ROUTING", &
474 48 : l_val=native_grid_atom_chunk_routing)
475 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNK_MAX_ROWS", &
476 48 : i_val=native_grid_atom_chunk_max_rows)
477 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_PARTITION", &
478 48 : i_val=native_grid_atom_partition)
479 22 : SELECT CASE (native_grid_atom_partition)
480 : CASE (1)
481 22 : native_grid_atom_partition = skala_gpw_atom_partition_hard
482 : CASE (2)
483 26 : native_grid_atom_partition = skala_gpw_atom_partition_smooth
484 : CASE DEFAULT
485 : CALL cp_abort(__LOCATION__, &
486 48 : "Unknown GAUXC%NATIVE_GRID_ATOM_PARTITION value.")
487 : END SELECT
488 48 : native_grid_atom_chunk_routing = native_grid_atom_chunk_routing .OR. native_grid_atom_chunks
489 48 : native_grid_atom_chunks = native_grid_atom_chunks .OR. native_grid_atom_chunk_routing
490 48 : IF (native_grid_atom_chunk_max_rows < -1) THEN
491 : CALL cp_abort(__LOCATION__, &
492 0 : "GAUXC%NATIVE_GRID_ATOM_CHUNK_MAX_ROWS must be -1, zero, or positive.")
493 : END IF
494 48 : IF (needs_atom_force .OR. compute_virial) THEN
495 8 : IF (native_grid_atom_partition == skala_gpw_atom_partition_hard) THEN
496 0 : native_grid_atom_partition = skala_gpw_atom_partition_smooth
497 : END IF
498 8 : native_grid_atom_chunk_routing = .FALSE.
499 8 : native_grid_atom_chunks = .FALSE.
500 : END IF
501 : ! The portable SKALA export used by the regtests builds ragged-index tensors on CPU.
502 48 : CALL torch_use_cuda(native_grid_use_cuda)
503 : selected_cuda_device = configure_native_grid_cuda( &
504 48 : native_grid_use_cuda, native_grid_cuda_device, rho_r(1)%pw_grid%para%group)
505 48 : CALL ensure_model_loaded(model_path, selected_cuda_device, xc_section)
506 :
507 48 : IF (lsd) THEN
508 6 : needs%rho_spin = .TRUE.
509 6 : needs%drho_spin = .TRUE.
510 6 : needs%tau_spin = .TRUE.
511 : ELSE
512 42 : needs%rho = .TRUE.
513 42 : needs%drho = .TRUE.
514 42 : needs%tau = .TRUE.
515 : END IF
516 :
517 48 : CALL section_vals_val_get(xc_section, "XC_GRID%XC_DERIV", i_val=xc_deriv_method_id)
518 48 : CALL section_vals_val_get(xc_section, "XC_GRID%XC_SMOOTH_RHO", i_val=xc_rho_smooth_id)
519 :
520 : CALL xc_rho_set_create(rho_set, &
521 : rho_r(1)%pw_grid%bounds_local, &
522 : rho_cutoff=section_get_rval(xc_section, "density_cutoff"), &
523 : drho_cutoff=section_get_rval(xc_section, "gradient_cutoff"), &
524 48 : tau_cutoff=section_get_rval(xc_section, "tau_cutoff"))
525 : CALL xc_rho_set_update(rho_set, rho_r, rho_g, tau, needs, &
526 48 : xc_deriv_method_id, xc_rho_smooth_id, pw_pool)
527 :
528 : CALL skala_gpw_feature_build(features, rho_set, rho_r, particle_set, cell, &
529 : requires_grad=(.NOT. my_just_energy), weights=weights, &
530 : requires_coordinate_grad=(needs_atom_force .OR. compute_virial), &
531 : requires_stress_grad=compute_virial, &
532 : use_atom_chunks=native_grid_atom_chunks, &
533 : route_atom_chunks=native_grid_atom_chunk_routing, &
534 88 : atom_partition=native_grid_atom_partition)
535 48 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_DIAGNOSTICS", l_val=native_grid_diagnostics)
536 48 : IF (native_grid_diagnostics) THEN
537 24 : CALL print_native_grid_diagnostics(features, rho_r(1)%pw_grid%para%group%mepos == 0)
538 : END IF
539 :
540 48 : IF (features%uses_atom_chunks .AND. native_grid_atom_chunk_max_rows == -1) THEN
541 : native_grid_atom_chunk_max_rows = auto_atom_chunk_max_rows(features, &
542 36 : rho_r(1)%pw_grid%para%group)
543 : END IF
544 48 : IF (native_grid_diagnostics .AND. features%uses_atom_chunks .AND. &
545 : rho_r(1)%pw_grid%para%group%mepos == 0) THEN
546 12 : iw = cp_logger_get_default_io_unit()
547 12 : IF (iw > 0) THEN
548 : WRITE (UNIT=iw, FMT="(T2,A,1X,I0)") &
549 12 : "SKALA_GPW| Native grid atom chunk max rows", native_grid_atom_chunk_max_rows
550 : END IF
551 : END IF
552 48 : native_grid_atom_subchunks = 1
553 48 : IF (features%uses_atom_chunks .AND. native_grid_atom_chunk_max_rows > 0) THEN
554 40 : native_grid_atom_subchunks = skala_gpw_atom_subchunk_count(native_grid_atom_chunk_max_rows)
555 40 : CALL rho_r(1)%pw_grid%para%group%max(native_grid_atom_subchunks)
556 : END IF
557 48 : use_atom_subchunks = features%uses_atom_chunks .AND. native_grid_atom_subchunks > 1
558 48 : has_atom_chunk_work = .NOT. features%uses_atom_chunks .OR. features%chunk_feature_count > 0
559 48 : exc = 0.0_dp
560 48 : IF (use_atom_subchunks) THEN
561 : CALL evaluate_atom_subchunks(features, rho_r(1)%pw_grid%para%group, &
562 : native_grid_atom_chunk_max_rows, &
563 : compute_grads=(.NOT. my_just_energy), exc=exc, &
564 : density_grad=density_grad, grad_grad=grad_grad, &
565 2 : kin_grad=kin_grad, collapse_spin_grads=(nspins == 1))
566 46 : ELSE IF (has_atom_chunk_work) THEN
567 : CALL skala_torch_model_get_exc(cached_model, features%inputs, &
568 46 : features%grid_weights_t, exc_tensor, exc)
569 : END IF
570 48 : IF (features%uses_atom_chunks) CALL rho_r(1)%pw_grid%para%group%sum(exc)
571 :
572 48 : IF (.NOT. my_just_energy) THEN
573 48 : IF (.NOT. use_atom_subchunks) THEN
574 46 : IF (has_atom_chunk_work) THEN
575 46 : CALL timeset("skala_gpw_backward", phase_handle)
576 46 : CALL torch_tensor_backward_scalar(exc_tensor)
577 46 : CALL timestop(phase_handle)
578 :
579 46 : IF (compute_virial) THEN
580 6 : IF (native_grid_diagnostics) virial_before = virial_xc
581 : CALL build_weight_virial(virial_xc, features, exc, grid_weight_grad_t, &
582 : atomic_grid_weight_grad_t, &
583 : rho_r(1)%pw_grid%para%group%mepos == 0, &
584 6 : native_grid_diagnostics)
585 6 : IF (native_grid_diagnostics) THEN
586 : CALL print_virial_delta("weight-residual", virial_xc - virial_before, &
587 0 : rho_r(1)%pw_grid%para%group%mepos == 0)
588 : END IF
589 : END IF
590 : END IF
591 :
592 46 : CALL timeset("skala_gpw_grad_fetch", phase_handle)
593 46 : IF (features%uses_atom_chunks) THEN
594 : CALL fetch_and_gather_atom_chunk_grads(features, rho_r(1)%pw_grid%para%group, &
595 38 : density_grad, grad_grad, kin_grad)
596 : ELSE
597 8 : CALL fetch_local_feature_grads(features, density_grad, grad_grad, kin_grad)
598 : END IF
599 46 : CALL timestop(phase_handle)
600 : END IF
601 48 : IF (needs_atom_force) THEN
602 : CALL add_explicit_coordinate_force(atom_force, features, atom_coord_grad_t, &
603 8 : rho_r(1)%pw_grid%para%group%mepos == 0)
604 8 : IF (features%atom_partition == skala_gpw_atom_partition_smooth) THEN
605 : CALL add_smooth_partition_force(atom_force, features, particle_set, cell, rho_r, &
606 8 : grid_weight_grad_t, atomic_grid_weight_grad_t)
607 : END IF
608 : have_atom_coord_grad = .TRUE.
609 : END IF
610 :
611 48 : CALL timeset("skala_gpw_vxc_unpack", phase_handle)
612 48 : IF (native_grid_diagnostics) THEN
613 240 : bo = rho_r(1)%pw_grid%bounds_local
614 24 : density_contraction = 0.0_dp
615 24 : tau_contraction = 0.0_dp
616 24 : ipt = 0
617 1004 : DO k = bo(1, 3), bo(2, 3)
618 41848 : DO j = bo(1, 2), bo(2, 2)
619 905054 : DO i = bo(1, 1), bo(2, 1)
620 863230 : ipt = ipt + 1
621 904074 : IF (nspins == 1) THEN
622 : density_contraction = density_contraction + rho_r(1)%array(i, j, k)* &
623 589855 : 0.5_dp*(density_grad(ipt, 1) + density_grad(ipt, 2))
624 : tau_contraction = tau_contraction + tau(1)%array(i, j, k)* &
625 589855 : 0.5_dp*(kin_grad(ipt, 1) + kin_grad(ipt, 2))
626 : ELSE
627 820125 : DO ispin = 1, nspins
628 : density_contraction = density_contraction + &
629 546750 : rho_r(ispin)%array(i, j, k)*density_grad(ipt, ispin)
630 : tau_contraction = tau_contraction + &
631 820125 : tau(ispin)%array(i, j, k)*kin_grad(ipt, ispin)
632 : END DO
633 : END IF
634 : END DO
635 : END DO
636 : END DO
637 24 : CALL rho_r(1)%pw_grid%para%group%sum(density_contraction)
638 24 : CALL rho_r(1)%pw_grid%para%group%sum(tau_contraction)
639 24 : IF (rho_r(1)%pw_grid%para%group%mepos == 0) THEN
640 12 : iw = cp_logger_get_default_io_unit()
641 12 : IF (iw > 0) THEN
642 : WRITE (iw, "(T2,A,1X,ES20.10)") &
643 12 : "SKALA_GPW| XC density-gradient contraction", density_contraction
644 : WRITE (iw, "(T2,A,1X,ES20.10)") &
645 12 : "SKALA_GPW| XC kinetic-gradient contraction", tau_contraction
646 : END IF
647 : END IF
648 : END IF
649 48 : IF (compute_virial) THEN
650 6 : IF (native_grid_diagnostics) virial_before = virial_xc
651 6 : CALL build_virial_from_feature_grads(virial_xc, rho_set, rho_r, grad_grad)
652 6 : IF (native_grid_diagnostics) THEN
653 : CALL print_virial_delta("feature-gradient", virial_xc - virial_before, &
654 0 : rho_r(1)%pw_grid%para%group%mepos == 0)
655 0 : virial_before = virial_xc
656 : END IF
657 6 : IF (.NOT. have_atom_coord_grad) THEN
658 0 : CALL torch_tensor_grad(features%coarse_0_atomic_coords_t, atom_coord_grad_t)
659 0 : have_atom_coord_grad = .TRUE.
660 : END IF
661 : CALL build_static_coordinate_virial(virial_xc, features, atom_coord_grad_t, &
662 : grid_coord_grad_t, &
663 : rho_r(1)%pw_grid%para%group%mepos == 0, &
664 6 : native_grid_diagnostics)
665 6 : IF (native_grid_diagnostics) THEN
666 : CALL print_virial_delta("static-coordinates", virial_xc - virial_before, &
667 0 : rho_r(1)%pw_grid%para%group%mepos == 0)
668 0 : virial_before = virial_xc
669 : END IF
670 6 : IF (features%atom_partition == skala_gpw_atom_partition_smooth) THEN
671 : CALL build_smooth_partition_virial(virial_xc, features, particle_set, cell, rho_r, &
672 6 : grid_weight_grad_t, atomic_grid_weight_grad_t)
673 6 : IF (native_grid_diagnostics) THEN
674 : CALL print_virial_delta("smooth-partition", virial_xc - virial_before, &
675 0 : rho_r(1)%pw_grid%para%group%mepos == 0)
676 : virial_before = virial_xc
677 : END IF
678 : END IF
679 : END IF
680 48 : CALL omp_set_num_threads(omp_max_threads_restore)
681 : CALL build_vxc_from_feature_grads(vxc_rho, vxc_tau, rho_r, pw_pool, &
682 : density_grad, grad_grad, kin_grad, &
683 48 : xc_deriv_method_id)
684 48 : CALL timestop(phase_handle)
685 :
686 48 : CALL timeset("skala_gpw_grad_release", phase_handle)
687 48 : DEALLOCATE (density_grad, grad_grad, kin_grad)
688 48 : IF (have_atom_coord_grad) CALL torch_tensor_release(atom_coord_grad_t)
689 48 : CALL timestop(phase_handle)
690 : END IF
691 :
692 48 : CALL timeset("skala_gpw_cleanup", phase_handle)
693 48 : IF (.NOT. use_atom_subchunks .AND. has_atom_chunk_work) CALL torch_tensor_release(exc_tensor)
694 48 : CALL skala_gpw_feature_release(features)
695 48 : CALL xc_rho_set_release(rho_set, pw_pool=pw_pool)
696 48 : CALL torch_use_cuda(.TRUE.)
697 48 : CALL omp_set_num_threads(omp_max_threads_restore)
698 48 : CALL timestop(phase_handle)
699 :
700 960 : END SUBROUTINE skala_gpw_eval
701 :
702 : ! **************************************************************************************************
703 : !> \brief Evaluate the derivative of native SKALA XC energy with respect to CP2K's external
704 : !> real-space integration-weight multiplier.
705 : !> \param weight_deriv_r ...
706 : !> \param rho_r ...
707 : !> \param rho_g ...
708 : !> \param tau ...
709 : !> \param xc_section ...
710 : !> \param weights ...
711 : !> \param pw_pool ...
712 : !> \param particle_set ...
713 : !> \param cell ...
714 : ! **************************************************************************************************
715 0 : SUBROUTINE skala_gpw_weight_derivative(weight_deriv_r, rho_r, rho_g, tau, xc_section, &
716 : weights, pw_pool, particle_set, cell)
717 : TYPE(pw_r3d_rs_type), INTENT(INOUT) :: weight_deriv_r
718 : TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: rho_r
719 : TYPE(pw_c1d_gs_type), DIMENSION(:), POINTER :: rho_g
720 : TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: tau
721 : TYPE(section_vals_type), POINTER :: xc_section
722 : TYPE(pw_r3d_rs_type), POINTER :: weights
723 : TYPE(pw_pool_type), POINTER :: pw_pool
724 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
725 : TYPE(cell_type), POINTER :: cell
726 :
727 : CHARACTER(len=default_path_length) :: model_path
728 : INTEGER :: feature_begin, feature_end, feature_pos, i, iatom, j, k, local_row, &
729 : native_grid_atom_partition, native_grid_cuda_device, natom, nspins, &
730 : omp_max_threads_restore, row, selected_cuda_device, xc_deriv_method_id, xc_rho_smooth_id
731 : INTEGER, DIMENSION(2, 3) :: bo
732 : LOGICAL :: lsd, native_grid_use_cuda
733 0 : LOGICAL, ALLOCATABLE, DIMENSION(:) :: included
734 : REAL(KIND=dp) :: exc, local_derivative
735 0 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: partition_weights
736 0 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: atom_coords_pbc
737 0 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: dweights_datom, dweights_dstrain
738 : REAL(KIND=dp), DIMENSION(3) :: grid_point
739 0 : REAL(KIND=dp), DIMENSION(:), POINTER :: grid_weight_grad
740 : TYPE(section_vals_type), POINTER :: gauxc_section
741 0 : TYPE(skala_gpw_feature_type) :: features
742 : TYPE(torch_tensor_type) :: exc_tensor, grid_weight_grad_t
743 : TYPE(xc_rho_cflags_type) :: needs
744 : TYPE(xc_rho_set_type) :: rho_set
745 :
746 0 : CPASSERT(ASSOCIATED(rho_r))
747 0 : CPASSERT(ASSOCIATED(rho_g))
748 0 : CPASSERT(ASSOCIATED(tau))
749 0 : CALL pw_zero(weight_deriv_r)
750 0 : omp_max_threads_restore = omp_get_max_threads()
751 :
752 0 : nspins = SIZE(rho_r)
753 0 : lsd = (nspins /= 1)
754 0 : CALL get_skala_model_path(xc_section, model_path)
755 0 : gauxc_section => get_gauxc_section(xc_section)
756 0 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_USE_CUDA", l_val=native_grid_use_cuda)
757 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_CUDA_DEVICE", &
758 0 : i_val=native_grid_cuda_device)
759 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_PARTITION", &
760 0 : i_val=native_grid_atom_partition)
761 0 : SELECT CASE (native_grid_atom_partition)
762 : CASE (1)
763 0 : native_grid_atom_partition = skala_gpw_atom_partition_hard
764 : CASE (2)
765 0 : native_grid_atom_partition = skala_gpw_atom_partition_smooth
766 : CASE DEFAULT
767 : CALL cp_abort(__LOCATION__, &
768 0 : "Unknown GAUXC%NATIVE_GRID_ATOM_PARTITION value.")
769 : END SELECT
770 0 : IF (native_grid_atom_partition == skala_gpw_atom_partition_hard) THEN
771 0 : native_grid_atom_partition = skala_gpw_atom_partition_smooth
772 : END IF
773 :
774 0 : CALL torch_use_cuda(native_grid_use_cuda)
775 : selected_cuda_device = configure_native_grid_cuda( &
776 0 : native_grid_use_cuda, native_grid_cuda_device, rho_r(1)%pw_grid%para%group)
777 0 : CALL ensure_model_loaded(model_path, selected_cuda_device, xc_section)
778 :
779 0 : IF (lsd) THEN
780 0 : needs%rho_spin = .TRUE.
781 0 : needs%drho_spin = .TRUE.
782 0 : needs%tau_spin = .TRUE.
783 : ELSE
784 0 : needs%rho = .TRUE.
785 0 : needs%drho = .TRUE.
786 0 : needs%tau = .TRUE.
787 : END IF
788 :
789 0 : CALL section_vals_val_get(xc_section, "XC_GRID%XC_DERIV", i_val=xc_deriv_method_id)
790 0 : CALL section_vals_val_get(xc_section, "XC_GRID%XC_SMOOTH_RHO", i_val=xc_rho_smooth_id)
791 :
792 : CALL xc_rho_set_create(rho_set, &
793 : rho_r(1)%pw_grid%bounds_local, &
794 : rho_cutoff=section_get_rval(xc_section, "density_cutoff"), &
795 : drho_cutoff=section_get_rval(xc_section, "gradient_cutoff"), &
796 0 : tau_cutoff=section_get_rval(xc_section, "tau_cutoff"))
797 : CALL xc_rho_set_update(rho_set, rho_r, rho_g, tau, needs, &
798 0 : xc_deriv_method_id, xc_rho_smooth_id, pw_pool)
799 :
800 : CALL skala_gpw_feature_build(features, rho_set, rho_r, particle_set, cell, &
801 : requires_grad=.FALSE., weights=weights, &
802 : requires_coordinate_grad=.FALSE., &
803 : requires_stress_grad=.TRUE., &
804 : use_atom_chunks=.FALSE., route_atom_chunks=.FALSE., &
805 0 : atom_partition=native_grid_atom_partition)
806 : CALL skala_torch_model_get_exc(cached_model, features%inputs, features%grid_weights_t, &
807 0 : exc_tensor, exc)
808 0 : CALL torch_tensor_backward_scalar(exc_tensor)
809 0 : NULLIFY (grid_weight_grad)
810 0 : CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
811 0 : CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
812 :
813 0 : natom = SIZE(particle_set)
814 0 : IF (native_grid_atom_partition == skala_gpw_atom_partition_smooth) THEN
815 : ALLOCATE (atom_coords_pbc(3, natom), included(natom), partition_weights(natom), &
816 0 : dweights_datom(3, natom, natom), dweights_dstrain(3, 3, natom))
817 0 : DO iatom = 1, natom
818 0 : atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.TRUE.)
819 : END DO
820 : END IF
821 :
822 0 : bo = rho_r(1)%pw_grid%bounds_local
823 0 : local_row = 0
824 0 : DO k = bo(1, 3), bo(2, 3)
825 0 : DO j = bo(1, 2), bo(2, 2)
826 0 : DO i = bo(1, 1), bo(2, 1)
827 0 : local_row = local_row + 1
828 0 : feature_begin = features%local_feature_offsets(local_row)
829 0 : feature_end = features%local_feature_offsets(local_row + 1) - 1
830 0 : local_derivative = 0.0_dp
831 0 : IF (native_grid_atom_partition == skala_gpw_atom_partition_hard) THEN
832 0 : CPASSERT(feature_end == feature_begin)
833 0 : row = features%local_feature_rows(feature_begin)
834 0 : local_derivative = grid_weight_grad(row)
835 : ELSE
836 0 : grid_point = native_grid_coordinate(rho_r(1)%pw_grid, [i, j, k])
837 : CALL skala_gpw_smooth_partition_derivatives( &
838 : grid_point, atom_coords_pbc, cell, partition_weights, included, &
839 0 : dweights_datom, dweights_dstrain)
840 0 : CPASSERT(feature_end - feature_begin + 1 == COUNT(included))
841 0 : feature_pos = feature_begin
842 0 : DO iatom = 1, natom
843 0 : IF (.NOT. included(iatom)) CYCLE
844 0 : row = features%local_feature_rows(feature_pos)
845 : local_derivative = local_derivative + &
846 0 : partition_weights(iatom)*grid_weight_grad(row)
847 0 : feature_pos = feature_pos + 1
848 : END DO
849 0 : CPASSERT(feature_pos == feature_end + 1)
850 : END IF
851 0 : weight_deriv_r%array(i, j, k) = local_derivative
852 : END DO
853 : END DO
854 : END DO
855 0 : CPASSERT(local_row == features%nflat_local)
856 :
857 0 : IF (ALLOCATED(atom_coords_pbc)) THEN
858 0 : DEALLOCATE (atom_coords_pbc, dweights_datom, dweights_dstrain, included, &
859 0 : partition_weights)
860 : END IF
861 0 : CALL torch_tensor_release(grid_weight_grad_t)
862 0 : CALL torch_tensor_release(exc_tensor)
863 0 : CALL skala_gpw_feature_release(features)
864 0 : CALL xc_rho_set_release(rho_set, pw_pool=pw_pool)
865 0 : CALL torch_use_cuda(.TRUE.)
866 0 : CALL omp_set_num_threads(omp_max_threads_restore)
867 :
868 0 : END SUBROUTINE skala_gpw_weight_derivative
869 :
870 : ! **************************************************************************************************
871 : !> \brief Evaluate SKALA on a GAPW one-center atomic grid.
872 : !> \param xc_section ...
873 : !> \param grid_atom ...
874 : !> \param group ...
875 : !> \param atom_coord ...
876 : !> \param rho ...
877 : !> \param drho ...
878 : !> \param tau ...
879 : !> \param weights ...
880 : !> \param lsd ...
881 : !> \param nspins ...
882 : !> \param na ...
883 : !> \param nr ...
884 : !> \param exc ...
885 : !> \param vxc ...
886 : !> \param vxg ...
887 : !> \param vtau ...
888 : !> \param energy_only ...
889 : !> \param atom_force ...
890 : !> \param atom_virial ...
891 : ! **************************************************************************************************
892 4 : SUBROUTINE skala_gapw_atom_vxc_of_r(xc_section, grid_atom, group, atom_coord, &
893 4 : rho, drho, tau, weights, lsd, nspins, na, nr, &
894 : exc, vxc, vxg, vtau, energy_only, atom_force, atom_virial)
895 : TYPE(section_vals_type), POINTER :: xc_section
896 : TYPE(grid_atom_type), POINTER :: grid_atom
897 :
898 : CLASS(mp_comm_type), INTENT(IN) :: group
899 : REAL(KIND=dp), DIMENSION(3), INTENT(IN) :: atom_coord
900 : REAL(KIND=dp), DIMENSION(:, :, :), POINTER :: rho, tau, vxc, vtau
901 : REAL(KIND=dp), DIMENSION(:, :, :, :), POINTER :: drho, vxg
902 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: weights
903 : LOGICAL, INTENT(IN) :: lsd
904 : INTEGER, INTENT(IN) :: nspins, na, nr
905 : REAL(KIND=dp), INTENT(OUT) :: exc
906 : LOGICAL, INTENT(IN), OPTIONAL :: energy_only
907 : REAL(KIND=dp), DIMENSION(3), INTENT(OUT), &
908 : OPTIONAL :: atom_force
909 : REAL(KIND=dp), DIMENSION(3, 3), INTENT(OUT), &
910 : OPTIONAL :: atom_virial
911 :
912 : CHARACTER(len=default_path_length) :: model_path
913 : INTEGER :: ia, idir, ir, native_grid_cuda_device, &
914 : jdir, nflat, omp_max_threads_restore, row, &
915 : selected_cuda_device
916 4 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: atomic_grid_sizes
917 4 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:, :) :: atomic_grid_size_bound_shape
918 : LOGICAL :: need_coord_grad, my_energy_only, native_grid_use_cuda
919 : REAL(KIND=dp) :: tmp
920 4 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: atomic_grid_weights, grid_weights
921 4 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: coarse_0_atomic_coords, density, &
922 4 : grid_coords, kin
923 4 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: grad
924 4 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: atom_coord_grad, density_grad, &
925 4 : grid_coord_grad, kin_grad
926 4 : REAL(KIND=dp), DIMENSION(:, :, :), POINTER :: grad_grad
927 : TYPE(section_vals_type), POINTER :: gauxc_section
928 : TYPE(torch_dict_type) :: inputs
929 : TYPE(torch_tensor_type) :: atomic_grid_size_bound_shape_t, &
930 : atomic_grid_sizes_t, &
931 : atomic_grid_weights_t, &
932 : atom_coord_grad_t, &
933 : coarse_0_atomic_coords_t, density_t, &
934 : density_grad_t, exc_tensor, grad_t, &
935 : grad_grad_t, grid_coord_grad_t, &
936 : grid_coords_t, grid_weights_t, kin_t, &
937 : kin_grad_t
938 :
939 0 : CPASSERT(ASSOCIATED(xc_section))
940 4 : CPASSERT(ASSOCIATED(grid_atom))
941 4 : CPASSERT(ASSOCIATED(rho))
942 4 : CPASSERT(ASSOCIATED(drho))
943 4 : CPASSERT(ASSOCIATED(tau))
944 4 : omp_max_threads_restore = omp_get_max_threads()
945 :
946 4 : my_energy_only = .FALSE.
947 4 : IF (PRESENT(energy_only)) my_energy_only = energy_only
948 4 : need_coord_grad = PRESENT(atom_force) .OR. PRESENT(atom_virial)
949 4 : exc = 0.0_dp
950 4 : IF (PRESENT(atom_force)) atom_force = 0.0_dp
951 4 : IF (PRESENT(atom_virial)) atom_virial = 0.0_dp
952 4 : IF (.NOT. my_energy_only) THEN
953 10208 : vxc = 0.0_dp
954 40208 : vxg = 0.0_dp
955 10208 : vtau = 0.0_dp
956 : END IF
957 :
958 4 : CALL get_skala_model_path(xc_section, model_path)
959 4 : gauxc_section => get_gauxc_section(xc_section)
960 4 : CPASSERT(ASSOCIATED(gauxc_section))
961 4 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_USE_CUDA", l_val=native_grid_use_cuda)
962 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_CUDA_DEVICE", &
963 4 : i_val=native_grid_cuda_device)
964 4 : CALL torch_use_cuda(native_grid_use_cuda)
965 : selected_cuda_device = configure_native_grid_cuda( &
966 4 : native_grid_use_cuda, native_grid_cuda_device, group)
967 4 : CALL ensure_model_loaded(model_path, selected_cuda_device, xc_section)
968 :
969 4 : nflat = na*nr
970 : ALLOCATE (density(nflat, 2), grad(nflat, 3, 2), kin(nflat, 2), &
971 : grid_coords(3, nflat), grid_weights(nflat), &
972 : atomic_grid_weights(nflat), atomic_grid_sizes(1), &
973 52 : coarse_0_atomic_coords(3, 1), atomic_grid_size_bound_shape(0, nflat))
974 4 : density = 0.0_dp
975 4 : grad = 0.0_dp
976 4 : kin = 0.0_dp
977 4 : grid_coords = 0.0_dp
978 4 : grid_weights = 0.0_dp
979 4 : atomic_grid_weights = 0.0_dp
980 4 : atomic_grid_sizes(1) = INT(nflat, KIND=int_8)
981 : atomic_grid_size_bound_shape = 0_int_8
982 16 : coarse_0_atomic_coords(:, 1) = atom_coord
983 :
984 : row = 0
985 204 : DO ir = 1, nr
986 10204 : DO ia = 1, na
987 10000 : row = row + 1
988 : grid_coords(1, row) = atom_coord(1) + grid_atom%rad(ir)* &
989 10000 : grid_atom%sin_pol(ia)*grid_atom%cos_azi(ia)
990 : grid_coords(2, row) = atom_coord(2) + grid_atom%rad(ir)* &
991 10000 : grid_atom%sin_pol(ia)*grid_atom%sin_azi(ia)
992 10000 : grid_coords(3, row) = atom_coord(3) + grid_atom%rad(ir)*grid_atom%cos_pol(ia)
993 10000 : grid_weights(row) = weights(ia, ir)
994 10000 : atomic_grid_weights(row) = weights(ia, ir)
995 10200 : IF (nspins == 1) THEN
996 30000 : density(row, :) = 0.5_dp*rho(ia, ir, 1)
997 40000 : DO idir = 1, 3
998 100000 : grad(row, idir, :) = 0.5_dp*drho(idir, ia, ir, 1)
999 : END DO
1000 30000 : kin(row, :) = 0.5_dp*tau(ia, ir, 1)
1001 : ELSE
1002 0 : density(row, :) = rho(ia, ir, 1:2)
1003 0 : DO idir = 1, 3
1004 0 : grad(row, idir, :) = drho(idir, ia, ir, 1:2)
1005 : END DO
1006 0 : kin(row, :) = tau(ia, ir, 1:2)
1007 : END IF
1008 : END DO
1009 : END DO
1010 :
1011 4 : CALL torch_tensor_from_array(grid_coords_t, grid_coords)
1012 4 : CALL torch_tensor_to_device_leaf(grid_coords_t, need_coord_grad)
1013 4 : CALL torch_tensor_from_array(grid_weights_t, grid_weights)
1014 4 : CALL torch_tensor_to_device_leaf(grid_weights_t, .FALSE.)
1015 4 : CALL torch_tensor_from_array(atomic_grid_weights_t, atomic_grid_weights)
1016 4 : CALL torch_tensor_to_device_leaf(atomic_grid_weights_t, .FALSE.)
1017 4 : CALL torch_tensor_from_array(atomic_grid_sizes_t, atomic_grid_sizes)
1018 4 : CALL torch_tensor_to_device_leaf(atomic_grid_sizes_t, .FALSE.)
1019 : CALL torch_tensor_from_array(atomic_grid_size_bound_shape_t, &
1020 4 : atomic_grid_size_bound_shape)
1021 4 : CALL torch_tensor_to_device_leaf(atomic_grid_size_bound_shape_t, .FALSE.)
1022 4 : CALL torch_tensor_from_array(coarse_0_atomic_coords_t, coarse_0_atomic_coords)
1023 4 : CALL torch_tensor_to_device_leaf(coarse_0_atomic_coords_t, need_coord_grad)
1024 4 : CALL torch_tensor_from_array(density_t, density)
1025 4 : CALL torch_tensor_to_device_leaf(density_t,.NOT. my_energy_only)
1026 4 : CALL torch_tensor_from_array(grad_t, grad)
1027 4 : CALL torch_tensor_to_device_leaf(grad_t,.NOT. my_energy_only)
1028 4 : CALL torch_tensor_from_array(kin_t, kin)
1029 4 : CALL torch_tensor_to_device_leaf(kin_t,.NOT. my_energy_only)
1030 :
1031 4 : CALL torch_dict_create(inputs)
1032 4 : CALL torch_dict_insert(inputs, "grid_coords", grid_coords_t)
1033 4 : CALL torch_dict_insert(inputs, "grid_weights", grid_weights_t)
1034 4 : CALL torch_dict_insert(inputs, "atomic_grid_weights", atomic_grid_weights_t)
1035 4 : CALL torch_dict_insert(inputs, "atomic_grid_sizes", atomic_grid_sizes_t)
1036 : CALL torch_dict_insert(inputs, "atomic_grid_size_bound_shape", &
1037 4 : atomic_grid_size_bound_shape_t)
1038 4 : CALL torch_dict_insert(inputs, "density", density_t)
1039 4 : CALL torch_dict_insert(inputs, "grad", grad_t)
1040 4 : CALL torch_dict_insert(inputs, "kin", kin_t)
1041 4 : CALL torch_dict_insert(inputs, "coarse_0_atomic_coords", coarse_0_atomic_coords_t)
1042 :
1043 4 : CALL skala_torch_model_get_exc(cached_model, inputs, grid_weights_t, exc_tensor, exc)
1044 :
1045 4 : IF (.NOT. my_energy_only) THEN
1046 4 : NULLIFY (atom_coord_grad, density_grad, grad_grad, grid_coord_grad, kin_grad)
1047 4 : CALL torch_tensor_backward_scalar(exc_tensor)
1048 4 : IF (need_coord_grad) THEN
1049 4 : CALL torch_tensor_grad(grid_coords_t, grid_coord_grad_t)
1050 4 : CALL torch_tensor_grad(coarse_0_atomic_coords_t, atom_coord_grad_t)
1051 4 : CALL torch_tensor_data_ptr(grid_coord_grad_t, grid_coord_grad)
1052 4 : CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
1053 4 : IF (PRESENT(atom_force)) THEN
1054 16 : atom_force(:) = atom_coord_grad(:, 1)
1055 10004 : DO row = 1, nflat
1056 40004 : atom_force(:) = atom_force(:) + grid_coord_grad(:, row)
1057 : END DO
1058 : END IF
1059 4 : IF (PRESENT(atom_virial)) THEN
1060 10004 : DO row = 1, nflat
1061 40004 : DO idir = 1, 3
1062 130000 : DO jdir = 1, 3
1063 90000 : tmp = grid_coord_grad(idir, row)*coarse_0_atomic_coords(jdir, 1)
1064 120000 : atom_virial(idir, jdir) = atom_virial(idir, jdir) + tmp
1065 : END DO
1066 : END DO
1067 : END DO
1068 16 : DO idir = 1, 3
1069 52 : DO jdir = 1, 3
1070 36 : tmp = atom_coord_grad(idir, 1)*coarse_0_atomic_coords(jdir, 1)
1071 48 : atom_virial(idir, jdir) = atom_virial(idir, jdir) + tmp
1072 : END DO
1073 : END DO
1074 : END IF
1075 : END IF
1076 : CALL torch_tensor_grad_batch3(density_t, grad_t, kin_t, density_grad_t, &
1077 4 : grad_grad_t, kin_grad_t)
1078 4 : CALL torch_tensor_data_ptr(density_grad_t, density_grad)
1079 4 : CALL torch_tensor_data_ptr(grad_grad_t, grad_grad)
1080 4 : CALL torch_tensor_data_ptr(kin_grad_t, kin_grad)
1081 :
1082 4 : row = 0
1083 204 : DO ir = 1, nr
1084 10204 : DO ia = 1, na
1085 10000 : row = row + 1
1086 10200 : IF (lsd) THEN
1087 0 : vxc(ia, ir, 1:2) = density_grad(row, 1:2)
1088 0 : DO idir = 1, 3
1089 0 : vxg(idir, ia, ir, 1:2) = grad_grad(row, idir, 1:2)
1090 : END DO
1091 0 : vtau(ia, ir, 1:2) = kin_grad(row, 1:2)
1092 : ELSE
1093 10000 : vxc(ia, ir, 1) = 0.5_dp*(density_grad(row, 1) + density_grad(row, 2))
1094 40000 : DO idir = 1, 3
1095 : vxg(idir, ia, ir, 1) = &
1096 40000 : 0.5_dp*(grad_grad(row, idir, 1) + grad_grad(row, idir, 2))
1097 : END DO
1098 10000 : vtau(ia, ir, 1) = 0.5_dp*(kin_grad(row, 1) + kin_grad(row, 2))
1099 : END IF
1100 : END DO
1101 : END DO
1102 :
1103 4 : CALL torch_tensor_release(density_grad_t)
1104 4 : CALL torch_tensor_release(grad_grad_t)
1105 4 : CALL torch_tensor_release(kin_grad_t)
1106 4 : IF (need_coord_grad) THEN
1107 4 : CALL torch_tensor_release(grid_coord_grad_t)
1108 4 : CALL torch_tensor_release(atom_coord_grad_t)
1109 : END IF
1110 : END IF
1111 :
1112 4 : CALL torch_tensor_release(exc_tensor)
1113 4 : CALL torch_tensor_release(density_t)
1114 4 : CALL torch_tensor_release(grad_t)
1115 4 : CALL torch_tensor_release(kin_t)
1116 4 : CALL torch_tensor_release(grid_coords_t)
1117 4 : CALL torch_tensor_release(grid_weights_t)
1118 4 : CALL torch_tensor_release(atomic_grid_weights_t)
1119 4 : CALL torch_tensor_release(atomic_grid_sizes_t)
1120 4 : CALL torch_tensor_release(atomic_grid_size_bound_shape_t)
1121 4 : CALL torch_tensor_release(coarse_0_atomic_coords_t)
1122 4 : CALL torch_dict_release(inputs)
1123 0 : DEALLOCATE (atomic_grid_size_bound_shape, atomic_grid_sizes, atomic_grid_weights, &
1124 4 : coarse_0_atomic_coords, density, grad, grid_coords, grid_weights, kin)
1125 4 : CALL torch_use_cuda(.TRUE.)
1126 4 : CALL omp_set_num_threads(omp_max_threads_restore)
1127 :
1128 12 : END SUBROUTINE skala_gapw_atom_vxc_of_r
1129 :
1130 : ! **************************************************************************************************
1131 : !> \brief Evaluate a rank-local set of complete atom blocks and sum their SKALA energies.
1132 : !> \param xc_section ...
1133 : !> \param group ...
1134 : !> \param density ...
1135 : !> \param grad ...
1136 : !> \param kin ...
1137 : !> \param grid_coords ...
1138 : !> \param grid_weights ...
1139 : !> \param atomic_grid_weights ...
1140 : !> \param atomic_grid_sizes ...
1141 : !> \param atomic_coords ...
1142 : !> \param exc ...
1143 : !> \param density_grad_out ...
1144 : !> \param grad_grad_out ...
1145 : !> \param kin_grad_out ...
1146 : !> \param grid_coord_grad_out ...
1147 : !> \param grid_weight_grad_out ...
1148 : !> \param atomic_grid_weight_grad_out ...
1149 : !> \param atom_coord_grad_out ...
1150 : ! **************************************************************************************************
1151 252 : SUBROUTINE skala_gapw_atom_composite_energy(xc_section, group, density, grad, kin, &
1152 : grid_coords, grid_weights, atomic_grid_weights, &
1153 : atomic_grid_sizes, atomic_coords, exc, &
1154 : density_grad_out, grad_grad_out, kin_grad_out, &
1155 : grid_coord_grad_out, grid_weight_grad_out, &
1156 : atomic_grid_weight_grad_out, atom_coord_grad_out)
1157 : TYPE(section_vals_type), POINTER :: xc_section
1158 :
1159 : CLASS(mp_comm_type), INTENT(IN) :: group
1160 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), INTENT(IN) :: density, grid_coords, &
1161 : atomic_coords, kin
1162 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :), INTENT(IN) :: grad
1163 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:), INTENT(IN) :: grid_weights, atomic_grid_weights
1164 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:), INTENT(IN) :: atomic_grid_sizes
1165 : REAL(KIND=dp), INTENT(OUT) :: exc
1166 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), INTENT(OUT), OPTIONAL :: density_grad_out, &
1167 : kin_grad_out, &
1168 : grid_coord_grad_out, &
1169 : atom_coord_grad_out
1170 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :), INTENT(OUT), OPTIONAL :: grad_grad_out
1171 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:), INTENT(OUT), OPTIONAL :: grid_weight_grad_out, &
1172 : atomic_grid_weight_grad_out
1173 :
1174 : CHARACTER(len=default_path_length) :: model_path
1175 : INTEGER :: atom_begin, atom_count, chunk_max_rows, chunk_row_count, chunk_row_start, &
1176 : ichunk, local_natom, local_nrow, max_grid_size, native_grid_cuda_device, &
1177 : nchunks, omp_max_threads_restore, phase_handle, selected_cuda_device
1178 252 : INTEGER, ALLOCATABLE, DIMENSION(:) :: chunk_atom_begin, chunk_atom_count, &
1179 252 : chunk_row_begin, chunk_row_counts
1180 252 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:, :) :: atomic_grid_size_bound_shape
1181 : LOGICAL :: active_rank, native_grid_atom_chunks, native_grid_use_cuda, &
1182 : need_coordinate_derivatives, need_derivatives
1183 : REAL(KIND=dp) :: chunk_exc
1184 252 : REAL(KIND=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
1185 252 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: atom_coord_grad, density_grad, &
1186 252 : grid_coord_grad, kin_grad
1187 252 : REAL(KIND=dp), DIMENSION(:, :, :), POINTER :: grad_grad
1188 : TYPE(section_vals_type), POINTER :: gauxc_section
1189 : TYPE(torch_dict_type) :: inputs
1190 : TYPE(torch_tensor_type) :: atomic_grid_size_bound_shape_t, &
1191 : atomic_grid_sizes_t, &
1192 : atomic_grid_sizes_chunk_t, &
1193 : atomic_grid_weight_grad_t, &
1194 : atomic_grid_weights_chunk_t, &
1195 : atom_coord_grad_t, atomic_coords_chunk_t, &
1196 : density_chunk_t, density_grad_t, &
1197 : exc_tensor, grad_chunk_t, grad_grad_t, &
1198 : grid_coord_grad_t, grid_coords_chunk_t, &
1199 : grid_weight_grad_t, grid_weights_chunk_t, &
1200 : kin_chunk_t, kin_grad_t
1201 : TYPE(torch_tensor_type), TARGET :: atomic_coords_local_t, &
1202 : atomic_grid_weights_local_t, &
1203 : density_local_t, grad_local_t, &
1204 : grid_coords_local_t, grid_weights_local_t, &
1205 : kin_local_t
1206 : TYPE(torch_tensor_type), POINTER :: atomic_coords_t, atomic_grid_weights_t, &
1207 : density_t, grad_t, grid_coords_t, &
1208 : grid_weights_t, kin_t
1209 :
1210 0 : CPASSERT(ASSOCIATED(xc_section))
1211 252 : CPASSERT(SIZE(density, 1) == SIZE(grid_weights))
1212 252 : CPASSERT(SIZE(density, 1) == SIZE(atomic_grid_weights))
1213 252 : CPASSERT(SIZE(density, 1) == SIZE(grid_coords, 2))
1214 252 : CPASSERT(SIZE(density, 1) == SIZE(grad, 1))
1215 252 : CPASSERT(SIZE(density, 1) == SIZE(kin, 1))
1216 252 : CPASSERT(SIZE(density, 2) == 2)
1217 252 : CPASSERT(SIZE(grad, 2) == 3)
1218 252 : CPASSERT(SIZE(grad, 3) == 2)
1219 252 : CPASSERT(SIZE(kin, 2) == 2)
1220 252 : CPASSERT(SIZE(atomic_grid_sizes) == SIZE(atomic_coords, 2))
1221 505 : CPASSERT(SUM(atomic_grid_sizes) == INT(SIZE(density, 1), KIND=int_8))
1222 252 : need_derivatives = PRESENT(density_grad_out)
1223 252 : CPASSERT(PRESENT(grad_grad_out) .EQV. need_derivatives)
1224 252 : CPASSERT(PRESENT(kin_grad_out) .EQV. need_derivatives)
1225 252 : need_coordinate_derivatives = PRESENT(grid_coord_grad_out)
1226 252 : CPASSERT(PRESENT(grid_weight_grad_out) .EQV. need_coordinate_derivatives)
1227 252 : CPASSERT(PRESENT(atomic_grid_weight_grad_out) .EQV. need_coordinate_derivatives)
1228 252 : CPASSERT(PRESENT(atom_coord_grad_out) .EQV. need_coordinate_derivatives)
1229 252 : CPASSERT((.NOT. need_coordinate_derivatives) .OR. need_derivatives)
1230 252 : omp_max_threads_restore = omp_get_max_threads()
1231 :
1232 252 : local_nrow = SIZE(density, 1)
1233 252 : local_natom = SIZE(atomic_grid_sizes)
1234 252 : active_rank = local_natom > 0
1235 252 : CPASSERT(active_rank .EQV. (local_nrow > 0))
1236 :
1237 252 : CALL get_skala_model_path(xc_section, model_path)
1238 252 : gauxc_section => get_gauxc_section(xc_section)
1239 252 : CPASSERT(ASSOCIATED(gauxc_section))
1240 252 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_USE_CUDA", l_val=native_grid_use_cuda)
1241 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_CUDA_DEVICE", &
1242 252 : i_val=native_grid_cuda_device)
1243 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNKS", &
1244 252 : l_val=native_grid_atom_chunks)
1245 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNK_MAX_ROWS", &
1246 252 : i_val=chunk_max_rows)
1247 252 : IF (chunk_max_rows < -1) THEN
1248 : CALL cp_abort(__LOCATION__, &
1249 0 : "GAUXC%NATIVE_GRID_ATOM_CHUNK_MAX_ROWS must be -1, zero, or positive.")
1250 : END IF
1251 252 : IF (.NOT. native_grid_atom_chunks) chunk_max_rows = 0
1252 252 : IF (chunk_max_rows == -1) THEN
1253 242 : chunk_max_rows = auto_atom_composite_chunk_max_rows(atomic_grid_sizes, group)
1254 : END IF
1255 252 : CALL torch_use_cuda(native_grid_use_cuda)
1256 : selected_cuda_device = configure_native_grid_cuda( &
1257 252 : native_grid_use_cuda, native_grid_cuda_device, group)
1258 252 : exc = 0.0_dp
1259 252 : IF (active_rank) THEN
1260 230 : CALL timeset("skala_atom_tensor_setup", phase_handle)
1261 230 : CALL ensure_model_loaded(model_path, selected_cuda_device, xc_section)
1262 : CALL atom_composite_chunk_layout(atomic_grid_sizes, chunk_max_rows, &
1263 : chunk_atom_begin, chunk_atom_count, &
1264 230 : chunk_row_begin, chunk_row_counts)
1265 230 : nchunks = SIZE(chunk_atom_begin)
1266 230 : CPASSERT(nchunks > 0)
1267 :
1268 230 : IF (native_grid_use_cuda) THEN
1269 0 : grid_coords_t => cached_atom_cuda_tensors%grid_coords_t
1270 0 : grid_weights_t => cached_atom_cuda_tensors%grid_weights_t
1271 0 : atomic_grid_weights_t => cached_atom_cuda_tensors%atomic_grid_weights_t
1272 0 : atomic_coords_t => cached_atom_cuda_tensors%atomic_coords_t
1273 0 : density_t => cached_atom_cuda_tensors%density_t
1274 0 : grad_t => cached_atom_cuda_tensors%grad_t
1275 0 : kin_t => cached_atom_cuda_tensors%kin_t
1276 : CALL torch_tensor_reset_from_array( &
1277 0 : grid_coords_t, grid_coords, need_coordinate_derivatives)
1278 : CALL torch_tensor_reset_from_array( &
1279 0 : grid_weights_t, grid_weights, need_coordinate_derivatives)
1280 : CALL torch_tensor_reset_from_array( &
1281 0 : atomic_grid_weights_t, atomic_grid_weights, need_coordinate_derivatives)
1282 : CALL torch_tensor_reset_from_array( &
1283 0 : atomic_coords_t, atomic_coords, need_coordinate_derivatives)
1284 0 : CALL torch_tensor_reset_from_array(density_t, density, need_derivatives)
1285 0 : CALL torch_tensor_reset_from_array(grad_t, grad, need_derivatives)
1286 0 : CALL torch_tensor_reset_from_array(kin_t, kin, need_derivatives)
1287 : ELSE
1288 230 : grid_coords_t => grid_coords_local_t
1289 230 : grid_weights_t => grid_weights_local_t
1290 230 : atomic_grid_weights_t => atomic_grid_weights_local_t
1291 230 : atomic_coords_t => atomic_coords_local_t
1292 230 : density_t => density_local_t
1293 230 : grad_t => grad_local_t
1294 230 : kin_t => kin_local_t
1295 230 : CALL torch_tensor_from_array(grid_coords_t, grid_coords)
1296 230 : CALL torch_tensor_to_device_leaf(grid_coords_t, need_coordinate_derivatives)
1297 230 : CALL torch_tensor_from_array(grid_weights_t, grid_weights)
1298 230 : CALL torch_tensor_to_device_leaf(grid_weights_t, need_coordinate_derivatives)
1299 230 : CALL torch_tensor_from_array(atomic_grid_weights_t, atomic_grid_weights)
1300 230 : CALL torch_tensor_to_device_leaf(atomic_grid_weights_t, need_coordinate_derivatives)
1301 230 : CALL torch_tensor_from_array(atomic_coords_t, atomic_coords)
1302 230 : CALL torch_tensor_to_device_leaf(atomic_coords_t, need_coordinate_derivatives)
1303 230 : CALL torch_tensor_from_array(density_t, density)
1304 230 : CALL torch_tensor_to_device_leaf(density_t, need_derivatives)
1305 230 : CALL torch_tensor_from_array(grad_t, grad)
1306 230 : CALL torch_tensor_to_device_leaf(grad_t, need_derivatives)
1307 230 : CALL torch_tensor_from_array(kin_t, kin)
1308 230 : CALL torch_tensor_to_device_leaf(kin_t, need_derivatives)
1309 : END IF
1310 230 : CALL torch_tensor_from_array(atomic_grid_sizes_t, atomic_grid_sizes)
1311 230 : CALL torch_tensor_to_device_leaf(atomic_grid_sizes_t, .FALSE.)
1312 230 : CALL timestop(phase_handle)
1313 :
1314 460 : DO ichunk = 1, nchunks
1315 230 : atom_begin = chunk_atom_begin(ichunk)
1316 230 : atom_count = chunk_atom_count(ichunk)
1317 230 : chunk_row_start = chunk_row_begin(ichunk)
1318 230 : chunk_row_count = chunk_row_counts(ichunk)
1319 : max_grid_size = INT(MAXVAL(atomic_grid_sizes( &
1320 483 : atom_begin:atom_begin + atom_count - 1)))
1321 460 : ALLOCATE (atomic_grid_size_bound_shape(0, max_grid_size))
1322 : atomic_grid_size_bound_shape = 0_int_8
1323 :
1324 : CALL torch_tensor_narrow(grid_coords_t, 0, chunk_row_start - 1, &
1325 230 : chunk_row_count, grid_coords_chunk_t)
1326 : CALL torch_tensor_narrow(grid_weights_t, 0, chunk_row_start - 1, &
1327 230 : chunk_row_count, grid_weights_chunk_t)
1328 : CALL torch_tensor_narrow(atomic_grid_weights_t, 0, chunk_row_start - 1, &
1329 230 : chunk_row_count, atomic_grid_weights_chunk_t)
1330 : CALL torch_tensor_narrow(atomic_grid_sizes_t, 0, atom_begin - 1, &
1331 230 : atom_count, atomic_grid_sizes_chunk_t)
1332 : CALL torch_tensor_narrow(atomic_coords_t, 0, atom_begin - 1, &
1333 230 : atom_count, atomic_coords_chunk_t)
1334 : CALL torch_tensor_narrow(density_t, 1, chunk_row_start - 1, &
1335 230 : chunk_row_count, density_chunk_t)
1336 : CALL torch_tensor_narrow(grad_t, 2, chunk_row_start - 1, &
1337 230 : chunk_row_count, grad_chunk_t)
1338 : CALL torch_tensor_narrow(kin_t, 1, chunk_row_start - 1, &
1339 230 : chunk_row_count, kin_chunk_t)
1340 : CALL torch_tensor_from_array(atomic_grid_size_bound_shape_t, &
1341 230 : atomic_grid_size_bound_shape)
1342 230 : CALL torch_tensor_to_device_leaf(atomic_grid_size_bound_shape_t, .FALSE.)
1343 :
1344 230 : CALL torch_dict_create(inputs)
1345 230 : CALL torch_dict_insert(inputs, "grid_coords", grid_coords_chunk_t)
1346 230 : CALL torch_dict_insert(inputs, "grid_weights", grid_weights_chunk_t)
1347 230 : CALL torch_dict_insert(inputs, "atomic_grid_weights", atomic_grid_weights_chunk_t)
1348 230 : CALL torch_dict_insert(inputs, "atomic_grid_sizes", atomic_grid_sizes_chunk_t)
1349 : CALL torch_dict_insert(inputs, "atomic_grid_size_bound_shape", &
1350 230 : atomic_grid_size_bound_shape_t)
1351 230 : CALL torch_dict_insert(inputs, "density", density_chunk_t)
1352 230 : CALL torch_dict_insert(inputs, "grad", grad_chunk_t)
1353 230 : CALL torch_dict_insert(inputs, "kin", kin_chunk_t)
1354 230 : CALL torch_dict_insert(inputs, "coarse_0_atomic_coords", atomic_coords_chunk_t)
1355 :
1356 230 : CALL timeset("skala_atom_torch_forward", phase_handle)
1357 : CALL skala_torch_model_get_exc(cached_model, inputs, grid_weights_chunk_t, &
1358 230 : exc_tensor, chunk_exc)
1359 230 : CALL timestop(phase_handle)
1360 230 : exc = exc + chunk_exc
1361 230 : IF (need_derivatives) THEN
1362 230 : CALL timeset("skala_atom_torch_backward", phase_handle)
1363 230 : CALL torch_tensor_backward_scalar(exc_tensor)
1364 230 : CALL timestop(phase_handle)
1365 : END IF
1366 :
1367 230 : CALL torch_tensor_release(exc_tensor)
1368 230 : CALL torch_dict_release(inputs)
1369 230 : CALL torch_tensor_release(grid_coords_chunk_t)
1370 230 : CALL torch_tensor_release(grid_weights_chunk_t)
1371 230 : CALL torch_tensor_release(atomic_grid_weights_chunk_t)
1372 230 : CALL torch_tensor_release(atomic_grid_sizes_chunk_t)
1373 230 : CALL torch_tensor_release(atomic_coords_chunk_t)
1374 230 : CALL torch_tensor_release(density_chunk_t)
1375 230 : CALL torch_tensor_release(grad_chunk_t)
1376 230 : CALL torch_tensor_release(kin_chunk_t)
1377 230 : CALL torch_tensor_release(atomic_grid_size_bound_shape_t)
1378 690 : DEALLOCATE (atomic_grid_size_bound_shape)
1379 : END DO
1380 : END IF
1381 252 : CALL group%sum(exc)
1382 :
1383 252 : IF (need_derivatives) THEN
1384 0 : ALLOCATE (density_grad_out(local_nrow, 2), grad_grad_out(local_nrow, 3, 2), &
1385 1698 : kin_grad_out(local_nrow, 2))
1386 252 : density_grad_out = 0.0_dp
1387 252 : grad_grad_out = 0.0_dp
1388 252 : kin_grad_out = 0.0_dp
1389 252 : IF (active_rank) THEN
1390 230 : NULLIFY (density_grad, grad_grad, kin_grad)
1391 230 : CALL timeset("skala_atom_grad_fetch", phase_handle)
1392 : CALL torch_tensor_grad_batch3(density_t, grad_t, kin_t, density_grad_t, &
1393 230 : grad_grad_t, kin_grad_t)
1394 230 : CALL torch_tensor_data_ptr(density_grad_t, density_grad)
1395 230 : CALL torch_tensor_data_ptr(grad_grad_t, grad_grad)
1396 230 : CALL torch_tensor_data_ptr(kin_grad_t, kin_grad)
1397 857210 : density_grad_out(:, :) = density_grad
1398 2571630 : grad_grad_out(:, :, :) = grad_grad
1399 857210 : kin_grad_out(:, :) = kin_grad
1400 230 : CALL timestop(phase_handle)
1401 : END IF
1402 :
1403 252 : IF (need_coordinate_derivatives) THEN
1404 58 : IF (active_rank) CALL timeset("skala_atom_coord_grad_fetch", phase_handle)
1405 0 : ALLOCATE (grid_coord_grad_out(3, local_nrow), grid_weight_grad_out(local_nrow), &
1406 420 : atomic_grid_weight_grad_out(local_nrow), atom_coord_grad_out(3, local_natom))
1407 58 : grid_coord_grad_out = 0.0_dp
1408 58 : grid_weight_grad_out = 0.0_dp
1409 58 : atomic_grid_weight_grad_out = 0.0_dp
1410 58 : atom_coord_grad_out = 0.0_dp
1411 58 : IF (active_rank) THEN
1412 47 : NULLIFY (atomic_grid_weight_grad, atom_coord_grad, grid_coord_grad, grid_weight_grad)
1413 47 : CALL torch_tensor_grad(grid_coords_t, grid_coord_grad_t)
1414 47 : CALL torch_tensor_grad(grid_weights_t, grid_weight_grad_t)
1415 47 : CALL torch_tensor_grad(atomic_grid_weights_t, atomic_grid_weight_grad_t)
1416 47 : CALL torch_tensor_grad(atomic_coords_t, atom_coord_grad_t)
1417 47 : CALL torch_tensor_data_ptr(grid_coord_grad_t, grid_coord_grad)
1418 47 : CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
1419 47 : CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
1420 47 : CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
1421 420207 : grid_coord_grad_out(:, :) = grid_coord_grad
1422 105087 : grid_weight_grad_out(:) = grid_weight_grad
1423 105087 : atomic_grid_weight_grad_out(:) = atomic_grid_weight_grad
1424 279 : atom_coord_grad_out(:, :) = atom_coord_grad
1425 47 : CALL torch_tensor_release(grid_coord_grad_t)
1426 47 : CALL torch_tensor_release(grid_weight_grad_t)
1427 47 : CALL torch_tensor_release(atomic_grid_weight_grad_t)
1428 47 : CALL torch_tensor_release(atom_coord_grad_t)
1429 47 : CALL timestop(phase_handle)
1430 : END IF
1431 : END IF
1432 241 : IF (active_rank) THEN
1433 230 : CALL torch_tensor_release(density_grad_t)
1434 230 : CALL torch_tensor_release(grad_grad_t)
1435 230 : CALL torch_tensor_release(kin_grad_t)
1436 : END IF
1437 : END IF
1438 :
1439 241 : IF (active_rank) THEN
1440 230 : CALL timeset("skala_atom_tensor_release", phase_handle)
1441 230 : IF (.NOT. native_grid_use_cuda) THEN
1442 230 : CALL torch_tensor_release(density_t)
1443 230 : CALL torch_tensor_release(grad_t)
1444 230 : CALL torch_tensor_release(kin_t)
1445 230 : CALL torch_tensor_release(grid_coords_t)
1446 230 : CALL torch_tensor_release(grid_weights_t)
1447 230 : CALL torch_tensor_release(atomic_grid_weights_t)
1448 230 : CALL torch_tensor_release(atomic_coords_t)
1449 : END IF
1450 230 : CALL torch_tensor_release(atomic_grid_sizes_t)
1451 230 : DEALLOCATE (chunk_atom_begin, chunk_atom_count, chunk_row_begin, chunk_row_counts)
1452 230 : CALL timestop(phase_handle)
1453 : END IF
1454 252 : CALL torch_use_cuda(.TRUE.)
1455 252 : CALL omp_set_num_threads(omp_max_threads_restore)
1456 :
1457 504 : END SUBROUTINE skala_gapw_atom_composite_energy
1458 :
1459 : ! **************************************************************************************************
1460 : !> \brief Select an automatic row cap for atom-composite Torch calls.
1461 : !> \param atomic_grid_sizes ...
1462 : !> \param group ...
1463 : !> \return ...
1464 : ! **************************************************************************************************
1465 242 : FUNCTION auto_atom_composite_chunk_max_rows(atomic_grid_sizes, group) RESULT(max_rows)
1466 : INTEGER(KIND=int_8), DIMENSION(:), INTENT(IN) :: atomic_grid_sizes
1467 :
1468 : CLASS(mp_comm_type), INTENT(IN) :: group
1469 :
1470 : INTEGER :: local_max_atom_rows, local_natom, &
1471 : local_padded_rows, max_rows
1472 :
1473 242 : local_padded_rows = 0
1474 242 : local_natom = SIZE(atomic_grid_sizes)
1475 242 : IF (local_natom > 0) THEN
1476 463 : local_max_atom_rows = INT(MAXVAL(atomic_grid_sizes))
1477 220 : IF (local_natom > atom_chunk_auto_max_rows .OR. &
1478 : local_max_atom_rows > atom_chunk_auto_max_rows/local_natom) THEN
1479 0 : local_padded_rows = atom_chunk_auto_max_rows + 1
1480 : ELSE
1481 220 : local_padded_rows = local_natom*local_max_atom_rows
1482 : END IF
1483 : END IF
1484 242 : CALL group%max(local_padded_rows)
1485 242 : IF (local_padded_rows <= atom_chunk_auto_max_rows) THEN
1486 242 : max_rows = 0
1487 242 : RETURN
1488 : END IF
1489 242 : max_rows = atom_chunk_auto_max_rows
1490 :
1491 : END FUNCTION auto_atom_composite_chunk_max_rows
1492 :
1493 : ! **************************************************************************************************
1494 : !> \brief Split contiguous atom-composite rows without dividing an atomic block.
1495 : !> \param atomic_grid_sizes ...
1496 : !> \param max_rows ...
1497 : !> \param atom_begin ...
1498 : !> \param atom_count ...
1499 : !> \param row_begin ...
1500 : !> \param row_count ...
1501 : ! **************************************************************************************************
1502 230 : SUBROUTINE atom_composite_chunk_layout(atomic_grid_sizes, max_rows, atom_begin, atom_count, &
1503 : row_begin, row_count)
1504 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:), &
1505 : INTENT(IN) :: atomic_grid_sizes
1506 : INTEGER, INTENT(IN) :: max_rows
1507 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: atom_begin, atom_count, row_begin, &
1508 : row_count
1509 :
1510 : INTEGER :: atom_begin_tmp, atom_count_tmp, atom_rows, chunk_atoms, chunk_max_atom_rows, &
1511 : iatom, insert_at, natom, nchunks, padded_rows_tmp, row_begin_tmp, row_count_tmp, rows, &
1512 : subchunk
1513 230 : INTEGER, ALLOCATABLE, DIMENSION(:) :: padded_rows
1514 :
1515 230 : natom = SIZE(atomic_grid_sizes)
1516 230 : CPASSERT(natom > 0)
1517 483 : CPASSERT(ALL(atomic_grid_sizes > 0_int_8))
1518 230 : IF (max_rows <= 0) THEN
1519 220 : ALLOCATE (atom_begin(1), atom_count(1), row_begin(1), row_count(1))
1520 220 : atom_begin(1) = 1
1521 220 : atom_count(1) = natom
1522 220 : row_begin(1) = 1
1523 463 : row_count(1) = INT(SUM(atomic_grid_sizes))
1524 : RETURN
1525 : END IF
1526 :
1527 : nchunks = 1
1528 : chunk_atoms = 0
1529 : chunk_max_atom_rows = 0
1530 20 : DO iatom = 1, natom
1531 10 : atom_rows = INT(atomic_grid_sizes(iatom))
1532 10 : IF (chunk_atoms > 0 .AND. &
1533 : MAX(chunk_max_atom_rows, atom_rows) > max_rows/(chunk_atoms + 1)) THEN
1534 0 : nchunks = nchunks + 1
1535 0 : chunk_atoms = 0
1536 0 : chunk_max_atom_rows = 0
1537 : END IF
1538 10 : chunk_atoms = chunk_atoms + 1
1539 20 : chunk_max_atom_rows = MAX(chunk_max_atom_rows, atom_rows)
1540 : END DO
1541 60 : ALLOCATE (atom_begin(nchunks), atom_count(nchunks), row_begin(nchunks), row_count(nchunks))
1542 :
1543 10 : subchunk = 1
1544 10 : atom_begin(subchunk) = 1
1545 10 : row_begin(subchunk) = 1
1546 10 : chunk_atoms = 0
1547 10 : chunk_max_atom_rows = 0
1548 10 : rows = 0
1549 20 : DO iatom = 1, natom
1550 10 : atom_rows = INT(atomic_grid_sizes(iatom))
1551 10 : IF (chunk_atoms > 0 .AND. &
1552 : MAX(chunk_max_atom_rows, atom_rows) > max_rows/(chunk_atoms + 1)) THEN
1553 0 : atom_count(subchunk) = iatom - atom_begin(subchunk)
1554 0 : row_count(subchunk) = rows
1555 0 : subchunk = subchunk + 1
1556 0 : atom_begin(subchunk) = iatom
1557 0 : row_begin(subchunk) = row_begin(subchunk - 1) + row_count(subchunk - 1)
1558 0 : chunk_atoms = 0
1559 0 : chunk_max_atom_rows = 0
1560 0 : rows = 0
1561 : END IF
1562 10 : chunk_atoms = chunk_atoms + 1
1563 10 : chunk_max_atom_rows = MAX(chunk_max_atom_rows, atom_rows)
1564 20 : rows = rows + atom_rows
1565 : END DO
1566 10 : atom_count(subchunk) = natom - atom_begin(subchunk) + 1
1567 10 : row_count(subchunk) = rows
1568 10 : CPASSERT(subchunk == nchunks)
1569 :
1570 20 : ALLOCATE (padded_rows(nchunks))
1571 20 : DO subchunk = 1, nchunks
1572 10 : iatom = atom_begin(subchunk) + atom_count(subchunk) - 1
1573 : padded_rows(subchunk) = atom_count(subchunk)* &
1574 30 : INT(MAXVAL(atomic_grid_sizes(atom_begin(subchunk):iatom)))
1575 : END DO
1576 :
1577 : ! Largest blocks run first so CUDA allocations can be reused by all following calls.
1578 10 : DO subchunk = 2, nchunks
1579 0 : padded_rows_tmp = padded_rows(subchunk)
1580 0 : atom_begin_tmp = atom_begin(subchunk)
1581 0 : atom_count_tmp = atom_count(subchunk)
1582 0 : row_begin_tmp = row_begin(subchunk)
1583 0 : row_count_tmp = row_count(subchunk)
1584 0 : insert_at = subchunk
1585 0 : DO WHILE (insert_at > 1 .AND. padded_rows(insert_at - 1) < padded_rows_tmp)
1586 0 : padded_rows(insert_at) = padded_rows(insert_at - 1)
1587 0 : atom_begin(insert_at) = atom_begin(insert_at - 1)
1588 0 : atom_count(insert_at) = atom_count(insert_at - 1)
1589 0 : row_begin(insert_at) = row_begin(insert_at - 1)
1590 0 : row_count(insert_at) = row_count(insert_at - 1)
1591 0 : insert_at = insert_at - 1
1592 : END DO
1593 0 : padded_rows(insert_at) = padded_rows_tmp
1594 0 : atom_begin(insert_at) = atom_begin_tmp
1595 0 : atom_count(insert_at) = atom_count_tmp
1596 0 : row_begin(insert_at) = row_begin_tmp
1597 10 : row_count(insert_at) = row_count_tmp
1598 : END DO
1599 10 : DEALLOCATE (padded_rows)
1600 :
1601 : END SUBROUTINE atom_composite_chunk_layout
1602 :
1603 : ! **************************************************************************************************
1604 : !> \brief Add the explicit SKALA derivative with respect to atom-center coordinates.
1605 : !> \param atom_force ...
1606 : !> \param features ...
1607 : !> \param atom_coord_grad_t ...
1608 : !> \param root_rank ...
1609 : ! **************************************************************************************************
1610 8 : SUBROUTINE add_explicit_coordinate_force(atom_force, features, atom_coord_grad_t, root_rank)
1611 : REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT) :: atom_force
1612 : TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1613 : TYPE(torch_tensor_type), INTENT(INOUT) :: atom_coord_grad_t
1614 : LOGICAL, INTENT(IN) :: root_rank
1615 :
1616 8 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: atom_coord_grad
1617 :
1618 8 : NULLIFY (atom_coord_grad)
1619 8 : CALL torch_tensor_grad(features%coarse_0_atomic_coords_t, atom_coord_grad_t)
1620 8 : IF (root_rank) THEN
1621 4 : CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
1622 4 : CPASSERT(SIZE(atom_force, 1) == SIZE(atom_coord_grad, 1))
1623 4 : CPASSERT(SIZE(atom_force, 2) == SIZE(atom_coord_grad, 2))
1624 36 : atom_force(:, :) = atom_force(:, :) + atom_coord_grad(:, :)
1625 : END IF
1626 :
1627 8 : END SUBROUTINE add_explicit_coordinate_force
1628 :
1629 : ! **************************************************************************************************
1630 : !> \brief Add the force from SMOOTH native-grid atom partition weights.
1631 : !> \param atom_force ...
1632 : !> \param features ...
1633 : !> \param particle_set ...
1634 : !> \param cell ...
1635 : !> \param rho_r ...
1636 : !> \param grid_weight_grad_t ...
1637 : !> \param atomic_grid_weight_grad_t ...
1638 : ! **************************************************************************************************
1639 8 : SUBROUTINE add_smooth_partition_force(atom_force, features, particle_set, cell, rho_r, &
1640 : grid_weight_grad_t, atomic_grid_weight_grad_t)
1641 : REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT) :: atom_force
1642 : TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1643 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
1644 : TYPE(cell_type), POINTER :: cell
1645 : TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: rho_r
1646 : TYPE(torch_tensor_type), INTENT(INOUT) :: grid_weight_grad_t, &
1647 : atomic_grid_weight_grad_t
1648 :
1649 : INTEGER :: feature_begin, feature_end, feature_pos, &
1650 : i, iatom, j, jatom, k, local_row, &
1651 : natom, row
1652 : INTEGER, DIMENSION(2, 3) :: bo
1653 : LOGICAL, ALLOCATABLE, DIMENSION(:) :: included
1654 : REAL(KIND=dp) :: grid_base_weight, weight_grad
1655 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: weights
1656 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: atom_coords_pbc
1657 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: dweights_datom, dweights_dstrain
1658 : REAL(KIND=dp), DIMENSION(3) :: grid_point
1659 8 : REAL(KIND=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
1660 :
1661 8 : NULLIFY (atomic_grid_weight_grad, grid_weight_grad)
1662 8 : CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
1663 8 : CALL torch_tensor_grad(features%atomic_grid_weights_t, atomic_grid_weight_grad_t)
1664 8 : CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
1665 8 : CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
1666 :
1667 8 : natom = SIZE(particle_set)
1668 8 : CPASSERT(SIZE(atom_force, 1) == 3)
1669 8 : CPASSERT(SIZE(atom_force, 2) == natom)
1670 : ALLOCATE (atom_coords_pbc(3, natom), included(natom), weights(natom), &
1671 96 : dweights_datom(3, natom, natom), dweights_dstrain(3, 3, natom))
1672 24 : DO iatom = 1, natom
1673 24 : atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.TRUE.)
1674 : END DO
1675 :
1676 80 : bo = rho_r(1)%pw_grid%bounds_local
1677 8 : local_row = 0
1678 182 : DO k = bo(1, 3), bo(2, 3)
1679 4088 : DO j = bo(1, 2), bo(2, 2)
1680 48927 : DO i = bo(1, 1), bo(2, 1)
1681 44847 : local_row = local_row + 1
1682 179388 : grid_point = native_grid_coordinate(rho_r(1)%pw_grid, [i, j, k])
1683 : CALL skala_gpw_smooth_partition_derivatives(grid_point, atom_coords_pbc, cell, &
1684 : weights, included, dweights_datom, &
1685 44847 : dweights_dstrain)
1686 44847 : feature_begin = features%local_feature_offsets(local_row)
1687 44847 : feature_end = features%local_feature_offsets(local_row + 1) - 1
1688 134541 : CPASSERT(feature_end - feature_begin + 1 == COUNT(included))
1689 44847 : grid_base_weight = 0.0_dp
1690 134287 : DO feature_pos = feature_begin, feature_end
1691 89440 : row = features%local_feature_rows(feature_pos)
1692 134287 : grid_base_weight = grid_base_weight + features%grid_weights(row)
1693 : END DO
1694 : feature_pos = feature_begin
1695 134541 : DO iatom = 1, natom
1696 89694 : IF (.NOT. included(iatom)) CYCLE
1697 89440 : row = features%local_feature_rows(feature_pos)
1698 : weight_grad = grid_base_weight*grid_weight_grad(row) + &
1699 : rho_r(1)%pw_grid%dvol*atomic_grid_weight_grad(row)* &
1700 89440 : smooth_partition_atomic_weight_scale_derivative(weights(iatom))
1701 268320 : DO jatom = 1, natom
1702 : atom_force(:, jatom) = atom_force(:, jatom) + &
1703 804960 : weight_grad*dweights_datom(:, jatom, iatom)
1704 : END DO
1705 134541 : feature_pos = feature_pos + 1
1706 : END DO
1707 48753 : CPASSERT(feature_pos == feature_end + 1)
1708 : END DO
1709 : END DO
1710 : END DO
1711 8 : CPASSERT(local_row == features%nflat_local)
1712 :
1713 8 : DEALLOCATE (atom_coords_pbc, dweights_datom, dweights_dstrain, included, weights)
1714 8 : CALL torch_tensor_release(grid_weight_grad_t)
1715 8 : CALL torch_tensor_release(atomic_grid_weight_grad_t)
1716 :
1717 8 : END SUBROUTINE add_smooth_partition_force
1718 :
1719 : ! **************************************************************************************************
1720 : !> \brief Add the virial from SMOOTH native-grid atom partition weights.
1721 : !> \param virial_xc ...
1722 : !> \param features ...
1723 : !> \param particle_set ...
1724 : !> \param cell ...
1725 : !> \param rho_r ...
1726 : !> \param grid_weight_grad_t ...
1727 : !> \param atomic_grid_weight_grad_t ...
1728 : ! **************************************************************************************************
1729 6 : SUBROUTINE build_smooth_partition_virial(virial_xc, features, particle_set, cell, rho_r, &
1730 : grid_weight_grad_t, atomic_grid_weight_grad_t)
1731 : REAL(KIND=dp), DIMENSION(3, 3), INTENT(INOUT) :: virial_xc
1732 : TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1733 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
1734 : TYPE(cell_type), POINTER :: cell
1735 : TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: rho_r
1736 : TYPE(torch_tensor_type), INTENT(INOUT) :: grid_weight_grad_t, &
1737 : atomic_grid_weight_grad_t
1738 :
1739 : INTEGER :: feature_begin, feature_end, feature_pos, &
1740 : i, iatom, idir, j, jdir, k, local_row, &
1741 : natom, row
1742 : INTEGER, DIMENSION(2, 3) :: bo
1743 : LOGICAL, ALLOCATABLE, DIMENSION(:) :: included
1744 : REAL(KIND=dp) :: grid_base_weight, tmp, weight_grad
1745 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: weights
1746 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: atom_coords_pbc
1747 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: dweights_datom, dweights_dstrain
1748 : REAL(KIND=dp), DIMENSION(3) :: grid_point
1749 6 : REAL(KIND=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
1750 :
1751 6 : NULLIFY (atomic_grid_weight_grad, grid_weight_grad)
1752 6 : CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
1753 6 : CALL torch_tensor_grad(features%atomic_grid_weights_t, atomic_grid_weight_grad_t)
1754 6 : CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
1755 6 : CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
1756 :
1757 6 : natom = SIZE(particle_set)
1758 : ALLOCATE (atom_coords_pbc(3, natom), included(natom), weights(natom), &
1759 72 : dweights_datom(3, natom, natom), dweights_dstrain(3, 3, natom))
1760 18 : DO iatom = 1, natom
1761 18 : atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.TRUE.)
1762 : END DO
1763 :
1764 60 : bo = rho_r(1)%pw_grid%bounds_local
1765 6 : local_row = 0
1766 150 : DO k = bo(1, 3), bo(2, 3)
1767 3606 : DO j = bo(1, 2), bo(2, 2)
1768 45072 : DO i = bo(1, 1), bo(2, 1)
1769 41472 : local_row = local_row + 1
1770 165888 : grid_point = native_grid_coordinate(rho_r(1)%pw_grid, [i, j, k])
1771 : CALL skala_gpw_smooth_partition_derivatives(grid_point, atom_coords_pbc, cell, &
1772 : weights, included, dweights_datom, &
1773 41472 : dweights_dstrain)
1774 41472 : feature_begin = features%local_feature_offsets(local_row)
1775 41472 : feature_end = features%local_feature_offsets(local_row + 1) - 1
1776 124416 : CPASSERT(feature_end - feature_begin + 1 == COUNT(included))
1777 41472 : grid_base_weight = 0.0_dp
1778 124206 : DO feature_pos = feature_begin, feature_end
1779 82734 : row = features%local_feature_rows(feature_pos)
1780 124206 : grid_base_weight = grid_base_weight + features%grid_weights(row)
1781 : END DO
1782 : feature_pos = feature_begin
1783 124416 : DO iatom = 1, natom
1784 82944 : IF (.NOT. included(iatom)) CYCLE
1785 82734 : row = features%local_feature_rows(feature_pos)
1786 : weight_grad = grid_base_weight*grid_weight_grad(row) + &
1787 : rho_r(1)%pw_grid%dvol*atomic_grid_weight_grad(row)* &
1788 82734 : smooth_partition_atomic_weight_scale_derivative(weights(iatom))
1789 330936 : DO idir = 1, 3
1790 827340 : DO jdir = 1, idir
1791 496404 : tmp = weight_grad*dweights_dstrain(idir, jdir, iatom)
1792 496404 : virial_xc(jdir, idir) = virial_xc(jdir, idir) + tmp
1793 744606 : IF (idir /= jdir) virial_xc(idir, jdir) = virial_xc(idir, jdir) + tmp
1794 : END DO
1795 : END DO
1796 124416 : feature_pos = feature_pos + 1
1797 : END DO
1798 44928 : CPASSERT(feature_pos == feature_end + 1)
1799 : END DO
1800 : END DO
1801 : END DO
1802 6 : CPASSERT(local_row == features%nflat_local)
1803 :
1804 6 : DEALLOCATE (atom_coords_pbc, dweights_datom, dweights_dstrain, included, weights)
1805 6 : CALL torch_tensor_release(grid_weight_grad_t)
1806 6 : CALL torch_tensor_release(atomic_grid_weight_grad_t)
1807 :
1808 6 : END SUBROUTINE build_smooth_partition_virial
1809 :
1810 : ! **************************************************************************************************
1811 : !> \brief Return the Cartesian coordinate of a regular GPW grid point.
1812 : !> \param pw_grid ...
1813 : !> \param index ...
1814 : !> \return ...
1815 : ! **************************************************************************************************
1816 86319 : FUNCTION native_grid_coordinate(pw_grid, index) RESULT(coord)
1817 : TYPE(pw_grid_type), POINTER :: pw_grid
1818 : INTEGER, DIMENSION(3), INTENT(IN) :: index
1819 : REAL(KIND=dp), DIMENSION(3) :: coord
1820 :
1821 : INTEGER, DIMENSION(3) :: relative_index
1822 :
1823 345276 : relative_index = index - pw_grid%bounds(1, :)
1824 : coord = REAL(relative_index(1), KIND=dp)*pw_grid%dh(:, 1) + &
1825 : REAL(relative_index(2), KIND=dp)*pw_grid%dh(:, 2) + &
1826 345276 : REAL(relative_index(3), KIND=dp)*pw_grid%dh(:, 3)
1827 :
1828 86319 : END FUNCTION native_grid_coordinate
1829 :
1830 : ! **************************************************************************************************
1831 : !> \brief Evaluate a rank-local atom chunk as multiple atom-contiguous Torch subchunks.
1832 : !> \param features ...
1833 : !> \param group ...
1834 : !> \param max_rows ...
1835 : !> \param compute_grads ...
1836 : !> \param exc ...
1837 : !> \param density_grad ...
1838 : !> \param grad_grad ...
1839 : !> \param kin_grad ...
1840 : !> \param collapse_spin_grads ...
1841 : ! **************************************************************************************************
1842 2 : SUBROUTINE evaluate_atom_subchunks(features, group, max_rows, compute_grads, exc, &
1843 : density_grad, grad_grad, kin_grad, collapse_spin_grads)
1844 : TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1845 :
1846 : CLASS(mp_comm_type), INTENT(IN) :: group
1847 : INTEGER, INTENT(IN) :: max_rows
1848 : LOGICAL, INTENT(IN) :: compute_grads, collapse_spin_grads
1849 : REAL(KIND=dp), INTENT(OUT) :: exc
1850 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), &
1851 : INTENT(OUT) :: density_grad, kin_grad
1852 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :), &
1853 : INTENT(OUT) :: grad_grad
1854 :
1855 : INTEGER :: isubchunk, nroute_grad_per_point, &
1856 : nroute_recv_points, nroute_send_points, &
1857 : nsubchunks, phase_handle, subphase_handle
1858 2 : INTEGER, ALLOCATABLE, DIMENSION(:) :: route_grad_return_recv_counts, &
1859 2 : route_grad_return_recv_displs, &
1860 2 : route_grad_return_send_counts, &
1861 2 : route_grad_return_send_displs, &
1862 2 : subchunk_atom_begin, &
1863 2 : subchunk_atom_count, &
1864 2 : subchunk_row_begin, &
1865 2 : subchunk_row_count
1866 : REAL(KIND=dp) :: subchunk_exc
1867 2 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: recv_grad_buffer, send_grad_buffer
1868 2 : TYPE(skala_gpw_feature_type) :: subchunk
1869 : TYPE(torch_tensor_type) :: subchunk_exc_tensor
1870 :
1871 0 : CPASSERT(features%uses_atom_chunks)
1872 2 : CPASSERT(max_rows > 0)
1873 : CALL skala_gpw_atom_subchunk_layout(max_rows, subchunk_atom_begin, subchunk_atom_count, &
1874 2 : subchunk_row_begin, subchunk_row_count)
1875 2 : nsubchunks = SIZE(subchunk_atom_begin)
1876 :
1877 2 : exc = 0.0_dp
1878 2 : IF (compute_grads) THEN
1879 2 : CPASSERT(features%uses_atom_chunk_routing)
1880 6 : nroute_recv_points = SUM(features%route_point_recv_counts)
1881 2 : nroute_send_points = SIZE(features%route_send_local_rows)
1882 6 : CPASSERT(SUM(features%route_point_send_counts) == nroute_send_points)
1883 2 : CPASSERT(SIZE(features%route_chunk_offsets) == nroute_recv_points + 1)
1884 2 : nroute_grad_per_point = ngrad_per_point
1885 2 : IF (collapse_spin_grads) nroute_grad_per_point = ncollapsed_grad_per_point
1886 : ALLOCATE (send_grad_buffer(MAX(1, nroute_grad_per_point*nroute_recv_points)), &
1887 : recv_grad_buffer(MAX(1, nroute_grad_per_point*nroute_send_points)), &
1888 : route_grad_return_send_counts(SIZE(features%route_point_recv_counts)), &
1889 : route_grad_return_send_displs(SIZE(features%route_point_recv_displs)), &
1890 : route_grad_return_recv_counts(SIZE(features%route_point_send_counts)), &
1891 26 : route_grad_return_recv_displs(SIZE(features%route_point_send_displs)))
1892 : route_grad_return_send_counts(:) = &
1893 6 : nroute_grad_per_point*features%route_point_recv_counts
1894 : route_grad_return_send_displs(:) = &
1895 6 : nroute_grad_per_point*features%route_point_recv_displs
1896 : route_grad_return_recv_counts(:) = &
1897 6 : nroute_grad_per_point*features%route_point_send_counts
1898 : route_grad_return_recv_displs(:) = &
1899 6 : nroute_grad_per_point*features%route_point_send_displs
1900 : END IF
1901 :
1902 2 : CALL timeset("skala_gpw_atom_subchunks", phase_handle)
1903 6 : DO isubchunk = 1, nsubchunks
1904 4 : CALL timeset("skala_gpw_atom_subchunk_build", subphase_handle)
1905 : CALL skala_gpw_feature_build_atom_subchunk_bounds(features, subchunk, &
1906 : subchunk_atom_begin(isubchunk), &
1907 : subchunk_atom_count(isubchunk), &
1908 : subchunk_row_begin(isubchunk), &
1909 : subchunk_row_count(isubchunk), &
1910 4 : compute_grads)
1911 4 : CALL timestop(subphase_handle)
1912 4 : CALL timeset("skala_gpw_atom_subchunk_forward", subphase_handle)
1913 : CALL skala_torch_model_get_exc(cached_model, subchunk%inputs, &
1914 : subchunk%grid_weights_t, subchunk_exc_tensor, &
1915 4 : subchunk_exc)
1916 4 : CALL timestop(subphase_handle)
1917 4 : exc = exc + subchunk_exc
1918 4 : IF (compute_grads) THEN
1919 4 : CALL timeset("skala_gpw_atom_subchunk_backward", subphase_handle)
1920 4 : CALL torch_tensor_backward_scalar(subchunk_exc_tensor)
1921 4 : CALL timestop(subphase_handle)
1922 : END IF
1923 4 : CALL timeset("skala_gpw_atom_subchunk_release", subphase_handle)
1924 4 : CALL torch_tensor_release(subchunk_exc_tensor)
1925 4 : CALL skala_gpw_feature_release(subchunk)
1926 18 : CALL timestop(subphase_handle)
1927 : END DO
1928 2 : IF (compute_grads .AND. features%chunk_feature_count > 0) THEN
1929 2 : CALL timeset("skala_gpw_atom_subchunk_grad_pack", subphase_handle)
1930 2 : CALL pack_atom_chunk_grads(features, send_grad_buffer, .TRUE., collapse_spin_grads)
1931 2 : CALL timestop(subphase_handle)
1932 : END IF
1933 2 : CALL timestop(phase_handle)
1934 :
1935 2 : IF (compute_grads) THEN
1936 2 : CALL timeset("skala_gpw_grad_route_comm", phase_handle)
1937 : CALL group%alltoall(send_grad_buffer, route_grad_return_send_counts, &
1938 : route_grad_return_send_displs, recv_grad_buffer, &
1939 2 : route_grad_return_recv_counts, route_grad_return_recv_displs)
1940 2 : CALL timestop(phase_handle)
1941 :
1942 2 : CALL timeset("skala_gpw_grad_route_scatter", phase_handle)
1943 : CALL scatter_routed_atom_chunk_grads(features, recv_grad_buffer, collapse_spin_grads, &
1944 2 : density_grad, grad_grad, kin_grad)
1945 2 : CALL timestop(phase_handle)
1946 :
1947 0 : DEALLOCATE (recv_grad_buffer, route_grad_return_recv_counts, &
1948 0 : route_grad_return_recv_displs, route_grad_return_send_counts, &
1949 2 : route_grad_return_send_displs, send_grad_buffer)
1950 : END IF
1951 2 : DEALLOCATE (subchunk_atom_begin, subchunk_atom_count, subchunk_row_begin, subchunk_row_count)
1952 :
1953 4 : END SUBROUTINE evaluate_atom_subchunks
1954 :
1955 : ! **************************************************************************************************
1956 : !> \brief Select an automatic atom-subchunk row cap.
1957 : !> \param features ...
1958 : !> \param group ...
1959 : !> \return ...
1960 : ! **************************************************************************************************
1961 36 : FUNCTION auto_atom_chunk_max_rows(features, group) RESULT(max_rows)
1962 : TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1963 :
1964 : CLASS(mp_comm_type), INTENT(IN) :: group
1965 : INTEGER :: max_rows
1966 :
1967 : INTEGER :: local_rows_max, target_rows
1968 :
1969 36 : local_rows_max = features%chunk_feature_count
1970 36 : CALL group%max(local_rows_max)
1971 36 : IF (local_rows_max <= 0) THEN
1972 36 : max_rows = 0
1973 : RETURN
1974 : END IF
1975 :
1976 36 : IF (group%num_pe > 1) THEN
1977 36 : target_rows = CEILING(REAL(local_rows_max, KIND=dp)/2.0_dp)
1978 : max_rows = atom_chunk_auto_row_quantum* &
1979 36 : ((target_rows + atom_chunk_auto_row_quantum - 1)/atom_chunk_auto_row_quantum)
1980 : ELSE
1981 0 : target_rows = NINT(REAL(local_rows_max, KIND=dp)/4.0_dp)
1982 : max_rows = atom_chunk_auto_row_quantum* &
1983 : MAX(1, NINT(REAL(target_rows, KIND=dp)/ &
1984 0 : REAL(atom_chunk_auto_row_quantum, KIND=dp)))
1985 : END IF
1986 36 : max_rows = MAX(atom_chunk_auto_min_rows, MIN(atom_chunk_auto_max_rows, max_rows))
1987 :
1988 36 : END FUNCTION auto_atom_chunk_max_rows
1989 :
1990 : ! **************************************************************************************************
1991 : !> \brief Map full Torch feature gradients back to this rank's local grid order.
1992 : !> \param features ...
1993 : !> \param density_grad ...
1994 : !> \param grad_grad ...
1995 : !> \param kin_grad ...
1996 : ! **************************************************************************************************
1997 8 : SUBROUTINE fetch_local_feature_grads(features, density_grad, grad_grad, kin_grad)
1998 : TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1999 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), &
2000 : INTENT(OUT) :: density_grad
2001 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :), &
2002 : INTENT(OUT) :: grad_grad
2003 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), &
2004 : INTENT(OUT) :: kin_grad
2005 :
2006 : INTEGER :: feature_pos, i, j, k, local_row, row
2007 8 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: density_grad_all, kin_grad_all
2008 8 : REAL(KIND=dp), DIMENSION(:, :, :), POINTER :: grad_grad_all
2009 : TYPE(torch_tensor_type) :: density_grad_t, grad_grad_t, kin_grad_t
2010 :
2011 8 : NULLIFY (density_grad_all, grad_grad_all, kin_grad_all)
2012 : CALL get_feature_grad_views(features, density_grad_t, grad_grad_t, kin_grad_t, &
2013 8 : density_grad_all, grad_grad_all, kin_grad_all)
2014 8 : CPASSERT(SIZE(density_grad_all, 1) == features%nflat)
2015 8 : CPASSERT(SIZE(density_grad_all, 2) == 2)
2016 8 : CPASSERT(SIZE(grad_grad_all, 1) == features%nflat)
2017 8 : CPASSERT(SIZE(grad_grad_all, 2) == 3)
2018 8 : CPASSERT(SIZE(grad_grad_all, 3) == 2)
2019 8 : CPASSERT(SIZE(kin_grad_all, 1) == features%nflat)
2020 8 : CPASSERT(SIZE(kin_grad_all, 2) == 2)
2021 :
2022 0 : ALLOCATE (density_grad(features%nflat_local, 2), &
2023 0 : grad_grad(features%nflat_local, 3, 2), &
2024 56 : kin_grad(features%nflat_local, 2))
2025 8 : density_grad = 0.0_dp
2026 8 : grad_grad = 0.0_dp
2027 8 : kin_grad = 0.0_dp
2028 8 : local_row = 0
2029 198 : DO k = LBOUND(features%feature_index, 3), UBOUND(features%feature_index, 3)
2030 4436 : DO j = LBOUND(features%feature_index, 2), UBOUND(features%feature_index, 2)
2031 56739 : DO i = LBOUND(features%feature_index, 1), UBOUND(features%feature_index, 1)
2032 44847 : local_row = local_row + 1
2033 134287 : DO feature_pos = features%local_feature_offsets(local_row), &
2034 48753 : features%local_feature_offsets(local_row + 1) - 1
2035 89440 : row = features%local_feature_rows(feature_pos)
2036 89440 : CPASSERT(row >= 1 .AND. row <= features%nflat)
2037 : density_grad(local_row, :) = density_grad(local_row, :) + &
2038 268320 : density_grad_all(row, :)
2039 : grad_grad(local_row, :, :) = grad_grad(local_row, :, :) + &
2040 804960 : grad_grad_all(row, :, :)
2041 313167 : kin_grad(local_row, :) = kin_grad(local_row, :) + kin_grad_all(row, :)
2042 : END DO
2043 : END DO
2044 : END DO
2045 : END DO
2046 8 : CPASSERT(local_row == features%nflat_local)
2047 :
2048 8 : CALL torch_tensor_release(density_grad_t)
2049 8 : CALL torch_tensor_release(grad_grad_t)
2050 8 : CALL torch_tensor_release(kin_grad_t)
2051 :
2052 8 : END SUBROUTINE fetch_local_feature_grads
2053 :
2054 : ! **************************************************************************************************
2055 : !> \brief Pack atom-chunk Torch gradients into CP2K communication buffers.
2056 : !> \param features ...
2057 : !> \param TARGET ...
2058 : !> \param route_to_return_positions ...
2059 : !> \param collapse_spin_grads ...
2060 : ! **************************************************************************************************
2061 40 : SUBROUTINE pack_atom_chunk_grads(features, TARGET, route_to_return_positions, &
2062 : collapse_spin_grads)
2063 : TYPE(skala_gpw_feature_type), INTENT(IN) :: features
2064 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:), &
2065 : INTENT(INOUT) :: target
2066 : LOGICAL, INTENT(IN) :: route_to_return_positions
2067 : LOGICAL, INTENT(IN), OPTIONAL :: collapse_spin_grads
2068 :
2069 : INTEGER :: base, feature_pos, irow, &
2070 : ngrad_buffer_per_point, point_pos, &
2071 : target_points
2072 : LOGICAL :: my_collapse_spin_grads
2073 40 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: chunk_density_grad, chunk_kin_grad
2074 40 : REAL(KIND=dp), DIMENSION(:, :, :), POINTER :: chunk_grad_grad
2075 : TYPE(torch_tensor_type) :: density_grad_t, grad_grad_t, kin_grad_t
2076 :
2077 40 : my_collapse_spin_grads = .FALSE.
2078 80 : IF (PRESENT(collapse_spin_grads)) my_collapse_spin_grads = collapse_spin_grads
2079 40 : ngrad_buffer_per_point = ngrad_per_point
2080 40 : IF (my_collapse_spin_grads) ngrad_buffer_per_point = ncollapsed_grad_per_point
2081 :
2082 40 : NULLIFY (chunk_density_grad, chunk_grad_grad, chunk_kin_grad)
2083 : CALL get_feature_grad_views(features, density_grad_t, grad_grad_t, kin_grad_t, &
2084 40 : chunk_density_grad, chunk_grad_grad, chunk_kin_grad)
2085 40 : CPASSERT(MOD(SIZE(TARGET), ngrad_buffer_per_point) == 0)
2086 40 : target_points = SIZE(TARGET)/ngrad_buffer_per_point
2087 40 : CPASSERT(SIZE(chunk_density_grad, 1) == features%chunk_feature_count)
2088 40 : CPASSERT(SIZE(chunk_grad_grad, 1) == features%chunk_feature_count)
2089 40 : CPASSERT(SIZE(chunk_grad_grad, 2) == 3)
2090 40 : CPASSERT(SIZE(chunk_kin_grad, 1) == features%chunk_feature_count)
2091 40 : IF (features%uses_collapsed_rks_dynamic) THEN
2092 34 : CPASSERT(my_collapse_spin_grads)
2093 34 : CPASSERT(SIZE(chunk_density_grad, 2) == 1)
2094 34 : CPASSERT(SIZE(chunk_grad_grad, 3) == 1)
2095 34 : CPASSERT(SIZE(chunk_kin_grad, 2) == 1)
2096 : ELSE
2097 6 : CPASSERT(SIZE(chunk_density_grad, 2) == 2)
2098 6 : CPASSERT(SIZE(chunk_grad_grad, 3) == 2)
2099 6 : CPASSERT(SIZE(chunk_kin_grad, 2) == 2)
2100 : END IF
2101 :
2102 40 : IF (route_to_return_positions) THEN
2103 40 : CPASSERT(target_points == SIZE(features%route_chunk_offsets) - 1)
2104 : !$OMP PARALLEL DO DEFAULT(NONE) &
2105 : !$OMP SHARED(chunk_density_grad, chunk_grad_grad, chunk_kin_grad, features, &
2106 : !$OMP my_collapse_spin_grads, ngrad_buffer_per_point, TARGET, target_points) &
2107 40 : !$OMP PRIVATE(base, feature_pos, irow, point_pos)
2108 : DO point_pos = 1, target_points
2109 : base = ngrad_buffer_per_point*(point_pos - 1)
2110 : TARGET(base + 1:base + ngrad_buffer_per_point) = 0.0_dp
2111 : DO feature_pos = features%route_chunk_offsets(point_pos), &
2112 : features%route_chunk_offsets(point_pos + 1) - 1
2113 : irow = features%route_chunk_rows(feature_pos)
2114 : CPASSERT(irow >= 1 .AND. irow <= features%chunk_feature_count)
2115 : IF (my_collapse_spin_grads) THEN
2116 : IF (features%uses_collapsed_rks_dynamic) THEN
2117 : TARGET(base + 1) = TARGET(base + 1) + &
2118 : 0.5_dp*chunk_density_grad(irow, 1)
2119 : TARGET(base + 2) = TARGET(base + 2) + &
2120 : 0.5_dp*chunk_grad_grad(irow, 1, 1)
2121 : TARGET(base + 3) = TARGET(base + 3) + &
2122 : 0.5_dp*chunk_grad_grad(irow, 2, 1)
2123 : TARGET(base + 4) = TARGET(base + 4) + &
2124 : 0.5_dp*chunk_grad_grad(irow, 3, 1)
2125 : TARGET(base + 5) = TARGET(base + 5) + &
2126 : 0.5_dp*chunk_kin_grad(irow, 1)
2127 : ELSE
2128 : TARGET(base + 1) = TARGET(base + 1) + &
2129 : 0.5_dp*(chunk_density_grad(irow, 1) + &
2130 : chunk_density_grad(irow, 2))
2131 : TARGET(base + 2) = TARGET(base + 2) + &
2132 : 0.5_dp*(chunk_grad_grad(irow, 1, 1) + &
2133 : chunk_grad_grad(irow, 1, 2))
2134 : TARGET(base + 3) = TARGET(base + 3) + &
2135 : 0.5_dp*(chunk_grad_grad(irow, 2, 1) + &
2136 : chunk_grad_grad(irow, 2, 2))
2137 : TARGET(base + 4) = TARGET(base + 4) + &
2138 : 0.5_dp*(chunk_grad_grad(irow, 3, 1) + &
2139 : chunk_grad_grad(irow, 3, 2))
2140 : TARGET(base + 5) = TARGET(base + 5) + &
2141 : 0.5_dp*(chunk_kin_grad(irow, 1) + &
2142 : chunk_kin_grad(irow, 2))
2143 : END IF
2144 : ELSE
2145 : TARGET(base + 1:base + 2) = TARGET(base + 1:base + 2) + &
2146 : chunk_density_grad(irow, :)
2147 : TARGET(base + 3) = TARGET(base + 3) + chunk_grad_grad(irow, 1, 1)
2148 : TARGET(base + 4) = TARGET(base + 4) + chunk_grad_grad(irow, 2, 1)
2149 : TARGET(base + 5) = TARGET(base + 5) + chunk_grad_grad(irow, 3, 1)
2150 : TARGET(base + 6) = TARGET(base + 6) + chunk_grad_grad(irow, 1, 2)
2151 : TARGET(base + 7) = TARGET(base + 7) + chunk_grad_grad(irow, 2, 2)
2152 : TARGET(base + 8) = TARGET(base + 8) + chunk_grad_grad(irow, 3, 2)
2153 : TARGET(base + 9:base + 10) = TARGET(base + 9:base + 10) + &
2154 : chunk_kin_grad(irow, :)
2155 : END IF
2156 : END DO
2157 : END DO
2158 : !$OMP END PARALLEL DO
2159 : ELSE
2160 0 : CPASSERT(target_points >= features%chunk_feature_count)
2161 : !$OMP PARALLEL DO DEFAULT(NONE) &
2162 : !$OMP SHARED(chunk_density_grad, chunk_grad_grad, chunk_kin_grad, features, &
2163 : !$OMP my_collapse_spin_grads, ngrad_buffer_per_point, TARGET) &
2164 0 : !$OMP PRIVATE(base, irow)
2165 : DO irow = 1, features%chunk_feature_count
2166 : base = ngrad_buffer_per_point*(irow - 1)
2167 : IF (my_collapse_spin_grads) THEN
2168 : IF (features%uses_collapsed_rks_dynamic) THEN
2169 : TARGET(base + 1) = 0.5_dp*chunk_density_grad(irow, 1)
2170 : TARGET(base + 2) = 0.5_dp*chunk_grad_grad(irow, 1, 1)
2171 : TARGET(base + 3) = 0.5_dp*chunk_grad_grad(irow, 2, 1)
2172 : TARGET(base + 4) = 0.5_dp*chunk_grad_grad(irow, 3, 1)
2173 : TARGET(base + 5) = 0.5_dp*chunk_kin_grad(irow, 1)
2174 : ELSE
2175 : TARGET(base + 1) = 0.5_dp*(chunk_density_grad(irow, 1) + &
2176 : chunk_density_grad(irow, 2))
2177 : TARGET(base + 2) = 0.5_dp*(chunk_grad_grad(irow, 1, 1) + &
2178 : chunk_grad_grad(irow, 1, 2))
2179 : TARGET(base + 3) = 0.5_dp*(chunk_grad_grad(irow, 2, 1) + &
2180 : chunk_grad_grad(irow, 2, 2))
2181 : TARGET(base + 4) = 0.5_dp*(chunk_grad_grad(irow, 3, 1) + &
2182 : chunk_grad_grad(irow, 3, 2))
2183 : TARGET(base + 5) = 0.5_dp*(chunk_kin_grad(irow, 1) + &
2184 : chunk_kin_grad(irow, 2))
2185 : END IF
2186 : ELSE
2187 : TARGET(base + 1:base + 2) = chunk_density_grad(irow, :)
2188 : TARGET(base + 3) = chunk_grad_grad(irow, 1, 1)
2189 : TARGET(base + 4) = chunk_grad_grad(irow, 2, 1)
2190 : TARGET(base + 5) = chunk_grad_grad(irow, 3, 1)
2191 : TARGET(base + 6) = chunk_grad_grad(irow, 1, 2)
2192 : TARGET(base + 7) = chunk_grad_grad(irow, 2, 2)
2193 : TARGET(base + 8) = chunk_grad_grad(irow, 3, 2)
2194 : TARGET(base + 9:base + 10) = chunk_kin_grad(irow, :)
2195 : END IF
2196 : END DO
2197 : !$OMP END PARALLEL DO
2198 : END IF
2199 :
2200 40 : CALL torch_tensor_release(density_grad_t)
2201 40 : CALL torch_tensor_release(grad_grad_t)
2202 40 : CALL torch_tensor_release(kin_grad_t)
2203 :
2204 40 : END SUBROUTINE pack_atom_chunk_grads
2205 :
2206 : ! **************************************************************************************************
2207 : !> \brief Scatter routed atom-chunk gradients into local grid-row order.
2208 : !> \param features ...
2209 : !> \param recv_grad_buffer ...
2210 : !> \param collapse_spin_grads ...
2211 : !> \param density_grad ...
2212 : !> \param grad_grad ...
2213 : !> \param kin_grad ...
2214 : ! **************************************************************************************************
2215 40 : SUBROUTINE scatter_routed_atom_chunk_grads(features, recv_grad_buffer, collapse_spin_grads, &
2216 : density_grad, grad_grad, kin_grad)
2217 : TYPE(skala_gpw_feature_type), INTENT(IN) :: features
2218 : REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: recv_grad_buffer
2219 : LOGICAL, INTENT(IN) :: collapse_spin_grads
2220 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), &
2221 : INTENT(OUT) :: density_grad
2222 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :), &
2223 : INTENT(OUT) :: grad_grad
2224 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), &
2225 : INTENT(OUT) :: kin_grad
2226 :
2227 : INTEGER :: base, local_row, nflat_local, &
2228 : nroute_grad_per_point, nroute_points, &
2229 : point_pos, row_route_pos
2230 :
2231 40 : nflat_local = features%nflat_local
2232 40 : nroute_points = SIZE(features%route_send_local_rows)
2233 40 : nroute_grad_per_point = ngrad_per_point
2234 40 : IF (collapse_spin_grads) nroute_grad_per_point = ncollapsed_grad_per_point
2235 40 : CPASSERT(SIZE(recv_grad_buffer) >= nroute_grad_per_point*nroute_points)
2236 0 : ALLOCATE (density_grad(nflat_local, 2), grad_grad(nflat_local, 3, 2), &
2237 280 : kin_grad(nflat_local, 2))
2238 :
2239 : !$OMP PARALLEL DO DEFAULT(NONE) &
2240 : !$OMP SHARED(collapse_spin_grads, density_grad, features, grad_grad, kin_grad, nflat_local, &
2241 : !$OMP nroute_grad_per_point, nroute_points, recv_grad_buffer) &
2242 40 : !$OMP PRIVATE(base, local_row, point_pos, row_route_pos)
2243 : DO local_row = 1, nflat_local
2244 : density_grad(local_row, :) = 0.0_dp
2245 : grad_grad(local_row, :, :) = 0.0_dp
2246 : kin_grad(local_row, :) = 0.0_dp
2247 : DO row_route_pos = features%route_row_offsets(local_row), &
2248 : features%route_row_offsets(local_row + 1) - 1
2249 : point_pos = features%route_row_positions(row_route_pos)
2250 : CPASSERT(point_pos >= 1 .AND. point_pos <= nroute_points)
2251 : base = nroute_grad_per_point*(point_pos - 1)
2252 : IF (collapse_spin_grads) THEN
2253 : density_grad(local_row, :) = density_grad(local_row, :) + &
2254 : recv_grad_buffer(base + 1)
2255 : grad_grad(local_row, 1, :) = grad_grad(local_row, 1, :) + &
2256 : recv_grad_buffer(base + 2)
2257 : grad_grad(local_row, 2, :) = grad_grad(local_row, 2, :) + &
2258 : recv_grad_buffer(base + 3)
2259 : grad_grad(local_row, 3, :) = grad_grad(local_row, 3, :) + &
2260 : recv_grad_buffer(base + 4)
2261 : kin_grad(local_row, :) = kin_grad(local_row, :) + recv_grad_buffer(base + 5)
2262 : ELSE
2263 : density_grad(local_row, :) = density_grad(local_row, :) + &
2264 : recv_grad_buffer(base + 1:base + 2)
2265 : grad_grad(local_row, 1, 1) = grad_grad(local_row, 1, 1) + &
2266 : recv_grad_buffer(base + 3)
2267 : grad_grad(local_row, 2, 1) = grad_grad(local_row, 2, 1) + &
2268 : recv_grad_buffer(base + 4)
2269 : grad_grad(local_row, 3, 1) = grad_grad(local_row, 3, 1) + &
2270 : recv_grad_buffer(base + 5)
2271 : grad_grad(local_row, 1, 2) = grad_grad(local_row, 1, 2) + &
2272 : recv_grad_buffer(base + 6)
2273 : grad_grad(local_row, 2, 2) = grad_grad(local_row, 2, 2) + &
2274 : recv_grad_buffer(base + 7)
2275 : grad_grad(local_row, 3, 2) = grad_grad(local_row, 3, 2) + &
2276 : recv_grad_buffer(base + 8)
2277 : kin_grad(local_row, :) = kin_grad(local_row, :) + &
2278 : recv_grad_buffer(base + 9:base + 10)
2279 : END IF
2280 : END DO
2281 : END DO
2282 : !$OMP END PARALLEL DO
2283 :
2284 40 : END SUBROUTINE scatter_routed_atom_chunk_grads
2285 :
2286 : ! **************************************************************************************************
2287 : !> \brief Return CPU views of autograd outputs for the SKALA dynamic feature tensors.
2288 : !> \param features ...
2289 : !> \param density_grad_t ...
2290 : !> \param grad_grad_t ...
2291 : !> \param kin_grad_t ...
2292 : !> \param density_grad ...
2293 : !> \param grad_grad ...
2294 : !> \param kin_grad ...
2295 : ! **************************************************************************************************
2296 48 : SUBROUTINE get_feature_grad_views(features, density_grad_t, grad_grad_t, kin_grad_t, &
2297 : density_grad, grad_grad, kin_grad)
2298 : TYPE(skala_gpw_feature_type), INTENT(IN) :: features
2299 : TYPE(torch_tensor_type), INTENT(INOUT) :: density_grad_t, grad_grad_t, kin_grad_t
2300 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: density_grad
2301 : REAL(KIND=dp), DIMENSION(:, :, :), POINTER :: grad_grad
2302 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: kin_grad
2303 :
2304 48 : NULLIFY (density_grad, grad_grad, kin_grad)
2305 : CALL torch_tensor_grad_batch3(features%density_t, features%grad_t, features%kin_t, &
2306 48 : density_grad_t, grad_grad_t, kin_grad_t)
2307 48 : CALL torch_tensor_data_ptr(density_grad_t, density_grad)
2308 48 : CALL torch_tensor_data_ptr(grad_grad_t, grad_grad)
2309 48 : CALL torch_tensor_data_ptr(kin_grad_t, kin_grad)
2310 :
2311 48 : END SUBROUTINE get_feature_grad_views
2312 :
2313 : ! **************************************************************************************************
2314 : !> \brief Fetch atom-chunk gradients and route them back to their local grid owners.
2315 : !> \param features ...
2316 : !> \param group ...
2317 : !> \param density_grad ...
2318 : !> \param grad_grad ...
2319 : !> \param kin_grad ...
2320 : ! **************************************************************************************************
2321 38 : SUBROUTINE fetch_and_gather_atom_chunk_grads(features, group, density_grad, grad_grad, &
2322 : kin_grad)
2323 : TYPE(skala_gpw_feature_type), INTENT(IN) :: features
2324 :
2325 : CLASS(mp_comm_type), INTENT(IN) :: group
2326 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), &
2327 : INTENT(OUT) :: density_grad, kin_grad
2328 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :), &
2329 : INTENT(OUT) :: grad_grad
2330 :
2331 : INTEGER :: base, feature_pos, i, j, k, local_row, &
2332 : nflat_local, nroute_grad_per_point, &
2333 : nroute_recv_points, nroute_send_points, &
2334 : phase_handle, row
2335 38 : INTEGER, ALLOCATABLE, DIMENSION(:) :: route_grad_return_recv_counts, &
2336 38 : route_grad_return_recv_displs, &
2337 38 : route_grad_return_send_counts, &
2338 38 : route_grad_return_send_displs
2339 38 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: chunk_grad_buffer, global_grad_buffer, &
2340 38 : recv_grad_buffer, send_grad_buffer
2341 :
2342 38 : CPASSERT(features%uses_atom_chunks)
2343 :
2344 38 : nflat_local = features%nflat_local
2345 38 : IF (features%uses_atom_chunk_routing) THEN
2346 114 : nroute_recv_points = SUM(features%route_point_recv_counts)
2347 38 : nroute_send_points = SIZE(features%route_send_local_rows)
2348 114 : CPASSERT(SUM(features%route_point_send_counts) == nroute_send_points)
2349 38 : CPASSERT(SIZE(features%route_chunk_offsets) == nroute_recv_points + 1)
2350 :
2351 38 : nroute_grad_per_point = ngrad_per_point
2352 38 : IF (features%uses_collapsed_rks_dynamic) THEN
2353 32 : nroute_grad_per_point = ncollapsed_grad_per_point
2354 : END IF
2355 : ALLOCATE (send_grad_buffer(MAX(1, nroute_grad_per_point*nroute_recv_points)), &
2356 : recv_grad_buffer(MAX(1, nroute_grad_per_point*nroute_send_points)), &
2357 : route_grad_return_send_counts(SIZE(features%route_point_recv_counts)), &
2358 : route_grad_return_send_displs(SIZE(features%route_point_recv_displs)), &
2359 : route_grad_return_recv_counts(SIZE(features%route_point_send_counts)), &
2360 494 : route_grad_return_recv_displs(SIZE(features%route_point_send_displs)))
2361 : route_grad_return_send_counts(:) = &
2362 114 : nroute_grad_per_point*features%route_point_recv_counts
2363 : route_grad_return_send_displs(:) = &
2364 114 : nroute_grad_per_point*features%route_point_recv_displs
2365 : route_grad_return_recv_counts(:) = &
2366 114 : nroute_grad_per_point*features%route_point_send_counts
2367 : route_grad_return_recv_displs(:) = &
2368 114 : nroute_grad_per_point*features%route_point_send_displs
2369 :
2370 38 : IF (features%chunk_feature_count > 0) THEN
2371 38 : CALL timeset("skala_gpw_grad_torch_pack", phase_handle)
2372 : CALL pack_atom_chunk_grads(features, send_grad_buffer, .TRUE., &
2373 38 : features%uses_collapsed_rks_dynamic)
2374 38 : CALL timestop(phase_handle)
2375 : END IF
2376 :
2377 38 : CALL timeset("skala_gpw_grad_route_comm", phase_handle)
2378 : CALL group%alltoall(send_grad_buffer, route_grad_return_send_counts, &
2379 : route_grad_return_send_displs, recv_grad_buffer, &
2380 38 : route_grad_return_recv_counts, route_grad_return_recv_displs)
2381 38 : CALL timestop(phase_handle)
2382 :
2383 38 : CALL timeset("skala_gpw_grad_route_scatter", phase_handle)
2384 : CALL scatter_routed_atom_chunk_grads(features, recv_grad_buffer, &
2385 : features%uses_collapsed_rks_dynamic, &
2386 38 : density_grad, grad_grad, kin_grad)
2387 38 : CALL timestop(phase_handle)
2388 :
2389 0 : DEALLOCATE (recv_grad_buffer, route_grad_return_recv_counts, &
2390 0 : route_grad_return_recv_displs, route_grad_return_send_counts, &
2391 114 : route_grad_return_send_displs, send_grad_buffer)
2392 : ELSE
2393 : ALLOCATE (chunk_grad_buffer(MAX(1, ngrad_per_point*features%chunk_feature_count)), &
2394 0 : global_grad_buffer(ngrad_per_point*features%nflat))
2395 0 : IF (features%chunk_feature_count > 0) THEN
2396 0 : CALL timeset("skala_gpw_grad_torch_pack", phase_handle)
2397 0 : CALL pack_atom_chunk_grads(features, chunk_grad_buffer, .FALSE.)
2398 0 : CALL timestop(phase_handle)
2399 : END IF
2400 :
2401 0 : CALL timeset("skala_gpw_grad_allgatherv", phase_handle)
2402 : CALL group%allgatherv(chunk_grad_buffer, global_grad_buffer, &
2403 0 : features%chunk_grad_counts, features%chunk_grad_displs)
2404 0 : CALL timestop(phase_handle)
2405 :
2406 0 : CALL timeset("skala_gpw_grad_scatter", phase_handle)
2407 0 : ALLOCATE (density_grad(nflat_local, 2), grad_grad(nflat_local, 3, 2), &
2408 0 : kin_grad(nflat_local, 2))
2409 0 : density_grad = 0.0_dp
2410 0 : grad_grad = 0.0_dp
2411 0 : kin_grad = 0.0_dp
2412 0 : local_row = 0
2413 0 : DO k = LBOUND(features%feature_index, 3), UBOUND(features%feature_index, 3)
2414 0 : DO j = LBOUND(features%feature_index, 2), UBOUND(features%feature_index, 2)
2415 0 : DO i = LBOUND(features%feature_index, 1), UBOUND(features%feature_index, 1)
2416 0 : local_row = local_row + 1
2417 0 : DO feature_pos = features%local_feature_offsets(local_row), &
2418 0 : features%local_feature_offsets(local_row + 1) - 1
2419 0 : row = features%local_feature_rows(feature_pos)
2420 0 : CPASSERT(row >= 1 .AND. row <= features%nflat)
2421 0 : base = ngrad_per_point*(row - 1)
2422 : density_grad(local_row, :) = density_grad(local_row, :) + &
2423 0 : global_grad_buffer(base + 1:base + 2)
2424 : grad_grad(local_row, 1, 1) = grad_grad(local_row, 1, 1) + &
2425 0 : global_grad_buffer(base + 3)
2426 : grad_grad(local_row, 2, 1) = grad_grad(local_row, 2, 1) + &
2427 0 : global_grad_buffer(base + 4)
2428 : grad_grad(local_row, 3, 1) = grad_grad(local_row, 3, 1) + &
2429 0 : global_grad_buffer(base + 5)
2430 : grad_grad(local_row, 1, 2) = grad_grad(local_row, 1, 2) + &
2431 0 : global_grad_buffer(base + 6)
2432 : grad_grad(local_row, 2, 2) = grad_grad(local_row, 2, 2) + &
2433 0 : global_grad_buffer(base + 7)
2434 : grad_grad(local_row, 3, 2) = grad_grad(local_row, 3, 2) + &
2435 0 : global_grad_buffer(base + 8)
2436 : kin_grad(local_row, :) = kin_grad(local_row, :) + &
2437 0 : global_grad_buffer(base + 9:base + 10)
2438 : END DO
2439 : END DO
2440 : END DO
2441 : END DO
2442 0 : CALL timestop(phase_handle)
2443 0 : DEALLOCATE (chunk_grad_buffer, global_grad_buffer)
2444 :
2445 : END IF
2446 :
2447 38 : END SUBROUTINE fetch_and_gather_atom_chunk_grads
2448 :
2449 : ! **************************************************************************************************
2450 : !> \brief Build the native SKALA XC virial from feature gradients.
2451 : !> \param virial_xc ...
2452 : !> \param rho_set ...
2453 : !> \param rho_r ...
2454 : !> \param grad_grad ...
2455 : ! **************************************************************************************************
2456 6 : SUBROUTINE build_virial_from_feature_grads(virial_xc, rho_set, rho_r, grad_grad)
2457 : REAL(KIND=dp), DIMENSION(3, 3), INTENT(INOUT) :: virial_xc
2458 : TYPE(xc_rho_set_type), INTENT(IN) :: rho_set
2459 : TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: rho_r
2460 : REAL(KIND=dp), DIMENSION(:, :, :), INTENT(IN) :: grad_grad
2461 :
2462 : INTEGER :: i, idir, ipt, ispin, j, jdir, k, nspins
2463 : INTEGER, DIMENSION(2, 3) :: bo
2464 : REAL(KIND=dp) :: grad_i, tmp
2465 72 : TYPE(cp_3d_r_cp_type), DIMENSION(3) :: drho, drhoa, drhob
2466 :
2467 6 : nspins = SIZE(rho_r)
2468 60 : bo = rho_r(1)%pw_grid%bounds_local
2469 6 : ipt = 0
2470 :
2471 6 : IF (nspins == 1) THEN
2472 6 : CALL xc_rho_set_get(rho_set, drho=drho)
2473 150 : DO k = bo(1, 3), bo(2, 3)
2474 3606 : DO j = bo(1, 2), bo(2, 2)
2475 45072 : DO i = bo(1, 1), bo(2, 1)
2476 41472 : ipt = ipt + 1
2477 169344 : DO idir = 1, 3
2478 124416 : grad_i = 0.5_dp*(grad_grad(ipt, idir, 1) + grad_grad(ipt, idir, 2))
2479 539136 : DO jdir = 1, 3
2480 373248 : tmp = -grad_i*drho(jdir)%array(i, j, k)
2481 497664 : virial_xc(jdir, idir) = virial_xc(jdir, idir) + tmp
2482 : END DO
2483 : END DO
2484 : END DO
2485 : END DO
2486 : END DO
2487 : ELSE
2488 0 : CALL xc_rho_set_get(rho_set, drhoa=drhoa, drhob=drhob)
2489 0 : DO k = bo(1, 3), bo(2, 3)
2490 0 : DO j = bo(1, 2), bo(2, 2)
2491 0 : DO i = bo(1, 1), bo(2, 1)
2492 0 : ipt = ipt + 1
2493 0 : DO idir = 1, 3
2494 0 : DO jdir = 1, 3
2495 : tmp = 0.0_dp
2496 0 : DO ispin = 1, 2
2497 0 : IF (ispin == 1) THEN
2498 0 : tmp = tmp - grad_grad(ipt, idir, ispin)*drhoa(jdir)%array(i, j, k)
2499 : ELSE
2500 0 : tmp = tmp - grad_grad(ipt, idir, ispin)*drhob(jdir)%array(i, j, k)
2501 : END IF
2502 : END DO
2503 0 : virial_xc(jdir, idir) = virial_xc(jdir, idir) + tmp
2504 : END DO
2505 : END DO
2506 : END DO
2507 : END DO
2508 : END DO
2509 : END IF
2510 :
2511 6 : END SUBROUTINE build_virial_from_feature_grads
2512 :
2513 : ! **************************************************************************************************
2514 : !> \brief Print a native SKALA XC virial contribution for diagnostics.
2515 : !> \param label ...
2516 : !> \param delta ...
2517 : !> \param root_rank ...
2518 : ! **************************************************************************************************
2519 0 : SUBROUTINE print_virial_delta(label, delta, root_rank)
2520 : CHARACTER(LEN=*), INTENT(IN) :: label
2521 : REAL(KIND=dp), DIMENSION(3, 3), INTENT(IN) :: delta
2522 : LOGICAL, INTENT(IN) :: root_rank
2523 :
2524 : INTEGER :: i, iw
2525 :
2526 0 : IF (.NOT. root_rank) RETURN
2527 0 : iw = cp_logger_get_default_io_unit()
2528 0 : IF (iw <= 0) RETURN
2529 0 : WRITE (iw, "(T2,A,1X,A)") "SKALA_GPW| XC virial contribution", TRIM(label)
2530 0 : DO i = 1, 3
2531 0 : WRITE (iw, "(T2,A,1X,3ES20.10)") "SKALA_GPW|", delta(i, 1:3)
2532 : END DO
2533 :
2534 : END SUBROUTINE print_virial_delta
2535 :
2536 : ! **************************************************************************************************
2537 : !> \brief Add explicit SKALA coordinate-feature contributions to the XC virial.
2538 : !> \param virial_xc ...
2539 : !> \param features ...
2540 : !> \param atom_coord_grad_t ...
2541 : !> \param grid_coord_grad_t ...
2542 : !> \param root_rank ...
2543 : !> \param print_components ...
2544 : ! **************************************************************************************************
2545 6 : SUBROUTINE build_static_coordinate_virial(virial_xc, features, atom_coord_grad_t, &
2546 : grid_coord_grad_t, root_rank, print_components)
2547 : REAL(KIND=dp), DIMENSION(3, 3), INTENT(INOUT) :: virial_xc
2548 : TYPE(skala_gpw_feature_type), INTENT(IN) :: features
2549 : TYPE(torch_tensor_type), INTENT(INOUT) :: atom_coord_grad_t, grid_coord_grad_t
2550 : LOGICAL, INTENT(IN) :: root_rank
2551 : LOGICAL, INTENT(IN), OPTIONAL :: print_components
2552 :
2553 : INTEGER :: feature_pos, i, iatom, idir, iw, j, &
2554 : jdir, k, local_row, row
2555 : LOGICAL :: my_print_components
2556 : REAL(KIND=dp) :: tmp
2557 : REAL(KIND=dp), DIMENSION(3) :: atom_grad_sum, grid_grad_sum
2558 : REAL(KIND=dp), DIMENSION(3, 3) :: atom_virial, grid_virial
2559 6 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: atom_coord_grad, grid_coord_grad
2560 :
2561 6 : my_print_components = .FALSE.
2562 6 : IF (PRESENT(print_components)) my_print_components = print_components
2563 :
2564 6 : NULLIFY (atom_coord_grad, grid_coord_grad)
2565 6 : CALL torch_tensor_grad(features%grid_coords_t, grid_coord_grad_t)
2566 6 : CALL torch_tensor_data_ptr(grid_coord_grad_t, grid_coord_grad)
2567 6 : CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
2568 :
2569 6 : grid_virial = 0.0_dp
2570 6 : atom_virial = 0.0_dp
2571 6 : grid_grad_sum = 0.0_dp
2572 6 : atom_grad_sum = 0.0_dp
2573 6 : local_row = 0
2574 162 : DO k = LBOUND(features%feature_index, 3), UBOUND(features%feature_index, 3)
2575 3894 : DO j = LBOUND(features%feature_index, 2), UBOUND(features%feature_index, 2)
2576 51984 : DO i = LBOUND(features%feature_index, 1), UBOUND(features%feature_index, 1)
2577 41472 : local_row = local_row + 1
2578 124206 : DO feature_pos = features%local_feature_offsets(local_row), &
2579 44928 : features%local_feature_offsets(local_row + 1) - 1
2580 82734 : row = features%local_feature_rows(feature_pos)
2581 330936 : grid_grad_sum(:) = grid_grad_sum(:) + grid_coord_grad(:, row)
2582 372408 : DO idir = 1, 3
2583 1075542 : DO jdir = 1, 3
2584 744606 : tmp = grid_coord_grad(idir, row)*features%grid_coords(jdir, row)
2585 744606 : grid_virial(idir, jdir) = grid_virial(idir, jdir) + tmp
2586 992808 : virial_xc(idir, jdir) = virial_xc(idir, jdir) + tmp
2587 : END DO
2588 : END DO
2589 : END DO
2590 : END DO
2591 : END DO
2592 : END DO
2593 6 : CPASSERT(local_row == features%nflat_local)
2594 :
2595 6 : IF (root_rank) THEN
2596 9 : DO iatom = 1, SIZE(features%coarse_0_atomic_coords, 2)
2597 24 : atom_grad_sum(:) = atom_grad_sum(:) + atom_coord_grad(:, iatom)
2598 27 : DO idir = 1, 3
2599 78 : DO jdir = 1, 3
2600 54 : tmp = atom_coord_grad(idir, iatom)*features%coarse_0_atomic_coords(jdir, iatom)
2601 54 : atom_virial(idir, jdir) = atom_virial(idir, jdir) + tmp
2602 72 : virial_xc(idir, jdir) = virial_xc(idir, jdir) + tmp
2603 : END DO
2604 : END DO
2605 : END DO
2606 : END IF
2607 :
2608 6 : IF (my_print_components .AND. root_rank) THEN
2609 0 : iw = cp_logger_get_default_io_unit()
2610 0 : IF (iw > 0) THEN
2611 0 : CALL print_virial_delta("static-grid", grid_virial, .TRUE.)
2612 0 : CALL print_virial_delta("static-atom", atom_virial, .TRUE.)
2613 0 : WRITE (iw, "(T2,A,1X,3ES20.10)") "SKALA_GPW| XC coordinate gradient grid sum", &
2614 0 : grid_grad_sum
2615 0 : WRITE (iw, "(T2,A,1X,3ES20.10)") "SKALA_GPW| XC coordinate gradient atom sum", &
2616 0 : atom_grad_sum
2617 0 : WRITE (iw, "(T2,A,1X,3ES20.10)") "SKALA_GPW| XC coordinate gradient total sum", &
2618 0 : grid_grad_sum + atom_grad_sum
2619 : END IF
2620 : END IF
2621 :
2622 6 : CALL torch_tensor_release(grid_coord_grad_t)
2623 :
2624 6 : END SUBROUTINE build_static_coordinate_virial
2625 :
2626 : ! **************************************************************************************************
2627 : !> \brief Add residual SKALA weight-feature contributions to the XC virial.
2628 : !> \param virial_xc ...
2629 : !> \param features ...
2630 : !> \param exc ...
2631 : !> \param grid_weight_grad_t ...
2632 : !> \param atomic_grid_weight_grad_t ...
2633 : !> \param root_rank ...
2634 : !> \param print_components ...
2635 : ! **************************************************************************************************
2636 6 : SUBROUTINE build_weight_virial(virial_xc, features, exc, grid_weight_grad_t, &
2637 : atomic_grid_weight_grad_t, root_rank, print_components)
2638 : REAL(KIND=dp), DIMENSION(3, 3), INTENT(INOUT) :: virial_xc
2639 : TYPE(skala_gpw_feature_type), INTENT(IN) :: features
2640 : REAL(KIND=dp), INTENT(IN) :: exc
2641 : TYPE(torch_tensor_type), INTENT(INOUT) :: grid_weight_grad_t, &
2642 : atomic_grid_weight_grad_t
2643 : LOGICAL, INTENT(IN) :: root_rank
2644 : LOGICAL, INTENT(IN), OPTIONAL :: print_components
2645 :
2646 : INTEGER :: feature_pos, i, idir, iw, j, k, &
2647 : local_row, row
2648 : LOGICAL :: my_print_components
2649 : REAL(KIND=dp) :: atomic_tmp, exc_tmp, grid_tmp, tmp
2650 6 : REAL(KIND=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
2651 :
2652 6 : my_print_components = .FALSE.
2653 6 : IF (PRESENT(print_components)) my_print_components = print_components
2654 :
2655 6 : NULLIFY (atomic_grid_weight_grad, grid_weight_grad)
2656 6 : CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
2657 6 : CALL torch_tensor_grad(features%atomic_grid_weights_t, atomic_grid_weight_grad_t)
2658 6 : CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
2659 6 : CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
2660 :
2661 6 : grid_tmp = 0.0_dp
2662 6 : atomic_tmp = 0.0_dp
2663 6 : local_row = 0
2664 162 : DO k = LBOUND(features%feature_index, 3), UBOUND(features%feature_index, 3)
2665 3894 : DO j = LBOUND(features%feature_index, 2), UBOUND(features%feature_index, 2)
2666 51984 : DO i = LBOUND(features%feature_index, 1), UBOUND(features%feature_index, 1)
2667 41472 : local_row = local_row + 1
2668 124206 : DO feature_pos = features%local_feature_offsets(local_row), &
2669 44928 : features%local_feature_offsets(local_row + 1) - 1
2670 82734 : row = features%local_feature_rows(feature_pos)
2671 82734 : grid_tmp = grid_tmp + grid_weight_grad(row)*features%grid_weights(row)
2672 : atomic_tmp = atomic_tmp + &
2673 124206 : atomic_grid_weight_grad(row)*features%atomic_grid_weights(row)
2674 : END DO
2675 : END DO
2676 : END DO
2677 : END DO
2678 6 : CPASSERT(local_row == features%nflat_local)
2679 6 : exc_tmp = 0.0_dp
2680 6 : IF (root_rank) exc_tmp = -exc
2681 6 : tmp = grid_tmp + atomic_tmp + exc_tmp
2682 :
2683 6 : IF (my_print_components .AND. root_rank) THEN
2684 0 : iw = cp_logger_get_default_io_unit()
2685 0 : IF (iw > 0) THEN
2686 0 : WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight grid", grid_tmp
2687 0 : WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight atomic", atomic_tmp
2688 0 : WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight final", exc_tmp
2689 0 : WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight residual", tmp
2690 : END IF
2691 : END IF
2692 :
2693 24 : DO idir = 1, 3
2694 24 : virial_xc(idir, idir) = virial_xc(idir, idir) + tmp
2695 : END DO
2696 :
2697 6 : CALL torch_tensor_release(grid_weight_grad_t)
2698 6 : CALL torch_tensor_release(atomic_grid_weight_grad_t)
2699 :
2700 6 : END SUBROUTINE build_weight_virial
2701 :
2702 : ! **************************************************************************************************
2703 : !> \brief Fill CP2K VXC real-space arrays from Torch feature gradients.
2704 : !> \param vxc_rho ...
2705 : !> \param vxc_tau ...
2706 : !> \param rho_r ...
2707 : !> \param pw_pool ...
2708 : !> \param density_grad ...
2709 : !> \param grad_grad ...
2710 : !> \param kin_grad ...
2711 : !> \param xc_deriv_method_id ...
2712 : !> \param global_grid_layout ...
2713 : ! **************************************************************************************************
2714 300 : SUBROUTINE build_vxc_from_feature_grads(vxc_rho, vxc_tau, rho_r, pw_pool, &
2715 300 : density_grad, grad_grad, kin_grad, &
2716 : xc_deriv_method_id, global_grid_layout)
2717 : TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: vxc_rho, vxc_tau, rho_r
2718 : TYPE(pw_pool_type), POINTER :: pw_pool
2719 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: density_grad
2720 : REAL(KIND=dp), DIMENSION(:, :, :), INTENT(IN) :: grad_grad
2721 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: kin_grad
2722 : INTEGER, INTENT(IN) :: xc_deriv_method_id
2723 : LOGICAL, INTENT(IN), OPTIONAL :: global_grid_layout
2724 :
2725 : INTEGER :: i, ipt, ispin, j, k, nspins
2726 : INTEGER, DIMENSION(2, 3) :: bo
2727 : LOGICAL :: my_global_grid_layout, valid_spin_shape
2728 : REAL(KIND=dp) :: dvol_inv
2729 : TYPE(pw_c1d_gs_type) :: tmp_g, vxc_g
2730 1200 : TYPE(pw_r3d_rs_type), DIMENSION(3) :: grad_pw
2731 :
2732 300 : nspins = SIZE(rho_r)
2733 3000 : bo = rho_r(1)%pw_grid%bounds_local
2734 300 : dvol_inv = 1.0_dp/rho_r(1)%pw_grid%dvol
2735 300 : my_global_grid_layout = .FALSE.
2736 300 : IF (PRESENT(global_grid_layout)) my_global_grid_layout = global_grid_layout
2737 252 : IF (my_global_grid_layout) THEN
2738 1008 : CPASSERT(SIZE(density_grad, 1) == PRODUCT(rho_r(1)%pw_grid%npts))
2739 : END IF
2740 300 : CPASSERT(SIZE(kin_grad, 2) == SIZE(density_grad, 2))
2741 300 : CPASSERT(SIZE(grad_grad, 3) == SIZE(density_grad, 2))
2742 300 : CPASSERT(SIZE(grad_grad, 2) == 3)
2743 : valid_spin_shape = SIZE(density_grad, 2) == nspins .OR. &
2744 300 : (nspins == 1 .AND. SIZE(density_grad, 2) == 2)
2745 0 : CPASSERT(valid_spin_shape)
2746 :
2747 1896 : ALLOCATE (vxc_rho(nspins), vxc_tau(nspins))
2748 648 : DO ispin = 1, nspins
2749 348 : CALL pw_pool%create_pw(vxc_rho(ispin))
2750 348 : CALL pw_pool%create_pw(vxc_tau(ispin))
2751 348 : CALL pw_zero(vxc_rho(ispin))
2752 648 : CALL pw_zero(vxc_tau(ispin))
2753 : END DO
2754 :
2755 300 : IF (xc_requires_tmp_g(xc_deriv_method_id) .OR. rho_r(1)%pw_grid%spherical) THEN
2756 300 : CALL pw_pool%create_pw(vxc_g)
2757 300 : IF (.NOT. rho_r(1)%pw_grid%spherical) CALL pw_pool%create_pw(tmp_g)
2758 : END IF
2759 :
2760 648 : DO ispin = 1, nspins
2761 1392 : DO i = 1, 3
2762 1044 : CALL pw_pool%create_pw(grad_pw(i))
2763 1392 : CALL pw_zero(grad_pw(i))
2764 : END DO
2765 :
2766 348 : ipt = 0
2767 7448 : DO k = bo(1, 3), bo(2, 3)
2768 176620 : DO j = bo(1, 2), bo(2, 2)
2769 2579354 : DO i = bo(1, 1), bo(2, 1)
2770 2403082 : IF (my_global_grid_layout) THEN
2771 : ipt = 1 + i - rho_r(1)%pw_grid%bounds(1, 1) + &
2772 : rho_r(1)%pw_grid%npts(1)*( &
2773 : j - rho_r(1)%pw_grid%bounds(1, 2) + &
2774 : rho_r(1)%pw_grid%npts(2)*( &
2775 1163283 : k - rho_r(1)%pw_grid%bounds(1, 3)))
2776 : ELSE
2777 1239799 : ipt = ipt + 1
2778 : END IF
2779 2572254 : IF (nspins == 1) THEN
2780 1714582 : IF (SIZE(density_grad, 2) == 1) THEN
2781 1021533 : vxc_rho(1)%array(i, j, k) = 0.5_dp*dvol_inv*density_grad(ipt, 1)
2782 1021533 : vxc_tau(1)%array(i, j, k) = 0.5_dp*dvol_inv*kin_grad(ipt, 1)
2783 1021533 : grad_pw(1)%array(i, j, k) = 0.5_dp*dvol_inv*grad_grad(ipt, 1, 1)
2784 1021533 : grad_pw(2)%array(i, j, k) = 0.5_dp*dvol_inv*grad_grad(ipt, 2, 1)
2785 1021533 : grad_pw(3)%array(i, j, k) = 0.5_dp*dvol_inv*grad_grad(ipt, 3, 1)
2786 : ELSE
2787 : vxc_rho(1)%array(i, j, k) = 0.5_dp*dvol_inv* &
2788 693049 : (density_grad(ipt, 1) + density_grad(ipt, 2))
2789 : vxc_tau(1)%array(i, j, k) = 0.5_dp*dvol_inv* &
2790 693049 : (kin_grad(ipt, 1) + kin_grad(ipt, 2))
2791 : grad_pw(1)%array(i, j, k) = 0.5_dp*dvol_inv* &
2792 693049 : (grad_grad(ipt, 1, 1) + grad_grad(ipt, 1, 2))
2793 : grad_pw(2)%array(i, j, k) = 0.5_dp*dvol_inv* &
2794 693049 : (grad_grad(ipt, 2, 1) + grad_grad(ipt, 2, 2))
2795 : grad_pw(3)%array(i, j, k) = 0.5_dp*dvol_inv* &
2796 693049 : (grad_grad(ipt, 3, 1) + grad_grad(ipt, 3, 2))
2797 : END IF
2798 : ELSE
2799 688500 : vxc_rho(ispin)%array(i, j, k) = dvol_inv*density_grad(ipt, ispin)
2800 688500 : vxc_tau(ispin)%array(i, j, k) = dvol_inv*kin_grad(ipt, ispin)
2801 688500 : grad_pw(1)%array(i, j, k) = dvol_inv*grad_grad(ipt, 1, ispin)
2802 688500 : grad_pw(2)%array(i, j, k) = dvol_inv*grad_grad(ipt, 2, ispin)
2803 688500 : grad_pw(3)%array(i, j, k) = dvol_inv*grad_grad(ipt, 3, ispin)
2804 : END IF
2805 : END DO
2806 : END DO
2807 : END DO
2808 :
2809 1392 : DO i = 1, 3
2810 1392 : CALL pw_scale(grad_pw(i), -1.0_dp)
2811 : END DO
2812 348 : CALL xc_pw_divergence(xc_deriv_method_id, grad_pw, tmp_g, vxc_g, vxc_rho(ispin))
2813 :
2814 1692 : DO i = 1, 3
2815 1392 : CALL pw_pool%give_back_pw(grad_pw(i))
2816 : END DO
2817 : END DO
2818 :
2819 300 : IF (ASSOCIATED(vxc_g%pw_grid)) CALL pw_pool%give_back_pw(vxc_g)
2820 300 : IF (ASSOCIATED(tmp_g%pw_grid)) CALL pw_pool%give_back_pw(tmp_g)
2821 :
2822 300 : END SUBROUTINE build_vxc_from_feature_grads
2823 :
2824 : ! **************************************************************************************************
2825 : !> \brief Print optional diagnostics for the CP2K-native SKALA GPW feature block.
2826 : !> \param features ...
2827 : !> \param print_active ...
2828 : ! **************************************************************************************************
2829 24 : SUBROUTINE print_native_grid_diagnostics(features, print_active)
2830 : TYPE(skala_gpw_feature_type), INTENT(IN) :: features
2831 : LOGICAL, INTENT(IN) :: print_active
2832 :
2833 : INTEGER :: atom_rows_max, atom_rows_min, &
2834 : chunk_rows_max, chunk_rows_min, iw
2835 : REAL(KIND=dp) :: chunk_imbalance
2836 :
2837 24 : IF (.NOT. print_active) RETURN
2838 :
2839 12 : iw = cp_logger_get_default_io_unit()
2840 12 : IF (iw <= 0) RETURN
2841 : WRITE (UNIT=iw, FMT="(/,T2,A,1X,ES19.11)") &
2842 12 : "SKALA_GPW| Native grid feature electrons", features%electron_count
2843 : WRITE (UNIT=iw, FMT="(T2,A,1X,ES19.11)") &
2844 12 : "SKALA_GPW| Native grid feature spin moment", features%spin_moment
2845 : WRITE (UNIT=iw, FMT="(T2,A,1X,ES19.11)") &
2846 12 : "SKALA_GPW| Native grid feature kinetic integral", features%kinetic_integral
2847 : WRITE (UNIT=iw, FMT="(T2,A,1X,ES19.11)") &
2848 12 : "SKALA_GPW| Native grid feature weight sum", features%grid_weight_sum
2849 12 : IF (ALLOCATED(features%atomic_grid_sizes)) THEN
2850 49 : atom_rows_min = INT(MINVAL(features%atomic_grid_sizes))
2851 49 : atom_rows_max = INT(MAXVAL(features%atomic_grid_sizes))
2852 : WRITE (UNIT=iw, FMT="(T2,A,1X,I0,1X,A,1X,I0,1X,A,1X,I0)") &
2853 12 : "SKALA_GPW| Native grid atom row range", atom_rows_min, "to", &
2854 61 : atom_rows_max, "sum", INT(SUM(features%atomic_grid_sizes))
2855 : END IF
2856 12 : IF (features%uses_atom_chunks) THEN
2857 : WRITE (UNIT=iw, FMT="(T2,A,1X,I0,1X,A,1X,I0)") &
2858 12 : "SKALA_GPW| Native grid atom chunk rows", features%chunk_feature_count, &
2859 24 : "of", features%nflat
2860 12 : IF (ALLOCATED(features%chunk_grad_counts)) THEN
2861 36 : chunk_rows_min = MINVAL(features%chunk_grad_counts)/ngrad_per_point
2862 36 : chunk_rows_max = MAXVAL(features%chunk_grad_counts)/ngrad_per_point
2863 12 : chunk_imbalance = REAL(chunk_rows_max, KIND=dp)/REAL(MAX(1, chunk_rows_min), KIND=dp)
2864 : WRITE (UNIT=iw, FMT="(T2,A,1X,I0,1X,A,1X,I0,1X,A,1X,ES12.5)") &
2865 12 : "SKALA_GPW| Native grid atom chunk row range", chunk_rows_min, &
2866 24 : "to", chunk_rows_max, "imbalance", chunk_imbalance
2867 : END IF
2868 : END IF
2869 :
2870 : END SUBROUTINE print_native_grid_diagnostics
2871 :
2872 : ! **************************************************************************************************
2873 : !> \brief Configure CUDA device selection for the native SKALA GPW Torch path.
2874 : !> \param use_cuda ...
2875 : !> \param requested_device ...
2876 : !> \param group ...
2877 : !> \return selected CUDA device, or -1 for CPU fallback/no visible CUDA device
2878 : ! **************************************************************************************************
2879 304 : FUNCTION configure_native_grid_cuda(use_cuda, requested_device, group) RESULT(selected_device)
2880 : LOGICAL, INTENT(IN) :: use_cuda
2881 : INTEGER, INTENT(IN) :: requested_device
2882 :
2883 : CLASS(mp_comm_type), INTENT(IN) :: group
2884 :
2885 : INTEGER :: cuda_device_count, iw, pe, selected_device
2886 304 : INTEGER, ALLOCATABLE, DIMENSION(:) :: selected_devices
2887 :
2888 304 : selected_device = -1
2889 :
2890 304 : IF (.NOT. use_cuda) RETURN
2891 :
2892 0 : IF (.NOT. torch_cuda_is_available()) THEN
2893 0 : cuda_device_count = 0
2894 : ELSE
2895 0 : cuda_device_count = torch_cuda_device_count()
2896 : END IF
2897 0 : IF (cuda_device_count > 0) THEN
2898 0 : IF (requested_device < 0) THEN
2899 0 : selected_device = MOD(group%mepos, cuda_device_count)
2900 : ELSE
2901 0 : selected_device = requested_device
2902 : END IF
2903 : END IF
2904 0 : IF (selected_device >= cuda_device_count) THEN
2905 : CALL cp_abort(__LOCATION__, &
2906 : "GAUXC%NATIVE_GRID_CUDA_DEVICE selects a CUDA device outside the visible "// &
2907 0 : "Torch CUDA device range.")
2908 : END IF
2909 0 : IF (selected_device >= 0) CALL offload_set_chosen_device(selected_device)
2910 :
2911 0 : ALLOCATE (selected_devices(group%num_pe))
2912 0 : CALL group%allgather(selected_device, selected_devices)
2913 :
2914 0 : IF (group%mepos /= 0) THEN
2915 0 : DEALLOCATE (selected_devices)
2916 0 : RETURN
2917 : END IF
2918 : IF (selected_device == logged_cuda_device .AND. &
2919 : cuda_device_count == logged_cuda_device_count .AND. &
2920 0 : group%num_pe == logged_cuda_nproc .AND. &
2921 : requested_device == logged_cuda_request) THEN
2922 0 : DEALLOCATE (selected_devices)
2923 0 : RETURN
2924 : END IF
2925 :
2926 0 : iw = cp_logger_get_default_io_unit()
2927 0 : IF (iw <= 0) THEN
2928 0 : DEALLOCATE (selected_devices)
2929 0 : RETURN
2930 : END IF
2931 0 : IF (selected_device >= 0) THEN
2932 : WRITE (UNIT=iw, FMT="(/,T2,A,1X,I0,1X,A,1X,I0,1X,A,1X,I0)") &
2933 0 : "SKALA_GPW| Native grid Torch CUDA device", selected_device, &
2934 0 : "of", cuda_device_count, "requested", requested_device
2935 : ELSE
2936 : WRITE (UNIT=iw, FMT="(/,T2,A)") &
2937 0 : "SKALA_GPW| Native grid Torch CUDA requested, but no Torch CUDA device is visible"
2938 : END IF
2939 : WRITE (UNIT=iw, FMT="(T2,A)", ADVANCE="NO") &
2940 0 : "SKALA_GPW| Native grid Torch CUDA rank devices"
2941 0 : DO pe = 1, group%num_pe
2942 0 : WRITE (UNIT=iw, FMT="(1X,I0,A,I0)", ADVANCE="NO") pe - 1, ":", selected_devices(pe)
2943 : END DO
2944 0 : WRITE (UNIT=iw, FMT=*)
2945 :
2946 0 : logged_cuda_device = selected_device
2947 0 : logged_cuda_device_count = cuda_device_count
2948 0 : logged_cuda_nproc = group%num_pe
2949 0 : logged_cuda_request = requested_device
2950 0 : DEALLOCATE (selected_devices)
2951 :
2952 304 : END FUNCTION configure_native_grid_cuda
2953 :
2954 : ! **************************************************************************************************
2955 : !> \brief Load and cache the TorchScript SKALA model.
2956 : !> \param model_path ...
2957 : !> \param cuda_device ...
2958 : !> \param xc_section XC settings, including the native model precision.
2959 : ! **************************************************************************************************
2960 564 : SUBROUTINE ensure_model_loaded(model_path, cuda_device, xc_section)
2961 : CHARACTER(len=*), INTENT(IN) :: model_path
2962 : INTEGER, INTENT(IN) :: cuda_device
2963 : TYPE(section_vals_type), POINTER :: xc_section
2964 :
2965 : LOGICAL :: fp64
2966 : TYPE(section_vals_type), POINTER :: gauxc_section
2967 :
2968 564 : gauxc_section => get_gauxc_section(xc_section)
2969 282 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_MODEL_FP64", l_val=fp64)
2970 :
2971 282 : IF (cached_model_loaded) THEN
2972 : IF (TRIM(cached_model_path) == TRIM(model_path) .AND. &
2973 265 : cached_model_cuda_device == cuda_device .AND. &
2974 265 : (cached_model_fp64 .EQV. fp64)) RETURN
2975 0 : CALL skala_torch_model_release(cached_model)
2976 0 : cached_model_loaded = .FALSE.
2977 : END IF
2978 :
2979 17 : CALL skala_torch_model_load(cached_model, TRIM(model_path), fp64=fp64)
2980 17 : cached_model_path = model_path
2981 17 : cached_model_cuda_device = cuda_device
2982 17 : cached_model_fp64 = fp64
2983 17 : cached_model_loaded = .TRUE.
2984 :
2985 282 : END SUBROUTINE ensure_model_loaded
2986 :
2987 : ! **************************************************************************************************
2988 : !> \brief Resolve the SKALA TorchScript model path from the GAUXC subsection.
2989 : !> \param xc_section ...
2990 : !> \param model_path ...
2991 : ! **************************************************************************************************
2992 304 : SUBROUTINE get_skala_model_path(xc_section, model_path)
2993 : TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
2994 : CHARACTER(len=default_path_length), INTENT(OUT) :: model_path
2995 :
2996 : CHARACTER(len=default_path_length) :: model_key
2997 : INTEGER :: env_status
2998 : LOGICAL :: native_grid_use_cuda
2999 : TYPE(section_vals_type), POINTER :: gauxc_section
3000 :
3001 304 : gauxc_section => get_gauxc_section(xc_section)
3002 304 : IF (.NOT. ASSOCIATED(gauxc_section)) THEN
3003 0 : CPABORT("Native SKALA GPW requires an XC_FUNCTIONAL%GAUXC section")
3004 : END IF
3005 :
3006 304 : CALL section_vals_val_get(gauxc_section, "MODEL", c_val=model_path)
3007 304 : model_key = ADJUSTL(model_path)
3008 304 : CALL uppercase(model_key)
3009 304 : IF (TRIM(model_key) == "NONE" .OR. TRIM(model_key) == "") THEN
3010 0 : CPABORT("Native SKALA GPW requires GAUXC%MODEL SKALA or a TorchScript model path")
3011 304 : ELSE IF (TRIM(model_key) == "SKALA") THEN
3012 304 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_USE_CUDA", l_val=native_grid_use_cuda)
3013 304 : IF (native_grid_use_cuda) THEN
3014 0 : CALL GET_ENVIRONMENT_VARIABLE("GAUXC_SKALA_CUDA_MODEL", model_path, STATUS=env_status)
3015 0 : IF (env_status == 0 .AND. LEN_TRIM(model_path) > 0) RETURN
3016 : END IF
3017 304 : CALL GET_ENVIRONMENT_VARIABLE("SKALA_MODEL", model_path, STATUS=env_status)
3018 304 : IF (env_status /= 0 .OR. LEN_TRIM(model_path) == 0) THEN
3019 0 : IF (native_grid_use_cuda) THEN
3020 : CALL cp_abort(__LOCATION__, &
3021 0 : "MODEL SKALA CUDA path requires GAUXC_SKALA_CUDA_MODEL or SKALA_MODEL")
3022 : ELSE
3023 : CALL cp_abort(__LOCATION__, &
3024 0 : "MODEL SKALA requires the SKALA_MODEL environment variable")
3025 : END IF
3026 : END IF
3027 : END IF
3028 :
3029 : END SUBROUTINE get_skala_model_path
3030 :
3031 : ! **************************************************************************************************
3032 : !> \brief Return the first GAUXC functional subsection, if present.
3033 : !> \param xc_section ...
3034 : !> \return ...
3035 : ! **************************************************************************************************
3036 735989 : FUNCTION get_gauxc_section(xc_section) RESULT(gauxc_section)
3037 : TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
3038 : TYPE(section_vals_type), POINTER :: gauxc_section
3039 :
3040 : INTEGER :: ifun
3041 : TYPE(section_vals_type), POINTER :: functionals, xc_fun
3042 :
3043 735989 : NULLIFY (gauxc_section)
3044 735989 : IF (.NOT. ASSOCIATED(xc_section)) RETURN
3045 :
3046 735989 : functionals => section_vals_get_subs_vals(xc_section, "XC_FUNCTIONAL")
3047 735989 : IF (.NOT. ASSOCIATED(functionals)) RETURN
3048 :
3049 735989 : ifun = 0
3050 : DO
3051 1424852 : ifun = ifun + 1
3052 1424852 : xc_fun => section_vals_get_subs_vals2(functionals, i_section=ifun)
3053 1424852 : IF (.NOT. ASSOCIATED(xc_fun)) EXIT
3054 1424852 : IF (xc_fun%section%name == "GAUXC") THEN
3055 : gauxc_section => xc_fun
3056 : EXIT
3057 : END IF
3058 : END DO
3059 :
3060 : END FUNCTION get_gauxc_section
3061 :
3062 0 : END MODULE skala_gpw_functional
|
