compiler: Pass along cgroup, not exprs

Author: FabioLuporini

devito/passes/clusters/aliases.py

@@ -7,11 +7,12 @@
 
 from devito.exceptions import CompilationError
 from devito.finite_differences import EvalDerivative, IndexDerivative, Weights
-from devito.ir import (SEQUENTIAL, PARALLEL_IF_PVT, SEPARABLE, Forward,
-                       IterationSpace, Interval, Cluster, ExprGeometry, Queue,
-                       IntervalGroup, LabeledVector, Vector, normalize_properties,
-                       relax_properties, unbounded, minimum, maximum, extrema,
-                       vmax, vmin)
+from devito.ir import (
+    SEQUENTIAL, PARALLEL_IF_PVT, SEPARABLE, Forward, IterationSpace, Interval,
+    Cluster, ClusterGroup, ExprGeometry, Queue, IntervalGroup, LabeledVector,
+    Vector, normalize_properties, relax_properties, unbounded, minimum, maximum,
+    extrema, vmax, vmin
+)
 from devito.passes.clusters.cse import _cse
 from devito.symbolics import (Uxmapper, estimate_cost, search, reuse_if_untouched,
                               retrieve_functions, uxreplace, sympy_dtype)
@@ -117,20 +118,18 @@ def __init__(self, sregistry, options, platform):
         self.opt_min_dtype = options['scalar-min-type']
         self.opt_multisubdomain = True
 
-    def _aliases_from_clusters(self, clusters, exclude, meta):
-        exprs = flatten([c.exprs for c in clusters])
-
+    def _aliases_from_clusters(self, cgroup, exclude, meta):
         # [Clusters]_n -> [Schedule]_m
         variants = []
-        for mapper in self._generate(exprs, exclude):
+        for mapper in self._generate(cgroup, exclude):
             # Clusters -> AliasList
             found = collect(mapper.extracted, meta.ispace, self.opt_minstorage)
-            pexprs, aliases = choose(found, exprs, mapper, self.opt_mingain)
+            exprs, aliases = choose(found, cgroup, mapper, self.opt_mingain)
 
             # AliasList -> Schedule
             schedule = lower_aliases(aliases, meta, self.opt_maxpar)
 
-            variants.append(Variant(schedule, pexprs))
+            variants.append(Variant(schedule, exprs))
 
         if not variants:
             return []
@@ -152,7 +151,7 @@ def _aliases_from_clusters(self, clusters, exclude, meta):
             processed = optimize_clusters_msds(processed)
 
         # [Clusters]_k -> [Clusters]_{k+n}
-        for c in clusters:
+        for c in cgroup:
             n = len(c.exprs)
             cexprs, exprs = exprs[:n], exprs[n:]
 
@@ -170,9 +169,9 @@ def _aliases_from_clusters(self, clusters, exclude, meta):
     def process(self, clusters):
         raise NotImplementedError
 
-    def _generate(self, exprs, exclude):
+    def _generate(self, cgroup, exclude):
         """
-        Generate one or more extractions from ``exprs``. An extraction is a
+        Generate one or more extractions from a ClusterGroup. An extraction is a
         set of CIRE candidates which may be turned into aliases. Two different
         extractions may contain overlapping sub-expressions and, therefore,
         should be processed and evaluated indipendently. An extraction won't
@@ -260,7 +259,7 @@ def callback(self, clusters, prefix, xtracted=None):
             if not g:
                 continue
 
-            made = self._aliases_from_clusters(g, exclude, ak)
+            made = self._aliases_from_clusters(ClusterGroup(g), exclude, ak)
 
             if made:
                 idx = processed.index(g[0])
@@ -285,7 +284,9 @@ def _select(self, variants):
 
 class CireInvariantsElementary(CireInvariants):
 
-    def _generate(self, exprs, exclude):
+    def _generate(self, cgroup, exclude):
+        exprs = cgroup.exprs
+
         # E.g., extract `sin(x)` and `sqrt(x)` from `a*sin(x)*sqrt(x)`
         rule = lambda e: e.is_Function or (e.is_Pow and e.exp.is_Number and 0 < e.exp < 1)
         cbk_search = lambda e: search(e, rule, 'all', 'bfs_first_hit')
@@ -308,7 +309,9 @@ def cbk_search(expr):
 
 class CireInvariantsDivs(CireInvariants):
 
-    def _generate(self, exprs, exclude):
+    def _generate(self, cgroup, exclude):
+        exprs = cgroup.exprs
+
         # E.g., extract `1/h_x`
         rule = lambda e: e.is_Pow and (not e.exp.is_Number or e.exp < 0)
         cbk_search = lambda e: search(e, rule, 'all', 'bfs_first_hit')
@@ -339,13 +342,17 @@ def process(self, clusters):
             # TODO: to process third- and higher-order derivatives, we could
             # extend this by calling `_aliases_from_clusters` repeatedly until
             # `made` is empty. To be investigated
-            made = self._aliases_from_clusters([c], exclude, self._lookup_key(c))
+            made = self._aliases_from_clusters(
+                ClusterGroup(c), exclude, self._lookup_key(c)
+            )
 
             processed.extend(flatten(made) or [c])
 
         return processed
 
-    def _generate(self, exprs, exclude):
+    def _generate(self, cgroup, exclude):
+        exprs = cgroup.exprs
+
         # E.g., extract `u.dx*a*b` and `u.dx*a*c` from
         # `[(u.dx*a*b).dy`, `(u.dx*a*c).dy]`
         basextr = self._do_generate(exprs, exclude, self._cbk_search,
@@ -600,14 +607,15 @@ def collect(extracted, ispace, minstorage):
     return aliases
 
 
-def choose(aliases, exprs, mapper, mingain):
+def choose(aliases, cgroup, mapper, mingain):
     """
     Analyze the detected aliases and, after applying a cost model to rule out
     the aliases with a bad memory/flops trade-off, inject them into the original
     expressions.
     """
-    aliases = AliasList(aliases)
+    exprs = cgroup.exprs
 
+    aliases = AliasList(aliases)
     if not aliases:
         return exprs, aliases