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