mpi: Polish _mpi_advanced_... utilities

Author: FabioLuporini

devito/data/data.py

@@ -261,13 +261,16 @@ def T(self):
 
     def _mpi_advanced_1d_target(self, glb_idx, axis):
         """
-        Return the raw local ndarray addressed by ``glb_idx`` without
-        advanced indexing along ``axis``.
+        Return a raw local view for MPI advanced-indexing communication.
 
-        The MPI advanced-indexing helper code in ``devito.data.utils`` owns
-        the communication pattern; this hook is kept on ``Data`` because only
-        the subclass can bypass its own ``__getitem__`` and obtain a plain
-        ndarray view.
+        The returned view is indexed by ``glb_idx`` on all dimensions except
+        ``axis``, which is replaced by ``slice(None)`` so the helper code in
+        ``devito.data.utils`` can pack or unpack the requested global integer
+        entries itself. ``target_axis`` is the position of ``axis`` in the
+        returned view after scalar-indexed dimensions have been dropped.
+
+        This hook is kept on ``Data`` because only the subclass can bypass its
+        own ``__getitem__`` and obtain a plain ndarray view.
         """
         target_idx = list(glb_idx)
         target_idx[axis] = slice(None)

devito/data/utils.py

@@ -218,10 +218,15 @@ def mpi_advanced_1d_set(data, glb_idx, val, axis, indices, decomposition,
 
 
 def _mpi_advanced_1d_error(data, error):
-    """Raise the first error reported by any rank, on every rank."""
+    """Raise rank-local advanced-indexing errors on every rank."""
     if data._distributor.nprocs > 1:
         errors = data._distributor.comm.allgather(error)
-        error = next((i for i in errors if i is not None), None)
+        errors = [
+            f"rank {rank}: {msg}"
+            for rank, msg in enumerate(errors)
+            if msg is not None
+        ]
+        error = "; ".join(errors) if errors else None
 
     if error is not None:
         raise ValueError(error)

tests/test_data.py

@@ -1499,16 +1499,16 @@ def test_advanced_indexing_errors(self, mode):
             np.empty(0, dtype=np.int64)
         duplicate_data = np.ones(duplicate_index.size, dtype=f.dtype)
 
-        with pytest.raises(ValueError, match="Duplicate global indices"):
+        with pytest.raises(ValueError, match="rank 0:.*Duplicate global indices"):
             f.data[duplicate_index] = duplicate_data
 
         oob_index = np.array([8]) if rank == 0 else np.empty(0, dtype=np.int64)
         oob_data = np.ones(oob_index.size, dtype=f.dtype)
 
-        with pytest.raises(ValueError, match="out-of-bounds"):
+        with pytest.raises(ValueError, match="rank 0:.*out-of-bounds"):
             f.data[oob_index] = oob_data
 
-        with pytest.raises(ValueError, match="out-of-bounds"):
+        with pytest.raises(ValueError, match="rank 0:.*out-of-bounds"):
             f.data[oob_index]
 
     @pytest.mark.parallel(mode=4)