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