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 23638 : SUBROUTINE allocate_matrix_buffers(send_count, recv_count, n_index, buffer_send, buffer_rec, &
44 23638 : 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 23638 : CPASSERT(PRESENT(send_blocks) .EQV. PRESENT(recv_blocks))
53 23638 : nprocs = SIZE(send_count)
54 0 : ALLOCATE (buffer_send%msg_offset(0:nprocs), buffer_rec%msg_offset(0:nprocs), &
55 141828 : buffer_send%indx_offset(0:nprocs), buffer_rec%indx_offset(0:nprocs))
56 23638 : buffer_send%msg_offset(0) = 0
57 23638 : buffer_rec%msg_offset(0) = 0
58 23638 : buffer_send%indx_offset(0) = 0
59 23638 : buffer_rec%indx_offset(0) = 0
60 70378 : DO peer = 0, nprocs - 1
61 46740 : buffer_send%msg_offset(peer + 1) = buffer_send%msg_offset(peer) + send_count(peer)
62 46740 : buffer_rec%msg_offset(peer + 1) = buffer_rec%msg_offset(peer) + recv_count(peer)
63 46740 : buffer_send%indx_offset(peer + 1) = buffer_send%indx_offset(peer) + send_count(peer)
64 46740 : buffer_rec%indx_offset(peer + 1) = buffer_rec%indx_offset(peer) + recv_count(peer)
65 70378 : IF (PRESENT(send_blocks)) THEN
66 40632 : buffer_send%indx_offset(peer + 1) = buffer_send%indx_offset(peer) + send_blocks(peer)
67 40632 : 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 249878 : buffer_rec%indx(n_index, buffer_rec%indx_offset(nprocs)))
74 23638 : 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 274472 : 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 274472 : TYPE(mp_request_type), ALLOCATABLE, DIMENSION(:) :: requests
94 :
95 274472 : CALL timeset(routineN, handle)
96 274472 : nprocs = comm%num_pe
97 274472 : my_do_indx = .TRUE.
98 274472 : IF (PRESENT(do_indx)) my_do_indx = do_indx
99 274472 : my_do_msg = .TRUE.
100 274472 : IF (PRESENT(do_msg)) my_do_msg = do_msg
101 :
102 : ! Copy the self segment directly, including the single-rank case.
103 274472 : peer = comm%mepos
104 274472 : IF (my_do_indx) THEN
105 : buffer_rec%indx(:, buffer_rec%indx_offset(peer) + 1:buffer_rec%indx_offset(peer + 1)) = &
106 53563151 : buffer_send%indx(:, buffer_send%indx_offset(peer) + 1:buffer_send%indx_offset(peer + 1))
107 : END IF
108 274472 : IF (my_do_msg) THEN
109 : buffer_rec%msg(buffer_rec%msg_offset(peer) + 1:buffer_rec%msg_offset(peer + 1)) = &
110 985382257 : buffer_send%msg(buffer_send%msg_offset(peer) + 1:buffer_send%msg_offset(peer + 1))
111 : END IF
112 :
113 274472 : IF (nprocs > 1) THEN
114 273592 : nreq = 0
115 273592 : IF (my_do_indx) THEN
116 : nreq = COUNT(buffer_rec%indx_offset(1:) > buffer_rec%indx_offset(:nprocs - 1)) + &
117 1367880 : COUNT(buffer_send%indx_offset(1:) > buffer_send%indx_offset(:nprocs - 1))
118 : END IF
119 273592 : IF (my_do_msg) THEN
120 : nreq = nreq + COUNT(buffer_rec%msg_offset(1:) > buffer_rec%msg_offset(:nprocs - 1)) + &
121 1367880 : COUNT(buffer_send%msg_offset(1:) > buffer_send%msg_offset(:nprocs - 1))
122 : END IF
123 2137696 : ALLOCATE (requests(nreq))
124 273592 : nreq = 0
125 :
126 : ! Post both channels before waiting so metadata and payload can overlap.
127 820776 : DO peer = 0, nprocs - 1
128 547184 : IF (peer == comm%mepos) CYCLE
129 273592 : IF (my_do_indx) THEN
130 273576 : first = buffer_rec%indx_offset(peer) + 1
131 273576 : last = buffer_rec%indx_offset(peer + 1)
132 273576 : IF (last >= first) THEN
133 145242 : nreq = nreq + 1
134 145242 : CALL comm%irecv(buffer_rec%indx(:, first:last), peer, requests(nreq), tag=4)
135 : END IF
136 : END IF
137 547184 : IF (my_do_msg .AND. buffer_rec%msg_offset(peer + 1) > buffer_rec%msg_offset(peer)) THEN
138 145242 : first = buffer_rec%msg_offset(peer) + 1
139 145242 : last = buffer_rec%msg_offset(peer + 1)
140 145242 : nreq = nreq + 1
141 145242 : CALL comm%irecv(buffer_rec%msg(first:last), peer, requests(nreq), tag=7)
142 : END IF
143 : END DO
144 :
145 820776 : DO peer = 0, nprocs - 1
146 547184 : IF (peer == comm%mepos) CYCLE
147 273592 : IF (my_do_indx) THEN
148 273576 : first = buffer_send%indx_offset(peer) + 1
149 273576 : last = buffer_send%indx_offset(peer + 1)
150 273576 : IF (last >= first) THEN
151 145242 : nreq = nreq + 1
152 145242 : CALL comm%isend(buffer_send%indx(:, first:last), peer, requests(nreq), tag=4)
153 : END IF
154 : END IF
155 547184 : IF (my_do_msg .AND. buffer_send%msg_offset(peer + 1) > buffer_send%msg_offset(peer)) THEN
156 145242 : first = buffer_send%msg_offset(peer) + 1
157 145242 : last = buffer_send%msg_offset(peer + 1)
158 145242 : nreq = nreq + 1
159 145242 : CALL comm%isend(buffer_send%msg(first:last), peer, requests(nreq), tag=7)
160 : END IF
161 : END DO
162 273592 : IF (nreq > 0) CALL mp_waitall(requests(:nreq))
163 : END IF
164 274472 : CALL timestop(handle)
165 274472 : END SUBROUTINE exchange_matrix_buffers
166 :
167 0 : END MODULE matrix_communication
|
