feat: added _bcast to DistributedMixIn and added comms as input for all methods

Author: mrava87SW

pylops_mpi/Distributed.py

@@ -1,14 +1,14 @@
 from mpi4py import MPI
 from pylops.utils import deps as pylops_deps  # avoid namespace crashes with pylops_mpi.utils
-from pylops_mpi.utils._mpi import mpi_allreduce, mpi_allgather, mpi_send, mpi_recv, _prepare_allgather_inputs, _unroll_allgather_recv
+from pylops_mpi.utils._mpi import mpi_allreduce, mpi_allgather, mpi_bcast, mpi_send, mpi_recv, _prepare_allgather_inputs, _unroll_allgather_recv
 from pylops_mpi.utils import deps
 
 cupy_message = pylops_deps.cupy_import("the DistributedArray module")
 nccl_message = deps.nccl_import("the DistributedArray module")
 
 if nccl_message is None and cupy_message is None:
     from pylops_mpi.utils._nccl import (
-        nccl_allgather, nccl_allreduce, nccl_send, nccl_recv
+        nccl_allgather, nccl_allreduce, nccl_bcast, nccl_send, nccl_recv
     )
 
 
@@ -22,39 +22,45 @@ class DistributedMixIn:
     MPI installation is available, the latter with CuPy arrays when a CUDA-Aware
     MPI installation is not available).
     """
-    def _allreduce(self, send_buf, recv_buf=None, op: MPI.Op = MPI.SUM):
+    def _allreduce(self, base_comm, base_comm_nccl, 
+                   send_buf, recv_buf=None, op: MPI.Op = MPI.SUM, 
+                   engine="numpy"):
         """Allreduce operation
         """
-        if deps.nccl_enabled and getattr(self, "base_comm_nccl"):
-            return nccl_allreduce(self.base_comm_nccl, send_buf, recv_buf, op)
+        if deps.nccl_enabled and base_comm_nccl is not None:
+            return nccl_allreduce(base_comm_nccl, send_buf, recv_buf, op)
         else:
-            return mpi_allreduce(self.base_comm, send_buf,
-                                 recv_buf, self.engine, op)
+            return mpi_allreduce(base_comm, send_buf,
+                                 recv_buf, engine, op)
 
-    def _allreduce_subcomm(self, send_buf, recv_buf=None, op: MPI.Op = MPI.SUM):
+    def _allreduce_subcomm(self, sub_comm, base_comm_nccl, 
+                           send_buf, recv_buf=None, op: MPI.Op = MPI.SUM,
+                           engine="numpy"):
         """Allreduce operation with subcommunicator
         """
-        if deps.nccl_enabled and getattr(self, "base_comm_nccl"):
-            return nccl_allreduce(self.sub_comm, send_buf, recv_buf, op)
+        if deps.nccl_enabled and base_comm_nccl is not None:
+            return nccl_allreduce(sub_comm, send_buf, recv_buf, op)
         else:
-            return mpi_allreduce(self.sub_comm, send_buf,
-                                 recv_buf, self.engine, op)
+            return mpi_allreduce(sub_comm, send_buf,
+                                 recv_buf, engine, op)
 
-    def _allgather(self, send_buf, recv_buf=None):
+    def _allgather(self, base_comm, base_comm_nccl, 
+                   send_buf, recv_buf=None,
+                   engine="numpy"):
         """Allgather operation
         """
-        if deps.nccl_enabled and self.base_comm_nccl:
+        if deps.nccl_enabled and base_comm_nccl is not None:
             if isinstance(send_buf, (tuple, list, int)):
-                return nccl_allgather(self.base_comm_nccl, send_buf, recv_buf)
+                return nccl_allgather(base_comm_nccl, send_buf, recv_buf)
             else:
-                send_shapes = self.base_comm.allgather(send_buf.shape)
+                send_shapes = base_comm.allgather(send_buf.shape)
                 (padded_send, padded_recv) = _prepare_allgather_inputs(send_buf, send_shapes, engine="cupy")
-                raw_recv = nccl_allgather(self.base_comm_nccl, padded_send, recv_buf if recv_buf else padded_recv)
+                raw_recv = nccl_allgather(base_comm_nccl, padded_send, recv_buf if recv_buf else padded_recv)
                 return _unroll_allgather_recv(raw_recv, padded_send.shape, send_shapes)
         else:
             if isinstance(send_buf, (tuple, list, int)):
-                return self.base_comm.allgather(send_buf)
-            return mpi_allgather(self.base_comm, send_buf, recv_buf, self.engine)
+                return base_comm.allgather(send_buf)
+            return mpi_allgather(base_comm, send_buf, recv_buf, engine)
 
     def _allgather_subcomm(self, send_buf, recv_buf=None):
         """Allgather operation with subcommunicator
@@ -70,6 +76,16 @@ def _allgather_subcomm(self, send_buf, recv_buf=None):
         else:
             return mpi_allgather(self.sub_comm, send_buf, recv_buf, self.engine)
 
+    def _bcast(self, local_array, index, value):
+        """BCast operation
+        """
+        if deps.nccl_enabled and getattr(self, "base_comm_nccl"):
+            nccl_bcast(self.base_comm_nccl, local_array, index, value)
+        else:
+            # self.local_array[index] = self.base_comm.bcast(value)
+            mpi_bcast(self.base_comm, self.rank, self.local_array, index, value,
+                      engine=self.engine)
+        
     def _send(self, send_buf, dest, count=None, tag=0):
         """Send operation
         """