(git:6344619)
Loading...
Searching...
No Matches
skala_gpw_functional.F
Go to the documentation of this file.
1!--------------------------------------------------------------------------------------------------!
2! CP2K: A general program to perform molecular dynamics simulations !
3! Copyright 2000-2026 CP2K developers group <https://cp2k.org> !
4! !
5! SPDX-License-Identifier: GPL-2.0-or-later !
6!--------------------------------------------------------------------------------------------------!
7
8! **************************************************************************************************
9!> \brief Experimental CP2K-native GPW real-space-grid path for SKALA TorchScript models.
10! **************************************************************************************************
11MODULE 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
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
109CONTAINS
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 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 uses_native_grid = .false.
123 gauxc_section => get_gauxc_section(xc_section)
124 IF (ASSOCIATED(gauxc_section)) THEN
125 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID", l_val=uses_native_grid)
126 END IF
127
128 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 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 IF (.NOT. xc_section_uses_gauxc_model(xc_section)) RETURN
143
144 gauxc_section => get_gauxc_section(xc_section)
145 cpassert(ASSOCIATED(gauxc_section))
146 CALL section_vals_val_get(gauxc_section, "PSEUDOPOTENTIAL_GAPW_REPRESENTATION", &
147 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 CONTINUE
153 CASE DEFAULT
154 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 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 skala_gapw_paw_one_center
171
172 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 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 use_composite_reference = .false.
187 native_grid = .false.
188 gauxc_section => get_gauxc_section(xc_section)
189 IF (ASSOCIATED(gauxc_section)) THEN
190 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 l_val=use_composite_reference)
193 END IF
194 use_composite_reference = native_grid .AND. use_composite_reference
195
196 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 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 use_atom_composite = .false.
211 native_grid = .false.
212 gauxc_section => get_gauxc_section(xc_section)
213 IF (ASSOCIATED(gauxc_section)) THEN
214 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 l_val=use_atom_composite)
218 END IF
219 use_atom_composite = native_grid .AND. use_atom_composite
220
221 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 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 use_atom_composite = .false.
236 composite_grid = skala_gapw_atom_composite_grid
237 gauxc_section => get_gauxc_section(xc_section)
238 IF (ASSOCIATED(gauxc_section)) THEN
239 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_LAYOUT", &
240 i_val=composite_grid)
241 END IF
242 SELECT CASE (composite_grid)
243 CASE (skala_gapw_common_grid)
244 CONTINUE
245 CASE (skala_gapw_atom_composite_grid)
246 use_atom_composite = .true.
247 CASE DEFAULT
248 CALL cp_abort(__location__, "Unknown native-grid layout.")
249 END SELECT
250
251 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 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 use_direct_ao = .false.
265 gauxc_section => get_gauxc_section(xc_section)
266 IF (ASSOCIATED(gauxc_section)) THEN
267 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_GAPW_COMPOSITE_DIRECT_AO", &
268 l_val=use_direct_ao)
269 END IF
270 use_direct_ao = native_skala_gapw_composite_reference(xc_section) .AND. use_direct_ao
271
272 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 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 uses_gauxc_model = .false.
287 gauxc_section => get_gauxc_section(xc_section)
288 IF (ASSOCIATED(gauxc_section)) THEN
289 CALL section_vals_val_get(gauxc_section, "MODEL", c_val=model_name)
290 CALL section_vals_val_get(gauxc_section, "FUNCTIONAL", c_val=xc_name)
291 model_key = adjustl(model_name)
292 xc_key = adjustl(xc_name)
293 CALL uppercase(model_key)
294 CALL uppercase(xc_key)
295 uses_gauxc_model = (trim(model_key) /= "" .AND. trim(model_key) /= "NONE" .AND. &
296 trim(model_key) /= trim(xc_key))
297 END IF
298
299 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 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 partition = skala_gapw_density_partition_hard_minus_soft
313 gauxc_section => get_gauxc_section(xc_section)
314 IF (ASSOCIATED(gauxc_section)) THEN
315 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_GAPW_DENSITY_PARTITION", &
316 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 CONTINUE
325 CASE DEFAULT
326 CALL cp_abort(__location__, &
327 "Unknown GAUXC%NATIVE_GRID_GAPW_DENSITY_PARTITION value.")
328 END SELECT
329
330 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 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 NULLIFY (gauxc_section)
345 IF (.NOT. ASSOCIATED(xc_section)) THEN
346 cpabort("Native SKALA GPW requires an XC section")
347 END IF
348
349 functionals => section_vals_get_subs_vals(xc_section, "XC_FUNCTIONAL")
350 IF (.NOT. ASSOCIATED(functionals)) THEN
351 cpabort("Native SKALA GPW requires an XC_FUNCTIONAL section")
352 END IF
353
354 nfun = 0
355 ifun = 0
356 DO
357 ifun = ifun + 1
358 xc_fun => section_vals_get_subs_vals2(functionals, i_section=ifun)
359 IF (.NOT. ASSOCIATED(xc_fun)) EXIT
360 nfun = nfun + 1
361 IF (xc_fun%section%name == "GAUXC") gauxc_section => xc_fun
362 END DO
363
364 IF (.NOT. ASSOCIATED(gauxc_section)) THEN
365 cpabort("Native SKALA GPW requires an XC_FUNCTIONAL%GAUXC section")
366 END IF
367 IF (nfun /= 1) THEN
368 cpabort("Native SKALA GPW requires GAUXC to be the only XC functional")
369 END IF
370
371 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID", l_val=native_grid)
372 IF (.NOT. native_grid .AND. &
373 .NOT. xc_section_uses_native_skala_evaluator(xc_section)) RETURN
374
375 CALL section_vals_val_get(gauxc_section, "MODEL", c_val=model_name)
376 model_key = adjustl(model_name)
377 CALL uppercase(model_key)
378 IF (trim(model_key) == "NONE" .OR. trim(model_key) == "") THEN
379 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 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 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 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: density_grad, kin_grad
432 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 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 virial_xc = 0.0_dp
443 exc = 0.0_dp
444 my_just_energy = .false.
445 IF (PRESENT(just_energy)) my_just_energy = just_energy
446 needs_atom_force = PRESENT(atom_force)
447 IF (needs_atom_force) atom_force = 0.0_dp
448 have_atom_coord_grad = .false.
449 omp_max_threads_restore = omp_get_max_threads()
450
451 IF (compute_virial .AND. my_just_energy) THEN
452 CALL cp_abort(__location__, &
453 "Native SKALA GPW stress/virial requires feature gradients.")
454 END IF
455 IF (.NOT. ASSOCIATED(rho_g)) THEN
456 CALL cp_abort(__location__, &
457 "Native SKALA GPW requires the reciprocal-space density to form density gradients.")
458 END IF
459 IF (.NOT. ASSOCIATED(tau)) THEN
460 CALL cp_abort(__location__, &
461 "Native SKALA GPW requires the kinetic-energy density.")
462 END IF
463
464 nspins = SIZE(rho_r)
465 lsd = (nspins /= 1)
466 CALL get_skala_model_path(xc_section, model_path)
467 gauxc_section => get_gauxc_section(xc_section)
468 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 i_val=native_grid_cuda_device)
471 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNKS", &
472 l_val=native_grid_atom_chunks)
473 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNK_ROUTING", &
474 l_val=native_grid_atom_chunk_routing)
475 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNK_MAX_ROWS", &
476 i_val=native_grid_atom_chunk_max_rows)
477 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_PARTITION", &
478 i_val=native_grid_atom_partition)
479 SELECT CASE (native_grid_atom_partition)
480 CASE (1)
481 native_grid_atom_partition = skala_gpw_atom_partition_hard
482 CASE (2)
483 native_grid_atom_partition = skala_gpw_atom_partition_smooth
484 CASE DEFAULT
485 CALL cp_abort(__location__, &
486 "Unknown GAUXC%NATIVE_GRID_ATOM_PARTITION value.")
487 END SELECT
488 native_grid_atom_chunk_routing = native_grid_atom_chunk_routing .OR. native_grid_atom_chunks
489 native_grid_atom_chunks = native_grid_atom_chunks .OR. native_grid_atom_chunk_routing
490 IF (native_grid_atom_chunk_max_rows < -1) THEN
491 CALL cp_abort(__location__, &
492 "GAUXC%NATIVE_GRID_ATOM_CHUNK_MAX_ROWS must be -1, zero, or positive.")
493 END IF
494 IF (needs_atom_force .OR. compute_virial) THEN
495 IF (native_grid_atom_partition == skala_gpw_atom_partition_hard) THEN
496 native_grid_atom_partition = skala_gpw_atom_partition_smooth
497 END IF
498 native_grid_atom_chunk_routing = .false.
499 native_grid_atom_chunks = .false.
500 END IF
501 ! The portable SKALA export used by the regtests builds ragged-index tensors on CPU.
502 CALL torch_use_cuda(native_grid_use_cuda)
503 selected_cuda_device = configure_native_grid_cuda( &
504 native_grid_use_cuda, native_grid_cuda_device, rho_r(1)%pw_grid%para%group)
505 CALL ensure_model_loaded(model_path, selected_cuda_device, xc_section)
506
507 IF (lsd) THEN
508 needs%rho_spin = .true.
509 needs%drho_spin = .true.
510 needs%tau_spin = .true.
511 ELSE
512 needs%rho = .true.
513 needs%drho = .true.
514 needs%tau = .true.
515 END IF
516
517 CALL section_vals_val_get(xc_section, "XC_GRID%XC_DERIV", i_val=xc_deriv_method_id)
518 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 tau_cutoff=section_get_rval(xc_section, "tau_cutoff"))
525 CALL xc_rho_set_update(rho_set, rho_r, rho_g, tau, needs, &
526 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 atom_partition=native_grid_atom_partition)
535 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_DIAGNOSTICS", l_val=native_grid_diagnostics)
536 IF (native_grid_diagnostics) THEN
537 CALL print_native_grid_diagnostics(features, rho_r(1)%pw_grid%para%group%mepos == 0)
538 END IF
539
540 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 rho_r(1)%pw_grid%para%group)
543 END IF
544 IF (native_grid_diagnostics .AND. features%uses_atom_chunks .AND. &
545 rho_r(1)%pw_grid%para%group%mepos == 0) THEN
546 iw = cp_logger_get_default_io_unit()
547 IF (iw > 0) THEN
548 WRITE (unit=iw, fmt="(T2,A,1X,I0)") &
549 "SKALA_GPW| Native grid atom chunk max rows", native_grid_atom_chunk_max_rows
550 END IF
551 END IF
552 native_grid_atom_subchunks = 1
553 IF (features%uses_atom_chunks .AND. native_grid_atom_chunk_max_rows > 0) THEN
554 native_grid_atom_subchunks = skala_gpw_atom_subchunk_count(native_grid_atom_chunk_max_rows)
555 CALL rho_r(1)%pw_grid%para%group%max(native_grid_atom_subchunks)
556 END IF
557 use_atom_subchunks = features%uses_atom_chunks .AND. native_grid_atom_subchunks > 1
558 has_atom_chunk_work = .NOT. features%uses_atom_chunks .OR. features%chunk_feature_count > 0
559 exc = 0.0_dp
560 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 kin_grad=kin_grad, collapse_spin_grads=(nspins == 1))
566 ELSE IF (has_atom_chunk_work) THEN
567 CALL skala_torch_model_get_exc(cached_model, features%inputs, &
568 features%grid_weights_t, exc_tensor, exc)
569 END IF
570 IF (features%uses_atom_chunks) CALL rho_r(1)%pw_grid%para%group%sum(exc)
571
572 IF (.NOT. my_just_energy) THEN
573 IF (.NOT. use_atom_subchunks) THEN
574 IF (has_atom_chunk_work) THEN
575 CALL timeset("skala_gpw_backward", phase_handle)
576 CALL torch_tensor_backward_scalar(exc_tensor)
577 CALL timestop(phase_handle)
578
579 IF (compute_virial) THEN
580 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 native_grid_diagnostics)
585 IF (native_grid_diagnostics) THEN
586 CALL print_virial_delta("weight-residual", virial_xc - virial_before, &
587 rho_r(1)%pw_grid%para%group%mepos == 0)
588 END IF
589 END IF
590 END IF
591
592 CALL timeset("skala_gpw_grad_fetch", phase_handle)
593 IF (features%uses_atom_chunks) THEN
594 CALL fetch_and_gather_atom_chunk_grads(features, rho_r(1)%pw_grid%para%group, &
595 density_grad, grad_grad, kin_grad)
596 ELSE
597 CALL fetch_local_feature_grads(features, density_grad, grad_grad, kin_grad)
598 END IF
599 CALL timestop(phase_handle)
600 END IF
601 IF (needs_atom_force) THEN
602 CALL add_explicit_coordinate_force(atom_force, features, atom_coord_grad_t, &
603 rho_r(1)%pw_grid%para%group%mepos == 0)
604 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 grid_weight_grad_t, atomic_grid_weight_grad_t)
607 END IF
608 have_atom_coord_grad = .true.
609 END IF
610
611 CALL timeset("skala_gpw_vxc_unpack", phase_handle)
612 IF (native_grid_diagnostics) THEN
613 bo = rho_r(1)%pw_grid%bounds_local
614 density_contraction = 0.0_dp
615 tau_contraction = 0.0_dp
616 ipt = 0
617 DO k = bo(1, 3), bo(2, 3)
618 DO j = bo(1, 2), bo(2, 2)
619 DO i = bo(1, 1), bo(2, 1)
620 ipt = ipt + 1
621 IF (nspins == 1) THEN
622 density_contraction = density_contraction + rho_r(1)%array(i, j, k)* &
623 0.5_dp*(density_grad(ipt, 1) + density_grad(ipt, 2))
624 tau_contraction = tau_contraction + tau(1)%array(i, j, k)* &
625 0.5_dp*(kin_grad(ipt, 1) + kin_grad(ipt, 2))
626 ELSE
627 DO ispin = 1, nspins
628 density_contraction = density_contraction + &
629 rho_r(ispin)%array(i, j, k)*density_grad(ipt, ispin)
630 tau_contraction = tau_contraction + &
631 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 CALL rho_r(1)%pw_grid%para%group%sum(density_contraction)
638 CALL rho_r(1)%pw_grid%para%group%sum(tau_contraction)
639 IF (rho_r(1)%pw_grid%para%group%mepos == 0) THEN
640 iw = cp_logger_get_default_io_unit()
641 IF (iw > 0) THEN
642 WRITE (iw, "(T2,A,1X,ES20.10)") &
643 "SKALA_GPW| XC density-gradient contraction", density_contraction
644 WRITE (iw, "(T2,A,1X,ES20.10)") &
645 "SKALA_GPW| XC kinetic-gradient contraction", tau_contraction
646 END IF
647 END IF
648 END IF
649 IF (compute_virial) THEN
650 IF (native_grid_diagnostics) virial_before = virial_xc
651 CALL build_virial_from_feature_grads(virial_xc, rho_set, rho_r, grad_grad)
652 IF (native_grid_diagnostics) THEN
653 CALL print_virial_delta("feature-gradient", virial_xc - virial_before, &
654 rho_r(1)%pw_grid%para%group%mepos == 0)
655 virial_before = virial_xc
656 END IF
657 IF (.NOT. have_atom_coord_grad) THEN
658 CALL torch_tensor_grad(features%coarse_0_atomic_coords_t, atom_coord_grad_t)
659 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 native_grid_diagnostics)
665 IF (native_grid_diagnostics) THEN
666 CALL print_virial_delta("static-coordinates", virial_xc - virial_before, &
667 rho_r(1)%pw_grid%para%group%mepos == 0)
668 virial_before = virial_xc
669 END IF
670 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 grid_weight_grad_t, atomic_grid_weight_grad_t)
673 IF (native_grid_diagnostics) THEN
674 CALL print_virial_delta("smooth-partition", virial_xc - virial_before, &
675 rho_r(1)%pw_grid%para%group%mepos == 0)
676 virial_before = virial_xc
677 END IF
678 END IF
679 END IF
680 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 xc_deriv_method_id)
684 CALL timestop(phase_handle)
685
686 CALL timeset("skala_gpw_grad_release", phase_handle)
687 DEALLOCATE (density_grad, grad_grad, kin_grad)
688 IF (have_atom_coord_grad) CALL torch_tensor_release(atom_coord_grad_t)
689 CALL timestop(phase_handle)
690 END IF
691
692 CALL timeset("skala_gpw_cleanup", phase_handle)
693 IF (.NOT. use_atom_subchunks .AND. has_atom_chunk_work) CALL torch_tensor_release(exc_tensor)
694 CALL skala_gpw_feature_release(features)
695 CALL xc_rho_set_release(rho_set, pw_pool=pw_pool)
696 CALL torch_use_cuda(.true.)
697 CALL omp_set_num_threads(omp_max_threads_restore)
698 CALL timestop(phase_handle)
699
700 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 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 LOGICAL, ALLOCATABLE, DIMENSION(:) :: included
734 REAL(kind=dp) :: exc, local_derivative
735 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: partition_weights
736 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: atom_coords_pbc
737 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :) :: dweights_datom, dweights_dstrain
738 REAL(kind=dp), DIMENSION(3) :: grid_point
739 REAL(kind=dp), DIMENSION(:), POINTER :: grid_weight_grad
740 TYPE(section_vals_type), POINTER :: gauxc_section
741 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 cpassert(ASSOCIATED(rho_r))
747 cpassert(ASSOCIATED(rho_g))
748 cpassert(ASSOCIATED(tau))
749 CALL pw_zero(weight_deriv_r)
750 omp_max_threads_restore = omp_get_max_threads()
751
752 nspins = SIZE(rho_r)
753 lsd = (nspins /= 1)
754 CALL get_skala_model_path(xc_section, model_path)
755 gauxc_section => get_gauxc_section(xc_section)
756 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 i_val=native_grid_cuda_device)
759 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_PARTITION", &
760 i_val=native_grid_atom_partition)
761 SELECT CASE (native_grid_atom_partition)
762 CASE (1)
763 native_grid_atom_partition = skala_gpw_atom_partition_hard
764 CASE (2)
765 native_grid_atom_partition = skala_gpw_atom_partition_smooth
766 CASE DEFAULT
767 CALL cp_abort(__location__, &
768 "Unknown GAUXC%NATIVE_GRID_ATOM_PARTITION value.")
769 END SELECT
770 IF (native_grid_atom_partition == skala_gpw_atom_partition_hard) THEN
771 native_grid_atom_partition = skala_gpw_atom_partition_smooth
772 END IF
773
774 CALL torch_use_cuda(native_grid_use_cuda)
775 selected_cuda_device = configure_native_grid_cuda( &
776 native_grid_use_cuda, native_grid_cuda_device, rho_r(1)%pw_grid%para%group)
777 CALL ensure_model_loaded(model_path, selected_cuda_device, xc_section)
778
779 IF (lsd) THEN
780 needs%rho_spin = .true.
781 needs%drho_spin = .true.
782 needs%tau_spin = .true.
783 ELSE
784 needs%rho = .true.
785 needs%drho = .true.
786 needs%tau = .true.
787 END IF
788
789 CALL section_vals_val_get(xc_section, "XC_GRID%XC_DERIV", i_val=xc_deriv_method_id)
790 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 tau_cutoff=section_get_rval(xc_section, "tau_cutoff"))
797 CALL xc_rho_set_update(rho_set, rho_r, rho_g, tau, needs, &
798 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 atom_partition=native_grid_atom_partition)
806 CALL skala_torch_model_get_exc(cached_model, features%inputs, features%grid_weights_t, &
807 exc_tensor, exc)
808 CALL torch_tensor_backward_scalar(exc_tensor)
809 NULLIFY (grid_weight_grad)
810 CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
811 CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
812
813 natom = SIZE(particle_set)
814 IF (native_grid_atom_partition == skala_gpw_atom_partition_smooth) THEN
815 ALLOCATE (atom_coords_pbc(3, natom), included(natom), partition_weights(natom), &
816 dweights_datom(3, natom, natom), dweights_dstrain(3, 3, natom))
817 DO iatom = 1, natom
818 atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.true.)
819 END DO
820 END IF
821
822 bo = rho_r(1)%pw_grid%bounds_local
823 local_row = 0
824 DO k = bo(1, 3), bo(2, 3)
825 DO j = bo(1, 2), bo(2, 2)
826 DO i = bo(1, 1), bo(2, 1)
827 local_row = local_row + 1
828 feature_begin = features%local_feature_offsets(local_row)
829 feature_end = features%local_feature_offsets(local_row + 1) - 1
830 local_derivative = 0.0_dp
831 IF (native_grid_atom_partition == skala_gpw_atom_partition_hard) THEN
832 cpassert(feature_end == feature_begin)
833 row = features%local_feature_rows(feature_begin)
834 local_derivative = grid_weight_grad(row)
835 ELSE
836 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 dweights_datom, dweights_dstrain)
840 cpassert(feature_end - feature_begin + 1 == count(included))
841 feature_pos = feature_begin
842 DO iatom = 1, natom
843 IF (.NOT. included(iatom)) cycle
844 row = features%local_feature_rows(feature_pos)
845 local_derivative = local_derivative + &
846 partition_weights(iatom)*grid_weight_grad(row)
847 feature_pos = feature_pos + 1
848 END DO
849 cpassert(feature_pos == feature_end + 1)
850 END IF
851 weight_deriv_r%array(i, j, k) = local_derivative
852 END DO
853 END DO
854 END DO
855 cpassert(local_row == features%nflat_local)
856
857 IF (ALLOCATED(atom_coords_pbc)) THEN
858 DEALLOCATE (atom_coords_pbc, dweights_datom, dweights_dstrain, included, &
859 partition_weights)
860 END IF
861 CALL torch_tensor_release(grid_weight_grad_t)
862 CALL torch_tensor_release(exc_tensor)
863 CALL skala_gpw_feature_release(features)
864 CALL xc_rho_set_release(rho_set, pw_pool=pw_pool)
865 CALL torch_use_cuda(.true.)
866 CALL omp_set_num_threads(omp_max_threads_restore)
867
868 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 SUBROUTINE skala_gapw_atom_vxc_of_r(xc_section, grid_atom, group, atom_coord, &
893 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 INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: atomic_grid_sizes
917 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 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: atomic_grid_weights, grid_weights
921 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: coarse_0_atomic_coords, density, &
922 grid_coords, kin
923 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :) :: grad
924 REAL(kind=dp), DIMENSION(:, :), POINTER :: atom_coord_grad, density_grad, &
925 grid_coord_grad, kin_grad
926 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 cpassert(ASSOCIATED(xc_section))
940 cpassert(ASSOCIATED(grid_atom))
941 cpassert(ASSOCIATED(rho))
942 cpassert(ASSOCIATED(drho))
943 cpassert(ASSOCIATED(tau))
944 omp_max_threads_restore = omp_get_max_threads()
945
946 my_energy_only = .false.
947 IF (PRESENT(energy_only)) my_energy_only = energy_only
948 need_coord_grad = PRESENT(atom_force) .OR. PRESENT(atom_virial)
949 exc = 0.0_dp
950 IF (PRESENT(atom_force)) atom_force = 0.0_dp
951 IF (PRESENT(atom_virial)) atom_virial = 0.0_dp
952 IF (.NOT. my_energy_only) THEN
953 vxc = 0.0_dp
954 vxg = 0.0_dp
955 vtau = 0.0_dp
956 END IF
957
958 CALL get_skala_model_path(xc_section, model_path)
959 gauxc_section => get_gauxc_section(xc_section)
960 cpassert(ASSOCIATED(gauxc_section))
961 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 i_val=native_grid_cuda_device)
964 CALL torch_use_cuda(native_grid_use_cuda)
965 selected_cuda_device = configure_native_grid_cuda( &
966 native_grid_use_cuda, native_grid_cuda_device, group)
967 CALL ensure_model_loaded(model_path, selected_cuda_device, xc_section)
968
969 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 coarse_0_atomic_coords(3, 1), atomic_grid_size_bound_shape(0, nflat))
974 density = 0.0_dp
975 grad = 0.0_dp
976 kin = 0.0_dp
977 grid_coords = 0.0_dp
978 grid_weights = 0.0_dp
979 atomic_grid_weights = 0.0_dp
980 atomic_grid_sizes(1) = int(nflat, kind=int_8)
981 atomic_grid_size_bound_shape = 0_int_8
982 coarse_0_atomic_coords(:, 1) = atom_coord
983
984 row = 0
985 DO ir = 1, nr
986 DO ia = 1, na
987 row = row + 1
988 grid_coords(1, row) = atom_coord(1) + grid_atom%rad(ir)* &
989 grid_atom%sin_pol(ia)*grid_atom%cos_azi(ia)
990 grid_coords(2, row) = atom_coord(2) + grid_atom%rad(ir)* &
991 grid_atom%sin_pol(ia)*grid_atom%sin_azi(ia)
992 grid_coords(3, row) = atom_coord(3) + grid_atom%rad(ir)*grid_atom%cos_pol(ia)
993 grid_weights(row) = weights(ia, ir)
994 atomic_grid_weights(row) = weights(ia, ir)
995 IF (nspins == 1) THEN
996 density(row, :) = 0.5_dp*rho(ia, ir, 1)
997 DO idir = 1, 3
998 grad(row, idir, :) = 0.5_dp*drho(idir, ia, ir, 1)
999 END DO
1000 kin(row, :) = 0.5_dp*tau(ia, ir, 1)
1001 ELSE
1002 density(row, :) = rho(ia, ir, 1:2)
1003 DO idir = 1, 3
1004 grad(row, idir, :) = drho(idir, ia, ir, 1:2)
1005 END DO
1006 kin(row, :) = tau(ia, ir, 1:2)
1007 END IF
1008 END DO
1009 END DO
1010
1011 CALL torch_tensor_from_array(grid_coords_t, grid_coords)
1012 CALL torch_tensor_to_device_leaf(grid_coords_t, need_coord_grad)
1013 CALL torch_tensor_from_array(grid_weights_t, grid_weights)
1014 CALL torch_tensor_to_device_leaf(grid_weights_t, .false.)
1015 CALL torch_tensor_from_array(atomic_grid_weights_t, atomic_grid_weights)
1016 CALL torch_tensor_to_device_leaf(atomic_grid_weights_t, .false.)
1017 CALL torch_tensor_from_array(atomic_grid_sizes_t, atomic_grid_sizes)
1018 CALL torch_tensor_to_device_leaf(atomic_grid_sizes_t, .false.)
1019 CALL torch_tensor_from_array(atomic_grid_size_bound_shape_t, &
1020 atomic_grid_size_bound_shape)
1021 CALL torch_tensor_to_device_leaf(atomic_grid_size_bound_shape_t, .false.)
1022 CALL torch_tensor_from_array(coarse_0_atomic_coords_t, coarse_0_atomic_coords)
1023 CALL torch_tensor_to_device_leaf(coarse_0_atomic_coords_t, need_coord_grad)
1024 CALL torch_tensor_from_array(density_t, density)
1025 CALL torch_tensor_to_device_leaf(density_t,.NOT. my_energy_only)
1026 CALL torch_tensor_from_array(grad_t, grad)
1027 CALL torch_tensor_to_device_leaf(grad_t,.NOT. my_energy_only)
1028 CALL torch_tensor_from_array(kin_t, kin)
1029 CALL torch_tensor_to_device_leaf(kin_t,.NOT. my_energy_only)
1030
1031 CALL torch_dict_create(inputs)
1032 CALL torch_dict_insert(inputs, "grid_coords", grid_coords_t)
1033 CALL torch_dict_insert(inputs, "grid_weights", grid_weights_t)
1034 CALL torch_dict_insert(inputs, "atomic_grid_weights", atomic_grid_weights_t)
1035 CALL torch_dict_insert(inputs, "atomic_grid_sizes", atomic_grid_sizes_t)
1036 CALL torch_dict_insert(inputs, "atomic_grid_size_bound_shape", &
1037 atomic_grid_size_bound_shape_t)
1038 CALL torch_dict_insert(inputs, "density", density_t)
1039 CALL torch_dict_insert(inputs, "grad", grad_t)
1040 CALL torch_dict_insert(inputs, "kin", kin_t)
1041 CALL torch_dict_insert(inputs, "coarse_0_atomic_coords", coarse_0_atomic_coords_t)
1042
1043 CALL skala_torch_model_get_exc(cached_model, inputs, grid_weights_t, exc_tensor, exc)
1044
1045 IF (.NOT. my_energy_only) THEN
1046 NULLIFY (atom_coord_grad, density_grad, grad_grad, grid_coord_grad, kin_grad)
1047 CALL torch_tensor_backward_scalar(exc_tensor)
1048 IF (need_coord_grad) THEN
1049 CALL torch_tensor_grad(grid_coords_t, grid_coord_grad_t)
1050 CALL torch_tensor_grad(coarse_0_atomic_coords_t, atom_coord_grad_t)
1051 CALL torch_tensor_data_ptr(grid_coord_grad_t, grid_coord_grad)
1052 CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
1053 IF (PRESENT(atom_force)) THEN
1054 atom_force(:) = atom_coord_grad(:, 1)
1055 DO row = 1, nflat
1056 atom_force(:) = atom_force(:) + grid_coord_grad(:, row)
1057 END DO
1058 END IF
1059 IF (PRESENT(atom_virial)) THEN
1060 DO row = 1, nflat
1061 DO idir = 1, 3
1062 DO jdir = 1, 3
1063 tmp = grid_coord_grad(idir, row)*coarse_0_atomic_coords(jdir, 1)
1064 atom_virial(idir, jdir) = atom_virial(idir, jdir) + tmp
1065 END DO
1066 END DO
1067 END DO
1068 DO idir = 1, 3
1069 DO jdir = 1, 3
1070 tmp = atom_coord_grad(idir, 1)*coarse_0_atomic_coords(jdir, 1)
1071 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 grad_grad_t, kin_grad_t)
1078 CALL torch_tensor_data_ptr(density_grad_t, density_grad)
1079 CALL torch_tensor_data_ptr(grad_grad_t, grad_grad)
1080 CALL torch_tensor_data_ptr(kin_grad_t, kin_grad)
1081
1082 row = 0
1083 DO ir = 1, nr
1084 DO ia = 1, na
1085 row = row + 1
1086 IF (lsd) THEN
1087 vxc(ia, ir, 1:2) = density_grad(row, 1:2)
1088 DO idir = 1, 3
1089 vxg(idir, ia, ir, 1:2) = grad_grad(row, idir, 1:2)
1090 END DO
1091 vtau(ia, ir, 1:2) = kin_grad(row, 1:2)
1092 ELSE
1093 vxc(ia, ir, 1) = 0.5_dp*(density_grad(row, 1) + density_grad(row, 2))
1094 DO idir = 1, 3
1095 vxg(idir, ia, ir, 1) = &
1096 0.5_dp*(grad_grad(row, idir, 1) + grad_grad(row, idir, 2))
1097 END DO
1098 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 CALL torch_tensor_release(density_grad_t)
1104 CALL torch_tensor_release(grad_grad_t)
1105 CALL torch_tensor_release(kin_grad_t)
1106 IF (need_coord_grad) THEN
1107 CALL torch_tensor_release(grid_coord_grad_t)
1108 CALL torch_tensor_release(atom_coord_grad_t)
1109 END IF
1110 END IF
1111
1112 CALL torch_tensor_release(exc_tensor)
1113 CALL torch_tensor_release(density_t)
1114 CALL torch_tensor_release(grad_t)
1115 CALL torch_tensor_release(kin_t)
1116 CALL torch_tensor_release(grid_coords_t)
1117 CALL torch_tensor_release(grid_weights_t)
1118 CALL torch_tensor_release(atomic_grid_weights_t)
1119 CALL torch_tensor_release(atomic_grid_sizes_t)
1120 CALL torch_tensor_release(atomic_grid_size_bound_shape_t)
1121 CALL torch_tensor_release(coarse_0_atomic_coords_t)
1122 CALL torch_dict_release(inputs)
1123 DEALLOCATE (atomic_grid_size_bound_shape, atomic_grid_sizes, atomic_grid_weights, &
1124 coarse_0_atomic_coords, density, grad, grid_coords, grid_weights, kin)
1125 CALL torch_use_cuda(.true.)
1126 CALL omp_set_num_threads(omp_max_threads_restore)
1127
1128 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 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 INTEGER, ALLOCATABLE, DIMENSION(:) :: chunk_atom_begin, chunk_atom_count, &
1179 chunk_row_begin, chunk_row_counts
1180 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 REAL(kind=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
1185 REAL(kind=dp), DIMENSION(:, :), POINTER :: atom_coord_grad, density_grad, &
1186 grid_coord_grad, kin_grad
1187 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 cpassert(ASSOCIATED(xc_section))
1211 cpassert(SIZE(density, 1) == SIZE(grid_weights))
1212 cpassert(SIZE(density, 1) == SIZE(atomic_grid_weights))
1213 cpassert(SIZE(density, 1) == SIZE(grid_coords, 2))
1214 cpassert(SIZE(density, 1) == SIZE(grad, 1))
1215 cpassert(SIZE(density, 1) == SIZE(kin, 1))
1216 cpassert(SIZE(density, 2) == 2)
1217 cpassert(SIZE(grad, 2) == 3)
1218 cpassert(SIZE(grad, 3) == 2)
1219 cpassert(SIZE(kin, 2) == 2)
1220 cpassert(SIZE(atomic_grid_sizes) == SIZE(atomic_coords, 2))
1221 cpassert(sum(atomic_grid_sizes) == int(SIZE(density, 1), kind=int_8))
1222 need_derivatives = PRESENT(density_grad_out)
1223 cpassert(PRESENT(grad_grad_out) .EQV. need_derivatives)
1224 cpassert(PRESENT(kin_grad_out) .EQV. need_derivatives)
1225 need_coordinate_derivatives = PRESENT(grid_coord_grad_out)
1226 cpassert(PRESENT(grid_weight_grad_out) .EQV. need_coordinate_derivatives)
1227 cpassert(PRESENT(atomic_grid_weight_grad_out) .EQV. need_coordinate_derivatives)
1228 cpassert(PRESENT(atom_coord_grad_out) .EQV. need_coordinate_derivatives)
1229 cpassert((.NOT. need_coordinate_derivatives) .OR. need_derivatives)
1230 omp_max_threads_restore = omp_get_max_threads()
1231
1232 local_nrow = SIZE(density, 1)
1233 local_natom = SIZE(atomic_grid_sizes)
1234 active_rank = local_natom > 0
1235 cpassert(active_rank .EQV. (local_nrow > 0))
1236
1237 CALL get_skala_model_path(xc_section, model_path)
1238 gauxc_section => get_gauxc_section(xc_section)
1239 cpassert(ASSOCIATED(gauxc_section))
1240 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 i_val=native_grid_cuda_device)
1243 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNKS", &
1244 l_val=native_grid_atom_chunks)
1245 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNK_MAX_ROWS", &
1246 i_val=chunk_max_rows)
1247 IF (chunk_max_rows < -1) THEN
1248 CALL cp_abort(__location__, &
1249 "GAUXC%NATIVE_GRID_ATOM_CHUNK_MAX_ROWS must be -1, zero, or positive.")
1250 END IF
1251 IF (.NOT. native_grid_atom_chunks) chunk_max_rows = 0
1252 IF (chunk_max_rows == -1) THEN
1253 chunk_max_rows = auto_atom_composite_chunk_max_rows(atomic_grid_sizes, group)
1254 END IF
1255 CALL torch_use_cuda(native_grid_use_cuda)
1256 selected_cuda_device = configure_native_grid_cuda( &
1257 native_grid_use_cuda, native_grid_cuda_device, group)
1258 exc = 0.0_dp
1259 IF (active_rank) THEN
1260 CALL timeset("skala_atom_tensor_setup", phase_handle)
1261 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 chunk_row_begin, chunk_row_counts)
1265 nchunks = SIZE(chunk_atom_begin)
1266 cpassert(nchunks > 0)
1267
1268 IF (native_grid_use_cuda) THEN
1269 grid_coords_t => cached_atom_cuda_tensors%grid_coords_t
1270 grid_weights_t => cached_atom_cuda_tensors%grid_weights_t
1271 atomic_grid_weights_t => cached_atom_cuda_tensors%atomic_grid_weights_t
1272 atomic_coords_t => cached_atom_cuda_tensors%atomic_coords_t
1273 density_t => cached_atom_cuda_tensors%density_t
1274 grad_t => cached_atom_cuda_tensors%grad_t
1275 kin_t => cached_atom_cuda_tensors%kin_t
1276 CALL torch_tensor_reset_from_array( &
1277 grid_coords_t, grid_coords, need_coordinate_derivatives)
1278 CALL torch_tensor_reset_from_array( &
1279 grid_weights_t, grid_weights, need_coordinate_derivatives)
1280 CALL torch_tensor_reset_from_array( &
1281 atomic_grid_weights_t, atomic_grid_weights, need_coordinate_derivatives)
1282 CALL torch_tensor_reset_from_array( &
1283 atomic_coords_t, atomic_coords, need_coordinate_derivatives)
1284 CALL torch_tensor_reset_from_array(density_t, density, need_derivatives)
1285 CALL torch_tensor_reset_from_array(grad_t, grad, need_derivatives)
1286 CALL torch_tensor_reset_from_array(kin_t, kin, need_derivatives)
1287 ELSE
1288 grid_coords_t => grid_coords_local_t
1289 grid_weights_t => grid_weights_local_t
1290 atomic_grid_weights_t => atomic_grid_weights_local_t
1291 atomic_coords_t => atomic_coords_local_t
1292 density_t => density_local_t
1293 grad_t => grad_local_t
1294 kin_t => kin_local_t
1295 CALL torch_tensor_from_array(grid_coords_t, grid_coords)
1296 CALL torch_tensor_to_device_leaf(grid_coords_t, need_coordinate_derivatives)
1297 CALL torch_tensor_from_array(grid_weights_t, grid_weights)
1298 CALL torch_tensor_to_device_leaf(grid_weights_t, need_coordinate_derivatives)
1299 CALL torch_tensor_from_array(atomic_grid_weights_t, atomic_grid_weights)
1300 CALL torch_tensor_to_device_leaf(atomic_grid_weights_t, need_coordinate_derivatives)
1301 CALL torch_tensor_from_array(atomic_coords_t, atomic_coords)
1302 CALL torch_tensor_to_device_leaf(atomic_coords_t, need_coordinate_derivatives)
1303 CALL torch_tensor_from_array(density_t, density)
1304 CALL torch_tensor_to_device_leaf(density_t, need_derivatives)
1305 CALL torch_tensor_from_array(grad_t, grad)
1306 CALL torch_tensor_to_device_leaf(grad_t, need_derivatives)
1307 CALL torch_tensor_from_array(kin_t, kin)
1308 CALL torch_tensor_to_device_leaf(kin_t, need_derivatives)
1309 END IF
1310 CALL torch_tensor_from_array(atomic_grid_sizes_t, atomic_grid_sizes)
1311 CALL torch_tensor_to_device_leaf(atomic_grid_sizes_t, .false.)
1312 CALL timestop(phase_handle)
1313
1314 DO ichunk = 1, nchunks
1315 atom_begin = chunk_atom_begin(ichunk)
1316 atom_count = chunk_atom_count(ichunk)
1317 chunk_row_start = chunk_row_begin(ichunk)
1318 chunk_row_count = chunk_row_counts(ichunk)
1319 max_grid_size = int(maxval(atomic_grid_sizes( &
1320 atom_begin:atom_begin + atom_count - 1)))
1321 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 chunk_row_count, grid_coords_chunk_t)
1326 CALL torch_tensor_narrow(grid_weights_t, 0, chunk_row_start - 1, &
1327 chunk_row_count, grid_weights_chunk_t)
1328 CALL torch_tensor_narrow(atomic_grid_weights_t, 0, chunk_row_start - 1, &
1329 chunk_row_count, atomic_grid_weights_chunk_t)
1330 CALL torch_tensor_narrow(atomic_grid_sizes_t, 0, atom_begin - 1, &
1331 atom_count, atomic_grid_sizes_chunk_t)
1332 CALL torch_tensor_narrow(atomic_coords_t, 0, atom_begin - 1, &
1333 atom_count, atomic_coords_chunk_t)
1334 CALL torch_tensor_narrow(density_t, 1, chunk_row_start - 1, &
1335 chunk_row_count, density_chunk_t)
1336 CALL torch_tensor_narrow(grad_t, 2, chunk_row_start - 1, &
1337 chunk_row_count, grad_chunk_t)
1338 CALL torch_tensor_narrow(kin_t, 1, chunk_row_start - 1, &
1339 chunk_row_count, kin_chunk_t)
1340 CALL torch_tensor_from_array(atomic_grid_size_bound_shape_t, &
1341 atomic_grid_size_bound_shape)
1342 CALL torch_tensor_to_device_leaf(atomic_grid_size_bound_shape_t, .false.)
1343
1344 CALL torch_dict_create(inputs)
1345 CALL torch_dict_insert(inputs, "grid_coords", grid_coords_chunk_t)
1346 CALL torch_dict_insert(inputs, "grid_weights", grid_weights_chunk_t)
1347 CALL torch_dict_insert(inputs, "atomic_grid_weights", atomic_grid_weights_chunk_t)
1348 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 atomic_grid_size_bound_shape_t)
1351 CALL torch_dict_insert(inputs, "density", density_chunk_t)
1352 CALL torch_dict_insert(inputs, "grad", grad_chunk_t)
1353 CALL torch_dict_insert(inputs, "kin", kin_chunk_t)
1354 CALL torch_dict_insert(inputs, "coarse_0_atomic_coords", atomic_coords_chunk_t)
1355
1356 CALL timeset("skala_atom_torch_forward", phase_handle)
1357 CALL skala_torch_model_get_exc(cached_model, inputs, grid_weights_chunk_t, &
1358 exc_tensor, chunk_exc)
1359 CALL timestop(phase_handle)
1360 exc = exc + chunk_exc
1361 IF (need_derivatives) THEN
1362 CALL timeset("skala_atom_torch_backward", phase_handle)
1363 CALL torch_tensor_backward_scalar(exc_tensor)
1364 CALL timestop(phase_handle)
1365 END IF
1366
1367 CALL torch_tensor_release(exc_tensor)
1368 CALL torch_dict_release(inputs)
1369 CALL torch_tensor_release(grid_coords_chunk_t)
1370 CALL torch_tensor_release(grid_weights_chunk_t)
1371 CALL torch_tensor_release(atomic_grid_weights_chunk_t)
1372 CALL torch_tensor_release(atomic_grid_sizes_chunk_t)
1373 CALL torch_tensor_release(atomic_coords_chunk_t)
1374 CALL torch_tensor_release(density_chunk_t)
1375 CALL torch_tensor_release(grad_chunk_t)
1376 CALL torch_tensor_release(kin_chunk_t)
1377 CALL torch_tensor_release(atomic_grid_size_bound_shape_t)
1378 DEALLOCATE (atomic_grid_size_bound_shape)
1379 END DO
1380 END IF
1381 CALL group%sum(exc)
1382
1383 IF (need_derivatives) THEN
1384 ALLOCATE (density_grad_out(local_nrow, 2), grad_grad_out(local_nrow, 3, 2), &
1385 kin_grad_out(local_nrow, 2))
1386 density_grad_out = 0.0_dp
1387 grad_grad_out = 0.0_dp
1388 kin_grad_out = 0.0_dp
1389 IF (active_rank) THEN
1390 NULLIFY (density_grad, grad_grad, kin_grad)
1391 CALL timeset("skala_atom_grad_fetch", phase_handle)
1392 CALL torch_tensor_grad_batch3(density_t, grad_t, kin_t, density_grad_t, &
1393 grad_grad_t, kin_grad_t)
1394 CALL torch_tensor_data_ptr(density_grad_t, density_grad)
1395 CALL torch_tensor_data_ptr(grad_grad_t, grad_grad)
1396 CALL torch_tensor_data_ptr(kin_grad_t, kin_grad)
1397 density_grad_out(:, :) = density_grad
1398 grad_grad_out(:, :, :) = grad_grad
1399 kin_grad_out(:, :) = kin_grad
1400 CALL timestop(phase_handle)
1401 END IF
1402
1403 IF (need_coordinate_derivatives) THEN
1404 IF (active_rank) CALL timeset("skala_atom_coord_grad_fetch", phase_handle)
1405 ALLOCATE (grid_coord_grad_out(3, local_nrow), grid_weight_grad_out(local_nrow), &
1406 atomic_grid_weight_grad_out(local_nrow), atom_coord_grad_out(3, local_natom))
1407 grid_coord_grad_out = 0.0_dp
1408 grid_weight_grad_out = 0.0_dp
1409 atomic_grid_weight_grad_out = 0.0_dp
1410 atom_coord_grad_out = 0.0_dp
1411 IF (active_rank) THEN
1412 NULLIFY (atomic_grid_weight_grad, atom_coord_grad, grid_coord_grad, grid_weight_grad)
1413 CALL torch_tensor_grad(grid_coords_t, grid_coord_grad_t)
1414 CALL torch_tensor_grad(grid_weights_t, grid_weight_grad_t)
1415 CALL torch_tensor_grad(atomic_grid_weights_t, atomic_grid_weight_grad_t)
1416 CALL torch_tensor_grad(atomic_coords_t, atom_coord_grad_t)
1417 CALL torch_tensor_data_ptr(grid_coord_grad_t, grid_coord_grad)
1418 CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
1419 CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
1420 CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
1421 grid_coord_grad_out(:, :) = grid_coord_grad
1422 grid_weight_grad_out(:) = grid_weight_grad
1423 atomic_grid_weight_grad_out(:) = atomic_grid_weight_grad
1424 atom_coord_grad_out(:, :) = atom_coord_grad
1425 CALL torch_tensor_release(grid_coord_grad_t)
1426 CALL torch_tensor_release(grid_weight_grad_t)
1427 CALL torch_tensor_release(atomic_grid_weight_grad_t)
1428 CALL torch_tensor_release(atom_coord_grad_t)
1429 CALL timestop(phase_handle)
1430 END IF
1431 END IF
1432 IF (active_rank) THEN
1433 CALL torch_tensor_release(density_grad_t)
1434 CALL torch_tensor_release(grad_grad_t)
1435 CALL torch_tensor_release(kin_grad_t)
1436 END IF
1437 END IF
1438
1439 IF (active_rank) THEN
1440 CALL timeset("skala_atom_tensor_release", phase_handle)
1441 IF (.NOT. native_grid_use_cuda) THEN
1442 CALL torch_tensor_release(density_t)
1443 CALL torch_tensor_release(grad_t)
1444 CALL torch_tensor_release(kin_t)
1445 CALL torch_tensor_release(grid_coords_t)
1446 CALL torch_tensor_release(grid_weights_t)
1447 CALL torch_tensor_release(atomic_grid_weights_t)
1448 CALL torch_tensor_release(atomic_coords_t)
1449 END IF
1450 CALL torch_tensor_release(atomic_grid_sizes_t)
1451 DEALLOCATE (chunk_atom_begin, chunk_atom_count, chunk_row_begin, chunk_row_counts)
1452 CALL timestop(phase_handle)
1453 END IF
1454 CALL torch_use_cuda(.true.)
1455 CALL omp_set_num_threads(omp_max_threads_restore)
1456
1457 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 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 local_padded_rows = 0
1474 local_natom = SIZE(atomic_grid_sizes)
1475 IF (local_natom > 0) THEN
1476 local_max_atom_rows = int(maxval(atomic_grid_sizes))
1477 IF (local_natom > atom_chunk_auto_max_rows .OR. &
1478 local_max_atom_rows > atom_chunk_auto_max_rows/local_natom) THEN
1479 local_padded_rows = atom_chunk_auto_max_rows + 1
1480 ELSE
1481 local_padded_rows = local_natom*local_max_atom_rows
1482 END IF
1483 END IF
1484 CALL group%max(local_padded_rows)
1485 IF (local_padded_rows <= atom_chunk_auto_max_rows) THEN
1486 max_rows = 0
1487 RETURN
1488 END IF
1489 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 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 INTEGER, ALLOCATABLE, DIMENSION(:) :: padded_rows
1514
1515 natom = SIZE(atomic_grid_sizes)
1516 cpassert(natom > 0)
1517 cpassert(all(atomic_grid_sizes > 0_int_8))
1518 IF (max_rows <= 0) THEN
1519 ALLOCATE (atom_begin(1), atom_count(1), row_begin(1), row_count(1))
1520 atom_begin(1) = 1
1521 atom_count(1) = natom
1522 row_begin(1) = 1
1523 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 DO iatom = 1, natom
1531 atom_rows = int(atomic_grid_sizes(iatom))
1532 IF (chunk_atoms > 0 .AND. &
1533 max(chunk_max_atom_rows, atom_rows) > max_rows/(chunk_atoms + 1)) THEN
1534 nchunks = nchunks + 1
1535 chunk_atoms = 0
1536 chunk_max_atom_rows = 0
1537 END IF
1538 chunk_atoms = chunk_atoms + 1
1539 chunk_max_atom_rows = max(chunk_max_atom_rows, atom_rows)
1540 END DO
1541 ALLOCATE (atom_begin(nchunks), atom_count(nchunks), row_begin(nchunks), row_count(nchunks))
1542
1543 subchunk = 1
1544 atom_begin(subchunk) = 1
1545 row_begin(subchunk) = 1
1546 chunk_atoms = 0
1547 chunk_max_atom_rows = 0
1548 rows = 0
1549 DO iatom = 1, natom
1550 atom_rows = int(atomic_grid_sizes(iatom))
1551 IF (chunk_atoms > 0 .AND. &
1552 max(chunk_max_atom_rows, atom_rows) > max_rows/(chunk_atoms + 1)) THEN
1553 atom_count(subchunk) = iatom - atom_begin(subchunk)
1554 row_count(subchunk) = rows
1555 subchunk = subchunk + 1
1556 atom_begin(subchunk) = iatom
1557 row_begin(subchunk) = row_begin(subchunk - 1) + row_count(subchunk - 1)
1558 chunk_atoms = 0
1559 chunk_max_atom_rows = 0
1560 rows = 0
1561 END IF
1562 chunk_atoms = chunk_atoms + 1
1563 chunk_max_atom_rows = max(chunk_max_atom_rows, atom_rows)
1564 rows = rows + atom_rows
1565 END DO
1566 atom_count(subchunk) = natom - atom_begin(subchunk) + 1
1567 row_count(subchunk) = rows
1568 cpassert(subchunk == nchunks)
1569
1570 ALLOCATE (padded_rows(nchunks))
1571 DO subchunk = 1, nchunks
1572 iatom = atom_begin(subchunk) + atom_count(subchunk) - 1
1573 padded_rows(subchunk) = atom_count(subchunk)* &
1574 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 DO subchunk = 2, nchunks
1579 padded_rows_tmp = padded_rows(subchunk)
1580 atom_begin_tmp = atom_begin(subchunk)
1581 atom_count_tmp = atom_count(subchunk)
1582 row_begin_tmp = row_begin(subchunk)
1583 row_count_tmp = row_count(subchunk)
1584 insert_at = subchunk
1585 DO WHILE (insert_at > 1 .AND. padded_rows(insert_at - 1) < padded_rows_tmp)
1586 padded_rows(insert_at) = padded_rows(insert_at - 1)
1587 atom_begin(insert_at) = atom_begin(insert_at - 1)
1588 atom_count(insert_at) = atom_count(insert_at - 1)
1589 row_begin(insert_at) = row_begin(insert_at - 1)
1590 row_count(insert_at) = row_count(insert_at - 1)
1591 insert_at = insert_at - 1
1592 END DO
1593 padded_rows(insert_at) = padded_rows_tmp
1594 atom_begin(insert_at) = atom_begin_tmp
1595 atom_count(insert_at) = atom_count_tmp
1596 row_begin(insert_at) = row_begin_tmp
1597 row_count(insert_at) = row_count_tmp
1598 END DO
1599 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 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 REAL(kind=dp), DIMENSION(:, :), POINTER :: atom_coord_grad
1617
1618 NULLIFY (atom_coord_grad)
1619 CALL torch_tensor_grad(features%coarse_0_atomic_coords_t, atom_coord_grad_t)
1620 IF (root_rank) THEN
1621 CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
1622 cpassert(SIZE(atom_force, 1) == SIZE(atom_coord_grad, 1))
1623 cpassert(SIZE(atom_force, 2) == SIZE(atom_coord_grad, 2))
1624 atom_force(:, :) = atom_force(:, :) + atom_coord_grad(:, :)
1625 END IF
1626
1627 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 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 REAL(kind=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
1660
1661 NULLIFY (atomic_grid_weight_grad, grid_weight_grad)
1662 CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
1663 CALL torch_tensor_grad(features%atomic_grid_weights_t, atomic_grid_weight_grad_t)
1664 CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
1665 CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
1666
1667 natom = SIZE(particle_set)
1668 cpassert(SIZE(atom_force, 1) == 3)
1669 cpassert(SIZE(atom_force, 2) == natom)
1670 ALLOCATE (atom_coords_pbc(3, natom), included(natom), weights(natom), &
1671 dweights_datom(3, natom, natom), dweights_dstrain(3, 3, natom))
1672 DO iatom = 1, natom
1673 atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.true.)
1674 END DO
1675
1676 bo = rho_r(1)%pw_grid%bounds_local
1677 local_row = 0
1678 DO k = bo(1, 3), bo(2, 3)
1679 DO j = bo(1, 2), bo(2, 2)
1680 DO i = bo(1, 1), bo(2, 1)
1681 local_row = local_row + 1
1682 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 dweights_dstrain)
1686 feature_begin = features%local_feature_offsets(local_row)
1687 feature_end = features%local_feature_offsets(local_row + 1) - 1
1688 cpassert(feature_end - feature_begin + 1 == count(included))
1689 grid_base_weight = 0.0_dp
1690 DO feature_pos = feature_begin, feature_end
1691 row = features%local_feature_rows(feature_pos)
1692 grid_base_weight = grid_base_weight + features%grid_weights(row)
1693 END DO
1694 feature_pos = feature_begin
1695 DO iatom = 1, natom
1696 IF (.NOT. included(iatom)) cycle
1697 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 smooth_partition_atomic_weight_scale_derivative(weights(iatom))
1701 DO jatom = 1, natom
1702 atom_force(:, jatom) = atom_force(:, jatom) + &
1703 weight_grad*dweights_datom(:, jatom, iatom)
1704 END DO
1705 feature_pos = feature_pos + 1
1706 END DO
1707 cpassert(feature_pos == feature_end + 1)
1708 END DO
1709 END DO
1710 END DO
1711 cpassert(local_row == features%nflat_local)
1712
1713 DEALLOCATE (atom_coords_pbc, dweights_datom, dweights_dstrain, included, weights)
1714 CALL torch_tensor_release(grid_weight_grad_t)
1715 CALL torch_tensor_release(atomic_grid_weight_grad_t)
1716
1717 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 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 REAL(kind=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
1750
1751 NULLIFY (atomic_grid_weight_grad, grid_weight_grad)
1752 CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
1753 CALL torch_tensor_grad(features%atomic_grid_weights_t, atomic_grid_weight_grad_t)
1754 CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
1755 CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
1756
1757 natom = SIZE(particle_set)
1758 ALLOCATE (atom_coords_pbc(3, natom), included(natom), weights(natom), &
1759 dweights_datom(3, natom, natom), dweights_dstrain(3, 3, natom))
1760 DO iatom = 1, natom
1761 atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.true.)
1762 END DO
1763
1764 bo = rho_r(1)%pw_grid%bounds_local
1765 local_row = 0
1766 DO k = bo(1, 3), bo(2, 3)
1767 DO j = bo(1, 2), bo(2, 2)
1768 DO i = bo(1, 1), bo(2, 1)
1769 local_row = local_row + 1
1770 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 dweights_dstrain)
1774 feature_begin = features%local_feature_offsets(local_row)
1775 feature_end = features%local_feature_offsets(local_row + 1) - 1
1776 cpassert(feature_end - feature_begin + 1 == count(included))
1777 grid_base_weight = 0.0_dp
1778 DO feature_pos = feature_begin, feature_end
1779 row = features%local_feature_rows(feature_pos)
1780 grid_base_weight = grid_base_weight + features%grid_weights(row)
1781 END DO
1782 feature_pos = feature_begin
1783 DO iatom = 1, natom
1784 IF (.NOT. included(iatom)) cycle
1785 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 smooth_partition_atomic_weight_scale_derivative(weights(iatom))
1789 DO idir = 1, 3
1790 DO jdir = 1, idir
1791 tmp = weight_grad*dweights_dstrain(idir, jdir, iatom)
1792 virial_xc(jdir, idir) = virial_xc(jdir, idir) + tmp
1793 IF (idir /= jdir) virial_xc(idir, jdir) = virial_xc(idir, jdir) + tmp
1794 END DO
1795 END DO
1796 feature_pos = feature_pos + 1
1797 END DO
1798 cpassert(feature_pos == feature_end + 1)
1799 END DO
1800 END DO
1801 END DO
1802 cpassert(local_row == features%nflat_local)
1803
1804 DEALLOCATE (atom_coords_pbc, dweights_datom, dweights_dstrain, included, weights)
1805 CALL torch_tensor_release(grid_weight_grad_t)
1806 CALL torch_tensor_release(atomic_grid_weight_grad_t)
1807
1808 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 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 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 REAL(relative_index(3), kind=dp)*pw_grid%dh(:, 3)
1827
1828 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 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 INTEGER, ALLOCATABLE, DIMENSION(:) :: route_grad_return_recv_counts, &
1859 route_grad_return_recv_displs, &
1860 route_grad_return_send_counts, &
1861 route_grad_return_send_displs, &
1862 subchunk_atom_begin, &
1863 subchunk_atom_count, &
1864 subchunk_row_begin, &
1865 subchunk_row_count
1866 REAL(kind=dp) :: subchunk_exc
1867 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: recv_grad_buffer, send_grad_buffer
1868 TYPE(skala_gpw_feature_type) :: subchunk
1869 TYPE(torch_tensor_type) :: subchunk_exc_tensor
1870
1871 cpassert(features%uses_atom_chunks)
1872 cpassert(max_rows > 0)
1873 CALL skala_gpw_atom_subchunk_layout(max_rows, subchunk_atom_begin, subchunk_atom_count, &
1874 subchunk_row_begin, subchunk_row_count)
1875 nsubchunks = SIZE(subchunk_atom_begin)
1876
1877 exc = 0.0_dp
1878 IF (compute_grads) THEN
1879 cpassert(features%uses_atom_chunk_routing)
1880 nroute_recv_points = sum(features%route_point_recv_counts)
1881 nroute_send_points = SIZE(features%route_send_local_rows)
1882 cpassert(sum(features%route_point_send_counts) == nroute_send_points)
1883 cpassert(SIZE(features%route_chunk_offsets) == nroute_recv_points + 1)
1884 nroute_grad_per_point = ngrad_per_point
1885 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 route_grad_return_recv_displs(SIZE(features%route_point_send_displs)))
1892 route_grad_return_send_counts(:) = &
1893 nroute_grad_per_point*features%route_point_recv_counts
1894 route_grad_return_send_displs(:) = &
1895 nroute_grad_per_point*features%route_point_recv_displs
1896 route_grad_return_recv_counts(:) = &
1897 nroute_grad_per_point*features%route_point_send_counts
1898 route_grad_return_recv_displs(:) = &
1899 nroute_grad_per_point*features%route_point_send_displs
1900 END IF
1901
1902 CALL timeset("skala_gpw_atom_subchunks", phase_handle)
1903 DO isubchunk = 1, nsubchunks
1904 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 compute_grads)
1911 CALL timestop(subphase_handle)
1912 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 subchunk_exc)
1916 CALL timestop(subphase_handle)
1917 exc = exc + subchunk_exc
1918 IF (compute_grads) THEN
1919 CALL timeset("skala_gpw_atom_subchunk_backward", subphase_handle)
1920 CALL torch_tensor_backward_scalar(subchunk_exc_tensor)
1921 CALL timestop(subphase_handle)
1922 END IF
1923 CALL timeset("skala_gpw_atom_subchunk_release", subphase_handle)
1924 CALL torch_tensor_release(subchunk_exc_tensor)
1925 CALL skala_gpw_feature_release(subchunk)
1926 CALL timestop(subphase_handle)
1927 END DO
1928 IF (compute_grads .AND. features%chunk_feature_count > 0) THEN
1929 CALL timeset("skala_gpw_atom_subchunk_grad_pack", subphase_handle)
1930 CALL pack_atom_chunk_grads(features, send_grad_buffer, .true., collapse_spin_grads)
1931 CALL timestop(subphase_handle)
1932 END IF
1933 CALL timestop(phase_handle)
1934
1935 IF (compute_grads) THEN
1936 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 route_grad_return_recv_counts, route_grad_return_recv_displs)
1940 CALL timestop(phase_handle)
1941
1942 CALL timeset("skala_gpw_grad_route_scatter", phase_handle)
1943 CALL scatter_routed_atom_chunk_grads(features, recv_grad_buffer, collapse_spin_grads, &
1944 density_grad, grad_grad, kin_grad)
1945 CALL timestop(phase_handle)
1946
1947 DEALLOCATE (recv_grad_buffer, route_grad_return_recv_counts, &
1948 route_grad_return_recv_displs, route_grad_return_send_counts, &
1949 route_grad_return_send_displs, send_grad_buffer)
1950 END IF
1951 DEALLOCATE (subchunk_atom_begin, subchunk_atom_count, subchunk_row_begin, subchunk_row_count)
1952
1953 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 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 local_rows_max = features%chunk_feature_count
1970 CALL group%max(local_rows_max)
1971 IF (local_rows_max <= 0) THEN
1972 max_rows = 0
1973 RETURN
1974 END IF
1975
1976 IF (group%num_pe > 1) THEN
1977 target_rows = ceiling(real(local_rows_max, kind=dp)/2.0_dp)
1978 max_rows = atom_chunk_auto_row_quantum* &
1979 ((target_rows + atom_chunk_auto_row_quantum - 1)/atom_chunk_auto_row_quantum)
1980 ELSE
1981 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 REAL(atom_chunk_auto_row_quantum, kind=dp)))
1985 END IF
1986 max_rows = max(atom_chunk_auto_min_rows, min(atom_chunk_auto_max_rows, max_rows))
1987
1988 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 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 REAL(kind=dp), DIMENSION(:, :), POINTER :: density_grad_all, kin_grad_all
2008 REAL(kind=dp), DIMENSION(:, :, :), POINTER :: grad_grad_all
2009 TYPE(torch_tensor_type) :: density_grad_t, grad_grad_t, kin_grad_t
2010
2011 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 density_grad_all, grad_grad_all, kin_grad_all)
2014 cpassert(SIZE(density_grad_all, 1) == features%nflat)
2015 cpassert(SIZE(density_grad_all, 2) == 2)
2016 cpassert(SIZE(grad_grad_all, 1) == features%nflat)
2017 cpassert(SIZE(grad_grad_all, 2) == 3)
2018 cpassert(SIZE(grad_grad_all, 3) == 2)
2019 cpassert(SIZE(kin_grad_all, 1) == features%nflat)
2020 cpassert(SIZE(kin_grad_all, 2) == 2)
2021
2022 ALLOCATE (density_grad(features%nflat_local, 2), &
2023 grad_grad(features%nflat_local, 3, 2), &
2024 kin_grad(features%nflat_local, 2))
2025 density_grad = 0.0_dp
2026 grad_grad = 0.0_dp
2027 kin_grad = 0.0_dp
2028 local_row = 0
2029 DO k = lbound(features%feature_index, 3), ubound(features%feature_index, 3)
2030 DO j = lbound(features%feature_index, 2), ubound(features%feature_index, 2)
2031 DO i = lbound(features%feature_index, 1), ubound(features%feature_index, 1)
2032 local_row = local_row + 1
2033 DO feature_pos = features%local_feature_offsets(local_row), &
2034 features%local_feature_offsets(local_row + 1) - 1
2035 row = features%local_feature_rows(feature_pos)
2036 cpassert(row >= 1 .AND. row <= features%nflat)
2037 density_grad(local_row, :) = density_grad(local_row, :) + &
2038 density_grad_all(row, :)
2039 grad_grad(local_row, :, :) = grad_grad(local_row, :, :) + &
2040 grad_grad_all(row, :, :)
2041 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 cpassert(local_row == features%nflat_local)
2047
2048 CALL torch_tensor_release(density_grad_t)
2049 CALL torch_tensor_release(grad_grad_t)
2050 CALL torch_tensor_release(kin_grad_t)
2051
2052 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 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 REAL(kind=dp), DIMENSION(:, :), POINTER :: chunk_density_grad, chunk_kin_grad
2074 REAL(kind=dp), DIMENSION(:, :, :), POINTER :: chunk_grad_grad
2075 TYPE(torch_tensor_type) :: density_grad_t, grad_grad_t, kin_grad_t
2076
2077 my_collapse_spin_grads = .false.
2078 IF (PRESENT(collapse_spin_grads)) my_collapse_spin_grads = collapse_spin_grads
2079 ngrad_buffer_per_point = ngrad_per_point
2080 IF (my_collapse_spin_grads) ngrad_buffer_per_point = ncollapsed_grad_per_point
2081
2082 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 chunk_density_grad, chunk_grad_grad, chunk_kin_grad)
2085 cpassert(mod(SIZE(TARGET), ngrad_buffer_per_point) == 0)
2086 target_points = SIZE(TARGET)/ngrad_buffer_per_point
2087 cpassert(SIZE(chunk_density_grad, 1) == features%chunk_feature_count)
2088 cpassert(SIZE(chunk_grad_grad, 1) == features%chunk_feature_count)
2089 cpassert(SIZE(chunk_grad_grad, 2) == 3)
2090 cpassert(SIZE(chunk_kin_grad, 1) == features%chunk_feature_count)
2091 IF (features%uses_collapsed_rks_dynamic) THEN
2092 cpassert(my_collapse_spin_grads)
2093 cpassert(SIZE(chunk_density_grad, 2) == 1)
2094 cpassert(SIZE(chunk_grad_grad, 3) == 1)
2095 cpassert(SIZE(chunk_kin_grad, 2) == 1)
2096 ELSE
2097 cpassert(SIZE(chunk_density_grad, 2) == 2)
2098 cpassert(SIZE(chunk_grad_grad, 3) == 2)
2099 cpassert(SIZE(chunk_kin_grad, 2) == 2)
2100 END IF
2101
2102 IF (route_to_return_positions) THEN
2103 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!$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 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!$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 CALL torch_tensor_release(density_grad_t)
2201 CALL torch_tensor_release(grad_grad_t)
2202 CALL torch_tensor_release(kin_grad_t)
2203
2204 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 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 nflat_local = features%nflat_local
2232 nroute_points = SIZE(features%route_send_local_rows)
2233 nroute_grad_per_point = ngrad_per_point
2234 IF (collapse_spin_grads) nroute_grad_per_point = ncollapsed_grad_per_point
2235 cpassert(SIZE(recv_grad_buffer) >= nroute_grad_per_point*nroute_points)
2236 ALLOCATE (density_grad(nflat_local, 2), grad_grad(nflat_local, 3, 2), &
2237 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!$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 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 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 NULLIFY (density_grad, grad_grad, kin_grad)
2305 CALL torch_tensor_grad_batch3(features%density_t, features%grad_t, features%kin_t, &
2306 density_grad_t, grad_grad_t, kin_grad_t)
2307 CALL torch_tensor_data_ptr(density_grad_t, density_grad)
2308 CALL torch_tensor_data_ptr(grad_grad_t, grad_grad)
2309 CALL torch_tensor_data_ptr(kin_grad_t, kin_grad)
2310
2311 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 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 INTEGER, ALLOCATABLE, DIMENSION(:) :: route_grad_return_recv_counts, &
2336 route_grad_return_recv_displs, &
2337 route_grad_return_send_counts, &
2338 route_grad_return_send_displs
2339 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: chunk_grad_buffer, global_grad_buffer, &
2340 recv_grad_buffer, send_grad_buffer
2341
2342 cpassert(features%uses_atom_chunks)
2343
2344 nflat_local = features%nflat_local
2345 IF (features%uses_atom_chunk_routing) THEN
2346 nroute_recv_points = sum(features%route_point_recv_counts)
2347 nroute_send_points = SIZE(features%route_send_local_rows)
2348 cpassert(sum(features%route_point_send_counts) == nroute_send_points)
2349 cpassert(SIZE(features%route_chunk_offsets) == nroute_recv_points + 1)
2350
2351 nroute_grad_per_point = ngrad_per_point
2352 IF (features%uses_collapsed_rks_dynamic) THEN
2353 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 route_grad_return_recv_displs(SIZE(features%route_point_send_displs)))
2361 route_grad_return_send_counts(:) = &
2362 nroute_grad_per_point*features%route_point_recv_counts
2363 route_grad_return_send_displs(:) = &
2364 nroute_grad_per_point*features%route_point_recv_displs
2365 route_grad_return_recv_counts(:) = &
2366 nroute_grad_per_point*features%route_point_send_counts
2367 route_grad_return_recv_displs(:) = &
2368 nroute_grad_per_point*features%route_point_send_displs
2369
2370 IF (features%chunk_feature_count > 0) THEN
2371 CALL timeset("skala_gpw_grad_torch_pack", phase_handle)
2372 CALL pack_atom_chunk_grads(features, send_grad_buffer, .true., &
2373 features%uses_collapsed_rks_dynamic)
2374 CALL timestop(phase_handle)
2375 END IF
2376
2377 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 route_grad_return_recv_counts, route_grad_return_recv_displs)
2381 CALL timestop(phase_handle)
2382
2383 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 density_grad, grad_grad, kin_grad)
2387 CALL timestop(phase_handle)
2388
2389 DEALLOCATE (recv_grad_buffer, route_grad_return_recv_counts, &
2390 route_grad_return_recv_displs, route_grad_return_send_counts, &
2391 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 global_grad_buffer(ngrad_per_point*features%nflat))
2395 IF (features%chunk_feature_count > 0) THEN
2396 CALL timeset("skala_gpw_grad_torch_pack", phase_handle)
2397 CALL pack_atom_chunk_grads(features, chunk_grad_buffer, .false.)
2398 CALL timestop(phase_handle)
2399 END IF
2400
2401 CALL timeset("skala_gpw_grad_allgatherv", phase_handle)
2402 CALL group%allgatherv(chunk_grad_buffer, global_grad_buffer, &
2403 features%chunk_grad_counts, features%chunk_grad_displs)
2404 CALL timestop(phase_handle)
2405
2406 CALL timeset("skala_gpw_grad_scatter", phase_handle)
2407 ALLOCATE (density_grad(nflat_local, 2), grad_grad(nflat_local, 3, 2), &
2408 kin_grad(nflat_local, 2))
2409 density_grad = 0.0_dp
2410 grad_grad = 0.0_dp
2411 kin_grad = 0.0_dp
2412 local_row = 0
2413 DO k = lbound(features%feature_index, 3), ubound(features%feature_index, 3)
2414 DO j = lbound(features%feature_index, 2), ubound(features%feature_index, 2)
2415 DO i = lbound(features%feature_index, 1), ubound(features%feature_index, 1)
2416 local_row = local_row + 1
2417 DO feature_pos = features%local_feature_offsets(local_row), &
2418 features%local_feature_offsets(local_row + 1) - 1
2419 row = features%local_feature_rows(feature_pos)
2420 cpassert(row >= 1 .AND. row <= features%nflat)
2421 base = ngrad_per_point*(row - 1)
2422 density_grad(local_row, :) = density_grad(local_row, :) + &
2423 global_grad_buffer(base + 1:base + 2)
2424 grad_grad(local_row, 1, 1) = grad_grad(local_row, 1, 1) + &
2425 global_grad_buffer(base + 3)
2426 grad_grad(local_row, 2, 1) = grad_grad(local_row, 2, 1) + &
2427 global_grad_buffer(base + 4)
2428 grad_grad(local_row, 3, 1) = grad_grad(local_row, 3, 1) + &
2429 global_grad_buffer(base + 5)
2430 grad_grad(local_row, 1, 2) = grad_grad(local_row, 1, 2) + &
2431 global_grad_buffer(base + 6)
2432 grad_grad(local_row, 2, 2) = grad_grad(local_row, 2, 2) + &
2433 global_grad_buffer(base + 7)
2434 grad_grad(local_row, 3, 2) = grad_grad(local_row, 3, 2) + &
2435 global_grad_buffer(base + 8)
2436 kin_grad(local_row, :) = kin_grad(local_row, :) + &
2437 global_grad_buffer(base + 9:base + 10)
2438 END DO
2439 END DO
2440 END DO
2441 END DO
2442 CALL timestop(phase_handle)
2443 DEALLOCATE (chunk_grad_buffer, global_grad_buffer)
2444
2445 END IF
2446
2447 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 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 TYPE(cp_3d_r_cp_type), DIMENSION(3) :: drho, drhoa, drhob
2466
2467 nspins = SIZE(rho_r)
2468 bo = rho_r(1)%pw_grid%bounds_local
2469 ipt = 0
2470
2471 IF (nspins == 1) THEN
2472 CALL xc_rho_set_get(rho_set, drho=drho)
2473 DO k = bo(1, 3), bo(2, 3)
2474 DO j = bo(1, 2), bo(2, 2)
2475 DO i = bo(1, 1), bo(2, 1)
2476 ipt = ipt + 1
2477 DO idir = 1, 3
2478 grad_i = 0.5_dp*(grad_grad(ipt, idir, 1) + grad_grad(ipt, idir, 2))
2479 DO jdir = 1, 3
2480 tmp = -grad_i*drho(jdir)%array(i, j, k)
2481 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 CALL xc_rho_set_get(rho_set, drhoa=drhoa, drhob=drhob)
2489 DO k = bo(1, 3), bo(2, 3)
2490 DO j = bo(1, 2), bo(2, 2)
2491 DO i = bo(1, 1), bo(2, 1)
2492 ipt = ipt + 1
2493 DO idir = 1, 3
2494 DO jdir = 1, 3
2495 tmp = 0.0_dp
2496 DO ispin = 1, 2
2497 IF (ispin == 1) THEN
2498 tmp = tmp - grad_grad(ipt, idir, ispin)*drhoa(jdir)%array(i, j, k)
2499 ELSE
2500 tmp = tmp - grad_grad(ipt, idir, ispin)*drhob(jdir)%array(i, j, k)
2501 END IF
2502 END DO
2503 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 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 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 IF (.NOT. root_rank) RETURN
2527 iw = cp_logger_get_default_io_unit()
2528 IF (iw <= 0) RETURN
2529 WRITE (iw, "(T2,A,1X,A)") "SKALA_GPW| XC virial contribution", trim(label)
2530 DO i = 1, 3
2531 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 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 REAL(kind=dp), DIMENSION(:, :), POINTER :: atom_coord_grad, grid_coord_grad
2560
2561 my_print_components = .false.
2562 IF (PRESENT(print_components)) my_print_components = print_components
2563
2564 NULLIFY (atom_coord_grad, grid_coord_grad)
2565 CALL torch_tensor_grad(features%grid_coords_t, grid_coord_grad_t)
2566 CALL torch_tensor_data_ptr(grid_coord_grad_t, grid_coord_grad)
2567 CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
2568
2569 grid_virial = 0.0_dp
2570 atom_virial = 0.0_dp
2571 grid_grad_sum = 0.0_dp
2572 atom_grad_sum = 0.0_dp
2573 local_row = 0
2574 DO k = lbound(features%feature_index, 3), ubound(features%feature_index, 3)
2575 DO j = lbound(features%feature_index, 2), ubound(features%feature_index, 2)
2576 DO i = lbound(features%feature_index, 1), ubound(features%feature_index, 1)
2577 local_row = local_row + 1
2578 DO feature_pos = features%local_feature_offsets(local_row), &
2579 features%local_feature_offsets(local_row + 1) - 1
2580 row = features%local_feature_rows(feature_pos)
2581 grid_grad_sum(:) = grid_grad_sum(:) + grid_coord_grad(:, row)
2582 DO idir = 1, 3
2583 DO jdir = 1, 3
2584 tmp = grid_coord_grad(idir, row)*features%grid_coords(jdir, row)
2585 grid_virial(idir, jdir) = grid_virial(idir, jdir) + tmp
2586 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 cpassert(local_row == features%nflat_local)
2594
2595 IF (root_rank) THEN
2596 DO iatom = 1, SIZE(features%coarse_0_atomic_coords, 2)
2597 atom_grad_sum(:) = atom_grad_sum(:) + atom_coord_grad(:, iatom)
2598 DO idir = 1, 3
2599 DO jdir = 1, 3
2600 tmp = atom_coord_grad(idir, iatom)*features%coarse_0_atomic_coords(jdir, iatom)
2601 atom_virial(idir, jdir) = atom_virial(idir, jdir) + tmp
2602 virial_xc(idir, jdir) = virial_xc(idir, jdir) + tmp
2603 END DO
2604 END DO
2605 END DO
2606 END IF
2607
2608 IF (my_print_components .AND. root_rank) THEN
2609 iw = cp_logger_get_default_io_unit()
2610 IF (iw > 0) THEN
2611 CALL print_virial_delta("static-grid", grid_virial, .true.)
2612 CALL print_virial_delta("static-atom", atom_virial, .true.)
2613 WRITE (iw, "(T2,A,1X,3ES20.10)") "SKALA_GPW| XC coordinate gradient grid sum", &
2614 grid_grad_sum
2615 WRITE (iw, "(T2,A,1X,3ES20.10)") "SKALA_GPW| XC coordinate gradient atom sum", &
2616 atom_grad_sum
2617 WRITE (iw, "(T2,A,1X,3ES20.10)") "SKALA_GPW| XC coordinate gradient total sum", &
2618 grid_grad_sum + atom_grad_sum
2619 END IF
2620 END IF
2621
2622 CALL torch_tensor_release(grid_coord_grad_t)
2623
2624 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 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 REAL(kind=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
2651
2652 my_print_components = .false.
2653 IF (PRESENT(print_components)) my_print_components = print_components
2654
2655 NULLIFY (atomic_grid_weight_grad, grid_weight_grad)
2656 CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
2657 CALL torch_tensor_grad(features%atomic_grid_weights_t, atomic_grid_weight_grad_t)
2658 CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
2659 CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
2660
2661 grid_tmp = 0.0_dp
2662 atomic_tmp = 0.0_dp
2663 local_row = 0
2664 DO k = lbound(features%feature_index, 3), ubound(features%feature_index, 3)
2665 DO j = lbound(features%feature_index, 2), ubound(features%feature_index, 2)
2666 DO i = lbound(features%feature_index, 1), ubound(features%feature_index, 1)
2667 local_row = local_row + 1
2668 DO feature_pos = features%local_feature_offsets(local_row), &
2669 features%local_feature_offsets(local_row + 1) - 1
2670 row = features%local_feature_rows(feature_pos)
2671 grid_tmp = grid_tmp + grid_weight_grad(row)*features%grid_weights(row)
2672 atomic_tmp = atomic_tmp + &
2673 atomic_grid_weight_grad(row)*features%atomic_grid_weights(row)
2674 END DO
2675 END DO
2676 END DO
2677 END DO
2678 cpassert(local_row == features%nflat_local)
2679 exc_tmp = 0.0_dp
2680 IF (root_rank) exc_tmp = -exc
2681 tmp = grid_tmp + atomic_tmp + exc_tmp
2682
2683 IF (my_print_components .AND. root_rank) THEN
2684 iw = cp_logger_get_default_io_unit()
2685 IF (iw > 0) THEN
2686 WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight grid", grid_tmp
2687 WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight atomic", atomic_tmp
2688 WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight final", exc_tmp
2689 WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight residual", tmp
2690 END IF
2691 END IF
2692
2693 DO idir = 1, 3
2694 virial_xc(idir, idir) = virial_xc(idir, idir) + tmp
2695 END DO
2696
2697 CALL torch_tensor_release(grid_weight_grad_t)
2698 CALL torch_tensor_release(atomic_grid_weight_grad_t)
2699
2700 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 SUBROUTINE build_vxc_from_feature_grads(vxc_rho, vxc_tau, rho_r, pw_pool, &
2715 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 TYPE(pw_r3d_rs_type), DIMENSION(3) :: grad_pw
2731
2732 nspins = SIZE(rho_r)
2733 bo = rho_r(1)%pw_grid%bounds_local
2734 dvol_inv = 1.0_dp/rho_r(1)%pw_grid%dvol
2735 my_global_grid_layout = .false.
2736 IF (PRESENT(global_grid_layout)) my_global_grid_layout = global_grid_layout
2737 IF (my_global_grid_layout) THEN
2738 cpassert(SIZE(density_grad, 1) == product(rho_r(1)%pw_grid%npts))
2739 END IF
2740 cpassert(SIZE(kin_grad, 2) == SIZE(density_grad, 2))
2741 cpassert(SIZE(grad_grad, 3) == SIZE(density_grad, 2))
2742 cpassert(SIZE(grad_grad, 2) == 3)
2743 valid_spin_shape = SIZE(density_grad, 2) == nspins .OR. &
2744 (nspins == 1 .AND. SIZE(density_grad, 2) == 2)
2745 cpassert(valid_spin_shape)
2746
2747 ALLOCATE (vxc_rho(nspins), vxc_tau(nspins))
2748 DO ispin = 1, nspins
2749 CALL pw_pool%create_pw(vxc_rho(ispin))
2750 CALL pw_pool%create_pw(vxc_tau(ispin))
2751 CALL pw_zero(vxc_rho(ispin))
2752 CALL pw_zero(vxc_tau(ispin))
2753 END DO
2754
2755 IF (xc_requires_tmp_g(xc_deriv_method_id) .OR. rho_r(1)%pw_grid%spherical) THEN
2756 CALL pw_pool%create_pw(vxc_g)
2757 IF (.NOT. rho_r(1)%pw_grid%spherical) CALL pw_pool%create_pw(tmp_g)
2758 END IF
2759
2760 DO ispin = 1, nspins
2761 DO i = 1, 3
2762 CALL pw_pool%create_pw(grad_pw(i))
2763 CALL pw_zero(grad_pw(i))
2764 END DO
2765
2766 ipt = 0
2767 DO k = bo(1, 3), bo(2, 3)
2768 DO j = bo(1, 2), bo(2, 2)
2769 DO i = bo(1, 1), bo(2, 1)
2770 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 k - rho_r(1)%pw_grid%bounds(1, 3)))
2776 ELSE
2777 ipt = ipt + 1
2778 END IF
2779 IF (nspins == 1) THEN
2780 IF (SIZE(density_grad, 2) == 1) THEN
2781 vxc_rho(1)%array(i, j, k) = 0.5_dp*dvol_inv*density_grad(ipt, 1)
2782 vxc_tau(1)%array(i, j, k) = 0.5_dp*dvol_inv*kin_grad(ipt, 1)
2783 grad_pw(1)%array(i, j, k) = 0.5_dp*dvol_inv*grad_grad(ipt, 1, 1)
2784 grad_pw(2)%array(i, j, k) = 0.5_dp*dvol_inv*grad_grad(ipt, 2, 1)
2785 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 (density_grad(ipt, 1) + density_grad(ipt, 2))
2789 vxc_tau(1)%array(i, j, k) = 0.5_dp*dvol_inv* &
2790 (kin_grad(ipt, 1) + kin_grad(ipt, 2))
2791 grad_pw(1)%array(i, j, k) = 0.5_dp*dvol_inv* &
2792 (grad_grad(ipt, 1, 1) + grad_grad(ipt, 1, 2))
2793 grad_pw(2)%array(i, j, k) = 0.5_dp*dvol_inv* &
2794 (grad_grad(ipt, 2, 1) + grad_grad(ipt, 2, 2))
2795 grad_pw(3)%array(i, j, k) = 0.5_dp*dvol_inv* &
2796 (grad_grad(ipt, 3, 1) + grad_grad(ipt, 3, 2))
2797 END IF
2798 ELSE
2799 vxc_rho(ispin)%array(i, j, k) = dvol_inv*density_grad(ipt, ispin)
2800 vxc_tau(ispin)%array(i, j, k) = dvol_inv*kin_grad(ipt, ispin)
2801 grad_pw(1)%array(i, j, k) = dvol_inv*grad_grad(ipt, 1, ispin)
2802 grad_pw(2)%array(i, j, k) = dvol_inv*grad_grad(ipt, 2, ispin)
2803 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 DO i = 1, 3
2810 CALL pw_scale(grad_pw(i), -1.0_dp)
2811 END DO
2812 CALL xc_pw_divergence(xc_deriv_method_id, grad_pw, tmp_g, vxc_g, vxc_rho(ispin))
2813
2814 DO i = 1, 3
2815 CALL pw_pool%give_back_pw(grad_pw(i))
2816 END DO
2817 END DO
2818
2819 IF (ASSOCIATED(vxc_g%pw_grid)) CALL pw_pool%give_back_pw(vxc_g)
2820 IF (ASSOCIATED(tmp_g%pw_grid)) CALL pw_pool%give_back_pw(tmp_g)
2821
2822 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 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 IF (.NOT. print_active) RETURN
2838
2839 iw = cp_logger_get_default_io_unit()
2840 IF (iw <= 0) RETURN
2841 WRITE (unit=iw, fmt="(/,T2,A,1X,ES19.11)") &
2842 "SKALA_GPW| Native grid feature electrons", features%electron_count
2843 WRITE (unit=iw, fmt="(T2,A,1X,ES19.11)") &
2844 "SKALA_GPW| Native grid feature spin moment", features%spin_moment
2845 WRITE (unit=iw, fmt="(T2,A,1X,ES19.11)") &
2846 "SKALA_GPW| Native grid feature kinetic integral", features%kinetic_integral
2847 WRITE (unit=iw, fmt="(T2,A,1X,ES19.11)") &
2848 "SKALA_GPW| Native grid feature weight sum", features%grid_weight_sum
2849 IF (ALLOCATED(features%atomic_grid_sizes)) THEN
2850 atom_rows_min = int(minval(features%atomic_grid_sizes))
2851 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 "SKALA_GPW| Native grid atom row range", atom_rows_min, "to", &
2854 atom_rows_max, "sum", int(sum(features%atomic_grid_sizes))
2855 END IF
2856 IF (features%uses_atom_chunks) THEN
2857 WRITE (unit=iw, fmt="(T2,A,1X,I0,1X,A,1X,I0)") &
2858 "SKALA_GPW| Native grid atom chunk rows", features%chunk_feature_count, &
2859 "of", features%nflat
2860 IF (ALLOCATED(features%chunk_grad_counts)) THEN
2861 chunk_rows_min = minval(features%chunk_grad_counts)/ngrad_per_point
2862 chunk_rows_max = maxval(features%chunk_grad_counts)/ngrad_per_point
2863 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 "SKALA_GPW| Native grid atom chunk row range", chunk_rows_min, &
2866 "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 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 INTEGER, ALLOCATABLE, DIMENSION(:) :: selected_devices
2887
2888 selected_device = -1
2889
2890 IF (.NOT. use_cuda) RETURN
2891
2892 IF (.NOT. torch_cuda_is_available()) THEN
2893 cuda_device_count = 0
2894 ELSE
2895 cuda_device_count = torch_cuda_device_count()
2896 END IF
2897 IF (cuda_device_count > 0) THEN
2898 IF (requested_device < 0) THEN
2899 selected_device = mod(group%mepos, cuda_device_count)
2900 ELSE
2901 selected_device = requested_device
2902 END IF
2903 END IF
2904 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 "Torch CUDA device range.")
2908 END IF
2909 IF (selected_device >= 0) CALL offload_set_chosen_device(selected_device)
2910
2911 ALLOCATE (selected_devices(group%num_pe))
2912 CALL group%allgather(selected_device, selected_devices)
2913
2914 IF (group%mepos /= 0) THEN
2915 DEALLOCATE (selected_devices)
2916 RETURN
2917 END IF
2918 IF (selected_device == logged_cuda_device .AND. &
2919 cuda_device_count == logged_cuda_device_count .AND. &
2920 group%num_pe == logged_cuda_nproc .AND. &
2921 requested_device == logged_cuda_request) THEN
2922 DEALLOCATE (selected_devices)
2923 RETURN
2924 END IF
2925
2926 iw = cp_logger_get_default_io_unit()
2927 IF (iw <= 0) THEN
2928 DEALLOCATE (selected_devices)
2929 RETURN
2930 END IF
2931 IF (selected_device >= 0) THEN
2932 WRITE (unit=iw, fmt="(/,T2,A,1X,I0,1X,A,1X,I0,1X,A,1X,I0)") &
2933 "SKALA_GPW| Native grid Torch CUDA device", selected_device, &
2934 "of", cuda_device_count, "requested", requested_device
2935 ELSE
2936 WRITE (unit=iw, fmt="(/,T2,A)") &
2937 "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 "SKALA_GPW| Native grid Torch CUDA rank devices"
2941 DO pe = 1, group%num_pe
2942 WRITE (unit=iw, fmt="(1X,I0,A,I0)", advance="NO") pe - 1, ":", selected_devices(pe)
2943 END DO
2944 WRITE (unit=iw, fmt=*)
2945
2946 logged_cuda_device = selected_device
2947 logged_cuda_device_count = cuda_device_count
2948 logged_cuda_nproc = group%num_pe
2949 logged_cuda_request = requested_device
2950 DEALLOCATE (selected_devices)
2951
2952 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 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 gauxc_section => get_gauxc_section(xc_section)
2969 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_MODEL_FP64", l_val=fp64)
2970
2971 IF (cached_model_loaded) THEN
2972 IF (trim(cached_model_path) == trim(model_path) .AND. &
2973 cached_model_cuda_device == cuda_device .AND. &
2974 (cached_model_fp64 .EQV. fp64)) RETURN
2975 CALL skala_torch_model_release(cached_model)
2976 cached_model_loaded = .false.
2977 END IF
2978
2979 CALL skala_torch_model_load(cached_model, trim(model_path), fp64=fp64)
2980 cached_model_path = model_path
2981 cached_model_cuda_device = cuda_device
2982 cached_model_fp64 = fp64
2983 cached_model_loaded = .true.
2984
2985 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 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 gauxc_section => get_gauxc_section(xc_section)
3002 IF (.NOT. ASSOCIATED(gauxc_section)) THEN
3003 cpabort("Native SKALA GPW requires an XC_FUNCTIONAL%GAUXC section")
3004 END IF
3005
3006 CALL section_vals_val_get(gauxc_section, "MODEL", c_val=model_path)
3007 model_key = adjustl(model_path)
3008 CALL uppercase(model_key)
3009 IF (trim(model_key) == "NONE" .OR. trim(model_key) == "") THEN
3010 cpabort("Native SKALA GPW requires GAUXC%MODEL SKALA or a TorchScript model path")
3011 ELSE IF (trim(model_key) == "SKALA") THEN
3012 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_USE_CUDA", l_val=native_grid_use_cuda)
3013 IF (native_grid_use_cuda) THEN
3014 CALL get_environment_variable("GAUXC_SKALA_CUDA_MODEL", model_path, status=env_status)
3015 IF (env_status == 0 .AND. len_trim(model_path) > 0) RETURN
3016 END IF
3017 CALL get_environment_variable("SKALA_MODEL", model_path, status=env_status)
3018 IF (env_status /= 0 .OR. len_trim(model_path) == 0) THEN
3019 IF (native_grid_use_cuda) THEN
3020 CALL cp_abort(__location__, &
3021 "MODEL SKALA CUDA path requires GAUXC_SKALA_CUDA_MODEL or SKALA_MODEL")
3022 ELSE
3023 CALL cp_abort(__location__, &
3024 "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 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 NULLIFY (gauxc_section)
3044 IF (.NOT. ASSOCIATED(xc_section)) RETURN
3045
3046 functionals => section_vals_get_subs_vals(xc_section, "XC_FUNCTIONAL")
3047 IF (.NOT. ASSOCIATED(functionals)) RETURN
3048
3049 ifun = 0
3050 DO
3051 ifun = ifun + 1
3052 xc_fun => section_vals_get_subs_vals2(functionals, i_section=ifun)
3053 IF (.NOT. ASSOCIATED(xc_fun)) EXIT
3054 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
3062END MODULE skala_gpw_functional
cell_types::pbc
Definition cell_types.F:105
pw_methods::pw_scale
Definition pw_methods.F:94
pw_methods::pw_zero
Definition pw_methods.F:83
torch_api::torch_tensor_data_ptr
Definition torch_api.F:65
torch_api::torch_tensor_from_array
Definition torch_api.F:44
torch_api::torch_tensor_reset_from_array
Definition torch_api.F:59
cell_types
Handles all functions related to the CELL.
Definition cell_types.F:15
cp_array_utils
various utilities that regard array of different kinds: output, allocation,... maybe it is not a good...
Definition cp_array_utils.F:18
cp_log_handling
various routines to log and control the output. The idea is that decisions about where to log should ...
Definition cp_log_handling.F:41
cp_log_handling::cp_logger_get_default_io_unit
integer function, public cp_logger_get_default_io_unit(logger)
returns the unit nr for the ionode (-1 on all other processors) skips as well checks if the procs cal...
Definition cp_log_handling.F:523
input_section_types
objects that represent the structure of input sections and the data contained in an input section
Definition input_section_types.F:15
input_section_types::section_get_rval
real(kind=dp) function, public section_get_rval(section_vals, keyword_name)
...
Definition input_section_types.F:953
input_section_types::section_vals_get_subs_vals2
type(section_vals_type) function, pointer, public section_vals_get_subs_vals2(section_vals, i_section, i_rep_section)
returns the values of the n-th non default subsection (null if no such section exists (not so many no...
Definition input_section_types.F:791
input_section_types::section_vals_get_subs_vals
recursive type(section_vals_type) function, pointer, public section_vals_get_subs_vals(section_vals, subsection_name, i_rep_section, can_return_null)
returns the values of the requested subsection
Definition input_section_types.F:735
input_section_types::section_vals_val_get
subroutine, public section_vals_val_get(section_vals, keyword_name, i_rep_section, i_rep_val, n_rep_val, val, l_val, i_val, r_val, c_val, l_vals, i_vals, r_vals, c_vals, explicit)
returns the requested value
Definition input_section_types.F:1052
kinds
Defines the basic variable types.
Definition kinds.F:23
kinds::int_8
integer, parameter, public int_8
Definition kinds.F:54
kinds::dp
integer, parameter, public dp
Definition kinds.F:34
kinds::default_path_length
integer, parameter, public default_path_length
Definition kinds.F:58
message_passing
Interface to the message passing library MPI.
Definition message_passing.F:27
offload_api
Fortran API for the offload package, which is written in C.
Definition offload_api.F:12
offload_api::offload_set_chosen_device
subroutine, public offload_set_chosen_device(device_id)
Selects the chosen device to be used.
Definition offload_api.F:132
particle_types
Define the data structure for the particle information.
Definition particle_types.F:19
pw_grid_types
Definition pw_grid_types.F:13
pw_methods
Definition pw_methods.F:25
pw_pool_types
Manages a pool of grids (to be used for example as tmp objects), but can also be used to instantiate ...
Definition pw_pool_types.F:24
pw_types
Definition pw_types.F:24
qs_grid_atom
Definition qs_grid_atom.F:8
skala_gpw_features
Build SKALA TorchScript feature dictionaries from CP2K GPW real-space grids.
Definition skala_gpw_features.F:11
skala_gpw_features::skala_gpw_atom_subchunk_layout
subroutine, public skala_gpw_atom_subchunk_layout(max_rows, atom_begin, atom_count, row_begin, row_count)
Build atom and row bounds for all atom-contiguous subchunks in one pass.
Definition skala_gpw_features.F:1871
skala_gpw_features::skala_gpw_atom_partition_hard
integer, parameter, public skala_gpw_atom_partition_hard
Definition skala_gpw_features.F:38
skala_gpw_features::smooth_partition_atomic_weight_scale_derivative
pure real(kind=dp) function, public smooth_partition_atomic_weight_scale_derivative(weight)
Derivative of the sparse atom-row quadrature taper with respect to partition weight.
Definition skala_gpw_features.F:3244
skala_gpw_features::skala_gpw_feature_release
subroutine, public skala_gpw_feature_release(features)
Release Torch objects and backing arrays owned by a feature bundle.
Definition skala_gpw_features.F:1751
skala_gpw_features::skala_gpw_atom_subchunk_count
integer function, public skala_gpw_atom_subchunk_count(max_rows)
Return how many atom-contiguous subchunks the cached rank chunk needs.
Definition skala_gpw_features.F:1834
skala_gpw_features::skala_gpw_feature_build
subroutine, public skala_gpw_feature_build(features, rho_set, rho_r, particle_set, cell, requires_grad, weights, requires_coordinate_grad, requires_stress_grad, use_atom_chunks, route_atom_chunks, atom_partition)
Build a flat SKALA molecular feature dictionary from a local GPW grid.
Definition skala_gpw_features.F:204
skala_gpw_features::skala_gpw_feature_build_atom_subchunk_bounds
subroutine, public skala_gpw_feature_build_atom_subchunk_bounds(parent, features, atom_begin, atom_count, row_begin, row_count, requires_grad)
Build an atom-contiguous subchunk feature bundle from precomputed bounds.
Definition skala_gpw_features.F:1996
skala_gpw_features::skala_gpw_smooth_partition_derivatives
subroutine, public skala_gpw_smooth_partition_derivatives(grid_point, atom_coords, cell, weights, included, dweights_datom, dweights_dstrain)
Build smooth atom weights and their atom/cell deformation derivatives.
Definition skala_gpw_features.F:3048
skala_gpw_features::skala_gpw_atom_partition_smooth
integer, parameter, public skala_gpw_atom_partition_smooth
Definition skala_gpw_features.F:38
skala_gpw_functional
Experimental CP2K-native GPW real-space-grid path for SKALA TorchScript models.
Definition skala_gpw_functional.F:11
skala_gpw_functional::skala_gapw_density_partition_soft_only
integer, parameter, public skala_gapw_density_partition_soft_only
Definition skala_gpw_functional.F:78
skala_gpw_functional::skala_gpw_weight_derivative
subroutine, public skala_gpw_weight_derivative(weight_deriv_r, rho_r, rho_g, tau, xc_section, weights, pw_pool, particle_set, cell)
Evaluate the derivative of native SKALA XC energy with respect to CP2K's external real-space integrat...
Definition skala_gpw_functional.F:717
skala_gpw_functional::skala_gapw_atom_composite_energy
subroutine, public skala_gapw_atom_composite_energy(xc_section, group, density, grad, kin, grid_coords, grid_weights, atomic_grid_weights, atomic_grid_sizes, atomic_coords, exc, density_grad_out, grad_grad_out, kin_grad_out, grid_coord_grad_out, grid_weight_grad_out, atomic_grid_weight_grad_out, atom_coord_grad_out)
Evaluate a rank-local set of complete atom blocks and sum their SKALA energies.
Definition skala_gpw_functional.F:1157
skala_gpw_functional::ensure_native_skala_grid_scope
subroutine, public ensure_native_skala_grid_scope(xc_section)
Enforce the currently implemented native SKALA GPW input scope.
Definition skala_gpw_functional.F:337
skala_gpw_functional::skala_gpw_eval
subroutine, public skala_gpw_eval(vxc_rho, vxc_tau, exc, rho_r, rho_g, tau, xc_section, weights, pw_pool, particle_set, cell, compute_virial, virial_xc, just_energy, atom_force)
Evaluate SKALA energy and first derivatives on a CP2K GPW grid.
Definition skala_gpw_functional.F:405
skala_gpw_functional::build_vxc_from_feature_grads
subroutine, public build_vxc_from_feature_grads(vxc_rho, vxc_tau, rho_r, pw_pool, density_grad, grad_grad, kin_grad, xc_deriv_method_id, global_grid_layout)
Fill CP2K VXC real-space arrays from Torch feature gradients.
Definition skala_gpw_functional.F:2717
skala_gpw_functional::skala_gapw_density_partition_none
integer, parameter, public skala_gapw_density_partition_none
Definition skala_gpw_functional.F:78
skala_gpw_functional::native_skala_gapw_composite_direct_ao
logical function, public native_skala_gapw_composite_direct_ao(xc_section)
Return true if the GAPW composite reference uses direct full-ORB collocation.
Definition skala_gpw_functional.F:259
skala_gpw_functional::xc_section_uses_gauxc_model
logical function, public xc_section_uses_gauxc_model(xc_section)
Return true if the GAUXC subsection requests a model evaluation.
Definition skala_gpw_functional.F:280
skala_gpw_functional::native_skala_uses_atom_composite_grid
logical function, public native_skala_uses_atom_composite_grid(xc_section)
Return true when native Skala uses atom-centered grids.
Definition skala_gpw_functional.F:229
skala_gpw_functional::xc_section_uses_native_skala_evaluator
logical function, public xc_section_uses_native_skala_evaluator(xc_section)
Return true when SKALA must be evaluated by the CP2K-native grid machinery.
Definition skala_gpw_functional.F:165
skala_gpw_functional::skala_gapw_density_partition_hard_minus_soft
integer, parameter, public skala_gapw_density_partition_hard_minus_soft
Definition skala_gpw_functional.F:78
skala_gpw_functional::skala_gapw_representation
integer function, public skala_gapw_representation(xc_section)
Return the pseudopotential GAPW representation selected for an active model.
Definition skala_gpw_functional.F:136
skala_gpw_functional::native_skala_gapw_density_partition
integer function, public native_skala_gapw_density_partition(xc_section)
Return the hard/soft GAPW one-center density partition for native SKALA.
Definition skala_gpw_functional.F:307
skala_gpw_functional::get_gauxc_section
type(section_vals_type) function, pointer, public get_gauxc_section(xc_section)
Return the first GAUXC functional subsection, if present.
Definition skala_gpw_functional.F:3037
skala_gpw_functional::skala_gapw_atom_vxc_of_r
subroutine, public skala_gapw_atom_vxc_of_r(xc_section, grid_atom, group, atom_coord, rho, drho, tau, weights, lsd, nspins, na, nr, exc, vxc, vxg, vtau, energy_only, atom_force, atom_virial)
Evaluate SKALA on a GAPW one-center atomic grid.
Definition skala_gpw_functional.F:895
skala_gpw_functional::native_skala_gapw_atom_composite_requested
logical function, public native_skala_gapw_atom_composite_requested(xc_section)
Return true when the explicit atom-centered composite reference is requested.
Definition skala_gpw_functional.F:204
skala_gpw_functional::cached_model
type(skala_torch_model_type), save cached_model
Definition skala_gpw_functional.F:98
skala_gpw_functional::native_skala_gapw_composite_reference
logical function, public native_skala_gapw_composite_reference(xc_section)
Return true if native SKALA should use the full GAPW ORB density on one common grid.
Definition skala_gpw_functional.F:180
skala_gpw_functional::xc_section_uses_native_skala_grid
logical function, public xc_section_uses_native_skala_grid(xc_section)
Return true if the GAUXC subsection requests the CP2K-native GPW grid path.
Definition skala_gpw_functional.F:117
skala_gpw_functional::skala_gapw_density_partition_hard_only
integer, parameter, public skala_gapw_density_partition_hard_only
Definition skala_gpw_functional.F:78
skala_torch_api
Small CP2K wrapper around the SKALA TorchScript functional protocol.
Definition skala_torch_api.F:11
skala_torch_api::skala_torch_model_release
subroutine, public skala_torch_model_release(model)
Release a loaded SKALA TorchScript model.
Definition skala_torch_api.F:88
skala_torch_api::skala_torch_model_load
subroutine, public skala_torch_model_load(model, filename, fp64)
Load a SKALA TorchScript model and its feature metadata.
Definition skala_torch_api.F:55
skala_torch_api::skala_torch_model_get_exc
subroutine, public skala_torch_model_get_exc(model, inputs, grid_weights, exc_tensor, exc)
Evaluate the weighted SKALA exchange-correlation energy.
Definition skala_torch_api.F:173
string_utilities
Utilities for string manipulations.
Definition string_utilities.F:16
string_utilities::uppercase
elemental subroutine, public uppercase(string)
Convert all lower case characters in a string to upper case.
Definition string_utilities.F:3362
torch_api
Definition torch_api.F:7
torch_api::torch_dict_release
subroutine, public torch_dict_release(dict)
Releases a Torch dictionary and all its ressources.
Definition torch_api.F:1843
torch_api::torch_use_cuda
subroutine, public torch_use_cuda(use_cuda)
Select whether Torch wrappers should use CUDA when available.
Definition torch_api.F:1555
torch_api::torch_tensor_backward_scalar
subroutine, public torch_tensor_backward_scalar(tensor)
Runs autograd on a scalar Torch tensor.
Definition torch_api.F:1504
torch_api::torch_tensor_narrow
subroutine, public torch_tensor_narrow(tensor, dim, start_index, length, result)
Creates a view of a contiguous tensor slice.
Definition torch_api.F:1431
torch_api::torch_tensor_to_device_leaf
subroutine, public torch_tensor_to_device_leaf(tensor, requires_grad)
Moves a tensor to the active Torch device and makes it an autograd leaf.
Definition torch_api.F:1527
torch_api::torch_dict_create
subroutine, public torch_dict_create(dict)
Creates an empty Torch dictionary.
Definition torch_api.F:1726
torch_api::torch_tensor_grad
subroutine, public torch_tensor_grad(tensor, grad)
Returns the gradient of a Torch tensor which was computed by autograd.
Definition torch_api.F:1576
torch_api::torch_tensor_grad_batch3
subroutine, public torch_tensor_grad_batch3(tensor1, tensor2, tensor3, grad1, grad2, grad3)
Copies three autograd gradients to CPU memory.
Definition torch_api.F:1604
torch_api::torch_cuda_device_count
integer function, public torch_cuda_device_count()
Return the number of CUDA devices visible to Torch.
Definition torch_api.F:2229
torch_api::torch_dict_insert
subroutine, public torch_dict_insert(dict, key, tensor)
Inserts a Torch tensor into a Torch dictionary.
Definition torch_api.F:1777
torch_api::torch_cuda_is_available
logical function, public torch_cuda_is_available()
Returns true iff the Torch CUDA backend is available.
Definition torch_api.F:2208
torch_api::torch_tensor_release
subroutine, public torch_tensor_release(tensor)
Releases a Torch tensor and all its ressources.
Definition torch_api.F:1702
xc_input_constants
input constants for xc
Definition xc_input_constants.F:10
xc_input_constants::skala_gapw_atom_composite_grid
integer, parameter, public skala_gapw_atom_composite_grid
Definition xc_input_constants.F:34
xc_input_constants::skala_gapw_direct_valence
integer, parameter, public skala_gapw_direct_valence
Definition xc_input_constants.F:29
xc_input_constants::skala_gapw_paw_one_center
integer, parameter, public skala_gapw_paw_one_center
Definition xc_input_constants.F:29
xc_input_constants::skala_gapw_common_grid
integer, parameter, public skala_gapw_common_grid
Definition xc_input_constants.F:34
xc_input_constants::skala_gapw_cp2k_default
integer, parameter, public skala_gapw_cp2k_default
Definition xc_input_constants.F:29
xc_input_constants::skala_gapw_paw_one_center_split
integer, parameter, public skala_gapw_paw_one_center_split
Definition xc_input_constants.F:29
xc_rho_cflags_types
contains the structure
Definition xc_rho_cflags_types.F:14
xc_rho_set_types
contains the structure
Definition xc_rho_set_types.F:14
xc_rho_set_types::xc_rho_set_create
subroutine, public xc_rho_set_create(rho_set, local_bounds, rho_cutoff, drho_cutoff, tau_cutoff)
allocates and does (minimal) initialization of a rho_set
Definition xc_rho_set_types.F:111
xc_rho_set_types::xc_rho_set_release
subroutine, public xc_rho_set_release(rho_set, pw_pool)
releases the given rho_set
Definition xc_rho_set_types.F:129
xc_rho_set_types::xc_rho_set_update
subroutine, public xc_rho_set_update(rho_set, rho_r, rho_g, tau, needs, xc_deriv_method_id, xc_rho_smooth_id, pw_pool, spinflip)
updates the given rho set with the density given by rho_r (and rho_g). The rho set will contain the c...
Definition xc_rho_set_types.F:691
xc_rho_set_types::xc_rho_set_get
subroutine, public xc_rho_set_get(rho_set, can_return_null, rho, drho, norm_drho, rhoa, rhob, norm_drhoa, norm_drhob, rho_1_3, rhoa_1_3, rhob_1_3, laplace_rho, laplace_rhoa, laplace_rhob, drhoa, drhob, rho_cutoff, drho_cutoff, tau_cutoff, tau, tau_a, tau_b, local_bounds)
returns the various attributes of rho_set
Definition xc_rho_set_types.F:281
xc_util
contains utility functions for the xc package
Definition xc_util.F:14
xc_util::xc_pw_divergence
subroutine, public xc_pw_divergence(xc_deriv_method_id, pw_to_deriv, tmp_g, vxc_g, vxc_r)
Calculates the divergence of pw_to_deriv.
Definition xc_util.F:253
xc_util::xc_requires_tmp_g
elemental logical function, public xc_requires_tmp_g(xc_deriv_id)
...
Definition xc_util.F:58
cell_types::cell_type
Type defining parameters related to the simulation cell.
Definition cell_types.F:60
cp_array_utils::cp_3d_r_cp_type
represent a pointer to a contiguous 3d array
Definition cp_array_utils.F:149
input_section_types::section_vals_type
stores the values of a section
Definition input_section_types.F:127
message_passing::mp_comm_type
Definition message_passing.F:179
particle_types::particle_type
Definition particle_types.F:35
pw_grid_types::pw_grid_type
Definition pw_grid_types.F:53
pw_pool_types::pw_pool_type
Manages a pool of grids (to be used for example as tmp objects), but can also be used to instantiate ...
Definition pw_pool_types.F:83
pw_types::pw_c1d_gs_type
Definition pw_types.F:127
pw_types::pw_r3d_rs_type
Definition pw_types.F:62
qs_grid_atom::grid_atom_type
Definition qs_grid_atom.F:44
skala_gpw_features::skala_gpw_feature_type
Definition skala_gpw_features.F:128
skala_torch_api::skala_torch_model_type
Definition skala_torch_api.F:38
torch_api::torch_dict_type
Definition torch_api.F:34
torch_api::torch_tensor_type
Definition torch_api.F:29
xc_rho_cflags_types::xc_rho_cflags_type
contains a flag for each component of xc_rho_set, so that you can use it to tell which components you...
Definition xc_rho_cflags_types.F:48
xc_rho_set_types::xc_rho_set_type
represent a density, with all the representation and data needed to perform a functional evaluation
Definition xc_rho_set_types.F:78