From c7bec04bb5b86ba03aa41639fbaded6e1761f046 Mon Sep 17 00:00:00 2001 From: mloubout Date: Fri, 6 Feb 2026 10:35:34 -0500 Subject: [PATCH 1/4] compiler: fix multi-cond async --- devito/finite_differences/finite_difference.py | 2 ++ devito/ir/equations/equation.py | 8 +++++++- devito/passes/clusters/asynchrony.py | 2 +- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/devito/finite_differences/finite_difference.py b/devito/finite_differences/finite_difference.py index 30199fb3d8..b2aa142e26 100644 --- a/devito/finite_differences/finite_difference.py +++ b/devito/finite_differences/finite_difference.py @@ -157,6 +157,8 @@ def first_derivative(expr, dim, fd_order, **kwargs): def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coefficients, expand, weights=None): + if deriv_order == 0 and not expr.is_Add: + print(expr, dim, fd_order) # Always expand time derivatives to avoid issue with buffering and streaming. # Time derivative are almost always short stencils and won't benefit from # unexpansion in the rare case the derivative is not evaluated for time stepping. diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 8f7d35e155..1c671333bd 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -229,7 +229,13 @@ def __new__(cls, *args, **kwargs): index = d.index if d.condition is not None and d in expr.free_symbols: index = index - relational_min(d.condition, d.parent) - expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor)}) + # If there is a condition we might access on a non-factor + # index and need to make sure we don't overwrite the previous + # index + num = index + d.symbolic_factor - 1 + else: + num = index + expr = uxreplace(expr, {d: IntDiv(num, d.symbolic_factor)}) conditionals = frozendict(conditionals) diff --git a/devito/passes/clusters/asynchrony.py b/devito/passes/clusters/asynchrony.py index e32190ddef..59cdcaaf48 100644 --- a/devito/passes/clusters/asynchrony.py +++ b/devito/passes/clusters/asynchrony.py @@ -240,7 +240,7 @@ def _actions_from_update_memcpy(c, d, clusters, actions, sregistry): fetch = e.rhs.indices[d] fshift = {Forward: 1, Backward: -1}.get(direction, 0) - findex = fetch + fshift if fetch.find(IntDiv) else fetch._subs(pd, pd + fshift) + findex = fetch._subs(pd, pd + fshift) # If fetching into e.g. `ub[t1]` we might need to prefetch into e.g. `ub[t0]` tindex0 = e.lhs.indices[d] From 342be17770a89672a50726e36cf957fe19d23b61 Mon Sep 17 00:00:00 2001 From: mloubout Date: Fri, 6 Feb 2026 12:49:44 -0500 Subject: [PATCH 2/4] compiler: fix multi-cond for multi-layer --- devito/ir/equations/equation.py | 11 ++++------- devito/passes/clusters/asynchrony.py | 2 +- devito/passes/clusters/buffering.py | 4 ++-- devito/symbolics/extended_sympy.py | 6 ++++++ devito/types/relational.py | 17 ++++++++++++++++- 5 files changed, 29 insertions(+), 11 deletions(-) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 1c671333bd..b8755f41cc 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -11,7 +11,7 @@ ) from devito.symbolics import IntDiv, limits_mapper, uxreplace from devito.tools import Pickable, Tag, frozendict -from devito.types import Eq, Inc, ReduceMax, ReduceMin, relational_min +from devito.types import Eq, Inc, ReduceMax, ReduceMin, relational_min, relational_shift __all__ = [ 'ClusterizedEq', @@ -229,13 +229,10 @@ def __new__(cls, *args, **kwargs): index = d.index if d.condition is not None and d in expr.free_symbols: index = index - relational_min(d.condition, d.parent) - # If there is a condition we might access on a non-factor - # index and need to make sure we don't overwrite the previous - # index - num = index + d.symbolic_factor - 1 + shift = relational_shift(d.condition, d.parent) else: - num = index - expr = uxreplace(expr, {d: IntDiv(num, d.symbolic_factor)}) + shift = 0 + expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor) + shift}) conditionals = frozendict(conditionals) diff --git a/devito/passes/clusters/asynchrony.py b/devito/passes/clusters/asynchrony.py index 59cdcaaf48..e32190ddef 100644 --- a/devito/passes/clusters/asynchrony.py +++ b/devito/passes/clusters/asynchrony.py @@ -240,7 +240,7 @@ def _actions_from_update_memcpy(c, d, clusters, actions, sregistry): fetch = e.rhs.indices[d] fshift = {Forward: 1, Backward: -1}.get(direction, 0) - findex = fetch._subs(pd, pd + fshift) + findex = fetch + fshift if fetch.find(IntDiv) else fetch._subs(pd, pd + fshift) # If fetching into e.g. `ub[t1]` we might need to prefetch into e.g. `ub[t0]` tindex0 = e.lhs.indices[d] diff --git a/devito/passes/clusters/buffering.py b/devito/passes/clusters/buffering.py index 2a8c78e7fc..57e0ad591a 100644 --- a/devito/passes/clusters/buffering.py +++ b/devito/passes/clusters/buffering.py @@ -3,7 +3,7 @@ from itertools import chain import numpy as np -from sympy import S +from sympy import S, simplify from devito.exceptions import CompilationError from devito.ir import ( @@ -775,7 +775,7 @@ def infer_buffer_size(f, dim, clusters): slots = [Vector(i) for i in slots] size = int((vmax(*slots) - vmin(*slots) + 1)[0]) - return size + return simplify(size) def offset_from_centre(d, indices): diff --git a/devito/symbolics/extended_sympy.py b/devito/symbolics/extended_sympy.py index 2c352d50d4..2ea8f45c93 100644 --- a/devito/symbolics/extended_sympy.py +++ b/devito/symbolics/extended_sympy.py @@ -17,6 +17,7 @@ ) from devito.types import Symbol from devito.types.basic import Basic +from devito.types.relational import Ge __all__ = ['CondEq', 'CondNe', 'BitwiseNot', 'BitwiseXor', 'BitwiseAnd', # noqa 'LeftShift', 'RightShift', 'IntDiv', 'CallFromPointer', @@ -46,6 +47,11 @@ def canonical(self): def negated(self): return CondNe(*self.args, evaluate=False) + @property + def _as_min(self): + from devito.symbolics.extended_dtypes import INT + return INT(Ge(*self.args)) + class CondNe(sympy.Ne): diff --git a/devito/types/relational.py b/devito/types/relational.py index 731ec29bc7..e841d3db96 100644 --- a/devito/types/relational.py +++ b/devito/types/relational.py @@ -3,7 +3,8 @@ import sympy -__all__ = ['Ge', 'Gt', 'Le', 'Lt', 'Ne', 'relational_max', 'relational_min'] +__all__ = ['Ge', 'Gt', 'Le', 'Lt', 'Ne', 'relational_max', 'relational_min', + 'relational_shift'] class AbstractRel: @@ -291,3 +292,17 @@ def _(expr, s): return expr.gts else: return sympy.S.Infinity + + +def relational_shift(expr, s): + """ + Infer shift incurred by the expression. Generally only + applies when a CondEq is used as it adds a single value. + """ + if not expr.has(s): + return 0 + + try: + return expr._as_min + except (TypeError, AttributeError): + return 0 From 54c5e49e2aa2b7fe4179d5af16288c5fa8d27748 Mon Sep 17 00:00:00 2001 From: mloubout Date: Thu, 12 Feb 2026 09:16:03 -0500 Subject: [PATCH 3/4] api: fix interpolate with complex dtype --- devito/finite_differences/finite_difference.py | 2 -- devito/ir/equations/equation.py | 13 +++++++++++++ devito/passes/clusters/cse.py | 14 +++++++++++++- tests/test_interpolation.py | 15 +++++++++++++++ 4 files changed, 41 insertions(+), 3 deletions(-) diff --git a/devito/finite_differences/finite_difference.py b/devito/finite_differences/finite_difference.py index b2aa142e26..30199fb3d8 100644 --- a/devito/finite_differences/finite_difference.py +++ b/devito/finite_differences/finite_difference.py @@ -157,8 +157,6 @@ def first_derivative(expr, dim, fd_order, **kwargs): def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coefficients, expand, weights=None): - if deriv_order == 0 and not expr.is_Add: - print(expr, dim, fd_order) # Always expand time derivatives to avoid issue with buffering and streaming. # Time derivative are almost always short stencils and won't benefit from # unexpansion in the rare case the derivative is not evaluated for time stepping. diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index b8755f41cc..3c54d72714 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -234,6 +234,19 @@ def __new__(cls, *args, **kwargs): shift = 0 expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor) + shift}) + # Merge conditionals when possible. E.g if we have an implicit_dim + # and there is a dimension with the same parent, we ca merged + # its condition + for d in input_expr.implicit_dims: + if d not in conditionals: + continue + for cd in dict(conditionals): + if cd.parent == d.parent and cd != d: + cond = conditionals.pop(d) + mode = cd.relation and d.relation + conditionals[cd] = mode(cond, conditionals[cd]) + break + conditionals = frozendict(conditionals) # Lower all Differentiable operations into SymPy operations diff --git a/devito/passes/clusters/cse.py b/devito/passes/clusters/cse.py index d4d7f0a8b8..ba2c089ee9 100644 --- a/devito/passes/clusters/cse.py +++ b/devito/passes/clusters/cse.py @@ -36,6 +36,18 @@ def retrieve_ctemps(exprs, mode='all'): return search(exprs, lambda expr: isinstance(expr, CTemp), mode, 'dfs') +def cse_dtype(exprdtype, cdtype): + """ + Return the dtype of a CSE temporary given the dtype of the expression to be + captured and the cluster's dtype. + """ + if np.issubdtype(cdtype, np.complexfloating): + return np.promote_types(exprdtype, cdtype(0).real.__class__).type + else: + # Real cluster, can safely promote to the largest precision + return np.promote_types(exprdtype, cdtype).type + + @cluster_pass def cse(cluster, sregistry=None, options=None, **kwargs): """ @@ -86,7 +98,7 @@ def cse(cluster, sregistry=None, options=None, **kwargs): if cluster.is_fence: return cluster - make_dtype = lambda e: np.promote_types(e.dtype, dtype).type + make_dtype = lambda e: cse_dtype(e.dtype, dtype) make = lambda e: CTemp(name=sregistry.make_name(), dtype=make_dtype(e)) exprs = _cse(cluster, make, min_cost=min_cost, mode=mode) diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 1cde2e13e2..8f59ca0076 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -855,6 +855,21 @@ def test_point_symbol_types(dtype, expected): assert point_symbol.dtype is expected +@pytest.mark.parametrize('dtype', [np.complex64, np.complex128]) +def test_interp_complex(dtype): + grid = Grid((11, 11, 11)) + + sc = SparseFunction(name="sc", grid=grid, npoint=1, dtype=dtype) + sc.coordinates.data[:] = [.5, .5, .5] + + fc = Function(name="fc", grid=grid, npoint=2, dtype=dtype) + fc.data[:] = np.random.randn(*grid.shape) + 1j * np.random.randn(*grid.shape) + opC = Operator([sc.interpolate(expr=fc)], name="OpC") + opC() + + assert np.isclose(sc.data[0], fc.data[5, 5, 5]) + + class SD0(SubDomain): name = 'sd0' From 4fd0296edce2c1b025701e990703bf85f5e3b2a1 Mon Sep 17 00:00:00 2001 From: mloubout Date: Mon, 16 Feb 2026 13:48:57 -0500 Subject: [PATCH 4/4] api: fix buffering with multiple conditions --- devito/ir/equations/equation.py | 18 +++++++------- devito/ir/support/vector.py | 3 +++ devito/types/relational.py | 21 +++++++++++++--- tests/test_buffering.py | 43 +++++++++++++++++++++++++++++---- 4 files changed, 68 insertions(+), 17 deletions(-) diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 3c54d72714..eb5c9e4de2 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -213,7 +213,7 @@ def __new__(cls, *args, **kwargs): relations=ordering.relations, mode='partial') ispace = IterationSpace(intervals, iterators) - # Construct the conditionals and replace the ConditionalDimensions in `expr` + # Construct the conditionals conditionals = {} for d in ordering: if not d.is_Conditional: @@ -225,14 +225,6 @@ def __new__(cls, *args, **kwargs): if d._factor is not None: cond = d.relation(cond, GuardFactor(d)) conditionals[d] = cond - # Replace dimension with index - index = d.index - if d.condition is not None and d in expr.free_symbols: - index = index - relational_min(d.condition, d.parent) - shift = relational_shift(d.condition, d.parent) - else: - shift = 0 - expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor) + shift}) # Merge conditionals when possible. E.g if we have an implicit_dim # and there is a dimension with the same parent, we ca merged @@ -249,6 +241,14 @@ def __new__(cls, *args, **kwargs): conditionals = frozendict(conditionals) + # Replace the ConditionalDimensions in `expr` + for d, cond in conditionals.items(): + # Replace dimension with index + index = d.index + index = index - relational_min(cond, d.parent) + shift = relational_shift(cond, d.parent) + expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor) + shift}) + # Lower all Differentiable operations into SymPy operations rhs = diff2sympy(expr.rhs) diff --git a/devito/ir/support/vector.py b/devito/ir/support/vector.py index 02e26e2a02..79d84605ba 100644 --- a/devito/ir/support/vector.py +++ b/devito/ir/support/vector.py @@ -128,6 +128,7 @@ def __lt__(self, other): return True elif q_positive(i): return False + raise TypeError("Non-comparable index functions") from e return False @@ -164,6 +165,7 @@ def __gt__(self, other): return True elif q_negative(i): return False + raise TypeError("Non-comparable index functions") from e return False @@ -203,6 +205,7 @@ def __le__(self, other): return True elif q_positive(i): return False + raise TypeError("Non-comparable index functions") from e # Note: unlike `__lt__`, if we end up here, then *it is* <=. For example, diff --git a/devito/types/relational.py b/devito/types/relational.py index e841d3db96..1fb3e19355 100644 --- a/devito/types/relational.py +++ b/devito/types/relational.py @@ -302,7 +302,22 @@ def relational_shift(expr, s): if not expr.has(s): return 0 - try: - return expr._as_min - except (TypeError, AttributeError): + return _relational_shift(expr, s) + + +@singledispatch +def _relational_shift(s, expr): + return 0 + + +@_relational_shift.register(sympy.Or) +@_relational_shift.register(sympy.And) +def _(expr, s): + return sum([_relational_shift(e, s) for e in expr.args]) + + +@_relational_shift.register(sympy.Eq) +def _(expr, s): + if isinstance(expr.lhs, sympy.Mod): return 0 + return expr._as_min diff --git a/tests/test_buffering.py b/tests/test_buffering.py index 0a0037c223..f89f7262d4 100644 --- a/tests/test_buffering.py +++ b/tests/test_buffering.py @@ -754,7 +754,7 @@ def test_buffer_reuse(): assert all(np.all(vsave.data[i-1] == i + 1) for i in range(1, nt + 1)) -def test_multi_cond(): +def test_multi_cond_v0(): grid = Grid((3, 3)) nt = 5 @@ -774,14 +774,47 @@ def test_multi_cond(): T = TimeFunction(grid=grid, name='T', time_order=0, space_order=0) eqs = [Eq(T, grid.time_dim)] - # this to save times from 0 to nt - 2 + # This saves + # - All subsampled times since ct1 is the dimension of f + # - The last time step (ntmod - 2) through ctend (since it's set as ct1 or ctend) + eqs.append(Eq(f, T, implicit_dims=ctend)) + + # run operator with buffering + op = Operator(eqs, opt='buffering') + op.apply(time_m=0, time_M=ntmod-2) + + for i in range(nt-1): + assert np.allclose(f.data[i], i*2) + assert np.allclose(f.data[nt-1], ntmod - 2) + + +def test_multi_cond_v1(): + grid = Grid((3, 3)) + nt = 5 + + x, y = grid.dimensions + + factor = 2 + ntmod = (nt - 1) * factor + 1 + + ct1 = ConditionalDimension(name="ct1", parent=grid.time_dim, + factor=factor, relation=Or, + condition=CondEq(grid.time_dim, ntmod - 2)) + + f = TimeFunction(grid=grid, name='f', time_order=0, + space_order=0, save=nt, time_dim=ct1) + T = TimeFunction(grid=grid, name='T', time_order=0, space_order=0) + + eqs = [Eq(T, grid.time_dim)] + # This saves + # - All subsampled times since ct1 is the dimension of f with factor 2 + # - The last time step (ntmod - 2) since ct1 also has the condition for ntmod - 2 eqs.append(Eq(f, T)) - # this to save the last time sample nt - 1 - eqs.append(Eq(f.forward, T+1, implicit_dims=ctend)) # run operator with buffering op = Operator(eqs, opt='buffering') op.apply(time_m=0, time_M=ntmod-2) - for i in range(nt): + for i in range(nt-1): assert np.allclose(f.data[i], i*2) + assert np.allclose(f.data[nt-1], ntmod - 2)