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 Build SKALA TorchScript feature dictionaries from CP2K GPW real-space grids.
10 : ! **************************************************************************************************
11 : MODULE skala_gpw_features
12 : USE cell_types, ONLY: cell_type,&
13 : pbc
14 : USE cp_array_utils, ONLY: cp_3d_r_cp_type
15 : USE kinds, ONLY: dp,&
16 : int_8
17 : USE message_passing, ONLY: mp_comm_type
18 : USE particle_types, ONLY: particle_type
19 : USE pw_grid_types, ONLY: pw_grid_type
20 : USE pw_types, ONLY: pw_r3d_rs_type
21 : USE torch_api, ONLY: &
22 : torch_dict_clone, torch_dict_create, torch_dict_insert, torch_dict_release, &
23 : torch_dict_type, torch_tensor_expand_dim, torch_tensor_from_array, torch_tensor_narrow, &
24 : torch_tensor_release, torch_tensor_reset_from_array, torch_tensor_to_device_leaf, &
25 : torch_tensor_type
26 : USE xc_rho_set_types, ONLY: xc_rho_set_get,&
27 : xc_rho_set_type
28 : #include "./base/base_uses.f90"
29 :
30 : IMPLICIT NONE
31 :
32 : PRIVATE
33 :
34 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'skala_gpw_features'
35 : REAL(KIND=dp), PARAMETER, PRIVATE :: layout_tol = 1.0E-12_dp
36 : INTEGER, PARAMETER, PRIVATE :: ndynamic_per_point = 10, nrks_dynamic_per_point = 5, &
37 : nstatic_per_point = 5, ngrad_per_point = 10
38 : INTEGER, PARAMETER, PUBLIC :: skala_gpw_atom_partition_hard = 1, &
39 : skala_gpw_atom_partition_smooth = 2
40 : REAL(KIND=dp), PARAMETER, PRIVATE :: smooth_partition_eps = 1.0E-12_dp
41 : REAL(KIND=dp), PARAMETER, PRIVATE :: smooth_partition_taper_end = &
42 : 10.0_dp*smooth_partition_eps
43 :
44 : PUBLIC :: skala_gpw_atom_subchunk_count, skala_gpw_atom_subchunk_layout, &
45 : skala_gpw_feature_build, skala_gpw_feature_build_atom_subchunk, &
46 : skala_gpw_feature_build_atom_subchunk_bounds, skala_gpw_feature_release, &
47 : skala_gpw_feature_type, skala_gpw_smooth_partition_derivatives, &
48 : smooth_atom_partition, smooth_partition_atomic_weight_scale_derivative
49 :
50 : TYPE skala_gpw_layout_cache_type
51 : INTEGER :: chunk_atom_begin = 1, chunk_atom_end = 0, &
52 : chunk_feature_begin = 1, &
53 : chunk_feature_count = 0, chunk_natom = 0, &
54 : natom = 0, nflat = 0, nflat_local = 0, &
55 : npoint = 0, nproc = 0, &
56 : atom_partition = skala_gpw_atom_partition_hard
57 : INTEGER, DIMENSION(2, 3) :: bo = 0, bounds = 0
58 : INTEGER, DIMENSION(3) :: npts = 0
59 : INTEGER, ALLOCATABLE, DIMENSION(:) :: dynamic_counts, dynamic_displs, &
60 : chunk_feature_counts, chunk_feature_displs, &
61 : chunk_grad_counts, chunk_grad_displs, &
62 : feature_counts, feature_displs, &
63 : feature_source_points, global_to_feature, &
64 : local_feature_counts, local_feature_offsets, &
65 : local_feature_points, local_feature_rows, &
66 : chunk_return_positions, route_chunk_offsets, &
67 : route_chunk_rows, route_row_offsets, &
68 : route_row_positions, route_send_local_rows, &
69 : route_point_recv_counts, &
70 : route_point_recv_displs, &
71 : route_point_send_counts, &
72 : route_point_send_displs
73 : INTEGER, ALLOCATABLE, DIMENSION(:, :, :) :: feature_index
74 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: atomic_grid_sizes, chunk_atomic_grid_sizes, &
75 : chunk_feature_indices
76 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: local_feature_indices
77 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:, :) :: atomic_grid_size_bound_shape, &
78 : chunk_atomic_grid_size_bound_shape
79 : TYPE(torch_dict_type) :: chunk_inputs
80 : TYPE(torch_dict_type) :: chunk_static_inputs
81 : TYPE(torch_dict_type) :: inputs
82 : TYPE(torch_dict_type) :: static_inputs
83 : TYPE(torch_tensor_type) :: atomic_grid_size_bound_shape_t
84 : TYPE(torch_tensor_type) :: atomic_grid_sizes_t
85 : TYPE(torch_tensor_type) :: atomic_grid_weights_t
86 : TYPE(torch_tensor_type) :: chunk_atomic_grid_size_bound_shape_t
87 : TYPE(torch_tensor_type) :: chunk_atomic_grid_sizes_t
88 : TYPE(torch_tensor_type) :: chunk_atomic_grid_weights_t
89 : TYPE(torch_tensor_type) :: chunk_coarse_0_atomic_coords_t
90 : TYPE(torch_tensor_type) :: chunk_density_t
91 : TYPE(torch_tensor_type) :: chunk_density_input_t
92 : TYPE(torch_tensor_type) :: chunk_feature_indices_t
93 : TYPE(torch_tensor_type) :: chunk_grad_t
94 : TYPE(torch_tensor_type) :: chunk_grad_input_t
95 : TYPE(torch_tensor_type) :: chunk_grid_coords_t
96 : TYPE(torch_tensor_type) :: chunk_grid_weights_t
97 : TYPE(torch_tensor_type) :: chunk_kin_t
98 : TYPE(torch_tensor_type) :: chunk_kin_input_t
99 : TYPE(torch_tensor_type) :: coarse_0_atomic_coords_t
100 : TYPE(torch_tensor_type) :: density_t
101 : TYPE(torch_tensor_type) :: grid_coords_t
102 : TYPE(torch_tensor_type) :: grid_weights_t
103 : TYPE(torch_tensor_type) :: grad_t
104 : TYPE(torch_tensor_type) :: kin_t
105 : TYPE(torch_tensor_type) :: local_feature_indices_t
106 : REAL(KIND=dp) :: dvol = 0.0_dp, weight_sum = 0.0_dp, &
107 : weight_sumsq = 0.0_dp
108 : REAL(KIND=dp), DIMENSION(3, 3) :: cell_hmat = 0.0_dp, dh = 0.0_dp
109 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: atomic_grid_weights, chunk_atomic_grid_weights, &
110 : chunk_grid_weights, grid_weights
111 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: atom_coords, chunk_coarse_0_atomic_coords, &
112 : chunk_grid_coords, coarse_0_atomic_coords, &
113 : grid_coords
114 : LOGICAL :: active = .FALSE., has_weights = .FALSE., &
115 : chunk_dynamic_input_views_active = .FALSE., &
116 : chunk_dynamic_tensors_active = .FALSE., &
117 : chunk_inputs_active = .FALSE., &
118 : chunk_inputs_use_collapsed_rks = .FALSE., &
119 : chunk_static_tensors_active = .FALSE., &
120 : dynamic_tensors_active = .FALSE., &
121 : inputs_active = .FALSE., &
122 : static_tensors_active = .FALSE.
123 : END TYPE skala_gpw_layout_cache_type
124 :
125 : TYPE skala_gpw_feature_type
126 : INTEGER :: chunk_feature_count = 0, nflat = 0, &
127 : nflat_local = 0, &
128 : atom_partition = skala_gpw_atom_partition_hard
129 : TYPE(torch_dict_type) :: inputs
130 : TYPE(torch_tensor_type) :: atomic_grid_size_bound_shape_t
131 : TYPE(torch_tensor_type) :: atomic_grid_sizes_t
132 : TYPE(torch_tensor_type) :: atomic_grid_weights_t
133 : TYPE(torch_tensor_type) :: coarse_0_atomic_coords_t
134 : TYPE(torch_tensor_type) :: density_input_t
135 : TYPE(torch_tensor_type) :: density_t
136 : TYPE(torch_tensor_type) :: grad_t
137 : TYPE(torch_tensor_type) :: grad_input_t
138 : TYPE(torch_tensor_type) :: grid_coords_t
139 : TYPE(torch_tensor_type) :: grid_weights_t
140 : TYPE(torch_tensor_type) :: kin_input_t
141 : TYPE(torch_tensor_type) :: kin_t
142 : TYPE(torch_tensor_type) :: local_feature_indices_t
143 : INTEGER, ALLOCATABLE, DIMENSION(:) :: chunk_grad_counts, chunk_grad_displs, &
144 : local_feature_counts, local_feature_offsets, &
145 : local_feature_rows
146 : INTEGER, POINTER, DIMENSION(:) :: route_point_recv_counts => NULL(), &
147 : route_point_recv_displs => NULL(), &
148 : route_point_send_counts => NULL(), &
149 : route_point_send_displs => NULL(), &
150 : route_chunk_offsets => NULL(), &
151 : route_chunk_rows => NULL(), &
152 : route_row_offsets => NULL(), &
153 : route_row_positions => NULL(), &
154 : route_send_local_rows => NULL()
155 : INTEGER, ALLOCATABLE, DIMENSION(:, :, :) :: feature_index
156 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: atomic_grid_sizes
157 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:, :) :: atomic_grid_size_bound_shape
158 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: atomic_grid_weights, grid_weights
159 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: chunk_density, chunk_kin, &
160 : coarse_0_atomic_coords, density, &
161 : grid_coords, kin
162 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: chunk_grad, grad
163 : REAL(KIND=dp) :: electron_count = 0.0_dp, &
164 : grid_weight_sum = 0.0_dp, &
165 : kinetic_integral = 0.0_dp, &
166 : spin_moment = 0.0_dp
167 : LOGICAL :: active = .FALSE., owns_coordinate_tensor = .FALSE., &
168 : owns_grid_coordinate_tensor = .FALSE., &
169 : owns_weight_tensors = .FALSE., &
170 : owns_dynamic_tensors = .TRUE., &
171 : owns_inputs = .TRUE., &
172 : owns_static_tensors = .TRUE., &
173 : uses_atom_chunk_routing = .FALSE., &
174 : uses_atom_chunks = .FALSE., &
175 : uses_collapsed_rks_dynamic = .FALSE.
176 : END TYPE skala_gpw_feature_type
177 :
178 : TYPE(skala_gpw_layout_cache_type), SAVE, TARGET :: cached_layout
179 :
180 : CONTAINS
181 :
182 : ! **************************************************************************************************
183 : !> \brief Build a flat SKALA molecular feature dictionary from a local GPW grid.
184 : !> \param features ...
185 : !> \param rho_set ...
186 : !> \param rho_r ...
187 : !> \param particle_set ...
188 : !> \param cell ...
189 : !> \param requires_grad ...
190 : !> \param weights ...
191 : !> \param requires_coordinate_grad ...
192 : !> \param requires_stress_grad ...
193 : !> \param use_atom_chunks ...
194 : !> \param route_atom_chunks ...
195 : !> \param atom_partition ...
196 : ! **************************************************************************************************
197 322 : SUBROUTINE skala_gpw_feature_build(features, rho_set, rho_r, particle_set, cell, &
198 : requires_grad, weights, requires_coordinate_grad, &
199 : requires_stress_grad, use_atom_chunks, route_atom_chunks, &
200 : atom_partition)
201 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
202 : TYPE(xc_rho_set_type), INTENT(IN) :: rho_set
203 : TYPE(pw_r3d_rs_type), DIMENSION(:), INTENT(IN) :: rho_r
204 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
205 : TYPE(cell_type), POINTER :: cell
206 : LOGICAL, INTENT(IN), OPTIONAL :: requires_grad
207 : TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
208 : LOGICAL, INTENT(IN), OPTIONAL :: requires_coordinate_grad, &
209 : requires_stress_grad, use_atom_chunks, &
210 : route_atom_chunks
211 : INTEGER, INTENT(IN), OPTIONAL :: atom_partition
212 :
213 : INTEGER :: handle, i, ipt, ispin, j, k, local_row, my_atom_partition, &
214 : ndynamic_local_per_point, nflat, nflat_local, nspins, phase_handle, real_base, row
215 : INTEGER, DIMENSION(2, 3) :: bo
216 : LOGICAL :: collapse_spin_dynamics, my_requires_coordinate_grad, my_requires_grad, &
217 : my_requires_stress_grad, my_route_atom_chunks, my_use_atom_chunks, &
218 : use_atom_chunk_protocol, use_atom_chunk_routing
219 322 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: global_dynamic, local_dynamic
220 322 : REAL(KIND=dp), DIMENSION(:, :, :), POINTER :: rho, rhoa, rhob, tau_a, tau_b, tau_total
221 3864 : TYPE(cp_3d_r_cp_type), DIMENSION(3) :: drho, drhoa, drhob
222 : TYPE(pw_grid_type), POINTER :: pw_grid
223 :
224 322 : CALL timeset("skala_gpw_feature_build", handle)
225 :
226 322 : my_requires_grad = .FALSE.
227 322 : IF (PRESENT(requires_grad)) my_requires_grad = requires_grad
228 322 : my_requires_coordinate_grad = .FALSE.
229 322 : IF (PRESENT(requires_coordinate_grad)) THEN
230 322 : my_requires_coordinate_grad = requires_coordinate_grad
231 : END IF
232 322 : my_requires_stress_grad = .FALSE.
233 322 : IF (PRESENT(requires_stress_grad)) my_requires_stress_grad = requires_stress_grad
234 322 : my_use_atom_chunks = .FALSE.
235 322 : IF (PRESENT(use_atom_chunks)) my_use_atom_chunks = use_atom_chunks
236 322 : my_route_atom_chunks = .FALSE.
237 322 : IF (PRESENT(route_atom_chunks)) my_route_atom_chunks = route_atom_chunks
238 322 : my_atom_partition = skala_gpw_atom_partition_hard
239 322 : IF (PRESENT(atom_partition)) my_atom_partition = atom_partition
240 322 : IF (my_atom_partition /= skala_gpw_atom_partition_hard .AND. &
241 : my_atom_partition /= skala_gpw_atom_partition_smooth) THEN
242 0 : CALL cp_abort(__LOCATION__, "Unknown native SKALA atom-partition mode.")
243 : END IF
244 322 : CPASSERT(ASSOCIATED(cell))
245 322 : CPASSERT(ASSOCIATED(particle_set))
246 322 : CPASSERT(SIZE(rho_r) == 1 .OR. SIZE(rho_r) == 2)
247 322 : CPASSERT(ASSOCIATED(rho_r(1)%pw_grid))
248 322 : pw_grid => rho_r(1)%pw_grid
249 :
250 322 : nspins = SIZE(rho_r)
251 3220 : bo = pw_grid%bounds_local
252 322 : nflat_local = pw_grid%ngpts_local
253 :
254 322 : CALL timeset("skala_gpw_pre_release", phase_handle)
255 322 : CALL skala_gpw_feature_release(features)
256 322 : CALL timestop(phase_handle)
257 :
258 : use_atom_chunk_protocol = my_use_atom_chunks .AND. &
259 322 : .NOT. (my_requires_coordinate_grad .OR. my_requires_stress_grad)
260 322 : use_atom_chunk_routing = use_atom_chunk_protocol .AND. my_route_atom_chunks
261 322 : CALL timeset("skala_gpw_layout_cache", phase_handle)
262 : CALL ensure_layout_cache(pw_grid, particle_set, cell, weights, my_atom_partition, &
263 322 : needs_full_static_tensors=.NOT. use_atom_chunk_protocol)
264 322 : CALL timestop(phase_handle)
265 322 : nflat = cached_layout%nflat
266 322 : collapse_spin_dynamics = nspins == 1 .AND. use_atom_chunk_routing
267 322 : ndynamic_local_per_point = ndynamic_per_point
268 322 : IF (collapse_spin_dynamics) ndynamic_local_per_point = nrks_dynamic_per_point
269 966 : ALLOCATE (local_dynamic(ndynamic_local_per_point*nflat_local))
270 322 : local_dynamic = 0.0_dp
271 :
272 322 : CALL timeset("skala_gpw_pack_local", phase_handle)
273 322 : IF (nspins == 1) THEN
274 274 : CALL xc_rho_set_get(rho_set, rho=rho, drho=drho, tau=tau_total)
275 : ELSE
276 : CALL xc_rho_set_get(rho_set, rhoa=rhoa, rhob=rhob, drhoa=drhoa, drhob=drhob, &
277 48 : tau_a=tau_a, tau_b=tau_b)
278 : END IF
279 :
280 322 : local_row = 0
281 6720 : DO k = bo(1, 3), bo(2, 3)
282 152510 : DO j = bo(1, 2), bo(2, 2)
283 2116655 : DO i = bo(1, 1), bo(2, 1)
284 1964467 : local_row = local_row + 1
285 1964467 : real_base = ndynamic_local_per_point*(local_row - 1)
286 :
287 2110257 : IF (nspins == 1) THEN
288 1620217 : IF (collapse_spin_dynamics) THEN
289 1301893 : local_dynamic(real_base + 1) = 0.5_dp*rho(i, j, k)
290 1301893 : local_dynamic(real_base + 2) = 0.5_dp*drho(1)%array(i, j, k)
291 1301893 : local_dynamic(real_base + 3) = 0.5_dp*drho(2)%array(i, j, k)
292 1301893 : local_dynamic(real_base + 4) = 0.5_dp*drho(3)%array(i, j, k)
293 1301893 : local_dynamic(real_base + 5) = 0.5_dp*tau_total(i, j, k)
294 : ELSE
295 318324 : local_dynamic(real_base + 1) = 0.5_dp*rho(i, j, k)
296 318324 : local_dynamic(real_base + 2) = 0.5_dp*rho(i, j, k)
297 954972 : DO ispin = 1, 2
298 : local_dynamic(real_base + 2 + 3*(ispin - 1) + 1) = &
299 636648 : 0.5_dp*drho(1)%array(i, j, k)
300 : local_dynamic(real_base + 2 + 3*(ispin - 1) + 2) = &
301 636648 : 0.5_dp*drho(2)%array(i, j, k)
302 : local_dynamic(real_base + 2 + 3*(ispin - 1) + 3) = &
303 636648 : 0.5_dp*drho(3)%array(i, j, k)
304 954972 : local_dynamic(real_base + 8 + ispin) = 0.5_dp*tau_total(i, j, k)
305 : END DO
306 : END IF
307 : ELSE
308 344250 : local_dynamic(real_base + 1) = rhoa(i, j, k)
309 344250 : local_dynamic(real_base + 2) = rhob(i, j, k)
310 344250 : local_dynamic(real_base + 3) = drhoa(1)%array(i, j, k)
311 344250 : local_dynamic(real_base + 4) = drhoa(2)%array(i, j, k)
312 344250 : local_dynamic(real_base + 5) = drhoa(3)%array(i, j, k)
313 344250 : local_dynamic(real_base + 6) = drhob(1)%array(i, j, k)
314 344250 : local_dynamic(real_base + 7) = drhob(2)%array(i, j, k)
315 344250 : local_dynamic(real_base + 8) = drhob(3)%array(i, j, k)
316 344250 : local_dynamic(real_base + 9) = tau_a(i, j, k)
317 344250 : local_dynamic(real_base + 10) = tau_b(i, j, k)
318 : END IF
319 : END DO
320 : END DO
321 : END DO
322 322 : CALL timestop(phase_handle)
323 :
324 322 : CALL timeset("skala_gpw_copy_layout", phase_handle)
325 : CALL copy_cached_layout(features, use_atom_chunk_routing, &
326 : my_requires_coordinate_grad .OR. my_requires_stress_grad, &
327 : my_requires_stress_grad .OR. &
328 : (my_atom_partition == skala_gpw_atom_partition_smooth .AND. &
329 580 : (my_requires_coordinate_grad .OR. my_requires_stress_grad)))
330 322 : CALL timestop(phase_handle)
331 :
332 322 : IF (use_atom_chunk_routing) THEN
333 258 : CALL timeset("skala_gpw_route_dyn", phase_handle)
334 : CALL route_atom_chunk_dynamics(features, local_dynamic, pw_grid%para%group, &
335 258 : collapse_spin_dynamics)
336 258 : features%uses_atom_chunk_routing = .TRUE.
337 258 : features%uses_atom_chunks = .TRUE.
338 258 : CALL timestop(phase_handle)
339 : ELSE
340 192 : ALLOCATE (global_dynamic(ndynamic_per_point*cached_layout%npoint))
341 64 : CALL timeset("skala_gpw_allgatherv", phase_handle)
342 : CALL pw_grid%para%group%allgatherv(local_dynamic, global_dynamic, &
343 : cached_layout%dynamic_counts, &
344 64 : cached_layout%dynamic_displs)
345 64 : CALL timestop(phase_handle)
346 :
347 64 : CALL timeset("skala_gpw_reorder_dyn", phase_handle)
348 0 : ALLOCATE (features%density(nflat, 2), features%grad(nflat, 3, 2), &
349 448 : features%kin(nflat, 2))
350 2567256 : features%density = 0.0_dp
351 7701768 : features%grad = 0.0_dp
352 2567256 : features%kin = 0.0_dp
353 :
354 1283596 : DO row = 1, nflat
355 1283532 : ipt = cached_layout%feature_source_points(row)
356 1283532 : real_base = ndynamic_per_point*(ipt - 1)
357 3850596 : features%density(row, :) = global_dynamic(real_base + 1:real_base + 2)
358 1283532 : features%grad(row, 1, 1) = global_dynamic(real_base + 3)
359 1283532 : features%grad(row, 2, 1) = global_dynamic(real_base + 4)
360 1283532 : features%grad(row, 3, 1) = global_dynamic(real_base + 5)
361 1283532 : features%grad(row, 1, 2) = global_dynamic(real_base + 6)
362 1283532 : features%grad(row, 2, 2) = global_dynamic(real_base + 7)
363 1283532 : features%grad(row, 3, 2) = global_dynamic(real_base + 8)
364 3850660 : features%kin(row, :) = global_dynamic(real_base + 9:real_base + 10)
365 : END DO
366 192 : CALL timestop(phase_handle)
367 : END IF
368 :
369 322 : CALL timeset("skala_gpw_feature_sums", phase_handle)
370 322 : IF (features%uses_atom_chunks) THEN
371 258 : features%electron_count = 0.0_dp
372 258 : features%kinetic_integral = 0.0_dp
373 258 : features%spin_moment = 0.0_dp
374 258 : IF (features%chunk_feature_count > 0) THEN
375 258 : IF (features%uses_collapsed_rks_dynamic) THEN
376 : features%electron_count = SUM(2.0_dp*features%chunk_density(:, 1)* &
377 1887014 : cached_layout%chunk_grid_weights)
378 : features%kinetic_integral = SUM(2.0_dp*features%chunk_kin(:, 1)* &
379 1887014 : cached_layout%chunk_grid_weights)
380 : ELSE
381 : features%electron_count = SUM((features%chunk_density(:, 1) + &
382 : features%chunk_density(:, 2))* &
383 408021 : cached_layout%chunk_grid_weights)
384 : features%kinetic_integral = SUM((features%chunk_kin(:, 1) + &
385 : features%chunk_kin(:, 2))* &
386 408021 : cached_layout%chunk_grid_weights)
387 : features%spin_moment = SUM((features%chunk_density(:, 1) - &
388 : features%chunk_density(:, 2))* &
389 408021 : cached_layout%chunk_grid_weights)
390 : END IF
391 : END IF
392 258 : CALL pw_grid%para%group%sum(features%electron_count)
393 258 : CALL pw_grid%para%group%sum(features%kinetic_integral)
394 258 : CALL pw_grid%para%group%sum(features%spin_moment)
395 : ELSE
396 : features%electron_count = SUM((features%density(:, 1) + features%density(:, 2))* &
397 1283596 : features%grid_weights)
398 : features%kinetic_integral = SUM((features%kin(:, 1) + features%kin(:, 2))* &
399 1283596 : features%grid_weights)
400 : features%spin_moment = SUM((features%density(:, 1) - features%density(:, 2))* &
401 1283596 : features%grid_weights)
402 : END IF
403 322 : IF (use_atom_chunk_routing) THEN
404 4589812 : features%grid_weight_sum = SUM(cached_layout%grid_weights)
405 : ELSE
406 1283596 : features%grid_weight_sum = SUM(features%grid_weights)
407 : END IF
408 322 : CALL timestop(phase_handle)
409 :
410 322 : CALL timeset("skala_gpw_tensor_update", phase_handle)
411 322 : IF (use_atom_chunk_protocol .AND. .NOT. features%uses_atom_chunks) THEN
412 0 : IF (features%chunk_feature_count > 0) CALL extract_atom_chunk_dynamics(features)
413 0 : features%uses_atom_chunks = .TRUE.
414 : END IF
415 322 : IF (.NOT. features%uses_atom_chunks .OR. features%chunk_feature_count > 0) THEN
416 : CALL add_feature_tensors(features, my_requires_grad, my_requires_coordinate_grad, &
417 : my_requires_stress_grad, &
418 : features%uses_atom_chunks, &
419 : requires_weight_grad= &
420 : (my_atom_partition == skala_gpw_atom_partition_smooth .AND. &
421 580 : (my_requires_coordinate_grad .OR. my_requires_stress_grad)))
422 : ELSE
423 : ! This rank participates in atom-chunk communication but owns no model input rows.
424 0 : features%owns_coordinate_tensor = .FALSE.
425 0 : features%owns_grid_coordinate_tensor = .FALSE.
426 0 : features%owns_weight_tensors = .FALSE.
427 0 : features%owns_dynamic_tensors = .FALSE.
428 0 : features%owns_inputs = .FALSE.
429 0 : features%owns_static_tensors = .FALSE.
430 : END IF
431 322 : CALL timestop(phase_handle)
432 322 : features%active = .TRUE.
433 :
434 322 : IF (ALLOCATED(global_dynamic)) DEALLOCATE (global_dynamic)
435 322 : DEALLOCATE (local_dynamic)
436 322 : CALL timestop(handle)
437 :
438 2576 : END SUBROUTINE skala_gpw_feature_build
439 :
440 : ! **************************************************************************************************
441 : !> \brief Ensure that static grid-to-atom layout data is cached for the current grid/geometry.
442 : !> \param pw_grid ...
443 : !> \param particle_set ...
444 : !> \param cell ...
445 : !> \param weights ...
446 : !> \param atom_partition ...
447 : !> \param needs_full_static_tensors ...
448 : ! **************************************************************************************************
449 454 : SUBROUTINE ensure_layout_cache(pw_grid, particle_set, cell, weights, atom_partition, &
450 : needs_full_static_tensors)
451 : TYPE(pw_grid_type), POINTER :: pw_grid
452 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
453 : TYPE(cell_type), POINTER :: cell
454 : TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
455 : INTEGER, INTENT(IN), OPTIONAL :: atom_partition
456 : LOGICAL, INTENT(IN) :: needs_full_static_tensors
457 :
458 : INTEGER :: my_atom_partition, phase_handle
459 : LOGICAL :: cache_matches
460 :
461 322 : my_atom_partition = skala_gpw_atom_partition_hard
462 322 : IF (PRESENT(atom_partition)) my_atom_partition = atom_partition
463 322 : IF (PRESENT(weights)) THEN
464 322 : CALL timeset("skala_gpw_layout_match", phase_handle)
465 : cache_matches = layout_cache_matches(pw_grid, particle_set, cell, weights, &
466 322 : my_atom_partition)
467 322 : CALL timestop(phase_handle)
468 : ELSE
469 0 : CALL timeset("skala_gpw_layout_match", phase_handle)
470 : cache_matches = layout_cache_matches(pw_grid, particle_set, cell, &
471 0 : atom_partition=my_atom_partition)
472 0 : CALL timestop(phase_handle)
473 : END IF
474 :
475 322 : IF (cache_matches) THEN
476 190 : IF (needs_full_static_tensors .AND. .NOT. cached_layout%static_tensors_active) THEN
477 64 : CALL timeset("skala_gpw_layout_tensors", phase_handle)
478 64 : CALL build_full_static_layout_tensors(cached_layout)
479 64 : CALL timestop(phase_handle)
480 : END IF
481 190 : RETURN
482 : END IF
483 :
484 132 : CALL timeset("skala_gpw_layout_rebuild", phase_handle)
485 132 : IF (PRESENT(weights)) THEN
486 : CALL rebuild_layout_cache(pw_grid, particle_set, cell, weights, my_atom_partition, &
487 132 : needs_full_static_tensors)
488 : ELSE
489 : CALL rebuild_layout_cache(pw_grid, particle_set, cell, &
490 : atom_partition=my_atom_partition, &
491 0 : needs_full_static_tensors=needs_full_static_tensors)
492 : END IF
493 132 : CALL timestop(phase_handle)
494 :
495 : END SUBROUTINE ensure_layout_cache
496 :
497 : ! **************************************************************************************************
498 : !> \brief Check whether the current static layout cache can be reused.
499 : !> \param pw_grid ...
500 : !> \param particle_set ...
501 : !> \param cell ...
502 : !> \param weights ...
503 : !> \param atom_partition ...
504 : !> \return ...
505 : ! **************************************************************************************************
506 322 : FUNCTION layout_cache_matches(pw_grid, particle_set, cell, weights, atom_partition) RESULT(matches)
507 : TYPE(pw_grid_type), POINTER :: pw_grid
508 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
509 : TYPE(cell_type), POINTER :: cell
510 : TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
511 : INTEGER, INTENT(IN), OPTIONAL :: atom_partition
512 : LOGICAL :: matches
513 :
514 : INTEGER :: iatom, my_atom_partition
515 : LOGICAL :: weights_match
516 :
517 322 : my_atom_partition = skala_gpw_atom_partition_hard
518 322 : IF (PRESENT(atom_partition)) my_atom_partition = atom_partition
519 322 : matches = .FALSE.
520 322 : IF (.NOT. cached_layout%active) RETURN
521 230 : IF (cached_layout%atom_partition /= my_atom_partition) RETURN
522 230 : IF (cached_layout%natom /= SIZE(particle_set)) RETURN
523 230 : IF (cached_layout%nflat_local /= pw_grid%ngpts_local) RETURN
524 230 : IF (cached_layout%nproc /= pw_grid%para%group%num_pe) RETURN
525 2300 : IF (ANY(cached_layout%bo /= pw_grid%bounds_local)) RETURN
526 2300 : IF (ANY(cached_layout%bounds /= pw_grid%bounds)) RETURN
527 920 : IF (ANY(cached_layout%npts /= pw_grid%npts)) RETURN
528 230 : IF (ABS(cached_layout%dvol - pw_grid%dvol) > layout_tol) RETURN
529 2626 : IF (ANY(ABS(cached_layout%dh - pw_grid%dh) > layout_tol)) RETURN
530 2626 : IF (ANY(ABS(cached_layout%cell_hmat - cell%hmat) > layout_tol)) RETURN
531 202 : IF (.NOT. ALLOCATED(cached_layout%atom_coords)) RETURN
532 :
533 594 : DO iatom = 1, SIZE(particle_set)
534 1794 : IF (ANY(ABS(cached_layout%atom_coords(:, iatom) - particle_set(iatom)%r) > layout_tol)) RETURN
535 : END DO
536 :
537 190 : IF (PRESENT(weights)) THEN
538 190 : weights_match = layout_weights_match(pw_grid, weights)
539 : ELSE
540 0 : weights_match = layout_weights_match(pw_grid)
541 : END IF
542 190 : IF (.NOT. weights_match) RETURN
543 :
544 322 : matches = .TRUE.
545 :
546 : END FUNCTION layout_cache_matches
547 :
548 : ! **************************************************************************************************
549 : !> \brief Check whether current optional integration weights match the cached static tensors.
550 : !> \param pw_grid ...
551 : !> \param weights ...
552 : !> \return ...
553 : ! **************************************************************************************************
554 190 : FUNCTION layout_weights_match(pw_grid, weights) RESULT(matches)
555 : TYPE(pw_grid_type), POINTER :: pw_grid
556 : TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
557 : LOGICAL :: matches
558 :
559 : LOGICAL :: has_weights
560 : REAL(KIND=dp) :: weight_sum, weight_sumsq
561 :
562 190 : matches = .FALSE.
563 : MARK_USED(pw_grid)
564 190 : IF (PRESENT(weights)) THEN
565 190 : CALL weights_signature(weights, has_weights, weight_sum, weight_sumsq)
566 : ELSE
567 : CALL weights_signature(has_weights=has_weights, weight_sum=weight_sum, &
568 0 : weight_sumsq=weight_sumsq)
569 : END IF
570 :
571 190 : IF (cached_layout%has_weights .NEQV. has_weights) RETURN
572 190 : IF (ABS(cached_layout%weight_sum - weight_sum) > layout_tol) RETURN
573 190 : IF (ABS(cached_layout%weight_sumsq - weight_sumsq) > layout_tol) RETURN
574 :
575 190 : matches = .TRUE.
576 :
577 : END FUNCTION layout_weights_match
578 :
579 : ! **************************************************************************************************
580 : !> \brief Build the static SKALA layout cache.
581 : !> \param pw_grid ...
582 : !> \param particle_set ...
583 : !> \param cell ...
584 : !> \param weights ...
585 : !> \param atom_partition ...
586 : !> \param needs_full_static_tensors ...
587 : ! **************************************************************************************************
588 132 : SUBROUTINE rebuild_layout_cache(pw_grid, particle_set, cell, weights, atom_partition, &
589 : needs_full_static_tensors)
590 : TYPE(pw_grid_type), POINTER :: pw_grid
591 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
592 : TYPE(cell_type), POINTER :: cell
593 : TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
594 : INTEGER, INTENT(IN), OPTIONAL :: atom_partition
595 : LOGICAL, INTENT(IN) :: needs_full_static_tensors
596 :
597 : INTEGER :: feature_local, feature_slot, i, iatom, ipt, j, jatom, k, local_feature, &
598 : local_row, max_grid_size, max_local_features, my_atom_partition, natom, nfeature_local, &
599 : nflat, nflat_local, npoint, nproc, nx, ny, owner, pe, pe_index, phase_handle, row, &
600 : source_global, source_local, static_base
601 132 : INTEGER, ALLOCATABLE, DIMENSION(:) :: atom_offset, atom_position, chunk_atom_begin, &
602 132 : chunk_atom_end, cursor, feature_counts, feature_displs, global_owner, &
603 132 : global_source_points, local_feature_counts_tmp, local_owner, local_source_global, &
604 132 : local_source_points, point_counts, point_displs, static_counts, static_displs
605 : INTEGER, DIMENSION(2, 3) :: bo
606 : LOGICAL :: has_weights, nonperiodic
607 : REAL(KIND=dp) :: atomic_base_weight, base_weight, &
608 : included_sum, partition_weight, &
609 : weight_sum, weight_sumsq
610 132 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: distances, global_static, local_static, &
611 132 : partition_weights
612 132 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: atom_coords_pbc, pair_distances, &
613 132 : partition_atom_coords
614 : REAL(KIND=dp), DIMENSION(3) :: grid_point
615 :
616 132 : CALL release_layout_cache(cached_layout)
617 :
618 132 : my_atom_partition = skala_gpw_atom_partition_hard
619 132 : IF (PRESENT(atom_partition)) my_atom_partition = atom_partition
620 132 : natom = SIZE(particle_set)
621 1320 : bo = pw_grid%bounds_local
622 132 : nflat_local = pw_grid%ngpts_local
623 132 : nproc = pw_grid%para%group%num_pe
624 132 : pe_index = pw_grid%para%group%mepos + 1
625 :
626 132 : IF (PRESENT(weights)) THEN
627 132 : CALL weights_signature(weights, has_weights, weight_sum, weight_sumsq)
628 : ELSE
629 : CALL weights_signature(has_weights=has_weights, weight_sum=weight_sum, &
630 0 : weight_sumsq=weight_sumsq)
631 : END IF
632 :
633 132 : max_local_features = nflat_local
634 132 : IF (my_atom_partition == skala_gpw_atom_partition_smooth) THEN
635 106 : max_local_features = nflat_local*natom
636 : END IF
637 0 : ALLOCATE (local_owner(max_local_features), &
638 0 : local_source_points(max_local_features), &
639 0 : local_static(nstatic_per_point*max_local_features), &
640 0 : local_feature_counts_tmp(nflat_local), feature_counts(nproc), &
641 : feature_displs(nproc), point_counts(nproc), point_displs(nproc), &
642 2244 : static_counts(nproc), static_displs(nproc), atom_coords_pbc(3, natom))
643 0 : ALLOCATE (cached_layout%feature_index(bo(1, 1):bo(2, 1), &
644 : bo(1, 2):bo(2, 2), &
645 660 : bo(1, 3):bo(2, 3)))
646 1478846 : cached_layout%feature_index = 0
647 132 : local_feature_counts_tmp = 0
648 424 : DO iatom = 1, natom
649 424 : atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.TRUE.)
650 : END DO
651 156 : nonperiodic = ALL(cell%perd == 0)
652 132 : IF (my_atom_partition == skala_gpw_atom_partition_smooth .AND. nonperiodic) THEN
653 32 : ALLOCATE (pair_distances(natom, natom))
654 8 : pair_distances = 0.0_dp
655 16 : DO iatom = 1, natom - 1
656 24 : DO jatom = iatom + 1, natom
657 : pair_distances(iatom, jatom) = &
658 32 : SQRT(SUM((atom_coords_pbc(:, iatom) - atom_coords_pbc(:, jatom))**2))
659 16 : pair_distances(jatom, iatom) = pair_distances(iatom, jatom)
660 : END DO
661 : END DO
662 : ELSE
663 124 : ALLOCATE (pair_distances(0, 0))
664 : END IF
665 :
666 132 : CALL timeset("skala_gpw_layout_local", phase_handle)
667 132 : nx = bo(2, 1) - bo(1, 1) + 1
668 132 : ny = bo(2, 2) - bo(1, 2) + 1
669 132 : IF (my_atom_partition == skala_gpw_atom_partition_hard) THEN
670 26 : nfeature_local = nflat_local
671 : !$OMP PARALLEL DO DEFAULT(NONE) COLLAPSE(3) &
672 : !$OMP SHARED(atom_coords_pbc, bo, cached_layout, cell, local_feature_counts_tmp, local_owner, &
673 : !$OMP local_source_points, local_static, nx, ny, pw_grid, weights) &
674 26 : !$OMP PRIVATE(atomic_base_weight, base_weight, grid_point, i, j, k, local_row, owner, static_base)
675 : DO k = bo(1, 3), bo(2, 3)
676 : DO j = bo(1, 2), bo(2, 2)
677 : DO i = bo(1, 1), bo(2, 1)
678 : local_row = i - bo(1, 1) + 1 + &
679 : nx*(j - bo(1, 2) + ny*(k - bo(1, 3)))
680 : grid_point = grid_coordinate(pw_grid, [i, j, k])
681 : atomic_base_weight = pw_grid%dvol
682 : base_weight = pw_grid%dvol
683 : IF (PRESENT(weights)) THEN
684 : IF (ASSOCIATED(weights)) base_weight = base_weight*weights%array(i, j, k)
685 : END IF
686 : cached_layout%feature_index(i, j, k) = local_row
687 : owner = nearest_atom(grid_point, atom_coords_pbc, cell)
688 : local_feature_counts_tmp(local_row) = 1
689 : local_owner(local_row) = owner
690 : local_source_points(local_row) = local_row
691 : static_base = nstatic_per_point*(local_row - 1)
692 : local_static(static_base + 1:static_base + 3) = grid_point
693 : local_static(static_base + 4) = base_weight
694 : local_static(static_base + 5) = atomic_base_weight
695 : END DO
696 : END DO
697 : END DO
698 : !$OMP END PARALLEL DO
699 : ELSE
700 : !$OMP PARALLEL DEFAULT(NONE) &
701 : !$OMP SHARED(atom_coords_pbc, bo, cached_layout, cell, local_feature_counts_tmp, local_owner, &
702 : !$OMP local_source_points, local_static, natom, nonperiodic, nx, ny, pair_distances, &
703 : !$OMP pw_grid, weights) &
704 : !$OMP PRIVATE(atomic_base_weight, base_weight, distances, feature_slot, grid_point, i, iatom, &
705 : !$OMP included_sum, j, k, local_feature, local_row, owner, partition_weight, &
706 106 : !$OMP partition_atom_coords, partition_weights, static_base)
707 : ! Keep point-local scratch outside the hot loop and private to each thread.
708 : ALLOCATE (distances(natom), partition_atom_coords(3, natom), partition_weights(natom))
709 : !$OMP DO COLLAPSE(3)
710 : DO k = bo(1, 3), bo(2, 3)
711 : DO j = bo(1, 2), bo(2, 2)
712 : DO i = bo(1, 1), bo(2, 1)
713 : local_row = i - bo(1, 1) + 1 + &
714 : nx*(j - bo(1, 2) + ny*(k - bo(1, 3)))
715 : grid_point = grid_coordinate(pw_grid, [i, j, k])
716 : atomic_base_weight = pw_grid%dvol
717 : base_weight = pw_grid%dvol
718 : IF (PRESENT(weights)) THEN
719 : IF (ASSOCIATED(weights)) base_weight = base_weight*weights%array(i, j, k)
720 : END IF
721 : cached_layout%feature_index(i, j, k) = local_row
722 : IF (nonperiodic) THEN
723 : CALL smooth_atom_partition(grid_point, atom_coords_pbc, cell, &
724 : partition_weights, partition_atom_coords, distances, &
725 : pair_distances)
726 : ELSE
727 : CALL smooth_atom_partition(grid_point, atom_coords_pbc, cell, &
728 : partition_weights, partition_atom_coords, distances)
729 : END IF
730 : included_sum = SUM(partition_weights, MASK=partition_weights > smooth_partition_eps)
731 : IF (included_sum <= 0.0_dp) THEN
732 : owner = nearest_atom(grid_point, atom_coords_pbc, cell)
733 : partition_weights = 0.0_dp
734 : partition_weights(owner) = 1.0_dp
735 : included_sum = 1.0_dp
736 : END IF
737 : local_feature = 0
738 : DO iatom = 1, natom
739 : IF (partition_weights(iatom) <= smooth_partition_eps) CYCLE
740 : partition_weight = partition_weights(iatom)/included_sum
741 : local_feature = local_feature + 1
742 : feature_slot = natom*(local_row - 1) + local_feature
743 : local_owner(feature_slot) = iatom
744 : local_source_points(feature_slot) = local_row
745 : static_base = nstatic_per_point*(feature_slot - 1)
746 : local_static(static_base + 1:static_base + 3) = grid_point
747 : local_static(static_base + 4) = base_weight*partition_weight
748 : local_static(static_base + 5) = atomic_base_weight* &
749 : smooth_partition_atomic_weight_scale( &
750 : partition_weight)
751 : END DO
752 : local_feature_counts_tmp(local_row) = local_feature
753 : END DO
754 : END DO
755 : END DO
756 : !$OMP END DO
757 : DEALLOCATE (distances, partition_atom_coords, partition_weights)
758 : !$OMP END PARALLEL
759 :
760 : ! Compact the conflict-free row slots while preserving the serial row/atom order.
761 106 : nfeature_local = 0
762 403129 : DO local_row = 1, nflat_local
763 403023 : CPASSERT(local_feature_counts_tmp(local_row) > 0)
764 1206969 : DO local_feature = 1, local_feature_counts_tmp(local_row)
765 803840 : nfeature_local = nfeature_local + 1
766 803840 : feature_slot = natom*(local_row - 1) + local_feature
767 803840 : IF (feature_slot == nfeature_local) CYCLE
768 677076 : local_owner(nfeature_local) = local_owner(feature_slot)
769 677076 : local_source_points(nfeature_local) = local_source_points(feature_slot)
770 : local_static(nstatic_per_point*(nfeature_local - 1) + 1: &
771 : nstatic_per_point*nfeature_local) = &
772 : local_static(nstatic_per_point*(feature_slot - 1) + 1: &
773 7977623 : nstatic_per_point*feature_slot)
774 : END DO
775 : END DO
776 : END IF
777 132 : CALL timestop(phase_handle)
778 :
779 : ! SKALA groups all grid points by atom. This ordering is static while the
780 : ! grid, cell, atom positions, and optional integration weights are unchanged.
781 132 : CALL timeset("skala_gpw_layout_gather", phase_handle)
782 132 : CALL pw_grid%para%group%allgather(nflat_local, point_counts)
783 132 : CALL counts_to_displs(point_counts, point_displs)
784 396 : npoint = SUM(point_counts)
785 132 : CALL pw_grid%para%group%allgather(nfeature_local, feature_counts)
786 132 : CALL counts_to_displs(feature_counts, feature_displs)
787 396 : DO pe = 1, nproc
788 264 : static_counts(pe) = nstatic_per_point*feature_counts(pe)
789 396 : static_displs(pe) = nstatic_per_point*feature_displs(pe)
790 : END DO
791 396 : nflat = SUM(feature_counts)
792 : ALLOCATE (global_owner(nflat), global_source_points(nflat), &
793 1056 : global_static(nstatic_per_point*nflat), local_source_global(nfeature_local))
794 1791159 : DO feature_local = 1, nfeature_local
795 1791159 : local_source_global(feature_local) = point_displs(pe_index) + local_source_points(feature_local)
796 : END DO
797 : CALL pw_grid%para%group%allgatherv(local_owner(1:nfeature_local), global_owner, feature_counts, &
798 132 : feature_displs)
799 : CALL pw_grid%para%group%allgatherv(local_source_global, global_source_points, feature_counts, &
800 132 : feature_displs)
801 : CALL pw_grid%para%group%allgatherv(local_static(1:nstatic_per_point*nfeature_local), &
802 : global_static, static_counts, &
803 132 : static_displs)
804 132 : CALL timestop(phase_handle)
805 :
806 0 : ALLOCATE (cached_layout%chunk_feature_counts(nproc), &
807 0 : cached_layout%chunk_feature_displs(nproc), &
808 0 : cached_layout%chunk_grad_counts(nproc), cached_layout%chunk_grad_displs(nproc), &
809 0 : cached_layout%feature_counts(nproc), cached_layout%feature_displs(nproc), &
810 0 : cached_layout%dynamic_counts(nproc), cached_layout%dynamic_displs(nproc), &
811 0 : cached_layout%route_point_recv_counts(nproc), &
812 0 : cached_layout%route_point_recv_displs(nproc), &
813 0 : cached_layout%route_point_send_counts(nproc), &
814 0 : cached_layout%route_point_send_displs(nproc), &
815 0 : cached_layout%feature_source_points(nflat), &
816 0 : cached_layout%global_to_feature(npoint), cached_layout%atomic_grid_sizes(natom), &
817 0 : cached_layout%local_feature_counts(nflat_local), &
818 0 : cached_layout%local_feature_offsets(nflat_local + 1), &
819 0 : cached_layout%local_feature_rows(nfeature_local), &
820 0 : cached_layout%local_feature_points(nfeature_local), &
821 0 : cached_layout%local_feature_indices(nfeature_local), atom_offset(natom + 1), &
822 : atom_position(natom), chunk_atom_begin(nproc), chunk_atom_end(nproc), &
823 4224 : cursor(nflat_local))
824 396 : cached_layout%feature_counts(:) = feature_counts
825 396 : cached_layout%feature_displs(:) = feature_displs
826 396 : cached_layout%dynamic_counts(:) = ndynamic_per_point*point_counts
827 396 : cached_layout%dynamic_displs(:) = ndynamic_per_point*point_displs
828 424 : cached_layout%atomic_grid_sizes = 0_int_8
829 2780552 : cached_layout%global_to_feature = 0
830 1390342 : cached_layout%local_feature_counts(:) = local_feature_counts_tmp
831 132 : cached_layout%local_feature_offsets(1) = 1
832 1390342 : DO local_row = 1, nflat_local
833 : cached_layout%local_feature_offsets(local_row + 1) = &
834 : cached_layout%local_feature_offsets(local_row) + &
835 1390342 : cached_layout%local_feature_counts(local_row)
836 : END DO
837 1390342 : cursor(:) = cached_layout%local_feature_offsets(1:nflat_local)
838 :
839 132 : CALL timeset("skala_gpw_layout_atom_sort", phase_handle)
840 3582186 : DO ipt = 1, nflat
841 : cached_layout%atomic_grid_sizes(global_owner(ipt)) = &
842 3582186 : cached_layout%atomic_grid_sizes(global_owner(ipt)) + 1_int_8
843 : END DO
844 132 : atom_offset(1) = 1
845 424 : DO iatom = 1, natom
846 424 : atom_offset(iatom + 1) = atom_offset(iatom) + INT(cached_layout%atomic_grid_sizes(iatom))
847 : END DO
848 424 : DO iatom = 1, natom
849 424 : atom_position(iatom) = atom_offset(iatom)
850 : END DO
851 424 : max_grid_size = MAXVAL(INT(cached_layout%atomic_grid_sizes))
852 : CALL build_atom_chunks(cached_layout%atomic_grid_sizes, atom_offset, nproc, &
853 : chunk_atom_begin, chunk_atom_end, &
854 : cached_layout%chunk_feature_counts, &
855 132 : cached_layout%chunk_feature_displs)
856 396 : cached_layout%chunk_grad_counts(:) = ngrad_per_point*cached_layout%chunk_feature_counts
857 396 : cached_layout%chunk_grad_displs(:) = ngrad_per_point*cached_layout%chunk_feature_displs
858 132 : cached_layout%chunk_atom_begin = chunk_atom_begin(pe_index)
859 132 : cached_layout%chunk_atom_end = chunk_atom_end(pe_index)
860 132 : cached_layout%chunk_feature_begin = cached_layout%chunk_feature_displs(pe_index) + 1
861 132 : cached_layout%chunk_feature_count = cached_layout%chunk_feature_counts(pe_index)
862 : cached_layout%chunk_natom = cached_layout%chunk_atom_end - &
863 132 : cached_layout%chunk_atom_begin + 1
864 :
865 0 : ALLOCATE (cached_layout%grid_coords(3, nflat), cached_layout%grid_weights(nflat), &
866 0 : cached_layout%atomic_grid_weights(nflat), &
867 0 : cached_layout%coarse_0_atomic_coords(3, natom), &
868 0 : cached_layout%atomic_grid_size_bound_shape(0, max_grid_size), &
869 1188 : cached_layout%atom_coords(3, natom))
870 14328348 : cached_layout%grid_coords = 0.0_dp
871 3582186 : cached_layout%grid_weights = 0.0_dp
872 3582186 : cached_layout%atomic_grid_weights = 0.0_dp
873 1584528 : cached_layout%atomic_grid_size_bound_shape = 0_int_8
874 :
875 424 : DO iatom = 1, natom
876 2044 : cached_layout%atom_coords(:, iatom) = particle_set(iatom)%r
877 1300 : cached_layout%coarse_0_atomic_coords(:, iatom) = atom_coords_pbc(:, iatom)
878 : END DO
879 :
880 3582186 : DO ipt = 1, nflat
881 3582054 : owner = global_owner(ipt)
882 3582054 : row = atom_position(owner)
883 3582054 : atom_position(owner) = atom_position(owner) + 1
884 3582054 : source_global = global_source_points(ipt)
885 3582054 : cached_layout%feature_source_points(row) = source_global
886 3582054 : IF (cached_layout%global_to_feature(source_global) == 0) THEN
887 2780420 : cached_layout%global_to_feature(source_global) = row
888 : END IF
889 3582054 : static_base = nstatic_per_point*(ipt - 1)
890 14328216 : cached_layout%grid_coords(:, row) = global_static(static_base + 1:static_base + 3)
891 3582054 : cached_layout%grid_weights(row) = global_static(static_base + 4)
892 3582054 : cached_layout%atomic_grid_weights(row) = global_static(static_base + 5)
893 3582054 : source_local = source_global - point_displs(pe_index)
894 3582186 : IF (source_local >= 1 .AND. source_local <= nflat_local) THEN
895 1791027 : feature_local = cursor(source_local)
896 1791027 : cursor(source_local) = cursor(source_local) + 1
897 1791027 : cached_layout%local_feature_rows(feature_local) = row
898 1791027 : cached_layout%local_feature_points(feature_local) = source_local
899 : END IF
900 : END DO
901 :
902 2780552 : CPASSERT(ALL(cached_layout%global_to_feature > 0))
903 1791159 : CPASSERT(ALL(cached_layout%local_feature_rows > 0))
904 1791159 : CPASSERT(ALL(cached_layout%local_feature_points > 0))
905 3212 : DO k = bo(1, 3), bo(2, 3)
906 88636 : DO j = bo(1, 2), bo(2, 2)
907 1478714 : DO i = bo(1, 1), bo(2, 1)
908 1390210 : local_row = cached_layout%feature_index(i, j, k)
909 : cached_layout%feature_index(i, j, k) = &
910 1475634 : cached_layout%local_feature_rows(cached_layout%local_feature_offsets(local_row))
911 : END DO
912 : END DO
913 : END DO
914 1791159 : DO feature_local = 1, nfeature_local
915 : cached_layout%local_feature_indices(feature_local) = &
916 1791159 : INT(cached_layout%local_feature_rows(feature_local) - 1, KIND=int_8)
917 : END DO
918 132 : CALL timestop(phase_handle)
919 132 : CALL timeset("skala_gpw_layout_chunk_routes", phase_handle)
920 : CALL build_atom_chunk_routes(cached_layout, cached_layout%local_feature_rows, &
921 132 : pw_grid%para%group)
922 132 : CALL build_atom_chunk_layout(cached_layout)
923 132 : CALL timestop(phase_handle)
924 :
925 132 : cached_layout%natom = natom
926 132 : cached_layout%nflat = nflat
927 132 : cached_layout%nflat_local = nflat_local
928 132 : cached_layout%npoint = npoint
929 132 : cached_layout%nproc = nproc
930 132 : cached_layout%atom_partition = my_atom_partition
931 1320 : cached_layout%bo = bo
932 2508 : cached_layout%bounds = pw_grid%bounds
933 924 : cached_layout%npts = pw_grid%npts
934 132 : cached_layout%dvol = pw_grid%dvol
935 3300 : cached_layout%dh = pw_grid%dh
936 3300 : cached_layout%cell_hmat = cell%hmat
937 132 : cached_layout%weight_sum = weight_sum
938 132 : cached_layout%weight_sumsq = weight_sumsq
939 132 : cached_layout%has_weights = has_weights
940 132 : CALL timeset("skala_gpw_layout_tensors", phase_handle)
941 132 : IF (needs_full_static_tensors) CALL build_full_static_layout_tensors(cached_layout)
942 132 : IF (cached_layout%chunk_feature_count > 0) THEN
943 132 : CALL build_chunk_static_layout_tensors(cached_layout)
944 : END IF
945 132 : CALL timestop(phase_handle)
946 132 : cached_layout%active = .TRUE.
947 :
948 0 : DEALLOCATE (atom_coords_pbc, atom_offset, atom_position, &
949 0 : chunk_atom_begin, chunk_atom_end, cursor, feature_counts, feature_displs, &
950 0 : global_owner, global_source_points, global_static, local_feature_counts_tmp, &
951 0 : local_owner, local_source_global, local_source_points, &
952 0 : local_static, pair_distances, point_counts, point_displs, static_counts, &
953 132 : static_displs)
954 :
955 792 : END SUBROUTINE rebuild_layout_cache
956 :
957 : ! **************************************************************************************************
958 : !> \brief Build cached full-layout Torch tensors for static SKALA inputs.
959 : !> \param cache ...
960 : ! **************************************************************************************************
961 64 : SUBROUTINE build_full_static_layout_tensors(cache)
962 : TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
963 :
964 64 : CPASSERT(.NOT. cache%static_tensors_active)
965 :
966 64 : CALL torch_tensor_from_array(cache%grid_coords_t, cache%grid_coords)
967 64 : CALL torch_tensor_to_device_leaf(cache%grid_coords_t, .FALSE.)
968 64 : CALL torch_tensor_from_array(cache%grid_weights_t, cache%grid_weights)
969 64 : CALL torch_tensor_to_device_leaf(cache%grid_weights_t, .FALSE.)
970 64 : CALL torch_tensor_from_array(cache%atomic_grid_weights_t, cache%atomic_grid_weights)
971 64 : CALL torch_tensor_to_device_leaf(cache%atomic_grid_weights_t, .FALSE.)
972 64 : CALL torch_tensor_from_array(cache%atomic_grid_sizes_t, cache%atomic_grid_sizes)
973 64 : CALL torch_tensor_to_device_leaf(cache%atomic_grid_sizes_t, .FALSE.)
974 64 : CALL torch_tensor_from_array(cache%coarse_0_atomic_coords_t, cache%coarse_0_atomic_coords)
975 64 : CALL torch_tensor_to_device_leaf(cache%coarse_0_atomic_coords_t, .FALSE.)
976 : CALL torch_tensor_from_array(cache%atomic_grid_size_bound_shape_t, &
977 64 : cache%atomic_grid_size_bound_shape)
978 64 : CALL torch_tensor_to_device_leaf(cache%atomic_grid_size_bound_shape_t, .FALSE.)
979 64 : CALL torch_tensor_from_array(cache%local_feature_indices_t, cache%local_feature_indices)
980 64 : CALL torch_tensor_to_device_leaf(cache%local_feature_indices_t, .FALSE.)
981 :
982 64 : CALL torch_dict_create(cache%static_inputs)
983 64 : CALL torch_dict_insert(cache%static_inputs, "grid_coords", cache%grid_coords_t)
984 64 : CALL torch_dict_insert(cache%static_inputs, "grid_weights", cache%grid_weights_t)
985 : CALL torch_dict_insert(cache%static_inputs, "atomic_grid_weights", &
986 64 : cache%atomic_grid_weights_t)
987 : CALL torch_dict_insert(cache%static_inputs, "atomic_grid_sizes", &
988 64 : cache%atomic_grid_sizes_t)
989 : CALL torch_dict_insert(cache%static_inputs, "atomic_grid_size_bound_shape", &
990 64 : cache%atomic_grid_size_bound_shape_t)
991 64 : cache%static_tensors_active = .TRUE.
992 :
993 64 : END SUBROUTINE build_full_static_layout_tensors
994 :
995 : ! **************************************************************************************************
996 : !> \brief Build cached atom-chunk Torch tensors for static SKALA inputs.
997 : !> \param cache ...
998 : ! **************************************************************************************************
999 132 : SUBROUTINE build_chunk_static_layout_tensors(cache)
1000 : TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
1001 :
1002 132 : CPASSERT(.NOT. cache%chunk_static_tensors_active)
1003 132 : CALL torch_tensor_from_array(cache%chunk_grid_coords_t, cache%chunk_grid_coords)
1004 132 : CALL torch_tensor_to_device_leaf(cache%chunk_grid_coords_t, .FALSE.)
1005 132 : CALL torch_tensor_from_array(cache%chunk_grid_weights_t, cache%chunk_grid_weights)
1006 132 : CALL torch_tensor_to_device_leaf(cache%chunk_grid_weights_t, .FALSE.)
1007 : CALL torch_tensor_from_array(cache%chunk_atomic_grid_weights_t, &
1008 132 : cache%chunk_atomic_grid_weights)
1009 132 : CALL torch_tensor_to_device_leaf(cache%chunk_atomic_grid_weights_t, .FALSE.)
1010 : CALL torch_tensor_from_array(cache%chunk_atomic_grid_sizes_t, &
1011 132 : cache%chunk_atomic_grid_sizes)
1012 132 : CALL torch_tensor_to_device_leaf(cache%chunk_atomic_grid_sizes_t, .FALSE.)
1013 : CALL torch_tensor_from_array(cache%chunk_coarse_0_atomic_coords_t, &
1014 132 : cache%chunk_coarse_0_atomic_coords)
1015 132 : CALL torch_tensor_to_device_leaf(cache%chunk_coarse_0_atomic_coords_t, .FALSE.)
1016 : CALL torch_tensor_from_array(cache%chunk_atomic_grid_size_bound_shape_t, &
1017 132 : cache%chunk_atomic_grid_size_bound_shape)
1018 132 : CALL torch_tensor_to_device_leaf(cache%chunk_atomic_grid_size_bound_shape_t, .FALSE.)
1019 132 : CALL torch_tensor_from_array(cache%chunk_feature_indices_t, cache%chunk_feature_indices)
1020 132 : CALL torch_tensor_to_device_leaf(cache%chunk_feature_indices_t, .FALSE.)
1021 :
1022 132 : CALL torch_dict_create(cache%chunk_static_inputs)
1023 : CALL torch_dict_insert(cache%chunk_static_inputs, "grid_coords", &
1024 132 : cache%chunk_grid_coords_t)
1025 : CALL torch_dict_insert(cache%chunk_static_inputs, "grid_weights", &
1026 132 : cache%chunk_grid_weights_t)
1027 : CALL torch_dict_insert(cache%chunk_static_inputs, "atomic_grid_weights", &
1028 132 : cache%chunk_atomic_grid_weights_t)
1029 : CALL torch_dict_insert(cache%chunk_static_inputs, "atomic_grid_sizes", &
1030 132 : cache%chunk_atomic_grid_sizes_t)
1031 : CALL torch_dict_insert(cache%chunk_static_inputs, "atomic_grid_size_bound_shape", &
1032 132 : cache%chunk_atomic_grid_size_bound_shape_t)
1033 132 : cache%chunk_static_tensors_active = .TRUE.
1034 :
1035 132 : END SUBROUTINE build_chunk_static_layout_tensors
1036 :
1037 : ! **************************************************************************************************
1038 : !> \brief Copy static cached layout arrays into a feature bundle.
1039 : !> \param features ...
1040 : !> \param use_atom_chunk_routing ...
1041 : !> \param needs_coordinate_array ...
1042 : !> \param needs_grid_coordinate_array ...
1043 : ! **************************************************************************************************
1044 322 : SUBROUTINE copy_cached_layout(features, use_atom_chunk_routing, needs_coordinate_array, &
1045 : needs_grid_coordinate_array)
1046 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1047 : LOGICAL, INTENT(IN) :: use_atom_chunk_routing, &
1048 : needs_coordinate_array, &
1049 : needs_grid_coordinate_array
1050 :
1051 322 : CPASSERT(cached_layout%active)
1052 :
1053 322 : IF (.NOT. use_atom_chunk_routing) THEN
1054 0 : ALLOCATE (features%feature_index(LBOUND(cached_layout%feature_index, 1): &
1055 : UBOUND(cached_layout%feature_index, 1), &
1056 : LBOUND(cached_layout%feature_index, 2): &
1057 : UBOUND(cached_layout%feature_index, 2), &
1058 : LBOUND(cached_layout%feature_index, 3): &
1059 320 : UBOUND(cached_layout%feature_index, 3)))
1060 192 : ALLOCATE (features%grid_weights(cached_layout%nflat))
1061 0 : ALLOCATE (features%local_feature_counts(cached_layout%nflat_local), &
1062 0 : features%local_feature_offsets(cached_layout%nflat_local + 1), &
1063 448 : features%local_feature_rows(SIZE(cached_layout%local_feature_rows)))
1064 :
1065 352039 : features%feature_index(:, :, :) = cached_layout%feature_index
1066 1283596 : features%grid_weights(:) = cached_layout%grid_weights
1067 321763 : features%local_feature_counts(:) = cached_layout%local_feature_counts
1068 321827 : features%local_feature_offsets(:) = cached_layout%local_feature_offsets
1069 641830 : features%local_feature_rows(:) = cached_layout%local_feature_rows
1070 : END IF
1071 322 : features%nflat = cached_layout%nflat
1072 322 : features%nflat_local = cached_layout%nflat_local
1073 322 : features%chunk_feature_count = cached_layout%chunk_feature_count
1074 322 : features%atom_partition = cached_layout%atom_partition
1075 966 : ALLOCATE (features%atomic_grid_sizes(cached_layout%natom))
1076 994 : features%atomic_grid_sizes(:) = cached_layout%atomic_grid_sizes
1077 322 : IF (needs_grid_coordinate_array) THEN
1078 192 : ALLOCATE (features%grid_coords(3, cached_layout%nflat))
1079 192 : ALLOCATE (features%atomic_grid_weights(cached_layout%nflat))
1080 5134192 : features%grid_coords(:, :) = cached_layout%grid_coords
1081 1283596 : features%atomic_grid_weights(:) = cached_layout%atomic_grid_weights
1082 : END IF
1083 0 : ALLOCATE (features%chunk_grad_counts(cached_layout%nproc), &
1084 1288 : features%chunk_grad_displs(cached_layout%nproc))
1085 966 : features%chunk_grad_counts(:) = cached_layout%chunk_grad_counts
1086 966 : features%chunk_grad_displs(:) = cached_layout%chunk_grad_displs
1087 322 : IF (use_atom_chunk_routing) THEN
1088 258 : features%route_point_recv_counts => cached_layout%route_point_recv_counts
1089 258 : features%route_point_recv_displs => cached_layout%route_point_recv_displs
1090 258 : features%route_point_send_counts => cached_layout%route_point_send_counts
1091 258 : features%route_point_send_displs => cached_layout%route_point_send_displs
1092 258 : features%route_chunk_offsets => cached_layout%route_chunk_offsets
1093 258 : features%route_chunk_rows => cached_layout%route_chunk_rows
1094 258 : features%route_row_offsets => cached_layout%route_row_offsets
1095 258 : features%route_row_positions => cached_layout%route_row_positions
1096 258 : features%route_send_local_rows => cached_layout%route_send_local_rows
1097 : END IF
1098 322 : IF (needs_coordinate_array) THEN
1099 192 : ALLOCATE (features%coarse_0_atomic_coords(3, cached_layout%natom))
1100 576 : features%coarse_0_atomic_coords(:, :) = cached_layout%coarse_0_atomic_coords
1101 : END IF
1102 :
1103 322 : END SUBROUTINE copy_cached_layout
1104 :
1105 : ! **************************************************************************************************
1106 : !> \brief Split the atom-ordered feature rows into contiguous atom chunks.
1107 : !> \param atomic_grid_sizes ...
1108 : !> \param atom_offset ...
1109 : !> \param nproc ...
1110 : !> \param chunk_atom_begin ...
1111 : !> \param chunk_atom_end ...
1112 : !> \param chunk_feature_counts ...
1113 : !> \param chunk_feature_displs ...
1114 : ! **************************************************************************************************
1115 132 : SUBROUTINE build_atom_chunks(atomic_grid_sizes, atom_offset, nproc, chunk_atom_begin, &
1116 132 : chunk_atom_end, chunk_feature_counts, chunk_feature_displs)
1117 : INTEGER(KIND=int_8), DIMENSION(:), INTENT(IN) :: atomic_grid_sizes
1118 : INTEGER, DIMENSION(:), INTENT(IN) :: atom_offset
1119 : INTEGER, INTENT(IN) :: nproc
1120 : INTEGER, DIMENSION(:), INTENT(OUT) :: chunk_atom_begin, chunk_atom_end, &
1121 : chunk_feature_counts, &
1122 : chunk_feature_displs
1123 :
1124 : INTEGER :: best_limit, count, displ, end_atom, lower_limit, max_end_atom, midpoint, natom, &
1125 : next_atom, next_count, pe, ranks_left, target_chunks, total_count, upper_limit
1126 :
1127 132 : natom = SIZE(atomic_grid_sizes)
1128 396 : chunk_atom_begin = natom + 1
1129 396 : chunk_atom_end = natom
1130 396 : chunk_feature_counts = 0
1131 396 : chunk_feature_displs = 0
1132 132 : IF (natom == 0) RETURN
1133 :
1134 132 : target_chunks = MIN(nproc, natom)
1135 132 : total_count = atom_offset(natom + 1) - 1
1136 424 : lower_limit = MAXVAL(INT(atomic_grid_sizes))
1137 132 : lower_limit = MAX(lower_limit, (total_count + target_chunks - 1)/target_chunks)
1138 132 : upper_limit = total_count
1139 132 : best_limit = upper_limit
1140 1778 : DO WHILE (lower_limit <= upper_limit)
1141 1646 : midpoint = (lower_limit + upper_limit)/2
1142 1778 : IF (atom_chunks_fit_limit(atomic_grid_sizes, midpoint, target_chunks)) THEN
1143 1526 : best_limit = midpoint
1144 1526 : upper_limit = midpoint - 1
1145 : ELSE
1146 120 : lower_limit = midpoint + 1
1147 : END IF
1148 : END DO
1149 :
1150 : displ = 0
1151 : next_atom = 1
1152 396 : DO pe = 1, nproc
1153 264 : chunk_feature_displs(pe) = displ
1154 264 : IF (pe > target_chunks .OR. next_atom > natom) CYCLE
1155 :
1156 264 : ranks_left = target_chunks - pe + 1
1157 264 : chunk_atom_begin(pe) = next_atom
1158 264 : max_end_atom = natom - ranks_left + 1
1159 264 : end_atom = next_atom
1160 264 : count = INT(atomic_grid_sizes(end_atom))
1161 292 : DO WHILE (end_atom < max_end_atom)
1162 38 : next_count = count + INT(atomic_grid_sizes(end_atom + 1))
1163 38 : IF (next_count > best_limit) EXIT
1164 : end_atom = end_atom + 1
1165 264 : count = next_count
1166 : END DO
1167 :
1168 264 : chunk_atom_end(pe) = end_atom
1169 264 : chunk_feature_counts(pe) = atom_offset(end_atom + 1) - atom_offset(next_atom)
1170 264 : displ = displ + chunk_feature_counts(pe)
1171 396 : next_atom = end_atom + 1
1172 : END DO
1173 :
1174 132 : CPASSERT(displ == atom_offset(natom + 1) - 1)
1175 :
1176 : END SUBROUTINE build_atom_chunks
1177 :
1178 : ! **************************************************************************************************
1179 : !> \brief Check if contiguous atom chunks can stay below a feature-count limit.
1180 : !> \param atomic_grid_sizes ...
1181 : !> \param limit ...
1182 : !> \param nchunks ...
1183 : !> \return ...
1184 : ! **************************************************************************************************
1185 1646 : FUNCTION atom_chunks_fit_limit(atomic_grid_sizes, limit, nchunks) RESULT(fits)
1186 : INTEGER(KIND=int_8), DIMENSION(:), INTENT(IN) :: atomic_grid_sizes
1187 : INTEGER, INTENT(IN) :: limit, nchunks
1188 : LOGICAL :: fits
1189 :
1190 : INTEGER :: atom_count, chunk_count, iatom, &
1191 : used_chunks
1192 :
1193 1646 : fits = .FALSE.
1194 1646 : IF (SIZE(atomic_grid_sizes) == 0) THEN
1195 1646 : fits = .TRUE.
1196 : RETURN
1197 : END IF
1198 :
1199 5358 : used_chunks = 1
1200 5358 : chunk_count = 0
1201 5358 : DO iatom = 1, SIZE(atomic_grid_sizes)
1202 3712 : atom_count = INT(atomic_grid_sizes(iatom))
1203 3712 : IF (atom_count > limit) RETURN
1204 5358 : IF (chunk_count + atom_count > limit) THEN
1205 1766 : used_chunks = used_chunks + 1
1206 1766 : chunk_count = atom_count
1207 : ELSE
1208 : chunk_count = chunk_count + atom_count
1209 : END IF
1210 : END DO
1211 1646 : fits = used_chunks <= nchunks
1212 :
1213 1646 : END FUNCTION atom_chunks_fit_limit
1214 :
1215 : ! **************************************************************************************************
1216 : !> \brief Return the MPI rank owning an atom-ordered feature row.
1217 : !> \param row ...
1218 : !> \param counts ...
1219 : !> \param displs ...
1220 : !> \return ...
1221 : ! **************************************************************************************************
1222 1791027 : FUNCTION feature_row_chunk_owner(row, counts, displs) RESULT(owner)
1223 : INTEGER, INTENT(IN) :: row
1224 : INTEGER, DIMENSION(:), INTENT(IN) :: counts, displs
1225 : INTEGER :: owner
1226 :
1227 : INTEGER :: pe
1228 :
1229 1791027 : owner = 0
1230 2645916 : DO pe = 1, SIZE(counts)
1231 2645916 : IF (row > displs(pe) .AND. row <= displs(pe) + counts(pe)) THEN
1232 1791027 : owner = pe
1233 : RETURN
1234 : END IF
1235 : END DO
1236 :
1237 : END FUNCTION feature_row_chunk_owner
1238 :
1239 : ! **************************************************************************************************
1240 : !> \brief Build zero-based displacement arrays from per-rank counts.
1241 : !> \param counts ...
1242 : !> \param displs ...
1243 : ! **************************************************************************************************
1244 792 : SUBROUTINE counts_to_displs(counts, displs)
1245 : INTEGER, DIMENSION(:), INTENT(IN) :: counts
1246 : INTEGER, DIMENSION(:), INTENT(OUT) :: displs
1247 :
1248 : INTEGER :: pe
1249 :
1250 792 : displs(1) = 0
1251 1584 : DO pe = 2, SIZE(counts)
1252 1584 : displs(pe) = displs(pe - 1) + counts(pe - 1)
1253 : END DO
1254 :
1255 792 : END SUBROUTINE counts_to_displs
1256 :
1257 : ! **************************************************************************************************
1258 : !> \brief Precompute all-to-all routing between local grid rows and atom chunks.
1259 : !> \param cache ...
1260 : !> \param local_to_global ...
1261 : !> \param group ...
1262 : ! **************************************************************************************************
1263 132 : SUBROUTINE build_atom_chunk_routes(cache, local_to_global, group)
1264 : TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
1265 : INTEGER, DIMENSION(:), INTENT(IN) :: local_to_global
1266 :
1267 : CLASS(mp_comm_type), INTENT(IN) :: group
1268 :
1269 : INTEGER :: chunk_row, dest, feature_begin, feature_end, &
1270 : feature_pos, local_feature, local_row, meta_pos, &
1271 : nflat_local, nproc, nrecv, nsend, pe, point_pos, &
1272 : route_index, row, row_position
1273 132 : INTEGER, ALLOCATABLE, DIMENSION(:) :: chunk_cursor, chunk_route_counts, &
1274 132 : feature_cursor, feature_recv_counts, &
1275 : feature_recv_displs, feature_send_counts, &
1276 132 : feature_send_displs, meta_recv_counts, &
1277 132 : meta_recv_displs, meta_send_counts, &
1278 132 : meta_send_displs, recv_meta, route_cursor, &
1279 132 : route_feature_dest, route_last_row, &
1280 132 : route_last_position, send_meta
1281 :
1282 132 : nflat_local = SIZE(cache%local_feature_offsets) - 1
1283 132 : nproc = SIZE(cache%route_point_send_counts)
1284 : ALLOCATE (route_feature_dest(SIZE(local_to_global)), &
1285 : feature_send_counts(nproc), feature_send_displs(nproc), &
1286 : feature_recv_counts(nproc), feature_recv_displs(nproc), &
1287 1320 : route_last_row(nproc), route_last_position(nproc))
1288 396 : cache%route_point_send_counts = 0
1289 132 : feature_send_counts = 0
1290 132 : route_feature_dest = 0
1291 132 : route_last_row = 0
1292 : ! Route each local grid row only once to each atom-chunk owner.
1293 1390342 : DO local_row = 1, nflat_local
1294 1390210 : feature_begin = cache%local_feature_offsets(local_row)
1295 1390210 : feature_end = cache%local_feature_offsets(local_row + 1) - 1
1296 3181369 : DO local_feature = feature_begin, feature_end
1297 : dest = feature_row_chunk_owner(local_to_global(local_feature), &
1298 : cache%chunk_feature_counts, &
1299 1791027 : cache%chunk_feature_displs)
1300 1791027 : CPASSERT(dest > 0)
1301 1791027 : route_feature_dest(local_feature) = dest
1302 1791027 : feature_send_counts(dest) = feature_send_counts(dest) + 1
1303 3181237 : IF (route_last_row(dest) /= local_row) THEN
1304 1791027 : cache%route_point_send_counts(dest) = cache%route_point_send_counts(dest) + 1
1305 1791027 : route_last_row(dest) = local_row
1306 : END IF
1307 : END DO
1308 : END DO
1309 132 : CALL counts_to_displs(cache%route_point_send_counts, cache%route_point_send_displs)
1310 132 : CALL counts_to_displs(feature_send_counts, feature_send_displs)
1311 132 : CALL group%alltoall(cache%route_point_send_counts, cache%route_point_recv_counts, 1)
1312 132 : CALL counts_to_displs(cache%route_point_recv_counts, cache%route_point_recv_displs)
1313 132 : CALL group%alltoall(feature_send_counts, feature_recv_counts, 1)
1314 132 : CALL counts_to_displs(feature_recv_counts, feature_recv_displs)
1315 :
1316 396 : nsend = SUM(cache%route_point_send_counts)
1317 396 : nrecv = SUM(cache%route_point_recv_counts)
1318 0 : ALLOCATE (cache%route_send_local_rows(nsend), &
1319 0 : cache%route_row_offsets(nflat_local + 1), &
1320 0 : cache%route_row_positions(nsend), &
1321 : route_cursor(nproc), feature_cursor(nproc), &
1322 : send_meta(2*SIZE(local_to_global)), &
1323 : recv_meta(2*cache%chunk_feature_count), &
1324 : meta_send_counts(nproc), meta_send_displs(nproc), &
1325 2244 : meta_recv_counts(nproc), meta_recv_displs(nproc))
1326 1791159 : cache%route_send_local_rows = 0
1327 1791159 : cache%route_row_positions = 0
1328 132 : send_meta = 0
1329 396 : route_cursor(:) = cache%route_point_send_displs + 1
1330 396 : feature_cursor(:) = feature_send_displs + 1
1331 132 : route_last_row = 0
1332 132 : route_last_position = 0
1333 132 : cache%route_row_offsets(1) = 1
1334 132 : row_position = 1
1335 1390342 : DO local_row = 1, nflat_local
1336 1390210 : feature_begin = cache%local_feature_offsets(local_row)
1337 1390210 : feature_end = cache%local_feature_offsets(local_row + 1) - 1
1338 3181237 : DO local_feature = feature_begin, feature_end
1339 1791027 : dest = route_feature_dest(local_feature)
1340 1791027 : IF (route_last_row(dest) /= local_row) THEN
1341 1791027 : point_pos = route_cursor(dest)
1342 1791027 : route_cursor(dest) = route_cursor(dest) + 1
1343 1791027 : route_last_row(dest) = local_row
1344 1791027 : route_last_position(dest) = point_pos
1345 1791027 : cache%route_send_local_rows(point_pos) = local_row
1346 1791027 : cache%route_row_positions(row_position) = point_pos
1347 1791027 : row_position = row_position + 1
1348 : ELSE
1349 0 : point_pos = route_last_position(dest)
1350 : END IF
1351 1791027 : feature_pos = feature_cursor(dest)
1352 1791027 : feature_cursor(dest) = feature_cursor(dest) + 1
1353 1791027 : send_meta(2*feature_pos - 1) = local_to_global(local_feature)
1354 3181237 : send_meta(2*feature_pos) = point_pos - cache%route_point_send_displs(dest)
1355 : END DO
1356 1390342 : cache%route_row_offsets(local_row + 1) = row_position
1357 : END DO
1358 132 : CPASSERT(row_position == nsend + 1)
1359 :
1360 : ! Tell each chunk owner which feature rows share a compressed route point.
1361 396 : meta_send_counts(:) = 2*feature_send_counts
1362 396 : meta_send_displs(:) = 2*feature_send_displs
1363 396 : meta_recv_counts(:) = 2*feature_recv_counts
1364 396 : meta_recv_displs(:) = 2*feature_recv_displs
1365 : CALL group%alltoall(send_meta, meta_send_counts, meta_send_displs, recv_meta, &
1366 132 : meta_recv_counts, meta_recv_displs)
1367 :
1368 : ! Store the route-to-feature expansion in CSR form for gradient reduction.
1369 0 : ALLOCATE (cache%chunk_return_positions(cache%chunk_feature_count), &
1370 0 : cache%route_chunk_offsets(nrecv + 1), &
1371 0 : cache%route_chunk_rows(cache%chunk_feature_count), &
1372 1188 : chunk_route_counts(nrecv), chunk_cursor(nrecv))
1373 1791159 : cache%chunk_return_positions = 0
1374 1791159 : cache%route_chunk_rows = 0
1375 132 : chunk_route_counts = 0
1376 396 : DO pe = 1, nproc
1377 1791291 : DO feature_pos = feature_recv_displs(pe) + 1, &
1378 396 : feature_recv_displs(pe) + feature_recv_counts(pe)
1379 1791027 : meta_pos = 2*feature_pos
1380 1791027 : row = recv_meta(meta_pos - 1)
1381 1791027 : route_index = recv_meta(meta_pos)
1382 1791027 : point_pos = cache%route_point_recv_displs(pe) + route_index
1383 1791027 : CPASSERT(point_pos >= 1 .AND. point_pos <= nrecv)
1384 1791027 : chunk_row = row - cache%chunk_feature_begin + 1
1385 1791027 : CPASSERT(chunk_row >= 1 .AND. chunk_row <= cache%chunk_feature_count)
1386 1791027 : cache%chunk_return_positions(chunk_row) = point_pos
1387 1791291 : chunk_route_counts(point_pos) = chunk_route_counts(point_pos) + 1
1388 : END DO
1389 : END DO
1390 132 : cache%route_chunk_offsets(1) = 1
1391 1791159 : DO point_pos = 1, nrecv
1392 : cache%route_chunk_offsets(point_pos + 1) = &
1393 1791159 : cache%route_chunk_offsets(point_pos) + chunk_route_counts(point_pos)
1394 : END DO
1395 1791159 : chunk_cursor(:) = cache%route_chunk_offsets(1:nrecv)
1396 396 : DO pe = 1, nproc
1397 1791291 : DO feature_pos = feature_recv_displs(pe) + 1, &
1398 396 : feature_recv_displs(pe) + feature_recv_counts(pe)
1399 1791027 : meta_pos = 2*feature_pos
1400 1791027 : row = recv_meta(meta_pos - 1)
1401 1791027 : chunk_row = row - cache%chunk_feature_begin + 1
1402 1791027 : point_pos = cache%chunk_return_positions(chunk_row)
1403 1791027 : cache%route_chunk_rows(chunk_cursor(point_pos)) = chunk_row
1404 1791291 : chunk_cursor(point_pos) = chunk_cursor(point_pos) + 1
1405 : END DO
1406 : END DO
1407 :
1408 396 : CPASSERT(SUM(feature_send_counts) == SIZE(local_to_global))
1409 396 : CPASSERT(SUM(feature_recv_counts) == cache%chunk_feature_count)
1410 396 : CPASSERT(SUM(cache%route_point_send_counts) == nsend)
1411 396 : CPASSERT(SUM(cache%route_point_recv_counts) == nrecv)
1412 1791159 : CPASSERT(ALL(cache%chunk_return_positions > 0))
1413 1791159 : CPASSERT(ALL(cache%route_chunk_rows > 0))
1414 1791159 : CPASSERT(ALL(cache%route_row_positions > 0))
1415 1791159 : CPASSERT(ALL(cache%route_send_local_rows > 0))
1416 :
1417 0 : DEALLOCATE (chunk_cursor, chunk_route_counts, feature_cursor, feature_recv_counts, &
1418 0 : feature_recv_displs, feature_send_counts, feature_send_displs, &
1419 0 : meta_recv_counts, meta_recv_displs, meta_send_counts, meta_send_displs, &
1420 0 : recv_meta, route_cursor, route_feature_dest, route_last_position, route_last_row, &
1421 132 : send_meta)
1422 :
1423 132 : END SUBROUTINE build_atom_chunk_routes
1424 :
1425 : ! **************************************************************************************************
1426 : !> \brief Materialize the current rank's atom chunk static layout.
1427 : !> \param cache ...
1428 : ! **************************************************************************************************
1429 132 : SUBROUTINE build_atom_chunk_layout(cache)
1430 : TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
1431 :
1432 : INTEGER :: irow, max_grid_size, row_begin, row_end
1433 :
1434 132 : IF (cache%chunk_feature_count <= 0 .OR. cache%chunk_natom <= 0) RETURN
1435 :
1436 132 : row_begin = cache%chunk_feature_begin
1437 132 : row_end = row_begin + cache%chunk_feature_count - 1
1438 0 : ALLOCATE (cache%chunk_grid_coords(3, cache%chunk_feature_count), &
1439 0 : cache%chunk_grid_weights(cache%chunk_feature_count), &
1440 0 : cache%chunk_atomic_grid_weights(cache%chunk_feature_count), &
1441 0 : cache%chunk_atomic_grid_sizes(cache%chunk_natom), &
1442 0 : cache%chunk_coarse_0_atomic_coords(3, cache%chunk_natom), &
1443 1452 : cache%chunk_feature_indices(cache%chunk_feature_count))
1444 7164240 : cache%chunk_grid_coords(:, :) = cache%grid_coords(:, row_begin:row_end)
1445 1791159 : cache%chunk_grid_weights(:) = cache%grid_weights(row_begin:row_end)
1446 1791159 : cache%chunk_atomic_grid_weights(:) = cache%atomic_grid_weights(row_begin:row_end)
1447 : cache%chunk_atomic_grid_sizes(:) = &
1448 278 : cache%atomic_grid_sizes(cache%chunk_atom_begin:cache%chunk_atom_end)
1449 : cache%chunk_coarse_0_atomic_coords(:, :) = &
1450 716 : cache%coarse_0_atomic_coords(:, cache%chunk_atom_begin:cache%chunk_atom_end)
1451 :
1452 278 : max_grid_size = MAXVAL(INT(cache%chunk_atomic_grid_sizes))
1453 264 : ALLOCATE (cache%chunk_atomic_grid_size_bound_shape(0, max_grid_size))
1454 1570168 : cache%chunk_atomic_grid_size_bound_shape = 0_int_8
1455 1791159 : DO irow = 1, cache%chunk_feature_count
1456 1791159 : cache%chunk_feature_indices(irow) = INT(irow - 1, KIND=int_8)
1457 : END DO
1458 :
1459 : END SUBROUTINE build_atom_chunk_layout
1460 :
1461 : ! **************************************************************************************************
1462 : !> \brief Send local dynamic feature rows to their atom-chunk owner ranks.
1463 : !> \param features ...
1464 : !> \param local_dynamic ...
1465 : !> \param group ...
1466 : !> \param collapse_spin_dynamics ...
1467 : ! **************************************************************************************************
1468 258 : SUBROUTINE route_atom_chunk_dynamics(features, local_dynamic, group, collapse_spin_dynamics)
1469 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1470 : REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: local_dynamic
1471 :
1472 : CLASS(mp_comm_type), INTENT(IN) :: group
1473 : LOGICAL, INTENT(IN) :: collapse_spin_dynamics
1474 :
1475 : INTEGER :: chunk_row, dyn_base, local_row, &
1476 : ndynamic_route_per_point, nrecv, nsend, &
1477 : point_pos, src_base
1478 : INTEGER, ALLOCATABLE, DIMENSION(:) :: recv_counts, recv_displs, send_counts, send_displs
1479 258 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: recv_dynamic, send_dynamic
1480 :
1481 258 : nsend = SIZE(cached_layout%route_send_local_rows)
1482 774 : nrecv = SUM(cached_layout%route_point_recv_counts)
1483 774 : CPASSERT(nsend == SUM(cached_layout%route_point_send_counts))
1484 258 : ndynamic_route_per_point = ndynamic_per_point
1485 258 : IF (collapse_spin_dynamics) ndynamic_route_per_point = nrks_dynamic_per_point
1486 :
1487 0 : ALLOCATE (send_dynamic(MAX(1, ndynamic_route_per_point*nsend)), &
1488 0 : recv_dynamic(MAX(1, ndynamic_route_per_point*nrecv)), &
1489 : send_counts(cached_layout%nproc), send_displs(cached_layout%nproc), &
1490 : recv_counts(cached_layout%nproc), &
1491 2580 : recv_displs(cached_layout%nproc))
1492 774 : send_counts(:) = ndynamic_route_per_point*cached_layout%route_point_send_counts
1493 774 : send_displs(:) = ndynamic_route_per_point*cached_layout%route_point_send_displs
1494 774 : recv_counts(:) = ndynamic_route_per_point*cached_layout%route_point_recv_counts
1495 774 : recv_displs(:) = ndynamic_route_per_point*cached_layout%route_point_recv_displs
1496 : !$OMP PARALLEL DO DEFAULT(NONE) &
1497 : !$OMP SHARED(cached_layout, local_dynamic, ndynamic_route_per_point, nsend, send_dynamic) &
1498 258 : !$OMP PRIVATE(dyn_base, local_row, point_pos, src_base)
1499 : DO point_pos = 1, nsend
1500 : dyn_base = ndynamic_route_per_point*(point_pos - 1)
1501 : local_row = cached_layout%route_send_local_rows(point_pos)
1502 : src_base = ndynamic_route_per_point*(local_row - 1)
1503 : send_dynamic(dyn_base + 1:dyn_base + ndynamic_route_per_point) = &
1504 : local_dynamic(src_base + 1:src_base + ndynamic_route_per_point)
1505 : END DO
1506 : !$OMP END PARALLEL DO
1507 :
1508 : CALL group%alltoall(send_dynamic, send_counts, send_displs, recv_dynamic, recv_counts, &
1509 258 : recv_displs)
1510 :
1511 258 : features%uses_collapsed_rks_dynamic = collapse_spin_dynamics
1512 258 : IF (cached_layout%chunk_feature_count > 0) THEN
1513 258 : IF (collapse_spin_dynamics) THEN
1514 0 : ALLOCATE (features%chunk_density(cached_layout%chunk_feature_count, 1), &
1515 0 : features%chunk_grad(cached_layout%chunk_feature_count, 3, 1), &
1516 1272 : features%chunk_kin(cached_layout%chunk_feature_count, 1))
1517 : ELSE
1518 0 : ALLOCATE (features%chunk_density(cached_layout%chunk_feature_count, 2), &
1519 0 : features%chunk_grad(cached_layout%chunk_feature_count, 3, 2), &
1520 322 : features%chunk_kin(cached_layout%chunk_feature_count, 2))
1521 : END IF
1522 :
1523 : !$OMP PARALLEL DO DEFAULT(NONE) &
1524 : !$OMP SHARED(cached_layout, collapse_spin_dynamics, features, ndynamic_route_per_point, nrecv, recv_dynamic) &
1525 258 : !$OMP PRIVATE(chunk_row, dyn_base, point_pos)
1526 : DO chunk_row = 1, cached_layout%chunk_feature_count
1527 : point_pos = cached_layout%chunk_return_positions(chunk_row)
1528 : CPASSERT(point_pos >= 1 .AND. point_pos <= nrecv)
1529 : dyn_base = ndynamic_route_per_point*(point_pos - 1)
1530 : IF (collapse_spin_dynamics) THEN
1531 : features%chunk_density(chunk_row, 1) = recv_dynamic(dyn_base + 1)
1532 : features%chunk_grad(chunk_row, 1, 1) = recv_dynamic(dyn_base + 2)
1533 : features%chunk_grad(chunk_row, 2, 1) = recv_dynamic(dyn_base + 3)
1534 : features%chunk_grad(chunk_row, 3, 1) = recv_dynamic(dyn_base + 4)
1535 : features%chunk_kin(chunk_row, 1) = recv_dynamic(dyn_base + 5)
1536 : ELSE
1537 : features%chunk_density(chunk_row, :) = recv_dynamic(dyn_base + 1:dyn_base + 2)
1538 : features%chunk_grad(chunk_row, 1, 1) = recv_dynamic(dyn_base + 3)
1539 : features%chunk_grad(chunk_row, 2, 1) = recv_dynamic(dyn_base + 4)
1540 : features%chunk_grad(chunk_row, 3, 1) = recv_dynamic(dyn_base + 5)
1541 : features%chunk_grad(chunk_row, 1, 2) = recv_dynamic(dyn_base + 6)
1542 : features%chunk_grad(chunk_row, 2, 2) = recv_dynamic(dyn_base + 7)
1543 : features%chunk_grad(chunk_row, 3, 2) = recv_dynamic(dyn_base + 8)
1544 : features%chunk_kin(chunk_row, :) = recv_dynamic(dyn_base + 9:dyn_base + 10)
1545 : END IF
1546 : END DO
1547 : !$OMP END PARALLEL DO
1548 : END IF
1549 :
1550 258 : DEALLOCATE (recv_counts, recv_displs, recv_dynamic, send_counts, send_displs, send_dynamic)
1551 :
1552 258 : END SUBROUTINE route_atom_chunk_dynamics
1553 :
1554 : ! **************************************************************************************************
1555 : !> \brief Extract the current rank's atom chunk from the global dynamic feature arrays.
1556 : !> \param features ...
1557 : ! **************************************************************************************************
1558 0 : SUBROUTINE extract_atom_chunk_dynamics(features)
1559 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1560 :
1561 : INTEGER :: row_begin, row_end
1562 :
1563 0 : CPASSERT(cached_layout%chunk_feature_count > 0)
1564 0 : row_begin = cached_layout%chunk_feature_begin
1565 0 : row_end = row_begin + cached_layout%chunk_feature_count - 1
1566 0 : ALLOCATE (features%chunk_density(cached_layout%chunk_feature_count, 2), &
1567 0 : features%chunk_grad(cached_layout%chunk_feature_count, 3, 2), &
1568 0 : features%chunk_kin(cached_layout%chunk_feature_count, 2))
1569 0 : features%chunk_density(:, :) = features%density(row_begin:row_end, :)
1570 0 : features%chunk_grad(:, :, :) = features%grad(row_begin:row_end, :, :)
1571 0 : features%chunk_kin(:, :) = features%kin(row_begin:row_end, :)
1572 :
1573 0 : END SUBROUTINE extract_atom_chunk_dynamics
1574 :
1575 : ! **************************************************************************************************
1576 : !> \brief Compute a local signature for optional integration weights.
1577 : !> \param weights ...
1578 : !> \param has_weights ...
1579 : !> \param weight_sum ...
1580 : !> \param weight_sumsq ...
1581 : ! **************************************************************************************************
1582 322 : SUBROUTINE weights_signature(weights, has_weights, weight_sum, weight_sumsq)
1583 : TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
1584 : LOGICAL, INTENT(OUT) :: has_weights
1585 : REAL(KIND=dp), INTENT(OUT) :: weight_sum, weight_sumsq
1586 :
1587 322 : has_weights = .FALSE.
1588 322 : weight_sum = 0.0_dp
1589 322 : weight_sumsq = 0.0_dp
1590 322 : IF (PRESENT(weights)) THEN
1591 322 : IF (ASSOCIATED(weights)) THEN
1592 0 : has_weights = .TRUE.
1593 0 : weight_sum = SUM(weights%array)
1594 0 : weight_sumsq = SUM(weights%array*weights%array)
1595 : END IF
1596 : END IF
1597 :
1598 322 : END SUBROUTINE weights_signature
1599 :
1600 : ! **************************************************************************************************
1601 : !> \brief Release cached layout arrays.
1602 : !> \param cache ...
1603 : ! **************************************************************************************************
1604 132 : SUBROUTINE release_layout_cache(cache)
1605 : TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
1606 :
1607 132 : IF (cache%inputs_active) THEN
1608 0 : CALL torch_dict_release(cache%inputs)
1609 0 : cache%inputs_active = .FALSE.
1610 : END IF
1611 :
1612 132 : IF (cache%chunk_inputs_active) THEN
1613 40 : CALL torch_dict_release(cache%chunk_inputs)
1614 40 : cache%chunk_inputs_active = .FALSE.
1615 : END IF
1616 :
1617 132 : IF (cache%dynamic_tensors_active) THEN
1618 10 : CALL torch_tensor_release(cache%density_t)
1619 10 : CALL torch_tensor_release(cache%grad_t)
1620 10 : CALL torch_tensor_release(cache%kin_t)
1621 10 : cache%dynamic_tensors_active = .FALSE.
1622 : END IF
1623 :
1624 132 : IF (cache%chunk_dynamic_tensors_active) THEN
1625 40 : IF (cache%chunk_dynamic_input_views_active) THEN
1626 36 : CALL torch_tensor_release(cache%chunk_density_input_t)
1627 36 : CALL torch_tensor_release(cache%chunk_grad_input_t)
1628 36 : CALL torch_tensor_release(cache%chunk_kin_input_t)
1629 36 : cache%chunk_dynamic_input_views_active = .FALSE.
1630 : END IF
1631 40 : CALL torch_tensor_release(cache%chunk_density_t)
1632 40 : CALL torch_tensor_release(cache%chunk_grad_t)
1633 40 : CALL torch_tensor_release(cache%chunk_kin_t)
1634 40 : cache%chunk_dynamic_tensors_active = .FALSE.
1635 : END IF
1636 :
1637 132 : IF (cache%static_tensors_active) THEN
1638 10 : CALL torch_tensor_release(cache%grid_coords_t)
1639 10 : CALL torch_tensor_release(cache%grid_weights_t)
1640 10 : CALL torch_tensor_release(cache%atomic_grid_weights_t)
1641 10 : CALL torch_tensor_release(cache%atomic_grid_sizes_t)
1642 10 : CALL torch_tensor_release(cache%coarse_0_atomic_coords_t)
1643 10 : CALL torch_tensor_release(cache%atomic_grid_size_bound_shape_t)
1644 10 : CALL torch_tensor_release(cache%local_feature_indices_t)
1645 10 : CALL torch_dict_release(cache%static_inputs)
1646 10 : cache%static_tensors_active = .FALSE.
1647 : END IF
1648 :
1649 132 : IF (cache%chunk_static_tensors_active) THEN
1650 40 : CALL torch_tensor_release(cache%chunk_grid_coords_t)
1651 40 : CALL torch_tensor_release(cache%chunk_grid_weights_t)
1652 40 : CALL torch_tensor_release(cache%chunk_atomic_grid_weights_t)
1653 40 : CALL torch_tensor_release(cache%chunk_atomic_grid_sizes_t)
1654 40 : CALL torch_tensor_release(cache%chunk_coarse_0_atomic_coords_t)
1655 40 : CALL torch_tensor_release(cache%chunk_atomic_grid_size_bound_shape_t)
1656 40 : CALL torch_tensor_release(cache%chunk_feature_indices_t)
1657 40 : CALL torch_dict_release(cache%chunk_static_inputs)
1658 : cache%chunk_static_tensors_active = .FALSE.
1659 : END IF
1660 :
1661 132 : IF (ALLOCATED(cache%chunk_feature_counts)) DEALLOCATE (cache%chunk_feature_counts)
1662 132 : IF (ALLOCATED(cache%chunk_feature_displs)) DEALLOCATE (cache%chunk_feature_displs)
1663 132 : IF (ALLOCATED(cache%chunk_grad_counts)) DEALLOCATE (cache%chunk_grad_counts)
1664 132 : IF (ALLOCATED(cache%chunk_grad_displs)) DEALLOCATE (cache%chunk_grad_displs)
1665 132 : IF (ALLOCATED(cache%chunk_return_positions)) DEALLOCATE (cache%chunk_return_positions)
1666 132 : IF (ALLOCATED(cache%route_chunk_offsets)) DEALLOCATE (cache%route_chunk_offsets)
1667 132 : IF (ALLOCATED(cache%route_chunk_rows)) DEALLOCATE (cache%route_chunk_rows)
1668 132 : IF (ALLOCATED(cache%route_point_recv_counts)) DEALLOCATE (cache%route_point_recv_counts)
1669 132 : IF (ALLOCATED(cache%route_point_recv_displs)) DEALLOCATE (cache%route_point_recv_displs)
1670 132 : IF (ALLOCATED(cache%route_point_send_counts)) DEALLOCATE (cache%route_point_send_counts)
1671 132 : IF (ALLOCATED(cache%route_point_send_displs)) DEALLOCATE (cache%route_point_send_displs)
1672 132 : IF (ALLOCATED(cache%route_row_offsets)) DEALLOCATE (cache%route_row_offsets)
1673 132 : IF (ALLOCATED(cache%route_row_positions)) DEALLOCATE (cache%route_row_positions)
1674 132 : IF (ALLOCATED(cache%route_send_local_rows)) DEALLOCATE (cache%route_send_local_rows)
1675 132 : IF (ALLOCATED(cache%dynamic_counts)) DEALLOCATE (cache%dynamic_counts)
1676 132 : IF (ALLOCATED(cache%dynamic_displs)) DEALLOCATE (cache%dynamic_displs)
1677 132 : IF (ALLOCATED(cache%feature_counts)) DEALLOCATE (cache%feature_counts)
1678 132 : IF (ALLOCATED(cache%feature_displs)) DEALLOCATE (cache%feature_displs)
1679 132 : IF (ALLOCATED(cache%feature_source_points)) DEALLOCATE (cache%feature_source_points)
1680 132 : IF (ALLOCATED(cache%global_to_feature)) DEALLOCATE (cache%global_to_feature)
1681 132 : IF (ALLOCATED(cache%feature_index)) DEALLOCATE (cache%feature_index)
1682 132 : IF (ALLOCATED(cache%atomic_grid_sizes)) DEALLOCATE (cache%atomic_grid_sizes)
1683 132 : IF (ALLOCATED(cache%chunk_atomic_grid_sizes)) DEALLOCATE (cache%chunk_atomic_grid_sizes)
1684 132 : IF (ALLOCATED(cache%chunk_feature_indices)) DEALLOCATE (cache%chunk_feature_indices)
1685 132 : IF (ALLOCATED(cache%local_feature_counts)) DEALLOCATE (cache%local_feature_counts)
1686 132 : IF (ALLOCATED(cache%local_feature_indices)) DEALLOCATE (cache%local_feature_indices)
1687 132 : IF (ALLOCATED(cache%local_feature_offsets)) DEALLOCATE (cache%local_feature_offsets)
1688 132 : IF (ALLOCATED(cache%local_feature_points)) DEALLOCATE (cache%local_feature_points)
1689 132 : IF (ALLOCATED(cache%local_feature_rows)) DEALLOCATE (cache%local_feature_rows)
1690 132 : IF (ALLOCATED(cache%atomic_grid_size_bound_shape)) THEN
1691 40 : DEALLOCATE (cache%atomic_grid_size_bound_shape)
1692 : END IF
1693 132 : IF (ALLOCATED(cache%chunk_atomic_grid_size_bound_shape)) THEN
1694 40 : DEALLOCATE (cache%chunk_atomic_grid_size_bound_shape)
1695 : END IF
1696 132 : IF (ALLOCATED(cache%atomic_grid_weights)) DEALLOCATE (cache%atomic_grid_weights)
1697 132 : IF (ALLOCATED(cache%chunk_atomic_grid_weights)) DEALLOCATE (cache%chunk_atomic_grid_weights)
1698 132 : IF (ALLOCATED(cache%chunk_grid_weights)) DEALLOCATE (cache%chunk_grid_weights)
1699 132 : IF (ALLOCATED(cache%grid_weights)) DEALLOCATE (cache%grid_weights)
1700 132 : IF (ALLOCATED(cache%atom_coords)) DEALLOCATE (cache%atom_coords)
1701 132 : IF (ALLOCATED(cache%chunk_coarse_0_atomic_coords)) THEN
1702 40 : DEALLOCATE (cache%chunk_coarse_0_atomic_coords)
1703 : END IF
1704 132 : IF (ALLOCATED(cache%coarse_0_atomic_coords)) DEALLOCATE (cache%coarse_0_atomic_coords)
1705 132 : IF (ALLOCATED(cache%chunk_grid_coords)) DEALLOCATE (cache%chunk_grid_coords)
1706 132 : IF (ALLOCATED(cache%grid_coords)) DEALLOCATE (cache%grid_coords)
1707 :
1708 132 : cache%chunk_atom_begin = 1
1709 132 : cache%chunk_atom_end = 0
1710 132 : cache%chunk_feature_begin = 1
1711 132 : cache%chunk_feature_count = 0
1712 132 : cache%chunk_natom = 0
1713 132 : cache%natom = 0
1714 132 : cache%nflat = 0
1715 132 : cache%nflat_local = 0
1716 132 : cache%npoint = 0
1717 132 : cache%nproc = 0
1718 132 : cache%atom_partition = skala_gpw_atom_partition_hard
1719 1320 : cache%bo = 0
1720 1320 : cache%bounds = 0
1721 528 : cache%npts = 0
1722 132 : cache%dvol = 0.0_dp
1723 132 : cache%weight_sum = 0.0_dp
1724 132 : cache%weight_sumsq = 0.0_dp
1725 1716 : cache%cell_hmat = 0.0_dp
1726 1716 : cache%dh = 0.0_dp
1727 132 : cache%active = .FALSE.
1728 132 : cache%has_weights = .FALSE.
1729 132 : cache%chunk_dynamic_tensors_active = .FALSE.
1730 132 : cache%chunk_dynamic_input_views_active = .FALSE.
1731 132 : cache%chunk_inputs_active = .FALSE.
1732 132 : cache%chunk_inputs_use_collapsed_rks = .FALSE.
1733 132 : cache%chunk_static_tensors_active = .FALSE.
1734 132 : cache%dynamic_tensors_active = .FALSE.
1735 132 : cache%inputs_active = .FALSE.
1736 132 : cache%static_tensors_active = .FALSE.
1737 :
1738 132 : END SUBROUTINE release_layout_cache
1739 :
1740 : ! **************************************************************************************************
1741 : !> \brief Release Torch objects and backing arrays owned by a feature bundle.
1742 : !> \param features ...
1743 : ! **************************************************************************************************
1744 652 : SUBROUTINE skala_gpw_feature_release(features)
1745 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1746 :
1747 652 : IF (features%active) THEN
1748 326 : IF (features%owns_dynamic_tensors) THEN
1749 4 : IF (features%uses_collapsed_rks_dynamic) THEN
1750 4 : CALL torch_tensor_release(features%density_input_t)
1751 4 : CALL torch_tensor_release(features%grad_input_t)
1752 4 : CALL torch_tensor_release(features%kin_input_t)
1753 : END IF
1754 4 : CALL torch_tensor_release(features%density_t)
1755 4 : CALL torch_tensor_release(features%grad_t)
1756 4 : CALL torch_tensor_release(features%kin_t)
1757 : END IF
1758 326 : IF (features%owns_static_tensors) THEN
1759 4 : CALL torch_tensor_release(features%grid_coords_t)
1760 4 : CALL torch_tensor_release(features%grid_weights_t)
1761 4 : CALL torch_tensor_release(features%atomic_grid_weights_t)
1762 4 : CALL torch_tensor_release(features%atomic_grid_sizes_t)
1763 4 : CALL torch_tensor_release(features%atomic_grid_size_bound_shape_t)
1764 : END IF
1765 326 : IF (features%owns_grid_coordinate_tensor) THEN
1766 54 : CALL torch_tensor_release(features%grid_coords_t)
1767 : END IF
1768 326 : IF (features%owns_weight_tensors) THEN
1769 64 : CALL torch_tensor_release(features%grid_weights_t)
1770 64 : CALL torch_tensor_release(features%atomic_grid_weights_t)
1771 : END IF
1772 326 : IF (features%owns_static_tensors .OR. features%owns_coordinate_tensor) THEN
1773 68 : CALL torch_tensor_release(features%coarse_0_atomic_coords_t)
1774 : END IF
1775 326 : IF (features%owns_inputs) CALL torch_dict_release(features%inputs)
1776 326 : features%active = .FALSE.
1777 326 : features%owns_coordinate_tensor = .FALSE.
1778 326 : features%owns_grid_coordinate_tensor = .FALSE.
1779 326 : features%owns_weight_tensors = .FALSE.
1780 326 : features%owns_dynamic_tensors = .TRUE.
1781 326 : features%owns_inputs = .TRUE.
1782 326 : features%owns_static_tensors = .TRUE.
1783 : features%uses_atom_chunk_routing = .FALSE.
1784 326 : features%uses_atom_chunks = .FALSE.
1785 : features%uses_collapsed_rks_dynamic = .FALSE.
1786 : END IF
1787 :
1788 652 : IF (ALLOCATED(features%chunk_density)) DEALLOCATE (features%chunk_density)
1789 652 : IF (ALLOCATED(features%chunk_grad)) DEALLOCATE (features%chunk_grad)
1790 652 : IF (ALLOCATED(features%chunk_kin)) DEALLOCATE (features%chunk_kin)
1791 652 : IF (ALLOCATED(features%density)) DEALLOCATE (features%density)
1792 652 : IF (ALLOCATED(features%grad)) DEALLOCATE (features%grad)
1793 652 : IF (ALLOCATED(features%kin)) DEALLOCATE (features%kin)
1794 652 : IF (ALLOCATED(features%chunk_grad_counts)) DEALLOCATE (features%chunk_grad_counts)
1795 652 : IF (ALLOCATED(features%chunk_grad_displs)) DEALLOCATE (features%chunk_grad_displs)
1796 652 : NULLIFY (features%route_point_recv_counts, features%route_point_recv_displs, &
1797 652 : features%route_point_send_counts, features%route_point_send_displs, &
1798 652 : features%route_chunk_offsets, features%route_chunk_rows, &
1799 652 : features%route_row_offsets, features%route_row_positions, &
1800 652 : features%route_send_local_rows)
1801 652 : IF (ALLOCATED(features%feature_index)) DEALLOCATE (features%feature_index)
1802 652 : IF (ALLOCATED(features%local_feature_counts)) DEALLOCATE (features%local_feature_counts)
1803 652 : IF (ALLOCATED(features%local_feature_offsets)) DEALLOCATE (features%local_feature_offsets)
1804 652 : IF (ALLOCATED(features%local_feature_rows)) DEALLOCATE (features%local_feature_rows)
1805 652 : IF (ALLOCATED(features%grid_coords)) DEALLOCATE (features%grid_coords)
1806 652 : IF (ALLOCATED(features%grid_weights)) DEALLOCATE (features%grid_weights)
1807 652 : IF (ALLOCATED(features%atomic_grid_weights)) DEALLOCATE (features%atomic_grid_weights)
1808 652 : IF (ALLOCATED(features%atomic_grid_sizes)) DEALLOCATE (features%atomic_grid_sizes)
1809 652 : IF (ALLOCATED(features%coarse_0_atomic_coords)) DEALLOCATE (features%coarse_0_atomic_coords)
1810 652 : IF (ALLOCATED(features%atomic_grid_size_bound_shape)) THEN
1811 4 : DEALLOCATE (features%atomic_grid_size_bound_shape)
1812 : END IF
1813 652 : features%chunk_feature_count = 0
1814 652 : features%nflat = 0
1815 652 : features%nflat_local = 0
1816 652 : features%atom_partition = skala_gpw_atom_partition_hard
1817 652 : features%uses_atom_chunk_routing = .FALSE.
1818 652 : features%uses_collapsed_rks_dynamic = .FALSE.
1819 :
1820 652 : END SUBROUTINE skala_gpw_feature_release
1821 :
1822 : ! **************************************************************************************************
1823 : !> \brief Return how many atom-contiguous subchunks the cached rank chunk needs.
1824 : !> \param max_rows ...
1825 : !> \return ...
1826 : ! **************************************************************************************************
1827 260 : FUNCTION skala_gpw_atom_subchunk_count(max_rows) RESULT(nsubchunks)
1828 : INTEGER, INTENT(IN) :: max_rows
1829 : INTEGER :: nsubchunks
1830 :
1831 : INTEGER :: atom_rows, iatom, rows
1832 :
1833 260 : nsubchunks = 0
1834 260 : IF (.NOT. cached_layout%active) RETURN
1835 260 : IF (cached_layout%chunk_natom <= 0) RETURN
1836 260 : IF (max_rows <= 0) THEN
1837 260 : nsubchunks = 1
1838 : RETURN
1839 : END IF
1840 :
1841 : rows = 0
1842 536 : DO iatom = 1, cached_layout%chunk_natom
1843 276 : atom_rows = INT(cached_layout%chunk_atomic_grid_sizes(iatom))
1844 276 : IF (rows > 0 .AND. rows + atom_rows > max_rows) THEN
1845 4 : nsubchunks = nsubchunks + 1
1846 4 : rows = 0
1847 : END IF
1848 536 : rows = rows + atom_rows
1849 : END DO
1850 260 : IF (rows > 0) nsubchunks = nsubchunks + 1
1851 260 : nsubchunks = MAX(1, nsubchunks)
1852 :
1853 260 : END FUNCTION skala_gpw_atom_subchunk_count
1854 :
1855 : ! **************************************************************************************************
1856 : !> \brief Build atom and row bounds for all atom-contiguous subchunks in one pass.
1857 : !> \param max_rows ...
1858 : !> \param atom_begin ...
1859 : !> \param atom_count ...
1860 : !> \param row_begin ...
1861 : !> \param row_count ...
1862 : ! **************************************************************************************************
1863 2 : SUBROUTINE skala_gpw_atom_subchunk_layout(max_rows, atom_begin, atom_count, row_begin, &
1864 : row_count)
1865 : INTEGER, INTENT(IN) :: max_rows
1866 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: atom_begin, atom_count, row_begin, &
1867 : row_count
1868 :
1869 : INTEGER :: atom_begin_tmp, atom_count_tmp, atom_rows, iatom, insert_at, nsubchunks, &
1870 : row_begin_tmp, row_count_tmp, rows, subchunk
1871 : INTEGER(KIND=int_8) :: padded_rows_tmp
1872 2 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: padded_rows
1873 :
1874 2 : nsubchunks = skala_gpw_atom_subchunk_count(max_rows)
1875 0 : ALLOCATE (atom_begin(nsubchunks), atom_count(nsubchunks), row_begin(nsubchunks), &
1876 12 : row_count(nsubchunks))
1877 2 : IF (nsubchunks == 0) RETURN
1878 :
1879 2 : CPASSERT(cached_layout%active)
1880 2 : CPASSERT(cached_layout%chunk_natom > 0)
1881 :
1882 2 : atom_begin = 0
1883 2 : atom_count = 0
1884 2 : row_begin = 0
1885 2 : row_count = 0
1886 :
1887 2 : IF (max_rows <= 0) THEN
1888 0 : atom_begin(1) = 1
1889 0 : atom_count(1) = cached_layout%chunk_natom
1890 0 : row_begin(1) = 1
1891 0 : row_count(1) = cached_layout%chunk_feature_count
1892 0 : RETURN
1893 : END IF
1894 :
1895 2 : subchunk = 1
1896 2 : atom_begin(subchunk) = 1
1897 2 : row_begin(subchunk) = 1
1898 2 : rows = 0
1899 6 : DO iatom = 1, cached_layout%chunk_natom
1900 4 : atom_rows = INT(cached_layout%chunk_atomic_grid_sizes(iatom))
1901 4 : IF (rows > 0 .AND. rows + atom_rows > max_rows) THEN
1902 2 : atom_count(subchunk) = iatom - atom_begin(subchunk)
1903 2 : row_count(subchunk) = rows
1904 2 : subchunk = subchunk + 1
1905 2 : atom_begin(subchunk) = iatom
1906 2 : row_begin(subchunk) = row_begin(subchunk - 1) + row_count(subchunk - 1)
1907 2 : rows = 0
1908 : END IF
1909 6 : rows = rows + atom_rows
1910 : END DO
1911 2 : atom_count(subchunk) = cached_layout%chunk_natom - atom_begin(subchunk) + 1
1912 2 : row_count(subchunk) = rows
1913 :
1914 2 : CPASSERT(subchunk == nsubchunks)
1915 :
1916 6 : ALLOCATE (padded_rows(nsubchunks))
1917 6 : DO subchunk = 1, nsubchunks
1918 4 : iatom = atom_begin(subchunk) + atom_count(subchunk) - 1
1919 : padded_rows(subchunk) = INT(atom_count(subchunk), KIND=int_8)* &
1920 : MAXVAL(cached_layout%chunk_atomic_grid_sizes( &
1921 10 : atom_begin(subchunk):iatom))
1922 : END DO
1923 :
1924 : ! Evaluate the largest padded block first so Torch can reuse its CUDA allocations.
1925 4 : DO subchunk = 2, nsubchunks
1926 2 : padded_rows_tmp = padded_rows(subchunk)
1927 2 : atom_begin_tmp = atom_begin(subchunk)
1928 2 : atom_count_tmp = atom_count(subchunk)
1929 2 : row_begin_tmp = row_begin(subchunk)
1930 2 : row_count_tmp = row_count(subchunk)
1931 2 : insert_at = subchunk
1932 2 : DO WHILE (insert_at > 1 .AND. padded_rows(insert_at - 1) < padded_rows_tmp)
1933 0 : padded_rows(insert_at) = padded_rows(insert_at - 1)
1934 0 : atom_begin(insert_at) = atom_begin(insert_at - 1)
1935 0 : atom_count(insert_at) = atom_count(insert_at - 1)
1936 0 : row_begin(insert_at) = row_begin(insert_at - 1)
1937 0 : row_count(insert_at) = row_count(insert_at - 1)
1938 2 : insert_at = insert_at - 1
1939 : END DO
1940 2 : padded_rows(insert_at) = padded_rows_tmp
1941 2 : atom_begin(insert_at) = atom_begin_tmp
1942 2 : atom_count(insert_at) = atom_count_tmp
1943 2 : row_begin(insert_at) = row_begin_tmp
1944 4 : row_count(insert_at) = row_count_tmp
1945 : END DO
1946 2 : DEALLOCATE (padded_rows)
1947 :
1948 : END SUBROUTINE skala_gpw_atom_subchunk_layout
1949 :
1950 : ! **************************************************************************************************
1951 : !> \brief Build an atom-contiguous subchunk feature bundle from a rank-local atom chunk.
1952 : !> \param parent ...
1953 : !> \param features ...
1954 : !> \param subchunk_index ...
1955 : !> \param max_rows ...
1956 : !> \param requires_grad ...
1957 : ! **************************************************************************************************
1958 0 : SUBROUTINE skala_gpw_feature_build_atom_subchunk(parent, features, subchunk_index, &
1959 : max_rows, requires_grad)
1960 : TYPE(skala_gpw_feature_type), INTENT(IN) :: parent
1961 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1962 : INTEGER, INTENT(IN) :: subchunk_index, max_rows
1963 : LOGICAL, INTENT(IN) :: requires_grad
1964 :
1965 : INTEGER :: atom_begin, atom_end, row_begin, row_end
1966 :
1967 0 : CPASSERT(parent%uses_atom_chunks)
1968 : CALL atom_subchunk_bounds(subchunk_index, max_rows, atom_begin, atom_end, &
1969 0 : row_begin, row_end)
1970 : CALL skala_gpw_feature_build_atom_subchunk_bounds(parent, features, atom_begin, &
1971 : atom_end - atom_begin + 1, row_begin, &
1972 0 : row_end - row_begin + 1, requires_grad)
1973 :
1974 0 : END SUBROUTINE skala_gpw_feature_build_atom_subchunk
1975 :
1976 : ! **************************************************************************************************
1977 : !> \brief Build an atom-contiguous subchunk feature bundle from precomputed bounds.
1978 : !> \param parent ...
1979 : !> \param features ...
1980 : !> \param atom_begin ...
1981 : !> \param atom_count ...
1982 : !> \param row_begin ...
1983 : !> \param row_count ...
1984 : !> \param requires_grad ...
1985 : ! **************************************************************************************************
1986 4 : SUBROUTINE skala_gpw_feature_build_atom_subchunk_bounds(parent, features, atom_begin, &
1987 : atom_count, row_begin, row_count, &
1988 : requires_grad)
1989 : TYPE(skala_gpw_feature_type), INTENT(IN) :: parent
1990 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1991 : INTEGER, INTENT(IN) :: atom_begin, atom_count, row_begin, &
1992 : row_count
1993 : LOGICAL, INTENT(IN) :: requires_grad
1994 :
1995 : INTEGER :: atom_end, max_grid_size, row_end
1996 :
1997 4 : CALL skala_gpw_feature_release(features)
1998 4 : CPASSERT(parent%uses_atom_chunks)
1999 4 : atom_end = atom_begin + atom_count - 1
2000 4 : row_end = row_begin + row_count - 1
2001 4 : CPASSERT(atom_count > 0)
2002 4 : CPASSERT(row_count > 0)
2003 : MARK_USED(requires_grad)
2004 8 : max_grid_size = MAXVAL(INT(cached_layout%chunk_atomic_grid_sizes(atom_begin:atom_end)))
2005 :
2006 8 : ALLOCATE (features%atomic_grid_size_bound_shape(0, max_grid_size))
2007 64004 : features%atomic_grid_size_bound_shape = 0_int_8
2008 :
2009 4 : features%chunk_feature_count = row_count
2010 4 : features%nflat = parent%nflat
2011 4 : features%nflat_local = parent%nflat_local
2012 64004 : features%grid_weight_sum = SUM(cached_layout%chunk_grid_weights(row_begin:row_end))
2013 4 : features%uses_atom_chunks = .TRUE.
2014 4 : features%uses_atom_chunk_routing = parent%uses_atom_chunk_routing
2015 : CALL add_subchunk_feature_tensors(parent, features, atom_begin, atom_count, row_begin, &
2016 4 : row_count)
2017 4 : features%active = .TRUE.
2018 :
2019 4 : END SUBROUTINE skala_gpw_feature_build_atom_subchunk_bounds
2020 :
2021 : ! **************************************************************************************************
2022 : !> \brief Return atom and row bounds for an atom-contiguous rank-local subchunk.
2023 : !> \param subchunk_index ...
2024 : !> \param max_rows ...
2025 : !> \param atom_begin ...
2026 : !> \param atom_end ...
2027 : !> \param row_begin ...
2028 : !> \param row_end ...
2029 : ! **************************************************************************************************
2030 0 : SUBROUTINE atom_subchunk_bounds(subchunk_index, max_rows, atom_begin, atom_end, &
2031 : row_begin, row_end)
2032 : INTEGER, INTENT(IN) :: subchunk_index, max_rows
2033 : INTEGER, INTENT(OUT) :: atom_begin, atom_end, row_begin, row_end
2034 :
2035 : INTEGER :: atom_rows, current_subchunk, iatom, &
2036 : row_cursor, rows
2037 :
2038 0 : CPASSERT(subchunk_index > 0)
2039 0 : CPASSERT(max_rows > 0)
2040 0 : CPASSERT(cached_layout%chunk_natom > 0)
2041 :
2042 0 : atom_begin = 1
2043 0 : atom_end = 0
2044 0 : row_begin = 1
2045 0 : row_end = 0
2046 0 : current_subchunk = 1
2047 0 : row_cursor = 1
2048 0 : rows = 0
2049 0 : DO iatom = 1, cached_layout%chunk_natom
2050 0 : atom_rows = INT(cached_layout%chunk_atomic_grid_sizes(iatom))
2051 0 : IF (rows > 0 .AND. rows + atom_rows > max_rows) THEN
2052 0 : IF (current_subchunk == subchunk_index) THEN
2053 0 : atom_end = iatom - 1
2054 0 : row_end = row_cursor - 1
2055 0 : RETURN
2056 : END IF
2057 0 : current_subchunk = current_subchunk + 1
2058 0 : atom_begin = iatom
2059 0 : row_begin = row_cursor
2060 0 : rows = 0
2061 : END IF
2062 0 : rows = rows + atom_rows
2063 0 : row_cursor = row_cursor + atom_rows
2064 : END DO
2065 :
2066 0 : IF (current_subchunk == subchunk_index) THEN
2067 0 : atom_end = cached_layout%chunk_natom
2068 0 : row_end = row_cursor - 1
2069 0 : RETURN
2070 : END IF
2071 :
2072 0 : CPABORT("Requested native SKALA atom subchunk does not exist.")
2073 :
2074 : END SUBROUTINE atom_subchunk_bounds
2075 :
2076 : ! **************************************************************************************************
2077 : !> \brief Insert a subchunk into a Torch dictionary using static views of the cached chunk tensors.
2078 : !> \param parent ...
2079 : !> \param features ...
2080 : !> \param atom_begin ...
2081 : !> \param atom_count ...
2082 : !> \param row_begin ...
2083 : !> \param row_count ...
2084 : ! **************************************************************************************************
2085 4 : SUBROUTINE add_subchunk_feature_tensors(parent, features, atom_begin, atom_count, row_begin, &
2086 : row_count)
2087 : TYPE(skala_gpw_feature_type), INTENT(IN) :: parent
2088 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
2089 : INTEGER, INTENT(IN) :: atom_begin, atom_count, row_begin, &
2090 : row_count
2091 :
2092 4 : CPASSERT(cached_layout%chunk_static_tensors_active)
2093 4 : CPASSERT(parent%active)
2094 4 : CPASSERT(ALLOCATED(features%atomic_grid_size_bound_shape))
2095 :
2096 4 : features%owns_coordinate_tensor = .FALSE.
2097 4 : features%owns_dynamic_tensors = .TRUE.
2098 4 : features%owns_inputs = .TRUE.
2099 4 : features%owns_static_tensors = .TRUE.
2100 4 : features%uses_collapsed_rks_dynamic = parent%uses_collapsed_rks_dynamic
2101 :
2102 : CALL torch_tensor_narrow(cached_layout%chunk_grid_coords_t, 0, row_begin - 1, &
2103 4 : row_count, features%grid_coords_t)
2104 : CALL torch_tensor_narrow(cached_layout%chunk_grid_weights_t, 0, row_begin - 1, &
2105 4 : row_count, features%grid_weights_t)
2106 : CALL torch_tensor_narrow(cached_layout%chunk_atomic_grid_weights_t, 0, row_begin - 1, &
2107 4 : row_count, features%atomic_grid_weights_t)
2108 : CALL torch_tensor_narrow(cached_layout%chunk_atomic_grid_sizes_t, 0, atom_begin - 1, &
2109 4 : atom_count, features%atomic_grid_sizes_t)
2110 : CALL torch_tensor_narrow(cached_layout%chunk_coarse_0_atomic_coords_t, 0, &
2111 4 : atom_begin - 1, atom_count, features%coarse_0_atomic_coords_t)
2112 : CALL torch_tensor_from_array(features%atomic_grid_size_bound_shape_t, &
2113 4 : features%atomic_grid_size_bound_shape)
2114 4 : CALL torch_tensor_to_device_leaf(features%atomic_grid_size_bound_shape_t, .FALSE.)
2115 : CALL torch_tensor_narrow(parent%density_t, 1, row_begin - 1, row_count, &
2116 4 : features%density_t)
2117 4 : CALL torch_tensor_narrow(parent%grad_t, 2, row_begin - 1, row_count, features%grad_t)
2118 4 : CALL torch_tensor_narrow(parent%kin_t, 1, row_begin - 1, row_count, features%kin_t)
2119 4 : IF (features%uses_collapsed_rks_dynamic) THEN
2120 4 : CALL torch_tensor_expand_dim(features%density_t, 0, 2, features%density_input_t)
2121 4 : CALL torch_tensor_expand_dim(features%grad_t, 0, 2, features%grad_input_t)
2122 4 : CALL torch_tensor_expand_dim(features%kin_t, 0, 2, features%kin_input_t)
2123 : END IF
2124 :
2125 4 : CALL torch_dict_create(features%inputs)
2126 4 : CALL torch_dict_insert(features%inputs, "grid_coords", features%grid_coords_t)
2127 4 : CALL torch_dict_insert(features%inputs, "grid_weights", features%grid_weights_t)
2128 : CALL torch_dict_insert(features%inputs, "atomic_grid_weights", &
2129 4 : features%atomic_grid_weights_t)
2130 : CALL torch_dict_insert(features%inputs, "atomic_grid_sizes", &
2131 4 : features%atomic_grid_sizes_t)
2132 : CALL torch_dict_insert(features%inputs, "atomic_grid_size_bound_shape", &
2133 4 : features%atomic_grid_size_bound_shape_t)
2134 4 : IF (features%uses_collapsed_rks_dynamic) THEN
2135 4 : CALL torch_dict_insert(features%inputs, "density", features%density_input_t)
2136 4 : CALL torch_dict_insert(features%inputs, "grad", features%grad_input_t)
2137 4 : CALL torch_dict_insert(features%inputs, "kin", features%kin_input_t)
2138 : ELSE
2139 0 : CALL torch_dict_insert(features%inputs, "density", features%density_t)
2140 0 : CALL torch_dict_insert(features%inputs, "grad", features%grad_t)
2141 0 : CALL torch_dict_insert(features%inputs, "kin", features%kin_t)
2142 : END IF
2143 : CALL torch_dict_insert(features%inputs, "coarse_0_atomic_coords", &
2144 4 : features%coarse_0_atomic_coords_t)
2145 :
2146 4 : END SUBROUTINE add_subchunk_feature_tensors
2147 :
2148 : ! **************************************************************************************************
2149 : !> \brief Insert owned subchunk arrays into a Torch dictionary.
2150 : !> \param features ...
2151 : !> \param requires_grad ...
2152 : ! **************************************************************************************************
2153 0 : SUBROUTINE add_owned_feature_tensors(features, requires_grad)
2154 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
2155 : LOGICAL, INTENT(IN) :: requires_grad
2156 :
2157 0 : CPASSERT(ALLOCATED(features%chunk_density))
2158 0 : CPASSERT(ALLOCATED(features%chunk_grad))
2159 0 : CPASSERT(ALLOCATED(features%chunk_kin))
2160 0 : CPASSERT(ALLOCATED(features%grid_coords))
2161 0 : CPASSERT(ALLOCATED(features%grid_weights))
2162 0 : CPASSERT(ALLOCATED(features%atomic_grid_weights))
2163 0 : CPASSERT(ALLOCATED(features%atomic_grid_sizes))
2164 0 : CPASSERT(ALLOCATED(features%atomic_grid_size_bound_shape))
2165 0 : CPASSERT(ALLOCATED(features%coarse_0_atomic_coords))
2166 :
2167 0 : features%owns_coordinate_tensor = .FALSE.
2168 0 : features%owns_dynamic_tensors = .TRUE.
2169 0 : features%owns_inputs = .TRUE.
2170 0 : features%owns_static_tensors = .TRUE.
2171 :
2172 0 : CALL torch_tensor_from_array(features%grid_coords_t, features%grid_coords)
2173 0 : CALL torch_tensor_to_device_leaf(features%grid_coords_t, .FALSE.)
2174 0 : CALL torch_tensor_from_array(features%grid_weights_t, features%grid_weights)
2175 0 : CALL torch_tensor_to_device_leaf(features%grid_weights_t, .FALSE.)
2176 0 : CALL torch_tensor_from_array(features%atomic_grid_weights_t, features%atomic_grid_weights)
2177 0 : CALL torch_tensor_to_device_leaf(features%atomic_grid_weights_t, .FALSE.)
2178 0 : CALL torch_tensor_from_array(features%atomic_grid_sizes_t, features%atomic_grid_sizes)
2179 0 : CALL torch_tensor_to_device_leaf(features%atomic_grid_sizes_t, .FALSE.)
2180 : CALL torch_tensor_from_array(features%coarse_0_atomic_coords_t, &
2181 0 : features%coarse_0_atomic_coords)
2182 0 : CALL torch_tensor_to_device_leaf(features%coarse_0_atomic_coords_t, .FALSE.)
2183 : CALL torch_tensor_from_array(features%atomic_grid_size_bound_shape_t, &
2184 0 : features%atomic_grid_size_bound_shape)
2185 0 : CALL torch_tensor_to_device_leaf(features%atomic_grid_size_bound_shape_t, .FALSE.)
2186 0 : CALL torch_tensor_from_array(features%density_t, features%chunk_density)
2187 0 : CALL torch_tensor_to_device_leaf(features%density_t, requires_grad)
2188 0 : CALL torch_tensor_from_array(features%grad_t, features%chunk_grad)
2189 0 : CALL torch_tensor_to_device_leaf(features%grad_t, requires_grad)
2190 0 : CALL torch_tensor_from_array(features%kin_t, features%chunk_kin)
2191 0 : CALL torch_tensor_to_device_leaf(features%kin_t, requires_grad)
2192 :
2193 0 : CALL torch_dict_create(features%inputs)
2194 0 : CALL torch_dict_insert(features%inputs, "grid_coords", features%grid_coords_t)
2195 0 : CALL torch_dict_insert(features%inputs, "grid_weights", features%grid_weights_t)
2196 : CALL torch_dict_insert(features%inputs, "atomic_grid_weights", &
2197 0 : features%atomic_grid_weights_t)
2198 : CALL torch_dict_insert(features%inputs, "atomic_grid_sizes", &
2199 0 : features%atomic_grid_sizes_t)
2200 : CALL torch_dict_insert(features%inputs, "atomic_grid_size_bound_shape", &
2201 0 : features%atomic_grid_size_bound_shape_t)
2202 0 : CALL torch_dict_insert(features%inputs, "density", features%density_t)
2203 0 : CALL torch_dict_insert(features%inputs, "grad", features%grad_t)
2204 0 : CALL torch_dict_insert(features%inputs, "kin", features%kin_t)
2205 : CALL torch_dict_insert(features%inputs, "coarse_0_atomic_coords", &
2206 0 : features%coarse_0_atomic_coords_t)
2207 :
2208 0 : END SUBROUTINE add_owned_feature_tensors
2209 :
2210 : ! **************************************************************************************************
2211 : !> \brief Insert all SKALA feature tensors into the Torch dictionary.
2212 : !> \param features ...
2213 : !> \param requires_grad ...
2214 : !> \param requires_coordinate_grad ...
2215 : !> \param requires_stress_grad ...
2216 : !> \param use_atom_chunks ...
2217 : !> \param requires_weight_grad ...
2218 : ! **************************************************************************************************
2219 322 : SUBROUTINE add_feature_tensors(features, requires_grad, requires_coordinate_grad, &
2220 : requires_stress_grad, use_atom_chunks, requires_weight_grad)
2221 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
2222 : LOGICAL, INTENT(IN) :: requires_grad, requires_coordinate_grad, &
2223 : requires_stress_grad, use_atom_chunks
2224 : LOGICAL, INTENT(IN), OPTIONAL :: requires_weight_grad
2225 :
2226 : LOGICAL :: my_requires_weight_grad
2227 :
2228 322 : my_requires_weight_grad = .FALSE.
2229 322 : IF (PRESENT(requires_weight_grad)) my_requires_weight_grad = requires_weight_grad
2230 :
2231 322 : features%owns_static_tensors = .FALSE.
2232 322 : features%owns_coordinate_tensor = .FALSE.
2233 322 : features%owns_grid_coordinate_tensor = .FALSE.
2234 322 : features%owns_weight_tensors = .FALSE.
2235 322 : features%owns_dynamic_tensors = .FALSE.
2236 322 : features%owns_inputs = .TRUE.
2237 322 : IF (use_atom_chunks) THEN
2238 258 : CPASSERT(.NOT. requires_coordinate_grad)
2239 258 : CPASSERT(.NOT. requires_stress_grad)
2240 258 : CPASSERT(.NOT. my_requires_weight_grad)
2241 258 : CPASSERT(cached_layout%chunk_static_tensors_active)
2242 258 : features%grid_coords_t = cached_layout%chunk_grid_coords_t
2243 258 : features%grid_weights_t = cached_layout%chunk_grid_weights_t
2244 258 : features%atomic_grid_weights_t = cached_layout%chunk_atomic_grid_weights_t
2245 258 : features%atomic_grid_sizes_t = cached_layout%chunk_atomic_grid_sizes_t
2246 : features%atomic_grid_size_bound_shape_t = &
2247 258 : cached_layout%chunk_atomic_grid_size_bound_shape_t
2248 258 : features%local_feature_indices_t = cached_layout%chunk_feature_indices_t
2249 :
2250 258 : IF (cached_layout%chunk_inputs_active .AND. &
2251 : (cached_layout%chunk_inputs_use_collapsed_rks .NEQV. &
2252 : features%uses_collapsed_rks_dynamic)) THEN
2253 0 : CALL torch_dict_release(cached_layout%chunk_inputs)
2254 0 : cached_layout%chunk_inputs_active = .FALSE.
2255 : END IF
2256 258 : IF (.NOT. features%uses_collapsed_rks_dynamic .AND. &
2257 : cached_layout%chunk_dynamic_input_views_active) THEN
2258 0 : CALL torch_tensor_release(cached_layout%chunk_density_input_t)
2259 0 : CALL torch_tensor_release(cached_layout%chunk_grad_input_t)
2260 0 : CALL torch_tensor_release(cached_layout%chunk_kin_input_t)
2261 0 : cached_layout%chunk_dynamic_input_views_active = .FALSE.
2262 : END IF
2263 :
2264 : CALL torch_tensor_reset_from_array(cached_layout%chunk_density_t, &
2265 258 : features%chunk_density, requires_grad=requires_grad)
2266 258 : features%density_t = cached_layout%chunk_density_t
2267 : CALL torch_tensor_reset_from_array(cached_layout%chunk_grad_t, features%chunk_grad, &
2268 258 : requires_grad=requires_grad)
2269 258 : features%grad_t = cached_layout%chunk_grad_t
2270 : CALL torch_tensor_reset_from_array(cached_layout%chunk_kin_t, features%chunk_kin, &
2271 258 : requires_grad=requires_grad)
2272 258 : features%kin_t = cached_layout%chunk_kin_t
2273 258 : cached_layout%chunk_dynamic_tensors_active = .TRUE.
2274 :
2275 258 : IF (features%uses_collapsed_rks_dynamic .AND. &
2276 : .NOT. cached_layout%chunk_dynamic_input_views_active) THEN
2277 : CALL torch_tensor_expand_dim(cached_layout%chunk_density_t, 0, 2, &
2278 120 : cached_layout%chunk_density_input_t)
2279 : CALL torch_tensor_expand_dim(cached_layout%chunk_grad_t, 0, 2, &
2280 120 : cached_layout%chunk_grad_input_t)
2281 : CALL torch_tensor_expand_dim(cached_layout%chunk_kin_t, 0, 2, &
2282 120 : cached_layout%chunk_kin_input_t)
2283 120 : cached_layout%chunk_dynamic_input_views_active = .TRUE.
2284 : END IF
2285 258 : IF (features%uses_collapsed_rks_dynamic) THEN
2286 212 : features%density_input_t = cached_layout%chunk_density_input_t
2287 212 : features%grad_input_t = cached_layout%chunk_grad_input_t
2288 212 : features%kin_input_t = cached_layout%chunk_kin_input_t
2289 : END IF
2290 :
2291 258 : IF (.NOT. cached_layout%chunk_inputs_active) THEN
2292 132 : CALL torch_dict_clone(cached_layout%chunk_static_inputs, cached_layout%chunk_inputs)
2293 132 : IF (features%uses_collapsed_rks_dynamic) THEN
2294 : CALL torch_dict_insert(cached_layout%chunk_inputs, "density", &
2295 120 : features%density_input_t)
2296 : CALL torch_dict_insert(cached_layout%chunk_inputs, "grad", &
2297 120 : features%grad_input_t)
2298 : CALL torch_dict_insert(cached_layout%chunk_inputs, "kin", &
2299 120 : features%kin_input_t)
2300 : ELSE
2301 : CALL torch_dict_insert(cached_layout%chunk_inputs, "density", &
2302 12 : cached_layout%chunk_density_t)
2303 : CALL torch_dict_insert(cached_layout%chunk_inputs, "grad", &
2304 12 : cached_layout%chunk_grad_t)
2305 : CALL torch_dict_insert(cached_layout%chunk_inputs, "kin", &
2306 12 : cached_layout%chunk_kin_t)
2307 : END IF
2308 : CALL torch_dict_insert(cached_layout%chunk_inputs, "coarse_0_atomic_coords", &
2309 132 : cached_layout%chunk_coarse_0_atomic_coords_t)
2310 132 : cached_layout%chunk_inputs_use_collapsed_rks = features%uses_collapsed_rks_dynamic
2311 132 : cached_layout%chunk_inputs_active = .TRUE.
2312 : END IF
2313 258 : features%inputs = cached_layout%chunk_inputs
2314 258 : features%owns_inputs = .FALSE.
2315 258 : features%coarse_0_atomic_coords_t = cached_layout%chunk_coarse_0_atomic_coords_t
2316 : ELSE
2317 64 : CPASSERT(cached_layout%static_tensors_active)
2318 64 : IF (.NOT. requires_stress_grad .AND. .NOT. my_requires_weight_grad) THEN
2319 0 : features%grid_coords_t = cached_layout%grid_coords_t
2320 0 : features%grid_weights_t = cached_layout%grid_weights_t
2321 0 : features%atomic_grid_weights_t = cached_layout%atomic_grid_weights_t
2322 : END IF
2323 64 : features%atomic_grid_sizes_t = cached_layout%atomic_grid_sizes_t
2324 64 : features%atomic_grid_size_bound_shape_t = cached_layout%atomic_grid_size_bound_shape_t
2325 64 : features%local_feature_indices_t = cached_layout%local_feature_indices_t
2326 :
2327 : CALL torch_tensor_reset_from_array(cached_layout%density_t, features%density, &
2328 64 : requires_grad=requires_grad)
2329 64 : features%density_t = cached_layout%density_t
2330 : CALL torch_tensor_reset_from_array(cached_layout%grad_t, features%grad, &
2331 64 : requires_grad=requires_grad)
2332 64 : features%grad_t = cached_layout%grad_t
2333 : CALL torch_tensor_reset_from_array(cached_layout%kin_t, features%kin, &
2334 64 : requires_grad=requires_grad)
2335 64 : features%kin_t = cached_layout%kin_t
2336 64 : cached_layout%dynamic_tensors_active = .TRUE.
2337 :
2338 64 : IF (requires_coordinate_grad .OR. requires_stress_grad .OR. my_requires_weight_grad) THEN
2339 64 : IF (requires_stress_grad .OR. my_requires_weight_grad) THEN
2340 64 : CALL torch_dict_create(features%inputs)
2341 64 : IF (requires_stress_grad) THEN
2342 54 : CALL torch_tensor_from_array(features%grid_coords_t, features%grid_coords)
2343 54 : CALL torch_tensor_to_device_leaf(features%grid_coords_t, .TRUE.)
2344 54 : CALL torch_dict_insert(features%inputs, "grid_coords", features%grid_coords_t)
2345 54 : features%owns_grid_coordinate_tensor = .TRUE.
2346 : ELSE
2347 10 : features%grid_coords_t = cached_layout%grid_coords_t
2348 10 : CALL torch_dict_insert(features%inputs, "grid_coords", features%grid_coords_t)
2349 : END IF
2350 64 : CALL torch_tensor_from_array(features%grid_weights_t, features%grid_weights)
2351 64 : CALL torch_tensor_to_device_leaf(features%grid_weights_t, .TRUE.)
2352 : CALL torch_tensor_from_array(features%atomic_grid_weights_t, &
2353 64 : features%atomic_grid_weights)
2354 64 : CALL torch_tensor_to_device_leaf(features%atomic_grid_weights_t, .TRUE.)
2355 64 : CALL torch_dict_insert(features%inputs, "grid_weights", features%grid_weights_t)
2356 : CALL torch_dict_insert(features%inputs, "atomic_grid_weights", &
2357 64 : features%atomic_grid_weights_t)
2358 : CALL torch_dict_insert(features%inputs, "atomic_grid_sizes", &
2359 64 : features%atomic_grid_sizes_t)
2360 : CALL torch_dict_insert(features%inputs, "atomic_grid_size_bound_shape", &
2361 64 : features%atomic_grid_size_bound_shape_t)
2362 64 : features%owns_weight_tensors = .TRUE.
2363 : ELSE
2364 0 : CALL torch_dict_clone(cached_layout%static_inputs, features%inputs)
2365 : END IF
2366 64 : CALL torch_dict_insert(features%inputs, "density", features%density_t)
2367 64 : CALL torch_dict_insert(features%inputs, "grad", features%grad_t)
2368 64 : CALL torch_dict_insert(features%inputs, "kin", features%kin_t)
2369 : ELSE
2370 0 : IF (.NOT. cached_layout%inputs_active) THEN
2371 0 : CALL torch_dict_clone(cached_layout%static_inputs, cached_layout%inputs)
2372 0 : CALL torch_dict_insert(cached_layout%inputs, "density", cached_layout%density_t)
2373 0 : CALL torch_dict_insert(cached_layout%inputs, "grad", cached_layout%grad_t)
2374 0 : CALL torch_dict_insert(cached_layout%inputs, "kin", cached_layout%kin_t)
2375 : CALL torch_dict_insert(cached_layout%inputs, "coarse_0_atomic_coords", &
2376 0 : cached_layout%coarse_0_atomic_coords_t)
2377 0 : cached_layout%inputs_active = .TRUE.
2378 : END IF
2379 0 : features%inputs = cached_layout%inputs
2380 0 : features%owns_inputs = .FALSE.
2381 0 : features%coarse_0_atomic_coords_t = cached_layout%coarse_0_atomic_coords_t
2382 : END IF
2383 : END IF
2384 :
2385 322 : IF (requires_coordinate_grad .OR. requires_stress_grad) THEN
2386 64 : CPASSERT(.NOT. use_atom_chunks)
2387 : CALL torch_tensor_from_array(features%coarse_0_atomic_coords_t, &
2388 64 : features%coarse_0_atomic_coords)
2389 64 : CALL torch_tensor_to_device_leaf(features%coarse_0_atomic_coords_t, .TRUE.)
2390 : CALL torch_dict_insert(features%inputs, "coarse_0_atomic_coords", &
2391 64 : features%coarse_0_atomic_coords_t)
2392 64 : features%owns_coordinate_tensor = .TRUE.
2393 : END IF
2394 :
2395 322 : END SUBROUTINE add_feature_tensors
2396 :
2397 : ! **************************************************************************************************
2398 : !> \brief Return the Cartesian coordinate of a regular GPW grid point.
2399 : !> \param pw_grid ...
2400 : !> \param index ...
2401 : !> \return ...
2402 : ! **************************************************************************************************
2403 1390210 : FUNCTION grid_coordinate(pw_grid, index) RESULT(coord)
2404 : TYPE(pw_grid_type), POINTER :: pw_grid
2405 : INTEGER, DIMENSION(3), INTENT(IN) :: index
2406 : REAL(KIND=dp), DIMENSION(3) :: coord
2407 :
2408 : INTEGER, DIMENSION(3) :: relative_index
2409 :
2410 5560840 : relative_index = index - pw_grid%bounds(1, :)
2411 : coord = REAL(relative_index(1), KIND=dp)*pw_grid%dh(:, 1) + &
2412 : REAL(relative_index(2), KIND=dp)*pw_grid%dh(:, 2) + &
2413 5560840 : REAL(relative_index(3), KIND=dp)*pw_grid%dh(:, 3)
2414 :
2415 1390210 : END FUNCTION grid_coordinate
2416 :
2417 : ! **************************************************************************************************
2418 : !> \brief Build Becke-like smooth atom weights for one native-grid point.
2419 : !> \param grid_point ...
2420 : !> \param atom_coords ...
2421 : !> \param cell ...
2422 : !> \param weights ...
2423 : !> \param partition_atom_coords ...
2424 : !> \param distances ...
2425 : !> \param pair_distances ...
2426 : ! **************************************************************************************************
2427 473023 : SUBROUTINE smooth_atom_partition(grid_point, atom_coords, cell, weights, partition_atom_coords, &
2428 473023 : distances, pair_distances)
2429 : REAL(KIND=dp), DIMENSION(3), INTENT(IN) :: grid_point
2430 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: atom_coords
2431 : TYPE(cell_type), POINTER :: cell
2432 : REAL(KIND=dp), DIMENSION(:), INTENT(OUT) :: weights
2433 : REAL(KIND=dp), DIMENSION(:, :), INTENT(OUT) :: partition_atom_coords
2434 : REAL(KIND=dp), DIMENSION(:), INTENT(OUT) :: distances
2435 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN), &
2436 : OPTIONAL :: pair_distances
2437 :
2438 : INTEGER :: iatom, jatom, natom
2439 : REAL(KIND=dp) :: mu, rab, rsum, switch, total
2440 : REAL(KIND=dp), DIMENSION(3) :: rij
2441 :
2442 473023 : natom = SIZE(atom_coords, 2)
2443 473023 : CPASSERT(SIZE(weights) == natom)
2444 473023 : CPASSERT(SIZE(partition_atom_coords, 1) == 3)
2445 473023 : CPASSERT(SIZE(partition_atom_coords, 2) == natom)
2446 473023 : CPASSERT(SIZE(distances) == natom)
2447 :
2448 473023 : IF (PRESENT(pair_distances)) THEN
2449 89424 : DO iatom = 1, natom
2450 238464 : rij = grid_point - atom_coords(:, iatom)
2451 268272 : distances(iatom) = SQRT(SUM(rij**2))
2452 : END DO
2453 : ELSE
2454 1329645 : DO iatom = 1, natom
2455 : partition_atom_coords(:, iatom) = &
2456 886430 : nearest_atom_image_coordinate(atom_coords(:, iatom), grid_point, cell)
2457 3545720 : rij = grid_point - partition_atom_coords(:, iatom)
2458 3988935 : distances(iatom) = SQRT(SUM(rij**2))
2459 : END DO
2460 : END IF
2461 :
2462 1419069 : weights = 1.0_dp
2463 946046 : DO iatom = 1, natom - 1
2464 1419069 : DO jatom = iatom + 1, natom
2465 473023 : IF (PRESENT(pair_distances)) THEN
2466 29808 : rab = pair_distances(iatom, jatom)
2467 : ELSE
2468 1772860 : rij = partition_atom_coords(:, iatom) - partition_atom_coords(:, jatom)
2469 1772860 : rab = SQRT(SUM(rij**2))
2470 : END IF
2471 473023 : IF (rab <= layout_tol) CYCLE
2472 473023 : mu = (distances(iatom) - distances(jatom))/rab
2473 473023 : mu = MAX(-1.0_dp, MIN(1.0_dp, mu))
2474 473023 : switch = 0.5_dp*(1.0_dp - becke_shape(mu))
2475 473023 : weights(iatom) = weights(iatom)*switch
2476 946046 : weights(jatom) = weights(jatom)*(1.0_dp - switch)
2477 : END DO
2478 : END DO
2479 :
2480 1419069 : total = SUM(weights)
2481 473023 : IF (total > 0.0_dp) THEN
2482 1419069 : weights = weights/total
2483 : ELSE
2484 : rsum = HUGE(1.0_dp)
2485 : jatom = 1
2486 0 : DO iatom = 1, natom
2487 0 : IF (distances(iatom) < rsum) THEN
2488 0 : rsum = distances(iatom)
2489 0 : jatom = iatom
2490 : END IF
2491 : END DO
2492 0 : weights = 0.0_dp
2493 0 : weights(jatom) = 1.0_dp
2494 : END IF
2495 :
2496 473023 : END SUBROUTINE smooth_atom_partition
2497 :
2498 : ! **************************************************************************************************
2499 : !> \brief Build smooth atom weights and their atom/cell deformation derivatives.
2500 : !> \param grid_point ...
2501 : !> \param atom_coords ...
2502 : !> \param cell ...
2503 : !> \param weights ...
2504 : !> \param included ...
2505 : !> \param dweights_datom ...
2506 : !> \param dweights_dstrain ...
2507 : ! **************************************************************************************************
2508 615625 : SUBROUTINE skala_gpw_smooth_partition_derivatives(grid_point, atom_coords, cell, &
2509 615625 : weights, included, dweights_datom, &
2510 615625 : dweights_dstrain)
2511 : REAL(KIND=dp), DIMENSION(3), INTENT(IN) :: grid_point
2512 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: atom_coords
2513 : TYPE(cell_type), POINTER :: cell
2514 : REAL(KIND=dp), DIMENSION(:), INTENT(OUT) :: weights
2515 : LOGICAL, DIMENSION(:), INTENT(OUT) :: included
2516 : REAL(KIND=dp), DIMENSION(:, :, :), INTENT(OUT) :: dweights_datom, dweights_dstrain
2517 :
2518 : INTEGER :: iatom, idir, jatom, jdir, natom
2519 : REAL(KIND=dp) :: dist_diff, ds_dmu, included_sum, mu, &
2520 : mu_raw, one_minus_switch, rab, rsum, &
2521 : switch, total
2522 : REAL(KIND=dp), DIMENSION(3) :: dmu_atom_i, dmu_atom_j, ds_atom_i, &
2523 : ds_atom_j, pair, unit_pair
2524 : REAL(KIND=dp), DIMENSION(3, 3) :: dmu_strain, ds_strain, mean_strain
2525 : REAL(KIND=dp), DIMENSION(3, SIZE(atom_coords, 2), &
2526 1231250 : SIZE(atom_coords, 2)) :: log_weight_atom
2527 1231250 : REAL(KIND=dp), DIMENSION(3, SIZE(atom_coords, 2)) :: mean_atom, partition_atom_coords, rvecs, &
2528 1231250 : unit_rvecs
2529 : REAL(KIND=dp), &
2530 1231250 : DIMENSION(3, 3, SIZE(atom_coords, 2)) :: log_weight_strain
2531 1231250 : REAL(KIND=dp), DIMENSION(SIZE(atom_coords, 2)) :: distances, normalized_weights, &
2532 615625 : raw_weights
2533 :
2534 615625 : natom = SIZE(atom_coords, 2)
2535 615625 : CPASSERT(SIZE(weights) == natom)
2536 615625 : CPASSERT(SIZE(included) == natom)
2537 615625 : CPASSERT(SIZE(dweights_datom, 1) == 3)
2538 615625 : CPASSERT(SIZE(dweights_datom, 2) == natom)
2539 615625 : CPASSERT(SIZE(dweights_datom, 3) == natom)
2540 615625 : CPASSERT(SIZE(dweights_dstrain, 1) == 3)
2541 615625 : CPASSERT(SIZE(dweights_dstrain, 2) == 3)
2542 615625 : CPASSERT(SIZE(dweights_dstrain, 3) == natom)
2543 :
2544 1846875 : weights = 0.0_dp
2545 1846875 : included = .FALSE.
2546 11696875 : dweights_datom = 0.0_dp
2547 16621875 : dweights_dstrain = 0.0_dp
2548 1846875 : raw_weights = 1.0_dp
2549 11696875 : log_weight_atom = 0.0_dp
2550 16621875 : log_weight_strain = 0.0_dp
2551 :
2552 1846875 : DO iatom = 1, natom
2553 : partition_atom_coords(:, iatom) = &
2554 1231250 : nearest_atom_image_coordinate(atom_coords(:, iatom), grid_point, cell)
2555 4925000 : rvecs(:, iatom) = grid_point - partition_atom_coords(:, iatom)
2556 4925000 : distances(iatom) = SQRT(SUM(rvecs(:, iatom)**2))
2557 1846875 : IF (distances(iatom) > layout_tol) THEN
2558 4925000 : unit_rvecs(:, iatom) = rvecs(:, iatom)/distances(iatom)
2559 : ELSE
2560 0 : unit_rvecs(:, iatom) = 0.0_dp
2561 : END IF
2562 : END DO
2563 :
2564 1231250 : DO iatom = 1, natom - 1
2565 1846875 : DO jatom = iatom + 1, natom
2566 2462500 : pair = partition_atom_coords(:, iatom) - partition_atom_coords(:, jatom)
2567 2462500 : rab = SQRT(SUM(pair**2))
2568 615625 : IF (rab <= layout_tol) CYCLE
2569 2462500 : unit_pair = pair/rab
2570 615625 : dist_diff = distances(iatom) - distances(jatom)
2571 615625 : mu_raw = dist_diff/rab
2572 615625 : mu = MAX(-1.0_dp, MIN(1.0_dp, mu_raw))
2573 615625 : switch = 0.5_dp*(1.0_dp - becke_shape(mu))
2574 615625 : one_minus_switch = 1.0_dp - switch
2575 :
2576 615625 : IF (ABS(mu_raw) < 1.0_dp) THEN
2577 614871 : ds_dmu = -0.5_dp*becke_shape_derivative(mu)
2578 : ELSE
2579 : ds_dmu = 0.0_dp
2580 : END IF
2581 614871 : IF (ABS(ds_dmu) > 0.0_dp .AND. switch > TINY(1.0_dp) .AND. &
2582 : one_minus_switch > TINY(1.0_dp)) THEN
2583 2451784 : dmu_atom_i = (-unit_rvecs(:, iatom)*rab - dist_diff*unit_pair)/rab**2
2584 2451784 : dmu_atom_j = (unit_rvecs(:, jatom)*rab + dist_diff*unit_pair)/rab**2
2585 2451784 : ds_atom_i = ds_dmu*dmu_atom_i
2586 2451784 : ds_atom_j = ds_dmu*dmu_atom_j
2587 : log_weight_atom(:, iatom, iatom) = &
2588 2451784 : log_weight_atom(:, iatom, iatom) + ds_atom_i/switch
2589 : log_weight_atom(:, iatom, jatom) = &
2590 2451784 : log_weight_atom(:, iatom, jatom) - ds_atom_i/one_minus_switch
2591 : log_weight_atom(:, jatom, iatom) = &
2592 2451784 : log_weight_atom(:, jatom, iatom) + ds_atom_j/switch
2593 : log_weight_atom(:, jatom, jatom) = &
2594 2451784 : log_weight_atom(:, jatom, jatom) - ds_atom_j/one_minus_switch
2595 :
2596 2451784 : DO idir = 1, 3
2597 7968298 : DO jdir = 1, 3
2598 : dmu_strain(idir, jdir) = &
2599 : ((unit_rvecs(idir, iatom)*rvecs(jdir, iatom) - &
2600 : unit_rvecs(idir, jatom)*rvecs(jdir, jatom))*rab - &
2601 7355352 : dist_diff*unit_pair(idir)*pair(jdir))/rab**2
2602 : END DO
2603 : END DO
2604 7968298 : ds_strain = ds_dmu*dmu_strain
2605 : log_weight_strain(:, :, iatom) = &
2606 7968298 : log_weight_strain(:, :, iatom) + ds_strain/switch
2607 : log_weight_strain(:, :, jatom) = &
2608 7968298 : log_weight_strain(:, :, jatom) - ds_strain/one_minus_switch
2609 : END IF
2610 :
2611 615625 : raw_weights(iatom) = raw_weights(iatom)*switch
2612 1231250 : raw_weights(jatom) = raw_weights(jatom)*one_minus_switch
2613 : END DO
2614 : END DO
2615 :
2616 1846875 : total = SUM(raw_weights)
2617 615625 : IF (total > 0.0_dp) THEN
2618 1846875 : normalized_weights = raw_weights/total
2619 1846875 : included = normalized_weights > smooth_partition_eps
2620 : ELSE
2621 : rsum = HUGE(1.0_dp)
2622 : jatom = 1
2623 0 : DO iatom = 1, natom
2624 0 : IF (distances(iatom) < rsum) THEN
2625 0 : rsum = distances(iatom)
2626 0 : jatom = iatom
2627 : END IF
2628 : END DO
2629 0 : included(jatom) = .TRUE.
2630 0 : weights(jatom) = 1.0_dp
2631 0 : RETURN
2632 : END IF
2633 :
2634 1846875 : included_sum = SUM(raw_weights, MASK=included)
2635 615625 : IF (included_sum <= 0.0_dp) THEN
2636 : rsum = HUGE(1.0_dp)
2637 : jatom = 1
2638 0 : DO iatom = 1, natom
2639 0 : IF (distances(iatom) < rsum) THEN
2640 0 : rsum = distances(iatom)
2641 0 : jatom = iatom
2642 : END IF
2643 : END DO
2644 0 : included = .FALSE.
2645 0 : included(jatom) = .TRUE.
2646 0 : weights = 0.0_dp
2647 0 : weights(jatom) = 1.0_dp
2648 0 : RETURN
2649 : END IF
2650 :
2651 1846875 : DO iatom = 1, natom
2652 1846875 : IF (included(iatom)) weights(iatom) = raw_weights(iatom)/included_sum
2653 : END DO
2654 :
2655 5540625 : mean_atom = 0.0_dp
2656 615625 : mean_strain = 0.0_dp
2657 1846875 : DO iatom = 1, natom
2658 1231250 : IF (.NOT. included(iatom)) CYCLE
2659 15930980 : mean_strain = mean_strain + weights(iatom)*log_weight_strain(:, :, iatom)
2660 4292005 : DO jatom = 1, natom
2661 : mean_atom(:, jatom) = mean_atom(:, jatom) + &
2662 11034930 : weights(iatom)*log_weight_atom(:, jatom, iatom)
2663 : END DO
2664 : END DO
2665 :
2666 1846875 : DO iatom = 1, natom
2667 1231250 : IF (.NOT. included(iatom)) CYCLE
2668 : dweights_dstrain(:, :, iatom) = &
2669 15930980 : weights(iatom)*(log_weight_strain(:, :, iatom) - mean_strain)
2670 4292005 : DO jatom = 1, natom
2671 : dweights_datom(:, jatom, iatom) = &
2672 11034930 : weights(iatom)*(log_weight_atom(:, jatom, iatom) - mean_atom(:, jatom))
2673 : END DO
2674 : END DO
2675 :
2676 : END SUBROUTINE skala_gpw_smooth_partition_derivatives
2677 :
2678 : ! **************************************************************************************************
2679 : !> \brief Smoothly suppress a sparse atom row's internal quadrature weight at the layout cutoff.
2680 : !> \param weight normalized smooth-partition weight
2681 : !> \return scale in the interval [0, 1]
2682 : ! **************************************************************************************************
2683 803840 : PURE FUNCTION smooth_partition_atomic_weight_scale(weight) RESULT(scale)
2684 : REAL(KIND=dp), INTENT(IN) :: weight
2685 : REAL(KIND=dp) :: scale
2686 :
2687 : REAL(KIND=dp) :: x
2688 :
2689 803840 : IF (weight <= smooth_partition_eps) THEN
2690 : scale = 0.0_dp
2691 803840 : ELSE IF (weight >= smooth_partition_taper_end) THEN
2692 : scale = 1.0_dp
2693 : ELSE
2694 : x = (weight - smooth_partition_eps)/ &
2695 884 : (smooth_partition_taper_end - smooth_partition_eps)
2696 884 : scale = x**3*(10.0_dp + x*(-15.0_dp + 6.0_dp*x))
2697 : END IF
2698 :
2699 803840 : END FUNCTION smooth_partition_atomic_weight_scale
2700 :
2701 : ! **************************************************************************************************
2702 : !> \brief Derivative of the sparse atom-row quadrature taper with respect to partition weight.
2703 : !> \param weight normalized smooth-partition weight
2704 : !> \return derivative of the internal quadrature-weight scale
2705 : ! **************************************************************************************************
2706 1208210 : PURE FUNCTION smooth_partition_atomic_weight_scale_derivative(weight) RESULT(derivative)
2707 : REAL(KIND=dp), INTENT(IN) :: weight
2708 : REAL(KIND=dp) :: derivative
2709 :
2710 : REAL(KIND=dp) :: x
2711 :
2712 1208210 : IF (weight <= smooth_partition_eps .OR. weight >= smooth_partition_taper_end) THEN
2713 : derivative = 0.0_dp
2714 : ELSE
2715 : x = (weight - smooth_partition_eps)/ &
2716 1488 : (smooth_partition_taper_end - smooth_partition_eps)
2717 : derivative = 30.0_dp*x**2*(1.0_dp - x)**2/ &
2718 1488 : (smooth_partition_taper_end - smooth_partition_eps)
2719 : END IF
2720 :
2721 1208210 : END FUNCTION smooth_partition_atomic_weight_scale_derivative
2722 :
2723 : ! **************************************************************************************************
2724 : !> \brief Becke fuzzy-cell shape function.
2725 : !> \param mu ...
2726 : !> \return ...
2727 : ! **************************************************************************************************
2728 1088648 : PURE FUNCTION becke_shape(mu) RESULT(val)
2729 : REAL(KIND=dp), INTENT(IN) :: mu
2730 : REAL(KIND=dp) :: val
2731 :
2732 : INTEGER :: iter
2733 :
2734 1088648 : val = mu
2735 4354592 : DO iter = 1, 3
2736 4354592 : val = 0.5_dp*val*(3.0_dp - val*val)
2737 : END DO
2738 :
2739 1088648 : END FUNCTION becke_shape
2740 :
2741 : ! **************************************************************************************************
2742 : !> \brief Derivative of the Becke fuzzy-cell shape function.
2743 : !> \param mu ...
2744 : !> \return ...
2745 : ! **************************************************************************************************
2746 614871 : PURE FUNCTION becke_shape_derivative(mu) RESULT(val)
2747 : REAL(KIND=dp), INTENT(IN) :: mu
2748 : REAL(KIND=dp) :: val
2749 :
2750 : INTEGER :: iter
2751 : REAL(KIND=dp) :: x
2752 :
2753 614871 : x = mu
2754 614871 : val = 1.0_dp
2755 2459484 : DO iter = 1, 3
2756 1844613 : val = val*1.5_dp*(1.0_dp - x*x)
2757 2459484 : x = 0.5_dp*x*(3.0_dp - x*x)
2758 : END DO
2759 :
2760 614871 : END FUNCTION becke_shape_derivative
2761 :
2762 : ! **************************************************************************************************
2763 : !> \brief Return the atom image nearest to a regular-grid point.
2764 : !> \param atom_coord ...
2765 : !> \param grid_point ...
2766 : !> \param cell ...
2767 : !> \return ...
2768 : ! **************************************************************************************************
2769 2117680 : FUNCTION nearest_atom_image_coordinate(atom_coord, grid_point, cell) RESULT(coord)
2770 : REAL(KIND=dp), DIMENSION(3), INTENT(IN) :: atom_coord, grid_point
2771 : TYPE(cell_type), POINTER :: cell
2772 : REAL(KIND=dp), DIMENSION(3) :: coord
2773 :
2774 : REAL(KIND=dp) :: dx, dy, dz
2775 :
2776 2117680 : IF (cell%orthorhombic) THEN
2777 2117680 : dx = atom_coord(1) - grid_point(1)
2778 2117680 : dy = atom_coord(2) - grid_point(2)
2779 2117680 : dz = atom_coord(3) - grid_point(3)
2780 2117680 : dx = dx - cell%hmat(1, 1)*cell%perd(1)*ANINT(cell%h_inv(1, 1)*dx)
2781 2117680 : dy = dy - cell%hmat(2, 2)*cell%perd(2)*ANINT(cell%h_inv(2, 2)*dy)
2782 2117680 : dz = dz - cell%hmat(3, 3)*cell%perd(3)*ANINT(cell%h_inv(3, 3)*dz)
2783 8470720 : coord = grid_point + [dx, dy, dz]
2784 : ELSE
2785 0 : coord = grid_point + pbc(grid_point, atom_coord, cell)
2786 : END IF
2787 :
2788 2117680 : END FUNCTION nearest_atom_image_coordinate
2789 :
2790 : ! **************************************************************************************************
2791 : !> \brief Assign a grid point to the nearest periodic atom.
2792 : !> \param grid_point ...
2793 : !> \param atom_coords ...
2794 : !> \param cell ...
2795 : !> \return ...
2796 : ! **************************************************************************************************
2797 987187 : FUNCTION nearest_atom(grid_point, atom_coords, cell) RESULT(owner)
2798 : REAL(KIND=dp), DIMENSION(3), INTENT(IN) :: grid_point
2799 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: atom_coords
2800 : TYPE(cell_type), POINTER :: cell
2801 : INTEGER :: owner
2802 :
2803 : INTEGER :: iatom
2804 : REAL(KIND=dp) :: best_r2, dx, dy, dz, r2
2805 : REAL(KIND=dp), DIMENSION(3) :: rij
2806 :
2807 987187 : owner = 1
2808 987187 : best_r2 = HUGE(1.0_dp)
2809 987187 : IF (cell%orthorhombic) THEN
2810 3886904 : DO iatom = 1, SIZE(atom_coords, 2)
2811 2899717 : dx = grid_point(1) - atom_coords(1, iatom)
2812 2899717 : dy = grid_point(2) - atom_coords(2, iatom)
2813 2899717 : dz = grid_point(3) - atom_coords(3, iatom)
2814 2899717 : dx = dx - cell%hmat(1, 1)*cell%perd(1)*ANINT(cell%h_inv(1, 1)*dx)
2815 2899717 : dy = dy - cell%hmat(2, 2)*cell%perd(2)*ANINT(cell%h_inv(2, 2)*dy)
2816 2899717 : dz = dz - cell%hmat(3, 3)*cell%perd(3)*ANINT(cell%h_inv(3, 3)*dz)
2817 2899717 : r2 = dx*dx + dy*dy + dz*dz
2818 3886904 : IF (r2 < best_r2) THEN
2819 1773819 : best_r2 = r2
2820 1773819 : owner = iatom
2821 : END IF
2822 : END DO
2823 : ELSE
2824 0 : DO iatom = 1, SIZE(atom_coords, 2)
2825 0 : rij = pbc(grid_point, atom_coords(:, iatom), cell)
2826 0 : r2 = SUM(rij**2)
2827 0 : IF (r2 < best_r2) THEN
2828 0 : best_r2 = r2
2829 0 : owner = iatom
2830 : END IF
2831 : END DO
2832 : END IF
2833 :
2834 987187 : END FUNCTION nearest_atom
2835 :
2836 0 : END MODULE skala_gpw_features
|
