Merge pull request #2912 from devitocodes/misc-sparse-fix

compiler: misc sparse bug fixes

Author: mloubout

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)
 
@@ -425,7 +419,7 @@ def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None
                                            subdomain=subdomain)
 
         # Accumulate point-wise contributions into a temporary
-        rhs = Symbol(name='sum', dtype=self.sfunction.dtype)
+        rhs = Symbol(name=f'sum{self.sfunction.name}', dtype=self.sfunction.dtype)
         summands = [Eq(rhs, 0., implicit_dims=implicit_dims)]
         # Substitute coordinate base symbols into the interpolation coefficients
         weights = self._weights(subdomain=subdomain)

devito/passes/iet/languages/CXX.py

@@ -19,6 +19,9 @@ def std_arith(prefix=''):
         prefix = prefix if prefix.endswith(' ') else f'{prefix} '
     return f"""
 #include <complex>
+#include <type_traits>
+
+// ---- scalar <op> complex<T>  (scalar promoted to T) --------------------
 
 template<typename _Tp, typename _Ti>
 {prefix}std::complex<_Tp> operator * (const _Ti & a, const std::complex<_Tp> & b){{
@@ -32,7 +35,7 @@ def std_arith(prefix=''):
 
 template<typename _Tp, typename _Ti>
 {prefix}std::complex<_Tp> operator / (const _Ti & a, const std::complex<_Tp> & b){{
-  _Tp denom = b.real() * b.real () + b.imag() * b.imag();
+  _Tp denom = b.real() * b.real() + b.imag() * b.imag();
   return std::complex<_Tp>(b.real() * a / denom, - b.imag() * a / denom);
 }}
 
@@ -53,14 +56,37 @@ def std_arith(prefix=''):
 
 template<typename _Tp, typename _Ti>
 {prefix}std::complex<_Tp> operator - (const _Ti & a, const std::complex<_Tp> & b){{
-  return std::complex<_Tp>(a = b.real(), b.imag());
+  return std::complex<_Tp>(a - b.real(), -b.imag());
 }}
 
 template<typename _Tp, typename _Ti>
 {prefix}std::complex<_Tp> operator - (const std::complex<_Tp> & b, const _Ti & a){{
   return std::complex<_Tp>(b.real() - a, b.imag());
 }}
 
+// ---- mixed-precision complex<T1> <op> complex<T2> ----------------------
+// Promote both sides to std::complex<common_type_t<T1,T2>> and delegate to
+// the standard library's same-type operator. The enable_if disables the
+// overload when T1 == T2 so we don't collide with std::complex's own ops.
+
+#define _MIXED_COMPLEX_OP(OP) \\
+template<typename _Tp1, typename _Tp2, \\
+         typename _Tr = std::common_type_t<_Tp1, _Tp2>, \\
+         typename = std::enable_if_t<!std::is_same<_Tp1, _Tp2>::value>> \\
+{prefix}std::complex<_Tr> \\
+operator OP (const std::complex<_Tp1> & a, \\
+             const std::complex<_Tp2> & b) {{ \\
+  return std::complex<_Tr>(a.real(), a.imag()) \\
+    OP std::complex<_Tr>(b.real(), b.imag()); \\
+}}
+
+_MIXED_COMPLEX_OP(*)
+_MIXED_COMPLEX_OP(/)
+_MIXED_COMPLEX_OP(+)
+_MIXED_COMPLEX_OP(-)
+
+#undef _MIXED_COMPLEX_OP
+
 """
 
 

devito/symbolics/inspection.py

@@ -316,11 +316,14 @@ def sympy_dtype(expr, base=None, default=None, smin=None):
     if expr is None:
         return default
 
-    dtypes = {base} - {None}
+    dtypes = set()
     for i in expr.free_symbols:
         with suppress(AttributeError):
             dtypes.add(i.dtype)
 
+    if not dtypes or not np.issubdtype(base, np.complexfloating):
+        dtypes.update({base} - {None})
+
     dtype = infer_dtype(dtypes)
 
     # Promote if we missed complex number, i.e f + I

devito/tools/dtypes_lowering.py

@@ -371,4 +371,5 @@ def extract_dtype(expr):
     """Extract the "winning" dtype from an expression"""
     dtypes = {getattr(e, 'dtype', None)
               for e in expr.free_symbols}
+
     return infer_dtype(dtypes - {None})

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

examples/userapi/06_sparse_operations.ipynb

@@ -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_dle.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_dtypes.py

@@ -7,7 +7,7 @@
 from conftest import assert_structure
 from devito import (
     DefaultDimension, Dimension, Eq, Function, Grid, MatrixSparseTimeFunction, Operator,
-    PrecomputedSparseFunction, PrecomputedSparseTimeFunction, SparseFunction,
+    PrecomputedSparseFunction, PrecomputedSparseTimeFunction, Real, SparseFunction,
     SparseTimeFunction, SubDomain, TimeFunction, switchconfig
 )
 from devito.operations.interpolators import LinearInterpolator, SincInterpolator
@@ -17,6 +17,11 @@
 from examples.seismic.acoustic import AcousticWaveSolver, acoustic_setup
 
 
+class SparseFirst(SparseFunction):
+    """ Custom sparse class with the sparse dimension as the first one"""
+    _sparse_position = 0
+
+
 def unit_box(name='a', shape=(11, 11), grid=None, space_order=1):
     """Create a field with value 0. to 1. in each dimension"""
     grid = grid or Grid(shape=shape)
@@ -698,11 +703,6 @@ def test_sparse_first():
     """
     Tests custom sprase function with sparse dimension as first index.
     """
-
-    class SparseFirst(SparseFunction):
-        """ Custom sparse class with the sparse dimension as the first one"""
-        _sparse_position = 0
-
     dr = Dimension("cd")
     ds = DefaultDimension("ps", default_value=3)
     grid = Grid((11, 11))
@@ -870,6 +870,28 @@ def test_interp_complex(dtype):
     assert np.isclose(sc.data[0], fc.data[5, 5, 5])
 
 
+@pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
+def test_interp_complex_and_real(dtype):
+    grid = Grid((11, 11, 11))
+
+    sc = SparseFunction(name="sc", grid=grid, npoint=1, dtype=dtype)
+    sc.coordinates.data[:] = [.5, .5, .5]
+    scre = SparseFunction(name="sce", grid=grid, npoint=1, coordinates=sc.coordinates)
+
+    fc = Function(name="fc", grid=grid, npoint=2, dtype=dtype)
+    fc.data[:] = np.random.randn(*grid.shape) + 1j * np.random.randn(*grid.shape)
+    exprs = sc.interpolate(expr=fc) + scre.interpolate(expr=Real(fc))
+    opC = Operator(exprs, name="OpC")
+    opC()
+
+    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,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):
     name = 'sd0'
 
@@ -989,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):
         """
@@ -1032,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):
         """
@@ -1080,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):
         """
@@ -1112,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)