Coverage for /dolfinx-env/lib/python3.14/site-packages/io4dolfinx/comm_helpers.py: 100%
144 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-21 11:50 +0000
1"""
2Helpers for sending and receiving values for checkpointing
3"""
5from __future__ import annotations
7from mpi4py import MPI
9import numpy as np
10import numpy.typing as npt
12from .utils import compute_insert_position, compute_local_range, valid_function_types
14__all__ = [
15 "send_dofmap_and_recv_values",
16 "send_and_recv_cell_perm",
17 "send_dofs_and_recv_values",
18 "neighbourhood_ranks",
19 "exchange_to_owners",
20 "numpy_to_mpi",
21 "all_to_all",
22]
24numpy_to_mpi = {
25 np.float64: MPI.DOUBLE,
26 np.float32: MPI.FLOAT,
27 np.complex64: MPI.COMPLEX,
28 np.complex128: MPI.DOUBLE_COMPLEX,
29 np.int64: MPI.INT64_T,
30 np.int32: MPI.INT32_T,
31}
34def all_to_all(comm, send_data, recv_data):
35 """
36 Exchange a single item with each neighbor in a distributed graph communicator.
38 Note:
39 The count is passed explicitly, and is 1 on every process. MPI-4.1 9.6.2 requires
40 the type signature of ``sendcount``/``sendtype`` at a process to equal that of
41 ``recvcount``/``recvtype`` at *any other* process in the communicator, not just at
42 its neighbors, so the count must be identical on every process whatever its degree.
43 Left implicit, mpi4py derives it as ``buffer size // degree`` and falls back to the
44 whole buffer when the degree is zero, making it rank-local: 1 where the degree is
45 nonzero, 0 where it is zero. Such a call is erroneous; Open MPI rejects it with
46 ``MPI_ERR_TRUNCATE`` while MPICH happens to accept it. See
47 https://github.com/open-mpi/ompi/issues/14452 for the discussion.
49 ``all_to_allv`` is not affected: the vector variant is only required to match
50 pairwise along each edge, so per-process counts may legitimately differ there.
51 """
52 dtype = numpy_to_mpi[send_data.dtype.type]
53 assert recv_data.dtype == send_data.dtype, (
54 f"Data types do not match, {recv_data.dtype} != {send_data.dtype}"
55 )
56 indegree, outdegree, _ = comm.Get_dist_neighbors_count()
57 assert (d_size := send_data.size) == outdegree, (
58 f"Number of send data {d_size} does not match number of destinations {outdegree}"
59 )
60 assert (d_size := recv_data.size) == indegree, (
61 f"Number of recv data {d_size} does not match number of sources {indegree}"
62 )
63 comm.Neighbor_alltoall([send_data, 1, dtype], [recv_data, 1, dtype])
66def send_dofmap_and_recv_values(
67 comm: MPI.Comm,
68 source_ranks: npt.NDArray[np.int32],
69 dest_ranks: npt.NDArray[np.int32],
70 output_owners: npt.NDArray[np.int32],
71 dest_size: npt.NDArray[np.int32],
72 input_cells: npt.NDArray[np.int64],
73 dofmap_pos: npt.NDArray[np.int32],
74 num_cells_global: int | np.int64,
75 values: npt.NDArray[valid_function_types],
76 dofmap_offsets: npt.NDArray[np.int32],
77) -> npt.NDArray[valid_function_types]:
78 """
79 Given a set of positions in input dofmap, give the global input index of this dofmap entry
80 in input file.
82 Args:
83 comm: The MPI communicator to create the Neighbourhood-communicator from
84 source_ranks: Ranks that will send dofmap indices to current process
85 dest_ranks: Ranks that will receive dofmap indices from current process
86 output_owners: The owners of each dofmap entry on this process. The unique set of
87 these entries should be the same as the dest_ranks.
88 dest_size: The number of entries sent to each owner
89 input_cells: A cell associated with the degree of freedom sent (global index).
90 dofmap_pos: The local position in the dofmap. I.e.
91 `dof = dofmap.links(input_cells)[dofmap_pos]`
92 num_cells_global: Number of global cells
93 values: Values currently held by this process. These are
94 ordered (num_cells_local, num_dofs_per_cell), flattened row-major.
95 dofmap_offsets: Local dofmap offsets to access the correct `values`.
97 Returns:
98 Values corresponding to the dofs owned by this process.
99 """
100 insert_position = compute_insert_position(output_owners, dest_ranks, dest_size)
102 # Pack the cells and dofmap position for all dofs this process is distributing
103 out_cells = np.zeros(len(output_owners), dtype=np.int64)
104 out_cells[insert_position] = input_cells
105 out_pos = np.zeros(len(output_owners), dtype=np.int32)
106 out_pos[insert_position] = dofmap_pos
108 # Compute map from the data index sent to each process and the local
109 # number on the current process
110 proc_to_dof = np.zeros_like(input_cells, dtype=np.int32)
111 proc_to_dof[insert_position] = np.arange(len(input_cells), dtype=np.int32)
112 del insert_position
114 # Send sizes to create data structures for receiving from NeighAlltoAllv
115 recv_size = np.zeros(len(source_ranks), dtype=np.int32)
116 assert isinstance(comm, MPI.Intracomm)
117 mesh_to_data_comm = comm.Create_dist_graph_adjacent(
118 source_ranks.tolist(), dest_ranks.tolist(), reorder=False
119 )
120 all_to_all(mesh_to_data_comm, dest_size, recv_size)
122 # Prepare data-structures for receiving
123 total_incoming = sum(recv_size)
124 inc_cells = np.zeros(total_incoming, dtype=np.int64)
125 inc_pos = np.zeros(total_incoming, dtype=np.intc)
127 # Compute incoming offset
128 inc_offsets = np.zeros(len(recv_size) + 1, dtype=np.intc)
129 inc_offsets[1:] = np.cumsum(recv_size)
131 # Send data
132 s_msg = [out_cells, dest_size, MPI.INT64_T]
133 r_msg = [inc_cells, recv_size, MPI.INT64_T]
134 mesh_to_data_comm.Neighbor_alltoallv(s_msg, r_msg)
136 s_msg = [out_pos, dest_size, MPI.INT32_T]
137 r_msg = [inc_pos, recv_size, MPI.INT32_T]
138 mesh_to_data_comm.Neighbor_alltoallv(s_msg, r_msg)
139 mesh_to_data_comm.Free()
141 local_input_range = compute_local_range(comm, num_cells_global)
142 values_to_distribute = np.zeros_like(inc_pos, dtype=values.dtype)
144 # Map values based on input cells and dofmap
145 local_cells = inc_cells - local_input_range[0]
146 values_to_distribute = values[dofmap_offsets[local_cells] + inc_pos]
148 # Send input dofs back to owning process
149 data_to_mesh_comm = comm.Create_dist_graph_adjacent(
150 dest_ranks.tolist(), source_ranks.tolist(), reorder=False
151 )
153 incoming_global_dofs = np.zeros(sum(dest_size), dtype=values.dtype)
154 s_msg = [values_to_distribute, recv_size, numpy_to_mpi[values.dtype.type]]
155 r_msg = [incoming_global_dofs, dest_size, numpy_to_mpi[values.dtype.type]]
156 data_to_mesh_comm.Neighbor_alltoallv(s_msg, r_msg)
158 # Sort incoming global dofs as they were inputted
159 assert len(incoming_global_dofs) == len(input_cells)
160 sorted_global_dofs = np.zeros_like(incoming_global_dofs, dtype=values.dtype)
161 sorted_global_dofs[proc_to_dof] = incoming_global_dofs
163 data_to_mesh_comm.Free()
164 return sorted_global_dofs
167def send_and_recv_cell_perm(
168 cells: npt.NDArray[np.int64],
169 perms: npt.NDArray[np.uint32],
170 cell_owners: npt.NDArray[np.int32],
171 comm: MPI.Comm,
172) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.uint32]]:
173 """
174 Send global cell index and permutation to corresponding entry in `dest_ranks`.
176 Args:
177 cells: The global input index of the cell
178 perms: The corresponding cell permutation of the cell
179 cell_owners: The rank to send the i-th entry of cells and perms to
180 comm: Rank of comm to generate neighbourhood communicator from
181 """
182 dest_ranks, _dest_size = np.unique(cell_owners, return_counts=True)
183 dest_size = _dest_size.astype(np.int32)
184 del _dest_size
185 assert isinstance(comm, MPI.Intracomm)
186 mesh_to_data = comm.Create_dist_graph(
187 [comm.rank], [len(dest_ranks)], dest_ranks.tolist(), reorder=False
188 )
189 source, dest, _ = mesh_to_data.Get_dist_neighbors()
190 assert np.allclose(dest, dest_ranks)
191 insert_position = compute_insert_position(cell_owners, dest_ranks.astype(np.int32), dest_size)
193 # Pack cells and permutations for sending
194 out_cells = np.zeros_like(cells, dtype=np.int64)
195 out_perm = np.zeros_like(perms, dtype=np.uint32)
196 out_cells[insert_position] = cells
197 out_perm[insert_position] = perms
198 del insert_position
200 # Send sizes to create data structures for receiving from NeighAlltoAllv
201 recv_size = np.zeros_like(source, dtype=np.int32)
202 all_to_all(mesh_to_data, dest_size, recv_size)
204 # Prepare data-structures for receiving
205 total_incoming = sum(recv_size)
206 inc_cells = np.zeros(total_incoming, dtype=np.int64)
207 inc_perm = np.zeros(total_incoming, dtype=np.uint32)
209 # Compute incoming offset
210 inc_offsets = np.zeros(len(recv_size) + 1, dtype=np.intc)
211 inc_offsets[1:] = np.cumsum(recv_size)
213 # Send data
214 s_msg = [out_cells, dest_size, MPI.INT64_T]
215 r_msg = [inc_cells, recv_size, MPI.INT64_T]
216 mesh_to_data.Neighbor_alltoallv(s_msg, r_msg)
218 s_msg = [out_perm, dest_size, MPI.UINT32_T]
219 r_msg = [inc_perm, recv_size, MPI.UINT32_T]
220 mesh_to_data.Neighbor_alltoallv(s_msg, r_msg)
221 mesh_to_data.Free()
222 return inc_cells, inc_perm
225def send_dofs_and_recv_values(
226 input_dofmap: npt.NDArray[np.int64],
227 dofmap_owners: npt.NDArray[np.int32],
228 comm: MPI.Comm,
229 input_array: npt.NDArray[valid_function_types],
230 array_start: int,
231):
232 """
233 Send a set of dofs (global index) to the process holding the DOF values to retrieve them.
235 Args:
236 input_dofmap: List of dofs (global index) that this process wants values for
237 dofmap_owners: The process currently holding the values this process want to get.
238 comm: MPI communicator
239 input_array: Values for dofs
240 array_start: The global starting index of `input_array`.
241 """
242 dest_ranks, _dest_size = np.unique(dofmap_owners, return_counts=True)
243 dest_size = _dest_size.astype(np.int32)
244 del _dest_size
246 assert isinstance(comm, MPI.Intracomm)
247 dofmap_to_values = comm.Create_dist_graph(
248 [comm.rank], [len(dest_ranks)], dest_ranks.tolist(), reorder=False
249 )
251 source, dest, _ = dofmap_to_values.Get_dist_neighbors()
252 assert np.allclose(dest_ranks, dest)
253 # Compute amount of data to send to each process
255 insert_position = compute_insert_position(dofmap_owners, dest_ranks, dest_size)
257 # Pack dofs for sending
258 out_dofs = np.zeros(len(dofmap_owners), dtype=np.int64)
259 out_dofs[insert_position] = input_dofmap
261 # Compute map from the data index sent to each process and the local number on
262 # the current process
263 proc_to_local = np.zeros_like(input_dofmap, dtype=np.int32)
264 proc_to_local[insert_position] = np.arange(len(input_dofmap), dtype=np.int32)
265 del insert_position
267 # Send sizes to create data structures for receiving from NeighAlltoAllv
268 recv_size = np.zeros_like(source, dtype=np.int32)
269 all_to_all(dofmap_to_values, dest_size, recv_size)
271 # Send input dofs to processes holding input array
272 inc_dofs = np.zeros(sum(recv_size), dtype=np.int64)
273 s_msg = [out_dofs, dest_size, MPI.INT64_T]
274 r_msg = [inc_dofs, recv_size, MPI.INT64_T]
275 dofmap_to_values.Neighbor_alltoallv(s_msg, r_msg)
276 dofmap_to_values.Free()
278 # Send back appropriate input values
279 if len(input_array) > 0:
280 sending_values = input_array[inc_dofs - array_start]
281 else:
282 sending_values = np.zeros(0, dtype=input_array.dtype)
284 values_to_dofmap = comm.Create_dist_graph_adjacent(dest, source, reorder=False)
285 inc_values = np.zeros_like(out_dofs, dtype=input_array.dtype)
286 s_msg_rev = [sending_values, recv_size, numpy_to_mpi[input_array.dtype.type]]
287 r_msg_rev = [inc_values, dest_size, numpy_to_mpi[input_array.dtype.type]]
288 values_to_dofmap.Neighbor_alltoallv(s_msg_rev, r_msg_rev)
289 values_to_dofmap.Free()
291 # Sort inputs according to local dof number (input process)
292 values = np.empty_like(inc_values, dtype=input_array.dtype)
293 values[proc_to_local] = inc_values
294 return values
297def neighbourhood_ranks(
298 comm: MPI.Intracomm, destinations: npt.NDArray[np.int32]
299) -> tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]:
300 """Find the ranks this process exchanges data with.
302 A process knows which ranks it must send to, but not which ranks will send
303 to it. Building a distributed graph over the whole communicator and reading
304 its neighbours back discovers the incoming side without an all-to-all.
306 Both lists are sorted ascending so that every process packs and unpacks
307 neighbour blocks in the same order.
309 Args:
310 comm: The MPI communicator
311 destinations: Ranks this process sends to (duplicates allowed)
313 Returns:
314 ``(sources, destinations)``, each sorted and without duplicates.
315 """
316 dest = np.unique(np.asarray(destinations, dtype=np.int32))
317 graph = comm.Create_dist_graph([comm.rank], [len(dest)], dest.tolist(), reorder=False)
318 sources, _, _ = graph.Get_dist_neighbors()
319 graph.Free()
320 return np.unique(np.asarray(sources, dtype=np.int32)), dest
323def exchange_to_owners(
324 comm: MPI.Intracomm,
325 owners: npt.NDArray[np.int32],
326 arrays: list[npt.NDArray],
327) -> tuple[
328 list[npt.NDArray],
329 npt.NDArray[np.int32],
330 npt.NDArray[np.int32],
331 npt.NDArray[np.int32],
332 npt.NDArray[np.int32],
333]:
334 """Send each entry of ``arrays`` to the rank naming it in ``owners``.
336 Entry ``i`` of every array in ``arrays`` travels together to ``owners[i]``.
337 Data is packed grouped by destination rank in ascending rank order and is
338 received grouped by source rank in ascending rank order, preserving the
339 sender's relative order within each group. Both properties matter to callers
340 that must pair a reply with the request that produced it.
342 Args:
343 comm: The MPI communicator
344 owners: Destination rank of each entry; ``-1`` is not allowed
345 arrays: Arrays to send, each of length ``len(owners)``. The leading axis
346 is the one distributed; trailing axes travel as a block.
348 Returns:
349 ``(received, sources, send_counts, recv_counts, insert_position)`` where
350 ``received`` holds the incoming counterpart of each entry of ``arrays``,
351 ``sources`` the ascending source ranks, ``send_counts`` and
352 ``recv_counts`` the number of entries exchanged with each destination and
353 source, and ``insert_position`` the index each local entry was packed at
354 (so a reply can be unpacked back into local order).
355 """
356 owners = np.asarray(owners, dtype=np.int32)
357 sources, destinations = neighbourhood_ranks(comm, owners)
359 send_counts = np.zeros(len(destinations), dtype=np.int32)
360 if len(owners) > 0:
361 present, counts = np.unique(owners, return_counts=True)
362 send_counts[np.searchsorted(destinations, present)] = counts
363 insert_position = compute_insert_position(owners, destinations, send_counts)
365 recv_counts = np.zeros(len(sources), dtype=np.int32)
366 forward = comm.Create_dist_graph_adjacent(
367 sources.tolist(), destinations.tolist(), reorder=False
368 )
369 all_to_all(forward, send_counts, recv_counts)
371 received = []
372 for array in arrays:
373 array = np.asarray(array)
374 block = int(np.prod(array.shape[1:], dtype=np.int64)) if array.ndim > 1 else 1
375 packed = np.zeros(array.shape, dtype=array.dtype)
376 packed[insert_position] = array
377 out = np.zeros((int(recv_counts.sum()), *array.shape[1:]), dtype=array.dtype)
378 mpi_type = numpy_to_mpi[array.dtype.type]
379 forward.Neighbor_alltoallv(
380 [packed.reshape(-1), send_counts * block, mpi_type],
381 [out.reshape(-1), recv_counts * block, mpi_type],
382 )
383 received.append(out)
384 forward.Free()
385 return received, sources, send_counts, recv_counts, insert_position