Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2024 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : ! **************************************************************************************************
9 : !> \brief Utility methods to build 3-center integral tensors of various types.
10 : ! **************************************************************************************************
11 :
12 : MODULE qs_tensors_types
13 : USE atomic_kind_types, ONLY: atomic_kind_type,&
14 : get_atomic_kind_set
15 : USE basis_set_types, ONLY: get_gto_basis_set,&
16 : gto_basis_set_p_type
17 : USE cp_array_utils, ONLY: cp_1d_i_p_type
18 : USE cp_blacs_env, ONLY: cp_blacs_env_create,&
19 : cp_blacs_env_release,&
20 : cp_blacs_env_type
21 : USE dbt_api, ONLY: dbt_create,&
22 : dbt_default_distvec,&
23 : dbt_distribution_destroy,&
24 : dbt_distribution_new,&
25 : dbt_distribution_type,&
26 : dbt_mp_environ_pgrid,&
27 : dbt_pgrid_type,&
28 : dbt_type
29 : USE distribution_2d_types, ONLY: distribution_2d_create_prv => distribution_2d_create,&
30 : distribution_2d_release,&
31 : distribution_2d_type
32 : USE message_passing, ONLY: mp_cart_type,&
33 : mp_comm_type,&
34 : mp_para_env_release,&
35 : mp_para_env_type
36 : USE particle_types, ONLY: particle_type
37 : USE qs_neighbor_list_types, ONLY: neighbor_list_iterator_p_type,&
38 : neighbor_list_set_p_type
39 : #include "./base/base_uses.f90"
40 :
41 : IMPLICIT NONE
42 :
43 : PRIVATE
44 :
45 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'qs_tensors_types'
46 :
47 : PUBLIC :: distribution_3d_type, neighbor_list_3c_type, neighbor_list_3c_iterator_type, &
48 : distribution_2d_create, distribution_3d_create, distribution_3d_destroy, &
49 : split_block_sizes, create_3c_tensor, create_2c_tensor, contiguous_tensor_dist, pgf_block_sizes, &
50 : create_tensor_batches
51 :
52 : INTEGER, PARAMETER, PUBLIC :: symmetric_none = 0, symmetric_ij = 1, symmetric_jk = 2, symmetrik_ik = 3, symmetric_ijk = 4
53 :
54 : INTEGER, PARAMETER, PUBLIC :: default_block_size = 64
55 : !! default block size for dense tensors, this block size should be covered by DBCSR/libcusmm
56 :
57 : TYPE distribution_3d_type
58 : TYPE(distribution_2d_type), POINTER :: dist_2d_1 => NULL(), dist_2d_2 => NULL()
59 : TYPE(mp_comm_type) :: comm_3d = mp_comm_type(), comm_2d_1 = mp_comm_type(), comm_2d_2 = mp_comm_type()
60 : LOGICAL :: owns_comm = .FALSE.
61 : END TYPE distribution_3d_type
62 :
63 : TYPE neighbor_list_3c_type
64 : TYPE(neighbor_list_set_p_type), DIMENSION(:), POINTER :: ij_list => NULL(), jk_list => NULL()
65 : INTEGER :: sym = symmetric_none
66 : TYPE(distribution_3d_type) :: dist_3d = distribution_3d_type()
67 : LOGICAL :: owns_dist = .FALSE.
68 : END TYPE
69 :
70 : TYPE neighbor_list_3c_iterator_type
71 : TYPE(neighbor_list_iterator_p_type), DIMENSION(:), POINTER :: iter_ij => NULL()
72 : TYPE(neighbor_list_iterator_p_type), DIMENSION(:), POINTER :: iter_jk => NULL()
73 : INTEGER :: iter_level = 0
74 : TYPE(neighbor_list_3c_type) :: ijk_nl = neighbor_list_3c_type()
75 : INTEGER, DIMENSION(2) :: bounds_i = 0, bounds_j = 0, bounds_k = 0
76 : END TYPE
77 :
78 : CONTAINS
79 : ! **************************************************************************************************
80 : !> \brief Create a 3d distribution
81 : !> \param dist_3d 3d distribution object
82 : !> \param dist1 distribution vector along 1st process grid dimension
83 : !> \param dist2 distribution vector along 2nd process grid dimension
84 : !> \param dist3 distribution vector along 3rd process grid dimension
85 : !> \param nkind ...
86 : !> \param particle_set ...
87 : !> \param mp_comm_3d MPI communicator with a 3d cartesian topology
88 : !> \param own_comm Whether mp_comm_3d should be owned by dist_3d (default false)
89 : ! **************************************************************************************************
90 770 : SUBROUTINE distribution_3d_create(dist_3d, dist1, dist2, dist3, nkind, particle_set, mp_comm_3d, own_comm)
91 : TYPE(distribution_3d_type) :: dist_3d
92 : INTEGER, DIMENSION(:), INTENT(IN) :: dist1, dist2, dist3
93 : INTEGER, INTENT(IN) :: nkind
94 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
95 : TYPE(mp_cart_type), INTENT(IN) :: mp_comm_3d
96 : LOGICAL, INTENT(IN), OPTIONAL :: own_comm
97 :
98 : CHARACTER(len=*), PARAMETER :: routineN = 'distribution_3d_create'
99 :
100 : INTEGER :: handle
101 : INTEGER, DIMENSION(2) :: mp_coor_1, mp_coor_2
102 770 : TYPE(mp_cart_type) :: comm_2d_1, comm_2d_2
103 :
104 770 : CALL timeset(routineN, handle)
105 :
106 770 : IF (PRESENT(own_comm)) THEN
107 770 : IF (own_comm) dist_3d%comm_3d = mp_comm_3d
108 770 : dist_3d%owns_comm = own_comm
109 : ELSE
110 0 : dist_3d%owns_comm = .FALSE.
111 : END IF
112 :
113 770 : CALL comm_2d_1%from_sub(mp_comm_3d, [.TRUE., .TRUE., .FALSE.])
114 770 : CALL comm_2d_2%from_sub(mp_comm_3d, [.FALSE., .TRUE., .TRUE.])
115 :
116 2310 : mp_coor_1 = comm_2d_1%mepos_cart
117 2310 : mp_coor_2 = comm_2d_2%mepos_cart
118 :
119 770 : CPASSERT(mp_coor_1(2) == mp_coor_2(1))
120 :
121 770 : CALL distribution_2d_create(dist_3d%dist_2d_1, dist1, dist2, nkind, particle_set, comm_2d_1)
122 770 : CALL distribution_2d_create(dist_3d%dist_2d_2, dist2, dist3, nkind, particle_set, comm_2d_2)
123 :
124 770 : dist_3d%comm_2d_1 = comm_2d_1
125 770 : dist_3d%comm_2d_2 = comm_2d_2
126 :
127 770 : CALL timestop(handle)
128 770 : END SUBROUTINE
129 :
130 : ! **************************************************************************************************
131 : !> \brief Destroy a 3d distribution
132 : !> \param dist ...
133 : ! **************************************************************************************************
134 770 : SUBROUTINE distribution_3d_destroy(dist)
135 : TYPE(distribution_3d_type) :: dist
136 :
137 : CHARACTER(len=*), PARAMETER :: routineN = 'distribution_3d_destroy'
138 :
139 : INTEGER :: handle
140 :
141 770 : CALL timeset(routineN, handle)
142 770 : CALL distribution_2d_release(dist%dist_2d_1)
143 770 : CALL distribution_2d_release(dist%dist_2d_2)
144 770 : CALL dist%comm_2d_1%free()
145 770 : CALL dist%comm_2d_2%free()
146 770 : IF (dist%owns_comm) CALL dist%comm_3d%free()
147 :
148 770 : NULLIFY (dist%dist_2d_1, dist%dist_2d_2)
149 :
150 770 : CALL timestop(handle)
151 770 : END SUBROUTINE
152 :
153 : ! **************************************************************************************************
154 : !> \brief Create a 2d distribution. This mainly wraps distribution_2d_create
155 : !> for consistency with distribution_3d_create.
156 : !> \param dist_2d 2d distribution object
157 : !> \param dist1 distribution vector along 1st process grid dimension
158 : !> \param dist2 distribution vector along 2nd process grid dimension
159 : !> \param nkind ...
160 : !> \param particle_set ...
161 : !> \param mp_comm_2d MPI communicator with a 3d cartesian topology
162 : !> \param blacs_env_ext ...
163 : ! **************************************************************************************************
164 9244 : SUBROUTINE distribution_2d_create(dist_2d, dist1, dist2, nkind, particle_set, mp_comm_2d, blacs_env_ext)
165 : TYPE(distribution_2d_type), POINTER :: dist_2d
166 : INTEGER, DIMENSION(:), INTENT(IN) :: dist1, dist2
167 : INTEGER, INTENT(IN) :: nkind
168 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
169 : TYPE(mp_cart_type), INTENT(IN), OPTIONAL :: mp_comm_2d
170 : TYPE(cp_blacs_env_type), OPTIONAL, POINTER :: blacs_env_ext
171 :
172 : CHARACTER(len=*), PARAMETER :: routineN = 'distribution_2d_create'
173 :
174 : INTEGER :: handle, iatom, ikind, n, natom
175 9244 : INTEGER, ALLOCATABLE, DIMENSION(:) :: nparticle_local_col, nparticle_local_row
176 : INTEGER, DIMENSION(2) :: mp_coor, mp_dims
177 9244 : INTEGER, DIMENSION(:, :), POINTER :: dist1_prv, dist2_prv
178 9244 : TYPE(cp_1d_i_p_type), DIMENSION(:), POINTER :: local_particle_col, local_particle_row
179 : TYPE(cp_blacs_env_type), POINTER :: blacs_env
180 : TYPE(mp_para_env_type), POINTER :: para_env
181 :
182 9244 : NULLIFY (blacs_env, local_particle_col, local_particle_row, para_env)
183 :
184 9244 : CALL timeset(routineN, handle)
185 :
186 9244 : CPASSERT(PRESENT(mp_comm_2d) .OR. PRESENT(blacs_env_ext))
187 :
188 9244 : IF (PRESENT(mp_comm_2d)) THEN
189 4620 : mp_dims = mp_comm_2d%num_pe_cart
190 4620 : mp_coor = mp_comm_2d%mepos_cart
191 1540 : ALLOCATE (para_env)
192 1540 : para_env = mp_comm_2d
193 : CALL cp_blacs_env_create(blacs_env, para_env, &
194 1540 : grid_2d=mp_dims)
195 :
196 1540 : CPASSERT(blacs_env%mepos(1) == mp_coor(1))
197 1540 : CPASSERT(blacs_env%mepos(2) == mp_coor(2))
198 1540 : CALL mp_para_env_release(para_env)
199 : END IF
200 :
201 9244 : IF (PRESENT(blacs_env_ext)) THEN
202 7704 : blacs_env => blacs_env_ext
203 7704 : mp_coor(1) = blacs_env%mepos(1)
204 7704 : mp_coor(2) = blacs_env%mepos(2)
205 : END IF
206 :
207 9244 : natom = SIZE(particle_set)
208 46220 : ALLOCATE (dist1_prv(natom, 2), dist2_prv(natom, 2))
209 29104 : dist1_prv(:, 1) = dist1
210 29104 : dist2_prv(:, 1) = dist2
211 :
212 75316 : ALLOCATE (local_particle_col(nkind), local_particle_row(nkind))
213 46220 : ALLOCATE (nparticle_local_row(nkind), nparticle_local_col(nkind))
214 47584 : nparticle_local_row = 0; nparticle_local_col = 0
215 :
216 29104 : DO iatom = 1, natom
217 19860 : ikind = particle_set(iatom)%atomic_kind%kind_number
218 :
219 19860 : IF (dist1_prv(iatom, 1) == mp_coor(1)) nparticle_local_row(ikind) = nparticle_local_row(ikind) + 1
220 29104 : IF (dist2_prv(iatom, 1) == mp_coor(2)) nparticle_local_col(ikind) = nparticle_local_col(ikind) + 1
221 : END DO
222 :
223 23792 : DO ikind = 1, nkind
224 14548 : n = nparticle_local_row(ikind)
225 42670 : ALLOCATE (local_particle_row(ikind)%array(n))
226 :
227 14548 : n = nparticle_local_col(ikind)
228 52830 : ALLOCATE (local_particle_col(ikind)%array(n))
229 : END DO
230 :
231 47584 : nparticle_local_row = 0; nparticle_local_col = 0
232 29104 : DO iatom = 1, natom
233 19860 : ikind = particle_set(iatom)%atomic_kind%kind_number
234 :
235 19860 : IF (dist1_prv(iatom, 1) == mp_coor(1)) THEN
236 18075 : nparticle_local_row(ikind) = nparticle_local_row(ikind) + 1
237 18075 : local_particle_row(ikind)%array(nparticle_local_row(ikind)) = iatom
238 : END IF
239 29104 : IF (dist2_prv(iatom, 1) == mp_coor(2)) THEN
240 19724 : nparticle_local_col(ikind) = nparticle_local_col(ikind) + 1
241 19724 : local_particle_col(ikind)%array(nparticle_local_col(ikind)) = iatom
242 : END IF
243 : END DO
244 :
245 : CALL distribution_2d_create_prv(dist_2d, row_distribution_ptr=dist1_prv, &
246 : col_distribution_ptr=dist2_prv, local_rows_ptr=local_particle_row, &
247 9244 : local_cols_ptr=local_particle_col, blacs_env=blacs_env)
248 :
249 9244 : IF (.NOT. PRESENT(blacs_env_ext)) THEN
250 1540 : CALL cp_blacs_env_release(blacs_env)
251 : END IF
252 :
253 9244 : CALL timestop(handle)
254 18488 : END SUBROUTINE
255 :
256 : ! **************************************************************************************************
257 : !> \brief contiguous distribution of weighted elements
258 : !> \param nel ...
259 : !> \param nbin ...
260 : !> \param weights ...
261 : !> \param limits_start ...
262 : !> \param limits_end ...
263 : !> \param dist ...
264 : ! **************************************************************************************************
265 1440 : SUBROUTINE contiguous_tensor_dist(nel, nbin, weights, limits_start, limits_end, dist)
266 : INTEGER, INTENT(IN) :: nel
267 : INTEGER, INTENT(INOUT) :: nbin
268 : INTEGER, DIMENSION(nel), INTENT(IN) :: weights
269 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT), &
270 : OPTIONAL :: limits_start, limits_end
271 : INTEGER, DIMENSION(nel), INTENT(OUT), OPTIONAL :: dist
272 :
273 : INTEGER :: el_end, el_start, end_weight, ibin, &
274 : nel_div, nel_rem, nel_split, nel_w, &
275 : w_partialsum
276 1440 : INTEGER, ALLOCATABLE, DIMENSION(:) :: lim_e, lim_s
277 :
278 7200 : ALLOCATE (lim_s(nbin), lim_e(nbin))
279 10736 : lim_s = 0; lim_e = 0
280 :
281 8714 : nel_w = SUM(weights)
282 1440 : nel_div = nel_w/nbin
283 1440 : nel_rem = MOD(nel_w, nbin)
284 :
285 1440 : w_partialsum = 0
286 1440 : el_end = 0
287 1440 : end_weight = 0
288 3702 : DO ibin = 1, nbin
289 3702 : nel_split = nel_div
290 3702 : IF (ibin <= nel_rem) THEN
291 936 : nel_split = nel_split + 1
292 : END IF
293 3702 : el_start = el_end + 1
294 3702 : el_end = el_start
295 3702 : w_partialsum = w_partialsum + weights(el_end)
296 3702 : end_weight = end_weight + nel_split
297 6592 : DO WHILE (w_partialsum < end_weight)
298 : !IF (ABS(w_partialsum + weights(el_end) - end_weight) > ABS(w_partialsum - end_weight)) EXIT
299 3572 : el_end = el_end + 1
300 3572 : w_partialsum = w_partialsum + weights(el_end)
301 6592 : IF (el_end == nel) EXIT
302 : END DO
303 :
304 3702 : IF (PRESENT(dist)) dist(el_start:el_end) = ibin - 1
305 3702 : lim_s(ibin) = el_start
306 3702 : lim_e(ibin) = el_end
307 :
308 3702 : IF (el_end == nel) EXIT
309 : END DO
310 :
311 1440 : IF (PRESENT(limits_start) .AND. PRESENT(limits_end)) THEN
312 8022 : ALLOCATE (limits_start(ibin)); limits_start(:ibin) = lim_s(:ibin)
313 8022 : ALLOCATE (limits_end(ibin)); limits_end(:ibin) = lim_e(:ibin)
314 : END IF
315 :
316 1440 : nbin = ibin
317 :
318 1440 : END SUBROUTINE contiguous_tensor_dist
319 :
320 : ! **************************************************************************************************
321 : !> \brief ...
322 : !> \param t3c Create 3-center tensor with load balanced default distribution.
323 : !> \param dist_1 ...
324 : !> \param dist_2 ...
325 : !> \param dist_3 ...
326 : !> \param pgrid ...
327 : !> \param sizes_1 ...
328 : !> \param sizes_2 ...
329 : !> \param sizes_3 ...
330 : !> \param map1 ...
331 : !> \param map2 ...
332 : !> \param name ...
333 : ! **************************************************************************************************
334 46300 : SUBROUTINE create_3c_tensor(t3c, dist_1, dist_2, dist_3, pgrid, sizes_1, sizes_2, sizes_3, map1, map2, name)
335 : TYPE(dbt_type), INTENT(OUT) :: t3c
336 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: dist_1, dist_2, dist_3
337 : TYPE(dbt_pgrid_type), INTENT(IN) :: pgrid
338 : INTEGER, DIMENSION(:), INTENT(IN) :: sizes_1, sizes_2, sizes_3, map1, map2
339 : CHARACTER(len=*), INTENT(IN) :: name
340 :
341 : CHARACTER(len=*), PARAMETER :: routineN = 'create_3c_tensor'
342 :
343 : INTEGER :: handle, size_1, size_2, size_3
344 : INTEGER, DIMENSION(3) :: pcoord, pdims
345 41670 : TYPE(dbt_distribution_type) :: dist
346 :
347 4630 : CALL timeset(routineN, handle)
348 :
349 4630 : CALL dbt_mp_environ_pgrid(pgrid, pdims, pcoord)
350 :
351 4630 : size_1 = SIZE(sizes_1)
352 4630 : size_2 = SIZE(sizes_2)
353 4630 : size_3 = SIZE(sizes_3)
354 :
355 13890 : ALLOCATE (dist_1(size_1))
356 13890 : ALLOCATE (dist_2(size_2))
357 13890 : ALLOCATE (dist_3(size_3))
358 :
359 4630 : CALL dbt_default_distvec(size_1, pdims(1), sizes_1, dist_1)
360 4630 : CALL dbt_default_distvec(size_2, pdims(2), sizes_2, dist_2)
361 4630 : CALL dbt_default_distvec(size_3, pdims(3), sizes_3, dist_3)
362 :
363 4630 : CALL dbt_distribution_new(dist, pgrid, dist_1, dist_2, dist_3)
364 4630 : CALL dbt_create(t3c, name, dist, map1, map2, sizes_1, sizes_2, sizes_3)
365 4630 : CALL dbt_distribution_destroy(dist)
366 :
367 4630 : CALL timestop(handle)
368 9260 : END SUBROUTINE
369 :
370 : ! **************************************************************************************************
371 : !> \brief ...
372 : !> \param t2c ...
373 : !> \param dist_1 ...
374 : !> \param dist_2 ...
375 : !> \param pgrid ...
376 : !> \param sizes_1 ...
377 : !> \param sizes_2 ...
378 : !> \param order ...
379 : !> \param name ...
380 : ! **************************************************************************************************
381 75520 : SUBROUTINE create_2c_tensor(t2c, dist_1, dist_2, pgrid, sizes_1, sizes_2, order, name)
382 : TYPE(dbt_type), INTENT(OUT) :: t2c
383 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: dist_1, dist_2
384 : TYPE(dbt_pgrid_type), INTENT(IN) :: pgrid
385 : INTEGER, DIMENSION(:), INTENT(IN) :: sizes_1, sizes_2
386 : INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL :: order
387 : CHARACTER(len=*), INTENT(IN) :: name
388 :
389 : CHARACTER(len=*), PARAMETER :: routineN = 'create_2c_tensor'
390 :
391 : INTEGER :: handle, size_1, size_2
392 : INTEGER, DIMENSION(2) :: order_in, pcoord, pdims
393 67968 : TYPE(dbt_distribution_type) :: dist
394 :
395 7552 : CALL timeset(routineN, handle)
396 :
397 7552 : IF (PRESENT(order)) THEN
398 0 : order_in = order
399 : ELSE
400 7552 : order_in = [1, 2]
401 : END IF
402 :
403 7552 : CALL dbt_mp_environ_pgrid(pgrid, pdims, pcoord)
404 :
405 7552 : size_1 = SIZE(sizes_1)
406 7552 : size_2 = SIZE(sizes_2)
407 :
408 22656 : ALLOCATE (dist_1(size_1))
409 22656 : ALLOCATE (dist_2(size_2))
410 :
411 7552 : CALL dbt_default_distvec(size_1, pdims(1), sizes_1, dist_1)
412 7552 : CALL dbt_default_distvec(size_2, pdims(2), sizes_2, dist_2)
413 :
414 7552 : CALL dbt_distribution_new(dist, pgrid, dist_1, dist_2)
415 22656 : CALL dbt_create(t2c, name, dist, [order_in(1)], [order_in(2)], sizes_1, sizes_2)
416 7552 : CALL dbt_distribution_destroy(dist)
417 :
418 7552 : CALL timestop(handle)
419 15104 : END SUBROUTINE
420 :
421 : ! **************************************************************************************************
422 : !> \brief ...
423 : !> \param blk_sizes ...
424 : !> \param blk_sizes_split ...
425 : !> \param max_size ...
426 : ! **************************************************************************************************
427 640 : SUBROUTINE split_block_sizes(blk_sizes, blk_sizes_split, max_size)
428 : INTEGER, DIMENSION(:), INTENT(IN) :: blk_sizes
429 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: blk_sizes_split
430 : INTEGER, INTENT(IN) :: max_size
431 :
432 : INTEGER :: blk_remainder, i, isplit, isplit_sum, &
433 : nsplit
434 :
435 640 : isplit_sum = 0
436 2178 : DO i = 1, SIZE(blk_sizes)
437 1538 : nsplit = (blk_sizes(i) + max_size - 1)/max_size
438 2178 : isplit_sum = isplit_sum + nsplit
439 : END DO
440 :
441 1920 : ALLOCATE (blk_sizes_split(isplit_sum))
442 :
443 640 : isplit_sum = 0
444 2178 : DO i = 1, SIZE(blk_sizes)
445 1538 : nsplit = (blk_sizes(i) + max_size - 1)/max_size
446 1538 : blk_remainder = blk_sizes(i)
447 3892 : DO isplit = 1, nsplit
448 1714 : isplit_sum = isplit_sum + 1
449 1714 : blk_sizes_split(isplit_sum) = MIN(max_size, blk_remainder)
450 3252 : blk_remainder = blk_remainder - max_size
451 : END DO
452 : END DO
453 :
454 640 : END SUBROUTINE split_block_sizes
455 :
456 : ! **************************************************************************************************
457 : !> \brief ...
458 : !> \param atomic_kind_set ...
459 : !> \param basis ...
460 : !> \param min_blk_size ...
461 : !> \param pgf_blk_sizes ...
462 : ! **************************************************************************************************
463 868 : SUBROUTINE pgf_block_sizes(atomic_kind_set, basis, min_blk_size, pgf_blk_sizes)
464 : TYPE(atomic_kind_type), DIMENSION(:), POINTER :: atomic_kind_set
465 : TYPE(gto_basis_set_p_type), DIMENSION(:), &
466 : INTENT(IN) :: basis
467 : INTEGER, INTENT(IN) :: min_blk_size
468 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: pgf_blk_sizes
469 :
470 : INTEGER :: blk_count, blk_count_prev, blk_size, &
471 : iatom, ikind, iset, natom, nblk, nset
472 868 : INTEGER, ALLOCATABLE, DIMENSION(:) :: kind_of, pgf_blk_sizes_tmp
473 868 : INTEGER, DIMENSION(:), POINTER :: nsgf_set
474 :
475 868 : CALL get_atomic_kind_set(atomic_kind_set, natom=natom, kind_of=kind_of)
476 :
477 868 : nblk = 0
478 3428 : DO iatom = 1, natom
479 2560 : ikind = kind_of(iatom)
480 2560 : CALL get_gto_basis_set(basis(ikind)%gto_basis_set, nset=nset)
481 3428 : nblk = nblk + nset
482 : END DO
483 :
484 13442 : ALLOCATE (pgf_blk_sizes_tmp(nblk)); pgf_blk_sizes_tmp = 0
485 :
486 : blk_count = 0
487 : blk_size = 0
488 3428 : DO iatom = 1, natom
489 2560 : blk_count_prev = blk_count
490 2560 : ikind = kind_of(iatom)
491 2560 : CALL get_gto_basis_set(basis(ikind)%gto_basis_set, nset=nset, nsgf_set=nsgf_set)
492 13398 : DO iset = 1, nset
493 10838 : blk_size = blk_size + nsgf_set(iset)
494 13398 : IF (blk_size >= min_blk_size) THEN
495 5772 : blk_count = blk_count + 1
496 5772 : pgf_blk_sizes_tmp(blk_count) = pgf_blk_sizes_tmp(blk_count) + blk_size
497 5772 : blk_size = 0
498 : END IF
499 : END DO
500 5988 : IF (blk_size > 0) THEN
501 812 : IF (blk_count == blk_count_prev) blk_count = blk_count + 1
502 812 : pgf_blk_sizes_tmp(blk_count) = pgf_blk_sizes_tmp(blk_count) + blk_size
503 812 : blk_size = 0
504 : END IF
505 : END DO
506 :
507 2604 : ALLOCATE (pgf_blk_sizes(blk_count))
508 6752 : pgf_blk_sizes(:) = pgf_blk_sizes_tmp(:blk_count)
509 1736 : END SUBROUTINE
510 :
511 : ! **************************************************************************************************
512 : !> \brief ...
513 : !> \param sizes ...
514 : !> \param nbatches ...
515 : !> \param starts_array ...
516 : !> \param ends_array ...
517 : !> \param starts_array_block ...
518 : !> \param ends_array_block ...
519 : ! **************************************************************************************************
520 1440 : SUBROUTINE create_tensor_batches(sizes, nbatches, starts_array, ends_array, &
521 : starts_array_block, ends_array_block)
522 : INTEGER, DIMENSION(:), INTENT(IN) :: sizes
523 : INTEGER, INTENT(INOUT) :: nbatches
524 : INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: starts_array, ends_array, &
525 : starts_array_block, ends_array_block
526 :
527 : INTEGER :: bsum, imem, nblocks
528 :
529 1440 : nblocks = SIZE(sizes)
530 :
531 1440 : CALL contiguous_tensor_dist(nblocks, nbatches, sizes, limits_start=starts_array_block, limits_end=ends_array_block)
532 :
533 4320 : ALLOCATE (starts_array(nbatches))
534 4320 : ALLOCATE (ends_array(nbatches))
535 :
536 1440 : bsum = 0
537 5142 : DO imem = 1, nbatches
538 3702 : starts_array(imem) = bsum + 1
539 10976 : bsum = bsum + SUM(sizes(starts_array_block(imem):ends_array_block(imem)))
540 5142 : ends_array(imem) = bsum
541 : END DO
542 1440 : END SUBROUTINE
543 :
544 0 : END MODULE
|
