Coverage for pySDC/implementations/controller_classes/controller_ParaDiag_MPI.py: 100%
93 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-19 17:32 +0000
1import numpy as np
2from mpi4py import MPI
4from pySDC.core.errors import ControllerError
5from pySDC.helpers.ParaDiagHelper import get_G_inv_matrix
6from pySDC.implementations.controller_classes.ParaDiag import ParaDiag
7from pySDC.implementations.controller_classes.controller_MPI import controller_MPI
10class controller_ParaDiag_MPI(ParaDiag, controller_MPI):
11 """
12 ParaDiag controller with MPI parallelism across time steps: one step per rank.
14 This is `controller_MPI` with a different iteration: where PFASST sweeps and cascades through
15 the levels, ParaDiag diagonalizes across the steps. Everything around the iteration -- blocks,
16 windowing, restarts, convergence -- is the driver it inherits, which is why the dispatcher it
17 uses is still called `pfasst`.
19 Everything here is written from a single processor's point of view. A rank owns exactly one step
20 and never inspects another rank's data; the places where ParaDiag genuinely needs information
21 from the whole block are expressed as communication:
23 - ``prepare_Jacobians`` -> Allreduce(SUM) over the step communicator
24 - ``compute_all_at_once_residual`` -> point-to-point exchange with the previous/next rank
25 - ``apply_matrix`` (the weighted FFT/iFFT in time) -> a ring reduction, see below
26 - convergence -> the inherited `it_check`, which allreduces because
27 `all_to_done` is forced on for ParaDiag
29 Note that ParaDiag steps can only converge together, so every rank always participates in every
30 iteration. That is not a policy choice: a rank that stopped early would never enter the
31 collectives below and the run would hang. `step_is_active` says the same thing about blocks --
32 a block is never run partially, so the driver's windowing never splits one.
33 """
35 def __init__(self, controller_params, description, comm=None):
36 """
37 Args:
38 controller_params: parameter set for the controller and the steps
39 description: all the parameters to set up the rest (levels, problems, transfer, ...)
40 comm: MPI communicator, one rank per time step
41 """
42 comm = MPI.COMM_WORLD if comm is None else comm
43 self.prepare_ParaDiag_params(controller_params, description)
45 self.sweeper_params = description['sweeper_params']
47 # each step needs its own G^-1, determined by where it sits in the block
48 self._G_inv_alpha = self.resolve_alpha(controller_params['alpha'], 0)
49 description['sweeper_params']['G_inv'] = get_G_inv_matrix(
50 comm.rank, comm.size, self._G_inv_alpha, description['sweeper_params']
51 )
53 super().__init__(controller_params, description, comm)
55 self.n_steps = comm.size
57 if len(self.S.levels) > 1:
58 raise ControllerError('Multi-level SDC not implemented in ParaDiag!')
60 # ------------------------------------------------------------------ collectives
62 def apply_matrix(self, mat, quantity):
63 """
64 Apply a square L x L matrix across the steps, where L is the number of ranks.
66 Each rank needs ``res_i = sum_j mat[i, j] * me_j`` but only holds ``me_i``. This is done as a
67 ring reduction: the values circulate once around the communicator and each rank accumulates
68 its own row as they pass. That keeps the working set at O(M) fields per rank, independent of
69 L -- an allgather would instead need O(L * M) fields on every rank, which is exactly the
70 gather this controller must not do.
72 The ring costs L - 1 rounds. A butterfly would need only log2(L), but only because the matrix
73 this is called with is a DFT; for a general matrix it needs the O(L * M) gather just ruled
74 out. That belongs with an FFT-shaped interface rather than this one, and pays off from about
75 L = 16 upwards.
77 Args:
78 mat: square matrix with as many rows as there are ranks
79 quantity (str): 'residual' or 'increment', the level attribute to transform in place
80 """
81 comm = self.comm
82 L, rank = comm.size, comm.rank
83 assert np.allclose(mat.shape, L), f'need a {L}x{L} matrix, got {mat.shape}'
85 lvl = self.S.levels[0]
86 M = lvl.sweep.params.num_nodes
87 prob = lvl.prob
89 if quantity == 'residual':
90 me = lvl.residual
91 elif quantity == 'increment':
92 me = lvl.increment
93 else:
94 raise NotImplementedError(f'Cannot apply matrix to {quantity!r}')
96 res = [prob.u_init for _ in range(M)]
97 # all M nodes travel in one contiguous buffer, so the ring costs L - 1 messages rather than
98 # M * (L - 1). Same volume, M times fewer message latencies.
99 held = np.array([me[m] for m in range(M)])
100 buf = np.empty_like(held)
102 nxt, prv = (rank + 1) % L, (rank - 1) % L
103 for k in range(L):
104 # after k rotations I am holding the value that started on rank (rank - k) % L
105 src = (rank - k) % L
106 for m in range(M):
107 res[m] += mat[rank, src] * held[m]
109 if k < L - 1:
110 comm.Sendrecv(held, dest=nxt, sendtag=k, recvbuf=buf, source=prv, recvtag=k)
111 held, buf = buf, held
113 for m in range(M):
114 me[m] = res[m]
116 def prepare_Jacobians(self):
117 """Average the solution across all steps, for constructing average Jacobians."""
118 if not self.params.average_jacobian:
119 return
121 lvl = self.S.levels[0]
122 M = lvl.sweep.coll.num_nodes
124 u_avg = []
125 for m in range(M):
126 contribution = lvl.prob.dtype_u(lvl.u[m + 1])
127 total = lvl.prob.dtype_u(lvl.prob.init, val=0)
128 self.comm.Allreduce(contribution, total, op=MPI.SUM)
129 u_avg.append(total / self.n_steps)
131 lvl.u_avg = u_avg
133 def compute_all_at_once_residual(self):
134 """
135 Compute the residual of the composite collocation problem.
137 Needs the previous step's end point as this step's initial condition, which is the only
138 point-to-point communication in a ParaDiag iteration.
139 """
140 S, comm = self.S, self.comm
141 lvl = S.levels[0]
143 lvl.sweep.compute_end_point()
145 for hook in self.hooks:
146 hook.pre_comm(step=S, level_number=0)
148 req = None
149 if not S.status.last:
150 req = lvl.uend.isend(dest=(comm.rank + 1) % comm.size, tag=S.status.iter, comm=comm)
151 if not S.status.first:
152 lvl.u[0].irecv(source=(comm.rank - 1) % comm.size, tag=S.status.iter, comm=comm).Wait()
153 if req is not None:
154 req.Wait()
156 for hook in self.hooks:
157 hook.post_comm(step=S, level_number=0, add_to_stats=True)
159 lvl.sweep.compute_residual()
161 def update_G_inv(self, k=0):
162 """
163 Rebuild this rank's G^-1 if alpha changed with the iteration.
165 Args:
166 k (int): 0-based ParaDiag iteration index
167 """
168 alpha = self.get_alpha(k)
169 if alpha == self._G_inv_alpha:
170 return
171 self._G_inv_alpha = alpha
172 self.S.levels[0].sweep.set_G_inv(get_G_inv_matrix(self.comm.rank, self.comm.size, alpha, self.sweeper_params))
174 def update_solution(self):
175 """Add the increment to get the next iterate. Purely local."""
176 lvl = self.S.levels[0]
177 for m in range(lvl.sweep.coll.num_nodes):
178 lvl.u[m + 1] += lvl.increment[m]
180 # ------------------------------------------------------------------ the ParaDiag iteration
182 def it_ParaDiag(self, comm, num_procs):
183 """A single ParaDiag iteration, from this rank's point of view."""
184 S = self.S
186 for hook in self.hooks:
187 hook.pre_sweep(step=S, level_number=0)
189 # `it_check` has already incremented the counter, so the first sweep is k = 0
190 k = max(S.status.iter - 1, 0)
191 self.update_G_inv(k)
193 self.prepare_Jacobians()
194 self.compute_all_at_once_residual()
196 self.FFT_in_time(quantity='residual', k=k)
197 S.levels[0].sweep.update_nodes() # local solve, embarrassingly parallel
198 self.iFFT_in_time(quantity='increment', k=k)
200 self.update_solution()
202 for hook in self.hooks:
203 hook.post_sweep(step=S, level_number=0)
205 S.status.stage = 'IT_CHECK'