compiler: Don't extract candidates guarded by a Dimension they don't span

An extraction is scheduled over the Dimensions it spans, but the guard
it inherits from its Cluster is kept as it is.  When the guard reads a
Dimension the extraction itself does not span, the temporary is computed
in a loop nest that does not define it, and the generated code does not
compile:

    for (int x = ...)
      if (sdf(x, y) >= 0)     /* y is not iterated here */
        r0[y] = ...

which is what an expression guarded by a ConditionalDimension over a
Function of all the Dimensions produces -- an immersed boundary
condition, say -- as soon as a sub-expression of it depends on fewer
Dimensions than the guard.  The added test is an MFE of exactly that,
and fails to compile with "'x' undeclared" without this change.

Rule those candidates out in `_do_generate`, next to the extractions
that would break a data dependency, so that they are never generated to
begin with.  They cannot go in `exclude` itself, whose entries reject a
candidate that *contains* them, whereas here a candidate is rejected for
what it *lacks*: excluding `x` above would take the whole `x, y`
extraction with it.

A SEQUENTIAL Dimension is exempt: its loop encloses the extraction's
own, so the guard is evaluated outside of them in any case, and
requiring the extraction to span it would cost the hoisting of a time
invariant guarded by a subsampled time Dimension
(`test_invariants_with_conditional`).

The requirement is carried on the transformer rather than passed down
through `_generate`, which is overridden downstream and whose signature
is therefore not ours to change.

The Dimensions a guard reads are what `Guards` is asked for here, so
give it a `dimensions` property rather than walking its values from the
outside.  `Cluster.guards_dimensions` was already doing that walk by
hand and now defers to it.  The property returns the Dimensions as they
are, not their roots: `guards_dimensions` reaches `expose_tuning_knobs`
through `used_dimensions`, which tests them for `is_Block`, and a root
is never a BlockDimension.

Asking `Properties` whether a Dimension is SEQUENTIAL requires it to
still be one: `CireInvariants._lookup_key` was rebuilding it as a plain
`frozendict`, losing the class and its API, where the other
`_lookup_key` passes it through untouched -- so the two CIRE variants
disagreed on the type.  `Properties` is a `frozendict` itself, so
preserving it there costs nothing.

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))