Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2026 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : ! **************************************************************************************************
9 : !> \brief Experimental CP2K-native GPW real-space-grid path for SKALA TorchScript models.
10 : ! **************************************************************************************************
11 : MODULE skala_gpw_functional
12 : USE OMP_LIB, ONLY: omp_get_max_threads,&
13 : omp_set_num_threads
14 : USE cell_types, ONLY: cell_type,&
15 : pbc
16 : USE cp_array_utils, ONLY: cp_3d_r_cp_type
17 : USE cp_log_handling, ONLY: cp_logger_get_default_io_unit
18 : USE input_section_types, ONLY: section_get_rval,&
19 : section_vals_get_subs_vals,&
20 : section_vals_get_subs_vals2,&
21 : section_vals_type,&
22 : section_vals_val_get
23 : USE kinds, ONLY: default_path_length,&
24 : dp,&
25 : int_8
26 : USE message_passing, ONLY: mp_comm_type
27 : USE offload_api, ONLY: offload_set_chosen_device
28 : USE particle_types, ONLY: particle_type
29 : USE pw_grid_types, ONLY: pw_grid_type
30 : USE pw_methods, ONLY: pw_scale,&
31 : pw_zero
32 : USE pw_pool_types, ONLY: pw_pool_type
33 : USE pw_types, ONLY: pw_c1d_gs_type,&
34 : pw_r3d_rs_type
35 : USE qs_grid_atom, ONLY: grid_atom_type
36 : USE skala_gpw_features, ONLY: &
37 : skala_gpw_atom_partition_hard, skala_gpw_atom_partition_smooth, &
38 : skala_gpw_atom_subchunk_count, skala_gpw_atom_subchunk_layout, skala_gpw_feature_build, &
39 : skala_gpw_feature_build_atom_subchunk_bounds, skala_gpw_feature_release, &
40 : skala_gpw_feature_type, skala_gpw_smooth_partition_derivatives, &
41 : smooth_partition_atomic_weight_scale_derivative
42 : USE skala_torch_api, ONLY: skala_torch_model_get_exc,&
43 : skala_torch_model_load,&
44 : skala_torch_model_release,&
45 : skala_torch_model_type
46 : USE string_utilities, ONLY: uppercase
47 : USE torch_api, ONLY: &
48 : torch_cuda_device_count, torch_cuda_is_available, torch_dict_create, torch_dict_insert, &
49 : torch_dict_release, torch_dict_type, torch_tensor_backward_scalar, torch_tensor_data_ptr, &
50 : torch_tensor_from_array, torch_tensor_grad, torch_tensor_grad_batch3, torch_tensor_narrow, &
51 : torch_tensor_release, torch_tensor_reset_from_array, torch_tensor_to_device_leaf, &
52 : torch_tensor_type, torch_use_cuda
53 : USE xc_input_constants, ONLY: skala_gapw_atom_composite_grid,&
54 : skala_gapw_common_grid,&
55 : skala_gapw_cp2k_default,&
56 : skala_gapw_direct_valence,&
57 : skala_gapw_paw_one_center,&
58 : skala_gapw_paw_one_center_split
59 : USE xc_rho_cflags_types, ONLY: xc_rho_cflags_type
60 : USE xc_rho_set_types, ONLY: xc_rho_set_create,&
61 : xc_rho_set_get,&
62 : xc_rho_set_release,&
63 : xc_rho_set_type,&
64 : xc_rho_set_update
65 : USE xc_util, ONLY: xc_pw_divergence,&
66 : xc_requires_tmp_g
67 : #include "./base/base_uses.f90"
68 :
69 : IMPLICIT NONE
70 :
71 : PRIVATE
72 :
73 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'skala_gpw_functional'
74 : INTEGER, PARAMETER, PRIVATE :: atom_chunk_auto_max_rows = 400000, &
75 : atom_chunk_auto_min_rows = 100000, &
76 : atom_chunk_auto_row_quantum = 100000, &
77 : ncollapsed_grad_per_point = 5, ngrad_per_point = 10
78 : INTEGER, PARAMETER, PUBLIC :: skala_gapw_density_partition_hard_minus_soft = 1, &
79 : skala_gapw_density_partition_hard_only = 2, &
80 : skala_gapw_density_partition_soft_only = 3, &
81 : skala_gapw_density_partition_none = 4
82 :
83 : TYPE skala_gapw_atom_cuda_tensor_cache_type
84 : TYPE(torch_tensor_type) :: atomic_coords_t, atomic_grid_weights_t, &
85 : density_t, grad_t, grid_coords_t, &
86 : grid_weights_t, kin_t
87 : END TYPE skala_gapw_atom_cuda_tensor_cache_type
88 :
89 : PUBLIC :: build_vxc_from_feature_grads, ensure_native_skala_grid_scope, get_gauxc_section, &
90 : skala_gapw_atom_composite_energy, skala_gapw_atom_vxc_of_r, &
91 : native_skala_gapw_atom_composite_requested, &
92 : native_skala_uses_atom_composite_grid, &
93 : native_skala_gapw_composite_direct_ao, native_skala_gapw_composite_reference, &
94 : native_skala_gapw_density_partition, skala_gpw_eval, skala_gpw_weight_derivative, &
95 : skala_gapw_representation, xc_section_uses_native_skala_evaluator, &
96 : xc_section_uses_native_skala_grid, xc_section_uses_gauxc_model
97 :
98 : TYPE(skala_torch_model_type), SAVE :: cached_model
99 : TYPE(skala_gapw_atom_cuda_tensor_cache_type), SAVE, TARGET :: cached_atom_cuda_tensors
100 : CHARACTER(len=default_path_length), SAVE :: cached_model_path = ""
101 : LOGICAL, SAVE :: cached_model_loaded = .FALSE.
102 : 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 :
108 : CONTAINS
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 293260 : 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 293260 : uses_native_grid = .FALSE.
122 293260 : gauxc_section => get_gauxc_section(xc_section)
123 293260 : IF (ASSOCIATED(gauxc_section)) THEN
124 1728 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID", l_val=uses_native_grid)
125 : END IF
126 :
127 293260 : 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 193057 : 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 192327 : IF (.NOT. xc_section_uses_gauxc_model(xc_section)) RETURN
142 :
143 730 : gauxc_section => get_gauxc_section(xc_section)
144 730 : CPASSERT(ASSOCIATED(gauxc_section))
145 : CALL section_vals_val_get(gauxc_section, "PSEUDOPOTENTIAL_GAPW_REPRESENTATION", &
146 730 : 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 0 : CONTINUE
152 : CASE DEFAULT
153 730 : 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 162027 : 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 162027 : skala_gapw_paw_one_center
170 :
171 162027 : 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 148703 : 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 148703 : use_composite_reference = .FALSE.
186 148703 : native_grid = .FALSE.
187 148703 : gauxc_section => get_gauxc_section(xc_section)
188 148703 : IF (ASSOCIATED(gauxc_section)) THEN
189 1668 : 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 1668 : l_val=use_composite_reference)
192 : END IF
193 148703 : use_composite_reference = native_grid .AND. use_composite_reference
194 :
195 148703 : 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 1826 : 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 1826 : use_atom_composite = .FALSE.
210 1826 : native_grid = .FALSE.
211 1826 : gauxc_section => get_gauxc_section(xc_section)
212 1826 : IF (ASSOCIATED(gauxc_section)) THEN
213 364 : 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 364 : l_val=use_atom_composite)
217 : END IF
218 1826 : use_atom_composite = native_grid .AND. use_atom_composite
219 :
220 1826 : 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 640 : 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 640 : use_atom_composite = .FALSE.
235 640 : composite_grid = skala_gapw_atom_composite_grid
236 640 : gauxc_section => get_gauxc_section(xc_section)
237 640 : IF (ASSOCIATED(gauxc_section)) THEN
238 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_LAYOUT", &
239 640 : i_val=composite_grid)
240 : END IF
241 : SELECT CASE (composite_grid)
242 : CASE (skala_gapw_common_grid)
243 576 : CONTINUE
244 : CASE (skala_gapw_atom_composite_grid)
245 576 : use_atom_composite = .TRUE.
246 : CASE DEFAULT
247 640 : CALL cp_abort(__LOCATION__, "Unknown native-grid layout.")
248 : END SELECT
249 :
250 640 : 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 8 : 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 8 : use_direct_ao = .FALSE.
264 8 : gauxc_section => get_gauxc_section(xc_section)
265 8 : IF (ASSOCIATED(gauxc_section)) THEN
266 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_GAPW_COMPOSITE_DIRECT_AO", &
267 8 : l_val=use_direct_ao)
268 : END IF
269 8 : use_direct_ao = native_skala_gapw_composite_reference(xc_section) .AND. use_direct_ao
270 :
271 8 : 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 225447 : 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 225447 : uses_gauxc_model = .FALSE.
286 225447 : gauxc_section => get_gauxc_section(xc_section)
287 225447 : IF (ASSOCIATED(gauxc_section)) THEN
288 1324 : CALL section_vals_val_get(gauxc_section, "MODEL", c_val=model_name)
289 1324 : CALL section_vals_val_get(gauxc_section, "FUNCTIONAL", c_val=xc_name)
290 1324 : model_key = ADJUSTL(model_name)
291 1324 : xc_key = ADJUSTL(xc_name)
292 1324 : CALL uppercase(model_key)
293 1324 : CALL uppercase(xc_key)
294 : uses_gauxc_model = (TRIM(model_key) /= "" .AND. TRIM(model_key) /= "NONE" .AND. &
295 1324 : TRIM(model_key) /= TRIM(xc_key))
296 : END IF
297 :
298 225447 : 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 282 : 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 282 : partition = skala_gapw_density_partition_hard_minus_soft
312 282 : gauxc_section => get_gauxc_section(xc_section)
313 282 : IF (ASSOCIATED(gauxc_section)) THEN
314 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_GAPW_DENSITY_PARTITION", &
315 282 : 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 0 : CONTINUE
324 : CASE DEFAULT
325 : CALL cp_abort(__LOCATION__, &
326 282 : "Unknown GAUXC%NATIVE_GRID_GAPW_DENSITY_PARTITION value.")
327 : END SELECT
328 :
329 282 : 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 624 : 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 312 : NULLIFY (gauxc_section)
344 312 : IF (.NOT. ASSOCIATED(xc_section)) THEN
345 0 : CPABORT("Native SKALA GPW requires an XC section")
346 : END IF
347 :
348 312 : functionals => section_vals_get_subs_vals(xc_section, "XC_FUNCTIONAL")
349 312 : IF (.NOT. ASSOCIATED(functionals)) THEN
350 0 : CPABORT("Native SKALA GPW requires an XC_FUNCTIONAL section")
351 : END IF
352 :
353 312 : nfun = 0
354 312 : ifun = 0
355 : DO
356 624 : ifun = ifun + 1
357 624 : xc_fun => section_vals_get_subs_vals2(functionals, i_section=ifun)
358 624 : IF (.NOT. ASSOCIATED(xc_fun)) EXIT
359 312 : nfun = nfun + 1
360 624 : IF (xc_fun%section%name == "GAUXC") gauxc_section => xc_fun
361 : END DO
362 :
363 312 : IF (.NOT. ASSOCIATED(gauxc_section)) THEN
364 0 : CPABORT("Native SKALA GPW requires an XC_FUNCTIONAL%GAUXC section")
365 : END IF
366 312 : IF (nfun /= 1) THEN
367 0 : CPABORT("Native SKALA GPW requires GAUXC to be the only XC functional")
368 : END IF
369 :
370 312 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID", l_val=native_grid)
371 312 : IF (.NOT. native_grid .AND. &
372 0 : .NOT. xc_section_uses_native_skala_evaluator(xc_section)) RETURN
373 :
374 312 : CALL section_vals_val_get(gauxc_section, "MODEL", c_val=model_name)
375 312 : model_key = ADJUSTL(model_name)
376 312 : CALL uppercase(model_key)
377 312 : IF (TRIM(model_key) == "NONE" .OR. TRIM(model_key) == "") THEN
378 0 : 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 48 : 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 48 : 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 48 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: density_grad, kin_grad
431 48 : 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 48 : 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 48 : virial_xc = 0.0_dp
442 48 : exc = 0.0_dp
443 48 : my_just_energy = .FALSE.
444 48 : IF (PRESENT(just_energy)) my_just_energy = just_energy
445 48 : needs_atom_force = PRESENT(atom_force)
446 112 : IF (needs_atom_force) atom_force = 0.0_dp
447 48 : have_atom_coord_grad = .FALSE.
448 48 : omp_max_threads_restore = omp_get_max_threads()
449 :
450 48 : IF (compute_virial .AND. my_just_energy) THEN
451 : CALL cp_abort(__LOCATION__, &
452 0 : "Native SKALA GPW stress/virial requires feature gradients.")
453 : END IF
454 48 : IF (.NOT. ASSOCIATED(rho_g)) THEN
455 : CALL cp_abort(__LOCATION__, &
456 0 : "Native SKALA GPW requires the reciprocal-space density to form density gradients.")
457 : END IF
458 48 : IF (.NOT. ASSOCIATED(tau)) THEN
459 : CALL cp_abort(__LOCATION__, &
460 0 : "Native SKALA GPW requires the kinetic-energy density.")
461 : END IF
462 :
463 48 : nspins = SIZE(rho_r)
464 48 : lsd = (nspins /= 1)
465 48 : CALL get_skala_model_path(xc_section, model_path)
466 48 : gauxc_section => get_gauxc_section(xc_section)
467 48 : 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 48 : i_val=native_grid_cuda_device)
470 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNKS", &
471 48 : l_val=native_grid_atom_chunks)
472 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNK_ROUTING", &
473 48 : l_val=native_grid_atom_chunk_routing)
474 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNK_MAX_ROWS", &
475 48 : i_val=native_grid_atom_chunk_max_rows)
476 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_PARTITION", &
477 48 : i_val=native_grid_atom_partition)
478 22 : SELECT CASE (native_grid_atom_partition)
479 : CASE (1)
480 22 : native_grid_atom_partition = skala_gpw_atom_partition_hard
481 : CASE (2)
482 26 : native_grid_atom_partition = skala_gpw_atom_partition_smooth
483 : CASE DEFAULT
484 : CALL cp_abort(__LOCATION__, &
485 48 : "Unknown GAUXC%NATIVE_GRID_ATOM_PARTITION value.")
486 : END SELECT
487 48 : native_grid_atom_chunk_routing = native_grid_atom_chunk_routing .OR. native_grid_atom_chunks
488 48 : native_grid_atom_chunks = native_grid_atom_chunks .OR. native_grid_atom_chunk_routing
489 48 : IF (native_grid_atom_chunk_max_rows < -1) THEN
490 : CALL cp_abort(__LOCATION__, &
491 0 : "GAUXC%NATIVE_GRID_ATOM_CHUNK_MAX_ROWS must be -1, zero, or positive.")
492 : END IF
493 48 : IF (needs_atom_force .OR. compute_virial) THEN
494 8 : IF (native_grid_atom_partition == skala_gpw_atom_partition_hard) THEN
495 0 : native_grid_atom_partition = skala_gpw_atom_partition_smooth
496 : END IF
497 8 : native_grid_atom_chunk_routing = .FALSE.
498 8 : native_grid_atom_chunks = .FALSE.
499 : END IF
500 : ! The portable SKALA export used by the regtests builds ragged-index tensors on CPU.
501 48 : CALL torch_use_cuda(native_grid_use_cuda)
502 : selected_cuda_device = configure_native_grid_cuda( &
503 48 : native_grid_use_cuda, native_grid_cuda_device, rho_r(1)%pw_grid%para%group)
504 48 : CALL ensure_model_loaded(model_path, selected_cuda_device)
505 :
506 48 : IF (lsd) THEN
507 6 : needs%rho_spin = .TRUE.
508 6 : needs%drho_spin = .TRUE.
509 6 : needs%tau_spin = .TRUE.
510 : ELSE
511 42 : needs%rho = .TRUE.
512 42 : needs%drho = .TRUE.
513 42 : needs%tau = .TRUE.
514 : END IF
515 :
516 48 : CALL section_vals_val_get(xc_section, "XC_GRID%XC_DERIV", i_val=xc_deriv_method_id)
517 48 : 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 48 : tau_cutoff=section_get_rval(xc_section, "tau_cutoff"))
524 : CALL xc_rho_set_update(rho_set, rho_r, rho_g, tau, needs, &
525 48 : 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 88 : atom_partition=native_grid_atom_partition)
534 48 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_DIAGNOSTICS", l_val=native_grid_diagnostics)
535 48 : IF (native_grid_diagnostics) THEN
536 24 : CALL print_native_grid_diagnostics(features, rho_r(1)%pw_grid%para%group%mepos == 0)
537 : END IF
538 :
539 48 : 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 36 : rho_r(1)%pw_grid%para%group)
542 : END IF
543 48 : IF (native_grid_diagnostics .AND. features%uses_atom_chunks .AND. &
544 : rho_r(1)%pw_grid%para%group%mepos == 0) THEN
545 12 : iw = cp_logger_get_default_io_unit()
546 12 : IF (iw > 0) THEN
547 : WRITE (UNIT=iw, FMT="(T2,A,1X,I0)") &
548 12 : "SKALA_GPW| Native grid atom chunk max rows", native_grid_atom_chunk_max_rows
549 : END IF
550 : END IF
551 48 : native_grid_atom_subchunks = 1
552 48 : IF (features%uses_atom_chunks .AND. native_grid_atom_chunk_max_rows > 0) THEN
553 40 : native_grid_atom_subchunks = skala_gpw_atom_subchunk_count(native_grid_atom_chunk_max_rows)
554 40 : CALL rho_r(1)%pw_grid%para%group%max(native_grid_atom_subchunks)
555 : END IF
556 48 : use_atom_subchunks = features%uses_atom_chunks .AND. native_grid_atom_subchunks > 1
557 48 : has_atom_chunk_work = .NOT. features%uses_atom_chunks .OR. features%chunk_feature_count > 0
558 48 : exc = 0.0_dp
559 48 : 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 2 : kin_grad=kin_grad, collapse_spin_grads=(nspins == 1))
565 46 : ELSE IF (has_atom_chunk_work) THEN
566 : CALL skala_torch_model_get_exc(cached_model, features%inputs, &
567 46 : features%grid_weights_t, exc_tensor, exc)
568 : END IF
569 48 : IF (features%uses_atom_chunks) CALL rho_r(1)%pw_grid%para%group%sum(exc)
570 :
571 48 : IF (.NOT. my_just_energy) THEN
572 48 : IF (.NOT. use_atom_subchunks) THEN
573 46 : IF (has_atom_chunk_work) THEN
574 46 : CALL timeset("skala_gpw_backward", phase_handle)
575 46 : CALL torch_tensor_backward_scalar(exc_tensor)
576 46 : CALL timestop(phase_handle)
577 :
578 46 : IF (compute_virial) THEN
579 6 : 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 6 : native_grid_diagnostics)
584 6 : IF (native_grid_diagnostics) THEN
585 : CALL print_virial_delta("weight-residual", virial_xc - virial_before, &
586 0 : rho_r(1)%pw_grid%para%group%mepos == 0)
587 : END IF
588 : END IF
589 : END IF
590 :
591 46 : CALL timeset("skala_gpw_grad_fetch", phase_handle)
592 46 : IF (features%uses_atom_chunks) THEN
593 : CALL fetch_and_gather_atom_chunk_grads(features, rho_r(1)%pw_grid%para%group, &
594 38 : density_grad, grad_grad, kin_grad)
595 : ELSE
596 8 : CALL fetch_local_feature_grads(features, density_grad, grad_grad, kin_grad)
597 : END IF
598 46 : CALL timestop(phase_handle)
599 : END IF
600 48 : IF (needs_atom_force) THEN
601 : CALL add_explicit_coordinate_force(atom_force, features, atom_coord_grad_t, &
602 8 : rho_r(1)%pw_grid%para%group%mepos == 0)
603 8 : 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 8 : grid_weight_grad_t, atomic_grid_weight_grad_t)
606 : END IF
607 : have_atom_coord_grad = .TRUE.
608 : END IF
609 :
610 48 : CALL timeset("skala_gpw_vxc_unpack", phase_handle)
611 48 : IF (native_grid_diagnostics) THEN
612 240 : bo = rho_r(1)%pw_grid%bounds_local
613 24 : density_contraction = 0.0_dp
614 24 : tau_contraction = 0.0_dp
615 24 : ipt = 0
616 1004 : DO k = bo(1, 3), bo(2, 3)
617 41848 : DO j = bo(1, 2), bo(2, 2)
618 905054 : DO i = bo(1, 1), bo(2, 1)
619 863230 : ipt = ipt + 1
620 904074 : IF (nspins == 1) THEN
621 : density_contraction = density_contraction + rho_r(1)%array(i, j, k)* &
622 589855 : 0.5_dp*(density_grad(ipt, 1) + density_grad(ipt, 2))
623 : tau_contraction = tau_contraction + tau(1)%array(i, j, k)* &
624 589855 : 0.5_dp*(kin_grad(ipt, 1) + kin_grad(ipt, 2))
625 : ELSE
626 820125 : DO ispin = 1, nspins
627 : density_contraction = density_contraction + &
628 546750 : rho_r(ispin)%array(i, j, k)*density_grad(ipt, ispin)
629 : tau_contraction = tau_contraction + &
630 820125 : 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 24 : CALL rho_r(1)%pw_grid%para%group%sum(density_contraction)
637 24 : CALL rho_r(1)%pw_grid%para%group%sum(tau_contraction)
638 24 : IF (rho_r(1)%pw_grid%para%group%mepos == 0) THEN
639 12 : iw = cp_logger_get_default_io_unit()
640 12 : IF (iw > 0) THEN
641 : WRITE (iw, "(T2,A,1X,ES20.10)") &
642 12 : "SKALA_GPW| XC density-gradient contraction", density_contraction
643 : WRITE (iw, "(T2,A,1X,ES20.10)") &
644 12 : "SKALA_GPW| XC kinetic-gradient contraction", tau_contraction
645 : END IF
646 : END IF
647 : END IF
648 48 : IF (compute_virial) THEN
649 6 : IF (native_grid_diagnostics) virial_before = virial_xc
650 6 : CALL build_virial_from_feature_grads(virial_xc, rho_set, rho_r, grad_grad)
651 6 : IF (native_grid_diagnostics) THEN
652 : CALL print_virial_delta("feature-gradient", virial_xc - virial_before, &
653 0 : rho_r(1)%pw_grid%para%group%mepos == 0)
654 0 : virial_before = virial_xc
655 : END IF
656 6 : IF (.NOT. have_atom_coord_grad) THEN
657 0 : CALL torch_tensor_grad(features%coarse_0_atomic_coords_t, atom_coord_grad_t)
658 0 : 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 6 : native_grid_diagnostics)
664 6 : IF (native_grid_diagnostics) THEN
665 : CALL print_virial_delta("static-coordinates", virial_xc - virial_before, &
666 0 : rho_r(1)%pw_grid%para%group%mepos == 0)
667 0 : virial_before = virial_xc
668 : END IF
669 6 : 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 6 : grid_weight_grad_t, atomic_grid_weight_grad_t)
672 6 : IF (native_grid_diagnostics) THEN
673 : CALL print_virial_delta("smooth-partition", virial_xc - virial_before, &
674 0 : rho_r(1)%pw_grid%para%group%mepos == 0)
675 : virial_before = virial_xc
676 : END IF
677 : END IF
678 : END IF
679 48 : 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 48 : xc_deriv_method_id)
683 48 : CALL timestop(phase_handle)
684 :
685 48 : CALL timeset("skala_gpw_grad_release", phase_handle)
686 48 : DEALLOCATE (density_grad, grad_grad, kin_grad)
687 48 : IF (have_atom_coord_grad) CALL torch_tensor_release(atom_coord_grad_t)
688 48 : CALL timestop(phase_handle)
689 : END IF
690 :
691 48 : CALL timeset("skala_gpw_cleanup", phase_handle)
692 48 : IF (.NOT. use_atom_subchunks .AND. has_atom_chunk_work) CALL torch_tensor_release(exc_tensor)
693 48 : CALL skala_gpw_feature_release(features)
694 48 : CALL xc_rho_set_release(rho_set, pw_pool=pw_pool)
695 48 : CALL torch_use_cuda(.TRUE.)
696 48 : CALL omp_set_num_threads(omp_max_threads_restore)
697 48 : CALL timestop(phase_handle)
698 :
699 960 : 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 0 : 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 0 : LOGICAL, ALLOCATABLE, DIMENSION(:) :: included
733 : REAL(KIND=dp) :: exc, local_derivative
734 0 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: partition_weights
735 0 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: atom_coords_pbc
736 0 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: dweights_datom, dweights_dstrain
737 : REAL(KIND=dp), DIMENSION(3) :: grid_point
738 0 : REAL(KIND=dp), DIMENSION(:), POINTER :: grid_weight_grad
739 : TYPE(section_vals_type), POINTER :: gauxc_section
740 0 : 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 0 : CPASSERT(ASSOCIATED(rho_r))
746 0 : CPASSERT(ASSOCIATED(rho_g))
747 0 : CPASSERT(ASSOCIATED(tau))
748 0 : CALL pw_zero(weight_deriv_r)
749 0 : omp_max_threads_restore = omp_get_max_threads()
750 :
751 0 : nspins = SIZE(rho_r)
752 0 : lsd = (nspins /= 1)
753 0 : CALL get_skala_model_path(xc_section, model_path)
754 0 : gauxc_section => get_gauxc_section(xc_section)
755 0 : 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 0 : i_val=native_grid_cuda_device)
758 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_PARTITION", &
759 0 : i_val=native_grid_atom_partition)
760 0 : SELECT CASE (native_grid_atom_partition)
761 : CASE (1)
762 0 : native_grid_atom_partition = skala_gpw_atom_partition_hard
763 : CASE (2)
764 0 : native_grid_atom_partition = skala_gpw_atom_partition_smooth
765 : CASE DEFAULT
766 : CALL cp_abort(__LOCATION__, &
767 0 : "Unknown GAUXC%NATIVE_GRID_ATOM_PARTITION value.")
768 : END SELECT
769 0 : IF (native_grid_atom_partition == skala_gpw_atom_partition_hard) THEN
770 0 : native_grid_atom_partition = skala_gpw_atom_partition_smooth
771 : END IF
772 :
773 0 : CALL torch_use_cuda(native_grid_use_cuda)
774 : selected_cuda_device = configure_native_grid_cuda( &
775 0 : native_grid_use_cuda, native_grid_cuda_device, rho_r(1)%pw_grid%para%group)
776 0 : CALL ensure_model_loaded(model_path, selected_cuda_device)
777 :
778 0 : IF (lsd) THEN
779 0 : needs%rho_spin = .TRUE.
780 0 : needs%drho_spin = .TRUE.
781 0 : needs%tau_spin = .TRUE.
782 : ELSE
783 0 : needs%rho = .TRUE.
784 0 : needs%drho = .TRUE.
785 0 : needs%tau = .TRUE.
786 : END IF
787 :
788 0 : CALL section_vals_val_get(xc_section, "XC_GRID%XC_DERIV", i_val=xc_deriv_method_id)
789 0 : 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 0 : tau_cutoff=section_get_rval(xc_section, "tau_cutoff"))
796 : CALL xc_rho_set_update(rho_set, rho_r, rho_g, tau, needs, &
797 0 : 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 0 : atom_partition=native_grid_atom_partition)
805 : CALL skala_torch_model_get_exc(cached_model, features%inputs, features%grid_weights_t, &
806 0 : exc_tensor, exc)
807 0 : CALL torch_tensor_backward_scalar(exc_tensor)
808 0 : NULLIFY (grid_weight_grad)
809 0 : CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
810 0 : CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
811 :
812 0 : natom = SIZE(particle_set)
813 0 : IF (native_grid_atom_partition == skala_gpw_atom_partition_smooth) THEN
814 : ALLOCATE (atom_coords_pbc(3, natom), included(natom), partition_weights(natom), &
815 0 : dweights_datom(3, natom, natom), dweights_dstrain(3, 3, natom))
816 0 : DO iatom = 1, natom
817 0 : atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.TRUE.)
818 : END DO
819 : END IF
820 :
821 0 : bo = rho_r(1)%pw_grid%bounds_local
822 0 : local_row = 0
823 0 : DO k = bo(1, 3), bo(2, 3)
824 0 : DO j = bo(1, 2), bo(2, 2)
825 0 : DO i = bo(1, 1), bo(2, 1)
826 0 : local_row = local_row + 1
827 0 : feature_begin = features%local_feature_offsets(local_row)
828 0 : feature_end = features%local_feature_offsets(local_row + 1) - 1
829 0 : local_derivative = 0.0_dp
830 0 : IF (native_grid_atom_partition == skala_gpw_atom_partition_hard) THEN
831 0 : CPASSERT(feature_end == feature_begin)
832 0 : row = features%local_feature_rows(feature_begin)
833 0 : local_derivative = grid_weight_grad(row)
834 : ELSE
835 0 : 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 0 : dweights_datom, dweights_dstrain)
839 0 : CPASSERT(feature_end - feature_begin + 1 == COUNT(included))
840 0 : feature_pos = feature_begin
841 0 : DO iatom = 1, natom
842 0 : IF (.NOT. included(iatom)) CYCLE
843 0 : row = features%local_feature_rows(feature_pos)
844 : local_derivative = local_derivative + &
845 0 : partition_weights(iatom)*grid_weight_grad(row)
846 0 : feature_pos = feature_pos + 1
847 : END DO
848 0 : CPASSERT(feature_pos == feature_end + 1)
849 : END IF
850 0 : weight_deriv_r%array(i, j, k) = local_derivative
851 : END DO
852 : END DO
853 : END DO
854 0 : CPASSERT(local_row == features%nflat_local)
855 :
856 0 : IF (ALLOCATED(atom_coords_pbc)) THEN
857 0 : DEALLOCATE (atom_coords_pbc, dweights_datom, dweights_dstrain, included, &
858 0 : partition_weights)
859 : END IF
860 0 : CALL torch_tensor_release(grid_weight_grad_t)
861 0 : CALL torch_tensor_release(exc_tensor)
862 0 : CALL skala_gpw_feature_release(features)
863 0 : CALL xc_rho_set_release(rho_set, pw_pool=pw_pool)
864 0 : CALL torch_use_cuda(.TRUE.)
865 0 : CALL omp_set_num_threads(omp_max_threads_restore)
866 :
867 0 : 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 4 : SUBROUTINE skala_gapw_atom_vxc_of_r(xc_section, grid_atom, group, atom_coord, &
892 4 : 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 4 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: atomic_grid_sizes
916 4 : 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 4 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: atomic_grid_weights, grid_weights
920 4 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: coarse_0_atomic_coords, density, &
921 4 : grid_coords, kin
922 4 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: grad
923 4 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: atom_coord_grad, density_grad, &
924 4 : grid_coord_grad, kin_grad
925 4 : 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 0 : CPASSERT(ASSOCIATED(xc_section))
939 4 : CPASSERT(ASSOCIATED(grid_atom))
940 4 : CPASSERT(ASSOCIATED(rho))
941 4 : CPASSERT(ASSOCIATED(drho))
942 4 : CPASSERT(ASSOCIATED(tau))
943 4 : omp_max_threads_restore = omp_get_max_threads()
944 :
945 4 : my_energy_only = .FALSE.
946 4 : IF (PRESENT(energy_only)) my_energy_only = energy_only
947 4 : need_coord_grad = PRESENT(atom_force) .OR. PRESENT(atom_virial)
948 4 : exc = 0.0_dp
949 4 : IF (PRESENT(atom_force)) atom_force = 0.0_dp
950 4 : IF (PRESENT(atom_virial)) atom_virial = 0.0_dp
951 4 : IF (.NOT. my_energy_only) THEN
952 10208 : vxc = 0.0_dp
953 40208 : vxg = 0.0_dp
954 10208 : vtau = 0.0_dp
955 : END IF
956 :
957 4 : CALL get_skala_model_path(xc_section, model_path)
958 4 : gauxc_section => get_gauxc_section(xc_section)
959 4 : CPASSERT(ASSOCIATED(gauxc_section))
960 4 : 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 4 : i_val=native_grid_cuda_device)
963 4 : CALL torch_use_cuda(native_grid_use_cuda)
964 : selected_cuda_device = configure_native_grid_cuda( &
965 4 : native_grid_use_cuda, native_grid_cuda_device, group)
966 4 : CALL ensure_model_loaded(model_path, selected_cuda_device)
967 :
968 4 : 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 52 : coarse_0_atomic_coords(3, 1), atomic_grid_size_bound_shape(0, nflat))
973 4 : density = 0.0_dp
974 4 : grad = 0.0_dp
975 4 : kin = 0.0_dp
976 4 : grid_coords = 0.0_dp
977 4 : grid_weights = 0.0_dp
978 4 : atomic_grid_weights = 0.0_dp
979 4 : atomic_grid_sizes(1) = INT(nflat, KIND=int_8)
980 : atomic_grid_size_bound_shape = 0_int_8
981 16 : coarse_0_atomic_coords(:, 1) = atom_coord
982 :
983 : row = 0
984 204 : DO ir = 1, nr
985 10204 : DO ia = 1, na
986 10000 : row = row + 1
987 : grid_coords(1, row) = atom_coord(1) + grid_atom%rad(ir)* &
988 10000 : grid_atom%sin_pol(ia)*grid_atom%cos_azi(ia)
989 : grid_coords(2, row) = atom_coord(2) + grid_atom%rad(ir)* &
990 10000 : grid_atom%sin_pol(ia)*grid_atom%sin_azi(ia)
991 10000 : grid_coords(3, row) = atom_coord(3) + grid_atom%rad(ir)*grid_atom%cos_pol(ia)
992 10000 : grid_weights(row) = weights(ia, ir)
993 10000 : atomic_grid_weights(row) = weights(ia, ir)
994 10200 : IF (nspins == 1) THEN
995 30000 : density(row, :) = 0.5_dp*rho(ia, ir, 1)
996 40000 : DO idir = 1, 3
997 100000 : grad(row, idir, :) = 0.5_dp*drho(idir, ia, ir, 1)
998 : END DO
999 30000 : kin(row, :) = 0.5_dp*tau(ia, ir, 1)
1000 : ELSE
1001 0 : density(row, :) = rho(ia, ir, 1:2)
1002 0 : DO idir = 1, 3
1003 0 : grad(row, idir, :) = drho(idir, ia, ir, 1:2)
1004 : END DO
1005 0 : kin(row, :) = tau(ia, ir, 1:2)
1006 : END IF
1007 : END DO
1008 : END DO
1009 :
1010 4 : CALL torch_tensor_from_array(grid_coords_t, grid_coords)
1011 4 : CALL torch_tensor_to_device_leaf(grid_coords_t, need_coord_grad)
1012 4 : CALL torch_tensor_from_array(grid_weights_t, grid_weights)
1013 4 : CALL torch_tensor_to_device_leaf(grid_weights_t, .FALSE.)
1014 4 : CALL torch_tensor_from_array(atomic_grid_weights_t, atomic_grid_weights)
1015 4 : CALL torch_tensor_to_device_leaf(atomic_grid_weights_t, .FALSE.)
1016 4 : CALL torch_tensor_from_array(atomic_grid_sizes_t, atomic_grid_sizes)
1017 4 : CALL torch_tensor_to_device_leaf(atomic_grid_sizes_t, .FALSE.)
1018 : CALL torch_tensor_from_array(atomic_grid_size_bound_shape_t, &
1019 4 : atomic_grid_size_bound_shape)
1020 4 : CALL torch_tensor_to_device_leaf(atomic_grid_size_bound_shape_t, .FALSE.)
1021 4 : CALL torch_tensor_from_array(coarse_0_atomic_coords_t, coarse_0_atomic_coords)
1022 4 : CALL torch_tensor_to_device_leaf(coarse_0_atomic_coords_t, need_coord_grad)
1023 4 : CALL torch_tensor_from_array(density_t, density)
1024 4 : CALL torch_tensor_to_device_leaf(density_t,.NOT. my_energy_only)
1025 4 : CALL torch_tensor_from_array(grad_t, grad)
1026 4 : CALL torch_tensor_to_device_leaf(grad_t,.NOT. my_energy_only)
1027 4 : CALL torch_tensor_from_array(kin_t, kin)
1028 4 : CALL torch_tensor_to_device_leaf(kin_t,.NOT. my_energy_only)
1029 :
1030 4 : CALL torch_dict_create(inputs)
1031 4 : CALL torch_dict_insert(inputs, "grid_coords", grid_coords_t)
1032 4 : CALL torch_dict_insert(inputs, "grid_weights", grid_weights_t)
1033 4 : CALL torch_dict_insert(inputs, "atomic_grid_weights", atomic_grid_weights_t)
1034 4 : CALL torch_dict_insert(inputs, "atomic_grid_sizes", atomic_grid_sizes_t)
1035 : CALL torch_dict_insert(inputs, "atomic_grid_size_bound_shape", &
1036 4 : atomic_grid_size_bound_shape_t)
1037 4 : CALL torch_dict_insert(inputs, "density", density_t)
1038 4 : CALL torch_dict_insert(inputs, "grad", grad_t)
1039 4 : CALL torch_dict_insert(inputs, "kin", kin_t)
1040 4 : CALL torch_dict_insert(inputs, "coarse_0_atomic_coords", coarse_0_atomic_coords_t)
1041 :
1042 4 : CALL skala_torch_model_get_exc(cached_model, inputs, grid_weights_t, exc_tensor, exc)
1043 :
1044 4 : IF (.NOT. my_energy_only) THEN
1045 4 : NULLIFY (atom_coord_grad, density_grad, grad_grad, grid_coord_grad, kin_grad)
1046 4 : CALL torch_tensor_backward_scalar(exc_tensor)
1047 4 : IF (need_coord_grad) THEN
1048 4 : CALL torch_tensor_grad(grid_coords_t, grid_coord_grad_t)
1049 4 : CALL torch_tensor_grad(coarse_0_atomic_coords_t, atom_coord_grad_t)
1050 4 : CALL torch_tensor_data_ptr(grid_coord_grad_t, grid_coord_grad)
1051 4 : CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
1052 4 : IF (PRESENT(atom_force)) THEN
1053 16 : atom_force(:) = atom_coord_grad(:, 1)
1054 10004 : DO row = 1, nflat
1055 40004 : atom_force(:) = atom_force(:) + grid_coord_grad(:, row)
1056 : END DO
1057 : END IF
1058 4 : IF (PRESENT(atom_virial)) THEN
1059 10004 : DO row = 1, nflat
1060 40004 : DO idir = 1, 3
1061 130000 : DO jdir = 1, 3
1062 90000 : tmp = grid_coord_grad(idir, row)*coarse_0_atomic_coords(jdir, 1)
1063 120000 : atom_virial(idir, jdir) = atom_virial(idir, jdir) + tmp
1064 : END DO
1065 : END DO
1066 : END DO
1067 16 : DO idir = 1, 3
1068 52 : DO jdir = 1, 3
1069 36 : tmp = atom_coord_grad(idir, 1)*coarse_0_atomic_coords(jdir, 1)
1070 48 : 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 4 : grad_grad_t, kin_grad_t)
1077 4 : CALL torch_tensor_data_ptr(density_grad_t, density_grad)
1078 4 : CALL torch_tensor_data_ptr(grad_grad_t, grad_grad)
1079 4 : CALL torch_tensor_data_ptr(kin_grad_t, kin_grad)
1080 :
1081 4 : row = 0
1082 204 : DO ir = 1, nr
1083 10204 : DO ia = 1, na
1084 10000 : row = row + 1
1085 10200 : IF (lsd) THEN
1086 0 : vxc(ia, ir, 1:2) = density_grad(row, 1:2)
1087 0 : DO idir = 1, 3
1088 0 : vxg(idir, ia, ir, 1:2) = grad_grad(row, idir, 1:2)
1089 : END DO
1090 0 : vtau(ia, ir, 1:2) = kin_grad(row, 1:2)
1091 : ELSE
1092 10000 : vxc(ia, ir, 1) = 0.5_dp*(density_grad(row, 1) + density_grad(row, 2))
1093 40000 : DO idir = 1, 3
1094 : vxg(idir, ia, ir, 1) = &
1095 40000 : 0.5_dp*(grad_grad(row, idir, 1) + grad_grad(row, idir, 2))
1096 : END DO
1097 10000 : 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 4 : CALL torch_tensor_release(density_grad_t)
1103 4 : CALL torch_tensor_release(grad_grad_t)
1104 4 : CALL torch_tensor_release(kin_grad_t)
1105 4 : IF (need_coord_grad) THEN
1106 4 : CALL torch_tensor_release(grid_coord_grad_t)
1107 4 : CALL torch_tensor_release(atom_coord_grad_t)
1108 : END IF
1109 : END IF
1110 :
1111 4 : CALL torch_tensor_release(exc_tensor)
1112 4 : CALL torch_tensor_release(density_t)
1113 4 : CALL torch_tensor_release(grad_t)
1114 4 : CALL torch_tensor_release(kin_t)
1115 4 : CALL torch_tensor_release(grid_coords_t)
1116 4 : CALL torch_tensor_release(grid_weights_t)
1117 4 : CALL torch_tensor_release(atomic_grid_weights_t)
1118 4 : CALL torch_tensor_release(atomic_grid_sizes_t)
1119 4 : CALL torch_tensor_release(atomic_grid_size_bound_shape_t)
1120 4 : CALL torch_tensor_release(coarse_0_atomic_coords_t)
1121 4 : CALL torch_dict_release(inputs)
1122 0 : DEALLOCATE (atomic_grid_size_bound_shape, atomic_grid_sizes, atomic_grid_weights, &
1123 4 : coarse_0_atomic_coords, density, grad, grid_coords, grid_weights, kin)
1124 4 : CALL torch_use_cuda(.TRUE.)
1125 4 : CALL omp_set_num_threads(omp_max_threads_restore)
1126 :
1127 12 : 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 264 : 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 264 : INTEGER, ALLOCATABLE, DIMENSION(:) :: chunk_atom_begin, chunk_atom_count, &
1178 264 : chunk_row_begin, chunk_row_counts
1179 264 : 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 264 : REAL(KIND=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
1184 264 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: atom_coord_grad, density_grad, &
1185 264 : grid_coord_grad, kin_grad
1186 264 : 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 0 : CPASSERT(ASSOCIATED(xc_section))
1210 264 : CPASSERT(SIZE(density, 1) == SIZE(grid_weights))
1211 264 : CPASSERT(SIZE(density, 1) == SIZE(atomic_grid_weights))
1212 264 : CPASSERT(SIZE(density, 1) == SIZE(grid_coords, 2))
1213 264 : CPASSERT(SIZE(density, 1) == SIZE(grad, 1))
1214 264 : CPASSERT(SIZE(density, 1) == SIZE(kin, 1))
1215 264 : CPASSERT(SIZE(density, 2) == 2)
1216 264 : CPASSERT(SIZE(grad, 2) == 3)
1217 264 : CPASSERT(SIZE(grad, 3) == 2)
1218 264 : CPASSERT(SIZE(kin, 2) == 2)
1219 264 : CPASSERT(SIZE(atomic_grid_sizes) == SIZE(atomic_coords, 2))
1220 529 : CPASSERT(SUM(atomic_grid_sizes) == INT(SIZE(density, 1), KIND=int_8))
1221 264 : need_derivatives = PRESENT(density_grad_out)
1222 264 : CPASSERT(PRESENT(grad_grad_out) .EQV. need_derivatives)
1223 264 : CPASSERT(PRESENT(kin_grad_out) .EQV. need_derivatives)
1224 264 : need_coordinate_derivatives = PRESENT(grid_coord_grad_out)
1225 264 : CPASSERT(PRESENT(grid_weight_grad_out) .EQV. need_coordinate_derivatives)
1226 264 : CPASSERT(PRESENT(atomic_grid_weight_grad_out) .EQV. need_coordinate_derivatives)
1227 264 : CPASSERT(PRESENT(atom_coord_grad_out) .EQV. need_coordinate_derivatives)
1228 264 : CPASSERT((.NOT. need_coordinate_derivatives) .OR. need_derivatives)
1229 264 : omp_max_threads_restore = omp_get_max_threads()
1230 :
1231 264 : local_nrow = SIZE(density, 1)
1232 264 : local_natom = SIZE(atomic_grid_sizes)
1233 264 : active_rank = local_natom > 0
1234 264 : CPASSERT(active_rank .EQV. (local_nrow > 0))
1235 :
1236 264 : CALL get_skala_model_path(xc_section, model_path)
1237 264 : gauxc_section => get_gauxc_section(xc_section)
1238 264 : CPASSERT(ASSOCIATED(gauxc_section))
1239 264 : 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 264 : i_val=native_grid_cuda_device)
1242 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNKS", &
1243 264 : l_val=native_grid_atom_chunks)
1244 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNK_MAX_ROWS", &
1245 264 : i_val=chunk_max_rows)
1246 264 : IF (chunk_max_rows < -1) THEN
1247 : CALL cp_abort(__LOCATION__, &
1248 0 : "GAUXC%NATIVE_GRID_ATOM_CHUNK_MAX_ROWS must be -1, zero, or positive.")
1249 : END IF
1250 264 : IF (.NOT. native_grid_atom_chunks) chunk_max_rows = 0
1251 264 : IF (chunk_max_rows == -1) THEN
1252 254 : chunk_max_rows = auto_atom_composite_chunk_max_rows(atomic_grid_sizes, group)
1253 : END IF
1254 264 : CALL torch_use_cuda(native_grid_use_cuda)
1255 : selected_cuda_device = configure_native_grid_cuda( &
1256 264 : native_grid_use_cuda, native_grid_cuda_device, group)
1257 264 : exc = 0.0_dp
1258 264 : IF (active_rank) THEN
1259 242 : CALL timeset("skala_atom_tensor_setup", phase_handle)
1260 242 : 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 242 : chunk_row_begin, chunk_row_counts)
1264 242 : nchunks = SIZE(chunk_atom_begin)
1265 242 : CPASSERT(nchunks > 0)
1266 :
1267 242 : IF (native_grid_use_cuda) THEN
1268 0 : grid_coords_t => cached_atom_cuda_tensors%grid_coords_t
1269 0 : grid_weights_t => cached_atom_cuda_tensors%grid_weights_t
1270 0 : atomic_grid_weights_t => cached_atom_cuda_tensors%atomic_grid_weights_t
1271 0 : atomic_coords_t => cached_atom_cuda_tensors%atomic_coords_t
1272 0 : density_t => cached_atom_cuda_tensors%density_t
1273 0 : grad_t => cached_atom_cuda_tensors%grad_t
1274 0 : kin_t => cached_atom_cuda_tensors%kin_t
1275 : CALL torch_tensor_reset_from_array( &
1276 0 : grid_coords_t, grid_coords, need_coordinate_derivatives)
1277 : CALL torch_tensor_reset_from_array( &
1278 0 : grid_weights_t, grid_weights, need_coordinate_derivatives)
1279 : CALL torch_tensor_reset_from_array( &
1280 0 : atomic_grid_weights_t, atomic_grid_weights, need_coordinate_derivatives)
1281 : CALL torch_tensor_reset_from_array( &
1282 0 : atomic_coords_t, atomic_coords, need_coordinate_derivatives)
1283 0 : CALL torch_tensor_reset_from_array(density_t, density, need_derivatives)
1284 0 : CALL torch_tensor_reset_from_array(grad_t, grad, need_derivatives)
1285 0 : CALL torch_tensor_reset_from_array(kin_t, kin, need_derivatives)
1286 : ELSE
1287 242 : grid_coords_t => grid_coords_local_t
1288 242 : grid_weights_t => grid_weights_local_t
1289 242 : atomic_grid_weights_t => atomic_grid_weights_local_t
1290 242 : atomic_coords_t => atomic_coords_local_t
1291 242 : density_t => density_local_t
1292 242 : grad_t => grad_local_t
1293 242 : kin_t => kin_local_t
1294 242 : CALL torch_tensor_from_array(grid_coords_t, grid_coords)
1295 242 : CALL torch_tensor_to_device_leaf(grid_coords_t, need_coordinate_derivatives)
1296 242 : CALL torch_tensor_from_array(grid_weights_t, grid_weights)
1297 242 : CALL torch_tensor_to_device_leaf(grid_weights_t, need_coordinate_derivatives)
1298 242 : CALL torch_tensor_from_array(atomic_grid_weights_t, atomic_grid_weights)
1299 242 : CALL torch_tensor_to_device_leaf(atomic_grid_weights_t, need_coordinate_derivatives)
1300 242 : CALL torch_tensor_from_array(atomic_coords_t, atomic_coords)
1301 242 : CALL torch_tensor_to_device_leaf(atomic_coords_t, need_coordinate_derivatives)
1302 242 : CALL torch_tensor_from_array(density_t, density)
1303 242 : CALL torch_tensor_to_device_leaf(density_t, need_derivatives)
1304 242 : CALL torch_tensor_from_array(grad_t, grad)
1305 242 : CALL torch_tensor_to_device_leaf(grad_t, need_derivatives)
1306 242 : CALL torch_tensor_from_array(kin_t, kin)
1307 242 : CALL torch_tensor_to_device_leaf(kin_t, need_derivatives)
1308 : END IF
1309 242 : CALL torch_tensor_from_array(atomic_grid_sizes_t, atomic_grid_sizes)
1310 242 : CALL torch_tensor_to_device_leaf(atomic_grid_sizes_t, .FALSE.)
1311 242 : CALL timestop(phase_handle)
1312 :
1313 484 : DO ichunk = 1, nchunks
1314 242 : atom_begin = chunk_atom_begin(ichunk)
1315 242 : atom_count = chunk_atom_count(ichunk)
1316 242 : chunk_row_start = chunk_row_begin(ichunk)
1317 242 : chunk_row_count = chunk_row_counts(ichunk)
1318 : max_grid_size = INT(MAXVAL(atomic_grid_sizes( &
1319 507 : atom_begin:atom_begin + atom_count - 1)))
1320 484 : 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 242 : chunk_row_count, grid_coords_chunk_t)
1325 : CALL torch_tensor_narrow(grid_weights_t, 0, chunk_row_start - 1, &
1326 242 : chunk_row_count, grid_weights_chunk_t)
1327 : CALL torch_tensor_narrow(atomic_grid_weights_t, 0, chunk_row_start - 1, &
1328 242 : chunk_row_count, atomic_grid_weights_chunk_t)
1329 : CALL torch_tensor_narrow(atomic_grid_sizes_t, 0, atom_begin - 1, &
1330 242 : atom_count, atomic_grid_sizes_chunk_t)
1331 : CALL torch_tensor_narrow(atomic_coords_t, 0, atom_begin - 1, &
1332 242 : atom_count, atomic_coords_chunk_t)
1333 : CALL torch_tensor_narrow(density_t, 1, chunk_row_start - 1, &
1334 242 : chunk_row_count, density_chunk_t)
1335 : CALL torch_tensor_narrow(grad_t, 2, chunk_row_start - 1, &
1336 242 : chunk_row_count, grad_chunk_t)
1337 : CALL torch_tensor_narrow(kin_t, 1, chunk_row_start - 1, &
1338 242 : chunk_row_count, kin_chunk_t)
1339 : CALL torch_tensor_from_array(atomic_grid_size_bound_shape_t, &
1340 242 : atomic_grid_size_bound_shape)
1341 242 : CALL torch_tensor_to_device_leaf(atomic_grid_size_bound_shape_t, .FALSE.)
1342 :
1343 242 : CALL torch_dict_create(inputs)
1344 242 : CALL torch_dict_insert(inputs, "grid_coords", grid_coords_chunk_t)
1345 242 : CALL torch_dict_insert(inputs, "grid_weights", grid_weights_chunk_t)
1346 242 : CALL torch_dict_insert(inputs, "atomic_grid_weights", atomic_grid_weights_chunk_t)
1347 242 : 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 242 : atomic_grid_size_bound_shape_t)
1350 242 : CALL torch_dict_insert(inputs, "density", density_chunk_t)
1351 242 : CALL torch_dict_insert(inputs, "grad", grad_chunk_t)
1352 242 : CALL torch_dict_insert(inputs, "kin", kin_chunk_t)
1353 242 : CALL torch_dict_insert(inputs, "coarse_0_atomic_coords", atomic_coords_chunk_t)
1354 :
1355 242 : CALL timeset("skala_atom_torch_forward", phase_handle)
1356 : CALL skala_torch_model_get_exc(cached_model, inputs, grid_weights_chunk_t, &
1357 242 : exc_tensor, chunk_exc)
1358 242 : CALL timestop(phase_handle)
1359 242 : exc = exc + chunk_exc
1360 242 : IF (need_derivatives) THEN
1361 242 : CALL timeset("skala_atom_torch_backward", phase_handle)
1362 242 : CALL torch_tensor_backward_scalar(exc_tensor)
1363 242 : CALL timestop(phase_handle)
1364 : END IF
1365 :
1366 242 : CALL torch_tensor_release(exc_tensor)
1367 242 : CALL torch_dict_release(inputs)
1368 242 : CALL torch_tensor_release(grid_coords_chunk_t)
1369 242 : CALL torch_tensor_release(grid_weights_chunk_t)
1370 242 : CALL torch_tensor_release(atomic_grid_weights_chunk_t)
1371 242 : CALL torch_tensor_release(atomic_grid_sizes_chunk_t)
1372 242 : CALL torch_tensor_release(atomic_coords_chunk_t)
1373 242 : CALL torch_tensor_release(density_chunk_t)
1374 242 : CALL torch_tensor_release(grad_chunk_t)
1375 242 : CALL torch_tensor_release(kin_chunk_t)
1376 242 : CALL torch_tensor_release(atomic_grid_size_bound_shape_t)
1377 726 : DEALLOCATE (atomic_grid_size_bound_shape)
1378 : END DO
1379 : END IF
1380 264 : CALL group%sum(exc)
1381 :
1382 264 : IF (need_derivatives) THEN
1383 0 : ALLOCATE (density_grad_out(local_nrow, 2), grad_grad_out(local_nrow, 3, 2), &
1384 1782 : kin_grad_out(local_nrow, 2))
1385 264 : density_grad_out = 0.0_dp
1386 264 : grad_grad_out = 0.0_dp
1387 264 : kin_grad_out = 0.0_dp
1388 264 : IF (active_rank) THEN
1389 242 : NULLIFY (density_grad, grad_grad, kin_grad)
1390 242 : CALL timeset("skala_atom_grad_fetch", phase_handle)
1391 : CALL torch_tensor_grad_batch3(density_t, grad_t, kin_t, density_grad_t, &
1392 242 : grad_grad_t, kin_grad_t)
1393 242 : CALL torch_tensor_data_ptr(density_grad_t, density_grad)
1394 242 : CALL torch_tensor_data_ptr(grad_grad_t, grad_grad)
1395 242 : CALL torch_tensor_data_ptr(kin_grad_t, kin_grad)
1396 917246 : density_grad_out(:, :) = density_grad
1397 2751738 : grad_grad_out(:, :, :) = grad_grad
1398 917246 : kin_grad_out(:, :) = kin_grad
1399 242 : CALL timestop(phase_handle)
1400 : END IF
1401 :
1402 264 : IF (need_coordinate_derivatives) THEN
1403 64 : IF (active_rank) CALL timeset("skala_atom_coord_grad_fetch", phase_handle)
1404 0 : ALLOCATE (grid_coord_grad_out(3, local_nrow), grid_weight_grad_out(local_nrow), &
1405 468 : atomic_grid_weight_grad_out(local_nrow), atom_coord_grad_out(3, local_natom))
1406 64 : grid_coord_grad_out = 0.0_dp
1407 64 : grid_weight_grad_out = 0.0_dp
1408 64 : atomic_grid_weight_grad_out = 0.0_dp
1409 64 : atom_coord_grad_out = 0.0_dp
1410 64 : IF (active_rank) THEN
1411 53 : NULLIFY (atomic_grid_weight_grad, atom_coord_grad, grid_coord_grad, grid_weight_grad)
1412 53 : CALL torch_tensor_grad(grid_coords_t, grid_coord_grad_t)
1413 53 : CALL torch_tensor_grad(grid_weights_t, grid_weight_grad_t)
1414 53 : CALL torch_tensor_grad(atomic_grid_weights_t, atomic_grid_weight_grad_t)
1415 53 : CALL torch_tensor_grad(atomic_coords_t, atom_coord_grad_t)
1416 53 : CALL torch_tensor_data_ptr(grid_coord_grad_t, grid_coord_grad)
1417 53 : CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
1418 53 : CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
1419 53 : CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
1420 480213 : grid_coord_grad_out(:, :) = grid_coord_grad
1421 120093 : grid_weight_grad_out(:) = grid_weight_grad
1422 120093 : atomic_grid_weight_grad_out(:) = atomic_grid_weight_grad
1423 309 : atom_coord_grad_out(:, :) = atom_coord_grad
1424 53 : CALL torch_tensor_release(grid_coord_grad_t)
1425 53 : CALL torch_tensor_release(grid_weight_grad_t)
1426 53 : CALL torch_tensor_release(atomic_grid_weight_grad_t)
1427 53 : CALL torch_tensor_release(atom_coord_grad_t)
1428 53 : CALL timestop(phase_handle)
1429 : END IF
1430 : END IF
1431 253 : IF (active_rank) THEN
1432 242 : CALL torch_tensor_release(density_grad_t)
1433 242 : CALL torch_tensor_release(grad_grad_t)
1434 242 : CALL torch_tensor_release(kin_grad_t)
1435 : END IF
1436 : END IF
1437 :
1438 253 : IF (active_rank) THEN
1439 242 : CALL timeset("skala_atom_tensor_release", phase_handle)
1440 242 : IF (.NOT. native_grid_use_cuda) THEN
1441 242 : CALL torch_tensor_release(density_t)
1442 242 : CALL torch_tensor_release(grad_t)
1443 242 : CALL torch_tensor_release(kin_t)
1444 242 : CALL torch_tensor_release(grid_coords_t)
1445 242 : CALL torch_tensor_release(grid_weights_t)
1446 242 : CALL torch_tensor_release(atomic_grid_weights_t)
1447 242 : CALL torch_tensor_release(atomic_coords_t)
1448 : END IF
1449 242 : CALL torch_tensor_release(atomic_grid_sizes_t)
1450 242 : DEALLOCATE (chunk_atom_begin, chunk_atom_count, chunk_row_begin, chunk_row_counts)
1451 242 : CALL timestop(phase_handle)
1452 : END IF
1453 264 : CALL torch_use_cuda(.TRUE.)
1454 264 : CALL omp_set_num_threads(omp_max_threads_restore)
1455 :
1456 528 : 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 254 : 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 254 : local_padded_rows = 0
1473 254 : local_natom = SIZE(atomic_grid_sizes)
1474 254 : IF (local_natom > 0) THEN
1475 487 : local_max_atom_rows = INT(MAXVAL(atomic_grid_sizes))
1476 232 : IF (local_natom > atom_chunk_auto_max_rows .OR. &
1477 : local_max_atom_rows > atom_chunk_auto_max_rows/local_natom) THEN
1478 0 : local_padded_rows = atom_chunk_auto_max_rows + 1
1479 : ELSE
1480 232 : local_padded_rows = local_natom*local_max_atom_rows
1481 : END IF
1482 : END IF
1483 254 : CALL group%max(local_padded_rows)
1484 254 : IF (local_padded_rows <= atom_chunk_auto_max_rows) THEN
1485 254 : max_rows = 0
1486 254 : RETURN
1487 : END IF
1488 254 : 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 242 : 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 242 : INTEGER, ALLOCATABLE, DIMENSION(:) :: padded_rows
1513 :
1514 242 : natom = SIZE(atomic_grid_sizes)
1515 242 : CPASSERT(natom > 0)
1516 507 : CPASSERT(ALL(atomic_grid_sizes > 0_int_8))
1517 242 : IF (max_rows <= 0) THEN
1518 232 : ALLOCATE (atom_begin(1), atom_count(1), row_begin(1), row_count(1))
1519 232 : atom_begin(1) = 1
1520 232 : atom_count(1) = natom
1521 232 : row_begin(1) = 1
1522 487 : 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 20 : DO iatom = 1, natom
1530 10 : atom_rows = INT(atomic_grid_sizes(iatom))
1531 10 : IF (chunk_atoms > 0 .AND. &
1532 : MAX(chunk_max_atom_rows, atom_rows) > max_rows/(chunk_atoms + 1)) THEN
1533 0 : nchunks = nchunks + 1
1534 0 : chunk_atoms = 0
1535 0 : chunk_max_atom_rows = 0
1536 : END IF
1537 10 : chunk_atoms = chunk_atoms + 1
1538 20 : chunk_max_atom_rows = MAX(chunk_max_atom_rows, atom_rows)
1539 : END DO
1540 60 : ALLOCATE (atom_begin(nchunks), atom_count(nchunks), row_begin(nchunks), row_count(nchunks))
1541 :
1542 10 : subchunk = 1
1543 10 : atom_begin(subchunk) = 1
1544 10 : row_begin(subchunk) = 1
1545 10 : chunk_atoms = 0
1546 10 : chunk_max_atom_rows = 0
1547 10 : rows = 0
1548 20 : DO iatom = 1, natom
1549 10 : atom_rows = INT(atomic_grid_sizes(iatom))
1550 10 : IF (chunk_atoms > 0 .AND. &
1551 : MAX(chunk_max_atom_rows, atom_rows) > max_rows/(chunk_atoms + 1)) THEN
1552 0 : atom_count(subchunk) = iatom - atom_begin(subchunk)
1553 0 : row_count(subchunk) = rows
1554 0 : subchunk = subchunk + 1
1555 0 : atom_begin(subchunk) = iatom
1556 0 : row_begin(subchunk) = row_begin(subchunk - 1) + row_count(subchunk - 1)
1557 0 : chunk_atoms = 0
1558 0 : chunk_max_atom_rows = 0
1559 0 : rows = 0
1560 : END IF
1561 10 : chunk_atoms = chunk_atoms + 1
1562 10 : chunk_max_atom_rows = MAX(chunk_max_atom_rows, atom_rows)
1563 20 : rows = rows + atom_rows
1564 : END DO
1565 10 : atom_count(subchunk) = natom - atom_begin(subchunk) + 1
1566 10 : row_count(subchunk) = rows
1567 10 : CPASSERT(subchunk == nchunks)
1568 :
1569 20 : ALLOCATE (padded_rows(nchunks))
1570 20 : DO subchunk = 1, nchunks
1571 10 : iatom = atom_begin(subchunk) + atom_count(subchunk) - 1
1572 : padded_rows(subchunk) = atom_count(subchunk)* &
1573 30 : 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 10 : DO subchunk = 2, nchunks
1578 0 : padded_rows_tmp = padded_rows(subchunk)
1579 0 : atom_begin_tmp = atom_begin(subchunk)
1580 0 : atom_count_tmp = atom_count(subchunk)
1581 0 : row_begin_tmp = row_begin(subchunk)
1582 0 : row_count_tmp = row_count(subchunk)
1583 0 : insert_at = subchunk
1584 0 : DO WHILE (insert_at > 1 .AND. padded_rows(insert_at - 1) < padded_rows_tmp)
1585 0 : padded_rows(insert_at) = padded_rows(insert_at - 1)
1586 0 : atom_begin(insert_at) = atom_begin(insert_at - 1)
1587 0 : atom_count(insert_at) = atom_count(insert_at - 1)
1588 0 : row_begin(insert_at) = row_begin(insert_at - 1)
1589 0 : row_count(insert_at) = row_count(insert_at - 1)
1590 0 : insert_at = insert_at - 1
1591 : END DO
1592 0 : padded_rows(insert_at) = padded_rows_tmp
1593 0 : atom_begin(insert_at) = atom_begin_tmp
1594 0 : atom_count(insert_at) = atom_count_tmp
1595 0 : row_begin(insert_at) = row_begin_tmp
1596 10 : row_count(insert_at) = row_count_tmp
1597 : END DO
1598 10 : 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 8 : 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 8 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: atom_coord_grad
1616 :
1617 8 : NULLIFY (atom_coord_grad)
1618 8 : CALL torch_tensor_grad(features%coarse_0_atomic_coords_t, atom_coord_grad_t)
1619 8 : IF (root_rank) THEN
1620 4 : CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
1621 4 : CPASSERT(SIZE(atom_force, 1) == SIZE(atom_coord_grad, 1))
1622 4 : CPASSERT(SIZE(atom_force, 2) == SIZE(atom_coord_grad, 2))
1623 36 : atom_force(:, :) = atom_force(:, :) + atom_coord_grad(:, :)
1624 : END IF
1625 :
1626 8 : 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 8 : 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 8 : REAL(KIND=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
1659 :
1660 8 : NULLIFY (atomic_grid_weight_grad, grid_weight_grad)
1661 8 : CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
1662 8 : CALL torch_tensor_grad(features%atomic_grid_weights_t, atomic_grid_weight_grad_t)
1663 8 : CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
1664 8 : CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
1665 :
1666 8 : natom = SIZE(particle_set)
1667 8 : CPASSERT(SIZE(atom_force, 1) == 3)
1668 8 : CPASSERT(SIZE(atom_force, 2) == natom)
1669 : ALLOCATE (atom_coords_pbc(3, natom), included(natom), weights(natom), &
1670 96 : dweights_datom(3, natom, natom), dweights_dstrain(3, 3, natom))
1671 24 : DO iatom = 1, natom
1672 24 : atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.TRUE.)
1673 : END DO
1674 :
1675 80 : bo = rho_r(1)%pw_grid%bounds_local
1676 8 : local_row = 0
1677 182 : DO k = bo(1, 3), bo(2, 3)
1678 4088 : DO j = bo(1, 2), bo(2, 2)
1679 48927 : DO i = bo(1, 1), bo(2, 1)
1680 44847 : local_row = local_row + 1
1681 179388 : 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 44847 : dweights_dstrain)
1685 44847 : feature_begin = features%local_feature_offsets(local_row)
1686 44847 : feature_end = features%local_feature_offsets(local_row + 1) - 1
1687 134541 : CPASSERT(feature_end - feature_begin + 1 == COUNT(included))
1688 44847 : grid_base_weight = 0.0_dp
1689 134287 : DO feature_pos = feature_begin, feature_end
1690 89440 : row = features%local_feature_rows(feature_pos)
1691 134287 : grid_base_weight = grid_base_weight + features%grid_weights(row)
1692 : END DO
1693 : feature_pos = feature_begin
1694 134541 : DO iatom = 1, natom
1695 89694 : IF (.NOT. included(iatom)) CYCLE
1696 89440 : 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 89440 : smooth_partition_atomic_weight_scale_derivative(weights(iatom))
1700 268320 : DO jatom = 1, natom
1701 : atom_force(:, jatom) = atom_force(:, jatom) + &
1702 804960 : weight_grad*dweights_datom(:, jatom, iatom)
1703 : END DO
1704 134541 : feature_pos = feature_pos + 1
1705 : END DO
1706 48753 : CPASSERT(feature_pos == feature_end + 1)
1707 : END DO
1708 : END DO
1709 : END DO
1710 8 : CPASSERT(local_row == features%nflat_local)
1711 :
1712 8 : DEALLOCATE (atom_coords_pbc, dweights_datom, dweights_dstrain, included, weights)
1713 8 : CALL torch_tensor_release(grid_weight_grad_t)
1714 8 : CALL torch_tensor_release(atomic_grid_weight_grad_t)
1715 :
1716 8 : 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 6 : 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 6 : REAL(KIND=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
1749 :
1750 6 : NULLIFY (atomic_grid_weight_grad, grid_weight_grad)
1751 6 : CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
1752 6 : CALL torch_tensor_grad(features%atomic_grid_weights_t, atomic_grid_weight_grad_t)
1753 6 : CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
1754 6 : CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
1755 :
1756 6 : natom = SIZE(particle_set)
1757 : ALLOCATE (atom_coords_pbc(3, natom), included(natom), weights(natom), &
1758 72 : dweights_datom(3, natom, natom), dweights_dstrain(3, 3, natom))
1759 18 : DO iatom = 1, natom
1760 18 : atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.TRUE.)
1761 : END DO
1762 :
1763 60 : bo = rho_r(1)%pw_grid%bounds_local
1764 6 : local_row = 0
1765 150 : DO k = bo(1, 3), bo(2, 3)
1766 3606 : DO j = bo(1, 2), bo(2, 2)
1767 45072 : DO i = bo(1, 1), bo(2, 1)
1768 41472 : local_row = local_row + 1
1769 165888 : 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 41472 : dweights_dstrain)
1773 41472 : feature_begin = features%local_feature_offsets(local_row)
1774 41472 : feature_end = features%local_feature_offsets(local_row + 1) - 1
1775 124416 : CPASSERT(feature_end - feature_begin + 1 == COUNT(included))
1776 41472 : grid_base_weight = 0.0_dp
1777 124206 : DO feature_pos = feature_begin, feature_end
1778 82734 : row = features%local_feature_rows(feature_pos)
1779 124206 : grid_base_weight = grid_base_weight + features%grid_weights(row)
1780 : END DO
1781 : feature_pos = feature_begin
1782 124416 : DO iatom = 1, natom
1783 82944 : IF (.NOT. included(iatom)) CYCLE
1784 82734 : 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 82734 : smooth_partition_atomic_weight_scale_derivative(weights(iatom))
1788 330936 : DO idir = 1, 3
1789 827340 : DO jdir = 1, idir
1790 496404 : tmp = weight_grad*dweights_dstrain(idir, jdir, iatom)
1791 496404 : virial_xc(jdir, idir) = virial_xc(jdir, idir) + tmp
1792 744606 : IF (idir /= jdir) virial_xc(idir, jdir) = virial_xc(idir, jdir) + tmp
1793 : END DO
1794 : END DO
1795 124416 : feature_pos = feature_pos + 1
1796 : END DO
1797 44928 : CPASSERT(feature_pos == feature_end + 1)
1798 : END DO
1799 : END DO
1800 : END DO
1801 6 : CPASSERT(local_row == features%nflat_local)
1802 :
1803 6 : DEALLOCATE (atom_coords_pbc, dweights_datom, dweights_dstrain, included, weights)
1804 6 : CALL torch_tensor_release(grid_weight_grad_t)
1805 6 : CALL torch_tensor_release(atomic_grid_weight_grad_t)
1806 :
1807 6 : 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 86319 : 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 345276 : 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 345276 : REAL(relative_index(3), KIND=dp)*pw_grid%dh(:, 3)
1826 :
1827 86319 : 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 2 : 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 2 : INTEGER, ALLOCATABLE, DIMENSION(:) :: route_grad_return_recv_counts, &
1858 2 : route_grad_return_recv_displs, &
1859 2 : route_grad_return_send_counts, &
1860 2 : route_grad_return_send_displs, &
1861 2 : subchunk_atom_begin, &
1862 2 : subchunk_atom_count, &
1863 2 : subchunk_row_begin, &
1864 2 : subchunk_row_count
1865 : REAL(KIND=dp) :: subchunk_exc
1866 2 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: recv_grad_buffer, send_grad_buffer
1867 2 : TYPE(skala_gpw_feature_type) :: subchunk
1868 : TYPE(torch_tensor_type) :: subchunk_exc_tensor
1869 :
1870 0 : CPASSERT(features%uses_atom_chunks)
1871 2 : CPASSERT(max_rows > 0)
1872 : CALL skala_gpw_atom_subchunk_layout(max_rows, subchunk_atom_begin, subchunk_atom_count, &
1873 2 : subchunk_row_begin, subchunk_row_count)
1874 2 : nsubchunks = SIZE(subchunk_atom_begin)
1875 :
1876 2 : exc = 0.0_dp
1877 2 : IF (compute_grads) THEN
1878 2 : CPASSERT(features%uses_atom_chunk_routing)
1879 6 : nroute_recv_points = SUM(features%route_point_recv_counts)
1880 2 : nroute_send_points = SIZE(features%route_send_local_rows)
1881 6 : CPASSERT(SUM(features%route_point_send_counts) == nroute_send_points)
1882 2 : CPASSERT(SIZE(features%route_chunk_offsets) == nroute_recv_points + 1)
1883 2 : nroute_grad_per_point = ngrad_per_point
1884 2 : 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 26 : route_grad_return_recv_displs(SIZE(features%route_point_send_displs)))
1891 : route_grad_return_send_counts(:) = &
1892 6 : nroute_grad_per_point*features%route_point_recv_counts
1893 : route_grad_return_send_displs(:) = &
1894 6 : nroute_grad_per_point*features%route_point_recv_displs
1895 : route_grad_return_recv_counts(:) = &
1896 6 : nroute_grad_per_point*features%route_point_send_counts
1897 : route_grad_return_recv_displs(:) = &
1898 6 : nroute_grad_per_point*features%route_point_send_displs
1899 : END IF
1900 :
1901 2 : CALL timeset("skala_gpw_atom_subchunks", phase_handle)
1902 6 : DO isubchunk = 1, nsubchunks
1903 4 : 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 4 : compute_grads)
1910 4 : CALL timestop(subphase_handle)
1911 4 : 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 4 : subchunk_exc)
1915 4 : CALL timestop(subphase_handle)
1916 4 : exc = exc + subchunk_exc
1917 4 : IF (compute_grads) THEN
1918 4 : CALL timeset("skala_gpw_atom_subchunk_backward", subphase_handle)
1919 4 : CALL torch_tensor_backward_scalar(subchunk_exc_tensor)
1920 4 : CALL timestop(subphase_handle)
1921 : END IF
1922 4 : CALL timeset("skala_gpw_atom_subchunk_release", subphase_handle)
1923 4 : CALL torch_tensor_release(subchunk_exc_tensor)
1924 4 : CALL skala_gpw_feature_release(subchunk)
1925 18 : CALL timestop(subphase_handle)
1926 : END DO
1927 2 : IF (compute_grads .AND. features%chunk_feature_count > 0) THEN
1928 2 : CALL timeset("skala_gpw_atom_subchunk_grad_pack", subphase_handle)
1929 2 : CALL pack_atom_chunk_grads(features, send_grad_buffer, .TRUE., collapse_spin_grads)
1930 2 : CALL timestop(subphase_handle)
1931 : END IF
1932 2 : CALL timestop(phase_handle)
1933 :
1934 2 : IF (compute_grads) THEN
1935 2 : 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 2 : route_grad_return_recv_counts, route_grad_return_recv_displs)
1939 2 : CALL timestop(phase_handle)
1940 :
1941 2 : CALL timeset("skala_gpw_grad_route_scatter", phase_handle)
1942 : CALL scatter_routed_atom_chunk_grads(features, recv_grad_buffer, collapse_spin_grads, &
1943 2 : density_grad, grad_grad, kin_grad)
1944 2 : CALL timestop(phase_handle)
1945 :
1946 0 : DEALLOCATE (recv_grad_buffer, route_grad_return_recv_counts, &
1947 0 : route_grad_return_recv_displs, route_grad_return_send_counts, &
1948 2 : route_grad_return_send_displs, send_grad_buffer)
1949 : END IF
1950 2 : DEALLOCATE (subchunk_atom_begin, subchunk_atom_count, subchunk_row_begin, subchunk_row_count)
1951 :
1952 4 : 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 36 : 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 36 : local_rows_max = features%chunk_feature_count
1969 36 : CALL group%max(local_rows_max)
1970 36 : IF (local_rows_max <= 0) THEN
1971 36 : max_rows = 0
1972 : RETURN
1973 : END IF
1974 :
1975 36 : IF (group%num_pe > 1) THEN
1976 36 : target_rows = CEILING(REAL(local_rows_max, KIND=dp)/2.0_dp)
1977 : max_rows = atom_chunk_auto_row_quantum* &
1978 36 : ((target_rows + atom_chunk_auto_row_quantum - 1)/atom_chunk_auto_row_quantum)
1979 : ELSE
1980 0 : 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 0 : REAL(atom_chunk_auto_row_quantum, KIND=dp)))
1984 : END IF
1985 36 : max_rows = MAX(atom_chunk_auto_min_rows, MIN(atom_chunk_auto_max_rows, max_rows))
1986 :
1987 36 : 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 8 : 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 8 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: density_grad_all, kin_grad_all
2007 8 : REAL(KIND=dp), DIMENSION(:, :, :), POINTER :: grad_grad_all
2008 : TYPE(torch_tensor_type) :: density_grad_t, grad_grad_t, kin_grad_t
2009 :
2010 8 : 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 8 : density_grad_all, grad_grad_all, kin_grad_all)
2013 8 : CPASSERT(SIZE(density_grad_all, 1) == features%nflat)
2014 8 : CPASSERT(SIZE(density_grad_all, 2) == 2)
2015 8 : CPASSERT(SIZE(grad_grad_all, 1) == features%nflat)
2016 8 : CPASSERT(SIZE(grad_grad_all, 2) == 3)
2017 8 : CPASSERT(SIZE(grad_grad_all, 3) == 2)
2018 8 : CPASSERT(SIZE(kin_grad_all, 1) == features%nflat)
2019 8 : CPASSERT(SIZE(kin_grad_all, 2) == 2)
2020 :
2021 0 : ALLOCATE (density_grad(features%nflat_local, 2), &
2022 0 : grad_grad(features%nflat_local, 3, 2), &
2023 56 : kin_grad(features%nflat_local, 2))
2024 8 : density_grad = 0.0_dp
2025 8 : grad_grad = 0.0_dp
2026 8 : kin_grad = 0.0_dp
2027 8 : local_row = 0
2028 198 : DO k = LBOUND(features%feature_index, 3), UBOUND(features%feature_index, 3)
2029 4436 : DO j = LBOUND(features%feature_index, 2), UBOUND(features%feature_index, 2)
2030 56739 : DO i = LBOUND(features%feature_index, 1), UBOUND(features%feature_index, 1)
2031 44847 : local_row = local_row + 1
2032 134287 : DO feature_pos = features%local_feature_offsets(local_row), &
2033 48753 : features%local_feature_offsets(local_row + 1) - 1
2034 89440 : row = features%local_feature_rows(feature_pos)
2035 89440 : CPASSERT(row >= 1 .AND. row <= features%nflat)
2036 : density_grad(local_row, :) = density_grad(local_row, :) + &
2037 268320 : density_grad_all(row, :)
2038 : grad_grad(local_row, :, :) = grad_grad(local_row, :, :) + &
2039 804960 : grad_grad_all(row, :, :)
2040 313167 : 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 8 : CPASSERT(local_row == features%nflat_local)
2046 :
2047 8 : CALL torch_tensor_release(density_grad_t)
2048 8 : CALL torch_tensor_release(grad_grad_t)
2049 8 : CALL torch_tensor_release(kin_grad_t)
2050 :
2051 8 : 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 40 : 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 40 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: chunk_density_grad, chunk_kin_grad
2073 40 : REAL(KIND=dp), DIMENSION(:, :, :), POINTER :: chunk_grad_grad
2074 : TYPE(torch_tensor_type) :: density_grad_t, grad_grad_t, kin_grad_t
2075 :
2076 40 : my_collapse_spin_grads = .FALSE.
2077 80 : IF (PRESENT(collapse_spin_grads)) my_collapse_spin_grads = collapse_spin_grads
2078 40 : ngrad_buffer_per_point = ngrad_per_point
2079 40 : IF (my_collapse_spin_grads) ngrad_buffer_per_point = ncollapsed_grad_per_point
2080 :
2081 40 : 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 40 : chunk_density_grad, chunk_grad_grad, chunk_kin_grad)
2084 40 : CPASSERT(MOD(SIZE(TARGET), ngrad_buffer_per_point) == 0)
2085 40 : target_points = SIZE(TARGET)/ngrad_buffer_per_point
2086 40 : CPASSERT(SIZE(chunk_density_grad, 1) == features%chunk_feature_count)
2087 40 : CPASSERT(SIZE(chunk_grad_grad, 1) == features%chunk_feature_count)
2088 40 : CPASSERT(SIZE(chunk_grad_grad, 2) == 3)
2089 40 : CPASSERT(SIZE(chunk_kin_grad, 1) == features%chunk_feature_count)
2090 40 : IF (features%uses_collapsed_rks_dynamic) THEN
2091 34 : CPASSERT(my_collapse_spin_grads)
2092 34 : CPASSERT(SIZE(chunk_density_grad, 2) == 1)
2093 34 : CPASSERT(SIZE(chunk_grad_grad, 3) == 1)
2094 34 : CPASSERT(SIZE(chunk_kin_grad, 2) == 1)
2095 : ELSE
2096 6 : CPASSERT(SIZE(chunk_density_grad, 2) == 2)
2097 6 : CPASSERT(SIZE(chunk_grad_grad, 3) == 2)
2098 6 : CPASSERT(SIZE(chunk_kin_grad, 2) == 2)
2099 : END IF
2100 :
2101 40 : IF (route_to_return_positions) THEN
2102 40 : 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 40 : !$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 0 : 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 0 : !$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 40 : CALL torch_tensor_release(density_grad_t)
2200 40 : CALL torch_tensor_release(grad_grad_t)
2201 40 : CALL torch_tensor_release(kin_grad_t)
2202 :
2203 40 : 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 40 : 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 40 : nflat_local = features%nflat_local
2231 40 : nroute_points = SIZE(features%route_send_local_rows)
2232 40 : nroute_grad_per_point = ngrad_per_point
2233 40 : IF (collapse_spin_grads) nroute_grad_per_point = ncollapsed_grad_per_point
2234 40 : CPASSERT(SIZE(recv_grad_buffer) >= nroute_grad_per_point*nroute_points)
2235 0 : ALLOCATE (density_grad(nflat_local, 2), grad_grad(nflat_local, 3, 2), &
2236 280 : 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 40 : !$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 40 : 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 48 : 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 48 : NULLIFY (density_grad, grad_grad, kin_grad)
2304 : CALL torch_tensor_grad_batch3(features%density_t, features%grad_t, features%kin_t, &
2305 48 : density_grad_t, grad_grad_t, kin_grad_t)
2306 48 : CALL torch_tensor_data_ptr(density_grad_t, density_grad)
2307 48 : CALL torch_tensor_data_ptr(grad_grad_t, grad_grad)
2308 48 : CALL torch_tensor_data_ptr(kin_grad_t, kin_grad)
2309 :
2310 48 : 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 38 : 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 38 : INTEGER, ALLOCATABLE, DIMENSION(:) :: route_grad_return_recv_counts, &
2335 38 : route_grad_return_recv_displs, &
2336 38 : route_grad_return_send_counts, &
2337 38 : route_grad_return_send_displs
2338 38 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: chunk_grad_buffer, global_grad_buffer, &
2339 38 : recv_grad_buffer, send_grad_buffer
2340 :
2341 38 : CPASSERT(features%uses_atom_chunks)
2342 :
2343 38 : nflat_local = features%nflat_local
2344 38 : IF (features%uses_atom_chunk_routing) THEN
2345 114 : nroute_recv_points = SUM(features%route_point_recv_counts)
2346 38 : nroute_send_points = SIZE(features%route_send_local_rows)
2347 114 : CPASSERT(SUM(features%route_point_send_counts) == nroute_send_points)
2348 38 : CPASSERT(SIZE(features%route_chunk_offsets) == nroute_recv_points + 1)
2349 :
2350 38 : nroute_grad_per_point = ngrad_per_point
2351 38 : IF (features%uses_collapsed_rks_dynamic) THEN
2352 32 : 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 494 : route_grad_return_recv_displs(SIZE(features%route_point_send_displs)))
2360 : route_grad_return_send_counts(:) = &
2361 114 : nroute_grad_per_point*features%route_point_recv_counts
2362 : route_grad_return_send_displs(:) = &
2363 114 : nroute_grad_per_point*features%route_point_recv_displs
2364 : route_grad_return_recv_counts(:) = &
2365 114 : nroute_grad_per_point*features%route_point_send_counts
2366 : route_grad_return_recv_displs(:) = &
2367 114 : nroute_grad_per_point*features%route_point_send_displs
2368 :
2369 38 : IF (features%chunk_feature_count > 0) THEN
2370 38 : CALL timeset("skala_gpw_grad_torch_pack", phase_handle)
2371 : CALL pack_atom_chunk_grads(features, send_grad_buffer, .TRUE., &
2372 38 : features%uses_collapsed_rks_dynamic)
2373 38 : CALL timestop(phase_handle)
2374 : END IF
2375 :
2376 38 : 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 38 : route_grad_return_recv_counts, route_grad_return_recv_displs)
2380 38 : CALL timestop(phase_handle)
2381 :
2382 38 : 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 38 : density_grad, grad_grad, kin_grad)
2386 38 : CALL timestop(phase_handle)
2387 :
2388 0 : DEALLOCATE (recv_grad_buffer, route_grad_return_recv_counts, &
2389 0 : route_grad_return_recv_displs, route_grad_return_send_counts, &
2390 114 : 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 0 : global_grad_buffer(ngrad_per_point*features%nflat))
2394 0 : IF (features%chunk_feature_count > 0) THEN
2395 0 : CALL timeset("skala_gpw_grad_torch_pack", phase_handle)
2396 0 : CALL pack_atom_chunk_grads(features, chunk_grad_buffer, .FALSE.)
2397 0 : CALL timestop(phase_handle)
2398 : END IF
2399 :
2400 0 : CALL timeset("skala_gpw_grad_allgatherv", phase_handle)
2401 : CALL group%allgatherv(chunk_grad_buffer, global_grad_buffer, &
2402 0 : features%chunk_grad_counts, features%chunk_grad_displs)
2403 0 : CALL timestop(phase_handle)
2404 :
2405 0 : CALL timeset("skala_gpw_grad_scatter", phase_handle)
2406 0 : ALLOCATE (density_grad(nflat_local, 2), grad_grad(nflat_local, 3, 2), &
2407 0 : kin_grad(nflat_local, 2))
2408 0 : density_grad = 0.0_dp
2409 0 : grad_grad = 0.0_dp
2410 0 : kin_grad = 0.0_dp
2411 0 : local_row = 0
2412 0 : DO k = LBOUND(features%feature_index, 3), UBOUND(features%feature_index, 3)
2413 0 : DO j = LBOUND(features%feature_index, 2), UBOUND(features%feature_index, 2)
2414 0 : DO i = LBOUND(features%feature_index, 1), UBOUND(features%feature_index, 1)
2415 0 : local_row = local_row + 1
2416 0 : DO feature_pos = features%local_feature_offsets(local_row), &
2417 0 : features%local_feature_offsets(local_row + 1) - 1
2418 0 : row = features%local_feature_rows(feature_pos)
2419 0 : CPASSERT(row >= 1 .AND. row <= features%nflat)
2420 0 : base = ngrad_per_point*(row - 1)
2421 : density_grad(local_row, :) = density_grad(local_row, :) + &
2422 0 : global_grad_buffer(base + 1:base + 2)
2423 : grad_grad(local_row, 1, 1) = grad_grad(local_row, 1, 1) + &
2424 0 : global_grad_buffer(base + 3)
2425 : grad_grad(local_row, 2, 1) = grad_grad(local_row, 2, 1) + &
2426 0 : global_grad_buffer(base + 4)
2427 : grad_grad(local_row, 3, 1) = grad_grad(local_row, 3, 1) + &
2428 0 : global_grad_buffer(base + 5)
2429 : grad_grad(local_row, 1, 2) = grad_grad(local_row, 1, 2) + &
2430 0 : global_grad_buffer(base + 6)
2431 : grad_grad(local_row, 2, 2) = grad_grad(local_row, 2, 2) + &
2432 0 : global_grad_buffer(base + 7)
2433 : grad_grad(local_row, 3, 2) = grad_grad(local_row, 3, 2) + &
2434 0 : global_grad_buffer(base + 8)
2435 : kin_grad(local_row, :) = kin_grad(local_row, :) + &
2436 0 : global_grad_buffer(base + 9:base + 10)
2437 : END DO
2438 : END DO
2439 : END DO
2440 : END DO
2441 0 : CALL timestop(phase_handle)
2442 0 : DEALLOCATE (chunk_grad_buffer, global_grad_buffer)
2443 :
2444 : END IF
2445 :
2446 38 : 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 6 : 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 72 : TYPE(cp_3d_r_cp_type), DIMENSION(3) :: drho, drhoa, drhob
2465 :
2466 6 : nspins = SIZE(rho_r)
2467 60 : bo = rho_r(1)%pw_grid%bounds_local
2468 6 : ipt = 0
2469 :
2470 6 : IF (nspins == 1) THEN
2471 6 : CALL xc_rho_set_get(rho_set, drho=drho)
2472 150 : DO k = bo(1, 3), bo(2, 3)
2473 3606 : DO j = bo(1, 2), bo(2, 2)
2474 45072 : DO i = bo(1, 1), bo(2, 1)
2475 41472 : ipt = ipt + 1
2476 169344 : DO idir = 1, 3
2477 124416 : grad_i = 0.5_dp*(grad_grad(ipt, idir, 1) + grad_grad(ipt, idir, 2))
2478 539136 : DO jdir = 1, 3
2479 373248 : tmp = -grad_i*drho(jdir)%array(i, j, k)
2480 497664 : 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 0 : CALL xc_rho_set_get(rho_set, drhoa=drhoa, drhob=drhob)
2488 0 : DO k = bo(1, 3), bo(2, 3)
2489 0 : DO j = bo(1, 2), bo(2, 2)
2490 0 : DO i = bo(1, 1), bo(2, 1)
2491 0 : ipt = ipt + 1
2492 0 : DO idir = 1, 3
2493 0 : DO jdir = 1, 3
2494 : tmp = 0.0_dp
2495 0 : DO ispin = 1, 2
2496 0 : IF (ispin == 1) THEN
2497 0 : tmp = tmp - grad_grad(ipt, idir, ispin)*drhoa(jdir)%array(i, j, k)
2498 : ELSE
2499 0 : tmp = tmp - grad_grad(ipt, idir, ispin)*drhob(jdir)%array(i, j, k)
2500 : END IF
2501 : END DO
2502 0 : 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 6 : 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 0 : 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 0 : IF (.NOT. root_rank) RETURN
2526 0 : iw = cp_logger_get_default_io_unit()
2527 0 : IF (iw <= 0) RETURN
2528 0 : WRITE (iw, "(T2,A,1X,A)") "SKALA_GPW| XC virial contribution", TRIM(label)
2529 0 : DO i = 1, 3
2530 0 : 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 6 : 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 6 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: atom_coord_grad, grid_coord_grad
2559 :
2560 6 : my_print_components = .FALSE.
2561 6 : IF (PRESENT(print_components)) my_print_components = print_components
2562 :
2563 6 : NULLIFY (atom_coord_grad, grid_coord_grad)
2564 6 : CALL torch_tensor_grad(features%grid_coords_t, grid_coord_grad_t)
2565 6 : CALL torch_tensor_data_ptr(grid_coord_grad_t, grid_coord_grad)
2566 6 : CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
2567 :
2568 6 : grid_virial = 0.0_dp
2569 6 : atom_virial = 0.0_dp
2570 6 : grid_grad_sum = 0.0_dp
2571 6 : atom_grad_sum = 0.0_dp
2572 6 : local_row = 0
2573 162 : DO k = LBOUND(features%feature_index, 3), UBOUND(features%feature_index, 3)
2574 3894 : DO j = LBOUND(features%feature_index, 2), UBOUND(features%feature_index, 2)
2575 51984 : DO i = LBOUND(features%feature_index, 1), UBOUND(features%feature_index, 1)
2576 41472 : local_row = local_row + 1
2577 124206 : DO feature_pos = features%local_feature_offsets(local_row), &
2578 44928 : features%local_feature_offsets(local_row + 1) - 1
2579 82734 : row = features%local_feature_rows(feature_pos)
2580 330936 : grid_grad_sum(:) = grid_grad_sum(:) + grid_coord_grad(:, row)
2581 372408 : DO idir = 1, 3
2582 1075542 : DO jdir = 1, 3
2583 744606 : tmp = grid_coord_grad(idir, row)*features%grid_coords(jdir, row)
2584 744606 : grid_virial(idir, jdir) = grid_virial(idir, jdir) + tmp
2585 992808 : 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 6 : CPASSERT(local_row == features%nflat_local)
2593 :
2594 6 : IF (root_rank) THEN
2595 9 : DO iatom = 1, SIZE(features%coarse_0_atomic_coords, 2)
2596 24 : atom_grad_sum(:) = atom_grad_sum(:) + atom_coord_grad(:, iatom)
2597 27 : DO idir = 1, 3
2598 78 : DO jdir = 1, 3
2599 54 : tmp = atom_coord_grad(idir, iatom)*features%coarse_0_atomic_coords(jdir, iatom)
2600 54 : atom_virial(idir, jdir) = atom_virial(idir, jdir) + tmp
2601 72 : virial_xc(idir, jdir) = virial_xc(idir, jdir) + tmp
2602 : END DO
2603 : END DO
2604 : END DO
2605 : END IF
2606 :
2607 6 : IF (my_print_components .AND. root_rank) THEN
2608 0 : iw = cp_logger_get_default_io_unit()
2609 0 : IF (iw > 0) THEN
2610 0 : CALL print_virial_delta("static-grid", grid_virial, .TRUE.)
2611 0 : CALL print_virial_delta("static-atom", atom_virial, .TRUE.)
2612 0 : WRITE (iw, "(T2,A,1X,3ES20.10)") "SKALA_GPW| XC coordinate gradient grid sum", &
2613 0 : grid_grad_sum
2614 0 : WRITE (iw, "(T2,A,1X,3ES20.10)") "SKALA_GPW| XC coordinate gradient atom sum", &
2615 0 : atom_grad_sum
2616 0 : WRITE (iw, "(T2,A,1X,3ES20.10)") "SKALA_GPW| XC coordinate gradient total sum", &
2617 0 : grid_grad_sum + atom_grad_sum
2618 : END IF
2619 : END IF
2620 :
2621 6 : CALL torch_tensor_release(grid_coord_grad_t)
2622 :
2623 6 : 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 6 : 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 6 : REAL(KIND=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
2650 :
2651 6 : my_print_components = .FALSE.
2652 6 : IF (PRESENT(print_components)) my_print_components = print_components
2653 :
2654 6 : NULLIFY (atomic_grid_weight_grad, grid_weight_grad)
2655 6 : CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
2656 6 : CALL torch_tensor_grad(features%atomic_grid_weights_t, atomic_grid_weight_grad_t)
2657 6 : CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
2658 6 : CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
2659 :
2660 6 : grid_tmp = 0.0_dp
2661 6 : atomic_tmp = 0.0_dp
2662 6 : local_row = 0
2663 162 : DO k = LBOUND(features%feature_index, 3), UBOUND(features%feature_index, 3)
2664 3894 : DO j = LBOUND(features%feature_index, 2), UBOUND(features%feature_index, 2)
2665 51984 : DO i = LBOUND(features%feature_index, 1), UBOUND(features%feature_index, 1)
2666 41472 : local_row = local_row + 1
2667 124206 : DO feature_pos = features%local_feature_offsets(local_row), &
2668 44928 : features%local_feature_offsets(local_row + 1) - 1
2669 82734 : row = features%local_feature_rows(feature_pos)
2670 82734 : grid_tmp = grid_tmp + grid_weight_grad(row)*features%grid_weights(row)
2671 : atomic_tmp = atomic_tmp + &
2672 124206 : atomic_grid_weight_grad(row)*features%atomic_grid_weights(row)
2673 : END DO
2674 : END DO
2675 : END DO
2676 : END DO
2677 6 : CPASSERT(local_row == features%nflat_local)
2678 6 : exc_tmp = 0.0_dp
2679 6 : IF (root_rank) exc_tmp = -exc
2680 6 : tmp = grid_tmp + atomic_tmp + exc_tmp
2681 :
2682 6 : IF (my_print_components .AND. root_rank) THEN
2683 0 : iw = cp_logger_get_default_io_unit()
2684 0 : IF (iw > 0) THEN
2685 0 : WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight grid", grid_tmp
2686 0 : WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight atomic", atomic_tmp
2687 0 : WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight final", exc_tmp
2688 0 : WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight residual", tmp
2689 : END IF
2690 : END IF
2691 :
2692 24 : DO idir = 1, 3
2693 24 : virial_xc(idir, idir) = virial_xc(idir, idir) + tmp
2694 : END DO
2695 :
2696 6 : CALL torch_tensor_release(grid_weight_grad_t)
2697 6 : CALL torch_tensor_release(atomic_grid_weight_grad_t)
2698 :
2699 6 : 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 312 : SUBROUTINE build_vxc_from_feature_grads(vxc_rho, vxc_tau, rho_r, pw_pool, &
2714 312 : 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 1248 : TYPE(pw_r3d_rs_type), DIMENSION(3) :: grad_pw
2730 :
2731 312 : nspins = SIZE(rho_r)
2732 3120 : bo = rho_r(1)%pw_grid%bounds_local
2733 312 : dvol_inv = 1.0_dp/rho_r(1)%pw_grid%dvol
2734 312 : my_global_grid_layout = .FALSE.
2735 312 : IF (PRESENT(global_grid_layout)) my_global_grid_layout = global_grid_layout
2736 264 : IF (my_global_grid_layout) THEN
2737 1056 : CPASSERT(SIZE(density_grad, 1) == PRODUCT(rho_r(1)%pw_grid%npts))
2738 : END IF
2739 312 : CPASSERT(SIZE(kin_grad, 2) == SIZE(density_grad, 2))
2740 312 : CPASSERT(SIZE(grad_grad, 3) == SIZE(density_grad, 2))
2741 312 : CPASSERT(SIZE(grad_grad, 2) == 3)
2742 : valid_spin_shape = SIZE(density_grad, 2) == nspins .OR. &
2743 312 : (nspins == 1 .AND. SIZE(density_grad, 2) == 2)
2744 0 : CPASSERT(valid_spin_shape)
2745 :
2746 1968 : ALLOCATE (vxc_rho(nspins), vxc_tau(nspins))
2747 672 : DO ispin = 1, nspins
2748 360 : CALL pw_pool%create_pw(vxc_rho(ispin))
2749 360 : CALL pw_pool%create_pw(vxc_tau(ispin))
2750 360 : CALL pw_zero(vxc_rho(ispin))
2751 672 : CALL pw_zero(vxc_tau(ispin))
2752 : END DO
2753 :
2754 312 : IF (xc_requires_tmp_g(xc_deriv_method_id) .OR. rho_r(1)%pw_grid%spherical) THEN
2755 312 : CALL pw_pool%create_pw(vxc_g)
2756 312 : IF (.NOT. rho_r(1)%pw_grid%spherical) CALL pw_pool%create_pw(tmp_g)
2757 : END IF
2758 :
2759 672 : DO ispin = 1, nspins
2760 1440 : DO i = 1, 3
2761 1080 : CALL pw_pool%create_pw(grad_pw(i))
2762 1440 : CALL pw_zero(grad_pw(i))
2763 : END DO
2764 :
2765 360 : ipt = 0
2766 7640 : DO k = bo(1, 3), bo(2, 3)
2767 179512 : DO j = bo(1, 2), bo(2, 2)
2768 2602484 : DO i = bo(1, 1), bo(2, 1)
2769 2423332 : 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 1183533 : k - rho_r(1)%pw_grid%bounds(1, 3)))
2775 : ELSE
2776 1239799 : ipt = ipt + 1
2777 : END IF
2778 2595204 : IF (nspins == 1) THEN
2779 1734832 : IF (SIZE(density_grad, 2) == 1) THEN
2780 1041783 : vxc_rho(1)%array(i, j, k) = 0.5_dp*dvol_inv*density_grad(ipt, 1)
2781 1041783 : vxc_tau(1)%array(i, j, k) = 0.5_dp*dvol_inv*kin_grad(ipt, 1)
2782 1041783 : grad_pw(1)%array(i, j, k) = 0.5_dp*dvol_inv*grad_grad(ipt, 1, 1)
2783 1041783 : grad_pw(2)%array(i, j, k) = 0.5_dp*dvol_inv*grad_grad(ipt, 2, 1)
2784 1041783 : 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 693049 : (density_grad(ipt, 1) + density_grad(ipt, 2))
2788 : vxc_tau(1)%array(i, j, k) = 0.5_dp*dvol_inv* &
2789 693049 : (kin_grad(ipt, 1) + kin_grad(ipt, 2))
2790 : grad_pw(1)%array(i, j, k) = 0.5_dp*dvol_inv* &
2791 693049 : (grad_grad(ipt, 1, 1) + grad_grad(ipt, 1, 2))
2792 : grad_pw(2)%array(i, j, k) = 0.5_dp*dvol_inv* &
2793 693049 : (grad_grad(ipt, 2, 1) + grad_grad(ipt, 2, 2))
2794 : grad_pw(3)%array(i, j, k) = 0.5_dp*dvol_inv* &
2795 693049 : (grad_grad(ipt, 3, 1) + grad_grad(ipt, 3, 2))
2796 : END IF
2797 : ELSE
2798 688500 : vxc_rho(ispin)%array(i, j, k) = dvol_inv*density_grad(ipt, ispin)
2799 688500 : vxc_tau(ispin)%array(i, j, k) = dvol_inv*kin_grad(ipt, ispin)
2800 688500 : grad_pw(1)%array(i, j, k) = dvol_inv*grad_grad(ipt, 1, ispin)
2801 688500 : grad_pw(2)%array(i, j, k) = dvol_inv*grad_grad(ipt, 2, ispin)
2802 688500 : 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 1440 : DO i = 1, 3
2809 1440 : CALL pw_scale(grad_pw(i), -1.0_dp)
2810 : END DO
2811 360 : CALL xc_pw_divergence(xc_deriv_method_id, grad_pw, tmp_g, vxc_g, vxc_rho(ispin))
2812 :
2813 1752 : DO i = 1, 3
2814 1440 : CALL pw_pool%give_back_pw(grad_pw(i))
2815 : END DO
2816 : END DO
2817 :
2818 312 : IF (ASSOCIATED(vxc_g%pw_grid)) CALL pw_pool%give_back_pw(vxc_g)
2819 312 : IF (ASSOCIATED(tmp_g%pw_grid)) CALL pw_pool%give_back_pw(tmp_g)
2820 :
2821 312 : 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 24 : 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 24 : IF (.NOT. print_active) RETURN
2837 :
2838 12 : iw = cp_logger_get_default_io_unit()
2839 12 : IF (iw <= 0) RETURN
2840 : WRITE (UNIT=iw, FMT="(/,T2,A,1X,ES19.11)") &
2841 12 : "SKALA_GPW| Native grid feature electrons", features%electron_count
2842 : WRITE (UNIT=iw, FMT="(T2,A,1X,ES19.11)") &
2843 12 : "SKALA_GPW| Native grid feature spin moment", features%spin_moment
2844 : WRITE (UNIT=iw, FMT="(T2,A,1X,ES19.11)") &
2845 12 : "SKALA_GPW| Native grid feature kinetic integral", features%kinetic_integral
2846 : WRITE (UNIT=iw, FMT="(T2,A,1X,ES19.11)") &
2847 12 : "SKALA_GPW| Native grid feature weight sum", features%grid_weight_sum
2848 12 : IF (ALLOCATED(features%atomic_grid_sizes)) THEN
2849 49 : atom_rows_min = INT(MINVAL(features%atomic_grid_sizes))
2850 49 : 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 12 : "SKALA_GPW| Native grid atom row range", atom_rows_min, "to", &
2853 61 : atom_rows_max, "sum", INT(SUM(features%atomic_grid_sizes))
2854 : END IF
2855 12 : IF (features%uses_atom_chunks) THEN
2856 : WRITE (UNIT=iw, FMT="(T2,A,1X,I0,1X,A,1X,I0)") &
2857 12 : "SKALA_GPW| Native grid atom chunk rows", features%chunk_feature_count, &
2858 24 : "of", features%nflat
2859 12 : IF (ALLOCATED(features%chunk_grad_counts)) THEN
2860 36 : chunk_rows_min = MINVAL(features%chunk_grad_counts)/ngrad_per_point
2861 36 : chunk_rows_max = MAXVAL(features%chunk_grad_counts)/ngrad_per_point
2862 12 : 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 12 : "SKALA_GPW| Native grid atom chunk row range", chunk_rows_min, &
2865 24 : "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 316 : 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 316 : INTEGER, ALLOCATABLE, DIMENSION(:) :: selected_devices
2886 :
2887 316 : selected_device = -1
2888 :
2889 316 : IF (.NOT. use_cuda) RETURN
2890 :
2891 0 : IF (.NOT. torch_cuda_is_available()) THEN
2892 0 : cuda_device_count = 0
2893 : ELSE
2894 0 : cuda_device_count = torch_cuda_device_count()
2895 : END IF
2896 0 : IF (cuda_device_count > 0) THEN
2897 0 : IF (requested_device < 0) THEN
2898 0 : selected_device = MOD(group%mepos, cuda_device_count)
2899 : ELSE
2900 0 : selected_device = requested_device
2901 : END IF
2902 : END IF
2903 0 : 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 0 : "Torch CUDA device range.")
2907 : END IF
2908 0 : IF (selected_device >= 0) CALL offload_set_chosen_device(selected_device)
2909 :
2910 0 : ALLOCATE (selected_devices(group%num_pe))
2911 0 : CALL group%allgather(selected_device, selected_devices)
2912 :
2913 0 : IF (group%mepos /= 0) THEN
2914 0 : DEALLOCATE (selected_devices)
2915 0 : RETURN
2916 : END IF
2917 : IF (selected_device == logged_cuda_device .AND. &
2918 : cuda_device_count == logged_cuda_device_count .AND. &
2919 0 : group%num_pe == logged_cuda_nproc .AND. &
2920 : requested_device == logged_cuda_request) THEN
2921 0 : DEALLOCATE (selected_devices)
2922 0 : RETURN
2923 : END IF
2924 :
2925 0 : iw = cp_logger_get_default_io_unit()
2926 0 : IF (iw <= 0) THEN
2927 0 : DEALLOCATE (selected_devices)
2928 0 : RETURN
2929 : END IF
2930 0 : IF (selected_device >= 0) THEN
2931 : WRITE (UNIT=iw, FMT="(/,T2,A,1X,I0,1X,A,1X,I0,1X,A,1X,I0)") &
2932 0 : "SKALA_GPW| Native grid Torch CUDA device", selected_device, &
2933 0 : "of", cuda_device_count, "requested", requested_device
2934 : ELSE
2935 : WRITE (UNIT=iw, FMT="(/,T2,A)") &
2936 0 : "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 0 : "SKALA_GPW| Native grid Torch CUDA rank devices"
2940 0 : DO pe = 1, group%num_pe
2941 0 : WRITE (UNIT=iw, FMT="(1X,I0,A,I0)", ADVANCE="NO") pe - 1, ":", selected_devices(pe)
2942 : END DO
2943 0 : WRITE (UNIT=iw, FMT=*)
2944 :
2945 0 : logged_cuda_device = selected_device
2946 0 : logged_cuda_device_count = cuda_device_count
2947 0 : logged_cuda_nproc = group%num_pe
2948 0 : logged_cuda_request = requested_device
2949 0 : DEALLOCATE (selected_devices)
2950 :
2951 316 : 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 294 : SUBROUTINE ensure_model_loaded(model_path, cuda_device)
2959 : CHARACTER(len=*), INTENT(IN) :: model_path
2960 : INTEGER, INTENT(IN) :: cuda_device
2961 :
2962 294 : IF (cached_model_loaded) THEN
2963 203 : IF (TRIM(cached_model_path) == TRIM(model_path) .AND. &
2964 : cached_model_cuda_device == cuda_device) RETURN
2965 0 : CALL skala_torch_model_release(cached_model)
2966 0 : cached_model_loaded = .FALSE.
2967 : END IF
2968 :
2969 91 : CALL skala_torch_model_load(cached_model, TRIM(model_path))
2970 91 : cached_model_path = model_path
2971 91 : cached_model_cuda_device = cuda_device
2972 91 : cached_model_loaded = .TRUE.
2973 :
2974 294 : 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 316 : 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 316 : gauxc_section => get_gauxc_section(xc_section)
2991 316 : IF (.NOT. ASSOCIATED(gauxc_section)) THEN
2992 0 : CPABORT("Native SKALA GPW requires an XC_FUNCTIONAL%GAUXC section")
2993 : END IF
2994 :
2995 316 : CALL section_vals_val_get(gauxc_section, "MODEL", c_val=model_path)
2996 316 : model_key = ADJUSTL(model_path)
2997 316 : CALL uppercase(model_key)
2998 316 : IF (TRIM(model_key) == "NONE" .OR. TRIM(model_key) == "") THEN
2999 0 : CPABORT("Native SKALA GPW requires GAUXC%MODEL SKALA or a TorchScript model path")
3000 316 : ELSE IF (TRIM(model_key) == "SKALA") THEN
3001 316 : CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_USE_CUDA", l_val=native_grid_use_cuda)
3002 316 : IF (native_grid_use_cuda) THEN
3003 0 : CALL GET_ENVIRONMENT_VARIABLE("GAUXC_SKALA_CUDA_MODEL", model_path, STATUS=env_status)
3004 0 : IF (env_status == 0 .AND. LEN_TRIM(model_path) > 0) RETURN
3005 : END IF
3006 316 : CALL GET_ENVIRONMENT_VARIABLE("GAUXC_SKALA_MODEL", model_path, STATUS=env_status)
3007 316 : IF (env_status /= 0 .OR. LEN_TRIM(model_path) == 0) THEN
3008 0 : IF (native_grid_use_cuda) THEN
3009 : CALL cp_abort(__LOCATION__, &
3010 0 : "MODEL SKALA CUDA path requires GAUXC_SKALA_CUDA_MODEL or GAUXC_SKALA_MODEL")
3011 : ELSE
3012 : CALL cp_abort(__LOCATION__, &
3013 0 : "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 701046 : 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 701046 : NULLIFY (gauxc_section)
3033 701046 : IF (.NOT. ASSOCIATED(xc_section)) RETURN
3034 :
3035 701046 : functionals => section_vals_get_subs_vals(xc_section, "XC_FUNCTIONAL")
3036 701046 : IF (.NOT. ASSOCIATED(functionals)) RETURN
3037 :
3038 701046 : ifun = 0
3039 : DO
3040 1353900 : ifun = ifun + 1
3041 1353900 : xc_fun => section_vals_get_subs_vals2(functionals, i_section=ifun)
3042 1353900 : IF (.NOT. ASSOCIATED(xc_fun)) EXIT
3043 1353900 : 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 :
3051 0 : END MODULE skala_gpw_functional
|
