Coverage for pySDC/implementations/controller_classes/controller_ParaDiag_MPI.py: 100%

1import numpy as np 

2from mpi4py import MPI 

3 

4from pySDC.helpers.ParaDiagHelper import get_G_inv_matrix 

5from pySDC.implementations.controller_classes.ParaDiag import ParaDiag 

6from pySDC.implementations.controller_classes.controller_MPI import controller_MPI 

7 

8 

9class controller_ParaDiag_MPI(ParaDiag, controller_MPI): 

10 """ 

11 ParaDiag controller with MPI parallelism across time steps: one step per rank. 

12 

13 This is `controller_MPI` with a different iteration: where PFASST sweeps and cascades through 

14 the levels, ParaDiag diagonalizes across the steps. Everything around the iteration -- blocks, 

15 windowing, restarts, convergence -- is the driver it inherits, which is why the dispatcher it 

16 uses is still called `pfasst`. 

17 

18 Everything here is written from a single processor's point of view. A rank owns exactly one step 

19 and never inspects another rank's data; the places where ParaDiag genuinely needs information 

20 from the whole block are expressed as communication: 

21 

22 - ``prepare_Jacobians`` -> Allreduce(SUM) over the step communicator 

23 - ``compute_all_at_once_residual`` -> point-to-point exchange with the previous/next rank 

24 - ``apply_matrix`` (the weighted FFT/iFFT in time) -> a ring reduction, see below 

25 - convergence -> the inherited `it_check`, which allreduces because 

26 `all_to_done` is forced on for ParaDiag 

27 

28 Note that ParaDiag steps can only converge together, so every rank always participates in every 

29 iteration. That is not a policy choice: a rank that stopped early would never enter the 

30 collectives below and the run would hang. `step_is_active` says the same thing about blocks -- 

31 a block is never run partially, so the driver's windowing never splits one. 

32 """ 

33 

34 def __init__(self, controller_params, description, comm=None): 

35 """ 

36 Args: 

37 controller_params: parameter set for the controller and the steps 

38 description: all the parameters to set up the rest (levels, problems, transfer, ...) 

39 comm: MPI communicator, one rank per time step 

40 """ 

41 comm = MPI.COMM_WORLD if comm is None else comm 

42 self.prepare_ParaDiag_params(controller_params, description) 

43 

44 self.sweeper_params = description['sweeper_params'] 

45 

46 # each step needs its own G^-1, determined by where it sits in the block 

47 self._G_inv_alpha = self.resolve_alpha(controller_params['alpha'], 0) 

48 description['sweeper_params']['G_inv'] = get_G_inv_matrix( 

49 comm.rank, comm.size, self._G_inv_alpha, description['sweeper_params'] 

50 ) 

51 

52 super().__init__(controller_params, description, comm) 

53 

54 self.n_steps = comm.size 

55 

56 # ------------------------------------------------------------------ collectives 

57 

58 def apply_matrix(self, mat, quantity): 

59 """ 

60 Apply a square L x L matrix across the steps, where L is the number of ranks. 

61 

62 Each rank needs ``res_i = sum_j mat[i, j] * me_j`` but only holds ``me_i``. This is done as a 

63 ring reduction: the values circulate once around the communicator and each rank accumulates 

64 its own row as they pass. That keeps the working set at O(M) fields per rank, independent of 

65 L -- an allgather would instead need O(L * M) fields on every rank, which is exactly the 

66 gather this controller must not do. 

67 

68 The ring costs L - 1 rounds. A butterfly would need only log2(L), but only because the matrix 

69 this is called with is a DFT; for a general matrix it needs the O(L * M) gather just ruled 

70 out. That belongs with an FFT-shaped interface rather than this one, and pays off from about 

71 L = 16 upwards. 

72 

73 Args: 

74 mat: square matrix with as many rows as there are ranks 

75 quantity (str): 'residual' or 'increment', the level attribute to transform in place 

76 """ 

77 comm = self.comm 

78 L, rank = comm.size, comm.rank 

79 assert np.allclose(mat.shape, L), f'need a {L}x{L} matrix, got {mat.shape}' 

80 

81 lvl = self.S.levels[0] 

82 M = lvl.sweep.params.num_nodes 

83 prob = lvl.prob 

84 

85 if quantity == 'residual': 

86 me = lvl.residual 

87 elif quantity == 'increment': 

88 me = lvl.increment 

89 else: 

90 raise NotImplementedError(f'Cannot apply matrix to {quantity!r}') 

91 

92 res = [prob.u_init for _ in range(M)] 

93 # all M nodes travel in one contiguous buffer, so the ring costs L - 1 messages rather than 

94 # M * (L - 1). Same volume, M times fewer message latencies. 

95 held = np.array([me[m] for m in range(M)]) 

96 buf = np.empty_like(held) 

97 

98 nxt, prv = (rank + 1) % L, (rank - 1) % L 

99 for k in range(L): 

100 # after k rotations I am holding the value that started on rank (rank - k) % L 

101 src = (rank - k) % L 

102 for m in range(M): 

103 res[m] += mat[rank, src] * held[m] 

104 

105 if k < L - 1: 

106 comm.Sendrecv(held, dest=nxt, sendtag=k, recvbuf=buf, source=prv, recvtag=k) 

107 held, buf = buf, held 

108 

109 for m in range(M): 

110 me[m] = res[m] 

111 

112 def prepare_Jacobians(self): 

113 """Average the solution across all steps, for constructing average Jacobians.""" 

114 if not self.params.average_jacobian: 

115 return 

116 

117 lvl = self.S.levels[0] 

118 M = lvl.sweep.coll.num_nodes 

119 

120 u_avg = [] 

121 for m in range(M): 

122 contribution = lvl.prob.dtype_u(lvl.u[m + 1]) 

123 total = lvl.prob.dtype_u(lvl.prob.init, val=0) 

124 self.comm.Allreduce(contribution, total, op=MPI.SUM) 

125 u_avg.append(total / self.n_steps) 

126 

127 lvl.u_avg = u_avg 

128 

129 def compute_all_at_once_residual(self): 

130 """ 

131 Compute the residual of the composite collocation problem. 

132 

133 Needs the previous step's end point as this step's initial condition, which is the only 

134 point-to-point communication in a ParaDiag iteration. 

135 """ 

136 S, comm = self.S, self.comm 

137 lvl = S.levels[0] 

138 

139 lvl.sweep.compute_end_point() 

140 

141 for hook in self.hooks: 

142 hook.pre_comm(step=S, level_number=0) 

143 

144 req = None 

145 if not S.status.last: 

146 req = lvl.uend.isend(dest=(comm.rank + 1) % comm.size, tag=S.status.iter, comm=comm) 

147 if not S.status.first: 

148 lvl.u[0].irecv(source=(comm.rank - 1) % comm.size, tag=S.status.iter, comm=comm).Wait() 

149 if req is not None: 

150 req.Wait() 

151 

152 for hook in self.hooks: 

153 hook.post_comm(step=S, level_number=0, add_to_stats=True) 

154 

155 lvl.sweep.compute_residual() 

156 

157 def update_G_inv(self, k=0): 

158 """ 

159 Rebuild this rank's G^-1 if alpha changed with the iteration. 

160 

161 Args: 

162 k (int): 0-based ParaDiag iteration index 

163 """ 

164 alpha = self.get_alpha(k) 

165 if alpha == self._G_inv_alpha: 

166 return 

167 self._G_inv_alpha = alpha 

168 self.S.levels[0].sweep.set_G_inv(get_G_inv_matrix(self.comm.rank, self.comm.size, alpha, self.sweeper_params)) 

169 

170 def update_solution(self): 

171 """Add the increment to get the next iterate. Purely local.""" 

172 lvl = self.S.levels[0] 

173 for m in range(lvl.sweep.coll.num_nodes): 

174 lvl.u[m + 1] += lvl.increment[m] 

175 

176 # ------------------------------------------------------------------ the ParaDiag iteration 

177 

178 def it_ParaDiag(self, comm, num_procs): 

179 """A single ParaDiag iteration, from this rank's point of view.""" 

180 S = self.S 

181 

182 for hook in self.hooks: 

183 hook.pre_sweep(step=S, level_number=0) 

184 

185 # `it_check` has already incremented the counter, so the first sweep is k = 0 

186 k = max(S.status.iter - 1, 0) 

187 self.update_G_inv(k) 

188 

189 self.prepare_Jacobians() 

190 self.compute_all_at_once_residual() 

191 

192 self.FFT_in_time(quantity='residual', k=k) 

193 S.levels[0].sweep.update_nodes() # local solve, embarrassingly parallel 

194 self.iFFT_in_time(quantity='increment', k=k) 

195 

196 self.update_solution() 

197 

198 for hook in self.hooks: 

199 hook.post_sweep(step=S, level_number=0) 

200 

201 S.status.stage = 'IT_CHECK'