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 Exchange numerical matrix payloads and integer metadata between MPI ranks.
10 : ! **************************************************************************************************
11 : MODULE matrix_communication
12 : USE kinds, ONLY: dp
13 : USE message_passing, ONLY: mp_comm_type,&
14 : mp_request_type,&
15 : mp_waitall
16 : #include "./base/base_uses.f90"
17 :
18 : IMPLICIT NONE
19 : PRIVATE
20 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'matrix_communication'
21 : PUBLIC :: matrix_buffer_type, allocate_matrix_buffers, exchange_matrix_buffers
22 :
23 : TYPE matrix_buffer_type
24 : REAL(KIND=dp), DIMENSION(:), ALLOCATABLE :: msg
25 : ! Metadata records are columns, contiguous within each rank's segment.
26 : INTEGER, DIMENSION(:, :), ALLOCATABLE :: indx
27 : ! Zero-based segment offsets, including the final total at index nprocs.
28 : INTEGER, DIMENSION(:), ALLOCATABLE :: msg_offset, indx_offset
29 : END TYPE matrix_buffer_type
30 :
31 : CONTAINS
32 :
33 : ! **************************************************************************************************
34 : !> \brief Allocate flat, rank-grouped matrix buffers and their zero-based offsets.
35 : !> \param send_count numerical payload length sent to each rank
36 : !> \param recv_count numerical payload length received from each rank
37 : !> \param n_index number of integer coordinates per metadata record
38 : !> \param buffer_send output send buffer
39 : !> \param buffer_rec output receive buffer
40 : !> \param send_blocks optional metadata-record counts, distinct from payload counts
41 : !> \param recv_blocks optional received metadata-record counts
42 : ! **************************************************************************************************
43 315709 : SUBROUTINE allocate_matrix_buffers(send_count, recv_count, n_index, buffer_send, buffer_rec, &
44 315709 : send_blocks, recv_blocks)
45 : INTEGER, DIMENSION(0:), INTENT(IN) :: send_count, recv_count
46 : INTEGER, INTENT(IN) :: n_index
47 : TYPE(matrix_buffer_type), INTENT(OUT) :: buffer_send, buffer_rec
48 : INTEGER, DIMENSION(0:), INTENT(IN), OPTIONAL :: send_blocks, recv_blocks
49 :
50 : INTEGER :: nprocs, peer
51 :
52 315709 : CPASSERT(PRESENT(send_blocks) .EQV. PRESENT(recv_blocks))
53 315709 : nprocs = SIZE(send_count)
54 0 : ALLOCATE (buffer_send%msg_offset(0:nprocs), buffer_rec%msg_offset(0:nprocs), &
55 1894254 : buffer_send%indx_offset(0:nprocs), buffer_rec%indx_offset(0:nprocs))
56 315709 : buffer_send%msg_offset(0) = 0
57 315709 : buffer_rec%msg_offset(0) = 0
58 315709 : buffer_send%indx_offset(0) = 0
59 315709 : buffer_rec%indx_offset(0) = 0
60 868486 : DO peer = 0, nprocs - 1
61 552777 : buffer_send%msg_offset(peer + 1) = buffer_send%msg_offset(peer) + send_count(peer)
62 552777 : buffer_rec%msg_offset(peer + 1) = buffer_rec%msg_offset(peer) + recv_count(peer)
63 552777 : buffer_send%indx_offset(peer + 1) = buffer_send%indx_offset(peer) + send_count(peer)
64 552777 : buffer_rec%indx_offset(peer + 1) = buffer_rec%indx_offset(peer) + recv_count(peer)
65 868486 : IF (PRESENT(send_blocks)) THEN
66 546321 : buffer_send%indx_offset(peer + 1) = buffer_send%indx_offset(peer) + send_blocks(peer)
67 546321 : buffer_rec%indx_offset(peer + 1) = buffer_rec%indx_offset(peer) + recv_blocks(peer)
68 : END IF
69 : END DO
70 0 : ALLOCATE (buffer_send%msg(buffer_send%msg_offset(nprocs)), &
71 0 : buffer_rec%msg(buffer_rec%msg_offset(nprocs)), &
72 0 : buffer_send%indx(n_index, buffer_send%indx_offset(nprocs)), &
73 3249351 : buffer_rec%indx(n_index, buffer_rec%indx_offset(nprocs)))
74 315709 : END SUBROUTINE allocate_matrix_buffers
75 :
76 : ! **************************************************************************************************
77 : !> \brief Synchronously exchange flat matrix buffers with nonempty peers.
78 : !> \param comm communicator covering all sending and receiving ranks
79 : !> \param buffer_rec preallocated receive buffer
80 : !> \param buffer_send populated send buffer
81 : !> \param do_indx whether to exchange integer metadata (default: true)
82 : !> \param do_msg whether to exchange numerical payloads (default: true)
83 : ! **************************************************************************************************
84 253523 : SUBROUTINE exchange_matrix_buffers(comm, buffer_rec, buffer_send, do_indx, do_msg)
85 : CLASS(mp_comm_type), INTENT(IN) :: comm
86 : TYPE(matrix_buffer_type), INTENT(INOUT) :: buffer_rec, buffer_send
87 : LOGICAL, INTENT(IN), OPTIONAL :: do_indx, do_msg
88 :
89 : CHARACTER(LEN=*), PARAMETER :: routineN = 'exchange_matrix_buffers'
90 :
91 : INTEGER :: first, handle, last, nprocs, nreq, peer
92 : LOGICAL :: my_do_indx, my_do_msg
93 253523 : TYPE(mp_request_type), ALLOCATABLE, DIMENSION(:) :: requests
94 :
95 253523 : CALL timeset(routineN, handle)
96 253523 : nprocs = comm%num_pe
97 253523 : my_do_indx = .TRUE.
98 253523 : IF (PRESENT(do_indx)) my_do_indx = do_indx
99 253523 : my_do_msg = .TRUE.
100 253523 : IF (PRESENT(do_msg)) my_do_msg = do_msg
101 :
102 : ! Copy the self segment directly, including the single-rank case.
103 253523 : peer = comm%mepos
104 253523 : IF (my_do_indx) THEN
105 : buffer_rec%indx(:, buffer_rec%indx_offset(peer) + 1:buffer_rec%indx_offset(peer + 1)) = &
106 50718192 : buffer_send%indx(:, buffer_send%indx_offset(peer) + 1:buffer_send%indx_offset(peer + 1))
107 : END IF
108 253523 : IF (my_do_msg) THEN
109 : buffer_rec%msg(buffer_rec%msg_offset(peer) + 1:buffer_rec%msg_offset(peer + 1)) = &
110 605857322 : buffer_send%msg(buffer_send%msg_offset(peer) + 1:buffer_send%msg_offset(peer + 1))
111 : END IF
112 :
113 253523 : IF (nprocs > 1) THEN
114 238504 : nreq = 0
115 238504 : IF (my_do_indx) THEN
116 : nreq = COUNT(buffer_rec%indx_offset(1:) > buffer_rec%indx_offset(:nprocs - 1)) + &
117 1185340 : COUNT(buffer_send%indx_offset(1:) > buffer_send%indx_offset(:nprocs - 1))
118 : END IF
119 238504 : IF (my_do_msg) THEN
120 : nreq = nreq + COUNT(buffer_rec%msg_offset(1:) > buffer_rec%msg_offset(:nprocs - 1)) + &
121 1192440 : COUNT(buffer_send%msg_offset(1:) > buffer_send%msg_offset(:nprocs - 1))
122 : END IF
123 1868218 : ALLOCATE (requests(nreq))
124 238504 : nreq = 0
125 :
126 : ! Post both channels before waiting so metadata and payload can overlap.
127 715512 : DO peer = 0, nprocs - 1
128 477008 : IF (peer == comm%mepos) CYCLE
129 238504 : IF (my_do_indx) THEN
130 237068 : first = buffer_rec%indx_offset(peer) + 1
131 237068 : last = buffer_rec%indx_offset(peer + 1)
132 237068 : IF (last >= first) THEN
133 122997 : nreq = nreq + 1
134 122997 : CALL comm%irecv(buffer_rec%indx(:, first:last), peer, requests(nreq), tag=4)
135 : END IF
136 : END IF
137 477008 : IF (my_do_msg .AND. buffer_rec%msg_offset(peer + 1) > buffer_rec%msg_offset(peer)) THEN
138 124308 : first = buffer_rec%msg_offset(peer) + 1
139 124308 : last = buffer_rec%msg_offset(peer + 1)
140 124308 : nreq = nreq + 1
141 124308 : CALL comm%irecv(buffer_rec%msg(first:last), peer, requests(nreq), tag=7)
142 : END IF
143 : END DO
144 :
145 715512 : DO peer = 0, nprocs - 1
146 477008 : IF (peer == comm%mepos) CYCLE
147 238504 : IF (my_do_indx) THEN
148 237068 : first = buffer_send%indx_offset(peer) + 1
149 237068 : last = buffer_send%indx_offset(peer + 1)
150 237068 : IF (last >= first) THEN
151 122997 : nreq = nreq + 1
152 122997 : CALL comm%isend(buffer_send%indx(:, first:last), peer, requests(nreq), tag=4)
153 : END IF
154 : END IF
155 477008 : IF (my_do_msg .AND. buffer_send%msg_offset(peer + 1) > buffer_send%msg_offset(peer)) THEN
156 124308 : first = buffer_send%msg_offset(peer) + 1
157 124308 : last = buffer_send%msg_offset(peer + 1)
158 124308 : nreq = nreq + 1
159 124308 : CALL comm%isend(buffer_send%msg(first:last), peer, requests(nreq), tag=7)
160 : END IF
161 : END DO
162 238504 : IF (nreq > 0) CALL mp_waitall(requests(:nreq))
163 : END IF
164 253523 : CALL timestop(handle)
165 253523 : END SUBROUTINE exchange_matrix_buffers
166 :
167 0 : END MODULE matrix_communication
|
