Change to ncp.cuda.Device().synchronize()

Author: rohanbabbar04

pylops_mpi/utils/_mpi.py

@@ -47,13 +47,13 @@ def mpi_allgather(base_comm: MPI.Comm,
         send_shapes = base_comm.allgather(send_buf.shape)
         (padded_send, padded_recv) = _prepare_allgather_inputs(send_buf, send_shapes, engine=engine)
         recv_buffer_to_use = recv_buf if recv_buf else padded_recv
-        _mpi_calls(base_comm.Allgather, padded_send, recv_buffer_to_use, engine=engine)
+        _mpi_calls(base_comm, "Allgather", padded_send, recv_buffer_to_use, engine=engine)
         return _unroll_allgather_recv(recv_buffer_to_use, padded_send.shape, send_shapes)
     else:
         # CuPy with non-CUDA-aware MPI
         if recv_buf is None:
-            return _mpi_calls(base_comm.allgather, send_buf)
-        _mpi_calls(base_comm.Allgather, send_buf, recv_buf)
+            return _mpi_calls(base_comm, "allgather", send_buf)
+        _mpi_calls(base_comm, "Allgather", send_buf, recv_buf)
         return recv_buf
 
 
@@ -92,14 +92,14 @@ def mpi_allreduce(base_comm: MPI.Comm,
     if deps.cuda_aware_mpi_enabled or engine == "numpy":
         ncp = get_module(engine)
         recv_buf = ncp.zeros(send_buf.size, dtype=send_buf.dtype)
-        _mpi_calls(base_comm.Allreduce, send_buf, recv_buf, op, engine=engine)
+        _mpi_calls(base_comm, "Allreduce", send_buf, recv_buf, op, engine=engine)
         return recv_buf
     else:
         # CuPy with non-CUDA-aware MPI
         if recv_buf is None:
-            return _mpi_calls(base_comm.allreduce, send_buf, op, engine=engine)
+            return _mpi_calls(base_comm, "allreduce", send_buf, op, engine=engine)
         # For MIN and MAX which require recv_buf
-        _mpi_calls(base_comm.Allreduce, send_buf, recv_buf, op, engine=engine)
+        _mpi_calls(base_comm, "Allreduce", send_buf, recv_buf, op, engine=engine)
         return recv_buf
 
 
@@ -131,11 +131,11 @@ def mpi_bcast(base_comm: MPI.Comm,
 
     """
     if deps.cuda_aware_mpi_enabled or engine == "numpy":
-        _mpi_calls(base_comm.Bcast, send_buf, engine=engine, root=root)
+        _mpi_calls(base_comm, "Bcast", send_buf, engine=engine, root=root)
         return send_buf
     # CuPy with non-CUDA-aware MPI: use object broadcast
     value = send_buf if base_comm.Get_rank() == root else None
-    return _mpi_calls(base_comm.bcast, value, engine=engine, root=root)
+    return _mpi_calls(base_comm, "bcast", value, engine=engine, root=root)
 
 
 def mpi_send(base_comm: MPI.Comm,
@@ -171,10 +171,10 @@ def mpi_send(base_comm: MPI.Comm,
         mpi_type = MPI._typedict[send_buf.dtype.char]
         if count is None:
             count = send_buf.size
-        _mpi_calls(base_comm.Send, [send_buf, count, mpi_type], engine=engine, dest=dest, tag=tag)
+        _mpi_calls(base_comm, "Send", [send_buf, count, mpi_type], engine=engine, dest=dest, tag=tag)
     else:
         # Uses CuPy without CUDA-aware MPI
-        _mpi_calls(base_comm.send, send_buf, dest, tag, engine=engine)
+        _mpi_calls(base_comm, "send", send_buf, dest, tag, engine=engine)
 
 
 def mpi_recv(base_comm: MPI.Comm,
@@ -219,23 +219,25 @@ def mpi_recv(base_comm: MPI.Comm,
             # dimension or shape-related integers are send/recv
             recv_buf = ncp.zeros(count, dtype=ncp.int32)
         mpi_type = MPI._typedict[recv_buf.dtype.char]
-        _mpi_calls(base_comm.Recv, [recv_buf, recv_buf.size, mpi_type], engine=engine, source=source, tag=tag)
+        _mpi_calls(base_comm, "Recv", [recv_buf, recv_buf.size, mpi_type], engine=engine, source=source, tag=tag)
     else:
         # Uses CuPy without CUDA-aware MPI
-        recv_buf = _mpi_calls(base_comm.recv, engine=engine, source=source, tag=tag)
+        recv_buf = _mpi_calls(base_comm, "recv", engine=engine, source=source, tag=tag)
     return recv_buf
 
 
-def _mpi_calls(call, *args, engine: Optional[str] = "numpy", **kwargs):
+def _mpi_calls(comm: MPI.Comm, func: str, *args, engine: Optional[str] = "numpy", **kwargs):
     """MPI Calls
     Wrapper around MPI comm calls with optional GPU synchronization for CuPy arrays.
 
     Parameters
     ----------
-    call: :obj:`MPI.Comm`
-        MPI Communicator function
+    comm: :obj:`MPI.Comm`
+        MPI Communicator
+    func
+        MPI Function to call.
     args
-        Arguments passed to the MPI call.
+        Arguments to pass to the function.
     engine: :obj:`str`, optional
         Engine used to store array (``numpy`` or ``cupy``)
     kwargs
@@ -246,6 +248,7 @@ def _mpi_calls(call, *args, engine: Optional[str] = "numpy", **kwargs):
         Result of the MPI call
     """
     if engine == "cupy":
-        cp = get_module(engine)
-        cp.cuda.runtime.deviceSynchronize()
-    return call(*args, **kwargs)
+        ncp = get_module(engine)
+        ncp.cuda.Device().synchronize()
+    mpi_func = getattr(comm, func)
+    return mpi_func(*args, **kwargs)