remove nccl_comm from Op constructor and take base_nccl_comm from operand x instead

Author: tharittk

pylops_mpi/basicoperators/HStack.py

@@ -5,7 +5,6 @@
 from pylops.utils import DTypeLike
 
 from pylops_mpi import DistributedArray, MPILinearOperator
-from pylops_mpi.DistributedArray import NcclCommunicatorType
 from .VStack import MPIVStack
 
 
@@ -90,15 +89,14 @@ class MPIHStack(MPILinearOperator):
 
     def __init__(self, ops: Sequence[LinearOperator],
                  base_comm: MPI.Comm = MPI.COMM_WORLD,
-                 base_comm_nccl: NcclCommunicatorType = None,
                  dtype: Optional[DTypeLike] = None):
         self.ops = ops
         nops = [oper.shape[0] for oper in self.ops]
         nops = np.concatenate(base_comm.allgather(nops), axis=0)
         if len(set(nops)) > 1:
             raise ValueError("Operators have different number of rows")
         hops = [oper.H for oper in self.ops]
-        self.HStack = MPIVStack(ops=hops, base_comm=base_comm, base_comm_nccl=base_comm_nccl, dtype=dtype).H
+        self.HStack = MPIVStack(ops=hops, base_comm=base_comm, dtype=dtype).H
         super().__init__(shape=self.HStack.shape, dtype=self.HStack.dtype, base_comm=base_comm)
 
     def _matvec(self, x: DistributedArray) -> DistributedArray:

pylops_mpi/basicoperators/VStack.py

@@ -16,7 +16,6 @@
     StackedDistributedArray
 )
 from pylops_mpi.utils.decorators import reshaped
-from pylops_mpi.DistributedArray import NcclCommunicatorType
 from pylops_mpi.utils import deps
 
 cupy_message = pylops_deps.cupy_import("the VStack module")
@@ -40,8 +39,6 @@ class MPIVStack(MPILinearOperator):
         One or more :class:`pylops.LinearOperator` to be vertically stacked.
     base_comm : :obj:`mpi4py.MPI.Comm`, optional
         Base MPI Communicator. Defaults to ``mpi4py.MPI.COMM_WORLD``.
-    base_comm_nccl : :obj:`cupy.cuda.nccl.NcclCommunicator`, optional
-        NCCL Communicator over which operators and arrays are distributed.
     dtype : :obj:`str`, optional
         Type of elements in input array.
 
@@ -110,10 +107,8 @@ class MPIVStack(MPILinearOperator):
 
     def __init__(self, ops: Sequence[LinearOperator],
                  base_comm: MPI.Comm = MPI.COMM_WORLD,
-                 base_comm_nccl: NcclCommunicatorType = None,
                  dtype: Optional[DTypeLike] = None):
         self.ops = ops
-        self.base_comm_nccl = base_comm_nccl
         nops = np.zeros(len(self.ops), dtype=np.int64)
         for iop, oper in enumerate(self.ops):
             nops[iop] = oper.shape[0]
@@ -152,8 +147,8 @@ def _rmatvec(self, x: DistributedArray) -> DistributedArray:
         for iop, oper in enumerate(self.ops):
             y1.append(oper.rmatvec(x.local_array[self.nnops[iop]: self.nnops[iop + 1]]))
         y1 = ncp.sum(ncp.vstack(y1), axis=0)
-        if deps.nccl_enabled and self.base_comm_nccl:
-            y[:] = nccl_allreduce(self.base_comm_nccl, y1, op=MPI.SUM)
+        if deps.nccl_enabled and x.base_comm_nccl:
+            y[:] = nccl_allreduce(x.base_comm_nccl, y1, op=MPI.SUM)
         else:
             y[:] = self.base_comm.allreduce(y1, op=MPI.SUM)
         return y

pylops_mpi/utils/decorators.py

@@ -54,6 +54,7 @@ def wrapper(self, x: DistributedArray):
                 local_shapes = None
                 global_shape = getattr(self, "dims")
             arr = DistributedArray(global_shape=global_shape,
+                                   base_comm_nccl=x.base_comm_nccl,
                                    local_shapes=local_shapes, axis=0,
                                    engine=x.engine, dtype=x.dtype)
             arr_local_shapes = np.asarray(arr.base_comm.allgather(np.prod(arr.local_shape)))

tests_nccl/test_stack_nccl.py

@@ -30,7 +30,7 @@ def test_vstack_nccl(par):
     rank = MPI.COMM_WORLD.Get_rank()
     A_gpu = cp.ones(shape=(par['ny'], par['nx'])) + par['imag'] * cp.ones(shape=(par['ny'], par['nx']))
     Op = pylops.MatrixMult(A=((rank + 1) * A_gpu).astype(par['dtype']))
-    VStack_MPI = pylops_mpi.MPIVStack(ops=[Op, ], base_comm_nccl=nccl_comm)
+    VStack_MPI = pylops_mpi.MPIVStack(ops=[Op, ], )
 
     # Broadcasted DistributedArray(global_shape == local_shape)
     x = pylops_mpi.DistributedArray(global_shape=par['nx'],
@@ -80,7 +80,7 @@ def test_stacked_vstack_nccl(par):
     rank = MPI.COMM_WORLD.Get_rank()
     A_gpu = cp.ones(shape=(par['ny'], par['nx'])) + par['imag'] * cp.ones(shape=(par['ny'], par['nx']))
     Op = pylops.MatrixMult(A=((rank + 1) * A_gpu).astype(par['dtype']))
-    VStack_MPI = pylops_mpi.MPIVStack(ops=[Op, ], base_comm_nccl=nccl_comm)
+    VStack_MPI = pylops_mpi.MPIVStack(ops=[Op, ], )
     StackedVStack_MPI = pylops_mpi.MPIStackedVStack([VStack_MPI, VStack_MPI])
 
     # Broadcasted DistributedArray(global_shape == local_shape)
@@ -121,13 +121,13 @@ def test_stacked_vstack_nccl(par):
 
 @pytest.mark.mpi(min_size=2)
 @pytest.mark.parametrize("par", [(par1), (par2)])
-def test_hstack(par):
+def test_hstack_nccl(par):
     """Test the MPIHStack operator with NCCL"""
     size = MPI.COMM_WORLD.Get_size()
     rank = MPI.COMM_WORLD.Get_rank()
     A_gpu = cp.ones(shape=(par['ny'], par['nx'])) + par['imag'] * cp.ones(shape=(par['ny'], par['nx']))
     Op = pylops.MatrixMult(A=((rank + 1) * A_gpu).astype(par['dtype']))
-    HStack_MPI = pylops_mpi.MPIHStack(ops=[Op, ], base_comm_nccl=nccl_comm)
+    HStack_MPI = pylops_mpi.MPIHStack(ops=[Op, ], )
 
     # Scattered DistributedArray
     x = pylops_mpi.DistributedArray(global_shape=size * par['nx'],