compiler: Linearize Indexeds nested in other Indexeds' indices

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*

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

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)