api: prevent inconsistency in interp radius dim by tighing it to the sparse dim

Author: mloubout

devito/ir/clusters/algorithms.py

@@ -484,8 +484,7 @@ def callback(self, clusters, prefix, seen=None):
             # `c` is scheduled
             index = 0
             for i in reversed(range(n)):
-                if not processed[i].ispace.is_subset(c.ispace) and \
-                   not processed[i].is_sparse:
+                if not processed[i].ispace.is_subset(c.ispace):
                     index = i + 1
                     break
             processed.insert(index, halo_touch)

devito/operations/interpolators.py

@@ -14,9 +14,7 @@
 from devito.logger import warning
 from devito.symbolics import INT, retrieve_function_carriers, retrieve_functions
 from devito.tools import Pickable, as_tuple, filter_ordered, flatten, memoized_meth
-from devito.types import (
-    ConditionalDimension, CustomDimension, Eq, Evaluable, Inc, SubFunction, Symbol
-)
+from devito.types import Eq, Evaluable, Inc, SubFunction, Symbol
 from devito.types.utils import DimensionTuple
 
 __all__ = ['LinearInterpolator', 'PrecomputedInterpolator', 'SincInterpolator']
@@ -239,13 +237,10 @@ def _weights(self, subdomain=None):
     def _gdims(self):
         return self.grid.dimensions
 
-    @cached_property
+    @property
     def _cdim(self):
         """Base CustomDimensions used to construct _rdim"""
-        parent = self.sfunction._sparse_dim
-        dims = [CustomDimension(f"r{self.sfunction.name}{d.name}",
-                                -self.r+1, self.r, 2*self.r, parent)
-                for d in self._gdims]
+        dims = [self.sfunction._crdim(d) for d in self._gdims]
         return dims
 
     @memoized_meth
@@ -274,8 +269,7 @@ def _rdim(self, subdomain=None):
                 rank_populated = subdomain.distributor.rank_populated
                 cond = sympy.And(rank_populated, cond)
 
-            rdims.append(ConditionalDimension(rd.name, rd, condition=cond,
-                                              indirect=True))
+            rdims.append(self.sfunction._cond_rdim(d.root, cond))
 
         return DimensionTuple(*rdims, getters=gdims)
 

devito/types/sparse.py

@@ -13,12 +13,13 @@
 )
 from devito.symbolics import indexify, retrieve_function_carriers
 from devito.tools import (
-    ReducerMap, as_tuple, dtype_to_mpidtype, filter_ordered, flatten, is_integer, prod
+    ReducerMap, as_tuple, dtype_to_mpidtype, filter_ordered, flatten, is_integer,
+    memoized_meth, prod
 )
 from devito.types.basic import Symbol
 from devito.types.dense import DiscreteFunction, SubFunction
 from devito.types.dimension import (
-    ConditionalDimension, DefaultDimension, Dimension, DynamicDimension
+    ConditionalDimension, CustomDimension, DefaultDimension, Dimension, DynamicDimension
 )
 from devito.types.dimension import dimensions as mkdims
 from devito.types.equation import Eq, Inc
@@ -386,6 +387,24 @@ def _position_map(self):
     def dist_origin(self):
         return self._dist_origin
 
+    @memoized_meth
+    def _crdim(self, dim):
+        """
+        The CustomDimension associated with the Dimension `dim` for
+        the radius of the interpolation/injection stencil
+        """
+        sname = self._sparse_dim.name
+        return CustomDimension(f"r{sname}{dim.name}", -self.r+1,
+                               self.r, 2*self.r, self._sparse_dim)
+
+    @memoized_meth
+    def _cond_rdim(self, dim, cond):
+        """
+        The interpolation/injection radius dimension with guard bounds
+        """
+        parent = self._crdim(dim)
+        return ConditionalDimension(parent.name, parent, condition=cond, indirect=True)
+
     def interpolate(self, *args, **kwargs):
         """
         Implement an interpolation operation from the grid onto the given sparse points

tests/test_dle.py

@@ -314,8 +314,8 @@ def test_prec_inject(self):
                                               'openmp': True,
                                               'par-collapse-ncores': 1}))
 
-        assert_structure(op, ['t', 't,p_s0_blk0,p_s,rsx,rsy'],
-                         't,p_s0_blk0,p_s,rsx,rsy')
+        assert_structure(op, ['t', 't,p_s0_blk0,p_s,rp_sx,rp_sy'],
+                         't,p_s0_blk0,p_s,rp_sx,rp_sy')
 
 
 class TestBlockingParTile:

tests/test_dtypes.py

@@ -311,9 +311,9 @@ def test_complex_reduction(dtypeu: np.dtype[np.complexfloating]) -> None:
         op()
 
         if op._options['linearize']:
-            ustr = 'uL0(t1, rsx + posx + 2, rsy + posy + 2)'
+            ustr = 'uL0(t1, rp_sx + posx + 2, rp_sy + posy + 2)'
         else:
-            ustr = 'u[t1][rsx + posx + 2][rsy + posy + 2]'
+            ustr = 'u[t1][rp_sx + posx + 2][rp_sy + posy + 2]'
 
         compiler = configuration['compiler']
         gnu = isinstance(compiler, GNUCompiler) or \

tests/test_interpolation.py

@@ -887,8 +887,9 @@ def test_interp_complex_and_real(dtype):
     assert np.isclose(sc.data[0], fc.data[5, 5, 5])
     assert np.isclose(scre.data[0], fc.data[5, 5, 5].real)
 
-    assert_structure(opC, ['p_sc', 'p_sc,rscx,rscy,rscz', 'p_sc,rscex,rscey,rscez'],
-                     'p_sc,rscx,rscy,rscz,rscex,rscey,rscez')
+    assert_structure(opC, ['p_sc', 'p_sc,rp_scx,rp_scy,rp_scz',
+                           'p_sc,rp_scx,rp_scy,rp_scz'],
+                     'p_sc,rp_scx,rp_scy,rp_scz,rp_scx,rp_scy,rp_scz')
 
 
 class SD0(SubDomain):
@@ -1010,9 +1011,9 @@ def test_interpolate_subdomain(self):
         assert np.all(np.isclose(sr1.data, check1))
         assert np.all(np.isclose(sr2.data, check2))
         assert_structure(op,
-                         ['p_sr0', 'p_sr0rsr0xrsr0y', 'p_sr1',
-                          'p_sr1rsr1xrsr1y', 'p_sr2', 'p_sr2rsr2xrsr2y'],
-                         'p_sr0rsr0xrsr0yp_sr1rsr1xrsr1yp_sr2rsr2xrsr2y')
+                         ['p_sr0', 'p_sr0rp_sr0xrp_sr0y', 'p_sr1',
+                          'p_sr1rp_sr1xrp_sr1y', 'p_sr2', 'p_sr2rp_sr2xrp_sr2y'],
+                         'p_sr0rp_sr0xrp_sr0yp_sr1rp_sr1xrp_sr1yp_sr2rp_sr2xrp_sr2y')
 
     def test_interpolate_subdomain_sinc(self):
         """
@@ -1053,9 +1054,9 @@ def test_interpolate_subdomain_sinc(self):
         assert np.all(np.isclose(sr0.data, sr2.data))
         assert np.all(np.isclose(sr1.data, sr2.data))
         assert_structure(op,
-                         ['p_sr0', 'p_sr0rsr0xrsr0y', 'p_sr1',
-                          'p_sr1rsr1xrsr1y', 'p_sr2', 'p_sr2rsr2xrsr2y'],
-                         'p_sr0rsr0xrsr0yp_sr1rsr1xrsr1yp_sr2rsr2xrsr2y')
+                         ['p_sr0', 'p_sr0rp_sr0xrp_sr0y', 'p_sr1',
+                          'p_sr1rp_sr1xrp_sr1y', 'p_sr2', 'p_sr2rp_sr2xrp_sr2y'],
+                         'p_sr0rp_sr0xrp_sr0yp_sr1rp_sr1xrp_sr1yp_sr2rp_sr2xrp_sr2y')
 
     def test_inject_subdomain(self):
         """
@@ -1101,8 +1102,8 @@ def test_inject_subdomain(self):
         assert np.all(np.isclose(f0.data, check0))
         assert np.all(np.isclose(f1.data, check1))
         assert_structure(op,
-                         ['p_sr0rsr0xrsr0y'],
-                         'p_sr0rsr0xrsr0y')
+                         ['p_sr0rp_sr0xrp_sr0y'],
+                         'p_sr0rp_sr0xrp_sr0y')
 
     def test_inject_subdomain_sinc(self):
         """
@@ -1133,8 +1134,8 @@ def test_inject_subdomain_sinc(self):
         assert np.all(np.isclose(f0.data, f2.data[:9, -9:]))
         assert np.all(np.isclose(f1.data, f2.data[1:-1, 1:-1]))
         assert_structure(op,
-                         ['p_sr0rsr0xrsr0y'],
-                         'p_sr0rsr0xrsr0y')
+                         ['p_sr0rp_sr0xrp_sr0y'],
+                         'p_sr0rp_sr0xrp_sr0y')
 
     @pytest.mark.xfail(reason="OOB issue")
     @pytest.mark.parallel(mode=4)