temporary use CPU buffer for CuPy + MPI in inf and -inf norm

Author: tharittk

pylops_mpi/DistributedArray.py

@@ -694,14 +694,25 @@ def _compute_vector_norm(self, local_array: NDArray,
             recv_buf = self._allreduce_subcomm(ncp.count_nonzero(local_array, axis=axis).astype(ncp.float64))
         elif ord == ncp.inf:
             # Calculate max followed by max reduction
-            recv_buf = self._allreduce_subcomm(ncp.max(ncp.abs(local_array), axis=axis).astype(ncp.float64),
-                                               recv_buf, op=MPI.MAX)
-            recv_buf = ncp.squeeze(recv_buf, axis=axis)
+            # TODO (tharitt): currently CuPy + MPI does not work well with buffered communication, particularly
+            # with MAX, MIN operator. Here we copy the array back to CPU, transfer, and copy them back to GPUs
+            send_buf = ncp.max(ncp.abs(local_array), axis=axis).astype(ncp.float64)
+            if self.engine=="cupy" and self.base_comm_nccl is None:
+                recv_buf = self._allreduce_subcomm(send_buf.get(), recv_buf.get(), op=MPI.MAX)
+                recv_buf = ncp.asarray(ncp.squeeze(recv_buf, axis=axis))
+            else:
+                recv_buf = self._allreduce_subcomm(send_buf, recv_buf, op=MPI.MAX)
+                recv_buf = ncp.squeeze(recv_buf, axis=axis)
         elif ord == -ncp.inf:
             # Calculate min followed by min reduction
-            recv_buf = self._allreduce_subcomm(ncp.min(ncp.abs(local_array), axis=axis).astype(ncp.float64),
-                                               recv_buf, op=MPI.MIN)
-            recv_buf = ncp.squeeze(recv_buf, axis=axis)
+            # TODO (tharitt): see the comment above in infinity norm
+            send_buf = ncp.min(ncp.abs(local_array), axis=axis).astype(ncp.float64)
+            if self.engine == "cupy" and self.base_comm_nccl is None:
+                recv_buf = self._allreduce_subcomm(send_buf.get(), recv_buf.get(), op=MPI.MIN)
+                recv_buf = ncp.asarray(ncp.squeeze(recv_buf, axis=axis))
+            else:
+                recv_buf = self._allreduce_subcomm(send_buf, recv_buf, op=MPI.MIN)
+                recv_buf = ncp.asarray(ncp.squeeze(recv_buf, axis=axis))
 
         else:
             recv_buf = self._allreduce_subcomm(ncp.sum(ncp.abs(ncp.float_power(local_array, ord)), axis=axis))

tests/test_distributedarray.py

@@ -205,8 +205,8 @@ def test_distributed_norm(par):
                     np.linalg.norm(par['x'], ord=1, axis=par['axis']), rtol=1e-14)
 
     # TODO (tharitt): FAIL with CuPy + MPI for inf norm
-    # assert_allclose(arr.norm(ord=np.inf, axis=par['axis']),
-                        # np.linalg.norm(par['x'], ord=np.inf, axis=par['axis']), rtol=1e-14)
+    assert_allclose(arr.norm(ord=np.inf, axis=par['axis']),
+                        np.linalg.norm(par['x'], ord=np.inf, axis=par['axis']), rtol=1e-14)
     assert_allclose(arr.norm(), np.linalg.norm(par['x'].flatten()), rtol=1e-13)
 
 
@@ -335,7 +335,7 @@ def test_distributed_maskednorm(par):
                     np.linalg.norm(par['x'], ord=1, axis=par['axis']) / nsub, rtol=1e-14)
 
     # TODO (tharitt): Fail with CuPy + MPI
-    # assert_allclose(arr.norm(ord=np.inf, axis=par['axis']),
-    #                 np.linalg.norm(par['x'], ord=np.inf, axis=par['axis']), rtol=1e-14)
+    assert_allclose(arr.norm(ord=np.inf, axis=par['axis']),
+                    np.linalg.norm(par['x'], ord=np.inf, axis=par['axis']), rtol=1e-14)
     assert_allclose(arr.norm(ord=2, axis=par['axis']),
                     np.linalg.norm(par['x'], ord=2, axis=par['axis']) / np.sqrt(nsub), rtol=1e-13)