fix groupStart/groupEnd
@@ -252,10 +252,10 @@ class NCCLBackend(_Backend):
stream = self._get_stream(stream)
idtype, icount = self._get_nccl_dtype_and_count(in_array)
odtype, ocount = self._get_nccl_dtype_and_count(out_array)
- self._comm.groupStart()
+ nccl.groupStart()
self._send(in_array, peer, idtype, icount, stream)
self._recv(out_array, peer, odtype, ocount, stream)
- self._comm.groupEnd()
+ nccl.groupEnd()
def scatter(self, in_array, out_array, root=0, stream=None):
"""Performs a scatter operation.
@@ -275,7 +275,7 @@ class NCCLBackend(_Backend):
self._check_contiguous(in_array)
self._check_contiguous(out_array)
stream = self._get_stream(stream)
- self._comm.groupStart()
+ nccl.groupStart()
if root == self.rank:
for i in range(self._n_devices):
array = in_array[i]
@@ -283,7 +283,7 @@ class NCCLBackend(_Backend):
self._send(array, i, idtype, icount, stream)
dtype, count = self._get_nccl_dtype_and_count(out_array)
self._recv(out_array, root, dtype, count, stream)
- self._comm.groupEnd()
+ nccl.groupEnd()
def gather(self, in_array, out_array, root=0, stream=None):
"""Performs a gather operation.
@@ -304,7 +304,7 @@ class NCCLBackend(_Backend):
self._check_contiguous(in_array)
self._check_contiguous(out_array)
stream = self._get_stream(stream)
- self._comm.groupStart()
+ nccl.groupStart()
if root == self.rank:
for i in range(self._n_devices):
array = out_array[i]
@@ -312,7 +312,7 @@ class NCCLBackend(_Backend):
self._recv(array, i, odtype, ocount, stream)
dtype, count = self._get_nccl_dtype_and_count(in_array)
self._send(in_array, root, dtype, count, stream)
- self._comm.groupEnd()
+ nccl.groupEnd()
def all_to_all(self, in_array, out_array, stream=None):
"""Performs an all to all operation.
@@ -340,11 +340,11 @@ class NCCLBackend(_Backend):
idtype, icount = self._get_nccl_dtype_and_count(in_array[0])
odtype, ocount = self._get_nccl_dtype_and_count(out_array[0])
# TODO check out dtypes are the same as in dtypes
- self._comm.groupStart()
+ nccl.groupStart()
for i in range(self._n_devices):
self._send(in_array[i], i, idtype, icount, stream)
self._recv(out_array[i], i, odtype, ocount, stream)
- self._comm.groupEnd()
+ nccl.groupEnd()
def barrier(self):
"""Performs a barrier operation.
