From 02e8d9eee94599ee0057f90fde4347916a0b9ebf Mon Sep 17 00:00:00 2001 From: mloubout Date: Fri, 21 Aug 2026 08:57:43 -0400 Subject: [PATCH] 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. --- devito/ir/clusters/cluster.py | 6 ++--- devito/ir/support/guards.py | 12 ++++++++- devito/passes/clusters/aliases.py | 44 +++++++++++++++++++++++-------- tests/test_dse.py | 39 ++++++++++++++++++++++++++- 4 files changed, 84 insertions(+), 17 deletions(-) diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index f45fa224e6..2919e63642 100644 --- a/devito/ir/clusters/cluster.py +++ b/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): diff --git a/devito/ir/support/guards.py b/devito/ir/support/guards.py index 212a8322aa..8a3842ddfd 100644 --- a/devito/ir/support/guards.py +++ b/devito/ir/support/guards.py @@ -5,7 +5,7 @@ """ 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 @@ -13,6 +13,7 @@ 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`. diff --git a/devito/passes/clusters/aliases.py b/devito/passes/clusters/aliases.py index ec549d89e1..17310e0998 100644 --- a/devito/passes/clusters/aliases.py +++ b/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) diff --git a/tests/test_dse.py b/tests/test_dse.py index 47b623f0cb..2c75f60efc 100644 --- a/tests/test_dse.py +++ b/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))