compiler: Tweak LocalSum

Author: FabioLuporini

devito/finite_differences/differentiable.py

@@ -994,23 +994,30 @@ class LocalSum(IndexSum, Pickable):
     """
 
     __rargs__ = ('expr',)
-    __rkwargs__ = ('cdims',)
+    __rkwargs__ = ('cdims', 'dtype')
 
-    def __new__(cls, expr, cdims=(), **kwargs):
+    def __new__(cls, expr, cdims=(), dtype=None, **kwargs):
         obj = sympy.Expr.__new__(cls, expr)
 
         obj._expr = expr
         obj._cdims = as_tuple(cdims)
+        obj._dtype = dtype
 
         return obj
 
     def _hashable_content(self):
-        return super()._hashable_content() + (self.cdims,)
+        return super()._hashable_content() + (self.cdims, self.dtype)
 
     @property
     def cdims(self):
         return self._cdims
 
+    @cached_property
+    def dtype(self):
+        if self._dtype is None:
+            return extract_dtype(self.expr)
+        return self._dtype
+
     @cached_property
     def dimensions(self):
         return tuple(d.parent for d in self.cdims)

devito/operations/interpolators.py

@@ -460,7 +460,7 @@ def _local_accumulator(self, expr, idx_subs, subdomain=None):
         weights = self._weights(subdomain=subdomain)
         rdims = self._rdim(subdomain=subdomain)
         summand = (weights * expr).xreplace(idx_subs)
-        return LocalSum(summand, cdims=rdims)
+        return LocalSum(summand, cdims=rdims, dtype=self.sfunction.dtype)
 
     @check_radius
     @check_coords

devito/passes/clusters/localsum.py

@@ -1,5 +1,3 @@
-from functools import partial
-
 from devito.ir import ClusterizedEq, Interval, IterationSpace
 from devito.symbolics import uxreplace
 from devito.tools import timed_pass
@@ -46,7 +44,7 @@ def lower_local_sums(clusters, sregistry=None, **kwargs):
                 processed.extend([init, update])
                 subs[reduction] = value
 
-            expr = e.apply(partial(uxreplace, rule=subs))
+            expr = uxreplace(e, subs)
             processed.append(c.rebuild(exprs=[expr]))
 
     return processed
@@ -56,7 +54,7 @@ def lower_local_sum(cluster, reduction, sregistry):
     """
     Construct the private initializer and guarded accumulation for one sum.
     """
-    value = Temp(name=sregistry.make_name(prefix='sum'), dtype=cluster.dtype)
+    value = Temp(name=sregistry.make_name(prefix='sum'), dtype=reduction.dtype)
 
     dims = reduction.dimensions
     inner = IterationSpace([Interval(d) for d in dims])

tests/test_interpolation.py

@@ -12,12 +12,13 @@
     switchconfig
 )
 from devito.finite_differences import LocalSum
-from devito.ir import LoweredEq
+from devito.ir import FindSymbols, LoweredEq
 from devito.operations.interpolators import (
     LinearInterpolator, SincInterpolator, _cell_indices
 )
 from devito.symbolics import uxreplace
 from devito.tools import as_tuple
+from devito.types import Temp
 from examples.seismic import (
     AcquisitionGeometry, Receiver, RickerSource, TimeAxis, demo_model
 )
@@ -85,6 +86,20 @@ def test_zero(self):
         op.apply()
         np.testing.assert_array_equal(rcv.data, 0.)
 
+    def test_dtype(self):
+        grid = Grid(shape=(17,), dtype=np.float64)
+        f = Function(name='f', grid=grid)
+        rcv = SparseFunction(name='rcv', grid=grid, npoint=3, dtype=np.float32)
+        exprs = rcv.interpolate(f)
+        reduction = exprs.evaluate[-1].rhs
+
+        assert reduction.dtype is rcv.dtype
+
+        op = Operator(exprs, name='MixedPrecisionSparseSum', opt='noop')
+        values = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
+        assert len(values) == 1
+        assert values[0].dtype is reduction.dtype
+
 
 # ---------------------------------------------------------------------------
 # Helpers