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

An alias 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 alias 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.

Discard those aliases in `collect`, as the pivot is formed and before
its distance vectors and cost are computed.  A SEQUENTIAL Dimension is
exempt: its loop encloses the alias' own, so the guard is evaluated
outside of them in any case, and requiring the alias to span it would
cost the hoisting of a time invariant guarded by a subsampled time
Dimension (`test_invariants_with_conditional`).

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 `collect` saw a
different type depending on which CIRE variant called it.  `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,14 @@ 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 {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, pull_dims, 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,
@@ -343,7 +344,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)
 
@@ -566,7 +567,22 @@ def collect(extracted, meta, minstorage):
         * a[i] + c[i] : because at least one of the operands differs
         * a[i+2] - b[i+2] : because at least one operation differs
         * a[i+2] + b[i] : because the distances along ``i`` differ (+2 and +0)
+
+    An aliasing expression is discarded if it does not span all of the
+    Dimensions its guard reads, since it would then be computed in a loop
+    nest that does not define its guard.
     """
+    # The Dimensions an alias must span for its guard to be evaluated within
+    # the alias' own loops: those the guards read, but for the SEQUENTIAL ones,
+    # along which the guard is evaluated outside of them anyway. Taken at their
+    # root, an alias and a guard not necessarily using the same derived
+    # Dimensions -- the alias' own are pulled with their ancestors below
+    if meta.guards:
+        guard_dims = {d.root for d in meta.guards.dimensions
+                      if not meta.properties.is_sequential(d._defines)}
+    else:
+        guard_dims = set()
+
     # Find the potential aliases
     found = []
     for expr in extracted:
@@ -657,6 +673,13 @@ def collect(extracted, meta, minstorage):
                     for i, b, v in zip(c.indexeds, c.bases, offsets, strict=True)}
             pivot = uxreplace(c.expr, subs)
 
+            # An alias is scheduled over the Dimensions it spans, but it
+            # inherits its guard as it is, so a guard reading a Dimension the
+            # alias does not span would be evaluated in a loop nest that does
+            # not define it
+            if not guard_dims <= pull_dims(pivot):
+                continue
+
             # Distance of each aliased expression from the basis alias
             aliaseds = []
             distances = []

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