Merge pull request #1824 from devitocodes/sparse-coord-order

mpi: Fix mask ordering for sparse gather

.github/workflows/pytest-core-mpi.yml

@@ -41,6 +41,7 @@ jobs:
 
     - name: Test with pytest
       run: |
+        python3 scripts/clear_devito_cache.py
         python3 -m pytest --cov --cov-config=.coveragerc --cov-report=xml -m parallel tests/
 
     - name: Upload coverage to Codecov

devito/types/sparse.py

@@ -151,37 +151,6 @@ def _dist_scatter_mask(self, dmap=None):
         ret[self._sparse_position] = mask
         return tuple(ret)
 
-    def _dist_subfunc_scatter_mask(self, dmap=None):
-        """
-        This method is analogous to :meth:`_dist_scatter_mask`, although
-        the mask is now suitable to index into self's SubFunctions, rather
-        than into ``self.data``.
-        """
-        return self._dist_scatter_mask(dmap=dmap)[self._sparse_position]
-
-    def _dist_gather_mask(self, dmap=None):
-        """
-        A mask to index into the ``data`` received upon returning from
-        ``self._dist_alltoall``. This mask creates a new data array in which
-        duplicate sparse data values have been discarded. The resulting data
-        array can thus be used to populate ``self.data``.
-        """
-        ret = list(self._dist_scatter_mask(dmap=dmap))
-        mask = ret[self._sparse_position]
-        inds = np.unique(mask, return_index=True)[1]
-        inds.sort()
-        ret[self._sparse_position] = inds.tolist()
-
-        return tuple(ret)
-
-    def _dist_subfunc_gather_mask(self, dmap=None):
-        """
-        This method is analogous to :meth:`_dist_subfunc_scatter_mask`, although
-        the mask is now suitable to index into self's SubFunctions, rather
-        than into ``self.data``.
-        """
-        return self._dist_gather_mask(dmap=dmap)[self._sparse_position]
-
     def _dist_count(self, dmap=None):
         """
         A 2-tuple of comm-sized iterables, which tells how many sparse points
@@ -261,13 +230,6 @@ def _dist_scatter(self):
         """
         raise NotImplementedError
 
-    def _dist_gather(self, data):
-        """
-        A ``numpy.ndarray`` containing up-to-date data and coordinate values
-        suitable for insertion into ``self.data``.
-        """
-        raise NotImplementedError
-
     def _arg_defaults(self, alias=None):
         key = alias or self
         mapper = {self: key}
@@ -698,9 +660,10 @@ def _dist_scatter(self, data=None):
 
         # Compute dist map only once
         dmap = self._dist_datamap
+        mask = self._dist_scatter_mask(dmap=dmap)
 
         # Pack sparse data values so that they can be sent out via an Alltoallv
-        data = data[self._dist_scatter_mask(dmap=dmap)]
+        data = data[mask]
         data = np.ascontiguousarray(np.transpose(data, self._dist_reorder_mask))
 
         # Send out the sparse point values
@@ -714,7 +677,7 @@ def _dist_scatter(self, data=None):
         data = np.ascontiguousarray(np.transpose(data, self._dist_reorder_mask))
 
         # Pack (reordered) coordinates so that they can be sent out via an Alltoallv
-        coords = self.coordinates.data._local[self._dist_subfunc_scatter_mask(dmap=dmap)]
+        coords = self.coordinates.data._local[mask[self._sparse_position]]
 
         # Send out the sparse point coordinates
         _, scount, sdisp, rshape, rcount, rdisp = self._dist_subfunc_alltoall(dmap=dmap)
@@ -739,6 +702,7 @@ def _dist_gather(self, data, coords):
 
         # Compute dist map only once
         dmap = self._dist_datamap
+        mask = self._dist_scatter_mask(dmap=dmap)
 
         # Pack sparse data values so that they can be sent out via an Alltoallv
         data = np.ascontiguousarray(np.transpose(data, self._dist_reorder_mask))
@@ -750,7 +714,7 @@ def _dist_gather(self, data, coords):
                        [gathered, scount, sdisp, mpitype])
         # Unpack data values so that they follow the expected storage layout
         gathered = np.ascontiguousarray(np.transpose(gathered, self._dist_reorder_mask))
-        self._data[:] = gathered[self._dist_gather_mask(dmap=dmap)]
+        self._data[mask] = gathered[:]
 
         if coords is not None:
             # Pack (reordered) coordinates so that they can be sent out via an Alltoallv
@@ -762,8 +726,7 @@ def _dist_gather(self, data, coords):
             mpitype = MPI._typedict[np.dtype(self.coordinates.dtype).char]
             comm.Alltoallv([coords, rcount, rdisp, mpitype],
                            [gathered, scount, sdisp, mpitype])
-            self._coordinates.data._local[:] = \
-                gathered[self._dist_subfunc_gather_mask(dmap=dmap)]
+            self._coordinates.data._local[mask[self._sparse_position]] = gathered[:]
 
         # Note: this method "mirrors" `_dist_scatter`: a sparse point that is sent
         # in `_dist_scatter` is here received; a sparse point that is received in

tests/test_gpu_common.py

@@ -20,8 +20,9 @@ class TestGPUInfo(object):
 
     def test_get_gpu_info(self):
         info = get_gpu_info()
+        known = ['nvidia', 'tesla', 'geforce', 'unspecified']
         try:
-            assert info['architecture'].lower() in ['tesla', 'geforce', 'unspecified']
+            assert info['architecture'].lower() in known
         except KeyError:
             # There might be than one GPUs, but for now we don't care
             # as we're not really exploiting this info yet...

tests/test_mpi.py

@@ -478,6 +478,22 @@ def test_sparse_coords(self):
                 coords_loc += sf.coordinates.data[i, 0]
             assert sf.data[i] == coords_loc
 
+    @pytest.mark.parallel(mode=4)
+    def test_sparse_coords_issue1823(self):
+        grid = Grid((101, 101, 101), extent=(1000, 1000, 1000))
+        coords = np.array([[1000., 0., 900.], [1000., 300., 700.],
+                           [1000., 500., 500.], [1000., 700., 300.],
+                           [1000., 900., 0.], [1000., 0., 850.]])
+        rec = SparseTimeFunction(name="s", grid=grid, coordinates=coords,
+                                 nt=10, npoint=6)
+        ref = SparseTimeFunction(name="s1", grid=grid, coordinates=coords,
+                                 nt=10, npoint=6)
+        u = TimeFunction(name="u", grid=grid, space_order=1)
+
+        Operator([Eq(u, u+1)]+rec.interpolate(u))()
+
+        assert np.allclose(rec.coordinates.data[:], ref.coordinates.data)
+
 
 class TestOperatorSimple(object):