nccl support for SecondDerivative and test_derivative_nccl for first and second order

Author: tharittk

pylops_mpi/DistributedArray.py

@@ -506,7 +506,7 @@ def _allgather(self, send_buf, recv_buf=None):
     def _send(self, send_buf, dest, count=None, tag=None):
         """ Send operation
         """
-        if deps.nccl_enabled and getattr(self, "base_comm_nccl"):
+        if deps.nccl_enabled and self.base_comm_nccl:
             if count is None:
                 # assuming sending the whole array
                 count = send_buf.size
@@ -519,7 +519,7 @@ def _recv(self, recv_buf=None, source=0, count=None, tag=None):
         """
         # NCCL must be called with recv_buf. Size cannot be inferred from
         # other arguments and thus cannot be dynamically allocated
-        if deps.nccl_enabled and getattr(self, "base_comm_nccl") and recv_buf is not None:
+        if deps.nccl_enabled and self.base_comm_nccl and recv_buf is not None:
             if recv_buf is not None:
                 if count is None:
                     # assuming data will take a space of the whole buffer
@@ -572,6 +572,7 @@ def add(self, dist_array):
         self._check_mask(dist_array)
         SumArray = DistributedArray(global_shape=self.global_shape,
                                     base_comm=self.base_comm,
+                                    base_comm_nccl=self.base_comm_nccl,
                                     dtype=self.dtype,
                                     partition=self.partition,
                                     local_shapes=self.local_shapes,
@@ -598,6 +599,7 @@ def multiply(self, dist_array):
 
         ProductArray = DistributedArray(global_shape=self.global_shape,
                                         base_comm=self.base_comm,
+                                        base_comm_nccl=self.base_comm_nccl,
                                         dtype=self.dtype,
                                         partition=self.partition,
                                         local_shapes=self.local_shapes,
@@ -748,6 +750,7 @@ def ravel(self, order: Optional[str] = "C"):
         """
         local_shapes = [(np.prod(local_shape, axis=-1), ) for local_shape in self.local_shapes]
         arr = DistributedArray(global_shape=np.prod(self.global_shape),
+                               base_comm_nccl=self.base_comm_nccl,
                                local_shapes=local_shapes,
                                mask=self.mask,
                                partition=self.partition,

pylops_mpi/basicoperators/SecondDerivative.py

@@ -112,20 +112,20 @@ def _register_multiplications(
     def _matvec(self, x: DistributedArray) -> DistributedArray:
         # If Partition.BROADCAST, then convert to Partition.SCATTER
         if x.partition is Partition.BROADCAST:
-            x = DistributedArray.to_dist(x=x.local_array)
+            x = DistributedArray.to_dist(x=x.local_array, base_comm_nccl=x.base_comm_nccl)
         return self._hmatvec(x)
 
     def _rmatvec(self, x: DistributedArray) -> DistributedArray:
         # If Partition.BROADCAST, then convert to Partition.SCATTER
         if x.partition is Partition.BROADCAST:
-            x = DistributedArray.to_dist(x=x.local_array)
+            x = DistributedArray.to_dist(x=x.local_array, base_comm_nccl=x.base_comm_nccl)
         return self._hrmatvec(x)
 
     @reshaped
     def _matvec_forward(self, x: DistributedArray) -> DistributedArray:
         ncp = get_module(x.engine)
-        y = DistributedArray(global_shape=x.global_shape, local_shapes=x.local_shapes,
-                             axis=x.axis, engine=x.engine, dtype=self.dtype)
+        y = DistributedArray(global_shape=x.global_shape, base_comm_nccl=x.base_comm_nccl,
+                             local_shapes=x.local_shapes, axis=x.axis, engine=x.engine, dtype=self.dtype)
         ghosted_x = x.add_ghost_cells(cells_back=2)
         y_forward = ghosted_x[2:] - 2 * ghosted_x[1:-1] + ghosted_x[:-2]
         if self.rank == self.size - 1:
@@ -136,7 +136,8 @@ def _matvec_forward(self, x: DistributedArray) -> DistributedArray:
     @reshaped
     def _rmatvec_forward(self, x: DistributedArray) -> DistributedArray:
         ncp = get_module(x.engine)
-        y = DistributedArray(global_shape=x.global_shape, local_shapes=x.local_shapes, axis=x.axis, dtype=self.dtype)
+        y = DistributedArray(global_shape=x.global_shape, base_comm_nccl=x.base_comm_nccl,
+                             local_shapes=x.local_shapes, axis=x.axis, engine=x.engine, dtype=self.dtype)
         y[:] = 0
         if self.rank == self.size - 1:
             y[:-2] += x[:-2]
@@ -162,8 +163,8 @@ def _rmatvec_forward(self, x: DistributedArray) -> DistributedArray:
     @reshaped
     def _matvec_backward(self, x: DistributedArray) -> DistributedArray:
         ncp = get_module(x.engine)
-        y = DistributedArray(global_shape=x.global_shape, local_shapes=x.local_shapes,
-                             axis=x.axis, engine=x.engine, dtype=self.dtype)
+        y = DistributedArray(global_shape=x.global_shape, base_comm_nccl=x.base_comm_nccl,
+                             local_shapes=x.local_shapes, axis=x.axis, engine=x.engine, dtype=self.dtype)
         ghosted_x = x.add_ghost_cells(cells_front=2)
         y_backward = ghosted_x[2:] - 2 * ghosted_x[1:-1] + ghosted_x[:-2]
         if self.rank == 0:
@@ -174,8 +175,8 @@ def _matvec_backward(self, x: DistributedArray) -> DistributedArray:
     @reshaped
     def _rmatvec_backward(self, x: DistributedArray) -> DistributedArray:
         ncp = get_module(x.engine)
-        y = DistributedArray(global_shape=x.global_shape, local_shapes=x.local_shapes,
-                             axis=x.axis, engine=x.engine, dtype=self.dtype)
+        y = DistributedArray(global_shape=x.global_shape, base_comm_nccl=x.base_comm_nccl,
+                             local_shapes=x.local_shapes, axis=x.axis, engine=x.engine, dtype=self.dtype)
         y[:] = 0
         ghosted_x = x.add_ghost_cells(cells_back=2)
         y_backward = ghosted_x[2:]
@@ -201,8 +202,8 @@ def _rmatvec_backward(self, x: DistributedArray) -> DistributedArray:
     @reshaped
     def _matvec_centered(self, x: DistributedArray) -> DistributedArray:
         ncp = get_module(x.engine)
-        y = DistributedArray(global_shape=x.global_shape, local_shapes=x.local_shapes,
-                             axis=x.axis, engine=x.engine, dtype=self.dtype)
+        y = DistributedArray(global_shape=x.global_shape, base_comm_nccl=x.base_comm_nccl,
+                             local_shapes=x.local_shapes, axis=x.axis, engine=x.engine, dtype=self.dtype)
         ghosted_x = x.add_ghost_cells(cells_front=1, cells_back=1)
         y_centered = ghosted_x[2:] - 2 * ghosted_x[1:-1] + ghosted_x[:-2]
         if self.rank == 0:
@@ -221,8 +222,8 @@ def _matvec_centered(self, x: DistributedArray) -> DistributedArray:
     @reshaped
     def _rmatvec_centered(self, x: DistributedArray) -> DistributedArray:
         ncp = get_module(x.engine)
-        y = DistributedArray(global_shape=x.global_shape, local_shapes=x.local_shapes,
-                             axis=x.axis, engine=x.engine, dtype=self.dtype)
+        y = DistributedArray(global_shape=x.global_shape, base_comm_nccl=x.base_comm_nccl,
+                             local_shapes=x.local_shapes, axis=x.axis, engine=x.engine, dtype=self.dtype)
         y[:] = 0
         ghosted_x = x.add_ghost_cells(cells_back=2)
         y_centered = ghosted_x[1:-1]