compiler: prevent temporary for local reductions

devito/builtins/utils.py

@@ -34,7 +34,7 @@ def __init__(self, *functions, op=dv.mpi.MPI.SUM, dtype=None):
         self.op = op
 
     def __enter__(self):
-        i = dv.Dimension(name='i',)
+        i = dv.Dimension(name='mri',)
         self.n = dv.Function(name='n', shape=(1,), dimensions=(i,),
                              grid=self.grid, dtype=self.dtype)
         self.n.data[0] = 0

devito/ir/clusters/algorithms.py

@@ -458,7 +458,7 @@ def normalize_reductions(cluster, sregistry, options):
 
     processed = []
     for e in cluster.exprs:
-        if e.is_Reduction and (e.lhs.is_Indexed or cluster.is_sparse):
+        if e.is_Reduction and e.lhs.is_Indexed and cluster.is_sparse:
             # Transform `e` such that we reduce into a scalar (ultimately via
             # atomic ops, though this part is carried out by a much later pass)
             # For example, given `i = m[p_src]` (i.e., indirection array), turn:
@@ -471,7 +471,8 @@ def normalize_reductions(cluster, sregistry, options):
             processed.extend([e.func(v, e.rhs, operation=None),
                               e.func(e.lhs, v)])
 
-        elif e.is_Reduction and e.lhs.is_Symbol and opt_mapify_reduce:
+        elif e.is_Reduction and e.lhs.is_Symbol and opt_mapify_reduce \
+                and not cluster.is_sparse:
             # Transform `e` into what is in essence an explicit map-reduce
             # For example, turn:
             # `s += f(u[x], v[x], ...)`

devito/ir/clusters/cluster.py

@@ -230,7 +230,18 @@ def is_dense(self):
 
     @property
     def is_sparse(self):
-        return not self.is_dense
+        """
+        A cluster is sparse if it represent a sparse operation i.e if both
+
+            * The cluster contains sparse functions
+            * The cluster uses dense functions
+
+        If only the first case is true, the cluster only contains operation on the sparse
+        function itself without indirection and therefore only contains dense operations.
+        """
+        return (any(f.is_SparseFunction for f in self.functions) and
+                len([f for f in self.functions
+                     if (f.is_Function and not f.is_SparseFunction)]) > 0)
 
     @property
     def is_halo_touch(self):

devito/ir/equations/equation.py

@@ -73,6 +73,11 @@ def apply(self, func):
         kwargs['conditionals'] = {k: func(v) for k, v in self.conditionals.items()}
         return self.func(*args, **kwargs)
 
+    def __repr__(self):
+        if not self.is_Reduction:
+            return super().__repr__()
+        else:
+            return '%s = %s(%s, %s)' % (self.lhs, self.operation, self.lhs, self.rhs)
     # Pickling support
     __reduce_ex__ = Pickable.__reduce_ex__
 

devito/types/sparse.py

@@ -338,11 +338,11 @@ def guard(self, expr=None):
                                   condition=condition, indirect=True)
 
         if expr is None:
-            out = self.indexify().xreplace({self._sparse_dim: cd})
+            out = self.indexify()._subs(self._sparse_dim, cd)
         else:
             functions = {f for f in retrieve_function_carriers(expr)
                          if f.is_SparseFunction}
-            out = indexify(expr).xreplace({f._sparse_dim: cd for f in functions})
+            out = indexify(expr).subs({f._sparse_dim: cd for f in functions})
 
         return out, temps
 

tests/test_dle.py

@@ -11,7 +11,7 @@
                     configuration, dimensions, info, cos)
 from devito.exceptions import InvalidArgument
 from devito.ir.iet import (Iteration, FindNodes, IsPerfectIteration,
-                           retrieve_iteration_tree)
+                           retrieve_iteration_tree, Expression)
 from devito.passes.iet.languages.openmp import Ompizer, OmpRegion
 from devito.tools import as_tuple
 from devito.types import Scalar
@@ -765,6 +765,29 @@ def test_array_sum_reduction(self, so, dim):
 
         assert np.allclose(f.data, 18)
 
+    def test_reduction_local(self):
+        grid = Grid((11, 11))
+        d = Dimension("i")
+        n = Function(name="n", dimensions=(d,), shape=(1,))
+        u = Function(name="u", grid=grid)
+        u.data.fill(1.)
+
+        op = Operator(Inc(n[0], u))
+        print(op)
+        op()
+
+        cond = FindNodes(Expression).visit(op)
+        iterations = FindNodes(Iteration).visit(op)
+        # Should not creat any temporary for the reduction
+        assert len(cond) == 1
+        if Ompizer._support_array_reduction(configuration['compiler']):
+            assert "reduction(+:n[0])" in iterations[0].pragmas[0].value
+        else:
+            # E.g. old GCC's
+            assert "atomic update" in str(iterations[-1])
+
+        assert n.data[0] == 11*11
+
     def test_array_max_reduction(self):
         """
         Test generation of OpenMP sum-reduction clauses involving Function's.

tests/test_dse.py

@@ -2669,13 +2669,20 @@ def test_sparse_const(self):
 
         u = TimeFunction(name="u", grid=grid)
         src = PrecomputedSparseTimeFunction(name="src", grid=grid, npoint=1, nt=11,
-                                            r=2, interpolation_coeffs=np.ones((1, 3, 2)))
+                                            r=2, interpolation_coeffs=np.ones((1, 3, 2)),
+                                            gridpoints=[[5, 5, 5]])
+        u.data.fill(1.)
+
         op = Operator(src.interpolate(u))
 
         cond = FindNodes(Conditional).visit(op)
         assert len(cond) == 1
+        assert len(cond[0].args['then_body'][0].exprs) == 1
         assert all(e.is_scalar for e in cond[0].args['then_body'][0].exprs)
 
+        op()
+        assert np.all(src.data == 8)
+
 
 class TestIsoAcoustic(object):