nccl for HStack Op

Author: tharittk

pylops_mpi/basicoperators/HStack.py

@@ -5,6 +5,7 @@
 from pylops.utils import DTypeLike
 
 from pylops_mpi import DistributedArray, MPILinearOperator
+from pylops_mpi.DistributedArray import NcclCommunicatorType
 from .VStack import MPIVStack
 
 
@@ -89,14 +90,15 @@ 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, dtype=dtype).H
+        self.HStack = MPIVStack(ops=hops, base_comm=base_comm, base_comm_nccl=base_comm_nccl, dtype=dtype).H
         super().__init__(shape=self.HStack.shape, dtype=self.HStack.dtype, base_comm=base_comm)
 
     def _matvec(self, x: DistributedArray) -> DistributedArray:

tests_nccl/test_stack_nccl.py

@@ -119,4 +119,46 @@ def test_stacked_vstack_nccl(par):
         assert_allclose(y_rmat_mpi.get(), y_rmat_np, rtol=1e-14)
 
 
-# TODO: Test of HStack
+@pytest.mark.mpi(min_size=2)
+@pytest.mark.parametrize("par", [(par1), (par2)])
+def test_hstack(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)
+
+    # Scattered DistributedArray
+    x = pylops_mpi.DistributedArray(global_shape=size * par['nx'],
+                                    base_comm_nccl=nccl_comm,
+                                    partition=pylops_mpi.Partition.SCATTER,
+                                    dtype=par['dtype'],
+                                    engine="cupy")
+    x[:] = cp.ones(shape=par['nx'], dtype=par['dtype'])
+    x_global = x.asarray()
+
+    # Broadcasted DistributedArray(global_shape == local_shape)
+    y = pylops_mpi.DistributedArray(global_shape=par['ny'],
+                                    base_comm_nccl=nccl_comm,
+                                    partition=pylops_mpi.Partition.BROADCAST,
+                                    dtype=par['dtype'],
+                                    engine="cupy")
+    y[:] = cp.ones(shape=par['ny'], dtype=par['dtype'])
+    y_global = y.asarray()
+
+    x_mat = HStack_MPI @ x
+    y_rmat = HStack_MPI.H @ y
+    assert isinstance(x_mat, pylops_mpi.DistributedArray)
+    assert isinstance(y_rmat, pylops_mpi.DistributedArray)
+
+    x_mat_mpi = x_mat.asarray()
+    y_rmat_mpi = y_rmat.asarray()
+
+    if rank == 0:
+        ops = [pylops.MatrixMult(A=((i + 1) * A_gpu.get()).astype(par['dtype'])) for i in range(size)]
+        HStack = pylops.HStack(ops=ops)
+        x_mat_np = HStack @ x_global.get()
+        y_rmat_np = HStack.H @ y_global.get()
+        assert_allclose(x_mat_mpi.get(), x_mat_np, rtol=1e-14)
+        assert_allclose(y_rmat_mpi.get(), y_rmat_np, rtol=1e-14)