Coverage for pySDC/helpers/NCCL_communicator.py: 90%

1from mpi4py import MPI 

2from cupy.cuda import nccl 

3import cupy as cp 

4import numpy as np 

5 

6 

7class NCCLComm(object): 

8 """ 

9 Wraps an MPI communicator and performs some calls to NCCL functions instead. 

10 """ 

11 

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 = {} 

22 

23 def __init__(self, comm): 

24 """ 

25 Args: 

26 comm (mpi4py.Intracomm): MPI communicator 

27 """ 

28 self.commMPI = comm 

29 

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)) 

37 

38 self.commNCCL = NCCLComm._communicators[key][1] 

39 

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. 

44 

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() 

50 

51 return getattr(self.commMPI, name) 

52 

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}!') 

69 

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 

79 

80 def get_op(self, MPI_op): 

81 if MPI_op == MPI.SUM: 

82 return nccl.NCCL_SUM 

83 elif MPI_op == MPI.PROD: 

84 return nccl.NCCL_PROD 

85 elif MPI_op == MPI.MAX: 

86 return nccl.NCCL_MAX 

87 elif MPI_op == MPI.MIN: 

88 return nccl.NCCL_MIN 

89 else: 

90 raise NotImplementedError('Don\'t know what NCCL operation to use to replace this MPI operation!') 

91 

92 def reduce(self, sendobj, op=MPI.SUM, root=0): 

93 sync = False 

94 if hasattr(sendobj, 'data'): 

95 if hasattr(sendobj.data, 'ptr'): 

96 sync = True 

97 if sync: 

98 cp.cuda.Device().synchronize() 

99 

100 return self.commMPI.reduce(sendobj, op=op, root=root) 

101 

102 def allreduce(self, sendobj, op=MPI.SUM): 

103 sync = False 

104 if hasattr(sendobj, 'data'): 

105 if hasattr(sendobj.data, 'ptr'): 

106 sync = True 

107 if sync: 

108 cp.cuda.Device().synchronize() 

109 

110 return self.commMPI.allreduce(sendobj, op=op) 

111 

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

113 if not hasattr(sendbuf.data, 'ptr'): 

114 return self.commMPI.Reduce(sendbuf=sendbuf, recvbuf=recvbuf, op=op, root=root) 

115 

116 dtype = self.get_dtype(sendbuf) 

117 count = self.get_count(sendbuf) 

118 op = self.get_op(op) 

119 recvbuf = cp.empty(1) if recvbuf is None else recvbuf 

120 stream = cp.cuda.get_current_stream() 

121 

122 self.commNCCL.reduce( 

123 sendbuf=sendbuf.data.ptr, 

124 recvbuf=recvbuf.data.ptr, 

125 count=count, 

126 datatype=dtype, 

127 op=op, 

128 root=root, 

129 stream=stream.ptr, 

130 ) 

131 

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

133 if not hasattr(sendbuf.data, 'ptr'): 

134 return self.commMPI.Allreduce(sendbuf=sendbuf, recvbuf=recvbuf, op=op) 

135 

136 dtype = self.get_dtype(sendbuf) 

137 count = self.get_count(sendbuf) 

138 op = self.get_op(op) 

139 stream = cp.cuda.get_current_stream() 

140 

141 self.commNCCL.allReduce( 

142 sendbuf=sendbuf.data.ptr, recvbuf=recvbuf.data.ptr, count=count, datatype=dtype, op=op, stream=stream.ptr 

143 ) 

144 

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

146 if not hasattr(buf.data, 'ptr'): 

147 return self.commMPI.Bcast(buf=buf, root=root) 

148 

149 dtype = self.get_dtype(buf) 

150 count = self.get_count(buf) 

151 stream = cp.cuda.get_current_stream() 

152 

153 self.commNCCL.bcast(buff=buf.data.ptr, count=count, datatype=dtype, root=root, stream=stream.ptr) 

154 

155 def Send(self, buf, dest, tag=0): 

156 """ 

157 Send a buffer to another rank through NCCL. 

158 

159 NCCL has no tags: a send is matched to whichever receive the destination posts next for 

160 this pair of ranks. That is fine when one message is in flight at a time, and wrong when 

161 several are, which is why the non-blocking `Issend`/`Irecv` that pySDC's time-parallel 

162 controller uses are left to fall through to MPI, where tags mean what they say. 

163 """ 

164 if not hasattr(buf.data, 'ptr'): 

165 return self.commMPI.Send(buf, dest=dest, tag=tag) 

166 

167 stream = cp.cuda.get_current_stream() 

168 self.commNCCL.send(buf.data.ptr, self.get_count(buf), self.get_dtype(buf), dest, stream.ptr) 

169 stream.synchronize() 

170 

171 def Recv(self, buf, source, tag=0): 

172 """ 

173 Receive a buffer from another rank through NCCL. See `Send` on the absence of tags. 

174 """ 

175 if not hasattr(buf.data, 'ptr'): 

176 return self.commMPI.Recv(buf, source=source, tag=tag) 

177 

178 stream = cp.cuda.get_current_stream() 

179 self.commNCCL.recv(buf.data.ptr, self.get_count(buf), self.get_dtype(buf), source, stream.ptr) 

180 stream.synchronize() 

181 

182 def Barrier(self): 

183 cp.cuda.get_current_stream().synchronize() 

184 self.commMPI.Barrier()