Merge pull request #3010 from devitocodes/cire-guarded-aliases

compiler: Don't collect aliases guarded by a Dimension they don't span

Author: mloubout

devito/ir/clusters/cluster.py

@@ -118,14 +118,12 @@ def exprs_dimensions(self):
         dims_implicit = {d for e in self.exprs for d in e.implicit_dims}
         return dims_explicit | dims_implicit
 
-    @cached_property
+    @property
     def guards_dimensions(self):
         """
         The Dimensions that appear explicitly in the guards.
         """
-        syms_guards = {d for e in self.guards.values() for d in e.free_symbols}
-        dims_guards = {i for i in syms_guards if i.is_Dimension}
-        return dims_guards
+        return self.guards.dimensions
 
     @cached_property
     def used_dimensions(self):

devito/ir/support/guards.py

@@ -5,14 +5,15 @@
 """
 
 from collections import Counter, defaultdict
-from functools import singledispatch
+from functools import cached_property, singledispatch
 from operator import ge, gt, le, lt
 
 import numpy as np
 from sympy import And, Expr, Ge, Gt, Le, Lt, Mul, true
 from sympy.logic.boolalg import BooleanFunction
 
 from devito.ir.support.space import Forward, IterationDirection
+from devito.ir.support.utils import pull_dims
 from devito.symbolics import CondEq, CondNe, IntDiv, search
 from devito.symbolics.manipulation import _uxreplace_handle, _uxreplace_registry
 from devito.tools import Pickable, as_tuple, frozendict, split
@@ -280,6 +281,15 @@ class Guards(frozendict):
     def get(self, d, v=true):
         return super().get(d, v)
 
+    @cached_property
+    def dimensions(self):
+        """
+        The Dimensions the guards read, that is those a guarded object must
+        be evaluated within.
+        """
+        return frozenset({d for v in self.values()
+                          for d in pull_dims(v, flag=False)})
+
     def has(self, d, cls):
         """
         True if the guard registered for `d` contains an instance of `cls`.

devito/passes/clusters/aliases.py

@@ -9,8 +9,9 @@
 from devito.finite_differences import EvalDerivative, IndexDerivative, Weights
 from devito.ir import (
     PARALLEL_IF_PVT, SEPARABLE, SEQUENTIAL, Cluster, ClusterGroup, ExprGeometry, Forward,
-    Interval, IntervalGroup, IterationSpace, LabeledVector, Queue, Vector, extrema,
-    maximum, minimum, normalize_properties, relax_properties, unbounded, vmax, vmin
+    Interval, IntervalGroup, IterationSpace, LabeledVector, Properties, Queue, Vector,
+    extrema, maximum, minimum, normalize_properties, relax_properties, unbounded, vmax,
+    vmin
 )
 from devito.passes.clusters.cse import _cse
 from devito.passes.clusters.utils import expose_tuning_knobs
@@ -19,8 +20,8 @@
     uxreplace
 )
 from devito.tools import (
-    Reconstructable, Stamp, as_mapper, as_tuple, flatten, frozendict, generator,
-    is_integer, split, timed_pass
+    Reconstructable, Stamp, as_mapper, as_tuple, flatten, generator, is_integer, split,
+    timed_pass
 )
 from devito.types import (
     CustomDimension, Eq, Hyperplane, IncrDimension, Indexed, ModuloDimension, Size,
@@ -288,6 +289,7 @@ def _do_generate(self, exprs, exclude, cbk_search, cbk_compose=None):
                     free_symbols = i.free_symbols
                 if {a.function for a in free_symbols} & exclude:
                     continue
+
                 mapper.add(i, make, terms)
 
         return mapper
@@ -304,6 +306,31 @@ def __init__(self, sregistry, options, platform):
     def process(self, clusters):
         return self._process_fatd(clusters, 1, xtracted=[])
 
+    @classmethod
+    def _make_exclude(cls, clusters, d, p):
+        """
+        The symbols an extraction must not touch.
+        """
+        # Rule out extractions that would break data dependencies
+        exclude = set().union(*[c.scope.writes for c in clusters])
+
+        # Rule out extractions that depend on the Dimension currently investigated,
+        # as they clearly wouldn't be invariants
+        exclude.update({d, *p.sub_iterators})
+
+        # An extraction is hoisted out of `d`, but it inherits its guard as it
+        # is, so it must not be hoisted past any Dimension the guard reads, or
+        # the guard would be evaluated where it is not defined. Excluding those
+        # Dimensions keeps the extraction inside the loops defining them. A
+        # SEQUENTIAL Dimension is exempt, its loop enclosing the extraction's own
+        for c in clusters:
+            exclude.update(
+                i for i in c.guards.dimensions
+                if not c.properties.is_sequential(i._defines)
+            )
+
+        return exclude
+
     def callback(self, clusters, prefix, xtracted=None):
         if not prefix:
             return clusters
@@ -314,12 +341,7 @@ def callback(self, clusters, prefix, xtracted=None):
         if d.is_Virtual:
             return clusters
 
-        # Rule out extractions that would break data dependencies
-        exclude = set().union(*[c.scope.writes for c in clusters])
-
-        # Rule out extractions that depend on the Dimension currently investigated,
-        # as they clearly wouldn't be invariants
-        exclude.update({d, *p.sub_iterators})
+        exclude = self._make_exclude(clusters, d, p)
 
         key = lambda c: self._lookup_key(c, d)
         processed = list(clusters)
@@ -343,7 +365,7 @@ def callback(self, clusters, prefix, xtracted=None):
     def _lookup_key(self, c, d):
         ispace = c.ispace.reset()
         intervals = c.ispace.intervals.drop(d).reset()
-        properties = frozendict({d: relax_properties(v) for d, v in c.properties.items()})
+        properties = Properties({d: relax_properties(v) for d, v in c.properties.items()})
 
         return AliasKey(ispace, intervals, c.dtype, c.guards, properties)
 

tests/test_dse.py

@@ -10,7 +10,7 @@
 )
 from devito import (  # noqa
     NODE, Abs, ConditionalDimension, Constant, DefaultDimension, Derivative, Dimension,
-    Eq, Function, Ge, Grid, Inc, Lt, Operator, SparseTimeFunction, SubDimension,
+    Eq, Function, Ge, Grid, Inc, Lt, Max, Operator, SparseTimeFunction, SubDimension,
     TimeFunction, configuration, cos, dimensions, div, exp, first_derivative, floor, grad,
     norm, sin, solve, sqrt, switchconfig, transpose
 )
@@ -2512,6 +2512,43 @@ def test_contraction_with_conditional(self):
         assert len(FindNodes(Conditional).visit(op)) == 1
         assert np.all(u.data[6:] == 1.42)
 
+    def test_no_extraction_guarded_by_unspanned_dimension(self):
+        """
+        An alias is scheduled over the Dimensions it spans, but it inherits its
+        guard as it is. A guard reading a Dimension the alias does not span
+        would then be evaluated in a loop nest that does not define it, giving
+        code that does not compile.
+
+        Here `sin(...)` depends on `y` alone while the condition reads both `x`
+        and `y`, which is what an immersed boundary condition produces.
+        """
+        grid = Grid(shape=(16, 16))
+        x, y = grid.dimensions
+
+        sdf = Function(name='sdf', grid=grid)
+        sdf.data[:] = 1.
+
+        cond = ConditionalDimension(name='inside', parent=y, condition=Ge(sdf, 0))
+
+        u = TimeFunction(name='u', grid=grid, space_order=4)
+        u.data[:] = 1.
+
+        # Expensive enough to be extracted, and `y`-only
+        prof = sin(Max(0., 1. - y)) + sin(Max(0., 1. + y))
+
+        eqn = Eq(u.forward, u.laplace + prof*u, implicit_dims=[cond])
+
+        op = Operator(eqn, opt=('advanced', {'cire-mingain': 0, 'openmp': False}))
+
+        # No temporary may be created over fewer Dimensions than the guard reads
+        for i in FindSymbols().visit(op):
+            if i.is_Array:
+                assert {x, y}.issubset(set(i.dimensions))
+
+        # Used to fail to compile with "'x' undeclared"
+        op.apply(time_M=2)
+        assert np.all(np.isfinite(u.data[:]))
+
     def test_collection_from_conditional(self):
         nt = 10
         grid = Grid(shape=(10, 10))