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 Redistribution of DBCSR blocks and distributed matrix slices.
10 : ! **************************************************************************************************
11 : MODULE cp_dbcsr_redistribution
12 : USE cp_dbcsr_api, ONLY: &
13 : dbcsr_clear, dbcsr_filter, dbcsr_finalize, dbcsr_get_block_p, dbcsr_get_info, &
14 : dbcsr_get_readonly_block_p, dbcsr_get_stored_coordinates, dbcsr_iterator_blocks_left, &
15 : dbcsr_iterator_next_block, dbcsr_iterator_readonly_start, dbcsr_iterator_start, &
16 : dbcsr_iterator_stop, dbcsr_iterator_type, dbcsr_p_type, dbcsr_reserve_blocks, dbcsr_set, &
17 : dbcsr_type, dbcsr_type_no_symmetry
18 : USE cp_dbcsr_contrib, ONLY: dbcsr_reserve_all_blocks
19 : USE cp_dbcsr_operations, ONLY: cp_dbcsr_m_by_n_from_template
20 : USE kinds, ONLY: dp,&
21 : int_8
22 : USE matrix_communication, ONLY: allocate_matrix_buffers,&
23 : exchange_matrix_buffers,&
24 : matrix_buffer_type
25 : USE message_passing, ONLY: mp_comm_type,&
26 : mp_para_env_type
27 : USE util, ONLY: sort
28 : #include "./base/base_uses.f90"
29 :
30 : IMPLICIT NONE
31 : PRIVATE
32 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'cp_dbcsr_redistribution'
33 : PUBLIC :: cp_dbcsr_to_subgroups, cp_dbcsr_from_subgroups, cp_array_to_dbcsr_slices, cp_array_to_dbcsr_blocks
34 :
35 : CONTAINS
36 :
37 : ! **************************************************************************************************
38 : !> \brief ...
39 : !> \param mat_global ...
40 : !> \param mat_local ...
41 : !> \param para_env ...
42 : !> \param num_pe_sub ...
43 : !> \param block_row_ranges ...
44 : ! **************************************************************************************************
45 12744 : SUBROUTINE cp_dbcsr_to_subgroups(mat_global, mat_local, para_env, num_pe_sub, block_row_ranges)
46 : TYPE(dbcsr_type) :: mat_global, mat_local
47 : TYPE(mp_para_env_type), INTENT(IN) :: para_env
48 : INTEGER :: num_pe_sub
49 : INTEGER, DIMENSION(:, :), OPTIONAL :: block_row_ranges
50 :
51 : CHARACTER(LEN=*), PARAMETER :: routineN = 'cp_dbcsr_to_subgroups'
52 :
53 : INTEGER :: block_offset, block_size, col, col_size, handle, handle2, i_block, igroup, imep, &
54 : imep_sub, msg_offset, ngroup, offset, row, row_size
55 : INTEGER, ALLOCATABLE :: counts_rec(:, :), counts_send(:, :)
56 12744 : INTEGER, ALLOCATABLE, DIMENSION(:) :: blk_counter, entry_counter
57 : LOGICAL :: found
58 12744 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: data_block
59 : TYPE(dbcsr_iterator_type) :: iter
60 12744 : TYPE(matrix_buffer_type) :: buffer_rec, buffer_send
61 :
62 12744 : CALL timeset(routineN, handle)
63 :
64 12744 : CALL timeset(routineN//"_count", handle2)
65 :
66 12744 : NULLIFY (data_block)
67 :
68 50976 : ALLOCATE (counts_send(2, 0:para_env%num_pe - 1), counts_rec(2, 0:para_env%num_pe - 1))
69 12744 : counts_send = 0
70 :
71 12744 : ngroup = para_env%num_pe/num_pe_sub
72 :
73 12744 : CALL dbcsr_iterator_start(iter, mat_global)
74 41327 : DO WHILE (dbcsr_iterator_blocks_left(iter))
75 :
76 : CALL dbcsr_iterator_next_block(iter, row, col, data_block, &
77 28583 : row_size=row_size, col_size=col_size)
78 :
79 28583 : CALL dbcsr_get_stored_coordinates(mat_local, row, col, imep_sub)
80 :
81 80413 : DO igroup = 0, ngroup - 1
82 :
83 39086 : IF (PRESENT(block_row_ranges)) THEN
84 18032 : IF (row < block_row_ranges(1, igroup + 1) .OR. row > block_row_ranges(2, igroup + 1)) CYCLE
85 : END IF
86 39086 : imep = imep_sub + igroup*num_pe_sub
87 :
88 39086 : counts_send(1, imep) = counts_send(1, imep) + row_size*col_size
89 67669 : counts_send(2, imep) = counts_send(2, imep) + 1
90 :
91 : END DO
92 :
93 : END DO
94 :
95 12744 : CALL dbcsr_iterator_stop(iter)
96 :
97 12744 : CALL timestop(handle2)
98 :
99 12744 : CALL timeset(routineN//"_exchange_counts", handle2)
100 :
101 12744 : IF (para_env%num_pe > 1) THEN
102 12744 : CALL para_env%alltoall(counts_send, counts_rec, 2)
103 : ELSE
104 0 : counts_rec(:, :) = counts_send
105 : END IF
106 :
107 12744 : CALL timestop(handle2)
108 :
109 12744 : CALL timeset(routineN//"_pack", handle2)
110 :
111 : CALL allocate_matrix_buffers(counts_send(1, :), counts_rec(1, :), 3, buffer_send, buffer_rec, &
112 12744 : counts_send(2, :), counts_rec(2, :))
113 :
114 38232 : ALLOCATE (entry_counter(0:para_env%num_pe - 1))
115 38232 : entry_counter(:) = buffer_send%msg_offset(:para_env%num_pe - 1)
116 :
117 25488 : ALLOCATE (blk_counter(0:para_env%num_pe - 1))
118 38232 : blk_counter(:) = buffer_send%indx_offset(:para_env%num_pe - 1)
119 :
120 12744 : CALL dbcsr_iterator_start(iter, mat_global)
121 41327 : DO WHILE (dbcsr_iterator_blocks_left(iter))
122 :
123 : CALL dbcsr_iterator_next_block(iter, row, col, data_block, &
124 28583 : row_size=row_size, col_size=col_size)
125 :
126 28583 : CALL dbcsr_get_stored_coordinates(mat_local, row, col, imep_sub)
127 :
128 80413 : DO igroup = 0, ngroup - 1
129 :
130 39086 : IF (PRESENT(block_row_ranges)) THEN
131 18032 : IF (row < block_row_ranges(1, igroup + 1) .OR. row > block_row_ranges(2, igroup + 1)) CYCLE
132 : END IF
133 :
134 39086 : imep = imep_sub + igroup*num_pe_sub
135 :
136 39086 : msg_offset = entry_counter(imep)
137 :
138 39086 : block_size = row_size*col_size
139 :
140 : buffer_send%msg(msg_offset + 1:msg_offset + block_size) = &
141 2698120 : RESHAPE(data_block(1:row_size, 1:col_size), [block_size])
142 :
143 39086 : entry_counter(imep) = entry_counter(imep) + block_size
144 :
145 39086 : blk_counter(imep) = blk_counter(imep) + 1
146 :
147 39086 : block_offset = blk_counter(imep)
148 :
149 : ! The payload offset in each record is relative to its peer's segment.
150 184927 : buffer_send%indx(:, block_offset) = [row, col, msg_offset - buffer_send%msg_offset(imep)]
151 :
152 : END DO
153 :
154 : END DO
155 :
156 12744 : CALL dbcsr_iterator_stop(iter)
157 :
158 12744 : CALL timestop(handle2)
159 :
160 12744 : CALL timeset(routineN//"_exchange_data", handle2)
161 :
162 12744 : CALL exchange_matrix_buffers(para_env, buffer_rec, buffer_send)
163 :
164 12744 : CALL timestop(handle2)
165 :
166 12744 : CALL timeset(routineN//"_reserve", handle2)
167 :
168 12744 : CALL dbcsr_set(mat_local, 0.0_dp)
169 12744 : CALL dbcsr_filter(mat_local, 1.0_dp)
170 12744 : CALL dbcsr_reserve_blocks(mat_local, rows=buffer_rec%indx(1, :), cols=buffer_rec%indx(2, :))
171 12744 : CALL dbcsr_finalize(mat_local)
172 12744 : CALL dbcsr_set(mat_local, 0.0_dp)
173 :
174 12744 : CALL timestop(handle2)
175 :
176 12744 : CALL timeset(routineN//"_unpack", handle2)
177 :
178 38232 : DO imep = 0, para_env%num_pe - 1
179 77318 : DO i_block = buffer_rec%indx_offset(imep) + 1, buffer_rec%indx_offset(imep + 1)
180 39086 : row = buffer_rec%indx(1, i_block)
181 39086 : col = buffer_rec%indx(2, i_block)
182 39086 : offset = buffer_rec%msg_offset(imep) + buffer_rec%indx(3, i_block)
183 : CALL dbcsr_get_block_p(mat_local, row, col, data_block, found=found, &
184 39086 : row_size=row_size, col_size=col_size)
185 39086 : CPASSERT(found)
186 : data_block(:, :) = RESHAPE(buffer_rec%msg(offset + 1:offset + row_size*col_size), &
187 181832 : [row_size, col_size])
188 : END DO
189 : END DO
190 :
191 12744 : CALL timestop(handle2)
192 :
193 12744 : CALL timestop(handle)
194 :
195 114696 : END SUBROUTINE cp_dbcsr_to_subgroups
196 :
197 : ! **************************************************************************************************
198 : !> \brief Sum subgroup blocks into the global layout, or copy mapped blocks.
199 : !> \param mat_local ...
200 : !> \param mat_global ...
201 : !> \param para_env communicator covering all source and destination ranks
202 : !> \param row_mapping optional source-to-destination block-row permutation
203 : !> \param copy_only copy without summation, filtering, or consuming the source (default: false)
204 : ! **************************************************************************************************
205 7384 : SUBROUTINE cp_dbcsr_from_subgroups(mat_local, mat_global, para_env, row_mapping, copy_only)
206 :
207 : TYPE(dbcsr_type) :: mat_local, mat_global
208 : TYPE(mp_para_env_type), INTENT(IN) :: para_env
209 : INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL :: row_mapping
210 : LOGICAL, INTENT(IN), OPTIONAL :: copy_only
211 :
212 : CHARACTER(LEN=*), PARAMETER :: routineN = 'cp_dbcsr_from_subgroups'
213 :
214 : INTEGER :: block_size, col, col_size, handle, &
215 : handle2, i_block, imep, mapped_row, &
216 : offset, row, row_size
217 : INTEGER, ALLOCATABLE :: counts_rec(:, :), counts_send(:, :)
218 7384 : INTEGER, ALLOCATABLE, DIMENSION(:) :: block_counter, entry_counter
219 : LOGICAL :: found, my_copy_only
220 7384 : REAL(KIND=dp), DIMENSION(:, :), POINTER :: data_block
221 : TYPE(dbcsr_iterator_type) :: iter
222 7384 : TYPE(matrix_buffer_type) :: buffer_rec, buffer_send
223 :
224 7384 : CALL timeset(routineN, handle)
225 7384 : my_copy_only = .FALSE.
226 7384 : IF (PRESENT(copy_only)) my_copy_only = copy_only
227 :
228 7384 : CALL timeset(routineN//"_prepare", handle2)
229 :
230 7384 : IF (.NOT. my_copy_only) THEN
231 7382 : CALL dbcsr_clear(mat_global)
232 7382 : CALL dbcsr_reserve_all_blocks(mat_global)
233 : END IF
234 :
235 29536 : ALLOCATE (counts_send(2, 0:para_env%num_pe - 1), counts_rec(2, 0:para_env%num_pe - 1))
236 7384 : counts_send = 0
237 :
238 7384 : CALL dbcsr_iterator_readonly_start(iter, mat_local)
239 22379 : DO WHILE (dbcsr_iterator_blocks_left(iter))
240 :
241 : CALL dbcsr_iterator_next_block(iter, row, col, data_block, &
242 14995 : row_size=row_size, col_size=col_size)
243 :
244 14995 : mapped_row = row
245 14995 : IF (PRESENT(row_mapping)) mapped_row = row_mapping(row)
246 14995 : CALL dbcsr_get_stored_coordinates(mat_global, mapped_row, col, imep)
247 :
248 14995 : counts_send(1, imep) = counts_send(1, imep) + row_size*col_size
249 14995 : counts_send(2, imep) = counts_send(2, imep) + 1
250 :
251 : END DO
252 :
253 7384 : CALL dbcsr_iterator_stop(iter)
254 :
255 7384 : CALL timestop(handle2)
256 :
257 7384 : CALL timeset(routineN//"_exchange_counts", handle2)
258 :
259 7384 : IF (para_env%num_pe > 1) THEN
260 7384 : CALL para_env%alltoall(counts_send, counts_rec, 2)
261 : ELSE
262 0 : counts_rec(:, :) = counts_send
263 : END IF
264 :
265 7384 : CALL timestop(handle2)
266 :
267 7384 : CALL timeset(routineN//"_pack", handle2)
268 :
269 : CALL allocate_matrix_buffers(counts_send(1, :), counts_rec(1, :), 3, buffer_send, buffer_rec, &
270 7384 : counts_send(2, :), counts_rec(2, :))
271 :
272 22152 : ALLOCATE (block_counter(0:para_env%num_pe - 1))
273 22152 : block_counter(:) = buffer_send%indx_offset(:para_env%num_pe - 1)
274 :
275 14768 : ALLOCATE (entry_counter(0:para_env%num_pe - 1))
276 22152 : entry_counter(:) = buffer_send%msg_offset(:para_env%num_pe - 1)
277 :
278 : ! fill buffer_send
279 7384 : CALL dbcsr_iterator_readonly_start(iter, mat_local)
280 22379 : DO WHILE (dbcsr_iterator_blocks_left(iter))
281 :
282 : CALL dbcsr_iterator_next_block(iter, row, col, data_block, &
283 14995 : row_size=row_size, col_size=col_size)
284 :
285 14995 : mapped_row = row
286 14995 : IF (PRESENT(row_mapping)) mapped_row = row_mapping(row)
287 14995 : CALL dbcsr_get_stored_coordinates(mat_global, mapped_row, col, imep)
288 :
289 14995 : block_size = row_size*col_size
290 :
291 14995 : offset = entry_counter(imep)
292 :
293 : buffer_send%msg(offset + 1:offset + block_size) = &
294 1832702 : RESHAPE(data_block(1:row_size, 1:col_size), [block_size])
295 :
296 14995 : i_block = block_counter(imep) + 1
297 :
298 : ! The receiver adds its own segment start to this peer-relative offset.
299 59980 : buffer_send%indx(:, i_block) = [mapped_row, col, offset - buffer_send%msg_offset(imep)]
300 :
301 14995 : entry_counter(imep) = entry_counter(imep) + block_size
302 :
303 14995 : block_counter(imep) = block_counter(imep) + 1
304 :
305 : END DO
306 :
307 7384 : CALL dbcsr_iterator_stop(iter)
308 :
309 7384 : CALL timestop(handle2)
310 :
311 7384 : CALL timeset(routineN//"_exchange_data", handle2)
312 :
313 7384 : CALL exchange_matrix_buffers(para_env, buffer_rec, buffer_send)
314 :
315 7384 : CALL timestop(handle2)
316 :
317 7384 : CALL timeset(routineN//"_unpack", handle2)
318 :
319 : ! Keep the rank and receive-buffer order of the block-wise accumulation.
320 22152 : DO imep = 0, para_env%num_pe - 1
321 37147 : DO i_block = buffer_rec%indx_offset(imep) + 1, buffer_rec%indx_offset(imep + 1)
322 14995 : row = buffer_rec%indx(1, i_block)
323 14995 : col = buffer_rec%indx(2, i_block)
324 14995 : offset = buffer_rec%msg_offset(imep) + buffer_rec%indx(3, i_block)
325 : CALL dbcsr_get_block_p(mat_global, row, col, data_block, found=found, &
326 14995 : row_size=row_size, col_size=col_size)
327 14995 : CPASSERT(found)
328 44758 : IF (my_copy_only) THEN
329 : data_block(:, :) = RESHAPE(buffer_rec%msg(offset + 1:offset + row_size*col_size), &
330 24 : [row_size, col_size])
331 : ELSE
332 : data_block(:, :) = data_block + &
333 : RESHAPE(buffer_rec%msg(offset + 1:offset + row_size*col_size), &
334 1966126 : [row_size, col_size])
335 : END IF
336 : END DO
337 : END DO
338 :
339 7384 : IF (.NOT. my_copy_only) THEN
340 : ! Remove numerically zero blocks from the summed matrix.
341 7382 : CALL dbcsr_filter(mat_global, 1.0E-30_dp)
342 : END IF
343 :
344 : IF (.NOT. my_copy_only) THEN
345 7382 : CALL dbcsr_set(mat_local, 0.0_dp)
346 7382 : CALL dbcsr_filter(mat_local, 1.0_dp)
347 : END IF
348 :
349 7384 : CALL timestop(handle2)
350 :
351 7384 : CALL timestop(handle)
352 :
353 66456 : END SUBROUTINE cp_dbcsr_from_subgroups
354 :
355 : ! **************************************************************************************************
356 : !> \brief Consume a distributed flattened array and create a DBCSR matrix per slice.
357 : !> \param array local flattened rows; all ranks have the same number of slices; deallocated on return
358 : !> \param nrows number of rows in each output matrix
359 : !> \param ncols number of columns in each output matrix
360 : !> \param para_env_sub communicator for the output matrices
361 : !> \param first_row first flattened row held locally
362 : !> \param last_row last flattened row held locally
363 : !> \param template DBCSR template defining output distribution and block sizes
364 : !> \param matrices preallocated array of output matrices
365 : !> \param matrix_ptrs preallocated array of output pointers; pointees are allocated here
366 : !> \note Supply exactly one output representation. Flattened rows use (i-1)*ncols+j.
367 : ! **************************************************************************************************
368 360 : SUBROUTINE cp_array_to_dbcsr_slices(array, nrows, ncols, para_env_sub, &
369 360 : first_row, last_row, template, matrices, matrix_ptrs)
370 : REAL(KIND=dp), ALLOCATABLE, INTENT(INOUT), TARGET :: array(:, :)
371 : INTEGER, INTENT(IN) :: nrows, ncols
372 : TYPE(mp_para_env_type), INTENT(IN), POINTER :: para_env_sub
373 : INTEGER, INTENT(IN) :: first_row, last_row
374 : TYPE(dbcsr_type), INTENT(INOUT) :: template
375 : TYPE(dbcsr_type), INTENT(INOUT), OPTIONAL, TARGET :: matrices(:)
376 : TYPE(dbcsr_p_type), INTENT(INOUT), OPTIONAL :: matrix_ptrs(:)
377 :
378 : CHARACTER(LEN=*), PARAMETER :: routineN = 'cp_array_to_dbcsr_slices'
379 :
380 : INTEGER :: handle, i, islice, j, nblock, nslice
381 360 : INTEGER, ALLOCATABLE :: blocks(:, :)
382 360 : INTEGER, POINTER :: local_cols(:), local_rows(:)
383 360 : REAL(KIND=dp), POINTER :: values(:, :, :)
384 360 : TYPE(dbcsr_p_type), ALLOCATABLE :: outputs(:)
385 :
386 360 : CALL timeset(routineN, handle)
387 360 : CPASSERT(PRESENT(matrices) .NEQV. PRESENT(matrix_ptrs))
388 360 : nslice = SIZE(array, 2)
389 17003 : ALLOCATE (outputs(nslice))
390 16283 : DO islice = 1, nslice
391 15923 : IF (PRESENT(matrices)) THEN
392 83 : outputs(islice)%matrix => matrices(islice)
393 : ELSE
394 15840 : ALLOCATE (matrix_ptrs(islice)%matrix)
395 15840 : outputs(islice)%matrix => matrix_ptrs(islice)%matrix
396 : END IF
397 : CALL cp_dbcsr_m_by_n_from_template(outputs(islice)%matrix, template=template, &
398 15923 : m=nrows, n=ncols, sym=dbcsr_type_no_symmetry)
399 15923 : CALL dbcsr_reserve_all_blocks(outputs(islice)%matrix)
400 16283 : CALL dbcsr_finalize(outputs(islice)%matrix)
401 : END DO
402 360 : IF (nslice > 0) THEN
403 360 : CALL dbcsr_get_info(outputs(1)%matrix, local_rows=local_rows, local_cols=local_cols)
404 1072 : ALLOCATE (blocks(3, SIZE(local_rows)*SIZE(local_cols)*nslice))
405 360 : nblock = 0
406 16283 : DO islice = 1, nslice
407 32206 : DO j = 1, SIZE(local_cols)
408 47051 : DO i = 1, SIZE(local_rows)
409 15205 : nblock = nblock + 1
410 76743 : blocks(:, nblock) = [local_rows(i), local_cols(j), islice]
411 : END DO
412 : END DO
413 : END DO
414 : ! A rank-remapped view preserves the row-major flattened ia convention without copying.
415 360 : values(1:SIZE(array, 1), 1:1, 1:nslice) => array
416 : CALL cp_array_to_dbcsr_blocks(values, first_row, last_row, para_env_sub, outputs, blocks, &
417 360 : row_major=.TRUE.)
418 : END IF
419 360 : DEALLOCATE (array)
420 360 : CALL timestop(handle)
421 720 : END SUBROUTINE cp_array_to_dbcsr_slices
422 :
423 : ! **************************************************************************************************
424 : !> \brief Fill existing DBCSR blocks directly from distributed flattened array slices.
425 : !> Requests split contiguous block fragments at source-rank boundaries. Sources combine
426 : !> components before replying, so each target element transfers only one real value.
427 : !> Slice tiles bound request/reply storage; existing sparsity is preserved.
428 : !> \param array local flattened entries, components, slices; borrowed, never deallocated
429 : !> \param first_entry first flattened entry held locally
430 : !> \param last_entry last flattened entry held locally; may precede first_entry on empty ranks
431 : !> \param comm communicator covering sources and output matrices, in any source-range order
432 : !> \param matrices existing output slices with identical dimensions and block sizes
433 : !> \param blocks requested (block row, block column, slice) columns; absent local blocks are skipped
434 : !> \param factors optional integer component multiplicities per requested block; default one component
435 : !> \param row_major flatten with column index fastest; default Fortran column-major ordering
436 : !> \note Source ranges must partition the whole flattened matrix. Requested blocks must be unique.
437 : !> Unrequested output entries are zeroed. The 64 MiB workspace target includes the borrowed
438 : !> array and descriptors; an irreducible single-slice exchange may exceed this target.
439 : ! **************************************************************************************************
440 360 : SUBROUTINE cp_array_to_dbcsr_blocks(array, first_entry, last_entry, comm, matrices, blocks, factors, row_major)
441 : REAL(KIND=dp), INTENT(IN) :: array(:, :, :)
442 : INTEGER, INTENT(IN) :: first_entry, last_entry, blocks(:, :)
443 :
444 : CLASS(mp_comm_type), INTENT(IN) :: comm
445 : TYPE(dbcsr_p_type), INTENT(INOUT) :: matrices(:)
446 : INTEGER, INTENT(IN), OPTIONAL :: factors(:, :)
447 : LOGICAL, INTENT(IN), OPTIONAL :: row_major
448 :
449 : CHARACTER(LEN=*), PARAMETER :: routineN = 'cp_array_to_dbcsr_blocks'
450 :
451 : INTEGER(KIND=int_8), PARAMETER :: max_bytes = 64_int_8*1024**2
452 :
453 : INTEGER :: a, b, col, component, element, first, handle, i, image, last, length, lo, &
454 : lower, middle, nc, nfast, nprocs, pass, peer, pos, rank, record, row, upper, width
455 : INTEGER(KIND=int_8) :: fixed_bytes, nel
456 360 : INTEGER, ALLOCATABLE :: counter(:), counts_rec(:, :), &
457 360 : counts_send(:, :), ends(:), order(:), &
458 360 : ranges(:, :), targets(:, :)
459 360 : INTEGER, POINTER :: fast_offset(:), fast_size(:), &
460 360 : slow_offset(:), slow_size(:)
461 : LOGICAL :: found, row_order
462 360 : REAL(KIND=dp), POINTER :: block(:, :)
463 360 : TYPE(matrix_buffer_type) :: requests_rec, requests_send
464 :
465 360 : CALL timeset(routineN, handle)
466 360 : IF (SIZE(matrices) == 0) THEN
467 0 : CALL timestop(handle)
468 : RETURN
469 : END IF
470 360 : row_order = .FALSE.
471 360 : IF (PRESENT(row_major)) row_order = row_major
472 : CALL dbcsr_get_info(matrices(1)%matrix, row_blk_size=fast_size, col_blk_size=slow_size, &
473 360 : row_blk_offset=fast_offset, col_blk_offset=slow_offset)
474 360 : IF (row_order) THEN
475 : CALL dbcsr_get_info(matrices(1)%matrix, row_blk_size=slow_size, col_blk_size=fast_size, &
476 360 : row_blk_offset=slow_offset, col_blk_offset=fast_offset)
477 : END IF
478 720 : nfast = SUM(fast_size)
479 720 : nel = INT(nfast, int_8)*SUM(slow_size)
480 360 : nc = SIZE(array, 2)
481 360 : CPASSERT(nel > 0 .AND. nel < HUGE(element))
482 360 : CPASSERT(nc >= 1)
483 360 : CPASSERT(SIZE(array, 1) == MAX(0, last_entry - first_entry + 1))
484 360 : CPASSERT(SIZE(array, 3) == SIZE(matrices) .AND. SIZE(blocks, 1) == 3)
485 360 : IF (PRESENT(factors)) THEN
486 0 : CPASSERT(SIZE(factors, 1) == nc .AND. SIZE(factors, 2) == SIZE(blocks, 2))
487 : ELSE
488 360 : CPASSERT(nc == 1)
489 : END IF
490 360 : nprocs = comm%num_pe
491 : ALLOCATE (ends(nprocs), order(nprocs), ranges(2, nprocs), counts_send(2, 0:nprocs - 1), &
492 3960 : counts_rec(2, 0:nprocs - 1), counter(0:nprocs - 1))
493 1080 : CALL comm%allgather([first_entry, last_entry], ranges)
494 736 : ends(:) = ranges(2, :)
495 736 : WHERE (ranges(2, :) < ranges(1, :)) ends = HUGE(element)
496 360 : CALL sort(ends, nprocs, order)
497 : ! Validate coverage, independently of communicator rank ordering and empty ranges.
498 360 : element = 0
499 736 : DO rank = 1, nprocs
500 376 : IF (ends(rank) == HUGE(element)) EXIT
501 376 : CPASSERT(ranges(1, order(rank)) == element + 1)
502 736 : element = ends(rank)
503 : END DO
504 360 : CPASSERT(INT(element, int_8) == nel)
505 2160 : fixed_bytes = 8_int_8*SIZE(array, KIND=int_8) + 4_int_8*SIZE(blocks, KIND=int_8) + 64_int_8*nprocs
506 360 : IF (PRESENT(factors)) fixed_bytes = fixed_bytes + 4_int_8*SIZE(factors, KIND=int_8)
507 : width = INT(MIN(INT(SIZE(matrices), int_8), &
508 360 : MAX(1_int_8, (max_bytes - fixed_bytes)/((80_int_8 + 8*nc)*nel))))
509 360 : CALL comm%min(width)
510 16283 : DO image = 1, SIZE(matrices)
511 16283 : CALL dbcsr_set(matrices(image)%matrix, 0.0_dp)
512 : END DO
513 360 : DO first = 1, SIZE(matrices), width
514 360 : last = MIN(SIZE(matrices), first + width - 1)
515 360 : counts_send(:, :) = 0
516 1080 : DO pass = 1, 2
517 720 : IF (pass == 2) THEN
518 360 : CALL comm%alltoall(counts_send, counts_rec, 2)
519 : CALL allocate_matrix_buffers(counts_send(1, :), counts_rec(1, :), 3 + nc, requests_send, requests_rec, &
520 360 : send_blocks=counts_send(2, :), recv_blocks=counts_rec(2, :))
521 1072 : ALLOCATE (targets(3, SIZE(requests_send%indx, 2)))
522 736 : counter(:) = requests_send%indx_offset(:nprocs - 1)
523 : END IF
524 31490 : DO i = 1, SIZE(blocks, 2)
525 30410 : image = blocks(3, i)
526 30410 : IF (image < first .OR. image > last) CYCLE
527 30410 : CALL dbcsr_get_readonly_block_p(matrices(image)%matrix, blocks(1, i), blocks(2, i), block, found=found)
528 30410 : IF (.NOT. found) CYCLE
529 30410 : a = blocks(1, i)
530 30410 : b = blocks(2, i)
531 30410 : IF (row_order) THEN
532 30410 : a = blocks(2, i)
533 30410 : b = blocks(1, i)
534 : END IF
535 174636 : DO col = 1, slow_size(b)
536 : row = 1
537 257042 : DO WHILE (row <= fast_size(a))
538 113536 : element = (slow_offset(b) + col - 2)*nfast + fast_offset(a) + row - 1
539 : ! Find the first nonempty source range whose upper bound covers this entry.
540 113536 : lower = 1
541 113536 : upper = nprocs
542 120600 : DO WHILE (lower < upper)
543 7064 : middle = (lower + upper)/2
544 120600 : IF (ends(middle) < element) THEN
545 3532 : lower = middle + 1
546 : ELSE
547 : upper = middle
548 : END IF
549 : END DO
550 113536 : peer = order(lower) - 1
551 113536 : length = MIN(fast_size(a) - row + 1, ends(lower) - element + 1)
552 113536 : IF (pass == 1) THEN
553 170304 : counts_send(:, peer) = counts_send(:, peer) + [length, 1]
554 : ELSE
555 56768 : counter(peer) = counter(peer) + 1
556 56768 : record = counter(peer)
557 227072 : requests_send%indx(1:3, record) = [element, length, image]
558 113536 : requests_send%indx(4:, record) = 1
559 56768 : IF (PRESENT(factors)) requests_send%indx(4:, record) = factors(:, i)
560 227072 : targets(:, record) = [i, row, col]
561 : END IF
562 113536 : row = row + length
563 : END DO
564 : END DO
565 : END DO
566 : END DO
567 360 : CALL exchange_matrix_buffers(comm, requests_rec, requests_send, do_msg=.FALSE.)
568 736 : DO peer = 0, nprocs - 1
569 376 : pos = requests_rec%msg_offset(peer)
570 57504 : DO record = requests_rec%indx_offset(peer) + 1, requests_rec%indx_offset(peer + 1)
571 56768 : lo = requests_rec%indx(1, record) - first_entry + 1
572 56768 : length = requests_rec%indx(2, record)
573 56768 : image = requests_rec%indx(3, record)
574 56768 : CPASSERT(lo >= 1 .AND. lo + length - 1 <= SIZE(array, 1))
575 : requests_rec%msg(pos + 1:pos + length) = &
576 948940 : REAL(requests_rec%indx(4, record), dp)*array(lo:lo + length - 1, 1, image)
577 56768 : DO component = 2, nc
578 : requests_rec%msg(pos + 1:pos + length) = requests_rec%msg(pos + 1:pos + length) + &
579 56768 : REAL(requests_rec%indx(3 + component, record), dp)*array(lo:lo + length - 1, component, image)
580 : END DO
581 57144 : pos = pos + length
582 : END DO
583 : END DO
584 360 : CALL exchange_matrix_buffers(comm, requests_send, requests_rec, do_indx=.FALSE.)
585 736 : DO peer = 0, nprocs - 1
586 376 : pos = requests_send%msg_offset(peer)
587 57504 : DO record = requests_send%indx_offset(peer) + 1, requests_send%indx_offset(peer + 1)
588 56768 : i = targets(1, record)
589 56768 : row = targets(2, record)
590 56768 : col = targets(3, record)
591 56768 : length = requests_send%indx(2, record)
592 56768 : CALL dbcsr_get_block_p(matrices(blocks(3, i))%matrix, blocks(1, i), blocks(2, i), block, found=found)
593 56768 : CPASSERT(found)
594 56768 : IF (row_order) THEN
595 948940 : block(col, row:row + length - 1) = requests_send%msg(pos + 1:pos + length)
596 : ELSE
597 0 : block(row:row + length - 1, col) = requests_send%msg(pos + 1:pos + length)
598 : END IF
599 113912 : pos = pos + length
600 : END DO
601 : END DO
602 360 : DEALLOCATE (targets)
603 : END DO
604 360 : CALL timestop(handle)
605 720 : END SUBROUTINE cp_array_to_dbcsr_blocks
606 :
607 : END MODULE cp_dbcsr_redistribution
|
