diff --git a/devito/passes/iet/linearization.py b/devito/passes/iet/linearization.py index e67500f7f1..a622a9d89e 100644 --- a/devito/passes/iet/linearization.py +++ b/devito/passes/iet/linearization.py @@ -9,6 +9,7 @@ ) from devito.passes.iet.engine import iet_pass from devito.passes.iet.parpragma import PragmaIteration +from devito.symbolics import retrieve_indexed, uxreplace from devito.tools import filter_ordered, flatten, prod from devito.types import Array, Bundle, FIndexed, Symbol, Wildcard from devito.types.dense import DiscreteFunction @@ -219,6 +220,12 @@ def linearize_accesses(iet, key0, tracker=None): if v is not None: subs[i] = v + # Uxreplace doesn't recurse into a replacement, so the Indexeds nested in + # another's indices are substituted here, innermost first. E.g., + # `T[0, I[0, 0, 0], 0] -> TL0(0, IL0(0, 0, 0), 0)` + for i in sorted(subs, key=lambda i: len(retrieve_indexed(i.indices, deep=True))): + subs[i] = uxreplace(subs[i], subs) + iet = Uxreplace(subs).visit(iet) # 2) What `iet` *offers* diff --git a/devito/passes/iet/misc.py b/devito/passes/iet/misc.py index 1b404daca1..ccf6b3f432 100644 --- a/devito/passes/iet/misc.py +++ b/devito/passes/iet/misc.py @@ -15,7 +15,8 @@ from devito.passes.iet.engine import iet_pass from devito.passes.iet.languages.C import CPrinter from devito.symbolics import ( - Cast, RoundUp, ValueLimit, evalrel, has_integer_args, limits_mapper + Cast, RoundUp, ValueLimit, evalrel, has_integer_args, limits_mapper, retrieve_indexed, + uxreplace ) from devito.tools import Bunch, as_mapper, as_tuple, filter_ordered, split from devito.types import FIndexed @@ -201,6 +202,10 @@ def _generate_macros_findexeds(iet, sregistry=None, tracker=None, **kwargs): subs[i] = v tracker[i.base] = Bunch(header=header, v=v) + # Nested FIndexeds, e.g. `TL0(0, IL0(0, 0, 0), 0)`, innermost first + for i in sorted(subs, key=lambda i: len(retrieve_indexed(i.indices, deep=True))): + subs[i] = uxreplace(subs[i], subs) + iet = Uxreplace(subs).visit(iet) return iet diff --git a/tests/test_linearize.py b/tests/test_linearize.py index c78880924d..9c6fca5273 100644 --- a/tests/test_linearize.py +++ b/tests/test_linearize.py @@ -3,8 +3,8 @@ import scipy.sparse from devito import ( - Eq, Function, Grid, Inc, MatrixSparseTimeFunction, Operator, SparseTimeFunction, - TimeFunction, configuration, sin, switchconfig + ConditionalDimension, Eq, Function, Grid, Inc, MatrixSparseTimeFunction, Operator, + SparseTimeFunction, TimeFunction, configuration, sin, switchconfig ) from devito.ir import Call, Callable, DummyExpr, Expression, FindNodes, SymbolRegistry from devito.passes import Graph, generate_macros, linearize @@ -737,3 +737,30 @@ def test_bundle_derived_dim_stride(): assert bundle.symbolic_shape[p] is not bundle.symbolic_shape[rp] assert bundle.symbolic_shape[rp] == 2 assert bundle.symbolic_shape.dindex(rp) == 1 + + +@pytest.mark.parametrize('opt', ['noop', ('advanced', {'linearize': True})]) +def test_nested_indexeds_in_guard(opt): + grid = Grid(shape=(4, 4)) + time = grid.time_dim + + T = TimeFunction(name='T', grid=grid, time_order=0, space_order=0) + ind = TimeFunction(name='I', grid=grid, time_order=0, space_order=0, + dtype=np.int32) + f = TimeFunction(name='f', grid=grid, time_order=0, space_order=0, save=4) + + # The guard reads `T` through an indirection + cd = ConditionalDimension(name='cd', parent=time, + condition=T[0, ind[0, 0, 0], 0] < 1.5) + eqns = [Eq(T, time), Eq(f, T + 1, implicit_dims=cd)] + + op = Operator(eqns, opt=opt, name='NestedGuard') + + if opt != 'noop': + assert 'TL0(0, IL0(0, 0, 0), 0) < 1.5' in str(op) + + ind.data[:] = 2 + expected = np.array([1, 2, 0, 0])[:, None, None] + + op.apply(time_M=3) + assert np.all(f.data == expected)