diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index e0281e1139..b28730d6c8 100644 --- a/devito/ir/equations/algorithms.py +++ b/devito/ir/equations/algorithms.py @@ -169,8 +169,8 @@ def concretize_subdims(exprs, **kwargs): across `exprs`, such as the thickness symbols. """ sregistry = kwargs.get('sregistry') + mapper = kwargs.get('concretize_mapper') - mapper = {} rebuilt = {} # Rebuilt implicit dims etc which are shared between dimensions _concretize_subdims(exprs, mapper, rebuilt, sregistry) @@ -228,9 +228,8 @@ def _(d, mapper, rebuilt, sregistry): # Already have a substitution for this dimension return - tkns = tuple(t._rebuild(name=sregistry.make_name(prefix=t.name)) for t in d.tkns) - mapper.update({tkn0: tkn1 for tkn0, tkn1 in zip(d.tkns, tkns)}) - mapper[d] = d._rebuild(thickness=tkns) + _concretize_subdims(d.tkns, mapper, rebuilt, sregistry) + mapper[d] = d._rebuild(thickness=tuple(mapper[tkn] for tkn in d.tkns)) @_concretize_subdims.register(ConditionalDimension) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index ccfa5a249f..a75aeb9e7b 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -262,6 +262,7 @@ def _lower(cls, expressions, **kwargs): """ # Create a symbol registry kwargs.setdefault('sregistry', SymbolRegistry()) + kwargs.setdefault('concretize_mapper', {}) expressions = as_tuple(expressions) diff --git a/tests/test_builtins.py b/tests/test_builtins.py index 102875d35a..483bd6a135 100644 --- a/tests/test_builtins.py +++ b/tests/test_builtins.py @@ -1,7 +1,10 @@ import pytest import numpy as np from scipy.ndimage import gaussian_filter -from scipy.misc import ascent +try: + from scipy.datasets import ascent +except ImportError: + from scipy.misc import ascent from devito import ConditionalDimension, Grid, Function, TimeFunction, switchconfig from devito.builtins import (assign, norm, gaussian_smooth, initialize_function, diff --git a/tests/test_dimension.py b/tests/test_dimension.py index 84df27ecef..2d2ffe9d7c 100644 --- a/tests/test_dimension.py +++ b/tests/test_dimension.py @@ -2050,8 +2050,31 @@ def test_correct_thicknesses(self): ix1 = SubDimension.right('x', x, 2) ix2 = SubDimension.middle('x', x, 2, 2) - rebuilt = concretize_subdims([ix0, ix1, ix2], sregistry=SymbolRegistry()) + rebuilt = concretize_subdims([ix0, ix1, ix2], sregistry=SymbolRegistry(), + concretize_mapper={}) assert rebuilt[0].is_left assert rebuilt[1].is_right assert rebuilt[2].is_middle + + def test_repeat_concretization(self): + """ + Ensure that SubDimensions are consistently concretized to the same object + across multiple calls to the function. This is necessary when using + `rcompile` on equations with SubDimensions. + """ + + grid = Grid((2, 2)) + + x = Dimension('x') + ix = SubDimension.middle('ix', x, 2, 2) + + u = Function(name='u', grid=grid) + eq = Eq(u, ix + ix.ltkn + ix.rtkn) + + kwargs = {'sregistry': SymbolRegistry(), 'concretize_mapper': {}} + + exprs1 = concretize_subdims([eq], **kwargs) + exprs2 = concretize_subdims([eq], **kwargs) + + assert exprs1 == exprs2 diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index 784dc1c3ad..db08f60de6 100644 --- a/tests/test_subdomains.py +++ b/tests/test_subdomains.py @@ -37,7 +37,8 @@ def define(self, dimensions): with timed_region('x'): # _lower_exprs expects a SymbolRegistry, so create one expr = Operator._lower_exprs([eq0], options={}, - sregistry=SymbolRegistry())[0] + sregistry=SymbolRegistry(), + concretize_mapper={})[0] assert str(expr.rhs) == 'ix*f[ix + 1, iy + 1] + iy' def test_multiple_middle(self):