Coverage for pySDC/helpers/firedrake_ensemble_communicator.py: 85%

1from mpi4py import MPI 

2import firedrake as fd 

3import numpy as np 

4 

5 

6class _Requests: 

7 """ 

8 The requests Firedrake returns for sending or receiving a function, one per subfunction, to be waited on together 

9 like a single request. 

10 """ 

11 

12 def __init__(self, requests): 

13 self.requests = requests 

14 

15 def Wait(self): 

16 MPI.Request.Waitall(self.requests) 

17 

18 wait = Wait 

19 

20 

21class FiredrakeEnsembleCommunicator: 

22 """ 

23 Ensemble communicator for performing multiple similar distributed simulations with Firedrake, see https://www.firedrakeproject.org/firedrake/parallelism.html 

24 This is intended to do space-time parallelism in pySDC. 

25 This class wraps the time communicator. All requests that are not overloaded are passed to the time communicator. For instance, `ensemble.rank` will return the rank in the time communicator. 

26 Some operations are overloaded to use the interface of the MPI communicator but handles communication with the ensemble communicator instead. 

27 """ 

28 

29 def __init__(self, comm, space_size): 

30 """ 

31 Args: 

32 comm (MPI.Intracomm): MPI communicator, which will be split into time and space communicators 

33 space_size (int): Size of the spatial communicators 

34 

35 Attributes: 

36 ensemble (firedrake.Ensemble): Ensemble communicator 

37 """ 

38 self.ensemble = fd.Ensemble(comm, space_size) 

39 self.comm_wold = comm 

40 

41 _owns_comm_wold = False # only communicators made by `Split` are freed by `Free` 

42 

43 def Split(self, *args, **kwargs): 

44 """ 

45 Split the world communicator with mpi4py's `Split` and return a new `FiredrakeEnsembleCommunicator` on the 

46 result, with the same spatial size. 

47 """ 

48 split = FiredrakeEnsembleCommunicator(self.comm_wold.Split(*args, **kwargs), space_size=self.space_comm.size) 

49 split._owns_comm_wold = True 

50 return split 

51 

52 @property 

53 def space_comm(self): 

54 """The spatial communicator of the ensemble, `ensemble.comm`.""" 

55 return self.ensemble.comm 

56 

57 @property 

58 def time_comm(self): 

59 """The time communicator of the ensemble, `ensemble.ensemble_comm`.""" 

60 return self.ensemble.ensemble_comm 

61 

62 def __getattr__(self, name): 

63 return getattr(self.time_comm, name) 

64 

65 def Reduce(self, sendbuf, recvbuf, op=MPI.SUM, root=0): 

66 """ 

67 Wrap `Reduce` on the time communicator for numpy arrays and `Ensemble.reduce` for Firedrake functions, which 

68 only takes `op=MPI.SUM` here. 

69 """ 

70 if type(sendbuf) in [np.ndarray]: 

71 self.ensemble.ensemble_comm.Reduce(sendbuf, recvbuf, op, root) 

72 else: 

73 assert op == MPI.SUM 

74 self.ensemble.reduce(sendbuf, recvbuf, root=root) 

75 

76 def Allreduce(self, sendbuf, recvbuf, op=MPI.SUM): 

77 """ 

78 Wrap `Allreduce` on the time communicator for numpy arrays and `Ensemble.allreduce` for Firedrake functions, 

79 which only takes `op=MPI.SUM` here. 

80 """ 

81 if type(sendbuf) in [np.ndarray]: 

82 self.ensemble.ensemble_comm.Allreduce(sendbuf, recvbuf, op) 

83 else: 

84 assert op == MPI.SUM 

85 self.ensemble.allreduce(sendbuf, recvbuf) 

86 

87 def Bcast(self, buf, root=0): 

88 """Wrap `Bcast` on the time communicator for numpy arrays and `Ensemble.bcast` for Firedrake functions.""" 

89 if type(buf) in [np.ndarray]: 

90 self.ensemble.ensemble_comm.Bcast(buf, root) 

91 else: 

92 self.ensemble.bcast(buf, root=root) 

93 

94 def Irecv(self, buf, source, tag=MPI.ANY_TAG): 

95 """ 

96 Wrap `Irecv` on the time communicator for numpy arrays and lists, and `Ensemble.irecv` for Firedrake functions, 

97 whose requests, one per subfunction, are returned to be waited on together. 

98 """ 

99 if type(buf) in [np.ndarray, list]: 

100 return self.ensemble.ensemble_comm.Irecv(buf=buf, source=source, tag=tag) 

101 return _Requests(self.ensemble.irecv(buf, source, tag=tag)) 

102 

103 def Isend(self, buf, dest, tag=0): 

104 """ 

105 Wrap `Isend` on the time communicator for numpy arrays and lists, and `Ensemble.isend` for Firedrake functions, 

106 whose requests, one per subfunction, are returned to be waited on together. 

107 """ 

108 if type(buf) in [np.ndarray, list]: 

109 return self.ensemble.ensemble_comm.Isend(buf=buf, dest=dest, tag=tag) 

110 return _Requests(self.ensemble.isend(buf, dest, tag=tag)) 

111 

112 def Free(self): 

113 """ 

114 Free the world communicator if `Split` made it. The ensemble frees its own spatial and time communicators when 

115 it is garbage collected. 

116 """ 

117 if self._owns_comm_wold: 

118 self.comm_wold.Free() 

119 

120 

121def get_ensemble(comm, space_size): 

122 """ 

123 Create a Firedrake ensemble. 

124 

125 Args: 

126 comm (MPI.Intracomm): MPI communicator, which will be split into time and space communicators 

127 space_size (int): Size of the spatial communicators 

128 

129 Returns: 

130 firedrake.Ensemble: The ensemble 

131 """ 

132 return fd.Ensemble(comm, space_size)