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 full matrices and distributed arrays.
10 : ! **************************************************************************************************
11 : MODULE cp_fm_redistribution
12 : USE cp_fm_struct, ONLY: cp_fm_struct_type
13 : USE cp_fm_types, ONLY: cp_fm_create,&
14 : cp_fm_get_info,&
15 : cp_fm_release,&
16 : cp_fm_type
17 : USE kinds, ONLY: dp
18 : USE matrix_communication, ONLY: allocate_matrix_buffers,&
19 : exchange_matrix_buffers,&
20 : matrix_buffer_type
21 : USE message_passing, ONLY: mp_para_env_type
22 : #include "../base/base_uses.f90"
23 :
24 : IMPLICIT NONE
25 : PRIVATE
26 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'cp_fm_redistribution'
27 : PUBLIC :: cp_array_to_fm, cp_fm_to_array, cp_array_columns_to_fm, &
28 : cp_fm_permute, cp_fm_to_replicated_array
29 :
30 : CONTAINS
31 :
32 : ! **************************************************************************************************
33 : !> \brief Add a permuted or restricted four-index tensor stored as a full matrix.
34 : !> \param fm_out destination; entries outside the mapped region remain unchanged
35 : !> \param fm_in source; each matrix coordinate encodes two tensor indices
36 : !> \param beta factor multiplying the source contribution
37 : !> \param nrow_secidx_in size of the second tensor index in source rows
38 : !> \param ncol_secidx_in size of the second tensor index in source columns
39 : !> \param nrow_secidx_out stride of the first tensor index in destination rows
40 : !> \param ncol_secidx_out stride of the first tensor index in destination columns
41 : !> \param reordering permutation of the four source indices
42 : !> \param row_offset zero-based destination row offset
43 : !> \param col_offset zero-based destination column offset
44 : !> \param first_index first retained value of each source index (default: one)
45 : !> \param last_index last retained value of each source index (default: full extent)
46 : !> \note The index map must be one-to-one and fit within the destination.
47 : ! **************************************************************************************************
48 814 : SUBROUTINE cp_fm_permute(fm_out, fm_in, beta, nrow_secidx_in, ncol_secidx_in, &
49 : nrow_secidx_out, ncol_secidx_out, reordering, row_offset, col_offset, &
50 : first_index, last_index)
51 : TYPE(cp_fm_type), INTENT(INOUT) :: fm_out
52 : TYPE(cp_fm_type), INTENT(IN) :: fm_in
53 : REAL(KIND=dp), INTENT(IN) :: beta
54 : INTEGER, INTENT(IN) :: nrow_secidx_in, ncol_secidx_in, &
55 : nrow_secidx_out, ncol_secidx_out
56 : INTEGER, DIMENSION(4), INTENT(IN) :: reordering
57 : INTEGER, INTENT(IN), OPTIONAL :: row_offset, col_offset
58 : INTEGER, DIMENSION(4), INTENT(IN), OPTIONAL :: first_index, last_index
59 :
60 : CHARACTER(LEN=*), PARAMETER :: routineN = 'cp_fm_permute'
61 :
62 : INTEGER :: handle, i, idx_col, idx_row, ientry, ii, &
63 : iproc, j, jj, my_col_offset, &
64 : my_row_offset, ncol_local, nprocs, &
65 : nrow_local, pass
66 814 : INTEGER, ALLOCATABLE, DIMENSION(:) :: entry_counter, num_entries_rec, &
67 : num_entries_send
68 : INTEGER, DIMENSION(4) :: first, indices, last, mapped
69 814 : INTEGER, DIMENSION(:), POINTER :: col_indices, row_indices
70 814 : TYPE(matrix_buffer_type) :: buffer_rec, buffer_send
71 : TYPE(mp_para_env_type), POINTER :: para_env
72 :
73 814 : CALL timeset(routineN, handle)
74 814 : para_env => fm_out%matrix_struct%para_env
75 814 : nprocs = para_env%num_pe
76 814 : my_row_offset = 0
77 814 : my_col_offset = 0
78 814 : IF (PRESENT(row_offset)) my_row_offset = row_offset
79 814 : IF (PRESENT(col_offset)) my_col_offset = col_offset
80 4070 : first = 1
81 : last = [fm_in%matrix_struct%nrow_global/nrow_secidx_in, nrow_secidx_in, &
82 4070 : fm_in%matrix_struct%ncol_global/ncol_secidx_in, ncol_secidx_in]
83 814 : IF (PRESENT(first_index)) first = first_index
84 814 : IF (PRESENT(last_index)) last = last_index
85 :
86 : CALL cp_fm_get_info(fm_in, nrow_local=nrow_local, ncol_local=ncol_local, &
87 814 : row_indices=row_indices, col_indices=col_indices)
88 : ALLOCATE (num_entries_send(0:nprocs - 1), num_entries_rec(0:nprocs - 1), &
89 4070 : entry_counter(0:nprocs - 1))
90 814 : num_entries_send = 0
91 :
92 : ! Count and then pack using the same coordinate map.
93 2442 : DO pass = 1, 2
94 1628 : IF (pass == 2) THEN
95 814 : CALL para_env%alltoall(num_entries_send, num_entries_rec, 1)
96 814 : CALL allocate_matrix_buffers(num_entries_send, num_entries_rec, 2, buffer_send, buffer_rec)
97 2434 : entry_counter(:) = buffer_send%msg_offset(:nprocs - 1)
98 : END IF
99 :
100 59962 : DO i = 1, nrow_local
101 57520 : indices(1) = (row_indices(i) - 1)/nrow_secidx_in + 1
102 57520 : indices(2) = MOD(row_indices(i) - 1, nrow_secidx_in) + 1
103 9128056 : DO j = 1, ncol_local
104 9068908 : indices(3) = (col_indices(j) - 1)/ncol_secidx_in + 1
105 9068908 : indices(4) = MOD(col_indices(j) - 1, ncol_secidx_in) + 1
106 90137172 : IF (ANY(indices < first) .OR. ANY(indices > last)) CYCLE
107 43715880 : mapped = indices - first + 1
108 : idx_row = my_row_offset + mapped(reordering(2)) + &
109 8743176 : (mapped(reordering(1)) - 1)*nrow_secidx_out
110 : idx_col = my_col_offset + mapped(reordering(4)) + &
111 8743176 : (mapped(reordering(3)) - 1)*ncol_secidx_out
112 : iproc = fm_out%matrix_struct%context%blacs2mpi( &
113 8743176 : fm_out%matrix_struct%g2p_row(idx_row), fm_out%matrix_struct%g2p_col(idx_col))
114 8800696 : IF (pass == 1) THEN
115 4371588 : num_entries_send(iproc) = num_entries_send(iproc) + 1
116 : ELSE
117 4371588 : entry_counter(iproc) = entry_counter(iproc) + 1
118 4371588 : ientry = entry_counter(iproc)
119 4371588 : buffer_send%msg(ientry) = fm_in%local_data(i, j)
120 13114764 : buffer_send%indx(:, ientry) = [idx_row, idx_col]
121 : END IF
122 : END DO
123 : END DO
124 : END DO
125 :
126 814 : CALL exchange_matrix_buffers(para_env, buffer_rec, buffer_send)
127 :
128 : !$OMP PARALLEL DO DEFAULT(NONE) &
129 : !$OMP SHARED(fm_out, buffer_rec, beta) &
130 814 : !$OMP PRIVATE(ientry, ii, jj)
131 : DO ientry = 1, SIZE(buffer_rec%msg)
132 : ii = fm_out%matrix_struct%g2l_row(buffer_rec%indx(1, ientry))
133 : jj = fm_out%matrix_struct%g2l_col(buffer_rec%indx(2, ientry))
134 : fm_out%local_data(ii, jj) = fm_out%local_data(ii, jj) + beta*buffer_rec%msg(ientry)
135 : END DO
136 : !$OMP END PARALLEL DO
137 814 : CALL timestop(handle)
138 2442 : END SUBROUTINE cp_fm_permute
139 :
140 : ! **************************************************************************************************
141 : !> \brief Assemble a full matrix from distributed rectangular array tiles.
142 : !> \param mat2D local tile; tiles on the FM communicator must not overlap
143 : !> \param fm_struct structure defining the destination layout and communicator
144 : !> \param my_start_row first global row of the tile
145 : !> \param my_end_row last global row of the tile
146 : !> \param my_start_col first global column of the tile
147 : !> \param my_end_col last global column of the tile
148 : !> \param fm_mat destination; entries outside the supplied tiles are zero
149 : !> \param do_release_mat whether to deallocate the input tile (default: true)
150 : ! **************************************************************************************************
151 1454 : SUBROUTINE cp_array_to_fm(mat2D, fm_struct, my_start_row, my_end_row, &
152 : my_start_col, my_end_col, fm_mat, do_release_mat)
153 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), &
154 : INTENT(INOUT) :: mat2D
155 : TYPE(cp_fm_struct_type), POINTER :: fm_struct
156 : INTEGER, INTENT(IN) :: my_start_row, my_end_row, my_start_col, &
157 : my_end_col
158 : TYPE(cp_fm_type), INTENT(OUT) :: fm_mat
159 : LOGICAL, INTENT(IN), OPTIONAL :: do_release_mat
160 :
161 : CHARACTER(LEN=*), PARAMETER :: routineN = 'cp_array_to_fm'
162 :
163 : INTEGER :: handle, i, ientry, j, pass, peer
164 1454 : INTEGER, ALLOCATABLE :: cursor(:), recv_count(:), send_count(:)
165 : LOGICAL :: release_mat
166 1454 : TYPE(matrix_buffer_type) :: buffer_rec, buffer_send
167 : TYPE(mp_para_env_type), POINTER :: para_env
168 :
169 1454 : CALL timeset(routineN, handle)
170 1454 : para_env => fm_struct%para_env
171 1454 : release_mat = .TRUE.
172 1454 : IF (PRESENT(do_release_mat)) release_mat = do_release_mat
173 1454 : CALL cp_fm_create(fm_mat, fm_struct, name="fm_mat", set_zero=.TRUE.)
174 : ALLOCATE (send_count(0:para_env%num_pe - 1), recv_count(0:para_env%num_pe - 1), &
175 7270 : cursor(0:para_env%num_pe - 1))
176 1454 : send_count = 0
177 :
178 4362 : DO pass = 1, 2
179 2908 : IF (pass == 2) THEN
180 1454 : CALL para_env%alltoall(send_count, recv_count, 1)
181 1454 : CALL allocate_matrix_buffers(send_count, recv_count, 2, buffer_send, buffer_rec)
182 4182 : cursor(:) = buffer_send%msg_offset(:para_env%num_pe - 1)
183 : END IF
184 428366 : DO j = my_start_col, my_end_col
185 27834934 : DO i = my_start_row, my_end_row
186 27408022 : peer = fm_struct%context%blacs2mpi(fm_struct%g2p_row(i), fm_struct%g2p_col(j))
187 27832026 : IF (pass == 1) THEN
188 13704011 : send_count(peer) = send_count(peer) + 1
189 : ELSE
190 13704011 : cursor(peer) = cursor(peer) + 1
191 13704011 : ientry = cursor(peer)
192 41112033 : buffer_send%indx(:, ientry) = [i, j]
193 13704011 : buffer_send%msg(ientry) = mat2D(i - my_start_row + 1, j - my_start_col + 1)
194 : END IF
195 : END DO
196 : END DO
197 : END DO
198 1454 : IF (release_mat) DEALLOCATE (mat2D)
199 1454 : CALL exchange_matrix_buffers(para_env, buffer_rec, buffer_send)
200 13705465 : DO ientry = 1, SIZE(buffer_rec%msg)
201 13704011 : i = fm_struct%g2l_row(buffer_rec%indx(1, ientry))
202 13704011 : j = fm_struct%g2l_col(buffer_rec%indx(2, ientry))
203 13705465 : fm_mat%local_data(i, j) = buffer_rec%msg(ientry)
204 : END DO
205 1454 : CALL timestop(handle)
206 2908 : END SUBROUTINE cp_array_to_fm
207 :
208 : ! **************************************************************************************************
209 : !> \brief Extract distributed rectangular array tiles from a full matrix.
210 : !> \param mat2D output tile, allocated here
211 : !> \param my_start_row first global row requested by this rank
212 : !> \param my_end_row last global row requested by this rank
213 : !> \param my_start_col first global column requested by this rank
214 : !> \param my_end_col last global column requested by this rank
215 : !> \param fm_mat source, released on return
216 : ! **************************************************************************************************
217 722 : SUBROUTINE cp_fm_to_array(mat2D, my_start_row, my_end_row, my_start_col, my_end_col, fm_mat)
218 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), &
219 : INTENT(OUT) :: mat2D
220 : INTEGER, INTENT(IN) :: my_start_row, my_end_row, my_start_col, &
221 : my_end_col
222 : TYPE(cp_fm_type), INTENT(INOUT) :: fm_mat
223 :
224 : CHARACTER(LEN=*), PARAMETER :: routineN = 'cp_fm_to_array'
225 :
226 : INTEGER :: handle, i, ientry, j, ncols, nrows, peer
227 : INTEGER, ALLOCATABLE :: bounds(:, :), recv_count(:), &
228 : send_count(:)
229 722 : INTEGER, POINTER :: col_indices(:), row_indices(:)
230 722 : TYPE(matrix_buffer_type) :: buffer_rec, buffer_send
231 : TYPE(mp_para_env_type), POINTER :: para_env
232 :
233 722 : CALL timeset(routineN, handle)
234 : CALL cp_fm_get_info(fm_mat, nrow_local=nrows, ncol_local=ncols, row_indices=row_indices, &
235 722 : col_indices=col_indices, para_env=para_env)
236 2888 : ALLOCATE (mat2D(my_end_row - my_start_row + 1, my_end_col - my_start_col + 1))
237 722 : mat2D = 0.0_dp
238 : ALLOCATE (bounds(4, 0:para_env%num_pe - 1), send_count(0:para_env%num_pe - 1), &
239 4332 : recv_count(0:para_env%num_pe - 1))
240 3610 : CALL para_env%allgather([my_start_row, my_end_row, my_start_col, my_end_col], bounds)
241 :
242 : ! The ordered FM row/column lists turn each rectangle into local index ranges.
243 2166 : DO peer = 0, para_env%num_pe - 1
244 : bounds(:, peer) = [COUNT(row_indices(:nrows) < bounds(1, peer)) + 1, &
245 : COUNT(row_indices(:nrows) <= bounds(2, peer)), &
246 : COUNT(col_indices(:ncols) < bounds(3, peer)) + 1, &
247 352720 : COUNT(col_indices(:ncols) <= bounds(4, peer))]
248 : send_count(peer) = MAX(0, bounds(2, peer) - bounds(1, peer) + 1)* &
249 2166 : MAX(0, bounds(4, peer) - bounds(3, peer) + 1)
250 : END DO
251 722 : CALL para_env%alltoall(send_count, recv_count, 1)
252 722 : CALL allocate_matrix_buffers(send_count, recv_count, 2, buffer_send, buffer_rec)
253 2166 : DO peer = 0, para_env%num_pe - 1
254 1444 : IF (send_count(peer) == 0) CYCLE
255 1444 : ientry = buffer_send%msg_offset(peer)
256 65912 : DO j = bounds(3, peer), bounds(4, peer)
257 2268180 : DO i = bounds(1, peer), bounds(2, peer)
258 2202990 : ientry = ientry + 1
259 6608970 : buffer_send%indx(:, ientry) = [row_indices(i), col_indices(j)]
260 2266736 : buffer_send%msg(ientry) = fm_mat%local_data(i, j)
261 : END DO
262 : END DO
263 : END DO
264 722 : CALL exchange_matrix_buffers(para_env, buffer_rec, buffer_send)
265 2203712 : DO ientry = 1, SIZE(buffer_rec%msg)
266 2202990 : i = buffer_rec%indx(1, ientry) - my_start_row + 1
267 2202990 : j = buffer_rec%indx(2, ientry) - my_start_col + 1
268 2203712 : mat2D(i, j) = buffer_rec%msg(ientry)
269 : END DO
270 722 : CALL cp_fm_release(fm_mat)
271 722 : CALL timestop(handle)
272 1444 : END SUBROUTINE cp_fm_to_array
273 :
274 : ! **************************************************************************************************
275 : !> \brief ...
276 : !> \param fm_S ...
277 : !> \param array_S ...
278 : !> \param weight ...
279 : !> \param add ...
280 : ! **************************************************************************************************
281 872 : SUBROUTINE cp_fm_to_replicated_array(fm_S, array_S, weight, add)
282 :
283 : TYPE(cp_fm_type), DIMENSION(:) :: fm_S
284 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: array_S
285 : REAL(KIND=dp), OPTIONAL :: weight
286 : LOGICAL, OPTIONAL :: add
287 :
288 : CHARACTER(LEN=*), PARAMETER :: routineN = 'cp_fm_to_replicated_array'
289 :
290 : INTEGER :: handle, i, i_row_local, img, j, &
291 : j_col_local, n_basis, ncol_local, &
292 : nimages, nrow_local
293 872 : INTEGER, DIMENSION(:), POINTER :: col_indices, row_indices
294 : LOGICAL :: my_add
295 : REAL(KIND=dp) :: my_weight
296 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: array_tmp
297 :
298 872 : CALL timeset(routineN, handle)
299 :
300 872 : my_weight = 1.0_dp
301 872 : IF (PRESENT(weight)) my_weight = weight
302 :
303 872 : my_add = .FALSE.
304 872 : IF (PRESENT(add)) my_add = add
305 :
306 872 : n_basis = SIZE(array_S, 1)
307 872 : nimages = SIZE(array_S, 3)
308 :
309 : ! checks
310 872 : CPASSERT(SIZE(array_S, 2) == n_basis)
311 872 : CPASSERT(SIZE(fm_S) == nimages)
312 872 : CPASSERT(LBOUND(array_S, 1) == 1)
313 872 : CPASSERT(LBOUND(array_S, 2) == 1)
314 872 : CPASSERT(LBOUND(array_S, 3) == 1)
315 :
316 : CALL cp_fm_get_info(matrix=fm_S(1), &
317 : nrow_local=nrow_local, &
318 : ncol_local=ncol_local, &
319 : row_indices=row_indices, &
320 872 : col_indices=col_indices)
321 :
322 872 : IF (.NOT. my_add) array_S(:, :, :) = 0.0_dp
323 4360 : ALLOCATE (array_tmp(SIZE(array_S, 1), SIZE(array_S, 2), SIZE(array_S, 3)))
324 872 : array_tmp(:, :, :) = 0.0_dp
325 :
326 8720 : DO img = 1, nimages
327 31238 : DO i_row_local = 1, nrow_local
328 22518 : i = row_indices(i_row_local)
329 160344 : DO j_col_local = 1, ncol_local
330 129978 : j = col_indices(j_col_local)
331 152496 : array_tmp(i, j, img) = fm_S(img)%local_data(i_row_local, j_col_local)
332 : END DO ! j_col_local
333 : END DO ! i_row_local
334 : END DO ! img
335 :
336 872 : CALL fm_S(1)%matrix_struct%para_env%sum(array_tmp)
337 :
338 313712 : array_S(:, :, :) = array_S(:, :, :) + my_weight*array_tmp(:, :, :)
339 :
340 872 : CALL timestop(handle)
341 :
342 2616 : END SUBROUTINE cp_fm_to_replicated_array
343 :
344 : ! **************************************************************************************************
345 : !> \brief Redistribute a column-distributed array into an existing full matrix.
346 : !> \param matrix destination; cleared before filling the received entries
347 : !> \param array local columns, with global rows and a one-based local column index
348 : !> \param para_env communicator covering all source and destination ranks
349 : !> \param first_col first global column held locally
350 : !> \param last_col last global column held locally
351 : !> \param nrows number of source rows to copy
352 : ! **************************************************************************************************
353 144 : SUBROUTINE cp_array_columns_to_fm(matrix, array, para_env, first_col, last_col, nrows)
354 : TYPE(cp_fm_type), INTENT(IN) :: matrix
355 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :), &
356 : INTENT(IN) :: array
357 : TYPE(mp_para_env_type), INTENT(IN) :: para_env
358 : INTEGER, INTENT(IN) :: first_col, last_col, nrows
359 :
360 : CHARACTER(LEN=*), PARAMETER :: routineN = 'cp_array_columns_to_fm'
361 :
362 : INTEGER :: handle, i, ientry, j, nprocs, pass, peer
363 144 : INTEGER, ALLOCATABLE :: counter(:), recv_count(:), send_count(:)
364 144 : TYPE(matrix_buffer_type) :: buffer_rec, buffer_send
365 :
366 144 : CALL timeset(routineN, handle)
367 144 : nprocs = para_env%num_pe
368 720 : ALLOCATE (send_count(0:nprocs - 1), recv_count(0:nprocs - 1), counter(0:nprocs - 1))
369 144 : send_count = 0
370 432 : DO pass = 1, 2
371 288 : IF (pass == 2) THEN
372 144 : CALL para_env%alltoall(send_count, recv_count, 1)
373 144 : CALL allocate_matrix_buffers(send_count, recv_count, 2, buffer_send, buffer_rec)
374 432 : counter(:) = buffer_send%msg_offset(:nprocs - 1)
375 : END IF
376 13200 : DO j = first_col, last_col
377 1191164 : DO i = 1, nrows
378 : peer = matrix%matrix_struct%context%blacs2mpi( &
379 1178108 : matrix%matrix_struct%g2p_row(i), matrix%matrix_struct%g2p_col(j))
380 1190876 : IF (pass == 1) THEN
381 589054 : send_count(peer) = send_count(peer) + 1
382 : ELSE
383 589054 : counter(peer) = counter(peer) + 1
384 589054 : ientry = counter(peer)
385 589054 : buffer_send%msg(ientry) = array(i, j - first_col + 1)
386 1767162 : buffer_send%indx(:, ientry) = [i, j]
387 : END IF
388 : END DO
389 : END DO
390 : END DO
391 :
392 144 : CALL exchange_matrix_buffers(para_env, buffer_rec, buffer_send)
393 573746 : matrix%local_data = 0.0_dp
394 589198 : DO ientry = 1, SIZE(buffer_rec%msg)
395 589054 : i = matrix%matrix_struct%g2l_row(buffer_rec%indx(1, ientry))
396 589054 : j = matrix%matrix_struct%g2l_col(buffer_rec%indx(2, ientry))
397 589198 : matrix%local_data(i, j) = buffer_rec%msg(ientry)
398 : END DO
399 144 : CALL timestop(handle)
400 288 : END SUBROUTINE cp_array_columns_to_fm
401 :
402 : END MODULE cp_fm_redistribution
|
