From 4440234df1d4f3a962550dc2b0de39c6ece36959 Mon Sep 17 00:00:00 2001 From: mloubout Date: Fri, 21 Aug 2026 08:57:43 -0400 Subject: [PATCH] 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. --- devito/ir/clusters/cluster.py | 6 ++--- devito/ir/support/guards.py | 11 ++++++++- devito/passes/clusters/aliases.py | 33 ++++++++++++++++++++++---- tests/test_dse.py | 39 ++++++++++++++++++++++++++++++- 4 files changed, 78 insertions(+), 11 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..9c279cb812 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,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`. diff --git a/devito/passes/clusters/aliases.py b/devito/passes/clusters/aliases.py index ec549d89e1..849f71d916 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, 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 = [] 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))