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