Coverage for pySDC/helpers/NCCL_communicator.py: 90%
98 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1from mpi4py import MPI
2from cupy.cuda import nccl
3import cupy as cp
4import numpy as np
7class NCCLComm(object):
8 """
9 Wraps an MPI communicator and performs some calls to NCCL functions instead.
10 """
12 #: One NCCL communicator per MPI communicator, keyed by the MPI handle.
13 #:
14 #: Creating one costs about 17 MB of device memory that is never released: NCCL has no
15 #: reference counting, and destroying a communicator is itself collective, so a `__del__`
16 #: would have ranks tearing them down whenever their garbage collectors happened to run --
17 #: in different orders, which deadlocks. Making them once and sharing them avoids both.
18 #:
19 #: The MPI communicator is kept alongside so it cannot be freed and have its handle reused
20 #: for a different one, which would hand out the wrong NCCL communicator.
21 _communicators = {}
23 def __init__(self, comm):
24 """
25 Args:
26 comm (mpi4py.Intracomm): MPI communicator
27 """
28 self.commMPI = comm
30 # `py2f` rather than the communicator itself: mpi4py defines `__eq__` without `__hash__`,
31 # so a communicator cannot be a dictionary key, and the handle is the same for any two
32 # Python wrappers around one communicator.
33 key = comm.py2f()
34 if key not in NCCLComm._communicators:
35 uid = comm.bcast(nccl.get_unique_id(), root=0)
36 NCCLComm._communicators[key] = (comm, nccl.NcclCommunicator(comm.size, uid, comm.rank))
38 self.commNCCL = NCCLComm._communicators[key][1]
40 def __getattr__(self, name):
41 """
42 Pass calls that are not explicitly overridden by NCCL functionality on to the MPI communicator.
43 When performing any operations that depend on data, we have to synchronize host and device beforehand.
45 Args:
46 name (str): Name of the requested attribute
47 """
48 if name not in ['size', 'rank', 'Get_rank', 'Get_size', 'Split', 'Create_cart', 'Is_inter', 'Get_topology']:
49 cp.cuda.get_current_stream().synchronize()
51 return getattr(self.commMPI, name)
53 @staticmethod
54 def get_dtype(data):
55 """
56 As NCCL doesn't support complex numbers, we have to act as if we're sending two real numbers if using complex.
57 """
58 dtype = data.dtype
59 if dtype in [np.dtype('float32'), np.dtype('complex64')]:
60 return nccl.NCCL_FLOAT32
61 elif dtype in [np.dtype('float64'), np.dtype('complex128')]:
62 return nccl.NCCL_FLOAT64
63 elif dtype in [np.dtype('int32')]:
64 return nccl.NCCL_INT32
65 elif dtype in [np.dtype('int64')]:
66 return nccl.NCCL_INT64
67 else:
68 raise NotImplementedError(f'Don\'t know what NCCL dtype to use to send data of dtype {data.dtype}!')
70 @staticmethod
71 def get_count(data):
72 """
73 As NCCL doesn't support complex numbers, we have to act as if we're sending two real numbers if using complex.
74 """
75 if cp.iscomplexobj(data):
76 return data.size * 2
77 else:
78 return data.size
80 def get_op(self, MPI_op, data=None):
81 """
82 Translate an MPI reduction operation to the NCCL one.
84 Complex data is reduced as twice as many reals, which is only right for sums: the product of two complex
85 numbers is not the pairwise product of their parts, and they have no order for `MAX` and `MIN`.
87 Args:
88 MPI_op (mpi4py.MPI.Op): One of `MPI.SUM`, `MPI.PROD`, `MPI.MAX` and `MPI.MIN`
89 data (cupy.ndarray, optional): The data to be reduced, to check that the operation works for its dtype
91 Returns:
92 NCCL reduction operation
93 """
94 if data is not None and cp.iscomplexobj(data) and MPI_op != MPI.SUM:
95 raise NotImplementedError('NCCL can only sum complex data, since it reduces it as pairs of reals!')
97 if MPI_op == MPI.SUM:
98 return nccl.NCCL_SUM
99 elif MPI_op == MPI.PROD:
100 return nccl.NCCL_PROD
101 elif MPI_op == MPI.MAX:
102 return nccl.NCCL_MAX
103 elif MPI_op == MPI.MIN:
104 return nccl.NCCL_MIN
105 else:
106 raise NotImplementedError('Don\'t know what NCCL operation to use to replace this MPI operation!')
108 def reduce(self, sendobj, op=MPI.SUM, root=0):
109 """Wrap mpi4py's pickle-based `reduce`, synchronizing the device first if `sendobj` holds CuPy data."""
110 sync = False
111 if hasattr(sendobj, 'data'):
112 if hasattr(sendobj.data, 'ptr'):
113 sync = True
114 if sync:
115 cp.cuda.Device().synchronize()
117 return self.commMPI.reduce(sendobj, op=op, root=root)
119 def allreduce(self, sendobj, op=MPI.SUM):
120 """Wrap mpi4py's pickle-based `allreduce`, synchronizing the device first if `sendobj` holds CuPy data."""
121 sync = False
122 if hasattr(sendobj, 'data'):
123 if hasattr(sendobj.data, 'ptr'):
124 sync = True
125 if sync:
126 cp.cuda.Device().synchronize()
128 return self.commMPI.allreduce(sendobj, op=op)
130 def Reduce(self, sendbuf, recvbuf, op=MPI.SUM, root=0):
131 """
132 Wrap mpi4py's `Reduce`, going through NCCL if `sendbuf` is a CuPy array and through MPI otherwise.
134 The NCCL reduction is enqueued on the current CuPy stream without synchronizing. It reduces complex data as
135 pairs of
136 reals, so `MPI.PROD` is not the complex product. `recvbuf` may be None on ranks other than `root`.
137 """
138 if not hasattr(sendbuf.data, 'ptr'):
139 return self.commMPI.Reduce(sendbuf=sendbuf, recvbuf=recvbuf, op=op, root=root)
141 dtype = self.get_dtype(sendbuf)
142 count = self.get_count(sendbuf)
143 op = self.get_op(op, sendbuf)
144 recvbuf = cp.empty(1) if recvbuf is None else recvbuf
145 stream = cp.cuda.get_current_stream()
147 self.commNCCL.reduce(
148 sendbuf=sendbuf.data.ptr,
149 recvbuf=recvbuf.data.ptr,
150 count=count,
151 datatype=dtype,
152 op=op,
153 root=root,
154 stream=stream.ptr,
155 )
157 def Allreduce(self, sendbuf, recvbuf, op=MPI.SUM):
158 """
159 Wrap mpi4py's `Allreduce`, going through NCCL if `sendbuf` is a CuPy array and through MPI otherwise.
161 The NCCL reduction is enqueued on the current CuPy stream without synchronizing. It reduces complex data as
162 pairs of
163 reals, so `MPI.PROD` is not the complex product.
164 """
165 if not hasattr(sendbuf.data, 'ptr'):
166 return self.commMPI.Allreduce(sendbuf=sendbuf, recvbuf=recvbuf, op=op)
168 dtype = self.get_dtype(sendbuf)
169 count = self.get_count(sendbuf)
170 op = self.get_op(op, sendbuf)
171 stream = cp.cuda.get_current_stream()
173 self.commNCCL.allReduce(
174 sendbuf=sendbuf.data.ptr, recvbuf=recvbuf.data.ptr, count=count, datatype=dtype, op=op, stream=stream.ptr
175 )
177 def Bcast(self, buf, root=0):
178 """
179 Wrap mpi4py's `Bcast`, going through NCCL on the current CuPy stream, without synchronizing, if `buf` is a CuPy
180 array and through MPI otherwise.
181 """
182 if not hasattr(buf.data, 'ptr'):
183 return self.commMPI.Bcast(buf=buf, root=root)
185 dtype = self.get_dtype(buf)
186 count = self.get_count(buf)
187 stream = cp.cuda.get_current_stream()
189 self.commNCCL.bcast(buff=buf.data.ptr, count=count, datatype=dtype, root=root, stream=stream.ptr)
191 def Send(self, buf, dest, tag=0):
192 """
193 Send a buffer to another rank through NCCL.
195 NCCL has no tags: a send is matched to whichever receive the destination posts next for
196 this pair of ranks. That is fine when one message is in flight at a time, and wrong when
197 several are, which is why the non-blocking `Issend`/`Irecv` that pySDC's time-parallel
198 controller uses are left to fall through to MPI, where tags mean what they say.
199 """
200 if not hasattr(buf.data, 'ptr'):
201 return self.commMPI.Send(buf, dest=dest, tag=tag)
203 stream = cp.cuda.get_current_stream()
204 self.commNCCL.send(buf.data.ptr, self.get_count(buf), self.get_dtype(buf), dest, stream.ptr)
205 stream.synchronize()
207 def Recv(self, buf, source, tag=0):
208 """
209 Receive a buffer from another rank through NCCL. See `Send` on the absence of tags.
210 """
211 if not hasattr(buf.data, 'ptr'):
212 return self.commMPI.Recv(buf, source=source, tag=tag)
214 stream = cp.cuda.get_current_stream()
215 self.commNCCL.recv(buf.data.ptr, self.get_count(buf), self.get_dtype(buf), source, stream.ptr)
216 stream.synchronize()
218 def Barrier(self):
219 """Synchronize the current CuPy stream, then wrap mpi4py's `Barrier`."""
220 cp.cuda.get_current_stream().synchronize()
221 self.commMPI.Barrier()