Coverage for dragonfly_sim/utils/mpi_utils.py: 75%

1import numpy as np 

2 

3 

4def allgather_numpy(comm, local_arr, global_size): 

5 """Allgatherv a rank-local float64 array into a global array on every rank. 

6 

7 Each rank contributes a (possibly different-length) slice of the global 

8 array. The slices are ordered by rank: rank 0 occupies [0, len_0), 

9 rank 1 occupies [len_0, len_0+len_1), etc. This ordering matches the 

10 contiguous DOF blocks produced by ``assign_global_dof_offset``. 

11 

12 Parameters 

13 ---------- 

14 comm : mpi4py.MPI.Comm 

15 local_arr : array-like, float64, 1-D 

16 The local contribution on this rank. Cast to C-contiguous float64 

17 before sending. 

18 global_size : int 

19 Total number of float64 elements summed across all ranks. Must be 

20 the same value on every rank. 

21 

22 Returns 

23 ------- 

24 np.ndarray, float64, shape (global_size,) 

25 The gathered array, identical on every rank. 

26 """ 

27 from mpi4py import MPI 

28 send = np.ascontiguousarray(local_arr, dtype=np.float64) 

29 

30 # Serial fast-path: skip MPI overhead entirely 

31 if comm.Get_size() == 1: 

32 recv = np.empty(global_size, dtype=np.float64) 

33 recv[:len(send)] = send 

34 return recv 

35 

36 recv = np.zeros(global_size, dtype=np.float64) 

37 sendcounts = np.array(comm.allgather(len(send)), dtype=np.int32) 

38 displs = np.concatenate([[0], np.cumsum(sendcounts[:-1])]).astype(np.int32) 

39 comm.Allgatherv( 

40 sendbuf=send, 

41 recvbuf=[recv, sendcounts, displs, MPI.DOUBLE], 

42 ) 

43 return recv 

44 

45 

46def assign_global_dof_offset(comm, local_count): 

47 """Compute the global DOF offset for this rank via exclusive prefix-sum. 

48 

49 Assigns a contiguous block of global indices to each rank: 

50 rank 0 → [0, local_count_0) 

51 rank 1 → [local_count_0, local_count_0 + local_count_1) 

52 ... 

53 

54 Implementation uses MPI ``Exscan`` (exclusive scan with SUM) to 

55 compute the offset for each rank in a single collective, then 

56 ``Allreduce`` on a scalar to obtain the global total. This avoids 

57 allocating an N-rank array and is the canonical MPI pattern for 

58 assigning contiguous global index ranges to processes. 

59 

60 Parameters 

61 ---------- 

62 comm : mpi4py.MPI.Comm 

63 local_count : int 

64 Number of local DOFs on this rank. 

65 

66 Returns 

67 ------- 

68 offset : int 

69 Index of the first global DOF owned by this rank. 

70 global_total : int 

71 Total DOF count across all ranks. 

72 """ 

73 from mpi4py import MPI 

74 

75 # Exclusive scan: offset[r] = sum of local_count for ranks 0 .. r-1 

76 # local_count_arr = np.array(local_count, dtype=np.int64) 

77 # offset_arr = np.zeros(1, dtype=np.int64) 

78 # offset_init = 0 

79 # Exscan is undefined on rank 0 — result stays 0, which is correct.assign_global_dof_offset 

80 offset = comm.exscan(local_count, op=MPI.SUM) 

81 if offset is None: 

82 # correct the offset on rank 0 

83 offset = 0 

84 # offset = offset_init 

85 

86 # Global total: reduce the local count to a single sum on all ranks 

87 # global_total_arr = np.array(local_count, dtype=np.int64) 

88 global_total = comm.allreduce(local_count, op=MPI.SUM) 

89 # global_total = int(global_total_arr[0]) 

90 

91 return offset, global_total