compiler: split normalization into sparse and dense

devito/ir/clusters/algorithms.py

@@ -7,7 +7,7 @@
 
 from devito.exceptions import InvalidOperator
 from devito.ir.support import (Any, Backward, Forward, IterationSpace,
-                               PARALLEL_IF_ATOMIC, pull_dims)
+                               pull_dims)
 from devito.ir.clusters.analysis import analyze
 from devito.ir.clusters.cluster import Cluster, ClusterGroup
 from devito.ir.clusters.visitors import Queue, QueueStateful, cluster_pass
@@ -402,7 +402,8 @@ def normalize(clusters, **kwargs):
     sregistry = kwargs['sregistry']
 
     clusters = normalize_nested_indexeds(clusters, sregistry)
-    clusters = normalize_reductions(clusters, sregistry, options)
+    clusters = normalize_reductions_dense(clusters, sregistry, options)
+    clusters = normalize_reductions_sparse(clusters, sregistry, options)
 
     return clusters
 
@@ -444,35 +445,25 @@ def pull_indexeds(expr, subs, mapper, parent=None):
     return cluster.rebuild(processed)
 
 
-@cluster_pass(mode='all')
-def normalize_reductions(cluster, sregistry, options):
+@cluster_pass(mode='dense')
+def normalize_reductions_dense(cluster, sregistry, options):
     """
     Extract the right-hand sides of reduction Eq's in to temporaries.
     """
+    if cluster.is_sparse:
+        return cluster
+
     opt_mapify_reduce = options['mapify-reduce']
 
-    dims = [d for d, v in cluster.properties.items() if PARALLEL_IF_ATOMIC in v]
+    dims = [d for d in cluster.properties.dimensions
+            if cluster.properties.is_parallel_atomic(d)]
 
     if not dims:
         return cluster
 
     processed = []
     for e in cluster.exprs:
-        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:
-            # `u[t, i] += f(u[t, i], src, ...)`
-            # into
-            # `s = f(u[t, i], src, ...)`
-            # `u[t, i] += s`
-            name = sregistry.make_name()
-            v = Symbol(name=name, dtype=e.dtype)
-            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 \
-                and not cluster.is_sparse:
+        if e.is_Reduction and e.lhs.is_Symbol and opt_mapify_reduce:
             # Transform `e` into what is in essence an explicit map-reduce
             # For example, turn:
             # `s += f(u[x], v[x], ...)`
@@ -485,7 +476,34 @@ def normalize_reductions(cluster, sregistry, options):
             a = Array(name=name, dtype=e.dtype, dimensions=dims)
             processed.extend([Eq(a.indexify(), e.rhs),
                               e.func(e.lhs, a.indexify())])
+        else:
+            processed.append(e)
+
+    return cluster.rebuild(processed)
 
+
+@cluster_pass(mode='sparse')
+def normalize_reductions_sparse(cluster, sregistry, options):
+    """
+    Extract the right-hand sides of reduction Eq's in to temporaries.
+    """
+    if not cluster.is_sparse:
+        return cluster
+
+    processed = []
+    for e in cluster.exprs:
+        if e.is_Reduction and e.lhs.is_Indexed:
+            # 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:
+            # `u[t, i] += f(u[t, i], src, ...)`
+            # into
+            # `s = f(u[t, i], src, ...)`
+            # `u[t, i] += s`
+            name = sregistry.make_name()
+            v = Symbol(name=name, dtype=e.dtype)
+            processed.extend([e.func(v, e.rhs, operation=None),
+                              e.func(e.lhs, v)])
         else:
             processed.append(e)
 

devito/ir/clusters/cluster.py

@@ -228,7 +228,7 @@ def is_dense(self):
                 not self.is_halo_touch and
                 all(a.is_regular for a in self.scope.accesses))
 
-    @property
+    @cached_property
     def is_sparse(self):
         """
         A cluster is sparse if it represent a sparse operation i.e if both

devito/ir/equations/equation.py

@@ -76,8 +76,11 @@ def apply(self, func):
     def __repr__(self):
         if not self.is_Reduction:
             return super().__repr__()
+        elif self.operation is OpInc:
+            return '%s += %s' % (self.lhs, self.rhs)
         else:
             return '%s = %s(%s, %s)' % (self.lhs, self.operation, self.lhs, self.rhs)
+
     # Pickling support
     __reduce_ex__ = Pickable.__reduce_ex__
 

devito/ir/support/properties.py

@@ -138,6 +138,9 @@ class Properties(frozendict):
     """
     A mapper {Dimension -> {properties}}.
     """
+    @property
+    def dimensions(self):
+        return tuple(self.keys())
 
     def add(self, dims, properties=None):
         m = dict(self)
@@ -205,6 +208,9 @@ def is_parallel(self, dims):
         return any(len(self[d] & {PARALLEL, PARALLEL_INDEP}) > 0
                    for d in as_tuple(dims))
 
+    def is_parallel_atomic(self, dims):
+        return any(len(self[d] & {PARALLEL_IF_ATOMIC}) > 0 for d in as_tuple(dims))
+
     def is_parallel_relaxed(self, dims):
         return any(len(self[d] & PARALLELS) > 0 for d in as_tuple(dims))